feat: restore pub UDAF/UDWF helpers and document inline encoding

Re-export `to_rust_accumulator`, `to_rust_partition_evaluator`, and
`PythonFunctionWindowUDF` (with a `MultiColumnWindowUDF` alias) by
promoting `udaf` and `udwf` to `pub mod` so prior downstream Rust
consumers keep their API surface after the inline-encoding refactor.

Adds an end-to-end window UDF pickle round-trip test that runs the
decoded evaluator over a real session, mirroring the aggregate test.

Documents the cloudpickle-based shipping behavior of Python aggregate
and window UDFs in the user-guide aggregations and windows pages.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: timsaucer

crates/core/src/lib.rs

@@ -59,11 +59,11 @@ mod array;
 #[cfg(feature = "substrait")]
 pub mod substrait;
 #[allow(clippy::borrow_deref_ref)]
-mod udaf;
+pub mod udaf;
 #[allow(clippy::borrow_deref_ref)]
 mod udf;
 pub mod udtf;
-mod udwf;
+pub mod udwf;
 
 #[cfg(feature = "mimalloc")]
 #[global_allocator]

crates/core/src/udaf.rs

@@ -25,7 +25,7 @@ use datafusion::common::ScalarValue;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion::logical_expr::{
-    Accumulator, AggregateUDF, AggregateUDFImpl, Signature, Volatility,
+    Accumulator, AccumulatorFactoryFunction, AggregateUDF, AggregateUDFImpl, Signature, Volatility,
 };
 use datafusion_ffi::udaf::FFI_AggregateUDF;
 use datafusion_python_util::parse_volatility;
@@ -154,6 +154,17 @@ fn instantiate_accumulator(accum: &Py<PyAny>) -> Result<Box<dyn Accumulator>> {
     Ok(Box::new(RustAccumulator::new(instance)))
 }
 
+/// Wrap a Python accumulator factory in an `AccumulatorFactoryFunction`.
+///
+/// Retained for downstream callers that previously consumed this
+/// helper to build a [`AccumulatorFactoryFunction`] for `create_udaf`
+/// or similar factory-based APIs. New in-crate code should construct
+/// a [`PythonFunctionAggregateUDF`] directly so the codec can downcast
+/// and ship it inline.
+pub fn to_rust_accumulator(accum: Py<PyAny>) -> AccumulatorFactoryFunction {
+    Arc::new(move |_args| instantiate_accumulator(&accum))
+}
+
 /// Named-struct `AggregateUDFImpl` for Python-defined aggregate UDFs.
 /// Holds the Python accumulator factory directly so the codec can
 /// downcast and cloudpickle it across process boundaries.

crates/core/src/udwf.rs

@@ -26,7 +26,7 @@ use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_expr::function::{PartitionEvaluatorArgs, WindowUDFFieldArgs};
 use datafusion::logical_expr::window_state::WindowAggState;
 use datafusion::logical_expr::{
-    PartitionEvaluator, Signature, Volatility, WindowUDF, WindowUDFImpl,
+    PartitionEvaluator, PartitionEvaluatorFactory, Signature, Volatility, WindowUDF, WindowUDFImpl,
 };
 use datafusion::scalar::ScalarValue;
 use datafusion_ffi::udwf::FFI_WindowUDF;
@@ -205,6 +205,17 @@ fn instantiate_partition_evaluator(evaluator: &Py<PyAny>) -> Result<Box<dyn Part
     Ok(Box::new(RustPartitionEvaluator::new(instance)))
 }
 
+/// Wrap a Python evaluator factory in a `PartitionEvaluatorFactory`.
+///
+/// Retained for downstream callers that previously consumed this
+/// helper to build a [`PartitionEvaluatorFactory`] for factory-based
+/// APIs. New in-crate code should construct a
+/// [`PythonFunctionWindowUDF`] directly so the codec can downcast and
+/// ship it inline.
+pub fn to_rust_partition_evaluator(evaluator: Py<PyAny>) -> PartitionEvaluatorFactory {
+    Arc::new(move || instantiate_partition_evaluator(&evaluator))
+}
+
 /// Represents an WindowUDF
 #[pyclass(
     from_py_object,
@@ -279,16 +290,26 @@ impl PyWindowUDF {
     }
 }
 
+/// `WindowUDFImpl` for Python-defined window UDFs.
+///
+/// Holds the Python evaluator factory directly so the codec can
+/// downcast and cloudpickle it across process boundaries. Replaces
+/// the prior factory-erased `MultiColumnWindowUDF`; the old name is
+/// kept as a type alias below for backward compatibility.
 #[derive(Debug)]
-pub(crate) struct PythonFunctionWindowUDF {
+pub struct PythonFunctionWindowUDF {
     name: String,
     evaluator: Py<PyAny>,
     signature: Signature,
     return_type: DataType,
 }
 
+/// Backward-compatible alias for downstream crates that referenced the
+/// previous struct name. New code should use [`PythonFunctionWindowUDF`].
+pub type MultiColumnWindowUDF = PythonFunctionWindowUDF;
+
 impl PythonFunctionWindowUDF {
-    pub(crate) fn new(
+    pub fn new(
         name: impl Into<String>,
         evaluator: Py<PyAny>,
         input_types: Vec<DataType>,

docs/source/user-guide/common-operations/aggregations.rst

@@ -434,3 +434,21 @@ The available aggregate functions are:
     - :py:meth:`datafusion.expr.GroupingSet.cube`
     - :py:meth:`datafusion.expr.GroupingSet.grouping_sets`
 
+User-Defined Aggregate Functions
+--------------------------------
+
+You can ship custom aggregations to the engine by subclassing
+:py:class:`~datafusion.user_defined.Accumulator` and registering it via
+:py:func:`~datafusion.udaf`. See :py:mod:`datafusion.user_defined` for
+the accumulator interface and worked examples.
+
+.. note:: Serialization
+
+   Python aggregate UDFs travel inline inside pickled or
+   :py:meth:`~datafusion.expr.Expr.to_bytes`-serialized expressions —
+   the accumulator class is captured by value via :mod:`cloudpickle`,
+   so worker processes do not need to pre-register the UDF. Any names
+   the accumulator resolves via ``import`` are captured **by reference**
+   and must be importable on the receiving worker. See
+   :py:mod:`datafusion.ipc` for the full IPC model and security caveats.
+

docs/source/user-guide/common-operations/windows.rst

@@ -213,3 +213,21 @@ The possible window functions are:
 
 3. Aggregate Functions
     - All :ref:`Aggregation Functions<aggregation>` can be used as window functions.
+
+User-Defined Window Functions
+-----------------------------
+
+You can ship custom window functions to the engine by subclassing
+:py:class:`~datafusion.user_defined.WindowEvaluator` and registering it
+via :py:func:`~datafusion.udwf`. See :py:mod:`datafusion.user_defined`
+for the evaluator interface and worked examples.
+
+.. note:: Serialization
+
+   Python window UDFs travel inline inside pickled or
+   :py:meth:`~datafusion.expr.Expr.to_bytes`-serialized expressions —
+   the evaluator class is captured by value via :mod:`cloudpickle`, so
+   worker processes do not need to pre-register the UDF. Any names the
+   evaluator resolves via ``import`` are captured **by reference** and
+   must be importable on the receiving worker. See
+   :py:mod:`datafusion.ipc` for the full IPC model and security caveats.

python/tests/test_pickle_expr.py

@@ -254,6 +254,23 @@ def test_window_udf_decodes_via_pickle_with_no_worker_ctx(self):
         decoded = pickle.loads(blob)  # noqa: S301
         assert "count_up" in decoded.canonical_name()
 
+    def test_window_udf_evaluates_after_roundtrip(self):
+        """End-to-end: decoded window UDF runs and emits per-row values
+        produced by the round-tripped evaluator factory."""
+        from datafusion.expr import WindowFrame
+
+        u = self._build_window_udf()
+        e = u(col("a"))
+        decoded = pickle.loads(pickle.dumps(e))  # noqa: S301
+
+        ctx = SessionContext()
+        df = ctx.from_pydict({"a": [1, 2, 3, 4, 5]})
+        framed = (
+            decoded.window_frame(WindowFrame("rows", None, None)).build().alias("c")
+        )
+        out = df.select(framed).to_pydict()
+        assert out["c"] == [0, 1, 2, 3, 4]
+
 
 class TestErrorPaths:
     def test_from_bytes_rejects_garbage(self):