Initial commit for window ffi integration

Author: timsaucer

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

@@ -0,0 +1,89 @@
+# 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, udwf
+from datafusion_ffi_example import MyRankUDF
+
+
+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([40, 10, 30, 20], type=pa.int64()),
+        ],
+        names=["a"],
+    )
+    ctx.register_record_batches("test_table", [[batch]])
+    return ctx
+
+
+def test_ffi_window_register():
+    ctx = setup_context_with_table()
+    my_udwf = udwf(MyRankUDF())
+    ctx.register_udwf(my_udwf)
+
+    result = ctx.sql(
+        "select a, my_custom_rank() over (order by a) from test_table"
+    ).collect()
+    assert len(result) == 1
+    assert result[0].num_columns == 2
+
+    results = [
+        (result[0][0][idx].as_py(), result[0][1][idx].as_py()) for idx in range(4)
+    ]
+    results.sort()
+
+    expected = [
+        (10, 1),
+        (20, 2),
+        (30, 3),
+        (40, 4),
+    ]
+    assert results == expected
+
+
+def test_ffi_window_call_directly():
+    ctx = setup_context_with_table()
+    my_udwf = udwf(MyRankUDF())
+
+    result = (
+        ctx.table("test_table")
+        .select(col("a"), my_udwf().order_by(col("a")).build())
+        .collect()
+    )
+
+    assert len(result) == 1
+    assert result[0].num_columns == 2
+
+    results = [
+        (result[0][0][idx].as_py(), result[0][1][idx].as_py()) for idx in range(4)
+    ]
+    results.sort()
+
+    expected = [
+        (10, 1),
+        (20, 2),
+        (30, 3),
+        (40, 4),
+    ]
+    assert results == expected

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

@@ -15,18 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow_array::{Array, BooleanArray};
-use arrow_schema::{DataType, FieldRef};
-use datafusion::common::ScalarValue;
+use arrow_schema::DataType;
 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::logical_expr::{Accumulator, AggregateUDF, AggregateUDFImpl, Signature};
 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;
@@ -35,7 +29,6 @@ use std::sync::Arc;
 #[pyclass(name = "MySumUDF", module = "datafusion_ffi_example", subclass)]
 #[derive(Debug, Clone)]
 pub(crate) struct MySumUDF {
-    signature: Signature,
     inner: Arc<Sum>,
 }
 
@@ -44,10 +37,6 @@ impl MySumUDF {
     #[new]
     fn new() -> Self {
         Self {
-            signature: Signature::new(
-                TypeSignature::Exact(vec![DataType::Int64]),
-                Volatility::Immutable,
-            ),
             inner: Arc::new(Sum::new()),
         }
     }

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

@@ -19,18 +19,21 @@ use crate::aggregate_udf::MySumUDF;
 use crate::scalar_udf::IsNullUDF;
 use crate::table_function::MyTableFunction;
 use crate::table_provider::MyTableProvider;
+use crate::window_udf::MyRankUDF;
 use pyo3::prelude::*;
 
 pub(crate) mod aggregate_udf;
 pub(crate) mod scalar_udf;
 pub(crate) mod table_function;
 pub(crate) mod table_provider;
+pub(crate) mod window_udf;
 
 #[pymodule]
 fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MyTableProvider>()?;
     m.add_class::<MyTableFunction>()?;
     m.add_class::<IsNullUDF>()?;
     m.add_class::<MySumUDF>()?;
+    m.add_class::<MyRankUDF>()?;
     Ok(())
 }

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

@@ -16,7 +16,7 @@
 // under the License.
 
 use arrow_array::{Array, BooleanArray};
-use arrow_schema::{DataType, FieldRef};
+use arrow_schema::DataType;
 use datafusion::common::ScalarValue;
 use datafusion::error::Result as DataFusionResult;
 use datafusion::logical_expr::{

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

@@ -0,0 +1,81 @@
+// 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_schema::{DataType, FieldRef};
+use datafusion::error::Result as DataFusionResult;
+use datafusion::functions_window::rank::rank_udwf;
+use datafusion::logical_expr::function::{PartitionEvaluatorArgs, WindowUDFFieldArgs};
+use datafusion::logical_expr::{PartitionEvaluator, Signature, WindowUDF, WindowUDFImpl};
+use datafusion_ffi::udwf::FFI_WindowUDF;
+use pyo3::types::PyCapsule;
+use pyo3::{pyclass, pymethods, Bound, PyResult, Python};
+use std::any::Any;
+use std::sync::Arc;
+
+#[pyclass(name = "MyRankUDF", module = "datafusion_ffi_example", subclass)]
+#[derive(Debug, Clone)]
+pub(crate) struct MyRankUDF {
+    inner: Arc<WindowUDF>,
+}
+
+#[pymethods]
+impl MyRankUDF {
+    #[new]
+    fn new() -> Self {
+        Self { inner: rank_udwf() }
+    }
+
+    fn __datafusion_window_udf__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = cr"datafusion_window_udf".into();
+
+        let func = Arc::new(WindowUDF::from(self.clone()));
+        let provider = FFI_WindowUDF::from(func);
+
+        PyCapsule::new(py, provider, Some(name))
+    }
+}
+
+impl WindowUDFImpl for MyRankUDF {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "my_custom_rank"
+    }
+
+    fn signature(&self) -> &Signature {
+        self.inner.signature()
+    }
+
+    fn partition_evaluator(
+        &self,
+        partition_evaluator_args: PartitionEvaluatorArgs,
+    ) -> DataFusionResult<Box<dyn PartitionEvaluator>> {
+        self.inner
+            .inner()
+            .partition_evaluator(partition_evaluator_args)
+    }
+
+    fn field(&self, field_args: WindowUDFFieldArgs) -> DataFusionResult<FieldRef> {
+        self.inner.inner().field(field_args)
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> DataFusionResult<Vec<DataType>> {
+        self.inner.coerce_types(arg_types)
+    }
+}

python/datafusion/user_defined.py

@@ -247,7 +247,7 @@ def from_pycapsule(func: ScalarUDFExportable) -> ScalarUDF:
         This function will instantiate a Scalar UDF that uses a DataFusion
         ScalarUDF that is exported via the FFI bindings.
         """
-        name = str(udf.__class__)
+        name = str(func.__class__)
         return ScalarUDF(
             name=name,
             func=func,
@@ -409,9 +409,9 @@ def udf4() -> Summarize:
             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.
+                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.
@@ -486,7 +486,7 @@ def from_pycapsule(func: AggregateUDFExportable) -> AggregateUDF:
         This function will instantiate a Aggregate UDF that uses a DataFusion
         AggregateUDF that is exported via the FFI bindings.
         """
-        name = str(udf.__class__)
+        name = str(func.__class__)
         return AggregateUDF(
             name=name,
             accumulator=func,
@@ -656,6 +656,12 @@ def include_rank(self) -> bool:
         return False
 
 
+class WindowUDFExportable(Protocol):
+    """Type hint for object that has __datafusion_window_udf__ PyCapsule."""
+
+    def __datafusion_window_udf__(self) -> object: ...  # noqa: D105
+
+
 class WindowUDF:
     """Class for performing window user-defined functions (UDF).
 
@@ -676,6 +682,9 @@ def __init__(
         See :py:func:`udwf` for a convenience function and argument
         descriptions.
         """
+        if hasattr(func, "__datafusion_window_udf__"):
+            self._udwf = df_internal.WindowUDF.from_pycapsule(func)
+            return
         self._udwf = df_internal.WindowUDF(
             name, func, input_types, return_type, str(volatility)
         )
@@ -751,7 +760,10 @@ def biased_numbers() -> BiasedNumbers:
 
         Args:
             func: Only needed when calling as a function. Skip this argument when
-                using ``udwf`` as a decorator.
+                using ``udwf`` as a decorator. If you have a Rust backed WindowUDF
+                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.
             return_type: The data type of the return value.
             volatility: See :py:class:`Volatility` for allowed values.
@@ -760,6 +772,9 @@ def biased_numbers() -> BiasedNumbers:
         Returns:
             A user-defined window function that can be used in window function calls.
         """
+        if hasattr(args[0], "__datafusion_window_udf__"):
+            return WindowUDF.from_pycapsule(args[0])
+
         if args and callable(args[0]):
             # Case 1: Used as a function, require the first parameter to be callable
             return WindowUDF._create_window_udf(*args, **kwargs)
@@ -827,6 +842,22 @@ def wrapper(*args: Any, **kwargs: Any) -> Expr:
 
         return decorator
 
+    @staticmethod
+    def from_pycapsule(func: WindowUDFExportable) -> WindowUDF:
+        """Create a Window UDF from WindowUDF PyCapsule object.
+
+        This function will instantiate a Window UDF that uses a DataFusion
+        WindowUDF that is exported via the FFI bindings.
+        """
+        name = str(func.__class__)
+        return WindowUDF(
+            name=name,
+            func=func,
+            input_types=None,
+            return_type=None,
+            volatility=None,
+        )
+
 
 class TableFunction:
     """Class for performing user-defined table functions (UDTF).

src/udwf.rs

@@ -27,16 +27,17 @@ use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 
 use crate::common::data_type::PyScalarValue;
-use crate::errors::to_datafusion_err;
+use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionResult};
 use crate::expr::PyExpr;
-use crate::utils::parse_volatility;
+use crate::utils::{parse_volatility, validate_pycapsule};
 use datafusion::arrow::datatypes::DataType;
 use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType, ToPyArrow};
 use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_expr::{
     PartitionEvaluator, PartitionEvaluatorFactory, Signature, Volatility, WindowUDF, WindowUDFImpl,
 };
-use pyo3::types::{PyList, PyTuple};
+use datafusion_ffi::udwf::{FFI_WindowUDF, ForeignWindowUDF};
+use pyo3::types::{PyCapsule, PyList, PyTuple};
 
 #[derive(Debug)]
 struct RustPartitionEvaluator {
@@ -245,6 +246,26 @@ impl PyWindowUDF {
         Ok(self.function.call(args).into())
     }
 
+    #[staticmethod]
+    pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
+        if func.hasattr("__datafusion_window_udf__")? {
+            let capsule = func.getattr("__datafusion_window_udf__")?.call0()?;
+            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+            validate_pycapsule(capsule, "datafusion_window_udf")?;
+
+            let udwf = unsafe { capsule.reference::<FFI_WindowUDF>() };
+            let udwf: ForeignWindowUDF = udwf.try_into()?;
+
+            Ok(Self {
+                function: udwf.into(),
+            })
+        } else {
+            Err(crate::errors::PyDataFusionError::Common(
+                "__datafusion_window_udf__ does not exist on WindowUDF object.".to_string(),
+            ))
+        }
+    }
+
     fn __repr__(&self) -> PyResult<String> {
         Ok(format!("WindowUDF({})", self.function.name()))
     }