Add aggregate udf via ffi

Author: timsaucer

examples/datafusion-ffi-example/python/tests/_test_aggregate_udf.py

@@ -0,0 +1,77 @@
+# 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.
+
+from __future__ import annotations
+
+import pyarrow as pa
+from datafusion import SessionContext, col, udaf
+from datafusion_ffi_example import MySumUDF
+
+
+def setup_context_with_table():
+    ctx = SessionContext()
+
+    # Pick numbers here so we get the same value in both groups
+    # since we cannot be certain of the output order of batches
+    batch = pa.RecordBatch.from_arrays(
+        [
+            pa.array([1, 2, 3, None], type=pa.int64()),
+            pa.array([1, 1, 2, 2], type=pa.int64()),
+        ],
+        names=["a", "b"],
+    )
+    ctx.register_record_batches("test_table", [[batch]])
+    return ctx
+
+
+def test_ffi_aggregate_register():
+    ctx = setup_context_with_table()
+    my_udaf = udaf(MySumUDF())
+    ctx.register_udaf(my_udaf)
+
+    result = ctx.sql("select my_custom_sum(a) from test_table group by b").collect()
+
+    assert len(result) == 2
+    assert result[0].num_columns == 1
+
+    result = [r.column(0) for r in result]
+    expected = [
+        pa.array([3], type=pa.int64()),
+        pa.array([3], type=pa.int64()),
+    ]
+
+    assert result == expected
+
+
+def test_ffi_aggregate_call_directly():
+    ctx = setup_context_with_table()
+    my_udaf = udaf(MySumUDF())
+
+    result = (
+        ctx.table("test_table").aggregate([col("b")], [my_udaf(col("a"))]).collect()
+    )
+
+    assert len(result) == 2
+    assert result[0].num_columns == 2
+
+    result = [r.column(1) for r in result]
+    expected = [
+        pa.array([3], type=pa.int64()),
+        pa.array([3], type=pa.int64()),
+    ]
+
+    assert result == expected

examples/datafusion-ffi-example/src/aggregate_udf.rs

@@ -0,0 +1,92 @@
+// 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 arrow_array::{Array, BooleanArray};
+use arrow_schema::{DataType, FieldRef};
+use datafusion::common::ScalarValue;
+use datafusion::error::Result as DataFusionResult;
+use datafusion::functions_aggregate::sum::Sum;
+use datafusion::logical_expr::function::AccumulatorArgs;
+use datafusion::logical_expr::{
+    Accumulator, AggregateUDF, AggregateUDFImpl, ColumnarValue, ScalarFunctionArgs, ScalarUDF,
+    ScalarUDFImpl, Signature, TypeSignature, Volatility,
+};
+use datafusion_ffi::udaf::FFI_AggregateUDF;
+use datafusion_ffi::udf::FFI_ScalarUDF;
+use pyo3::types::PyCapsule;
+use pyo3::{pyclass, pymethods, Bound, PyResult, Python};
+use std::any::Any;
+use std::sync::Arc;
+
+#[pyclass(name = "MySumUDF", module = "datafusion_ffi_example", subclass)]
+#[derive(Debug, Clone)]
+pub(crate) struct MySumUDF {
+    signature: Signature,
+    inner: Arc<Sum>,
+}
+
+#[pymethods]
+impl MySumUDF {
+    #[new]
+    fn new() -> Self {
+        Self {
+            signature: Signature::new(
+                TypeSignature::Exact(vec![DataType::Int64]),
+                Volatility::Immutable,
+            ),
+            inner: Arc::new(Sum::new()),
+        }
+    }
+
+    fn __datafusion_aggregate_udf__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = cr"datafusion_aggregate_udf".into();
+
+        let func = Arc::new(AggregateUDF::from(self.clone()));
+        let provider = FFI_AggregateUDF::from(func);
+
+        PyCapsule::new(py, provider, Some(name))
+    }
+}
+
+impl AggregateUDFImpl for MySumUDF {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "my_custom_sum"
+    }
+
+    fn signature(&self) -> &Signature {
+        self.inner.signature()
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
+        self.inner.return_type(arg_types)
+    }
+
+    fn accumulator(&self, acc_args: AccumulatorArgs) -> DataFusionResult<Box<dyn Accumulator>> {
+        self.inner.accumulator(acc_args)
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> DataFusionResult<Vec<DataType>> {
+        self.inner.coerce_types(arg_types)
+    }
+}

examples/datafusion-ffi-example/src/lib.rs

@@ -16,13 +16,15 @@
 // under the License.
 
 use crate::catalog_provider::MyCatalogProvider;
+use crate::aggregate_udf::MySumUDF;
 use crate::scalar_udf::IsNullUDF;
 use crate::table_function::MyTableFunction;
 use crate::table_provider::MyTableProvider;
 use pyo3::prelude::*;
 
 pub(crate) mod catalog_provider;
-mod scalar_udf;
+pub(crate) mod aggregate_udf;
+pub(crate) mod scalar_udf;
 pub(crate) mod table_function;
 pub(crate) mod table_provider;
 
@@ -32,5 +34,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MyTableFunction>()?;
     m.add_class::<MyCatalogProvider>()?;
     m.add_class::<IsNullUDF>()?;
+    m.add_class::<MySumUDF>()?;
     Ok(())
 }

python/datafusion/user_defined.py

@@ -277,6 +277,12 @@ def evaluate(self) -> pa.Scalar:
         """Return the resultant value."""
 
 
+class AggregateUDFExportable(Protocol):
+    """Type hint for object that has __datafusion_aggregate_udf__ PyCapsule."""
+
+    def __datafusion_aggregate_udf__(self) -> object: ...  # noqa: D105
+
+
 class AggregateUDF:
     """Class for performing scalar user-defined functions (UDF).
 
@@ -298,6 +304,9 @@ def __init__(
         See :py:func:`udaf` for a convenience function and argument
         descriptions.
         """
+        if hasattr(accumulator, "__datafusion_aggregate_udf__"):
+            self._udaf = df_internal.AggregateUDF.from_pycapsule(accumulator)
+            return
         self._udaf = df_internal.AggregateUDF(
             name,
             accumulator,
@@ -342,7 +351,7 @@ def udaf(
     ) -> AggregateUDF: ...
 
     @staticmethod
-    def udaf(*args: Any, **kwargs: Any):  # noqa: D417
+    def udaf(*args: Any, **kwargs: Any):  # noqa: D417, C901
         """Create a new User-Defined Aggregate Function (UDAF).
 
         This class allows you to define an aggregate function that can be used in
@@ -399,6 +408,10 @@ def udf4() -> Summarize:
         Args:
             accum: The accumulator python function. Only needed when calling as a
                 function. Skip this argument when using ``udaf`` as a decorator.
+                If you have a Rust backed AggregateUDF within a PyCapsule, you can
+                pass this parameter and ignore the rest. They will be determined directly
+                from the underlying function. See the online documentation for more
+                information.
             input_types: The data types of the arguments to ``accum``.
             return_type: The data type of the return value.
             state_type: The data types of the intermediate accumulation.
@@ -457,12 +470,32 @@ def wrapper(*args: Any, **kwargs: Any) -> Expr:
 
             return decorator
 
+        if hasattr(args[0], "__datafusion_aggregate_udf__"):
+            return AggregateUDF.from_pycapsule(args[0])
+
         if args and callable(args[0]):
             # Case 1: Used as a function, require the first parameter to be callable
             return _function(*args, **kwargs)
         # Case 2: Used as a decorator with parameters
         return _decorator(*args, **kwargs)
 
+    @staticmethod
+    def from_pycapsule(func: AggregateUDFExportable) -> AggregateUDF:
+        """Create an Aggregate UDF from AggregateUDF PyCapsule object.
+
+        This function will instantiate a Aggregate UDF that uses a DataFusion
+        AggregateUDF that is exported via the FFI bindings.
+        """
+        name = str(udf.__class__)
+        return AggregateUDF(
+            name=name,
+            accumulator=func,
+            input_types=None,
+            return_type=None,
+            state_type=None,
+            volatility=None,
+        )
+
 
 class WindowEvaluator:
     """Evaluator class for user-defined window functions (UDWF).

src/udaf.rs

@@ -19,6 +19,10 @@ use std::sync::Arc;
 
 use pyo3::{prelude::*, types::PyTuple};
 
+use crate::common::data_type::PyScalarValue;
+use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionResult};
+use crate::expr::PyExpr;
+use crate::utils::{parse_volatility, validate_pycapsule};
 use datafusion::arrow::array::{Array, ArrayRef};
 use datafusion::arrow::datatypes::DataType;
 use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow};
@@ -27,11 +31,8 @@ use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_expr::{
     create_udaf, Accumulator, AccumulatorFactoryFunction, AggregateUDF,
 };
-
-use crate::common::data_type::PyScalarValue;
-use crate::errors::to_datafusion_err;
-use crate::expr::PyExpr;
-use crate::utils::parse_volatility;
+use datafusion_ffi::udaf::{FFI_AggregateUDF, ForeignAggregateUDF};
+use pyo3::types::PyCapsule;
 
 #[derive(Debug)]
 struct RustAccumulator {
@@ -183,6 +184,26 @@ impl PyAggregateUDF {
         Ok(Self { function })
     }
 
+    #[staticmethod]
+    pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
+        if func.hasattr("__datafusion_aggregate_udf__")? {
+            let capsule = func.getattr("__datafusion_aggregate_udf__")?.call0()?;
+            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+            validate_pycapsule(capsule, "datafusion_aggregate_udf")?;
+
+            let udaf = unsafe { capsule.reference::<FFI_AggregateUDF>() };
+            let udaf: ForeignAggregateUDF = udaf.try_into()?;
+
+            Ok(Self {
+                function: udaf.into(),
+            })
+        } else {
+            Err(crate::errors::PyDataFusionError::Common(
+                "__datafusion_aggregate_udf__ does not exist on AggregateUDF object.".to_string(),
+            ))
+        }
+    }
+
     /// creates a new PyExpr with the call of the udf
     #[pyo3(signature = (*args))]
     fn __call__(&self, args: Vec<PyExpr>) -> PyResult<PyExpr> {