feat: Python UDFs: per-session inlining toggle and strict refusal setting

.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,39 @@ 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
     }
+
+    /// Toggle inline encoding of Python UDFs. See
+    /// `SessionContext.with_python_udf_inlining` (Python) for full
+    /// behavior and use cases.
+    ///
+    /// 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.
+    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 +324,104 @@ 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 {
+            refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "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 {
+            refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "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 {
+            refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?;
         }
         self.inner.try_decode_udwf(name, buf)
     }
 }
 
+/// Strict-mode gate: if `buf` is a well-framed inline payload for
+/// `family`, return the strict-refusal error; otherwise return
+/// `Ok(())` so the caller can delegate to its `inner` codec.
+///
+/// Routing through [`read_framed_payload`] (rather than a bare
+/// `starts_with` probe) means malformed inline bytes — wrong
+/// wire-format version, mismatched Python version, truncated header —
+/// surface *their* diagnostic instead of the strict-mode message.
+/// The strict message implies sender intent ("inlining is disabled"),
+/// so it should fire only when the bytes really would have decoded.
+///
+/// Fast path: short-circuit on the family-magic prefix before
+/// acquiring the GIL. Plans with many non-Python UDFs would otherwise
+/// pay a GIL acquisition per decode call just to confirm "not a
+/// Python UDF". `read_framed_payload` itself rejects buffers that
+/// don't start with `family`, so this is purely an optimization.
+fn refuse_if_inline(buf: &[u8], family: &[u8], kind: &str, name: &str) -> Result<()> {
+    if !buf.starts_with(family) {
+        return Ok(());
+    }
+    Python::attach(|py| match read_framed_payload(py, buf, family, kind)? {
+        Some(_) => Err(refuse_inline_payload(kind, name)),
+        None => Ok(()),
+    })
+}
+
+/// 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. Two remediations: \
+         (1) ask the sender to re-encode with inlining disabled so '{name}' \
+         travels by name, and register '{name}' on this receiver; or \
+         (2) enable inlining on this receiver (accepts the cloudpickle \
+         execution risk on inbound payloads). 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 +437,33 @@ 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
     }
+
+    /// Toggle inline encoding of Python UDFs on this physical codec.
+    ///
+    /// Mirrors [`PythonLogicalCodec::with_python_udf_inlining`]; see
+    /// that method for the full security and portability discussion.
+    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 +487,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 {
+            refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?;
         }
         self.inner.try_decode_udf(name, buf)
     }
@@ -417,29 +517,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 {
+            refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "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 {
+            refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?;
         }
         self.inner.try_decode_udwf(name, buf)
     }
@@ -476,6 +584,9 @@ pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec<u8>)
 /// 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.starts_with(PY_SCALAR_UDF_FAMILY) {
+        return Ok(None);
+    }
     Python::attach(|py| -> Result<Option<Arc<ScalarUDF>>> {
         let Some(payload) = read_framed_payload(py, buf, PY_SCALAR_UDF_FAMILY, "scalar UDF")?
         else {
@@ -732,6 +843,9 @@ pub(crate) fn try_encode_python_udwf(node: &WindowUDF, buf: &mut Vec<u8>) -> Res
 }
 
 pub(crate) fn try_decode_python_udwf(buf: &[u8]) -> Result<Option<Arc<WindowUDF>>> {
+    if !buf.starts_with(PY_WINDOW_UDF_FAMILY) {
+        return Ok(None);
+    }
     Python::attach(|py| -> Result<Option<Arc<WindowUDF>>> {
         let Some(payload) = read_framed_payload(py, buf, PY_WINDOW_UDF_FAMILY, "window UDF")?
         else {
@@ -814,6 +928,9 @@ pub(crate) fn try_encode_python_udaf(node: &AggregateUDF, buf: &mut Vec<u8>) ->
 }
 
 pub(crate) fn try_decode_python_udaf(buf: &[u8]) -> Result<Option<Arc<AggregateUDF>>> {
+    if !buf.starts_with(PY_AGG_UDF_FAMILY) {
+        return Ok(None);
+    }
     Python::attach(|py| -> Result<Option<Arc<AggregateUDF>>> {
         let Some(payload) = read_framed_payload(py, buf, PY_AGG_UDF_FAMILY, "aggregate UDF")?
         else {

crates/core/src/context.rs

@@ -1453,6 +1453,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 {

pyproject.toml

@@ -203,6 +203,7 @@ dev = [
   "pyarrow>=19.0.0",
   "pygithub==2.5.0",
   "pytest-asyncio>=0.23.3",
+  "pytest-timeout>=2.3.1",
   "pytest>=7.4.4",
   "pyyaml>=6.0.3",
   "ruff>=0.15.1",

python/datafusion/context.py

@@ -1963,3 +1963,61 @@ def with_physical_extension_codec(
         new = SessionContext.__new__(SessionContext)
         new.ctx = new_internal
         return new
+
+    def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext:
+        """Control whether Python UDFs are embedded in serialized expressions.
+
+        ``enabled`` is keyword-only and required: callers must pick a
+        mode explicitly. Fresh sessions inline UDFs (``enabled=True``
+        behavior) until this method overrides the toggle.
+
+        With ``enabled=True``, serialized expressions carry the Python
+        code for any scalar, aggregate, or window UDFs they reference.
+        The receiver rebuilds the UDFs from those bytes and does not
+        need to register them first.
+
+        With ``enabled=False``, serialized expressions store only the
+        UDF names. This has two uses:
+
+        * **Cross-language portability.** The bytes can be decoded by a
+          non-Python receiver, which must already have UDFs registered
+          under matching names.
+        * **Safer deserialization.** :meth:`Expr.from_bytes` will refuse
+          to rebuild Python UDFs rather than call ``cloudpickle.loads``
+          on untrusted input.
+
+        The setting affects :meth:`Expr.to_bytes` and
+        :meth:`Expr.from_bytes` whenever this session is passed as the
+        ``ctx`` argument. :func:`pickle.dumps` and :func:`pickle.loads`
+        do not pass a context, so to apply the setting through pickle,
+        register this session with
+        :func:`datafusion.ipc.set_sender_ctx` on the sender and
+        :func:`datafusion.ipc.set_worker_ctx` on the receiver.
+
+        .. warning:: Security
+            This setting narrows only :meth:`Expr.from_bytes`. Calling
+            :func:`pickle.loads` on untrusted bytes remains unsafe
+            regardless of the toggle.
+
+        Returns a new :class:`SessionContext` with the toggle applied;
+        the original session is unchanged.
+
+        Examples:
+            >>> import pyarrow as pa
+            >>> from datafusion import SessionContext, Expr, col, udf
+            >>> ctx = SessionContext()
+            >>> identity = udf(lambda a: a, [pa.int64()], pa.int64(),
+            ...                volatility="immutable", name="identity_demo")
+            >>> ctx.register_udf(identity)
+            >>> blob = identity(col("x")).to_bytes(ctx)
+            >>> strict = SessionContext().with_python_udf_inlining(enabled=False)
+            >>> try:
+            ...     Expr.from_bytes(blob, strict)
+            ... except Exception as e:
+            ...     print("Refusing to deserialize" in str(e))
+            True
+        """
+        new_internal = self.ctx.with_python_udf_inlining(enabled)
+        new = SessionContext.__new__(SessionContext)
+        new.ctx = new_internal
+        return new