feat: inline encoding for Python window UDFs

timsaucer · claude · timsaucer · commit 89d119f190ca · 2026-05-15T09:03:04.000-04:00
Window UDFs no longer need worker-side pre-registration. The codec
serializes the Python evaluator factory into the wire format and the
receiver reconstructs the UDF from bytes alone, same as scalar UDFs.

Refactor `MultiColumnWindowUDF` to store the Python evaluator
callable directly (`evaluator: Py&lt;PyAny&gt;`) instead of a
`PartitionEvaluatorFactory` closure. The factory closure was a
boxed `Fn` that captured the Python state opaquely, with nothing for
the codec to downcast back to. Now the named struct holds the
`Py&lt;PyAny&gt;` and builds a partition evaluator inside
`partition_evaluator()` on demand.

`PyWindowUDF::new` constructs `MultiColumnWindowUDF` directly with
the evaluator. `to_rust_partition_evaluator` is replaced by
`instantiate_partition_evaluator`, called from the trait method.

Codec wiring:
* `crates/core/src/codec.rs` adds `try_encode_python_window_udf` /
  `try_decode_python_window_udf` plus the `DFPYUDW1` magic prefix.
* `PythonLogicalCodec.try_encode_udwf` / `try_decode_udwf` and the
  matching `PythonPhysicalCodec` methods consult the helpers first
  and fall back to `inner` for non-Python window UDFs.

Test coverage in `test_pickle_expr.py::TestWindowUDFCodec` mirrors
the scalar UDF cases: self-contained blob, decode into fresh
context, decode via pickle with no worker context.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs
@@ -86,6 +86,7 @@ use datafusion::datasource::file_format::FileFormatFactory;
 use datafusion::execution::TaskContext;
 use datafusion::logical_expr::{
     AggregateUDF, Extension, LogicalPlan, ScalarUDF, ScalarUDFImpl, TypeSignature, WindowUDF,
+    WindowUDFImpl,
 };
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion::physical_plan::ExecutionPlan;
@@ -95,13 +96,19 @@ use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyTuple};
 
 use crate::udf::PythonFunctionScalarUDF;
+use crate::udwf::MultiColumnWindowUDF;
 
 /// Wire-format prefix that tags a `fun_definition` payload as an
 /// inlined Python scalar UDF (cloudpickled tuple of name, callable,
 /// input schema, return field, volatility). Defined once here so
 /// the encoder and decoder cannot drift.
 pub(crate) const PY_SCALAR_UDF_MAGIC: &[u8] = b"DFPYUDF1";
 
+/// Wire-format prefix for an inlined Python window UDF (cloudpickled
+/// tuple of name, evaluator factory, input schema, return type,
+/// volatility).
+pub(crate) const PY_WINDOW_UDF_MAGIC: &[u8] = b"DFPYUDW1";
+
 /// `LogicalExtensionCodec` parked on every `SessionContext`. Holds
 /// the Python-aware encoding hooks for logical-layer types
 /// (`LogicalPlan`, `Expr`) and delegates everything it does not
@@ -206,10 +213,16 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
     }
 
     fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> {
+        if try_encode_python_window_udf(node, buf)? {
+            return Ok(());
+        }
         self.inner.try_encode_udwf(node, buf)
     }
 
     fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> {
+        if let Some(udwf) = try_decode_python_window_udf(buf)? {
+            return Ok(udwf);
+        }
         self.inner.try_decode_udwf(name, buf)
     }
 }
@@ -296,10 +309,16 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
     }
 
     fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> {
+        if try_encode_python_window_udf(node, buf)? {
+            return Ok(());
+        }
         self.inner.try_encode_udwf(node, buf)
     }
 
     fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> {
+        if let Some(udwf) = try_decode_python_window_udf(buf)? {
+            return Ok(udwf);
+        }
         self.inner.try_decode_udwf(name, buf)
     }
 }
@@ -471,3 +490,126 @@ fn schema_from_ipc_bytes(bytes: &[u8]) -> arrow::error::Result<Schema> {
     let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)?;
     Ok(reader.schema().as_ref().clone())
 }
+
+// =============================================================================
+// Shared Python window UDF encode / decode helpers
+//
+// Cloudpickle tuple shape: `(name, evaluator_factory, input_schema_bytes,
+// return_schema_bytes, volatility_str)`. The evaluator factory is the
+// Python callable that produces a new evaluator instance per partition.
+// =============================================================================
+
+pub(crate) fn try_encode_python_window_udf(node: &WindowUDF, buf: &mut Vec<u8>) -> Result<bool> {
+    let Some(py_udf) = node.inner().as_any().downcast_ref::<MultiColumnWindowUDF>() else {
+        return Ok(false);
+    };
+
+    Python::attach(|py| -> Result<bool> {
+        let bytes = encode_python_window_udf(py, py_udf)
+            .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
+        buf.extend_from_slice(PY_WINDOW_UDF_MAGIC);
+        buf.extend_from_slice(&bytes);
+        Ok(true)
+    })
+}
+
+pub(crate) fn try_decode_python_window_udf(buf: &[u8]) -> Result<Option<Arc<WindowUDF>>> {
+    if buf.is_empty() || !buf.starts_with(PY_WINDOW_UDF_MAGIC) {
+        return Ok(None);
+    }
+    let payload = &buf[PY_WINDOW_UDF_MAGIC.len()..];
+
+    Python::attach(|py| -> Result<Option<Arc<WindowUDF>>> {
+        let udf = decode_python_window_udf(py, payload)
+            .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
+        Ok(Some(Arc::new(WindowUDF::new_from_impl(udf))))
+    })
+}
+
+fn encode_python_window_udf(py: Python<'_>, udf: &MultiColumnWindowUDF) -> PyResult<Vec<u8>> {
+    let cloudpickle = py.import("cloudpickle")?;
+
+    let signature = WindowUDFImpl::signature(udf);
+    let input_dtypes: Vec<arrow::datatypes::DataType> = match &signature.type_signature {
+        TypeSignature::Exact(types) => types.clone(),
+        other => {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "MultiColumnWindowUDF expected Signature::Exact, got {other:?}"
+            )));
+        }
+    };
+    let input_fields: Vec<Field> = input_dtypes
+        .into_iter()
+        .enumerate()
+        .map(|(i, dt)| Field::new(format!("arg_{i}"), dt, true))
+        .collect();
+    let input_schema = Schema::new(input_fields);
+    let input_schema_bytes = schema_to_ipc_bytes(&input_schema)
+        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
+
+    let return_schema = Schema::new(vec![Field::new("result", udf.return_type().clone(), true)]);
+    let return_schema_bytes = schema_to_ipc_bytes(&return_schema)
+        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
+
+    let volatility = format!("{:?}", signature.volatility).to_lowercase();
+
+    let payload = PyTuple::new(
+        py,
+        [
+            WindowUDFImpl::name(udf).into_pyobject(py)?.into_any(),
+            udf.evaluator().bind(py).clone().into_any(),
+            PyBytes::new(py, &input_schema_bytes).into_any(),
+            PyBytes::new(py, &return_schema_bytes).into_any(),
+            volatility.into_pyobject(py)?.into_any(),
+        ],
+    )?;
+
+    let blob = cloudpickle.call_method1("dumps", (payload,))?;
+    blob.extract::<Vec<u8>>()
+}
+
+fn decode_python_window_udf(py: Python<'_>, payload: &[u8]) -> PyResult<MultiColumnWindowUDF> {
+    let cloudpickle = py.import("cloudpickle")?;
+
+    let tuple = cloudpickle
+        .call_method1("loads", (PyBytes::new(py, payload),))?
+        .cast_into::<PyTuple>()?;
+
+    let name: String = tuple.get_item(0)?.extract()?;
+    let evaluator: Py<PyAny> = tuple.get_item(1)?.unbind();
+    let input_schema_bytes: Vec<u8> = tuple.get_item(2)?.extract()?;
+    let return_schema_bytes: Vec<u8> = tuple.get_item(3)?.extract()?;
+    let volatility_str: String = tuple.get_item(4)?.extract()?;
+
+    let input_schema = schema_from_ipc_bytes(&input_schema_bytes)
+        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
+    let input_types: Vec<arrow::datatypes::DataType> = input_schema
+        .fields()
+        .iter()
+        .map(|f| f.data_type().clone())
+        .collect();
+
+    let return_schema = schema_from_ipc_bytes(&return_schema_bytes)
+        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
+    let return_type = return_schema
+        .fields()
+        .first()
+        .ok_or_else(|| {
+            pyo3::exceptions::PyValueError::new_err(
+                "MultiColumnWindowUDF return schema must contain exactly one field",
+            )
+        })?
+        .data_type()
+        .clone();
+
+    let volatility = datafusion_python_util::parse_volatility(&volatility_str)
+        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
+
+    Ok(MultiColumnWindowUDF::from_parts(
+        name,
+        evaluator,
+        input_types,
+        return_type,
+        volatility,
+    ))
+}
diff --git a/crates/core/src/udwf.rs b/crates/core/src/udwf.rs
@@ -25,10 +25,9 @@ use datafusion::arrow::datatypes::DataType;
 use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType, ToPyArrow};
 use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_expr::function::{PartitionEvaluatorArgs, WindowUDFFieldArgs};
-use datafusion::logical_expr::ptr_eq::PtrEq;
 use datafusion::logical_expr::window_state::WindowAggState;
 use datafusion::logical_expr::{
-    PartitionEvaluator, PartitionEvaluatorFactory, Signature, Volatility, WindowUDF, WindowUDFImpl,
+    PartitionEvaluator, Signature, Volatility, WindowUDF, WindowUDFImpl,
 };
 use datafusion::scalar::ScalarValue;
 use datafusion_ffi::udwf::FFI_WindowUDF;
@@ -198,15 +197,13 @@ impl PartitionEvaluator for RustPartitionEvaluator {
     }
 }
 
-pub fn to_rust_partition_evaluator(evaluator: Py<PyAny>) -> PartitionEvaluatorFactory {
-    Arc::new(move || -> Result<Box<dyn PartitionEvaluator>> {
-        let evaluator = Python::attach(|py| {
-            evaluator
-                .call0(py)
-                .map_err(|e| DataFusionError::Execution(e.to_string()))
-        })?;
-        Ok(Box::new(RustPartitionEvaluator::new(evaluator)))
-    })
+fn instantiate_partition_evaluator(evaluator: &Py<PyAny>) -> Result<Box<dyn PartitionEvaluator>> {
+    let instance = Python::attach(|py| {
+        evaluator
+            .call0(py)
+            .map_err(|e| DataFusionError::Execution(e.to_string()))
+    })?;
+    Ok(Box::new(RustPartitionEvaluator::new(instance)))
 }
 
 /// Represents an WindowUDF
@@ -234,14 +231,14 @@ impl PyWindowUDF {
         volatility: &str,
     ) -> PyResult<Self> {
         let return_type = return_type.0;
-        let input_types = input_types.into_iter().map(|t| t.0).collect();
+        let input_types: Vec<DataType> = input_types.into_iter().map(|t| t.0).collect();
 
         let function = WindowUDF::from(MultiColumnWindowUDF::new(
             name,
+            evaluator,
             input_types,
             return_type,
             parse_volatility(volatility)?,
-            to_rust_partition_evaluator(evaluator),
         ));
         Ok(Self { function })
     }
@@ -278,42 +275,79 @@ impl PyWindowUDF {
     }
 }
 
-#[derive(Hash, Eq, PartialEq)]
+#[derive(Debug)]
 pub struct MultiColumnWindowUDF {
     name: String,
+    evaluator: Py<PyAny>,
     signature: Signature,
     return_type: DataType,
-    partition_evaluator_factory: PtrEq<PartitionEvaluatorFactory>,
-}
-
-impl std::fmt::Debug for MultiColumnWindowUDF {
-    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        f.debug_struct("WindowUDF")
-            .field("name", &self.name)
-            .field("signature", &self.signature)
-            .field("return_type", &"<func>")
-            .field("partition_evaluator_factory", &"<FUNC>")
-            .finish()
-    }
 }
 
 impl MultiColumnWindowUDF {
     pub fn new(
         name: impl Into<String>,
+        evaluator: Py<PyAny>,
         input_types: Vec<DataType>,
         return_type: DataType,
         volatility: Volatility,
-        partition_evaluator_factory: PartitionEvaluatorFactory,
     ) -> Self {
         let name = name.into();
         let signature = Signature::exact(input_types, volatility);
         Self {
             name,
+            evaluator,
             signature,
             return_type,
-            partition_evaluator_factory: partition_evaluator_factory.into(),
         }
     }
+
+    /// Stored Python callable that produces a fresh partition
+    /// evaluator instance per partition. Consumed by the codec to
+    /// cloudpickle the evaluator factory across process boundaries.
+    pub(crate) fn evaluator(&self) -> &Py<PyAny> {
+        &self.evaluator
+    }
+
+    pub(crate) fn return_type(&self) -> &DataType {
+        &self.return_type
+    }
+
+    pub(crate) fn from_parts(
+        name: String,
+        evaluator: Py<PyAny>,
+        input_types: Vec<DataType>,
+        return_type: DataType,
+        volatility: Volatility,
+    ) -> Self {
+        Self::new(name, evaluator, input_types, return_type, volatility)
+    }
+}
+
+impl Eq for MultiColumnWindowUDF {}
+impl PartialEq for MultiColumnWindowUDF {
+    fn eq(&self, other: &Self) -> bool {
+        self.name == other.name
+            && self.signature == other.signature
+            && self.return_type == other.return_type
+            && Python::attach(|py| {
+                self.evaluator
+                    .bind(py)
+                    .eq(other.evaluator.bind(py))
+                    .unwrap_or(false)
+            })
+    }
+}
+
+impl std::hash::Hash for MultiColumnWindowUDF {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+        self.signature.hash(state);
+        self.return_type.hash(state);
+        Python::attach(|py| {
+            let py_hash = self.evaluator.bind(py).hash().unwrap_or(0);
+            state.write_isize(py_hash);
+        });
+    }
 }
 
 impl WindowUDFImpl for MultiColumnWindowUDF {
@@ -339,7 +373,6 @@ impl WindowUDFImpl for MultiColumnWindowUDF {
         &self,
         _partition_evaluator_args: PartitionEvaluatorArgs,
     ) -> Result<Box<dyn PartitionEvaluator>> {
-        let _ = _partition_evaluator_args;
-        (self.partition_evaluator_factory)()
+        instantiate_partition_evaluator(&self.evaluator)
     }
 }
diff --git a/python/tests/test_pickle_expr.py b/python/tests/test_pickle_expr.py
@@ -130,6 +130,49 @@ def fn(arr):
         assert decoded.canonical_name() == e.canonical_name()
 
 
+class TestWindowUDFCodec:
+    """Python window UDFs travel inline like scalar UDFs."""
+
+    def _build_window_udf(self):
+        from datafusion import udwf
+        from datafusion.user_defined import WindowEvaluator
+
+        class CountUpEvaluator(WindowEvaluator):
+            def evaluate_all(self, values, num_rows):
+                return pa.array(list(range(num_rows)))
+
+        return udwf(
+            CountUpEvaluator,
+            [pa.int64()],
+            pa.int64(),
+            "immutable",
+            name="count_up",
+        )
+
+    def test_window_udf_self_contained_blob(self):
+        u = self._build_window_udf()
+        e = u(col("a"))
+        blob = pickle.dumps(e)
+        assert len(blob) > 200
+
+    def test_window_udf_decodes_into_fresh_ctx(self):
+        u = self._build_window_udf()
+        e = u(col("a"))
+        blob = e.to_bytes()
+        fresh = SessionContext()
+        from datafusion import Expr
+
+        decoded = Expr.from_bytes(blob, ctx=fresh)
+        assert "count_up" in decoded.canonical_name()
+
+    def test_window_udf_decodes_via_pickle_with_no_worker_ctx(self):
+        u = self._build_window_udf()
+        e = u(col("a"))
+        blob = pickle.dumps(e)
+        decoded = pickle.loads(blob)
+        assert "count_up" in decoded.canonical_name()
+
+
 class TestWorkerCtxLifecycle:
     def test_set_and_clear(self):
         assert get_worker_ctx() is None
