Pass Field information back and forth when using scalar UDFs

Author: timsaucer

python/datafusion/user_defined.py

@@ -34,7 +34,7 @@
     from _typeshed import CapsuleType as _PyCapsule
 
     _R = TypeVar("_R", bound=pa.DataType)
-    from collections.abc import Callable
+    from collections.abc import Callable, Sequence
 
 
 class Volatility(Enum):
@@ -81,6 +81,27 @@ def __str__(self) -> str:
         return self.name.lower()
 
 
+def data_type_or_field_to_field(value: pa.DataType | pa.Field, name: str) -> pa.Field:
+    """Helper function to return a Field from either a Field or DataType."""
+    if isinstance(value, pa.Field):
+        return value
+    return pa.field(name, type=value)
+
+
+def data_types_or_fields_to_field_list(
+    inputs: Sequence[pa.Field | pa.DataType] | pa.Field | pa.DataType,
+) -> list[pa.Field]:
+    """Helper function to return a list of Fields."""
+    if isinstance(inputs, pa.DataType):
+        return [pa.field("value", type=inputs)]
+    if isinstance(inputs, pa.Field):
+        return [inputs]
+
+    return [
+        data_type_or_field_to_field(v, f"value_{idx}") for (idx, v) in enumerate(inputs)
+    ]
+
+
 class ScalarUDFExportable(Protocol):
     """Type hint for object that has __datafusion_scalar_udf__ PyCapsule."""
 
@@ -103,7 +124,7 @@ def __init__(
         self,
         name: str,
         func: Callable[..., _R],
-        input_types: pa.DataType | list[pa.DataType],
+        input_types: list[pa.Field],
         return_type: _R,
         volatility: Volatility | str,
     ) -> None:
@@ -136,8 +157,8 @@ def __call__(self, *args: Expr) -> Expr:
     @overload
     @staticmethod
     def udf(
-        input_types: list[pa.DataType],
-        return_type: _R,
+        input_types: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field,
+        return_type: pa.DataType | pa.Field,
         volatility: Volatility | str,
         name: str | None = None,
     ) -> Callable[..., ScalarUDF]: ...
@@ -146,8 +167,8 @@ def udf(
     @staticmethod
     def udf(
         func: Callable[..., _R],
-        input_types: list[pa.DataType],
-        return_type: _R,
+        input_types: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field,
+        return_type: pa.DataType | pa.Field,
         volatility: Volatility | str,
         name: str | None = None,
     ) -> ScalarUDF: ...
@@ -200,8 +221,8 @@ def double_udf(x):
 
         def _function(
             func: Callable[..., _R],
-            input_types: list[pa.DataType],
-            return_type: _R,
+            input_types: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field,
+            return_type: pa.DataType | pa.Field,
             volatility: Volatility | str,
             name: str | None = None,
         ) -> ScalarUDF:
@@ -213,6 +234,8 @@ def _function(
                     name = func.__qualname__.lower()
                 else:
                     name = func.__class__.__name__.lower()
+            input_types = data_types_or_fields_to_field_list(input_types)
+            return_type = data_type_or_field_to_field(return_type, "value")
             return ScalarUDF(
                 name=name,
                 func=func,

src/udf.rs

@@ -15,67 +15,159 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::any::Any;
+use std::hash::{Hash, Hasher};
+use std::ptr::addr_of;
 use std::sync::Arc;
 
+use arrow::datatypes::{Field, FieldRef};
+use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema};
 use datafusion::arrow::array::{make_array, Array, ArrayData, ArrayRef};
 use datafusion::arrow::datatypes::DataType;
-use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType, ToPyArrow};
+use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType};
 use datafusion::error::DataFusionError;
-use datafusion::logical_expr::function::ScalarFunctionImplementation;
-use datafusion::logical_expr::{create_udf, ColumnarValue, ScalarUDF, ScalarUDFImpl};
+use datafusion::logical_expr::{
+    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
+    Volatility,
+};
 use datafusion_ffi::udf::FFI_ScalarUDF;
+use pyo3::ffi::Py_uintptr_t;
 use pyo3::prelude::*;
 use pyo3::types::{PyCapsule, PyTuple};
 
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionResult};
 use crate::expr::PyExpr;
 use crate::utils::{parse_volatility, validate_pycapsule};
 
-/// Create a Rust callable function from a python function that expects pyarrow arrays
-fn pyarrow_function_to_rust(
+/// This struct holds the Python written function that is a
+/// ScalarUDF.
+#[derive(Debug)]
+struct PythonFunctionScalarUDF {
+    name: String,
     func: Py<PyAny>,
-) -> impl Fn(&[ArrayRef]) -> Result<ArrayRef, DataFusionError> {
-    move |args: &[ArrayRef]| -> Result<ArrayRef, DataFusionError> {
+    signature: Signature,
+    return_field: FieldRef,
+}
+
+impl PythonFunctionScalarUDF {
+    fn new(
+        name: String,
+        func: Py<PyAny>,
+        input_fields: Vec<Field>,
+        return_field: Field,
+        volatility: Volatility,
+    ) -> Self {
+        let input_types = input_fields.iter().map(|f| f.data_type().clone()).collect();
+        let signature = Signature::exact(input_types, volatility);
+        Self {
+            name,
+            func,
+            signature,
+            return_field: Arc::new(return_field),
+        }
+    }
+}
+
+impl Eq for PythonFunctionScalarUDF {}
+impl PartialEq for PythonFunctionScalarUDF {
+    fn eq(&self, other: &Self) -> bool {
+        self.name == other.name
+            && self.signature == other.signature
+            && self.return_field == other.return_field
+            && Python::attach(|py| self.func.bind(py).eq(other.func.bind(py)).unwrap_or(false))
+    }
+}
+
+impl Hash for PythonFunctionScalarUDF {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+        self.signature.hash(state);
+        self.return_field.hash(state);
+
+        Python::attach(|py| {
+            let py_hash = self.func.bind(py).hash().unwrap_or(0); // Handle unhashable objects
+
+            state.write_isize(py_hash);
+        });
+    }
+}
+
+fn array_to_pyarrow_with_field(
+    py: Python,
+    array: ArrayRef,
+    field: &FieldRef,
+) -> PyResult<Py<PyAny>> {
+    let array = FFI_ArrowArray::new(&array.to_data());
+    let schema = FFI_ArrowSchema::try_from(field).map_err(py_datafusion_err)?;
+
+    let module = py.import("pyarrow")?;
+    let class = module.getattr("Array")?;
+    let array = class.call_method1(
+        "_import_from_c",
+        (
+            addr_of!(array) as Py_uintptr_t,
+            addr_of!(schema) as Py_uintptr_t,
+        ),
+    )?;
+    Ok(array.unbind())
+}
+
+impl ScalarUDFImpl for PythonFunctionScalarUDF {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        &self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> datafusion::common::Result<DataType> {
+        unimplemented!()
+    }
+
+    fn return_field_from_args(
+        &self,
+        _args: ReturnFieldArgs,
+    ) -> datafusion::common::Result<FieldRef> {
+        Ok(Arc::clone(&self.return_field))
+    }
+
+    fn invoke_with_args(
+        &self,
+        args: ScalarFunctionArgs,
+    ) -> datafusion::common::Result<ColumnarValue> {
+        let num_rows = args.number_rows;
         Python::attach(|py| {
             // 1. cast args to Pyarrow arrays
             let py_args = args
-                .iter()
-                .map(|arg| {
-                    arg.into_data()
-                        .to_pyarrow(py)
-                        .map_err(|e| DataFusionError::Execution(format!("{e:?}")))
+                .args
+                .into_iter()
+                .zip(args.arg_fields)
+                .map(|(arg, field)| {
+                    let array = arg.to_array(num_rows)?;
+                    array_to_pyarrow_with_field(py, array, &field).map_err(to_datafusion_err)
                 })
                 .collect::<Result<Vec<_>, _>>()?;
             let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?;
 
             // 2. call function
-            let value = func
+            let value = self
+                .func
                 .call(py, py_args, None)
                 .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
 
             // 3. cast to arrow::array::Array
             let array_data = ArrayData::from_pyarrow_bound(value.bind(py))
                 .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
-            Ok(make_array(array_data))
+            Ok(ColumnarValue::Array(make_array(array_data)))
         })
     }
 }
 
-/// Create a DataFusion's UDF implementation from a python function
-/// that expects pyarrow arrays. This is more efficient as it performs
-/// a zero-copy of the contents.
-fn to_scalar_function_impl(func: Py<PyAny>) -> ScalarFunctionImplementation {
-    // Make the python function callable from rust
-    let pyarrow_func = pyarrow_function_to_rust(func);
-
-    // Convert input/output from datafusion ColumnarValue to arrow arrays
-    Arc::new(move |args: &[ColumnarValue]| {
-        let array_refs = ColumnarValue::values_to_arrays(args)?;
-        let array_result = pyarrow_func(&array_refs)?;
-        Ok(array_result.into())
-    })
-}
-
 /// Represents a PyScalarUDF
 #[pyclass(frozen, name = "ScalarUDF", module = "datafusion", subclass)]
 #[derive(Debug, Clone)]
@@ -88,19 +180,21 @@ impl PyScalarUDF {
     #[new]
     #[pyo3(signature=(name, func, input_types, return_type, volatility))]
     fn new(
-        name: &str,
+        name: String,
         func: Py<PyAny>,
-        input_types: PyArrowType<Vec<DataType>>,
-        return_type: PyArrowType<DataType>,
+        input_types: PyArrowType<Vec<Field>>,
+        return_type: PyArrowType<Field>,
         volatility: &str,
     ) -> PyResult<Self> {
-        let function = create_udf(
+        let py_function = PythonFunctionScalarUDF::new(
             name,
+            func,
             input_types.0,
             return_type.0,
             parse_volatility(volatility)?,
-            to_scalar_function_impl(func),
         );
+        let function = ScalarUDF::new_from_impl(py_function);
+
         Ok(Self { function })
     }