fix: keep the UDF inlining setting when installing a codec

`with_logical_extension_codec` and `with_physical_extension_codec` built
the replacement wrapper with `Python{Logical,Physical}Codec::new`, whose
constructor defaults `python_udf_inlining` to true. Installing a codec on
a context that had opted out therefore turned inlining back on without
saying so:

    ctx = SessionContext().with_python_udf_inlining(enabled=False)
    ctx = ctx.with_logical_extension_codec(codec)   # inlining silently back on

That matters beyond a stale flag. Inlining is what embeds a cloudpickled
callable in the wire format, and it is opt-out precisely because that is
not portable across interpreters and not something every deployment wants
to ship. A codec install is not a request to change it.

Carry the receiver's setting across instead. Both new tests fail on the
prior build with `DFPYUDF` reappearing in the blob, and the paired
`..._preserves_inlining_when_enabled` case pins the default-on direction so
the fix cannot degenerate into hard-coding it off.

The physical case is covered in `test_plans.py` rather than alongside the
logical one: `Expr.to_bytes` only routes through the logical codec, so an
assertion there would pass with the physical bug still present. It takes an
`ExecutionPlan` to observe.

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

Author: timsaucer

crates/core/src/context.rs

@@ -1445,9 +1445,16 @@ impl PySessionContext {
     ) -> PyDataFusionResult<Self> {
         let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?;
         let inner: Arc<dyn LogicalExtensionCodec> = (&inner_ffi).into();
-        let logical_codec = Arc::new(PythonLogicalCodec::new(inner));
 
         let this = slf.borrow();
+        // Carry the receiver's inlining setting over. `PythonLogicalCodec::new`
+        // defaults it to on, so building the replacement without this would
+        // silently re-enable inline Python UDF encoding on a context that had
+        // opted out with `with_python_udf_inlining(enabled=False)`.
+        let logical_codec = Arc::new(
+            PythonLogicalCodec::new(inner)
+                .with_python_udf_inlining(this.logical_codec.python_udf_inlining()),
+        );
         let derived = Self {
             ctx: Arc::clone(&this.ctx),
             logical_codec,
@@ -1476,9 +1483,13 @@ impl PySessionContext {
     ) -> PyDataFusionResult<Self> {
         let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?;
         let inner: Arc<dyn PhysicalExtensionCodec> = (&inner_ffi).into();
-        let physical_codec = Arc::new(PythonPhysicalCodec::new(inner));
 
         let this = slf.borrow();
+        // See `with_logical_extension_codec` for why the flag is carried over.
+        let physical_codec = Arc::new(
+            PythonPhysicalCodec::new(inner)
+                .with_python_udf_inlining(this.physical_codec.python_udf_inlining()),
+        );
         let derived = Self {
             ctx: Arc::clone(&this.ctx),
             logical_codec: Arc::clone(&this.logical_codec),

python/tests/test_pickle_expr.py

@@ -387,6 +387,45 @@ def test_toggle_off_then_on_restores_inline_encoding(self):
         decoded = Expr.from_bytes(blob_toggled, ctx=SessionContext())
         assert "double" in decoded.canonical_name()
 
+    def test_installing_a_codec_preserves_strict_mode(self):
+        """Installing an extension codec must not re-enable inlining.
+
+        `with_{logical,physical}_extension_codec` builds a replacement
+        `Python{Logical,Physical}Codec` around the imported one, and the
+        constructor defaults inlining to on. A context that opted out has to
+        keep its setting, or a codec install silently starts shipping
+        cloudpickled callables again.
+
+        The session's own capsule getters are a convenient stand-in for a real
+        extension codec here: they return a genuine FFI codec, so the import
+        path under test is the same one an extension library exercises.
+
+        This covers the logical codec. The physical one is checked in
+        ``test_plans.py``, since it takes an ``ExecutionPlan`` to observe.
+        """
+        strict = SessionContext().with_python_udf_inlining(enabled=False)
+        e = self._build_double_udf()(col("a"))
+        assert b"DFPYUDF" not in e.to_bytes(strict)
+
+        with_logical = strict.with_logical_extension_codec(
+            strict.__datafusion_logical_extension_codec__()
+        )
+        assert b"DFPYUDF" not in e.to_bytes(with_logical)
+
+    def test_installing_a_codec_preserves_inlining_when_enabled(self):
+        """The converse: the default stays on across a codec install.
+
+        Guards against 'fixing' the above by hard-coding inlining off.
+        """
+        ctx = SessionContext()
+        e = self._build_double_udf()(col("a"))
+        assert b"DFPYUDF" in e.to_bytes(ctx)
+
+        installed = ctx.with_logical_extension_codec(
+            ctx.__datafusion_logical_extension_codec__()
+        )
+        assert b"DFPYUDF" in e.to_bytes(installed)
+
     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

python/tests/test_plans.py

@@ -17,13 +17,16 @@
 
 import datetime
 
+import pyarrow as pa
 import pytest
 from datafusion import (
     ExecutionPlan,
     LogicalPlan,
     Metric,
     MetricsSet,
     SessionContext,
+    col,
+    udf,
 )
 
 
@@ -92,6 +95,46 @@ def test_session_with_logical_extension_codec_roundtrip(ctx, df) -> None:
     assert df.collect() == df_round_trip.collect()
 
 
+def test_installing_a_physical_codec_preserves_strict_mode() -> None:
+    """Installing a physical extension codec must not re-enable inlining.
+
+    `with_physical_extension_codec` builds a replacement `PythonPhysicalCodec`
+    around the imported one, and the constructor defaults inlining to on. A
+    context that opted out via `with_python_udf_inlining(enabled=False)` has to
+    keep its setting, or installing a codec silently starts embedding
+    cloudpickled callables in serialized execution plans.
+
+    The logical counterpart lives in `test_pickle_expr.py`; this one needs an
+    `ExecutionPlan` because only the physical codec encodes it. `DFPYUDF` is
+    the scalar Python-UDF family prefix, shared by both layers; see
+    `PY_SCALAR_UDF_FAMILY` in crates/core/src/codec.rs.
+    """
+    identity = udf(
+        lambda arr: arr,
+        [pa.string()],
+        pa.string(),
+        volatility="immutable",
+        name="identity_str",
+    )
+
+    def plan_bytes(ctx: SessionContext) -> bytes:
+        df = ctx.read_csv(path="testing/data/csv/aggregate_test_100.csv").select(
+            identity(col("c1"))
+        )
+        return df.execution_plan().to_bytes(ctx)
+
+    # The inlining default is on, so the strict blob is what has to differ.
+    assert b"DFPYUDF" in plan_bytes(SessionContext())
+
+    strict = SessionContext().with_python_udf_inlining(enabled=False)
+    assert b"DFPYUDF" not in plan_bytes(strict)
+
+    installed = strict.with_physical_extension_codec(
+        strict.__datafusion_physical_extension_codec__()
+    )
+    assert b"DFPYUDF" not in plan_bytes(installed)
+
+
 def test_session_codec_capsule_getters(ctx) -> None:
     """SessionContext exposes both logical and physical codec capsules."""
     logical = ctx.ctx.__datafusion_logical_extension_codec__()