Add three-library FFI planner example

AI Disclosure: This code was written in part by an AI agent.:

Author: timsaucer

.github/workflows/build.yml

@@ -186,7 +186,7 @@ jobs:
           manylinux: "2_28"
 
       # FFI test wheel only needs to be built once per platform; gate to abi3.
-      - name: Build FFI test library
+      - name: Build FFI provider test library
         if: matrix.python-tag == 'abi3'
         uses: PyO3/maturin-action@v1
         with:
@@ -196,6 +196,16 @@ jobs:
           args: --out dist
           rustup-components: rust-std
 
+      - name: Build FFI query planner test library
+        if: matrix.python-tag == 'abi3'
+        uses: PyO3/maturin-action@v1
+        with:
+          target: x86_64-unknown-linux-gnu
+          manylinux: "2_28"
+          working-directory: examples/datafusion-ffi-query-planner-example
+          args: --out dist
+          rustup-components: rust-std
+
       - name: Archive wheels
         uses: actions/upload-artifact@v7
         with:
@@ -207,7 +217,9 @@ jobs:
         uses: actions/upload-artifact@v7
         with:
           name: test-ffi-manylinux-x86_64
-          path: examples/datafusion-ffi-example/dist/*
+          path: |
+            examples/datafusion-ffi-example/dist/*
+            examples/datafusion-ffi-query-planner-example/dist/*
 
   # ============================================
   # Build - Linux ARM64

.github/workflows/test.yml

@@ -121,6 +121,8 @@ jobs:
         run: |
           cd examples/datafusion-ffi-example
           uv run --no-project pytest python/tests/_test*.py
+          cd ../datafusion-ffi-query-planner-example
+          uv run --no-project pytest python/tests/_test*.py
 
       - name: Run tpchgen-cli to create 1 Gb dataset
         if: matrix.wheel-tag == 'abi3'

Cargo.lock

@@ -27,7 +27,12 @@ edition = "2024"
 rust-version = "1.88"
 
 [workspace]
-members = ["crates/core", "crates/util", "examples/datafusion-ffi-example"]
+members = [
+  "crates/core",
+  "crates/util",
+  "examples/datafusion-ffi-example",
+  "examples/datafusion-ffi-query-planner-example",
+]
 resolver = "3"
 
 [workspace.dependencies]

Cargo.toml

@@ -227,9 +227,34 @@ impl PySessionConfig {
     }
 }
 
+#[derive(Debug)]
+struct PlanningTaskContextProvider(Arc<TaskContext>);
+
+impl TaskContextProvider for PlanningTaskContextProvider {
+    fn task_ctx(&self) -> Arc<TaskContext> {
+        Arc::clone(&self.0)
+    }
+}
+
 #[derive(Debug, Clone)]
 struct PythonQueryPlanner {
     planner: FFI_QueryPlanner,
+    logical_codec: Arc<PythonLogicalCodec>,
+    physical_codec: Arc<PythonPhysicalCodec>,
+}
+
+impl PythonQueryPlanner {
+    fn with_codecs(
+        &self,
+        logical_codec: Arc<PythonLogicalCodec>,
+        physical_codec: Arc<PythonPhysicalCodec>,
+    ) -> Self {
+        Self {
+            planner: self.planner.clone(),
+            logical_codec,
+            physical_codec,
+        }
+    }
 }
 
 #[async_trait]
@@ -240,9 +265,35 @@ impl QueryPlanner for PythonQueryPlanner {
         session: &dyn Session,
     ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
         let runtime = get_tokio_runtime().handle().clone();
-        self.planner
+        let ctx_provider = Arc::new(PlanningTaskContextProvider(session.task_ctx()));
+        let (logical_codec, physical_codec) = {
+            let dyn_ctx_provider: Arc<dyn TaskContextProvider> = ctx_provider.clone();
+            let logical_codec: Arc<dyn LogicalExtensionCodec> =
+                Arc::clone(&self.logical_codec) as Arc<dyn LogicalExtensionCodec>;
+            let physical_codec: Arc<dyn PhysicalExtensionCodec + Send> =
+                Arc::clone(&self.physical_codec) as Arc<dyn PhysicalExtensionCodec + Send>;
+            (
+                FFI_LogicalExtensionCodec::new(
+                    logical_codec,
+                    Some(runtime.clone()),
+                    &dyn_ctx_provider,
+                ),
+                FFI_PhysicalExtensionCodec::new(
+                    physical_codec,
+                    Some(runtime.clone()),
+                    &dyn_ctx_provider,
+                ),
+            )
+        };
+
+        let mut planner = self.planner.clone();
+        planner.logical_codec = logical_codec;
+        planner.physical_codec = physical_codec;
+        let result = planner
             .create_physical_plan_with_session_runtime(logical_plan, session, Some(runtime))
-            .await
+            .await;
+        drop(ctx_provider);
+        result
     }
 }
 
@@ -1237,10 +1288,12 @@ impl PySessionContext {
     }
 
     pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
-        let mut planner = ffi_query_planner_from_pycapsule(&planner)?;
-        planner.logical_codec = self.ffi_logical_codec().as_ref().clone();
-        planner.physical_codec = self.ffi_physical_codec().as_ref().clone();
-        let planner = Arc::new(PythonQueryPlanner { planner });
+        let planner = ffi_query_planner_from_pycapsule(&planner)?;
+        let planner = Arc::new(PythonQueryPlanner {
+            planner,
+            logical_codec: Arc::clone(&self.logical_codec),
+            physical_codec: Arc::clone(&self.physical_codec),
+        });
         let state = SessionStateBuilder::new_from_existing(self.ctx.state())
             .with_query_planner(planner)
             .build();
@@ -1448,10 +1501,13 @@ impl PySessionContext {
         let inner: Arc<dyn LogicalExtensionCodec> = (&inner_ffi).into();
         let logical_codec = Arc::new(PythonLogicalCodec::new(inner));
 
+        let physical_codec = Arc::clone(&self.physical_codec);
+        let ctx = self
+            .ctx_with_query_planner_codecs(Arc::clone(&logical_codec), Arc::clone(&physical_codec));
         Ok(Self {
-            ctx: Arc::clone(&self.ctx),
+            ctx,
             logical_codec,
-            physical_codec: Arc::clone(&self.physical_codec),
+            physical_codec,
         })
     }
 
@@ -1470,9 +1526,12 @@ impl PySessionContext {
         let inner = physical_codec_from_pycapsule(&codec)?;
         let physical_codec = Arc::new(PythonPhysicalCodec::new(inner));
 
+        let logical_codec = Arc::clone(&self.logical_codec);
+        let ctx = self
+            .ctx_with_query_planner_codecs(Arc::clone(&logical_codec), Arc::clone(&physical_codec));
         Ok(Self {
-            ctx: Arc::clone(&self.ctx),
-            logical_codec: Arc::clone(&self.logical_codec),
+            ctx,
+            logical_codec,
             physical_codec,
         })
     }
@@ -1486,15 +1545,36 @@ impl PySessionContext {
             PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner()))
                 .with_python_udf_inlining(enabled),
         );
+        let ctx = self
+            .ctx_with_query_planner_codecs(Arc::clone(&logical_codec), Arc::clone(&physical_codec));
         Self {
-            ctx: Arc::clone(&self.ctx),
+            ctx,
             logical_codec,
             physical_codec,
         }
     }
 }
 
 impl PySessionContext {
+    fn ctx_with_query_planner_codecs(
+        &self,
+        logical_codec: Arc<PythonLogicalCodec>,
+        physical_codec: Arc<PythonPhysicalCodec>,
+    ) -> Arc<SessionContext> {
+        let state = self.ctx.state();
+        let query_planner = state.query_planner();
+        let planner_any: &dyn std::any::Any = query_planner.as_ref();
+        let Some(planner) = planner_any.downcast_ref::<PythonQueryPlanner>() else {
+            return Arc::clone(&self.ctx);
+        };
+
+        let planner = Arc::new(planner.with_codecs(logical_codec, physical_codec));
+        let state = SessionStateBuilder::new_from_existing(self.ctx.state())
+            .with_query_planner(planner)
+            .build();
+        Arc::new(SessionContext::new_with_state(state))
+    }
+
     async fn _table(&self, name: &str) -> datafusion::common::Result<DataFrame> {
         self.ctx.table(name).await
     }

crates/core/src/context.rs

@@ -232,6 +232,31 @@ extension that has been written using this approach and the most thoroughly impl
 As we continue to expose more of the DataFusion features, we intend to follow this same
 design pattern.
 
+## Query Planners Across Multiple Libraries
+
+A query can involve three independent native libraries: `datafusion-python`, a library
+that owns table providers or functions, and a library that owns the query planner. The
+examples use two separate extension crates so each role has a distinct shared-library
+identity:
+
+- [`datafusion-ffi-example`] owns providers, functions, and their codecs.
+- [`datafusion-ffi-query-planner-example`] owns the planner and its configuration.
+
+The `SessionContext` owns the codecs used for the exchange and supplies them to the
+foreign planner. This lets the planner decode provider-owned objects and lets
+`datafusion-python` decode the physical plan returned by the planner. The examples use
+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.
+
+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.
+
 ## Alternative Approach
 
 Suppose you needed to expose some other features of DataFusion and you could not wait
@@ -257,3 +282,5 @@ At the time of this writing, the FFI features are under active development. To s
 the latest status, we recommend reviewing the code in the [datafusion-ffi] crate.
 
 [datafusion-ffi]: https://crates.io/crates/datafusion-ffi
+[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example
+[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example

docs/source/contributor-guide/ffi.md

@@ -49,6 +49,15 @@ Here is a direct link to the file used in the examples:
 - [Fan out distinct expressions to a multiprocessing pool](./multiprocessing_pickle_expr.py)
 - [Distribute expression evaluation across Ray actors](./ray_pickle_expr.py)
 
+### Rust FFI Extensions
+
+- [Table providers, functions, and codecs](./datafusion-ffi-example/)
+- [Independent query planner and planner configuration](./datafusion-ffi-query-planner-example/)
+
+These two crates form a three-library interoperability example with
+`datafusion-python`. They are separate shared libraries so the tests exercise real FFI
+type and codec boundaries rather than same-library Rust downcasts.
+
 ### Substrait Support
 
 - [Serialize query plans using Substrait](./substrait.py)

examples/README.md

@@ -34,7 +34,6 @@ datafusion-functions-window = { workspace = true }
 datafusion-expr = { workspace = true }
 datafusion-ffi = { workspace = true }
 datafusion-proto = { workspace = true }
-datafusion-session = { workspace = true }
 
 arrow = { workspace = true }
 arrow-array = { workspace = true }

examples/datafusion-ffi-example/Cargo.toml

@@ -0,0 +1,31 @@
+# DataFusion Python FFI provider example
+
+This crate is the **provider library** in the three-library query-planning example. It exports table providers, functions, and the logical and physical codecs needed to serialize objects owned by this library. The companion planner is in [`../datafusion-ffi-query-planner-example`](../datafusion-ffi-query-planner-example/).
+
+The example intentionally uses separate `cdylib` crates for these roles:
+
+1. **A — `datafusion-python`:** owns the `SessionContext` and executes the result.
+2. **B — this crate:** owns table providers, functions, and provider execution plans.
+3. **C — the planner crate:** receives the logical plan and returns a physical plan.
+
+Separate shared libraries guarantee distinct DataFusion library markers. This catches type-identity mistakes that a planner and provider compiled into one shared library would hide.
+
+## Codec behavior
+
+`MyLogicalExtensionCodec` serializes this example's in-memory table providers, and `MyPhysicalExtensionCodec` serializes provider-owned memory scans and opaque FFI wrappers around them. Both use documented, process-local, one-shot token registries. The registries make ownership and callback routing visible without pretending to be a portable format. They assume trusted in-process payloads and consume each token during decoding. A production provider should instead encode durable metadata from which its provider and plans can be reconstructed.
+
+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.
+
+Register both provider codecs before installing the planner:
+
+```python
+ctx = ctx.with_logical_extension_codec(provider_logical_codec)
+ctx = ctx.with_physical_extension_codec(provider_physical_codec)
+ctx = ctx.with_query_planner(planner)
+```
+
+Derived contexts also rebind an installed planner when codecs change, but planner-last order is recommended because it states the ownership flow clearly.
+
+Arbitrary custom `LogicalPlan::Extension` nodes are not supported by the current DataFusion FFI logical codec. This example covers foreign table providers, UDFs, and physical execution plans only.

examples/datafusion-ffi-example/README.md

@@ -23,7 +23,6 @@ use crate::config::MyConfig;
 use crate::logical_extension_codec::MyLogicalExtensionCodec;
 use crate::physical_extension_codec::MyPhysicalExtensionCodec;
 use crate::physical_optimizer::MyPhysicalOptimizerRule;
-use crate::query_planner::MyQueryPlanner;
 use crate::scalar_udf::IsNullUDF;
 use crate::table_function::MyTableFunction;
 use crate::table_provider::MyTableProvider;
@@ -36,7 +35,6 @@ pub(crate) mod config;
 pub(crate) mod logical_extension_codec;
 pub(crate) mod physical_extension_codec;
 pub(crate) mod physical_optimizer;
-pub(crate) mod query_planner;
 pub(crate) mod scalar_udf;
 pub(crate) mod table_function;
 pub(crate) mod table_provider;
@@ -60,6 +58,5 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MyLogicalExtensionCodec>()?;
     m.add_class::<MyPhysicalExtensionCodec>()?;
     m.add_class::<MyPhysicalOptimizerRule>()?;
-    m.add_class::<MyQueryPlanner>()?;
     Ok(())
 }