Make extension codecs composable

crates/core/src/context.rs

@@ -1472,7 +1472,9 @@ impl PySessionContext {
     ) -> PyDataFusionResult<Self> {
         let inner_ffi = ffi_logical_codec_from_pycapsule(codec)?;
         let inner: Arc<dyn LogicalExtensionCodec> = (&inner_ffi).into();
-        let logical_codec = Arc::new(PythonLogicalCodec::new(inner));
+        // Prepend rather than replace: previously installed codecs stay
+        // active, with the most recently installed one consulted first.
+        let logical_codec = Arc::new(self.logical_codec.with_additional_codec(inner));
 
         let physical_codec = Arc::clone(&self.physical_codec);
         let ctx = self
@@ -1497,7 +1499,9 @@ impl PySessionContext {
         codec: Bound<'py, PyAny>,
     ) -> PyDataFusionResult<Self> {
         let inner = physical_codec_from_pycapsule(&codec)?;
-        let physical_codec = Arc::new(PythonPhysicalCodec::new(inner));
+        // Prepend rather than replace: previously installed codecs stay
+        // active, with the most recently installed one consulted first.
+        let physical_codec = Arc::new(self.physical_codec.with_additional_codec(inner));
 
         let logical_codec = Arc::clone(&self.logical_codec);
         let ctx = self
@@ -1511,11 +1515,15 @@ impl PySessionContext {
 
     pub fn with_python_udf_inlining(&self, enabled: bool) -> Self {
         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 ctx = self

docs/source/contributor-guide/ffi.md

@@ -248,15 +248,65 @@ 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.
 
 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
 local build commands.
 
+### One planner per session, with explicit fallback
+
+Unlike codecs, a `SessionState` holds exactly one query planner — installing another
+replaces it. Planner layering is therefore explicit: a planner that wants to handle
+only some queries should accept a fallback planner and delegate the rest to it. The
+current planner can be exported for that purpose with
+`ctx.__datafusion_query_planner__()`.
+
+One ordering rule applies: a planner capsule captures the session's codecs at export
+time and cannot be rebound afterward. Codec changes made after installing a single
+planner are rebound automatically, but a planner wrapped inside another planner as a
+fallback is opaque and keeps the codecs it was exported with. **Install all extension
+codecs before exporting or chaining planners.**
+
+Putting it together for a session using two extension libraries that each provide
+tables, functions, and a query planner:
+
+```python
+ctx = SessionContext(config)
+
+# 1. 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())
+
+# 2. Planners, innermost fallback first. Library A's planner falls back to
+#    DataFusion's default planner; library B's planner falls back to A's.
+ctx = ctx.with_query_planner(lib_a.Planner())
+ctx = ctx.with_query_planner(
+    lib_b.Planner(fallback=ctx.__datafusion_query_planner__())
+)
+
+# 3. Tables and functions — any time before the first query.
+ctx.register_table("t", lib_a.TableProvider())
+ctx.register_udf(udf(lib_b.SomeUDF()))
+```
+
 ## Alternative Approach
 
 Suppose you needed to expose some other features of DataFusion and you could not wait

examples/datafusion-ffi-example/README.md

@@ -35,7 +35,7 @@ Separate shared libraries guarantee distinct DataFusion library markers. This ca
 
 The example codecs do not inspect the callback `TaskContext`. A production codec that depends on session configuration or registered functions must ensure its exported FFI codec is bound to, and retains, the appropriate host `TaskContextProvider`.
 
-The current Python API installs one external logical codec and one external physical codec. It does not yet compose codecs from several independent plugin owners. This example therefore makes the provider library the sole external codec owner; the planner uses built-in physical nodes and receives the provider codecs from the host.
+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). In this example the provider library is the only 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:
 

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

@@ -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_physical_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-query-planner-example/README.md

@@ -55,6 +55,6 @@ ctx = ctx.with_query_planner(MyQueryPlanner())
 
 `PlannerConfig` 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 the derived context is created and is also used to decode the returned physical plan in `datafusion-python`. The API currently supports one external codec owner rather than a registry of independently composed codecs, so this planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; derived contexts rebind codecs after planner installation, but planner-last order is easier to audit.
+The provider's codec pair is attached to the planner when the derived context is created 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. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install codecs before the planner; derived contexts rebind codecs after a planner is installed directly, but a planner exported as a fallback for another planner keeps the codecs captured at export time.
 
 The pinned FFI logical codec cannot encode arbitrary custom `LogicalPlan::Extension` nodes. The example therefore demonstrates table-provider, UDF, and physical-plan interoperability without claiming custom logical extension support.

examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py

@@ -118,3 +118,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 = ctx.with_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

python/datafusion/context.py

@@ -1779,6 +1779,15 @@ def with_query_planner(
         its logical and physical extension codec settings. Codec changes made on
         a derived context are rebound to the planner before planning.
 
+        A session holds exactly one query planner; installing another replaces
+        it. To layer planners, construct the new planner with the current
+        planner as its fallback (export it via
+        :py:meth:`__datafusion_query_planner__`) before installing. A planner
+        exported this way captures the codecs installed at export time and
+        cannot be rebound afterward, so install all extension codecs before
+        chaining planners. See the FFI extensions guide for the full
+        multi-library registration recipe.
+
         Args:
             planner: Object exposing ``__datafusion_query_planner__`` or a raw
                 ``datafusion_query_planner`` PyCapsule.
@@ -2229,11 +2238,19 @@ def __datafusion_query_planner__(self) -> Any:
     def with_logical_extension_codec(
         self, codec: LogicalExtensionCodecExportable | _PyCapsule
     ) -> SessionContext:
-        """Create a new session context with specified codec.
+        """Create a new session context with an additional logical codec.
 
         Only FFI codecs are supported. Pass any object implementing
         ``__datafusion_logical_extension_codec__`` (see
         :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`).
+
+        Codecs compose: each call 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. Codecs signal "not mine" by returning an error, so
+        extension codecs should only answer for payloads they own —
+        typically identified by a distinct byte prefix.
         """
         new_internal = self.ctx.with_logical_extension_codec(codec)
         new = SessionContext.__new__(SessionContext)
@@ -2247,11 +2264,16 @@ def __datafusion_physical_extension_codec__(self) -> Any:
     def with_physical_extension_codec(
         self, codec: PhysicalExtensionCodecExportable | _PyCapsule
     ) -> SessionContext:
-        """Create a new session context with the specified physical codec.
+        """Create a new session context with an additional physical codec.
 
         Only FFI codecs are supported. Pass any object implementing
         ``__datafusion_physical_extension_codec__`` (see
         :py:class:`~datafusion.user_defined.PhysicalExtensionCodecExportable`).
+
+        Codecs compose the same way as in
+        :py:meth:`with_logical_extension_codec`: each call prepends to the
+        session's codec chain, and the most recently installed codec is
+        consulted first.
         """
         new_internal = self.ctx.with_physical_extension_codec(codec)
         new = SessionContext.__new__(SessionContext)