Merge pull request #450 from influxdata/crepererum/numpy

feat: numpy

Author: crepererum

guests/python/Justfile

@@ -331,15 +331,24 @@ link-lib profile: (build-lib profile) download-wasi-sdk download-wasi-adapter
         -L{{PYTHON_SDK_DIR}} \
         -lpython{{PYTHON_VERSION_HALF}}
 
-    # Link everything together.
-    wasm-tools component link \
+    # Link everything together:
+    #
+    # - Python compiled extensions (mapped to the correct place inside the guest file system)
+    # - libdl stub/mock for `dlopen`/`dlsym`
+    # - our Rust-based payload
+    # - CPython
+    # - WASI SDK libc/libc++ and emulated functionality/stubs
+    find "{{PYTHON_SITE_PACKAGES}}" -iname '*.so' | sed -E 's|^({{PYTHON_SDK_DIR}})(/.*)$|--dl-openable=\2=\1\2|' | xargs -- \
+        wasm-tools component link \
         --adapt={{WASI_ADAPTER}} \
         --output=$target/datafusion_udf_wasm_python.full.wasm \
         --use-built-in-libdl \
         --validate=true \
         $target/libdatafusion_udf_wasm_python.so \
         {{PYTHON_SDK_DIR}}/libpython{{PYTHON_VERSION_HALF}}.so \
         {{WASI_SYSROOT}}/lib/wasm32-wasip2/libc.so \
+        {{WASI_SYSROOT}}/lib/wasm32-wasip2/libc++.so \
+        {{WASI_SYSROOT}}/lib/wasm32-wasip2/libc++abi.so \
         {{WASI_SYSROOT}}/lib/wasm32-wasip2/libwasi-emulated-getpid.so \
         {{WASI_SYSROOT}}/lib/wasm32-wasip2/libwasi-emulated-process-clocks.so \
         {{WASI_SYSROOT}}/lib/wasm32-wasip2/libwasi-emulated-signal.so

guests/python/README.md

@@ -73,6 +73,7 @@ Currently we bundle the following libraries:
 
 - [`certifi`] (not really used though, see ["I/O > HTTP"](#http))
 - [`charset-normalizer`]
+- [`numpy`]
 - [`requests`]
 - [`urllib3`]
 
@@ -261,6 +262,7 @@ There is NO other I/O available that escapes the sandbox.
 [`Int64`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#variant.Int64
 [`None`]: https://docs.python.org/3/library/constants.html#None
 [`Null`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#variant.Null
+[`numpy`]: https://numpy.org/
 [`time`]: https://docs.python.org/3/library/datetime.html#datetime.time
 [`Time64`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#variant.Time64
 [`timedelta`]: https://docs.python.org/3/library/datetime.html#datetime.timedelta

guests/python/build.rs

@@ -11,7 +11,10 @@
 use std::{fs::File, io::Write, path::PathBuf};
 
 /// File endings that should be skipped when bundling the up the Python lib.
-const SKIP_ENDINGS: &[&str] = &[".a", ".pyc", ".so", ".wasm"];
+const SKIP_ENDINGS: &[&str] = &[".a", ".pyc", ".wasm"];
+
+/// File endings that are mocked as empty files.
+const MOCK_ENDINGS: &[&str] = &[".so"];
 
 fn main() {
     println!("cargo:rerun-if-changed=build.rs");
@@ -48,7 +51,17 @@ fn bundle_python_lib() {
             continue;
         }
 
-        archive.append_path_with_name(path_abs, path_rel).unwrap();
+        if MOCK_ENDINGS.iter().any(|ending| path_str.ends_with(ending)) {
+            const MOCK: &[u8] = b"";
+
+            let mut header = tar::Header::new_gnu();
+            header.set_size(MOCK.len() as _);
+            header.set_cksum();
+
+            archive.append_data(&mut header, path_rel, MOCK).unwrap();
+        } else {
+            archive.append_path_with_name(path_abs, path_rel).unwrap();
+        }
     }
     archive.finish().unwrap();
     archive.into_inner().unwrap().flush().unwrap();

guests/python/requirements.txt

@@ -5,3 +5,6 @@ requests==2.33.0 --hash=sha256:3324635456fa185245e24865e810cecec7b4caf933d7eb133
 
 # See https://github.com/influxdata/urllib3/pull/1 .
 urllib3 @ https://github.com/influxdata/urllib3/releases/download/2.6.103/urllib3-2.6.103-py3-none-any.whl --hash=sha256:bd06baffa5df72a88f8b0edaeb13cf55b1880bcbcd67c62b391f3ec124add2c3
+
+# our own WASI-native wheels
+numpy @ https://github.com/influxdata/python-wasm-wheels/releases/download/v1/numpy-1.26.0-py2.py3-none-any.whl --hash=sha256:ae15276be55efbeaf1de0c82468940aa799024f87f0b679d28351362ff96c466

host/tests/integration_tests/python/runtime/dependencies.rs

@@ -81,3 +81,36 @@ def foo(x: int) -> int:
         &Int64Array::from_iter([Some(11), Some(22), Some(11)]) as &dyn Array,
     );
 }
+
+#[tokio::test]
+async fn numpy() {
+    const CODE: &str = "
+import numpy as np
+
+def foo(x: int) -> int:
+    x = np.int64(x)
+    x = np.sqrt(x)
+    return int(x)
+";
+
+    let udf = python_scalar_udf(CODE).await.unwrap();
+    let array = udf
+        .invoke_async_with_args(ScalarFunctionArgs {
+            args: vec![ColumnarValue::Array(Arc::new(Int64Array::from_iter([
+                Some(1),
+                Some(4),
+                Some(9),
+            ])))],
+            arg_fields: vec![Arc::new(Field::new("a1", DataType::Int64, true))],
+            number_rows: 3,
+            return_field: Arc::new(Field::new("r", DataType::Int64, true)),
+            config_options: Arc::new(ConfigOptions::default()),
+        })
+        .await
+        .unwrap()
+        .unwrap_array();
+    assert_eq!(
+        array.as_ref(),
+        &Int64Array::from_iter([Some(1), Some(2), Some(3)]) as &dyn Array,
+    );
+}

host/tests/integration_tests/python/runtime/fs.rs

@@ -39,7 +39,7 @@ def listdir(cwd: str | None, dir: str) -> str:
         TestCase {
             cwd: None,
             dir: "/",
-            results: &["LICENSE", "lib"],
+            results: &["LICENSE", "lib", "libpython3.14.so"],
         },
         TestCase {
             cwd: None,
@@ -82,7 +82,7 @@ def listdir(cwd: str | None, dir: str) -> str:
         TestCase {
             cwd: Some("/lib"),
             dir: "..",
-            results: &["LICENSE", "lib"],
+            results: &["LICENSE", "lib", "libpython3.14.so"],
         },
     ];
 

host/tests/integration_tests/python/test_utils.rs

@@ -8,8 +8,8 @@ use crate::integration_tests::test_utils::FullError;
 
 /// Memory limit in bytes.
 ///
-/// 100MB.
-const MEMORY_LIMIT: usize = 100_000_000;
+/// 500MB.
+const MEMORY_LIMIT: usize = 500_000_000;
 
 /// Static precompiled Python WASM component for tests
 static COMPONENT: OnceCell<WasmComponentPrecompiled> = OnceCell::const_new();