refactor: address review nits for UDF inlining toggle + sender ctx

* codec: strict refusal routes through `read_framed_payload` so
  malformed inline bytes surface their own diagnostic; the
  "inlining is disabled" message now fires only when the payload
  would have decoded.
* codec: add summary line above `PythonPhysicalCodec::with_python_udf_inlining`
  cross-link for rustdoc rendering.
* expr: hoist `get_sender_ctx` import to module top; note that
  `__reduce__` also drives `copy.copy` / `copy.deepcopy`.
* context: accept `with_python_udf_inlining` positionally or as
  kwarg (drop `*,`).
* tests: replace size-ratio heuristic with semantic check for the
  `DFPYUDF` family prefix; switch single-batch closure test to
  `pool.apply`.

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

Author: timsaucer

crates/core/src/codec.rs

@@ -340,8 +340,8 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
             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));
+        } else {
+            refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?;
         }
         self.inner.try_decode_udf(name, buf)
     }
@@ -358,8 +358,8 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
             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));
+        } else {
+            refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?;
         }
         self.inner.try_decode_udaf(name, buf)
     }
@@ -376,13 +376,30 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
             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));
+        } 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.
+fn refuse_if_inline(buf: &[u8], family: &[u8], kind: &str, name: &str) -> Result<()> {
+    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 {
@@ -429,7 +446,10 @@ impl PythonPhysicalCodec {
         &self.inner
     }
 
-    /// See [`PythonLogicalCodec::with_python_udf_inlining`].
+    /// 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
@@ -472,8 +492,8 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
             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));
+        } else {
+            refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?;
         }
         self.inner.try_decode_udf(name, buf)
     }
@@ -502,8 +522,8 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
             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));
+        } else {
+            refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?;
         }
         self.inner.try_decode_udaf(name, buf)
     }
@@ -520,8 +540,8 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
             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));
+        } else {
+            refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?;
         }
         self.inner.try_decode_udwf(name, buf)
     }

python/datafusion/context.py

@@ -1770,7 +1770,7 @@ def with_physical_extension_codec(self, codec: Any) -> SessionContext:
         new.ctx = new_internal
         return new
 
-    def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext:
+    def with_python_udf_inlining(self, enabled: bool) -> SessionContext:
         """Control whether Python UDFs are embedded in serialized expressions.
 
         When ``enabled=True`` (the default), serialized expressions carry

python/datafusion/expr.py

@@ -53,6 +53,7 @@
 
 from ._internal import expr as expr_internal
 from ._internal import functions as functions_internal
+from .ipc import get_sender_ctx
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -594,10 +595,10 @@ def __reduce__(self) -> tuple[Callable[[bytes], Expr], tuple[bytes]]:
         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``.
+        through ``pickle.dumps``. The sender context is read by
+        ``__reduce__``, so :func:`copy.copy` and :func:`copy.deepcopy`
+        — which also go through ``__reduce__`` — pick it up too.
         """
-        from datafusion.ipc import get_sender_ctx
-
         return (Expr._reconstruct, (self.to_bytes(get_sender_ctx()),))
 
     @classmethod

python/tests/test_pickle_expr.py

@@ -337,9 +337,13 @@ def _build_double_udf(self):
             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."""
+    def test_strict_encoder_omits_inline_payload(self):
+        """Strict mode emits the by-name wire form: no `DFPYUDF` magic
+        in the blob, no cloudpickled callable. Semantic check is
+        sharper than a size-ratio heuristic — a renamed UDF or a
+        smaller-than-expected closure would still flip the magic
+        bytes, but might not move the size by 4x.
+        """
         ctx_inline = SessionContext()
         ctx_strict = ctx_inline.with_python_udf_inlining(enabled=False)
         u = self._build_double_udf()
@@ -348,7 +352,10 @@ def test_strict_encoder_emits_smaller_blob(self):
         blob_inline = e.to_bytes(ctx_inline)
         blob_strict = e.to_bytes(ctx_strict)
 
-        assert len(blob_strict) < len(blob_inline) // 4
+        # `DFPYUDF` is the scalar Python-UDF family prefix; see
+        # `PY_SCALAR_UDF_FAMILY` in crates/core/src/codec.rs.
+        assert b"DFPYUDF" in blob_inline
+        assert b"DFPYUDF" not in blob_strict
 
     def test_toggle_off_then_on_restores_inline_encoding(self):
         """`with_python_udf_inlining` is per-call clone semantics:
@@ -417,10 +424,10 @@ def test_strict_decoder_refuses_inline_payload(self):
     def test_sender_ctx_propagates_through_pickle(self):
         """`set_sender_ctx` makes `pickle.dumps` use a strict codec.
 
-        Without a sender context, pickle defaults to the inline codec and
-        the blob is large. With a strict sender context installed, the
-        blob shrinks because the Python callable is encoded by name
-        instead of cloudpickled.
+        Without a sender context, pickle defaults to the inline codec
+        and the blob contains the `DFPYUDF` family prefix. With a
+        strict sender context installed, the callable encodes by name
+        and the prefix is absent.
         """
         u = self._build_double_udf()
         e = u(col("a"))
@@ -434,7 +441,8 @@ def test_sender_ctx_propagates_through_pickle(self):
         finally:
             clear_sender_ctx()
 
-        assert len(blob_strict) < len(blob_default) // 4
+        assert b"DFPYUDF" in blob_default
+        assert b"DFPYUDF" not in blob_strict
 
     def test_sender_ctx_strict_roundtrip_via_pickle(self):
         """End-to-end pickle round-trip with strict mode on both sides.

python/tests/test_pickle_multiprocessing.py

@@ -123,9 +123,6 @@ def test_closure_capturing_udf_via_pool(start_method):
 
     ctx = mp.get_context(start_method)
     with ctx.Pool(processes=2) as pool:
-        results = pool.starmap(
-            helpers.unpickle_and_evaluate,
-            [(blob, [1, 2, 3])],
-        )
+        result = pool.apply(helpers.unpickle_and_evaluate, (blob, [1, 2, 3]))
 
-    assert results[0] == [7, 14, 21]
+    assert result == [7, 14, 21]