feat: pickle support for Expr via PythonLogicalCodec inline encoding

Builds on the codec consistency work in feat/proto-codecs. Python
scalar UDFs are cloudpickled inline into the proto `fun_definition`
field by PythonLogicalCodec / PythonPhysicalCodec, so a pickled Expr
that references a Python `udf()` reconstructs on the receiver with
no pre-registration. UDAFs, UDWFs, and FFI-imported UDFs still
resolve through the receiver's session.

Rust:
* `PythonFunctionScalarUDF` regains the `func()` / `input_fields()` /
  `return_field()` / `volatility()` / `from_parts()` accessors the
  codec needs.
* `crates/core/src/codec.rs` adds shared
  `try_encode_python_scalar_udf` / `try_decode_python_scalar_udf`
  helpers built on cloudpickle + pyarrow IPC for the input schema.
  Both `PythonLogicalCodec.try_encode_udf` and
  `PythonPhysicalCodec.try_encode_udf` consult the helper first and
  fall back to `inner` for non-Python UDFs (and the receiver's
  function registry on decode if the prefix does not match).

Python:
* `datafusion.ipc` module: thread-local `set_worker_ctx` /
  `clear_worker_ctx` / `get_worker_ctx` for installing a receiver
  `SessionContext` on a worker process. `_resolve_ctx` returns
  explicit > worker > fresh.
* `Expr.__reduce__` returns `(Expr._reconstruct, (self.to_bytes(),))`.
  `_reconstruct` calls `Expr.from_bytes(buf, ctx=None)` which
  consults the worker context.
* `Expr.from_bytes` signature switches to `(buf, ctx=None)` (was
  `(ctx, buf)`); no callers in main, only PR1 tests which are
  updated.
* `datafusion.ipc` exported from the top-level package.

Dependencies:
* `cloudpickle>=2.0` added as a runtime dep. Lazy-imported on the
  encode / decode hot paths — users who never pickle a plan or
  expression pay only the install footprint, not import-time cost.
* ruff `S301` added to the test-suite + examples ignore lists
  (legitimate `pickle.loads` use).

Tests:
* `test_pickle_expr.py` — 11 cases covering built-in expr pickle,
  scalar UDF self-contained blobs, closure-capturing UDFs, worker
  ctx lifecycle, thread-local isolation.
* `test_pickle_multiprocessing.py` + `_pickle_multiprocessing_helpers.py`
  — parametrized over `fork`/`forkserver`/`spawn` start methods. 9
  cases. Auto-skip when the sandbox blocks semaphore creation; CI
  runs the full matrix.
* `test_expr.py` — existing `from_bytes` tests updated to new
  signature.

1088 root tests pass (up from 1077), 13 skipped (up from 4, the new
mp cases skip locally under sandboxed semaphores).

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

Author: timsaucer

crates/core/src/codec.rs

@@ -77,22 +77,30 @@
 
 use std::sync::Arc;
 
-use arrow::datatypes::SchemaRef;
+use arrow::datatypes::{Field, Schema, SchemaRef};
+use arrow::pyarrow::ToPyArrow;
+use datafusion::arrow::pyarrow::FromPyArrow;
 use datafusion::common::{Result, TableReference};
 use datafusion::datasource::TableProvider;
 use datafusion::datasource::file_format::FileFormatFactory;
 use datafusion::execution::TaskContext;
-use datafusion::logical_expr::{AggregateUDF, Extension, LogicalPlan, ScalarUDF, WindowUDF};
+use datafusion::logical_expr::{
+    AggregateUDF, Extension, LogicalPlan, ScalarUDF, ScalarUDFImpl, WindowUDF,
+};
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec};
 use datafusion_proto::physical_plan::{DefaultPhysicalExtensionCodec, PhysicalExtensionCodec};
+use pyo3::BoundObject;
+use pyo3::prelude::*;
+use pyo3::types::{PyBytes, PyTuple};
+
+use crate::udf::PythonFunctionScalarUDF;
 
 /// 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.
-#[allow(dead_code)]
 pub(crate) const PY_SCALAR_UDF_MAGIC: &[u8] = b"DFPYUDF1";
 
 /// `LogicalExtensionCodec` parked on every `SessionContext`. Holds
@@ -177,10 +185,16 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
     }
 
     fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
+        if try_encode_python_scalar_udf(node, buf)? {
+            return Ok(());
+        }
         self.inner.try_encode_udf(node, buf)
     }
 
     fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+        if let Some(udf) = try_decode_python_scalar_udf(buf)? {
+            return Ok(udf);
+        }
         self.inner.try_decode_udf(name, buf)
     }
 
@@ -249,10 +263,16 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
     }
 
     fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
+        if try_encode_python_scalar_udf(node, buf)? {
+            return Ok(());
+        }
         self.inner.try_encode_udf(node, buf)
     }
 
     fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+        if let Some(udf) = try_decode_python_scalar_udf(buf)? {
+            return Ok(udf);
+        }
         self.inner.try_decode_udf(name, buf)
     }
 
@@ -284,3 +304,126 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
         self.inner.try_decode_udwf(name, buf)
     }
 }
+
+// =============================================================================
+// Shared Python scalar UDF encode / decode helpers
+//
+// Both `PythonLogicalCodec` and `PythonPhysicalCodec` consult these on
+// every `try_encode_udf` / `try_decode_udf` call. Same wire format on
+// both layers — a Python `ScalarUDF` referenced inside a `LogicalPlan`
+// or an `ExecutionPlan` round-trips identically.
+// =============================================================================
+
+/// Encode a Python scalar UDF inline if `node` is one. Returns
+/// `Ok(true)` when the payload (`DFPYUDF1` prefix + cloudpickled
+/// tuple) was written and the caller should skip its inner codec.
+/// Returns `Ok(false)` for any non-Python UDF, signalling the caller
+/// to delegate to its `inner`.
+pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<bool> {
+    let Some(py_udf) = node
+        .inner()
+        .as_any()
+        .downcast_ref::<PythonFunctionScalarUDF>()
+    else {
+        return Ok(false);
+    };
+
+    Python::attach(|py| -> Result<bool> {
+        let bytes = encode_python_scalar_udf(py, py_udf)
+            .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
+        buf.extend_from_slice(PY_SCALAR_UDF_MAGIC);
+        buf.extend_from_slice(&bytes);
+        Ok(true)
+    })
+}
+
+/// Decode an inline Python scalar UDF payload. Returns `Ok(None)`
+/// when `buf` does not carry the `DFPYUDF1` prefix, signalling the
+/// caller to delegate to its `inner` codec (and eventually the
+/// `FunctionRegistry`).
+pub(crate) fn try_decode_python_scalar_udf(buf: &[u8]) -> Result<Option<Arc<ScalarUDF>>> {
+    if buf.is_empty() || !buf.starts_with(PY_SCALAR_UDF_MAGIC) {
+        return Ok(None);
+    }
+    let payload = &buf[PY_SCALAR_UDF_MAGIC.len()..];
+
+    Python::attach(|py| -> Result<Option<Arc<ScalarUDF>>> {
+        let udf = decode_python_scalar_udf(py, payload)
+            .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
+        Ok(Some(Arc::new(ScalarUDF::new_from_impl(udf))))
+    })
+}
+
+/// Build the cloudpickle payload for a `PythonFunctionScalarUDF`.
+///
+/// Layout: `cloudpickle.dumps((name, func, input_schema_bytes,
+/// return_field, volatility_str))`. Input fields ride along as an
+/// IPC-encoded pyarrow Schema so they round-trip without extra
+/// plumbing.
+fn encode_python_scalar_udf(py: Python<'_>, udf: &PythonFunctionScalarUDF) -> PyResult<Vec<u8>> {
+    let cloudpickle = py.import("cloudpickle")?;
+
+    let input_schema = Schema::new(udf.input_fields().to_vec());
+    let pa_schema_obj = input_schema.to_pyarrow(py)?;
+    let pa_schema = pa_schema_obj.into_bound();
+    let schema_bytes: Vec<u8> = pa_schema
+        .call_method0("serialize")?
+        .call_method0("to_pybytes")?
+        .extract()?;
+
+    let return_field_obj = udf.return_field().as_ref().to_pyarrow(py)?;
+    let volatility = format!("{:?}", udf.volatility()).to_lowercase();
+
+    let payload = PyTuple::new(
+        py,
+        [
+            udf.name().into_pyobject(py)?.into_any(),
+            udf.func().bind(py).clone().into_any(),
+            PyBytes::new(py, &schema_bytes).into_any(),
+            return_field_obj.into_bound(),
+            volatility.into_pyobject(py)?.into_any(),
+        ],
+    )?;
+
+    let blob = cloudpickle.call_method1("dumps", (payload,))?;
+    blob.extract::<Vec<u8>>()
+}
+
+/// Inverse of [`encode_python_scalar_udf`].
+fn decode_python_scalar_udf(py: Python<'_>, payload: &[u8]) -> PyResult<PythonFunctionScalarUDF> {
+    let cloudpickle = py.import("cloudpickle")?;
+    let pyarrow = py.import("pyarrow")?;
+
+    let tuple = cloudpickle
+        .call_method1("loads", (PyBytes::new(py, payload),))?
+        .cast_into::<PyTuple>()?;
+
+    let name: String = tuple.get_item(0)?.extract()?;
+    let func: Py<PyAny> = tuple.get_item(1)?.unbind();
+    let schema_bytes: Vec<u8> = tuple.get_item(2)?.extract()?;
+    let return_field_py = tuple.get_item(3)?;
+    let volatility_str: String = tuple.get_item(4)?.extract()?;
+
+    let buffer = pyarrow.call_method1("py_buffer", (PyBytes::new(py, &schema_bytes),))?;
+    let pa_schema = pyarrow
+        .getattr("ipc")?
+        .call_method1("read_schema", (buffer,))?;
+
+    let schema = Schema::from_pyarrow_bound(&pa_schema)
+        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
+    let input_fields: Vec<Field> = schema.fields().iter().map(|f| f.as_ref().clone()).collect();
+
+    let return_field = Field::from_pyarrow_bound(&return_field_py)
+        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
+
+    let volatility = datafusion_python_util::parse_volatility(&volatility_str)
+        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
+
+    Ok(PythonFunctionScalarUDF::from_parts(
+        name,
+        func,
+        input_fields,
+        return_field,
+        volatility,
+    ))
+}

crates/core/src/udf.rs

@@ -43,11 +43,13 @@ use crate::expr::PyExpr;
 /// This struct holds the Python written function that is a
 /// ScalarUDF.
 #[derive(Debug)]
-struct PythonFunctionScalarUDF {
+pub(crate) struct PythonFunctionScalarUDF {
     name: String,
     func: Py<PyAny>,
-    signature: Signature,
+    input_fields: Vec<Field>,
     return_field: FieldRef,
+    signature: Signature,
+    volatility: Volatility,
 }
 
 impl PythonFunctionScalarUDF {
@@ -63,10 +65,40 @@ impl PythonFunctionScalarUDF {
         Self {
             name,
             func,
-            signature,
+            input_fields,
             return_field: Arc::new(return_field),
+            signature,
+            volatility,
         }
     }
+
+    /// Stored Python callable. Consumed by the codec to cloudpickle
+    /// the function body across process boundaries.
+    pub(crate) fn func(&self) -> &Py<PyAny> {
+        &self.func
+    }
+
+    pub(crate) fn input_fields(&self) -> &[Field] {
+        &self.input_fields
+    }
+
+    pub(crate) fn return_field(&self) -> &FieldRef {
+        &self.return_field
+    }
+
+    pub(crate) fn volatility(&self) -> Volatility {
+        self.volatility
+    }
+
+    pub(crate) fn from_parts(
+        name: String,
+        func: Py<PyAny>,
+        input_fields: Vec<Field>,
+        return_field: Field,
+        volatility: Volatility,
+    ) -> Self {
+        Self::new(name, func, input_fields, return_field, volatility)
+    }
 }
 
 impl Eq for PythonFunctionScalarUDF {}

pyproject.toml

@@ -44,6 +44,12 @@ classifiers = [
   "Programming Language :: Rust",
 ]
 dependencies = [
+  # cloudpickle is invoked by the Rust-side PythonLogicalCodec /
+  # PythonPhysicalCodec via pyo3 to serialize Python scalar UDF
+  # callables into the proto wire format. Lazy-imported on the encode
+  # / decode hot paths, so users who never serialize a plan or
+  # expression incur no runtime cost beyond the install footprint.
+  "cloudpickle>=2.0",
   "pyarrow>=16.0.0;python_version<'3.14'",
   "pyarrow>=22.0.0;python_version>='3.14'",
   "typing-extensions;python_version<'3.13'",
@@ -120,6 +126,7 @@ extend-allowed-calls = ["datafusion.lit", "lit"]
   "PT011",
   "RUF015",
   "S101",
+  "S301",
   "S608",
   "SLF",
 ]
@@ -133,6 +140,7 @@ extend-allowed-calls = ["datafusion.lit", "lit"]
   "PLR2004",
   "RUF015",
   "S101",
+  "S301",
   "T201",
   "W505",
 ]

python/datafusion/__init__.py

@@ -65,7 +65,7 @@
     import importlib_metadata  # type: ignore[import]
 
 # Public submodules
-from . import functions, object_store, substrait, unparser
+from . import functions, ipc, object_store, substrait, unparser
 
 # The following imports are okay to remain as opaque to the user.
 from ._internal import Config
@@ -142,6 +142,7 @@
     "configure_formatter",
     "expr",
     "functions",
+    "ipc",
     "lit",
     "literal",
     "object_store",

python/datafusion/expr.py

@@ -436,21 +436,51 @@ def variant_name(self) -> str:
     def to_bytes(self, ctx: SessionContext | None = None) -> bytes:
         """Serialize this expression to protobuf bytes.
 
-        When ``ctx`` is supplied, encoding routes through the session's
-        installed :class:`LogicalExtensionCodec`. Without ``ctx`` a
-        default codec is used.
+        Python scalar UDFs are cloudpickled inline by
+        :class:`PythonLogicalCodec`, so the returned blob is
+        self-contained for scalar UDFs. Aggregate / window / FFI UDFs
+        are stored by name only; the receiver must have them
+        registered.
+
+        When ``ctx`` is supplied, encoding also routes through the
+        session's installed codec stack.
         """
         ctx_arg = ctx.ctx if ctx is not None else None
-        return self.expr.to_bytes(ctx_arg)
+        return bytes(self.expr.to_bytes(ctx_arg))
 
-    @staticmethod
-    def from_bytes(ctx: SessionContext, data: bytes) -> Expr:
+    @classmethod
+    def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr:
         """Decode an expression from serialized protobuf bytes.
 
-        ``ctx`` provides the function registry for resolving UDF
-        references and the logical codec for in-band Python payloads.
+        ``ctx`` is the receiver :class:`SessionContext` used to resolve
+        function references not inlined by the codec (aggregate UDFs,
+        window UDFs, FFI UDFs). When ``ctx`` is ``None`` the worker
+        context set via :func:`datafusion.ipc.set_worker_ctx` is
+        consulted; if no worker context is set, a fresh
+        :class:`SessionContext` is used.
+        """
+        from datafusion.ipc import _resolve_ctx
+
+        resolved = _resolve_ctx(ctx)
+        return cls(expr_internal.RawExpr.from_bytes(resolved.ctx, buf))
+
+    def __reduce__(self) -> tuple:
+        """Pickle protocol hook.
+
+        :class:`PythonLogicalCodec` cloudpickles referenced Python
+        scalar UDFs directly into the proto wire format, so the
+        returned blob is self-contained. On unpickle the bytes are
+        decoded against the worker context set via
+        :func:`datafusion.ipc.set_worker_ctx` (or a fresh
+        :class:`SessionContext` if none) for any remaining
+        registry-resolved references.
         """
-        return Expr(expr_internal.RawExpr.from_bytes(ctx.ctx, data))
+        return (Expr._reconstruct, (self.to_bytes(),))
+
+    @classmethod
+    def _reconstruct(cls, proto_bytes: bytes) -> Expr:
+        """Internal entry point used by :meth:`__reduce__` on unpickle."""
+        return cls.from_bytes(proto_bytes)
 
     def __richcmp__(self, other: Expr, op: int) -> Expr:
         """Comparison operator."""