feat: per-session toggle for Python UDF inline encoding

Adds `SessionContext.with_python_udf_inlining(enabled)` for two
related use cases:

* **Cross-language portability.** With inlining disabled, the codec
  no longer emits `DFPYUDF1` / `DFPYUDA1` / `DFPYUDW1` cloudpickle
  blobs. Python UDFs travel by name only, the same way FFI-capsule
  UDFs do. Bytes round-trip through a non-Python decoder.
* **Untrusted-source decode.** `Expr.from_bytes` on bytes from a
  misbehaving sender no longer invokes `cloudpickle.loads`. Inline
  payloads received by a strict decoder raise a clear error.

`PythonLogicalCodec` and `PythonPhysicalCodec` gain a
`python_udf_inlining: bool` field (default `true`) and a builder
method `with_python_udf_inlining(enabled)`. The six UDF
encode/decode dispatchers consult the flag before calling the inline
helpers. Strict decoders that see a magic-prefix payload return a
clear `Plan` error rather than silently failing through to the inner
codec (which would otherwise produce "LogicalExtensionCodec is not
provided" — accurate but unhelpful).

`PySessionContext::with_python_udf_inlining(enabled)` rebuilds both
codecs with the new setting; Python wrapper at
`SessionContext.with_python_udf_inlining` mirrors. Test coverage:
encoder size delta, strict roundtrip via registry,
clear-error-on-inline-payload-when-strict.

`pickle.loads` on untrusted bytes remains unsafe regardless of this
flag; the toggle only governs the `to_bytes` / `from_bytes` codec
path. User guide documents both use cases plus the limitation.

1097 root tests pass (up from 1094 with 3 new strict-mode cases).

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

Author: timsaucer

crates/core/src/codec.rs

@@ -128,16 +128,34 @@ pub(crate) const PY_WINDOW_UDF_MAGIC: &[u8] = b"DFPYUDW1";
 #[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.
+    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 {
@@ -197,48 +215,72 @@ 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_MAGIC) {
+            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_agg_udf(node, buf)? {
+        if self.python_udf_inlining && try_encode_python_agg_udf(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_agg_udf(buf)? {
-            return Ok(udaf);
+        if self.python_udf_inlining {
+            if let Some(udaf) = try_decode_python_agg_udf(buf)? {
+                return Ok(udaf);
+            }
+        } else if buf.starts_with(PY_AGG_UDF_MAGIC) {
+            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_window_udf(node, buf)? {
+        if self.python_udf_inlining && 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);
+        if self.python_udf_inlining {
+            if let Some(udwf) = try_decode_python_window_udf(buf)? {
+                return Ok(udwf);
+            }
+        } else if buf.starts_with(PY_WINDOW_UDF_MAGIC) {
+            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 {
+    datafusion::error::DataFusionError::Plan(format!(
+        "Refusing to deserialize inline Python {kind} '{name}': Python UDF \
+         inlining is disabled on this session. Re-encode the bytes with \
+         inlining enabled, or register '{name}' on the sender's session \
+         before encode (and on the receiver before decode) so the UDF \
+         travels by name."
+    ))
+}
+
 /// `PhysicalExtensionCodec` mirror of [`PythonLogicalCodec`] parked
 /// on the same `SessionContext`. Carries the Python-aware encoding
 /// hooks for physical-layer types (`ExecutionPlan`, `PhysicalExpr`)
@@ -254,16 +296,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 {
@@ -287,15 +343,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_MAGIC) {
+            return Err(refuse_inline_payload("scalar UDF", name));
         }
         self.inner.try_decode_udf(name, buf)
     }
@@ -313,29 +373,37 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
     }
 
     fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> {
-        if try_encode_python_agg_udf(node, buf)? {
+        if self.python_udf_inlining && try_encode_python_agg_udf(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_agg_udf(buf)? {
-            return Ok(udaf);
+        if self.python_udf_inlining {
+            if let Some(udaf) = try_decode_python_agg_udf(buf)? {
+                return Ok(udaf);
+            }
+        } else if buf.starts_with(PY_AGG_UDF_MAGIC) {
+            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_window_udf(node, buf)? {
+        if self.python_udf_inlining && 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);
+        if self.python_udf_inlining {
+            if let Some(udwf) = try_decode_python_window_udf(buf)? {
+                return Ok(udwf);
+            }
+        } else if buf.starts_with(PY_WINDOW_UDF_MAGIC) {
+            return Err(refuse_inline_payload("window UDF", name));
         }
         self.inner.try_decode_udwf(name, buf)
     }

crates/core/src/context.rs

@@ -1407,6 +1407,28 @@ impl PySessionContext {
             physical_codec,
         })
     }
+
+    /// Toggle inline encoding of Python-defined UDFs on this session's
+    /// codec stack. Disable when producing bytes that must round-trip
+    /// through a non-Python decoder, or when reconstructing bytes from
+    /// an untrusted source via `Expr.from_bytes` (cloudpickle.loads
+    /// will not be invoked on the receiver). Pickle remains unsafe on
+    /// untrusted input regardless of this flag.
+    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 {

docs/source/user-guide/io/distributing_work.rst

@@ -173,6 +173,37 @@ Practical considerations
   the captured state is large, mutable, or not portable to the
   worker's environment.
 
+Disabling Python UDF inlining
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+For a stricter wire format, call
+:py:meth:`SessionContext.with_python_udf_inlining(False)
+<datafusion.SessionContext.with_python_udf_inlining>` on the session
+producing or consuming the bytes. With inlining disabled, Python
+UDFs travel by name only — the same way FFI-capsule UDFs do — and
+the receiver must have a matching registration.
+
+Two use cases:
+
+* **Cross-language portability.** A non-Python decoder cannot
+  reconstruct a cloudpickled payload. Senders aimed at Java, C++,
+  or another Rust binary disable inlining and rely on the receiver
+  having compatible UDF registrations.
+* **Untrusted-source decode.** With inlining disabled,
+  :py:meth:`Expr.from_bytes` never calls ``cloudpickle.loads`` on
+  the incoming bytes — an inline payload from a misbehaving sender
+  raises a clear error instead of executing arbitrary Python code.
+
+Mismatched configurations raise a descriptive error: an inline blob
+fed to a strict receiver fails fast rather than silently dropping
+into ``cloudpickle.loads``.
+
+Note that :py:func:`pickle.loads` itself remains unsafe on untrusted
+input regardless of this setting — an attacker producing the outer
+pickle envelope can execute arbitrary code before the codec ever
+sees the bytes. The toggle only protects the
+:py:meth:`Expr.from_bytes` API surface.
+
 Security
 ~~~~~~~~
 
@@ -182,8 +213,10 @@ Security
    arbitrary Python code on the receiver — pickle is doing the work
    under the hood and pickle is unsafe on untrusted input. Only
    accept expressions from trusted sources. For untrusted-source
-   workflows, restrict senders to built-in functions and
-   pre-registered Rust-side UDFs.
+   workflows, disable Python UDF inlining (see above), restrict
+   senders to built-in functions and pre-registered Rust-side UDFs,
+   and avoid :py:func:`pickle.loads` on externally supplied bytes
+   entirely.
 
 Query-level distribution via datafusion-distributed
 ---------------------------------------------------

python/datafusion/context.py

@@ -1769,3 +1769,28 @@ 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.
+
+        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.
+
+        ``pickle.loads`` on untrusted bytes remains unsafe regardless of
+        this setting; only the ``to_bytes`` / ``from_bytes`` API is
+        affected.
+        """
+        new_internal = self.ctx.with_python_udf_inlining(enabled)
+        new = SessionContext.__new__(SessionContext)
+        new.ctx = new_internal
+        return new

python/tests/test_pickle_expr.py

@@ -229,6 +229,62 @@ def test_window_udf_decodes_via_pickle_with_no_worker_ctx(self):
         assert "count_up" in decoded.canonical_name()
 
 
+class TestPythonUdfInliningToggle:
+    """`SessionContext.with_python_udf_inlining(False)` opts out of
+    inline Python UDF encoding for both encode and decode paths."""
+
+    def _build_double_udf(self):
+        return udf(
+            lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]),
+            [pa.int64()],
+            pa.int64(),
+            volatility="immutable",
+            name="double",
+        )
+
+    def test_strict_encoder_emits_smaller_blob(self):
+        """Strict mode skips cloudpickle of the Python callable, so the
+        encoded bytes are dramatically smaller than the inline form."""
+        ctx_inline = SessionContext()
+        ctx_strict = ctx_inline.with_python_udf_inlining(False)
+        u = self._build_double_udf()
+        e = u(col("a"))
+
+        blob_inline = e.to_bytes(ctx_inline)
+        blob_strict = e.to_bytes(ctx_strict)
+
+        assert len(blob_strict) < len(blob_inline) // 4
+
+    def test_strict_roundtrip_via_registry(self):
+        """When both sender and receiver disable inlining, the UDF
+        travels by name only and the receiver resolves it from its
+        registered functions."""
+        from datafusion import Expr
+
+        strict_sender = SessionContext().with_python_udf_inlining(False)
+        u = self._build_double_udf()
+        blob = u(col("a")).to_bytes(strict_sender)
+
+        receiver = SessionContext().with_python_udf_inlining(False)
+        receiver.register_udf(u)
+        restored = Expr.from_bytes(blob, ctx=receiver)
+        assert "double" in restored.canonical_name()
+
+    def test_strict_decoder_refuses_inline_payload(self):
+        """An inline-encoded blob fed to a strict receiver raises with a
+        clear error rather than silently invoking cloudpickle.loads."""
+        from datafusion import Expr
+
+        sender = SessionContext()
+        u = self._build_double_udf()
+        blob = u(col("a")).to_bytes(sender)
+
+        strict_receiver = SessionContext().with_python_udf_inlining(False)
+        strict_receiver.register_udf(u)
+        with pytest.raises(Exception, match="inlining is disabled"):
+            Expr.from_bytes(blob, ctx=strict_receiver)
+
+
 class TestWorkerCtxLifecycle:
     def test_set_and_clear(self):
         assert get_worker_ctx() is None