Make extension codecs composable

Installing a logical or physical extension codec now prepends it to a
codec chain instead of replacing the prior codec. The most recently
installed codec is consulted first, falling through codec by codec to
the default codec. This lets multiple independent extension libraries
install codecs on the same session, and removes the codec registration
ordering requirement between libraries.

Chain dispatch treats a codec error as "not mine". Encoding runs each
codec against a scratch buffer so failed attempts leave no partial
bytes, and treats Ok-with-no-bytes (encode by name) as no opinion so
later codecs still get a chance. When every codec fails, the errors
are aggregated so the owning codec's diagnostic is not masked by the
default codec's generic error.

Also preserves the python_udf_inlining setting when installing a
codec; previously it was silently reset to enabled.

Documents the remaining planner constraint: a session holds one query
planner, layering is explicit via fallback capsules, and codecs must
be installed before exporting or chaining planners because a planner
capsule captures the codecs at export time.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: timsaucer

crates/core/src/codec.rs

@@ -1465,14 +1465,13 @@ impl PySessionContext {
         let inner: Arc<dyn LogicalExtensionCodec> = (&inner_ffi).into();
 
         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()),
-        );
+        // Prepend rather than replace: previously installed codecs stay active,
+        // with the most recently installed one consulted first. Prepending also
+        // carries the receiver's inlining setting over, which a fresh
+        // `PythonLogicalCodec::new` would not — it defaults inlining to on, and
+        // 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(this.logical_codec.with_additional_codec(inner));
         let derived = Self {
             ctx: Arc::clone(&this.ctx),
             logical_codec,
@@ -1503,11 +1502,9 @@ impl PySessionContext {
         let inner: Arc<dyn PhysicalExtensionCodec> = (&inner_ffi).into();
 
         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()),
-        );
+        // See `with_logical_extension_codec` for why this prepends rather than
+        // replaces, and why that is also what carries the inlining flag over.
+        let physical_codec = Arc::new(this.physical_codec.with_additional_codec(inner));
         let derived = Self {
             ctx: Arc::clone(&this.ctx),
             logical_codec: Arc::clone(&this.logical_codec),
@@ -1525,7 +1522,7 @@ impl PySessionContext {
         // already inlines would otherwise rebind the session's planner to this
         // handle's codecs, and callers routinely discard the result. Returning
         // the codecs as-is is observationally equivalent to the rebuild below,
-        // which wraps the same inner codec in a fresh `Python*Codec`.
+        // which clones the same codec chain and only flips the flag.
         if self.logical_codec.python_udf_inlining() == enabled
             && self.physical_codec.python_udf_inlining() == enabled
         {
@@ -1537,11 +1534,15 @@ impl PySessionContext {
         }
 
         let logical_codec = Arc::new(
-            PythonLogicalCodec::new(Arc::clone(self.logical_codec.inner()))
+            self.logical_codec
+                .as_ref()
+                .clone()
                 .with_python_udf_inlining(enabled),
         );
         let physical_codec = Arc::new(
-            PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner()))
+            self.physical_codec
+                .as_ref()
+                .clone()
                 .with_python_udf_inlining(enabled),
         );
         let derived = Self {

crates/core/src/context.rs

@@ -248,10 +248,43 @@ foreign planner. This lets the planner decode provider-owned objects and lets
 process-local tokens to demonstrate ownership; production codecs should serialize
 durable metadata instead.
 
-The current Python API has one external logical codec and one external physical codec.
-Installing another codec replaces the prior codec rather than composing a registry.
-The example therefore has one external codec owner, and the planner uses built-in
-physical nodes. Install the provider codecs before the planner where possible.
+### Composable codecs
+
+Extension codecs compose. Each call to `with_logical_extension_codec` or
+`with_physical_extension_codec` adds the codec to the front of the session's codec
+chain rather than replacing prior codecs. During encoding and decoding, the most
+recently installed codec is consulted first, falling through codec by codec to
+DataFusion's default codec. A codec signals "not mine" by returning an error, which
+sends the chain on to the next codec. Two conventions keep this dispatch sound:
+
+- Frame your payloads with a distinct byte prefix (pick a `DF` namespace plus a
+  crate-specific suffix) and only decode payloads carrying your prefix.
+- Return an error for objects and payloads you do not own. A codec that answers
+  success for objects outside its family shadows every codec installed before it.
+
+Because dispatch keys off payload prefixes rather than install position, codec
+registration order between independent libraries does not matter. Two libraries that
+each own tables, functions, and a planner register like this:
+
+```python
+ctx = SessionContext(config)
+
+# Codecs from both libraries. Order between libraries does not matter.
+ctx = ctx.with_logical_extension_codec(lib_a.codec())
+ctx = ctx.with_logical_extension_codec(lib_b.codec())
+ctx = ctx.with_physical_extension_codec(lib_a.physical_codec())
+ctx = ctx.with_physical_extension_codec(lib_b.physical_codec())
+
+# A session holds one planner, so layering is explicit delegation. Install the
+# codecs first: the fallback captured here keeps the codecs it was exported
+# with. See "Rebinding a planner's codecs is one level deep" below.
+ctx.set_query_planner(lib_a.Planner())
+ctx.set_query_planner(lib_b.Planner(fallback=ctx.__datafusion_query_planner__()))
+
+# Tables and functions — any time before the first query.
+ctx.register_table("t", lib_a.TableProvider())
+ctx.register_udf(udf(lib_b.SomeUDF()))
+```
 
 The current FFI logical codec supports providers and UDFs but not arbitrary custom
 `LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and
@@ -347,7 +380,7 @@ side is visible to both.
 so it is a property of the session rather than of a handle on it, and installing one is
 visible to every context sharing that session — including ones a `with_*` call returned
 earlier. Installing a codec on a session that already has a foreign planner rebuilds
-that planner against the new codec for the same reason: there is one planner, and it has
+that planner against the new chain for the same reason: there is one planner, and it has
 to carry the codecs currently in force. This happens on the shared session, so it takes
 effect even if the returned context is discarded — `ctx.with_python_udf_inlining(...)`
 whose result is thrown away still leaves the session's planner carrying the codecs of
@@ -361,15 +394,17 @@ that surprises people:
 > installed one. Every other path — `Expr.to_bytes(ctx)`, `ExecutionPlan.to_bytes(ctx)`,
 > registering a provider — uses the codecs of the handle you call it on.
 
-Those can be different handles, and then one session has two codecs in effect at once:
+Those can be different handles, and then one session has two codec chains in effect at
+once:
 
 ```python
 ctx = ctx.with_logical_extension_codec(codec_a)
 ctx.set_query_planner(planner)
 ctx.with_logical_extension_codec(codec_b)  # discarded
 
-Expr.to_bytes(expr, ctx)   # encodes with codec_a -- ctx's own field
-ctx.sql(...).collect()     # plans with codec_b -- installed via the discarded handle
+Expr.to_bytes(expr, ctx)   # encodes with [codec_a, default] -- ctx's own field
+ctx.sql(...).collect()     # plans with [codec_b, codec_a, default] -- the discarded
+                           # handle's chain, installed on the shared session
 ```
 
 Chaining `ctx = ctx.with_...(...)`, as the example below does, keeps the two in step.

docs/source/contributor-guide/ffi.md

@@ -35,7 +35,9 @@ Separate shared libraries guarantee distinct DataFusion library markers. This ca
 
 Both codec getters take the `SessionContext` they are being installed on and pull the `TaskContextProvider` off it, so decode callbacks resolve session configuration and registered functions against the session that is running the query. Passing `require_udf_on_decode` to either constructor makes every decode call resolve a named scalar function out of that context, which is how the tests check where the registry came from.
 
-This example makes the provider library the sole external codec owner. Register both provider codecs before installing the planner:
+Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call prepends the codec to the session's codec chain, with the most recently installed codec consulted first and DataFusion's default codec as the terminal fallback. A codec signals "not mine" by returning an error, so several independent plugin libraries can install codecs on the same session as long as each only answers for payloads it owns (frame them with a distinct byte prefix). This example keeps the provider library as the sole external codec owner; the planner uses built-in physical nodes and receives the provider codecs from the host.
+
+Register both provider codecs before installing the planner:
 
 ```python
 ctx = ctx.with_logical_extension_codec(provider_logical_codec)
@@ -45,4 +47,4 @@ ctx.set_query_planner(planner)
 
 Installing a codec after the planner rebuilds the planner against it, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly. The exception is a planner that wraps a fallback: the rebuild reaches the installed planner only, not the fallback inside it, so codecs-first is a requirement there. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep), which also covers why re-installing a planner rebinds the session to the codecs of whichever handle it was installed on.
 
-For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.
+For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.

examples/datafusion-ffi-example/README.md

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 from datafusion import LogicalPlan, SessionContext
-from datafusion_ffi_example import MyLogicalExtensionCodec
+from datafusion_ffi_example import MyLogicalExtensionCodec, MyTableProvider
 
 
 def _setup_session_with_codec() -> tuple[SessionContext, MyLogicalExtensionCodec]:
@@ -80,3 +80,28 @@ def test_ffi_logical_codec_roundtrip():
     restored = LogicalPlan.from_bytes(ctx, blob)
     df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
     assert df.collect() == df_round_trip.collect()
+
+
+def test_ffi_logical_codec_composes_with_later_install():
+    """Codecs compose: installing a second codec prepends it to the
+    session's codec chain instead of replacing the first. The second
+    codec here (a default-backed codec exported from a fresh session)
+    cannot encode this library's table provider, so encoding falls
+    through to the user codec installed first. Under replace semantics
+    this test fails with `LogicalExtensionCodec is not provided`."""
+    ctx, codec = _setup_session_with_codec()
+    ctx = ctx.with_logical_extension_codec(
+        SessionContext().__datafusion_logical_extension_codec__()
+    )
+
+    ctx.register_table("numbers", MyTableProvider(1, 4, 1))
+    df = ctx.sql('SELECT "A" FROM numbers')
+    plan = df.logical_plan()
+
+    before = codec.table_provider_encode_calls()
+    blob = plan.to_bytes(ctx)
+    assert codec.table_provider_encode_calls() > before
+
+    restored = LogicalPlan.from_bytes(ctx, blob)
+    df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
+    assert df.collect() == df_round_trip.collect()

examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py

@@ -76,3 +76,26 @@ def test_ffi_physical_codec_roundtrip():
 
     restored = ExecutionPlan.from_bytes(ctx, blob)
     assert str(original) == str(restored)
+
+
+def test_ffi_physical_codec_composes_with_later_install():
+    """Codecs compose: a second install prepends to the chain instead
+    of replacing the first codec. The second codec here (default-backed
+    export from a fresh session) encodes UDFs by name without writing
+    bytes, which the chain treats as "no opinion" — so the user codec
+    installed first is still consulted. Under replace semantics its
+    counter stays at zero."""
+    ctx, codec = _setup_session_with_codec()
+    ctx = ctx.with_physical_extension_codec(
+        SessionContext().__datafusion_physical_extension_codec__()
+    )
+
+    df = ctx.sql("SELECT abs(a) AS x FROM t")
+    original = df.execution_plan()
+
+    before = codec.encode_udf_calls()
+    blob = original.to_bytes(ctx)
+    assert codec.encode_udf_calls() > before
+
+    restored = ExecutionPlan.from_bytes(ctx, blob)
+    assert str(original) == str(restored)

examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py

@@ -55,6 +55,6 @@ ctx.set_query_planner(MyQueryPlanner())
 
 `MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit.
 
-The provider's codec pair is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. This planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against it, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep).
+The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call prepends to the session's codec chain, so several libraries can install codecs on the same session and the order between them does not matter. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep).
 
-For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.
+For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.

examples/datafusion-ffi-query-planner-example/README.md

@@ -667,3 +667,26 @@ def test_query_planner_rejects_invalid_config(max_rows: str):
 
     with pytest.raises(Exception, match=r"max_rows|Invalid value"):
         ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect()
+
+
+def test_composed_codecs_with_query_planner():
+    """A second pair of codecs installed on top of the provider codecs
+    composes with them instead of replacing them. The extra codecs
+    (default-backed exports from a fresh session) decline everything,
+    so planner-driven encode/decode falls through to the provider
+    codecs and the query still succeeds end to end."""
+    ctx, logical_codec, physical_codec = configured_context(max_rows=2)
+    other = SessionContext()
+    ctx = ctx.with_logical_extension_codec(
+        other.__datafusion_logical_extension_codec__()
+    )
+    ctx = ctx.with_physical_extension_codec(
+        other.__datafusion_physical_extension_codec__()
+    )
+    ctx.set_query_planner(MyQueryPlanner())
+
+    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+    assert batches[0].column(0).to_pylist() == [0, 1]
+    assert logical_codec.table_provider_encode_calls() > 0
+    assert logical_codec.table_provider_decode_calls() > 0
+    assert physical_codec.execution_plan_decode_calls() > 0