feat: per-session Python UDF inlining toggle + sender ctx + strict refusal

Adds a per-session toggle that turns inline Python UDF encoding on or
off, plus the supporting plumbing to make it usable through
pickle.dumps.

Codec layer:
  * PythonLogicalCodec / PythonPhysicalCodec gain a python_udf_inlining
    bool (default true) and a with_python_udf_inlining(enabled) builder.
    Each try_encode_udf{,af,wf} short-circuits to inner when the toggle
    is off; each try_decode_udf{,af,wf} that recognizes a DFPY* magic
    on a strict codec returns a clean Execution error instead of
    invoking cloudpickle.loads. The refusal message names the UDF and
    the wire family so an operator can see at a glance whether to
    re-encode the bytes or register the UDF on the receiver.

Session layer:
  * PySessionContext::with_python_udf_inlining(enabled) returns a new
    session whose stacked logical + physical codecs both carry the
    toggle. The Arc<SessionState> is cloned (cheap), only the codec
    pair is rebuilt, so registrations and config stay attached.
  * SessionContext.with_python_udf_inlining(*, enabled) is the Python
    wrapper. enabled is keyword-only because positional booleans at
    the call site read as opaque.

Sender-side context:
  * datafusion.ipc gains set_sender_ctx / get_sender_ctx /
    clear_sender_ctx thread-locals. Expr.__reduce__ now consults
    get_sender_ctx() to pick the codec for outbound pickles, which is
    the only path through which a strict session affects pickle.dumps
    (the protocol calls __reduce__ with no arguments). Without a
    sender context the default codec is used.

Tests:
  * test_pickle_expr.py picks up TestPythonUdfInliningToggle (covers
    both directions of the toggle plus the explicit-ctx fast path),
    TestWorkerCtxLifecycle (set/clear/threading), and
    TestSenderCtxLifecycle.
  * New test_pickle_multiprocessing.py + helpers exercise the full
    driver -> worker round-trip on a multiprocessing.Pool with set_*_ctx
    installed in the worker initializer.
  * CI workflow gets a 30-minute timeout-minutes backstop so a hung
    pickle worker can't block the matrix indefinitely.

User-guide docs and the runnable examples land in PR4 of this series.

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

Author: timsaucer

.github/workflows/test.yml

@@ -29,6 +29,9 @@ env:
 jobs:
   test-matrix:
     runs-on: ubuntu-latest
+    # Backstop: a hung multiprocessing worker (e.g. during a pickle regression)
+    # should not block CI longer than this.
+    timeout-minutes: 30
     strategy:
       fail-fast: false
       matrix:

crates/core/src/codec.rs

@@ -232,16 +232,44 @@ fn strip_wire_header<'a>(
 #[derive(Debug)]
 pub struct PythonLogicalCodec {
     inner: Arc<dyn LogicalExtensionCodec>,
+    python_udf_inlining: bool,
 }
 
 impl PythonLogicalCodec {
     pub fn new(inner: Arc<dyn LogicalExtensionCodec>) -> Self {
-        Self { inner }
+        Self {
+            inner,
+            python_udf_inlining: true,
+        }
     }
 
     pub fn inner(&self) -> &Arc<dyn LogicalExtensionCodec> {
         &self.inner
     }
+
+    /// Whether Python-defined UDFs are encoded inline (and decoded
+    /// from cloudpickle blobs). Defaults to `true`. Set to `false`
+    /// when the codec sits on a session that must produce
+    /// cross-language wire bytes, or reject `cloudpickle.loads` on
+    /// untrusted `from_bytes` input.
+    ///
+    /// Security scope: strict mode (`false`) narrows only the codec
+    /// layer — it stops `Expr::from_bytes` from invoking
+    /// `cloudpickle.loads` on the inline `DFPY*` payload. It does
+    /// **not** make `pickle.loads(untrusted_bytes)` safe; treat every
+    /// `pickle.loads` on untrusted input as unsafe regardless of this
+    /// setting. See Python's [pickle module security warning][1] for
+    /// why `pickle.loads` is unsafe in general.
+    ///
+    /// [1]: https://docs.python.org/3/library/pickle.html#module-pickle
+    pub fn with_python_udf_inlining(mut self, enabled: bool) -> Self {
+        self.python_udf_inlining = enabled;
+        self
+    }
+
+    pub fn python_udf_inlining(&self) -> bool {
+        self.python_udf_inlining
+    }
 }
 
 impl Default for PythonLogicalCodec {
@@ -301,48 +329,76 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
     }
 
     fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
-        if try_encode_python_scalar_udf(node, buf)? {
+        if self.python_udf_inlining && 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);
+        if self.python_udf_inlining {
+            if let Some(udf) = try_decode_python_scalar_udf(buf)? {
+                return Ok(udf);
+            }
+        } else if buf.starts_with(PY_SCALAR_UDF_FAMILY) {
+            return Err(refuse_inline_payload("scalar UDF", name));
         }
         self.inner.try_decode_udf(name, buf)
     }
 
     fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> {
-        if try_encode_python_udaf(node, buf)? {
+        if self.python_udf_inlining && try_encode_python_udaf(node, buf)? {
             return Ok(());
         }
         self.inner.try_encode_udaf(node, buf)
     }
 
     fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> {
-        if let Some(udaf) = try_decode_python_udaf(buf)? {
-            return Ok(udaf);
+        if self.python_udf_inlining {
+            if let Some(udaf) = try_decode_python_udaf(buf)? {
+                return Ok(udaf);
+            }
+        } else if buf.starts_with(PY_AGG_UDF_FAMILY) {
+            return Err(refuse_inline_payload("aggregate UDF", name));
         }
         self.inner.try_decode_udaf(name, buf)
     }
 
     fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> {
-        if try_encode_python_udwf(node, buf)? {
+        if self.python_udf_inlining && try_encode_python_udwf(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_udwf(buf)? {
-            return Ok(udwf);
+        if self.python_udf_inlining {
+            if let Some(udwf) = try_decode_python_udwf(buf)? {
+                return Ok(udwf);
+            }
+        } else if buf.starts_with(PY_WINDOW_UDF_FAMILY) {
+            return Err(refuse_inline_payload("window UDF", name));
         }
         self.inner.try_decode_udwf(name, buf)
     }
 }
 
+/// Build the error returned by a strict codec when it receives an
+/// inline Python-UDF payload it has been told not to deserialize.
+fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusionError {
+    // `Execution`, not `Plan`: this is a wire-format decode refusal at
+    // codec time, not a planner-stage failure. Downstream error
+    // classification keys off the variant — surfacing this as a planner
+    // error would mis-route it into "fix your SQL" buckets.
+    datafusion::error::DataFusionError::Execution(format!(
+        "Refusing to deserialize inline Python {kind} '{name}': Python UDF \
+         inlining is disabled on this session. Ask the sender to re-encode \
+         with inlining disabled (so the UDF travels by name), or register \
+         '{name}' on this receiver's session and enable inlining on both \
+         sides — receivers cannot re-encode bytes they did not produce."
+    ))
+}
+
 /// `PhysicalExtensionCodec` mirror of [`PythonLogicalCodec`] parked
 /// on the same `SessionContext`. Carries the Python-aware encoding
 /// hooks for physical-layer types (`ExecutionPlan`, `PhysicalExpr`)
@@ -358,16 +414,30 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
 #[derive(Debug)]
 pub struct PythonPhysicalCodec {
     inner: Arc<dyn PhysicalExtensionCodec>,
+    python_udf_inlining: bool,
 }
 
 impl PythonPhysicalCodec {
     pub fn new(inner: Arc<dyn PhysicalExtensionCodec>) -> Self {
-        Self { inner }
+        Self {
+            inner,
+            python_udf_inlining: true,
+        }
     }
 
     pub fn inner(&self) -> &Arc<dyn PhysicalExtensionCodec> {
         &self.inner
     }
+
+    /// See [`PythonLogicalCodec::with_python_udf_inlining`].
+    pub fn with_python_udf_inlining(mut self, enabled: bool) -> Self {
+        self.python_udf_inlining = enabled;
+        self
+    }
+
+    pub fn python_udf_inlining(&self) -> bool {
+        self.python_udf_inlining
+    }
 }
 
 impl Default for PythonPhysicalCodec {
@@ -391,15 +461,19 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
     }
 
     fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
-        if try_encode_python_scalar_udf(node, buf)? {
+        if self.python_udf_inlining && 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);
+        if self.python_udf_inlining {
+            if let Some(udf) = try_decode_python_scalar_udf(buf)? {
+                return Ok(udf);
+            }
+        } else if buf.starts_with(PY_SCALAR_UDF_FAMILY) {
+            return Err(refuse_inline_payload("scalar UDF", name));
         }
         self.inner.try_decode_udf(name, buf)
     }
@@ -417,29 +491,37 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
     }
 
     fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> {
-        if try_encode_python_udaf(node, buf)? {
+        if self.python_udf_inlining && try_encode_python_udaf(node, buf)? {
             return Ok(());
         }
         self.inner.try_encode_udaf(node, buf)
     }
 
     fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> {
-        if let Some(udaf) = try_decode_python_udaf(buf)? {
-            return Ok(udaf);
+        if self.python_udf_inlining {
+            if let Some(udaf) = try_decode_python_udaf(buf)? {
+                return Ok(udaf);
+            }
+        } else if buf.starts_with(PY_AGG_UDF_FAMILY) {
+            return Err(refuse_inline_payload("aggregate UDF", name));
         }
         self.inner.try_decode_udaf(name, buf)
     }
 
     fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> {
-        if try_encode_python_udwf(node, buf)? {
+        if self.python_udf_inlining && try_encode_python_udwf(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_udwf(buf)? {
-            return Ok(udwf);
+        if self.python_udf_inlining {
+            if let Some(udwf) = try_decode_python_udwf(buf)? {
+                return Ok(udwf);
+            }
+        } else if buf.starts_with(PY_WINDOW_UDF_FAMILY) {
+            return Err(refuse_inline_payload("window UDF", name));
         }
         self.inner.try_decode_udwf(name, buf)
     }

crates/core/src/context.rs

@@ -1404,6 +1404,22 @@ impl PySessionContext {
             physical_codec,
         })
     }
+
+    pub fn with_python_udf_inlining(&self, enabled: bool) -> Self {
+        let logical_codec = Arc::new(
+            PythonLogicalCodec::new(Arc::clone(self.logical_codec.inner()))
+                .with_python_udf_inlining(enabled),
+        );
+        let physical_codec = Arc::new(
+            PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner()))
+                .with_python_udf_inlining(enabled),
+        );
+        Self {
+            ctx: Arc::clone(&self.ctx),
+            logical_codec,
+            physical_codec,
+        }
+    }
 }
 
 impl PySessionContext {

python/datafusion/context.py

@@ -1769,3 +1769,46 @@ def with_physical_extension_codec(self, codec: Any) -> SessionContext:
         new = SessionContext.__new__(SessionContext)
         new.ctx = new_internal
         return new
+
+    def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext:
+        """Toggle inline encoding of Python-defined UDFs on this session.
+
+        ``enabled`` is keyword-only:
+        ``with_python_udf_inlining(enabled=False)`` reads at the call
+        site as the inverse of
+        ``with_python_udf_inlining(enabled=True)``, where a positional
+        ``True`` / ``False`` would not.
+
+        When ``True`` (the default), Python scalar, aggregate, and window
+        UDFs travel inside the serialized expression and are
+        reconstructed on the receiver without pre-registration.
+
+        Set ``False`` to:
+
+        * Produce serialized bytes that round-trip through a non-Python
+          decoder (cross-language portability). UDFs are stored by name
+          only; the receiver must have matching registrations.
+        * Refuse to reconstruct Python UDFs from
+          :meth:`Expr.from_bytes` input that may come from an untrusted
+          source — ``cloudpickle.loads`` will not be invoked.
+
+        The toggle applies directly to :meth:`Expr.to_bytes` /
+        :meth:`Expr.from_bytes` calls that pass this session as their
+        ``ctx`` argument. To make the toggle apply through
+        :func:`pickle.dumps` (which calls :meth:`Expr.to_bytes` with no
+        context), install this session as the driver's sender context
+        via :func:`datafusion.ipc.set_sender_ctx` — and install it as
+        the worker's context via
+        :func:`datafusion.ipc.set_worker_ctx` for the corresponding
+        :func:`pickle.loads`.
+
+        For the full security model, see
+        :doc:`/user-guide/io/distributing_work` (Security section). In
+        short: this toggle narrows only the :meth:`Expr.from_bytes`
+        surface; :func:`pickle.loads` on untrusted bytes remains
+        unsafe regardless of the toggle.
+        """
+        new_internal = self.ctx.with_python_udf_inlining(enabled)
+        new = SessionContext.__new__(SessionContext)
+        new.ctx = new_internal
+        return new

python/datafusion/expr.py

@@ -446,13 +446,16 @@ def to_bytes(self, ctx: SessionContext | None = None) -> bytes:
         worker process for distributed evaluation.
 
         When ``ctx`` is supplied, encoding routes through that session's
-        installed :class:`LogicalExtensionCodec`. When ``ctx`` is
-        ``None``, the default codec is used.
+        installed :class:`LogicalExtensionCodec` (so settings like
+        :meth:`SessionContext.with_python_udf_inlining` take effect).
+        When ``ctx`` is ``None``, the default codec is used (Python UDF
+        inlining on, no user-installed extension codec).
 
         Built-in functions and Python UDFs (scalar, aggregate, window)
         travel inside the returned bytes; the worker does not need to
         pre-register them. UDFs imported via the FFI capsule protocol
-        travel by name only and must be registered on the worker.
+        travel by name only and must be registered on the worker. See
+        :doc:`/user-guide/io/distributing_work`.
 
         .. warning:: Security
             Bytes returned here may embed a cloudpickled Python
@@ -522,7 +525,9 @@ def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr:
 
         Accepts output of :meth:`to_bytes` or :func:`pickle.dumps`.
         ``ctx`` is the :class:`SessionContext` used to resolve any
-        function references that travel by name (e.g. FFI UDFs). When
+        function references that travel by name (e.g. FFI UDFs, or
+        Python UDFs sent with inlining disabled via
+        :meth:`SessionContext.with_python_udf_inlining`). When
         ``ctx`` is ``None`` the worker context installed via
         :func:`datafusion.ipc.set_worker_ctx` is consulted; if no worker
         context is installed, the global :class:`SessionContext` is used
@@ -586,8 +591,15 @@ def __reduce__(self) -> tuple[Callable[[bytes], Expr], tuple[bytes]]:
             >>> e = col("a") * lit(2)
             >>> pickle.loads(pickle.dumps(e)).canonical_name()
             'a * Int64(2)'
+
+        The encoding side honors a driver-side sender context installed
+        via :func:`datafusion.ipc.set_sender_ctx` — that is how
+        :meth:`SessionContext.with_python_udf_inlining` propagates
+        through ``pickle.dumps``.
         """
-        return (Expr._reconstruct, (self.to_bytes(),))
+        from datafusion.ipc import get_sender_ctx
+
+        return (Expr._reconstruct, (self.to_bytes(get_sender_ctx()),))
 
     @classmethod
     def _reconstruct(cls, proto_bytes: bytes) -> Expr: