Add integration test with scalar udf

Author: timsaucer

examples/ffi-library/Cargo.lock

@@ -24,13 +24,18 @@ edition = "2021"
 datafusion = { version = "45.0.0" }
 datafusion-ffi = { version = "45.0.0" }
 pyo3 = { version = "0.23", features = ["extension-module", "abi3", "abi3-py38"] }
-arrow = { version = "54" }
-arrow-array = { version = "54" }
-arrow-schema = { version = "54" }
+arrow = { version = "54.2.0" }
+arrow-array = { version = "54.2.0" }
+arrow-schema = { version = "54.2.0" }
 
 [build-dependencies]
 pyo3-build-config = "0.23"
 
 [lib]
 name = "datafusion_ffi_library"
 crate-type = ["cdylib", "rlib"]
+
+# TODO remove once we update datafusion versions to 46
+[patch.crates-io]
+datafusion = { git = "https://github.com/apache/datafusion.git", rev = "8ab0661a39bd69783b31b949e7a768fb518629e7", features = ["avro", "unicode_expressions"] }
+datafusion-ffi = { git = "https://github.com/apache/datafusion.git", rev = "8ab0661a39bd69783b31b949e7a768fb518629e7" }

examples/ffi-library/Cargo.toml

@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pyarrow as pa
+from datafusion import SessionContext, col, ffi_udf
+from datafusion_ffi_library import IsEvenFunction
+
+
+def test_table_loading():
+    ctx = SessionContext()
+    df = ctx.from_pydict({"a": [-3, -2, None, 0, 1, 2]})
+
+    is_even = ffi_udf(IsEvenFunction())
+
+    result = df.select(is_even(col("a"))).collect()
+    df.with_column("is_even", is_even(col("a"))).show()
+    print(result)
+
+    assert len(result) == 1
+    assert result[0].num_columns == 1
+
+    result = [r.column(0) for r in result]
+    expected = [
+        pa.array([False, True, None, None, False, True], type=pa.bool_()),
+    ]
+
+    assert result == expected

examples/ffi-library/python/tests/_test_scalar_udf.py

@@ -16,11 +16,13 @@
 // under the License.
 
 use pyo3::prelude::*;
-use table_provider::MyTableProvider;
+mod scalar_udf;
 mod table_provider;
 
 #[pymodule]
 fn datafusion_ffi_library(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<MyTableProvider>()?;
+    m.add_class::<table_provider::MyTableProvider>()?;
+    m.add_class::<scalar_udf::IsEvenFunction>()?;
+
     Ok(())
 }

examples/ffi-library/src/lib.rs

@@ -0,0 +1,70 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::{ffi::CString, sync::Arc};
+
+use arrow::array::BooleanArray;
+use arrow_array::ArrayRef;
+use datafusion::common::cast::as_int64_array;
+use datafusion::logical_expr::create_udf;
+use datafusion::logical_expr::Volatility;
+use datafusion::physical_plan::ColumnarValue;
+use datafusion::{arrow::datatypes::DataType, error::Result};
+use datafusion_ffi::udf::FFI_ScalarUDF;
+use pyo3::{prelude::*, types::PyCapsule};
+
+#[pyclass(name = "IsEvenFunction", module = "datafusion_ffi_library", subclass)]
+#[derive(Clone)]
+pub struct IsEvenFunction {}
+
+fn is_even(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+    assert_eq!(args.len(), 1);
+    let args = ColumnarValue::values_to_arrays(args)?;
+
+    let values = as_int64_array(&args[0]).expect("cast failed");
+
+    let array = values
+        .iter()
+        .map(|value| value.and_then(|v| if v == 0 { None } else { Some(v % 2 == 0) }))
+        .collect::<BooleanArray>();
+
+    Ok(ColumnarValue::from(Arc::new(array) as ArrayRef))
+}
+
+#[pymethods]
+impl IsEvenFunction {
+    #[new]
+    fn new() -> Self {
+        Self {}
+    }
+
+    fn __datafusion_scalar_udf__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = CString::new("datafusion_scalar_udf").unwrap();
+
+        let func = create_udf(
+            "is_even",
+            vec![DataType::Int64],
+            DataType::Boolean,
+            Volatility::Immutable,
+            Arc::new(is_even),
+        );
+
+        let ffi_func: FFI_ScalarUDF = (Arc::new(func)).try_into()?;
+
+        PyCapsule::new(py, ffi_func, Some(name))
+    }
+}

examples/ffi-library/src/scalar_udf.rs

@@ -31,7 +31,7 @@ use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyCapsule};
 
 /// In order to provide a test that demonstrates different sized record batches,
 /// the first batch will have num_rows, the second batch num_rows+1, and so on.
-#[pyclass(name = "MyTableProvider", module = "ffi_table_provider", subclass)]
+#[pyclass(name = "MyTableProvider", module = "datafusion_ffi_library", subclass)]
 #[derive(Clone)]
 pub struct MyTableProvider {
     num_cols: usize,

examples/ffi-library/src/table_provider.rs

@@ -118,6 +118,7 @@ def lit(value):
 
 
 udf = ScalarUDF.udf
+ffi_udf = ScalarUDF.ffi_udf
 
 udaf = AggregateUDF.udaf
 

python/datafusion/__init__.py

@@ -123,9 +123,9 @@ def __call__(self, *args: Expr) -> Expr:
         return Expr(self._udf.__call__(*args_raw))
 
     @staticmethod
-    def from_ffi(func: ScalarUDFExportable) -> ScalarUDF:
+    def ffi_udf(func: ScalarUDFExportable) -> ScalarUDF:
         """Create a User-Defined Function from a provided PyCapsule."""
-        udf = df_internal.ScalarUDF.from_ffi(func)
+        udf = df_internal.ScalarUDF.ffi_udf(func)
         return ScalarUDF(None, udf, None, None, None)
 
     @staticmethod

python/datafusion/udf.py

@@ -108,7 +108,7 @@ impl PyScalarUDF {
     }
 
     #[staticmethod]
-    fn from_ffi(func: Bound<PyAny>) -> PyResult<Self> {
+    fn ffi_udf(func: Bound<PyAny>) -> PyResult<Self> {
         if func.hasattr("__datafusion_scalar_udf__")? {
             let capsule = func.getattr("__datafusion_scalar_udf__")?.call0()?;
             let capsule = capsule.downcast::<PyCapsule>()?;
@@ -122,7 +122,7 @@ impl PyScalarUDF {
             })
         } else {
             Err(py_datafusion_err(
-                "__datafusion_table_provider__ does not exist on Table Provider object.",
+                "__datafusion_scalar_udf__ does not exist on Scalar UDF object.",
             ))
         }
     }