diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml index 40c212f..7e6c86e 100644 --- a/.github/workflows/bench.yml +++ b/.github/workflows/bench.yml @@ -3,9 +3,6 @@ on: push: branches: - main - pull_request: - branches: - - main permissions: contents: write diff --git a/Cargo.toml b/Cargo.toml index 9b6e735..b94c7fa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,15 +11,15 @@ exclude = ["arrow-udf-duckdb-example"] [workspace.dependencies] anyhow = "1" -arrow-arith = "54" -arrow-array = "54" -arrow-buffer = "54" -arrow-cast = "54" -arrow-schema = "54" -arrow-select = "54" -arrow-ipc = "54" -arrow-data = "54" -arrow-flight = "54" +arrow-arith = "58.1.0" +arrow-array = "58.1.0" +arrow-buffer = "58.1.0" +arrow-cast = "58.1.0" +arrow-schema = "58.1.0" +arrow-select = "58.1.0" +arrow-ipc = "58.1.0" +arrow-data = "58.1.0" +arrow-flight = "58.1.0" expect-test = "1" serde_json = "1" tokio = "1" diff --git a/arrow-udf-duckdb-example/Makefile b/arrow-udf-duckdb-example/Makefile index 5a84d37..4943f5d 100644 --- a/arrow-udf-duckdb-example/Makefile +++ b/arrow-udf-duckdb-example/Makefile @@ -12,7 +12,13 @@ all: configure debug include extension-ci-tools/makefiles/c_api_extensions/base.Makefile include extension-ci-tools/makefiles/c_api_extensions/rust.Makefile +TEST_RUNNER=$(PYTHON_VENV_BIN) $(PROJ_DIR)run_sqllogictest.py +TEST_RUNNER_BASE=$(TEST_RUNNER) --duckdb-root-dir $(PROJ_DIR) --test-dir test/sql $(EXTRA_EXTENSIONS_PARAM) +TEST_RUNNER_DEBUG=$(TEST_RUNNER_BASE) --build-dir build/debug +TEST_RUNNER_RELEASE=$(TEST_RUNNER_BASE) --build-dir build/release + configure: venv platform extension_version +configure: install_test_runner debug: build_extension_library_debug build_extension_with_metadata_debug release: build_extension_library_release build_extension_with_metadata_release @@ -21,5 +27,8 @@ test: test_debug test_debug: test_extension_debug test_release: test_extension_release +install_test_runner: venv + $(PYTHON_VENV_BIN) -m pip install pytest + clean: clean_build clean_rust clean_all: clean_configure clean diff --git a/arrow-udf-duckdb-example/run_sqllogictest.py b/arrow-udf-duckdb-example/run_sqllogictest.py new file mode 100644 index 0000000..9b942cf --- /dev/null +++ b/arrow-udf-duckdb-example/run_sqllogictest.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 + +import sys +from pathlib import Path + +import pytest +import sqllogic.test_sqllogic as sqllogic_runner + + +def main() -> int: + runner_path = Path(sqllogic_runner.__file__).resolve() + return pytest.main( + ["--noconftest", "-p", "sqllogic.conftest", str(runner_path), *sys.argv[1:]], + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/arrow-udf-runtime/Cargo.toml b/arrow-udf-runtime/Cargo.toml index 2cf23dc..d47fa70 100644 --- a/arrow-udf-runtime/Cargo.toml +++ b/arrow-udf-runtime/Cargo.toml @@ -42,7 +42,7 @@ wasmtime = { version = "27", optional = true } genawaiter2 = { version = "0.100.1", optional = true } tempfile = { version = "3", optional = true } -pyo3 = { version = "0.24.1", optional = true, features = ["auto-initialize"] } +pyo3 = { version = "0.28", optional = true, features = ["auto-initialize"] } atomic-time = { version = "0.1", optional = true } rquickjs = { version = "0.6", features = [ @@ -56,11 +56,11 @@ reqwest = { version = "0.12", features = ["json"], optional = true } serde_json = { version = "1", optional = true } arrow-flight = { workspace = true, optional = true } -tonic = { version = "0.12", optional = true } +tonic = { version = "0.14", optional = true } tracing = { version = "0.1", optional = true } [build-dependencies] -pyo3-build-config = { version = "0.24", features = ["resolve-config"] } +pyo3-build-config = { version = "0.28", features = ["resolve-config"] } [dev-dependencies] arrow-cast = { workspace = true, features = ["prettyprint"] } diff --git a/arrow-udf-runtime/src/javascript/mod.rs b/arrow-udf-runtime/src/javascript/mod.rs index f87c18a..48d55eb 100644 --- a/arrow-udf-runtime/src/javascript/mod.rs +++ b/arrow-udf-runtime/src/javascript/mod.rs @@ -568,18 +568,17 @@ impl Runtime { CallMode::ReturnNullOnNullInput => { // This is a bit tricky. We build input arrays without nulls, call user_fn on them, // and then add back null results to form the final result. - let n_cols = input.num_columns(); let n_rows = input.num_rows(); // 1. Build a bitmap of which rows have nulls let mut bitmap = Vec::with_capacity(n_rows); - for i in 0..n_rows { - let has_null = (0..n_cols).any(|j| js_columns[j][i].is_null()); + for row_idx in 0..n_rows { + let has_null = js_columns.iter().any(|column| column[row_idx].is_null()); bitmap.push(!has_null); } // 2. Build new inputs with only the rows that don't have nulls - let mut filtered_columns = Vec::with_capacity(n_cols); + let mut filtered_columns = Vec::with_capacity(js_columns.len()); for js_values in js_columns { let filtered_js_values: Vec<_> = js_values .into_iter() diff --git a/arrow-udf-runtime/src/python/interpreter.rs b/arrow-udf-runtime/src/python/interpreter.rs index 6fbaa7f..f493ed2 100644 --- a/arrow-udf-runtime/src/python/interpreter.rs +++ b/arrow-udf-runtime/src/python/interpreter.rs @@ -42,7 +42,7 @@ impl Interpreter { // XXX: import the `decimal` module in the interpreter before calling anything else // otherwise it will cause `SIGABRT: pointer being freed was not allocated` // when importing decimal in the second sub-interpreter. - Python::with_gil(|py| { + Python::attach(|py| { py.import("decimal").unwrap(); }); }); @@ -59,7 +59,7 @@ impl Interpreter { where F: for<'py> FnOnce(Python<'py>) -> Result, { - Python::with_gil(f) + Python::attach(f) } /// Run Python code in the sub-interpreter. diff --git a/arrow-udf-runtime/src/python/mod.rs b/arrow-udf-runtime/src/python/mod.rs index 80b96be..5e2d46d 100644 --- a/arrow-udf-runtime/src/python/mod.rs +++ b/arrow-udf-runtime/src/python/mod.rs @@ -23,12 +23,14 @@ use arrow_array::builder::{ArrayBuilder, Int32Builder, StringBuilder}; use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch}; use arrow_schema::{DataType, Field, FieldRef, Schema, SchemaRef}; use pyo3::types::{PyAnyMethods, PyIterator, PyModule, PyTuple}; -use pyo3::{Py, PyObject}; +use pyo3::{Py, PyAny}; use std::collections::HashMap; use std::ffi::CString; use std::fmt::Debug; use std::sync::Arc; +type PyObject = Py; + // #[cfg(Py_3_12)] mod interpreter; mod pyarrow; @@ -99,9 +101,17 @@ struct Aggregate { /// A builder for `Runtime`. #[derive(Default, Debug)] -pub struct Builder {} +pub struct Builder { + safe_codes: Option, +} impl Builder { + /// Run initialization code before user-defined functions are registered. + pub fn safe_codes(mut self, code: String) -> Self { + self.safe_codes = Some(code); + self + } + /// Build the `Runtime`. pub fn build(self) -> Result { let interpreter = Interpreter::new()?; @@ -117,6 +127,9 @@ class Struct: pass "#, )?; + if let Some(code) = self.safe_codes { + interpreter.run(&code)?; + } Ok(Runtime { interpreter, functions: HashMap::new(), diff --git a/arrow-udf-runtime/src/python/pyarrow.rs b/arrow-udf-runtime/src/python/pyarrow.rs index e897c91..ace91e7 100644 --- a/arrow-udf-runtime/src/python/pyarrow.rs +++ b/arrow-udf-runtime/src/python/pyarrow.rs @@ -22,10 +22,12 @@ use pyo3::{ ffi::c_str, prelude::PyDictMethods, types::{PyAnyMethods, PyDict}, - IntoPyObject, PyObject, PyResult, Python, + IntoPyObject, Py, PyAny, PyResult, Python, }; use std::{borrow::Cow, ffi::CString, sync::Arc}; +type PyObject = Py; + macro_rules! get_pyobject { ($array_type: ty, $py:expr, $array:expr, $i:expr) => {{ let array = $array.as_any().downcast_ref::<$array_type>().unwrap(); @@ -403,7 +405,7 @@ impl Converter { for val in values { if !val.is_none(py) { let py_any = val.bind(py); - let dict = py_any.downcast::()?; + let dict = py_any.cast::()?; flatten_keys.reserve(dict.len()); flatten_values.reserve(dict.len()); for key in dict.keys() { diff --git a/arrow-udf-runtime/tests/wasm_build.rs b/arrow-udf-runtime/tests/wasm_build.rs index f4f995d..0981e34 100644 --- a/arrow-udf-runtime/tests/wasm_build.rs +++ b/arrow-udf-runtime/tests/wasm_build.rs @@ -32,7 +32,14 @@ fn gcd(mut a: i32, mut b: i32) -> i32 { #[test] fn test_build_error() { let err = build("??", "").unwrap_err(); - assert!(err.to_string().contains("invalid key")); + let err = err.to_string(); + assert!(err.contains("failed to build wasm")); + assert!(err.contains("Cargo.toml")); + assert!( + err.contains("invalid key") + || err.contains("key with no value") + || err.contains("expected `=`") + ); } fn test_build_offline() { diff --git a/arrow-udf/arrow-udf-macros/src/gen.rs b/arrow-udf/arrow-udf-macros/src/gen.rs index ac3063b..edfc90f 100644 --- a/arrow-udf/arrow-udf-macros/src/gen.rs +++ b/arrow-udf/arrow-udf-macros/src/gen.rs @@ -673,7 +673,10 @@ fn transform_input(input: &Ident, ty: &str) -> TokenStream2 { if ty == "decimal" { return quote! { #input.parse::().expect("invalid decimal") }; } else if ty == "date32" { - return quote! { arrow_array::types::Date32Type::to_naive_date(#input) }; + return quote! { + arrow_array::types::Date32Type::to_naive_date_opt(#input) + .expect("invalid date32 value") + }; } else if ty == "time64" { return quote! { arrow_array::temporal_conversions::as_time::(#input).expect("invalid time") }; } else if ty == "timestamp" {