feat: accept a bare PyCapsule in ScalarUDF and WindowUDF

AggregateUDF already accepted the capsule returned by
__datafusion_aggregate_udf__ as well as an object exposing it (#1277).
Extend the same to ScalarUDF and WindowUDF, with matching overloads on
udf and udwf, so the three UDF kinds import the same way. Add FFI
example tests for all three, including the previously untested
AggregateUDF path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: timsaucer

crates/core/src/udf.rs

@@ -209,6 +209,16 @@ impl ScalarUDFImpl for PythonFunctionScalarUDF {
     }
 }
 
+fn scalar_udf_from_capsule(capsule: &Bound<'_, PyCapsule>) -> PyDataFusionResult<ScalarUDF> {
+    let data: NonNull<FFI_ScalarUDF> = capsule
+        .pointer_checked(Some(c"datafusion_scalar_udf"))?
+        .cast();
+    let udf = unsafe { data.as_ref() };
+    let udf: Arc<dyn ScalarUDFImpl> = udf.into();
+
+    Ok(ScalarUDF::new_from_shared_impl(udf))
+}
+
 /// Represents a PyScalarUDF
 #[pyclass(
     from_py_object,
@@ -247,22 +257,21 @@ impl PyScalarUDF {
 
     #[staticmethod]
     pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
+        if func.is_instance_of::<PyCapsule>() {
+            let capsule = func.cast::<PyCapsule>().map_err(to_datafusion_err)?;
+            let function = scalar_udf_from_capsule(capsule)?;
+            return Ok(Self { function });
+        }
+
         if func.hasattr("__datafusion_scalar_udf__")? {
             let capsule = call_capsule_getter(
                 func.clone(),
                 "__datafusion_scalar_udf__",
                 CapsuleGetterArg::None,
             )?;
             let capsule = capsule.cast::<PyCapsule>().map_err(to_datafusion_err)?;
-            let data: NonNull<FFI_ScalarUDF> = capsule
-                .pointer_checked(Some(c"datafusion_scalar_udf"))?
-                .cast();
-            let udf = unsafe { data.as_ref() };
-            let udf: Arc<dyn ScalarUDFImpl> = udf.into();
-
-            Ok(Self {
-                function: ScalarUDF::new_from_shared_impl(udf),
-            })
+            let function = scalar_udf_from_capsule(capsule)?;
+            Ok(Self { function })
         } else {
             Err(crate::errors::PyDataFusionError::Common(
                 "__datafusion_scalar_udf__ does not exist on ScalarUDF object.".to_string(),

crates/core/src/udwf.rs

@@ -216,6 +216,16 @@ pub fn to_rust_partition_evaluator(evaluator: Py<PyAny>) -> PartitionEvaluatorFa
     Arc::new(move || instantiate_partition_evaluator(&evaluator))
 }
 
+fn window_udf_from_capsule(capsule: &Bound<'_, PyCapsule>) -> PyDataFusionResult<WindowUDF> {
+    let data: NonNull<FFI_WindowUDF> = capsule
+        .pointer_checked(Some(c"datafusion_window_udf"))?
+        .cast();
+    let udwf = unsafe { data.as_ref() };
+    let udwf: Arc<dyn WindowUDFImpl> = udwf.into();
+
+    Ok(WindowUDF::new_from_shared_impl(udwf))
+}
+
 /// Represents an WindowUDF
 #[pyclass(
     from_py_object,
@@ -262,19 +272,17 @@ impl PyWindowUDF {
 
     #[staticmethod]
     pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
+        if func.is_instance_of::<PyCapsule>() {
+            let capsule = func.cast::<PyCapsule>().map_err(to_datafusion_err)?;
+            let function = window_udf_from_capsule(capsule)?;
+            return Ok(Self { function });
+        }
+
         let capsule =
             call_capsule_getter(func, "__datafusion_window_udf__", CapsuleGetterArg::None)?;
-
         let capsule = capsule.cast::<PyCapsule>().map_err(to_datafusion_err)?;
-        let data: NonNull<FFI_WindowUDF> = capsule
-            .pointer_checked(Some(c"datafusion_window_udf"))?
-            .cast();
-        let udwf = unsafe { data.as_ref() };
-        let udwf: Arc<dyn WindowUDFImpl> = udwf.into();
-
-        Ok(Self {
-            function: WindowUDF::new_from_shared_impl(udwf),
-        })
+        let function = window_udf_from_capsule(capsule)?;
+        Ok(Self { function })
     }
 
     fn __repr__(&self) -> PyResult<String> {

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

@@ -75,3 +75,11 @@ def test_ffi_aggregate_call_directly():
     ]
 
     assert result == expected
+
+
+def test_ffi_aggregate_from_bare_capsule():
+    ctx = setup_context_with_table()
+    my_udaf = udaf(MySumUDF().__datafusion_aggregate_udf__())
+
+    result = ctx.table("test_table").aggregate([], [my_udaf(col("a")).alias("r")])
+    assert result.collect_column("r").to_pylist() == [6]

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

@@ -68,3 +68,11 @@ def test_ffi_scalar_call_directly():
     ]
 
     assert result == expected
+
+
+def test_ffi_scalar_from_bare_capsule():
+    ctx = setup_context_with_table()
+    my_udf = udf(IsNullUDF().__datafusion_scalar_udf__())
+
+    result = ctx.table("test_table").select(my_udf(col("a")).alias("r"))
+    assert result.collect_column("r").to_pylist() == [False, False, False, True]

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

@@ -87,3 +87,15 @@ def test_ffi_window_call_directly():
         (40, 4),
     ]
     assert results == expected
+
+
+def test_ffi_window_from_bare_capsule():
+    ctx = setup_context_with_table()
+    my_udwf = udwf(MyRankUDF().__datafusion_window_udf__())
+
+    result = (
+        ctx.table("test_table")
+        .select(col("a"), my_udwf().order_by(col("a")).build().alias("r"))
+        .sort(col("a"))
+    )
+    assert result.collect_column("r").to_pylist() == [1, 2, 3, 4]

python/datafusion/user_defined.py

@@ -283,6 +283,10 @@ def udf(
     @staticmethod
     def udf(func: ScalarUDFExportable) -> ScalarUDF: ...
 
+    @overload
+    @staticmethod
+    def udf(func: _PyCapsule) -> ScalarUDF: ...
+
     @staticmethod
     def udf(*args: Any, **kwargs: Any):  # noqa: D417
         """Create a new User-Defined Function (UDF).
@@ -388,7 +392,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Callable:
 
             return decorator
 
-        if hasattr(args[0], "__datafusion_scalar_udf__"):
+        if hasattr(args[0], "__datafusion_scalar_udf__") or _is_pycapsule(args[0]):
             return ScalarUDF.from_pycapsule(args[0])
 
         if args and callable(args[0]):
@@ -398,12 +402,18 @@ def wrapper(*args: Any, **kwargs: Any) -> Callable:
         return _decorator(*args, **kwargs)
 
     @staticmethod
-    def from_pycapsule(func: ScalarUDFExportable) -> ScalarUDF:
+    def from_pycapsule(func: ScalarUDFExportable | _PyCapsule) -> ScalarUDF:
         """Create a Scalar UDF from ScalarUDF PyCapsule object.
 
         This function will instantiate a Scalar UDF that uses a DataFusion
         ScalarUDF that is exported via the FFI bindings.
         """
+        if _is_pycapsule(func):
+            scalar = cast("ScalarUDF", object.__new__(ScalarUDF))
+            scalar._udf = df_internal.ScalarUDF.from_pycapsule(func)
+            return scalar
+
+        func = cast("ScalarUDFExportable", func)
         name = str(func.__class__)
         return ScalarUDF(
             name=name,
@@ -1011,6 +1021,14 @@ def udwf(
         name: str | None = None,
     ) -> WindowUDF: ...
 
+    @overload
+    @staticmethod
+    def udwf(func: WindowUDFExportable) -> WindowUDF: ...
+
+    @overload
+    @staticmethod
+    def udwf(func: _PyCapsule) -> WindowUDF: ...
+
     @staticmethod
     def udwf(*args: Any, **kwargs: Any):  # noqa: D417
         """Create a new User-Defined Window Function (UDWF).
@@ -1075,7 +1093,7 @@ def udwf(*args: Any, **kwargs: Any):  # noqa: D417
         Returns:
             A user-defined window function that can be used in window function calls.
         """
-        if hasattr(args[0], "__datafusion_window_udf__"):
+        if hasattr(args[0], "__datafusion_window_udf__") or _is_pycapsule(args[0]):
             return WindowUDF.from_pycapsule(args[0])
 
         if args and callable(args[0]):
@@ -1146,12 +1164,18 @@ def wrapper(*args: Any, **kwargs: Any) -> Expr:
         return decorator
 
     @staticmethod
-    def from_pycapsule(func: WindowUDFExportable) -> WindowUDF:
+    def from_pycapsule(func: WindowUDFExportable | _PyCapsule) -> 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.
         """
+        if _is_pycapsule(func):
+            window = cast("WindowUDF", object.__new__(WindowUDF))
+            window._udwf = df_internal.WindowUDF.from_pycapsule(func)
+            return window
+
+        func = cast("WindowUDFExportable", func)
         name = str(func.__class__)
         return WindowUDF(
             name=name,