Add FFI query planner support

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

Author: timsaucer

Cargo.lock

@@ -50,6 +50,7 @@ datafusion-functions-aggregate = { version = "54.1.0" }
 datafusion-functions-window = { version = "54.1.0" }
 datafusion-spark = { version = "54.1.0" }
 datafusion-expr = { version = "54.1.0" }
+datafusion-session = { version = "54.1.0" }
 prost = "0.14.3"
 serde_json = "1"
 uuid = { version = "1.23" }
@@ -72,13 +73,14 @@ codegen-units = 2
 # We cannot publish to crates.io with any patches in the below section. Developers
 # must remove any entries in this section before creating a release candidate.
 [patch.crates-io]
-datafusion = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
-datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
-datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
-datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
-datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
-datafusion-common = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
-datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
-datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
-datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
-datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-substrait = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-proto = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-ffi = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-catalog = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-common = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-functions-aggregate = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-functions-window = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-spark = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-expr = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }
+datafusion-session = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" }

Cargo.toml

@@ -54,6 +54,7 @@ datafusion-substrait = { workspace = true, optional = true }
 datafusion-proto = { workspace = true }
 datafusion-ffi = { workspace = true }
 datafusion-spark = { workspace = true, features = ["core"] }
+datafusion-session = { workspace = true }
 prost = { workspace = true } # keep in line with `datafusion-substrait`
 serde_json = { workspace = true }
 uuid = { workspace = true, features = ["v4"] }

crates/core/Cargo.toml

@@ -24,6 +24,7 @@ use std::sync::Arc;
 use arrow::array::RecordBatchReader;
 use arrow::ffi_stream::ArrowArrayStreamReader;
 use arrow::pyarrow::FromPyArrow;
+use async_trait::async_trait;
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
@@ -36,14 +37,16 @@ use datafusion::datasource::listing::{
 };
 use datafusion::datasource::{MemTable, TableProvider};
 use datafusion::execution::context::{
-    DataFilePaths, SQLOptions, SessionConfig, SessionContext, TaskContext,
+    DataFilePaths, QueryPlanner, SQLOptions, SessionConfig, SessionContext, TaskContext,
 };
 use datafusion::execution::disk_manager::DiskManagerMode;
 use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, UnboundedMemoryPool};
 use datafusion::execution::options::{ArrowReadOptions, ReadOptions};
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::execution::session_state::SessionStateBuilder;
 use datafusion::execution::{FunctionRegistry, TaskContextProvider};
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::{
     AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions,
 };
@@ -53,15 +56,18 @@ use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
 use datafusion_ffi::execution::FFI_TaskContextProvider;
 use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use datafusion_ffi::query_planner::FFI_QueryPlanner;
 use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory;
 use datafusion_proto::logical_plan::LogicalExtensionCodec;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_python_util::{
     create_logical_extension_capsule, create_physical_extension_capsule,
-    ffi_logical_codec_from_pycapsule, get_global_ctx, get_tokio_runtime,
+    create_query_planner_capsule, ffi_logical_codec_from_pycapsule,
+    ffi_query_planner_from_pycapsule, get_global_ctx, get_tokio_runtime,
     physical_codec_from_pycapsule, physical_optimizer_rule_from_pycapsule, spawn_future,
     wait_for_future,
 };
+use datafusion_session::Session;
 use object_store::ObjectStore;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError};
@@ -221,6 +227,25 @@ impl PySessionConfig {
     }
 }
 
+#[derive(Debug, Clone)]
+struct PythonQueryPlanner {
+    planner: FFI_QueryPlanner,
+}
+
+#[async_trait]
+impl QueryPlanner for PythonQueryPlanner {
+    async fn create_physical_plan(
+        &self,
+        logical_plan: &LogicalPlan,
+        session: &dyn Session,
+    ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+        let runtime = get_tokio_runtime().handle().clone();
+        self.planner
+            .create_physical_plan_with_session_runtime(logical_plan, session, Some(runtime))
+            .await
+    }
+}
+
 /// Runtime options for a SessionContext
 #[pyclass(
     from_py_object,
@@ -1211,6 +1236,23 @@ impl PySessionContext {
         Ok(())
     }
 
+    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 state = SessionStateBuilder::new_from_existing(self.ctx.state())
+            .with_query_planner(planner)
+            .build();
+        let ctx = Arc::new(SessionContext::new_with_state(state));
+
+        Ok(Self {
+            ctx,
+            logical_codec: Arc::clone(&self.logical_codec),
+            physical_codec: Arc::clone(&self.physical_codec),
+        })
+    }
+
     pub fn table_provider(&self, name: &str, py: Python) -> PyResult<PyTable> {
         let provider = wait_for_future(py, self.ctx.table_provider(name))
             // Outer error: runtime/async failure
@@ -1385,6 +1427,19 @@ impl PySessionContext {
         create_logical_extension_capsule(py, ffi.as_ref())
     }
 
+    pub fn __datafusion_query_planner__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let planner = Arc::clone(self.ctx.state().query_planner());
+        let ffi = FFI_QueryPlanner::new_with_ffi_codecs(
+            planner,
+            self.ffi_logical_codec().as_ref().clone(),
+            self.ffi_physical_codec().as_ref().clone(),
+        );
+        create_query_planner_capsule(py, &ffi)
+    }
+
     pub fn with_logical_extension_codec<'py>(
         &self,
         codec: Bound<'py, PyAny>,

crates/core/src/context.rs

@@ -29,6 +29,7 @@ use datafusion_ffi::execution::FFI_TaskContextProvider;
 use datafusion_ffi::physical_optimizer::FFI_PhysicalOptimizerRule;
 use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use datafusion_ffi::query_planner::FFI_QueryPlanner;
 use datafusion_ffi::table_provider::FFI_TableProvider;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use pyo3::exceptions::{PyImportError, PyTypeError, PyValueError};
@@ -231,6 +232,38 @@ pub fn ffi_logical_codec_from_pycapsule(obj: Bound<PyAny>) -> PyResult<FFI_Logic
     Ok(codec.clone())
 }
 
+pub fn create_query_planner_capsule<'py>(
+    py: Python<'py>,
+    planner: &FFI_QueryPlanner,
+) -> PyResult<Bound<'py, PyCapsule>> {
+    PyCapsule::new_with_value(py, planner.clone(), cr"datafusion_query_planner")
+}
+
+pub fn ffi_query_planner_from_pycapsule(obj: &Bound<PyAny>) -> PyResult<FFI_QueryPlanner> {
+    let attr_name = "__datafusion_query_planner__";
+    let capsule = if obj.hasattr(attr_name)? {
+        obj.getattr(attr_name)?.call0()?
+    } else {
+        obj.clone()
+    };
+
+    let capsule = capsule.cast::<PyCapsule>()?;
+    validate_pycapsule(capsule, "datafusion_query_planner")?;
+    let data: NonNull<FFI_QueryPlanner> = capsule
+        .pointer_checked(Some(c"datafusion_query_planner"))?
+        .cast();
+    let planner = unsafe { data.as_ref() };
+    let planner_version = unsafe { (planner.version)() };
+    let expected_version = datafusion_ffi::version();
+    if planner_version != expected_version {
+        return Err(PyImportError::new_err(format!(
+            "Incompatible DataFusion query planner version {planner_version}; expected major version {expected_version}."
+        )));
+    }
+
+    Ok(planner.clone())
+}
+
 pub fn create_physical_extension_capsule<'py>(
     py: Python<'py>,
     codec: &FFI_PhysicalExtensionCodec,

crates/util/src/lib.rs

@@ -34,6 +34,7 @@ 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,39 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import pytest
+from datafusion import SessionContext
+from datafusion_ffi_example import MyQueryPlanner
+
+
+@pytest.mark.parametrize("raw_capsule", [False, True])
+def test_ffi_query_planner_runs_during_planning(raw_capsule: bool):
+    """A query planner imported from another library creates the physical plan."""
+    planner = MyQueryPlanner()
+    exported_planner = (
+        planner.__datafusion_query_planner__() if raw_capsule else planner
+    )
+    ctx = SessionContext().with_query_planner(exported_planner)
+
+    before = planner.plan_calls()
+    result = ctx.sql("SELECT 1 AS value").collect()
+    after = planner.plan_calls()
+
+    assert after > before
+    assert result == []

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

@@ -23,6 +23,7 @@ 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;
@@ -35,6 +36,7 @@ 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;
@@ -58,5 +60,6 @@ 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(())
 }

examples/datafusion-ffi-example/src/lib.rs

@@ -0,0 +1,92 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use arrow::datatypes::{DataType, Field, Schema};
+use async_trait::async_trait;
+use datafusion::execution::TaskContextProvider;
+use datafusion::execution::context::{QueryPlanner, SessionContext};
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::empty::EmptyExec;
+use datafusion_ffi::query_planner::FFI_QueryPlanner;
+use datafusion_python_util::get_tokio_runtime;
+use datafusion_session::Session;
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+
+#[derive(Debug)]
+struct CountingQueryPlanner {
+    plan_calls: Arc<AtomicUsize>,
+}
+
+#[async_trait]
+impl QueryPlanner for CountingQueryPlanner {
+    async fn create_physical_plan(
+        &self,
+        _logical_plan: &LogicalPlan,
+        _session: &dyn Session,
+    ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+        self.plan_calls.fetch_add(1, Ordering::SeqCst);
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "value",
+            DataType::Int64,
+            false,
+        )]));
+        Ok(Arc::new(EmptyExec::new(schema)))
+    }
+}
+
+/// Python-visible query planner used to test planning across a real FFI boundary.
+#[pyclass(
+    from_py_object,
+    name = "MyQueryPlanner",
+    module = "datafusion_ffi_example",
+    subclass
+)]
+#[derive(Debug, Default, Clone)]
+pub(crate) struct MyQueryPlanner {
+    plan_calls: Arc<AtomicUsize>,
+}
+
+#[pymethods]
+impl MyQueryPlanner {
+    #[new]
+    fn new() -> Self {
+        Self::default()
+    }
+
+    fn plan_calls(&self) -> usize {
+        self.plan_calls.load(Ordering::SeqCst)
+    }
+
+    fn __datafusion_query_planner__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let planner: Arc<dyn QueryPlanner + Send + Sync> = Arc::new(CountingQueryPlanner {
+            plan_calls: Arc::clone(&self.plan_calls),
+        });
+        let runtime = get_tokio_runtime().handle().clone();
+        let ctx_provider = Arc::new(SessionContext::new()) as Arc<dyn TaskContextProvider>;
+        let ffi = FFI_QueryPlanner::new(planner, Some(runtime), &ctx_provider, None, None);
+
+        PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner")
+    }
+}