This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new e08aed1e5d feat: Implement FFI_QueryPlanner (#24028)
e08aed1e5d is described below
commit e08aed1e5de41dcf81d529140dae07723b942a5e
Author: Tim Saucer <[email protected]>
AuthorDate: Fri Aug 7 07:17:17 2026 -0400
feat: Implement FFI_QueryPlanner (#24028)
## Which issue does this PR close?
- Unblocks https://github.com/apache/datafusion-python/issues/1612
## Rationale for this change
This is the last in a series of PRs that would enable FFI `Session` to
support a `QueryPlanner`. The prior work was in
- https://github.com/apache/datafusion/pull/23649
- https://github.com/apache/datafusion/pull/23703
- https://github.com/apache/datafusion/pull/23842
With those changes in place we now have the dependencies correct that we
can expose a `FFI_QueryPlanner` on a `FFI_Session`. With this we can
enable foreign libraries such as `datafusion-distributed` and `ballista`
to provide a query planner in Python and connect it directly to a
`datafusion-python`'s `SessionContext`.
## What changes are included in this PR?
Addition only. Adds these functions to `FFI_Session` and their
supporting structures:
- `query_planner()`
- `optimize()`
- `physical_optimizers()`
## Are these changes tested?
Unit and integration tests are provided.
## Are there any user-facing changes?
This is addition, but it does break the FFI ABI, which is already
evolving in DF55.
---------
Co-authored-by: Claude Opus 5 (1M context) <[email protected]>
---
datafusion/datasource-parquet/src/file_format.rs | 1 +
.../datasource-parquet/src/opener/encryption.rs | 1 +
datafusion/datasource-parquet/src/sink.rs | 1 +
datafusion/ffi/src/lib.rs | 1 +
.../ffi/src/proto/logical_extension_codec.rs | 8 +-
.../ffi/src/proto/physical_extension_codec.rs | 8 +-
datafusion/ffi/src/query_planner.rs | 445 +++++++++++++++++++++
datafusion/ffi/src/session/mod.rs | 267 ++++++++++++-
datafusion/ffi/src/table_provider.rs | 4 +-
datafusion/ffi/src/table_provider_factory.rs | 2 +-
datafusion/ffi/src/tests/mod.rs | 13 +
datafusion/ffi/src/tests/query_planner.rs | 172 ++++++++
datafusion/ffi/src/tests/utils.rs | 43 +-
datafusion/ffi/src/udtf.rs | 2 +-
datafusion/ffi/tests/ffi_query_planner.rs | 348 ++++++++++++++++
docs/source/library-user-guide/upgrading/55.0.0.md | 23 +-
16 files changed, 1299 insertions(+), 40 deletions(-)
diff --git a/datafusion/datasource-parquet/src/file_format.rs
b/datafusion/datasource-parquet/src/file_format.rs
index 29083ebfb2..8d19bedeb7 100644
--- a/datafusion/datasource-parquet/src/file_format.rs
+++ b/datafusion/datasource-parquet/src/file_format.rs
@@ -297,6 +297,7 @@ async fn get_file_decryption_properties(
}
#[cfg(not(feature = "parquet_encryption"))]
+#[expect(clippy::unused_async)]
async fn get_file_decryption_properties(
_state: &dyn Session,
_options: &TableParquetOptions,
diff --git a/datafusion/datasource-parquet/src/opener/encryption.rs
b/datafusion/datasource-parquet/src/opener/encryption.rs
index b725198237..498fe8acf7 100644
--- a/datafusion/datasource-parquet/src/opener/encryption.rs
+++ b/datafusion/datasource-parquet/src/opener/encryption.rs
@@ -76,6 +76,7 @@ impl EncryptionContext {
#[cfg(not(feature = "parquet_encryption"))]
#[expect(dead_code)]
+#[expect(clippy::unused_async)]
impl EncryptionContext {
pub(super) async fn get_file_decryption_properties(
&self,
diff --git a/datafusion/datasource-parquet/src/sink.rs
b/datafusion/datasource-parquet/src/sink.rs
index e11f1d29d7..53f6f1e6b4 100644
--- a/datafusion/datasource-parquet/src/sink.rs
+++ b/datafusion/datasource-parquet/src/sink.rs
@@ -243,6 +243,7 @@ async fn set_writer_encryption_properties(
}
#[cfg(not(feature = "parquet_encryption"))]
+#[expect(clippy::unused_async)]
async fn set_writer_encryption_properties(
builder: WriterPropertiesBuilder,
_runtime: &Arc<RuntimeEnv>,
diff --git a/datafusion/ffi/src/lib.rs b/datafusion/ffi/src/lib.rs
index fd2ac58576..de8f8cba9c 100644
--- a/datafusion/ffi/src/lib.rs
+++ b/datafusion/ffi/src/lib.rs
@@ -39,6 +39,7 @@ pub mod physical_optimizer;
pub mod placement;
pub mod plan_properties;
pub mod proto;
+pub mod query_planner;
pub mod record_batch_stream;
pub mod schema_provider;
pub mod session;
diff --git a/datafusion/ffi/src/proto/logical_extension_codec.rs
b/datafusion/ffi/src/proto/logical_extension_codec.rs
index 97aa5c901a..ed2c594f1b 100644
--- a/datafusion/ffi/src/proto/logical_extension_codec.rs
+++ b/datafusion/ffi/src/proto/logical_extension_codec.rs
@@ -99,7 +99,7 @@ pub struct FFI_LogicalExtensionCodec {
try_encode_udwf:
unsafe extern "C" fn(&Self, node: FFI_WindowUDF) ->
FFI_Result<SVec<u8>>,
- pub task_ctx_provider: FFI_TaskContextProvider,
+ pub(crate) task_ctx_provider: FFI_TaskContextProvider,
/// Used to create a clone on the provider of the execution plan. This
should
/// only need to be called by the receiver of the plan.
@@ -295,7 +295,7 @@ impl Drop for FFI_LogicalExtensionCodec {
impl FFI_LogicalExtensionCodec {
/// Creates a new [`FFI_LogicalExtensionCodec`].
pub fn new(
- codec: Arc<dyn LogicalExtensionCodec + Send>,
+ codec: Arc<dyn LogicalExtensionCodec>,
runtime: Option<Handle>,
task_ctx_provider: impl Into<FFI_TaskContextProvider>,
) -> Self {
@@ -712,14 +712,12 @@ mod tests {
#[test]
fn ffi_logical_extension_codec_local_bypass() {
- let codec =
- Arc::new(TestExtensionCodec {}) as Arc<dyn LogicalExtensionCodec +
Send>;
+ let codec = Arc::new(TestExtensionCodec {}) as Arc<dyn
LogicalExtensionCodec>;
let (_ctx, task_ctx_provider) =
crate::util::tests::test_session_and_ctx();
let mut ffi_codec =
FFI_LogicalExtensionCodec::new(Arc::clone(&codec), None,
task_ctx_provider);
- let codec = codec as Arc<dyn LogicalExtensionCodec>;
// Verify local libraries can be downcast to their original
let foreign_codec: Arc<dyn LogicalExtensionCodec> =
(&ffi_codec).into();
assert!(arc_ptr_eq(&foreign_codec, &codec));
diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs
b/datafusion/ffi/src/proto/physical_extension_codec.rs
index 9e64df82e3..95d2ed68a6 100644
--- a/datafusion/ffi/src/proto/physical_extension_codec.rs
+++ b/datafusion/ffi/src/proto/physical_extension_codec.rs
@@ -92,7 +92,7 @@ pub struct FFI_PhysicalExtensionCodec {
unsafe extern "C" fn(&Self, node: FFI_WindowUDF) ->
FFI_Result<SVec<u8>>,
/// Access the current [`TaskContext`].
- task_ctx_provider: FFI_TaskContextProvider,
+ pub(crate) task_ctx_provider: FFI_TaskContextProvider,
/// Used to create a clone on the provider of the execution plan. This
should
/// only need to be called by the receiver of the plan.
@@ -281,7 +281,7 @@ impl Drop for FFI_PhysicalExtensionCodec {
impl FFI_PhysicalExtensionCodec {
/// Creates a new [`FFI_PhysicalExtensionCodec`].
pub fn new(
- codec: Arc<dyn PhysicalExtensionCodec + Send>,
+ codec: Arc<dyn PhysicalExtensionCodec>,
runtime: Option<Handle>,
task_ctx_provider: impl Into<FFI_TaskContextProvider>,
) -> Self {
@@ -695,14 +695,12 @@ pub(crate) mod tests {
#[test]
fn ffi_physical_extension_codec_local_bypass() {
- let codec =
- Arc::new(TestExtensionCodec {}) as Arc<dyn PhysicalExtensionCodec
+ Send>;
+ let codec = Arc::new(TestExtensionCodec {}) as Arc<dyn
PhysicalExtensionCodec>;
let (_ctx, task_ctx_provider) =
crate::util::tests::test_session_and_ctx();
let mut ffi_codec =
FFI_PhysicalExtensionCodec::new(Arc::clone(&codec), None,
task_ctx_provider);
- let codec = codec as Arc<dyn PhysicalExtensionCodec>;
// Verify local libraries can be downcast to their original
let foreign_codec: Arc<dyn PhysicalExtensionCodec> =
(&ffi_codec).into();
assert!(arc_ptr_eq(&foreign_codec, &codec));
diff --git a/datafusion/ffi/src/query_planner.rs
b/datafusion/ffi/src/query_planner.rs
new file mode 100644
index 0000000000..6d895d65c5
--- /dev/null
+++ b/datafusion/ffi/src/query_planner.rs
@@ -0,0 +1,445 @@
+// 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.
+
+//! FFI support for [`QueryPlanner`].
+//!
+//! A typical deployment has three libraries. Library A (for example,
+//! `datafusion-python`) owns the [`Session`] and codec registry. Library B
owns
+//! a custom table provider and its extension nodes. Library C (for example,
+//! Ballista or `datafusion-distributed`) owns the query planner. A serializes
a
+//! logical plan and invokes C, while `FFI_SessionRef` lets C call session
+//! services in A. C deserializes the logical plan, creates a physical plan,
+//! serializes that result, and returns it for A to deserialize. The logical
and
+//! physical extension codecs preserve nodes supplied by B.
+//!
+//! The physical result is serialized instead of returned as an
+//! [`crate::execution_plan::FFI_ExecutionPlan`]. An FFI execution-plan handle
is
+//! a foreign trait-object proxy, so even a built-in plan created in C cannot
be
+//! downcast to its concrete
+//! type in A. Serialization reconstructs known plan nodes with A's local Rust
+//! type identities, allowing A's optimizers and other consumers to downcast
+//! them. Extension codecs control how custom nodes are reconstructed.
+//!
+//! A node returned by B while C is planning is still foreign to C unless a
+//! codec boundary reconstructs it in C. The query-planner boundary guarantees
+//! that C-local serializable nodes, and extension nodes understood by the
+//! configured codecs, are reconstructed for A when the completed plan returns.
+//!
+//! # Delegating back to library A
+//!
+//! C commonly wants A's built-in planning as a starting point, then rewrites
the
+//! result. A must export its planner *before* installing C's planner on the
+//! session, and C must retain that handle: after the swap,
+//! [`Session::query_planner`] reports C's own planner, and
+//! [`Session::create_physical_plan`] dispatches to it, so either one is a
+//! self-call. Delegating to the retained handle is safe, because DataFusion's
+//! built-in physical planner never re-dispatches through [`Session`].
+//!
+//! Retain the planner rather than the session. [`FFI_QueryPlanner`] owns a
+//! reference-counted planner, so it outlives A's original session, whereas
+//! `FFI_SessionRef` borrows its session with the lifetime erased.
+
+use std::ffi::c_void;
+use std::sync::Arc;
+
+use async_ffi::{FfiFuture, FutureExt};
+use async_trait::async_trait;
+use datafusion_common::error::{DataFusionError, Result};
+use datafusion_expr::LogicalPlan;
+use datafusion_physical_plan::ExecutionPlan;
+use datafusion_proto::bytes::{
+ logical_plan_from_bytes_with_extension_codec,
+ logical_plan_to_bytes_with_extension_codec,
+ physical_plan_from_bytes_with_extension_codec,
+ physical_plan_to_bytes_with_extension_codec,
+};
+use datafusion_proto::logical_plan::LogicalExtensionCodec;
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use datafusion_session::{QueryPlanner, Session};
+use stabby::vec::Vec as SVec;
+use tokio::runtime::Handle;
+
+use crate::execution::FFI_TaskContextProvider;
+use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use crate::session::{FFI_SessionRef, ForeignSession};
+use crate::util::FFI_Result;
+use crate::{df_result, sresult_return};
+
+/// An ABI-stable handle to a [`QueryPlanner`] owned by another library.
+///
+/// The Rust-facing adapters serialize the input [`LogicalPlan`] and resulting
+/// [`ExecutionPlan`]; callers do not invoke the byte-oriented function pointer
+/// directly.
+#[repr(C)]
+#[derive(Debug)]
+pub struct FFI_QueryPlanner {
+ create_physical_plan: unsafe extern "C" fn(
+ &Self,
+ logical_plan_serialized: SVec<u8>,
+ session: FFI_SessionRef,
+ ) -> FfiFuture<FFI_Result<SVec<u8>>>,
+
+ /// Codec used to encode and decode logical plans and extension nodes.
+ logical_codec: FFI_LogicalExtensionCodec,
+
+ /// Codec used to encode and decode physical plans and extension nodes.
+ physical_codec: FFI_PhysicalExtensionCodec,
+
+ /// Used to create a clone of the query planner.
+ clone: unsafe extern "C" fn(planner: &Self) -> Self,
+
+ /// Release the memory of the private data when it is no longer being used.
+ release: unsafe extern "C" fn(arg: &mut Self),
+
+ /// Return the major DataFusion version number of this planner.
+ pub version: unsafe extern "C" fn() -> u64,
+
+ /// Internal data. This is only to be accessed by the provider of the
planner.
+ /// A [`ForeignQueryPlanner`] should never attempt to access this data.
+ private_data: *mut c_void,
+
+ /// Utility to identify when FFI objects are accessed locally through
+ /// the foreign interface. See [`crate::get_library_marker_id`].
+ pub library_marker_id: extern "C" fn() -> usize,
+}
+
+unsafe impl Send for FFI_QueryPlanner {}
+unsafe impl Sync for FFI_QueryPlanner {}
+
+struct QueryPlannerPrivateData {
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+}
+
+impl FFI_QueryPlanner {
+ fn inner(&self) -> &Arc<dyn QueryPlanner + Send + Sync> {
+ let private_data = self.private_data as *const QueryPlannerPrivateData;
+ unsafe { &(*private_data).planner }
+ }
+}
+
+unsafe extern "C" fn create_physical_plan_fn_wrapper(
+ planner: &FFI_QueryPlanner,
+ logical_plan_serialized: SVec<u8>,
+ session: FFI_SessionRef,
+) -> FfiFuture<FFI_Result<SVec<u8>>> {
+ let internal_planner = Arc::clone(planner.inner());
+ let logical_codec: Arc<dyn LogicalExtensionCodec> =
(&planner.logical_codec).into();
+ let physical_codec: Arc<dyn PhysicalExtensionCodec> =
+ (&planner.physical_codec).into();
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&dyn Session, DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session =
Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
+
+ let logical_plan =
sresult_return!(logical_plan_from_bytes_with_extension_codec(
+ logical_plan_serialized.as_slice(),
+ session.task_ctx().as_ref(),
+ logical_codec.as_ref(),
+ ));
+
+ let physical_plan = sresult_return!(
+ internal_planner
+ .create_physical_plan(&logical_plan, session)
+ .await
+ );
+ let physical_plan =
sresult_return!(physical_plan_to_bytes_with_extension_codec(
+ physical_plan,
+ physical_codec.as_ref(),
+ ));
+
+ FFI_Result::Ok(SVec::from(physical_plan.as_ref()))
+ }
+ .into_ffi()
+}
+
+unsafe extern "C" fn release_fn_wrapper(planner: &mut FFI_QueryPlanner) {
+ unsafe {
+ debug_assert!(!planner.private_data.is_null());
+ let private_data =
+ Box::from_raw(planner.private_data as *mut
QueryPlannerPrivateData);
+ drop(private_data);
+ planner.private_data = std::ptr::null_mut();
+ }
+}
+
+unsafe extern "C" fn clone_fn_wrapper(planner: &FFI_QueryPlanner) ->
FFI_QueryPlanner {
+ let old_planner = Arc::clone(planner.inner());
+
+ let private_data = Box::into_raw(Box::new(QueryPlannerPrivateData {
+ planner: old_planner,
+ })) as *mut c_void;
+
+ FFI_QueryPlanner {
+ create_physical_plan: create_physical_plan_fn_wrapper,
+ logical_codec: planner.logical_codec.clone(),
+ physical_codec: planner.physical_codec.clone(),
+ clone: clone_fn_wrapper,
+ release: release_fn_wrapper,
+ version: super::version,
+ private_data,
+ library_marker_id: crate::get_library_marker_id,
+ }
+}
+
+impl Drop for FFI_QueryPlanner {
+ fn drop(&mut self) {
+ unsafe { (self.release)(self) }
+ }
+}
+
+impl Clone for FFI_QueryPlanner {
+ fn clone(&self) -> Self {
+ unsafe { (self.clone)(self) }
+ }
+}
+
+impl FFI_QueryPlanner {
+ /// Creates an [`FFI_QueryPlanner`] with native extension codecs.
+ ///
+ /// Both codecs are required so that the caller states which extension
nodes
+ /// survive the boundary. Pass
+ ///
[`DefaultLogicalExtensionCodec`](datafusion_proto::logical_plan::DefaultLogicalExtensionCodec)
+ /// and
+ ///
[`DefaultPhysicalExtensionCodec`](datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec)
+ /// when no custom nodes are involved. `runtime` and `task_ctx_provider`
+ /// support codec callbacks across the FFI boundary.
+ pub fn new(
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+ runtime: Option<Handle>,
+ task_ctx_provider: impl Into<FFI_TaskContextProvider>,
+ logical_codec: Arc<dyn LogicalExtensionCodec>,
+ physical_codec: Arc<dyn PhysicalExtensionCodec>,
+ ) -> Self {
+ let task_ctx_provider = task_ctx_provider.into();
+ let logical_codec = FFI_LogicalExtensionCodec::new(
+ logical_codec,
+ runtime.clone(),
+ task_ctx_provider.clone(),
+ );
+ let physical_codec =
+ FFI_PhysicalExtensionCodec::new(physical_codec, runtime,
task_ctx_provider);
+ Self::new_with_ffi_codecs(planner, logical_codec, physical_codec)
+ }
+
+ /// Creates an [`FFI_QueryPlanner`] using prebuilt FFI extension codecs.
+ ///
+ /// If `planner` is already foreign, this re-exports its original FFI
handle
+ /// rather than adding another wrapper layer. The handle still adopts the
+ /// codecs supplied here, so they are never silently discarded.
+ pub fn new_with_ffi_codecs(
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+ logical_codec: FFI_LogicalExtensionCodec,
+ physical_codec: FFI_PhysicalExtensionCodec,
+ ) -> Self {
+ let any_ref: &dyn std::any::Any = planner.as_ref();
+ if let Some(planner) = any_ref.downcast_ref::<ForeignQueryPlanner>() {
+ let mut planner = planner.0.clone();
+ planner.logical_codec = logical_codec;
+ planner.physical_codec = physical_codec;
+ return planner;
+ }
+
+ let private_data = Box::new(QueryPlannerPrivateData { planner });
+
+ Self {
+ create_physical_plan: create_physical_plan_fn_wrapper,
+ logical_codec,
+ physical_codec,
+ clone: clone_fn_wrapper,
+ release: release_fn_wrapper,
+ version: super::version,
+ private_data: Box::into_raw(private_data) as *mut c_void,
+ library_marker_id: crate::get_library_marker_id,
+ }
+ }
+
+ /// Creates a physical plan through this planner's FFI interface.
+ ///
+ /// This serializes `logical_plan`, exports `session` as an
+ /// `FFI_SessionRef`, invokes the planner's owning library, and
+ /// deserializes its physical-plan response. `session_runtime` is attached
+ /// to the exported session for callbacks that need its Tokio runtime.
+ ///
+ /// The [`QueryPlanner`] implementation for [`ForeignQueryPlanner`] cannot
+ /// obtain the session owner's runtime from the trait API, so it calls this
+ /// method with `None`. Embedders that own the runtime and need session
+ /// callbacks to enter it must call this method directly with
`Some(handle)`.
+ pub async fn create_physical_plan_with_session_runtime(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ session_runtime: Option<Handle>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let codec: Arc<dyn LogicalExtensionCodec> =
(&self.logical_codec).into();
+ let logical_plan =
+ logical_plan_to_bytes_with_extension_codec(logical_plan,
codec.as_ref())?;
+ let logical_plan = SVec::from(logical_plan.as_ref());
+ let task_ctx = session.task_ctx();
+ let session = FFI_SessionRef::new_with_ffi_codecs(
+ session,
+ session_runtime,
+ self.logical_codec.clone(),
+ self.physical_codec.clone(),
+ );
+
+ let physical_plan = unsafe {
+ df_result!((self.create_physical_plan)(self, logical_plan,
session).await)?
+ };
+ let physical_codec: Arc<dyn PhysicalExtensionCodec> =
+ (&self.physical_codec).into();
+
+ physical_plan_from_bytes_with_extension_codec(
+ physical_plan.as_slice(),
+ task_ctx.as_ref(),
+ physical_codec.as_ref(),
+ )
+ }
+}
+
+/// Consumer-side [`QueryPlanner`] adapter for an [`FFI_QueryPlanner`].
+///
+/// Calls serialize the logical plan, invoke the producing library, and
+/// deserialize its physical-plan response.
+#[derive(Debug)]
+pub struct ForeignQueryPlanner(pub FFI_QueryPlanner);
+
+unsafe impl Send for ForeignQueryPlanner {}
+unsafe impl Sync for ForeignQueryPlanner {}
+
+impl From<&FFI_QueryPlanner> for Arc<dyn QueryPlanner + Send + Sync> {
+ fn from(planner: &FFI_QueryPlanner) -> Self {
+ if (planner.library_marker_id)() == crate::get_library_marker_id() {
+ Arc::clone(planner.inner())
+ } else {
+ Arc::new(ForeignQueryPlanner(planner.clone()))
+ }
+ }
+}
+
+#[async_trait]
+impl QueryPlanner for ForeignQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.0
+ .create_physical_plan_with_session_runtime(logical_plan, session,
None)
+ .await
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::prelude::SessionContext;
+ use datafusion_common::Result;
+ use datafusion_execution::TaskContextProvider;
+ use datafusion_expr::LogicalPlanBuilder;
+ use datafusion_physical_plan::empty::EmptyExec;
+ use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec;
+ use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
+
+ use super::*;
+
+ #[derive(Debug)]
+ struct EmptyQueryPlanner;
+
+ #[async_trait]
+ impl QueryPlanner for EmptyQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ _logical_plan: &LogicalPlan,
+ _session: &dyn Session,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let schema =
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Int32,
true)]));
+ Ok(Arc::new(EmptyExec::new(schema)))
+ }
+ }
+
+ fn create_ffi_query_planner(ctx: Arc<SessionContext>) -> FFI_QueryPlanner {
+ let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn
TaskContextProvider>;
+ FFI_QueryPlanner::new(
+ Arc::new(EmptyQueryPlanner),
+ None,
+ &task_ctx_provider,
+ Arc::new(DefaultLogicalExtensionCodec {}),
+ Arc::new(DefaultPhysicalExtensionCodec {}),
+ )
+ }
+
+ #[test]
+ fn test_ffi_query_planner_local_bypass() {
+ let ctx = Arc::new(SessionContext::new());
+ let ffi_planner = create_ffi_query_planner(ctx);
+ let planner: Arc<dyn QueryPlanner + Send + Sync> =
(&ffi_planner).into();
+ let any_ref: &dyn std::any::Any = planner.as_ref();
+ assert!(any_ref.downcast_ref::<EmptyQueryPlanner>().is_some());
+ }
+
+ #[tokio::test]
+ async fn test_round_trip_ffi_query_planner_create_physical_plan() ->
Result<()> {
+ let ctx = Arc::new(SessionContext::new());
+ let mut ffi_planner = create_ffi_query_planner(Arc::clone(&ctx));
+ ffi_planner.library_marker_id = crate::mock_foreign_marker_id;
+
+ let planner: Arc<dyn QueryPlanner + Send + Sync> =
(&ffi_planner).into();
+ let any_ref: &dyn std::any::Any = planner.as_ref();
+ assert!(any_ref.downcast_ref::<ForeignQueryPlanner>().is_some());
+
+ let logical_plan = LogicalPlanBuilder::empty(false).build()?;
+ let state = ctx.state();
+ let physical_plan = planner.create_physical_plan(&logical_plan,
&state).await?;
+ assert_eq!(physical_plan.name(), "EmptyExec");
+ assert!(physical_plan.is::<EmptyExec>());
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn test_create_physical_plan_with_session_runtime() -> Result<()> {
+ let ctx = Arc::new(SessionContext::new());
+ let ffi_planner = create_ffi_query_planner(Arc::clone(&ctx));
+ let logical_plan = LogicalPlanBuilder::empty(false).build()?;
+ let state = ctx.state();
+
+ let physical_plan = ffi_planner
+ .create_physical_plan_with_session_runtime(
+ &logical_plan,
+ &state,
+ Some(Handle::current()),
+ )
+ .await?;
+
+ assert_eq!(physical_plan.name(), "EmptyExec");
+ assert!(physical_plan.is::<EmptyExec>());
+
+ Ok(())
+ }
+}
diff --git a/datafusion/ffi/src/session/mod.rs
b/datafusion/ffi/src/session/mod.rs
index 519384379e..83f842508a 100644
--- a/datafusion/ffi/src/session/mod.rs
+++ b/datafusion/ffi/src/session/mod.rs
@@ -15,10 +15,25 @@
// specific language governing permissions and limitations
// under the License.
+//! FFI support for [`Session`].
+//!
+//! # Delegating physical planning
+//!
+//! Consider a session owned by library A that uses a query planner owned by
+//! library C. After A installs C's planner, [`ForeignSession::query_planner`]
+//! returns C's planner and [`ForeignSession::create_physical_plan`] dispatches
+//! to C's planner. C must not call `create_physical_plan`, or invoke the
planner
+//! returned by `query_planner`, to delegate planning back to A. Repeating
either
+//! self-call recurses until the stack is exhausted.
+//!
+//! To delegate safely, A must export its original planner before installing
C's
+//! planner, and C must retain and invoke that planner directly. See the
+//! [`crate::query_planner`] module for details.
+
use std::any::Any;
use std::collections::HashMap;
use std::ffi::c_void;
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
use arrow_schema::SchemaRef;
use arrow_schema::ffi::FFI_ArrowSchema;
@@ -37,12 +52,18 @@ use datafusion_expr::{
};
use datafusion_physical_expr::PhysicalExpr;
use datafusion_physical_plan::ExecutionPlan;
-use datafusion_proto::bytes::{logical_plan_from_bytes, logical_plan_to_bytes};
+use datafusion_proto::bytes::{
+ logical_plan_from_bytes, logical_plan_from_bytes_with_extension_codec,
+ logical_plan_to_bytes, logical_plan_to_bytes_with_extension_codec,
+};
use datafusion_proto::logical_plan::LogicalExtensionCodec;
use datafusion_proto::logical_plan::from_proto::parse_expr;
use datafusion_proto::logical_plan::to_proto::serialize_expr;
+use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
use datafusion_proto::protobuf::LogicalExprNode;
-use datafusion_session::{CatalogProviderList, Session};
+use datafusion_session::{
+ CatalogProviderList, PhysicalOptimizerRule, QueryPlanner, Session,
+};
use prost::Message;
use stabby::str::Str as SStr;
@@ -55,7 +76,10 @@ use crate::catalog_provider_list::FFI_CatalogProviderList;
use crate::execution::FFI_TaskContext;
use crate::execution_plan::FFI_ExecutionPlan;
use crate::physical_expr::FFI_PhysicalExpr;
+use crate::physical_optimizer::FFI_PhysicalOptimizerRule;
use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use crate::query_planner::FFI_QueryPlanner;
use crate::session::config::FFI_SessionConfig;
use crate::udaf::FFI_AggregateUDF;
use crate::udf::FFI_ScalarUDF;
@@ -86,6 +110,13 @@ pub(crate) struct FFI_SessionRef {
catalog_list: unsafe extern "C" fn(&Self) -> FFI_CatalogProviderList,
+ query_planner: unsafe extern "C" fn(&Self) -> FFI_QueryPlanner,
+
+ optimize: unsafe extern "C" fn(
+ &Self,
+ logical_plan_serialized: SVec<u8>,
+ ) -> FFI_Result<SVec<u8>>,
+
create_physical_plan:
unsafe extern "C" fn(
&Self,
@@ -110,8 +141,12 @@ pub(crate) struct FFI_SessionRef {
task_ctx: unsafe extern "C" fn(&Self) -> FFI_TaskContext,
+ physical_optimizers: unsafe extern "C" fn(&Self) ->
SVec<FFI_PhysicalOptimizerRule>,
+
logical_codec: FFI_LogicalExtensionCodec,
+ physical_codec: FFI_PhysicalExtensionCodec,
+
/// Used to create a clone on the provider of the registry. This should
/// only need to be called by the receiver of the plan.
clone: unsafe extern "C" fn(plan: &Self) -> Self,
@@ -135,12 +170,12 @@ unsafe impl Send for FFI_SessionRef {}
unsafe impl Sync for FFI_SessionRef {}
struct SessionPrivateData<'a> {
- session: &'a (dyn Session + Send + Sync),
+ session: &'a dyn Session,
runtime: Option<Handle>,
}
impl FFI_SessionRef {
- fn inner(&self) -> &(dyn Session + Send + Sync) {
+ fn inner(&self) -> &dyn Session {
let private_data = self.private_data as *const SessionPrivateData;
unsafe { (*private_data).session }
}
@@ -173,6 +208,36 @@ unsafe extern "C" fn catalog_list_fn_wrapper(
)
}
+unsafe extern "C" fn query_planner_fn_wrapper(
+ session: &FFI_SessionRef,
+) -> FFI_QueryPlanner {
+ FFI_QueryPlanner::new_with_ffi_codecs(
+ session.inner().query_planner(),
+ session.logical_codec.clone(),
+ session.physical_codec.clone(),
+ )
+}
+
+unsafe extern "C" fn optimize_fn_wrapper(
+ session: &FFI_SessionRef,
+ logical_plan_serialized: SVec<u8>,
+) -> FFI_Result<SVec<u8>> {
+ let logical_codec: Arc<dyn LogicalExtensionCodec> =
(&session.logical_codec).into();
+ let inner = session.inner();
+ let logical_plan =
sresult_return!(logical_plan_from_bytes_with_extension_codec(
+ logical_plan_serialized.as_slice(),
+ inner.task_ctx().as_ref(),
+ logical_codec.as_ref(),
+ ));
+ let optimized_plan = sresult_return!(inner.optimize(&logical_plan));
+ let optimized_plan =
sresult_return!(logical_plan_to_bytes_with_extension_codec(
+ &optimized_plan,
+ logical_codec.as_ref(),
+ ));
+
+ FFI_Result::Ok(SVec::from(optimized_plan.as_ref()))
+}
+
unsafe extern "C" fn create_physical_plan_fn_wrapper(
session: &FFI_SessionRef,
logical_plan_serialized: SVec<u8>,
@@ -303,6 +368,18 @@ unsafe extern "C" fn task_ctx_fn_wrapper(session:
&FFI_SessionRef) -> FFI_TaskCo
session.inner().task_ctx().into()
}
+unsafe extern "C" fn physical_optimizers_fn_wrapper(
+ session: &FFI_SessionRef,
+) -> SVec<FFI_PhysicalOptimizerRule> {
+ let runtime = unsafe { session.runtime().clone() };
+ session
+ .inner()
+ .physical_optimizers()
+ .iter()
+ .map(|rule| FFI_PhysicalOptimizerRule::new(Arc::clone(rule),
runtime.clone()))
+ .collect()
+}
+
unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_SessionRef) {
unsafe {
let private_data =
@@ -324,6 +401,8 @@ unsafe extern "C" fn clone_fn_wrapper(provider:
&FFI_SessionRef) -> FFI_SessionR
session_id: session_id_fn_wrapper,
config: config_fn_wrapper,
catalog_list: catalog_list_fn_wrapper,
+ query_planner: query_planner_fn_wrapper,
+ optimize: optimize_fn_wrapper,
create_physical_plan: create_physical_plan_fn_wrapper,
create_physical_expr: create_physical_expr_fn_wrapper,
scalar_functions: scalar_functions_fn_wrapper,
@@ -332,7 +411,9 @@ unsafe extern "C" fn clone_fn_wrapper(provider:
&FFI_SessionRef) -> FFI_SessionR
table_options: table_options_fn_wrapper,
default_table_options: default_table_options_fn_wrapper,
task_ctx: task_ctx_fn_wrapper,
+ physical_optimizers: physical_optimizers_fn_wrapper,
logical_codec: provider.logical_codec.clone(),
+ physical_codec: provider.physical_codec.clone(),
clone: clone_fn_wrapper,
release: release_fn_wrapper,
@@ -350,14 +431,60 @@ impl Drop for FFI_SessionRef {
}
impl FFI_SessionRef {
- /// Creates a new [`FFI_SessionRef`].
+ /// Creates a new [`FFI_SessionRef`] with a default physical extension
codec.
+ ///
+ /// The synthesized [`DefaultPhysicalExtensionCodec`] supports built-in
physical
+ /// nodes only. A query planner obtained through this session reference
therefore
+ /// cannot encode or decode custom physical extension nodes. Use
+ /// [`Self::new_with_ffi_codecs`] with matching logical and physical
codecs when
+ /// custom physical nodes must cross the FFI boundary.
+ ///
+ /// The physical codec wrapper requires a
+ ///
[`FFI_TaskContextProvider`](crate::execution::FFI_TaskContextProvider), but this
+ /// constructor has only a session reference and a logical codec. It
therefore
+ /// reuses the logical codec's provider. The provider may be owned by
another
+ /// library; this is safe, but it must remain live and return the task
context
+ /// intended for codec callbacks. The default physical codec does not
successfully
+ /// decode extension nodes, so callers that need such callbacks must
instead use
+ /// [`Self::new_with_ffi_codecs`] with an explicitly configured physical
codec and
+ /// task context provider.
pub fn new(
- session: &(dyn Session + Send + Sync),
+ session: &dyn Session,
runtime: Option<Handle>,
logical_codec: FFI_LogicalExtensionCodec,
+ ) -> Self {
+ // `Session` provides a TaskContext but not the reference-counted
+ // TaskContextProvider needed by the FFI codec. Reuse the provider
associated
+ // with the logical codec under the assumptions documented above.
+ let physical_codec = FFI_PhysicalExtensionCodec::new(
+ Arc::new(DefaultPhysicalExtensionCodec {}),
+ runtime.clone(),
+ logical_codec.task_ctx_provider.clone(),
+ );
+ Self::new_with_ffi_codecs(session, runtime, logical_codec,
physical_codec)
+ }
+
+ /// Creates a new [`FFI_SessionRef`] using existing FFI codecs.
+ ///
+ /// The codecs must form a matching pair that can round-trip every logical
and
+ /// physical extension node exposed through the session. Their task context
+ /// providers must remain live and return contexts appropriate for their
decode
+ /// callbacks.
+ ///
+ /// If `session` is already foreign, this re-exports its original FFI
handle
+ /// rather than adding another wrapper layer. The handle adopts the codecs
+ /// supplied here while retaining its original private data and runtime.
+ pub fn new_with_ffi_codecs(
+ session: &dyn Session,
+ runtime: Option<Handle>,
+ logical_codec: FFI_LogicalExtensionCodec,
+ physical_codec: FFI_PhysicalExtensionCodec,
) -> Self {
if let Some(session) =
session.as_any().downcast_ref::<ForeignSession>() {
- return session.session.clone();
+ let mut session = session.session.clone();
+ session.logical_codec = logical_codec;
+ session.physical_codec = physical_codec;
+ return session;
}
let private_data = Box::new(SessionPrivateData { session, runtime });
@@ -366,6 +493,8 @@ impl FFI_SessionRef {
session_id: session_id_fn_wrapper,
config: config_fn_wrapper,
catalog_list: catalog_list_fn_wrapper,
+ query_planner: query_planner_fn_wrapper,
+ optimize: optimize_fn_wrapper,
create_physical_plan: create_physical_plan_fn_wrapper,
create_physical_expr: create_physical_expr_fn_wrapper,
scalar_functions: scalar_functions_fn_wrapper,
@@ -374,7 +503,9 @@ impl FFI_SessionRef {
table_options: table_options_fn_wrapper,
default_table_options: default_table_options_fn_wrapper,
task_ctx: task_ctx_fn_wrapper,
+ physical_optimizers: physical_optimizers_fn_wrapper,
logical_codec,
+ physical_codec,
clone: clone_fn_wrapper,
release: release_fn_wrapper,
@@ -387,8 +518,17 @@ impl FFI_SessionRef {
/// This wrapper struct exists on the receiver side of the FFI interface, so
it has
/// no guarantees about being able to access the data in `private_data`. Any
functions
-/// defined on this struct must only use the stable functions provided in
-/// FFI_Session to interact with the foreign table provider.
+/// defined on this struct must use only the stable function pointers in
+/// `FFI_SessionRef` to interact with the foreign session.
+///
+/// # Query planner delegation
+///
+/// If the session owner installed the current foreign query planner,
+/// [`Session::create_physical_plan`] dispatches back to that planner and
+/// [`Session::query_planner`] returns that planner. The planner must retain
and
+/// invoke the session owner's previous planner instead of using either method
to
+/// delegate back to the session. Otherwise, repeated delegation exhausts the
+/// stack. See [`crate::query_planner`] for details.
#[derive(Debug)]
pub struct ForeignSession {
session: FFI_SessionRef,
@@ -402,13 +542,15 @@ pub struct ForeignSession {
table_options: TableOptions,
runtime_env: Arc<RuntimeEnv>,
props: ExecutionProps,
+ query_planner: OnceLock<Arc<dyn QueryPlanner + Send + Sync>>,
+ physical_optimizers: OnceLock<Vec<Arc<dyn PhysicalOptimizerRule + Send +
Sync>>>,
}
unsafe impl Send for ForeignSession {}
unsafe impl Sync for ForeignSession {}
impl FFI_SessionRef {
- pub fn as_local(&self) -> Option<&(dyn Session + Send + Sync)> {
+ pub fn as_local(&self) -> Option<&dyn Session> {
if (self.library_marker_id)() == crate::get_library_marker_id() {
return Some(self.inner());
}
@@ -462,7 +604,6 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession {
)
})
.collect();
-
Ok(Self {
session: session.clone(),
config,
@@ -475,6 +616,8 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession {
extension_types:
Arc::new(MemoryExtensionTypeRegistry::default()),
runtime_env: Default::default(),
props: Default::default(),
+ query_planner: OnceLock::new(),
+ physical_optimizers: OnceLock::new(),
})
}
}
@@ -573,6 +716,31 @@ impl Session for ForeignSession {
Arc::clone(&self.catalog_list)
}
+ fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync> {
+ Arc::clone(self.query_planner.get_or_init(|| unsafe {
+ let planner = (self.session.query_planner)(&self.session);
+ (&planner).into()
+ }))
+ }
+
+ fn optimize(&self, plan: &LogicalPlan) ->
datafusion_common::Result<LogicalPlan> {
+ unsafe {
+ let codec: Arc<dyn LogicalExtensionCodec> =
+ (&self.session.logical_codec).into();
+ let logical_plan =
+ logical_plan_to_bytes_with_extension_codec(plan,
codec.as_ref())?;
+ let optimized_plan = df_result!((self.session.optimize)(
+ &self.session,
+ SVec::from(logical_plan.as_ref()),
+ ))?;
+ logical_plan_from_bytes_with_extension_codec(
+ optimized_plan.as_slice(),
+ self.task_ctx().as_ref(),
+ codec.as_ref(),
+ )
+ }
+ }
+
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
@@ -613,6 +781,15 @@ impl Session for ForeignSession {
}
}
+ fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send +
Sync>] {
+ self.physical_optimizers.get_or_init(|| unsafe {
+ (self.session.physical_optimizers)(&self.session)
+ .into_iter()
+ .map(|rule| (&rule).into())
+ .collect()
+ })
+ }
+
fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>> {
&self.scalar_functions
}
@@ -672,6 +849,7 @@ impl Session for ForeignSession {
#[cfg(test)]
mod tests {
use std::sync::Arc;
+ use std::sync::atomic::{AtomicUsize, Ordering};
use arrow_schema::{DataType, Field, Schema};
use datafusion::catalog::MemoryCatalogProvider;
@@ -683,6 +861,59 @@ mod tests {
use super::*;
+ static QUERY_PLANNER_CALLS: AtomicUsize = AtomicUsize::new(0);
+ static PHYSICAL_OPTIMIZER_CALLS: AtomicUsize = AtomicUsize::new(0);
+
+ unsafe extern "C" fn counting_query_planner(
+ session: &FFI_SessionRef,
+ ) -> FFI_QueryPlanner {
+ QUERY_PLANNER_CALLS.fetch_add(1, Ordering::Relaxed);
+ unsafe { query_planner_fn_wrapper(session) }
+ }
+
+ unsafe extern "C" fn counting_physical_optimizers(
+ session: &FFI_SessionRef,
+ ) -> SVec<FFI_PhysicalOptimizerRule> {
+ PHYSICAL_OPTIMIZER_CALLS.fetch_add(1, Ordering::Relaxed);
+ unsafe { physical_optimizers_fn_wrapper(session) }
+ }
+
+ #[test]
+ fn test_foreign_session_lazily_loads_planning_state() -> Result<(),
DataFusionError> {
+ QUERY_PLANNER_CALLS.store(0, Ordering::Relaxed);
+ PHYSICAL_OPTIMIZER_CALLS.store(0, Ordering::Relaxed);
+
+ let (ctx, task_ctx_provider) =
crate::util::tests::test_session_and_ctx();
+ let logical_codec = FFI_LogicalExtensionCodec::new(
+ Arc::new(DefaultLogicalExtensionCodec {}),
+ None,
+ task_ctx_provider,
+ );
+ let state = ctx.state();
+ let mut local_session = FFI_SessionRef::new(&state, None,
logical_codec);
+ local_session.query_planner = counting_query_planner;
+ local_session.physical_optimizers = counting_physical_optimizers;
+
+ let mut foreign_session = ForeignSession::try_from(&local_session)?;
+ assert_eq!(QUERY_PLANNER_CALLS.load(Ordering::Relaxed), 0);
+ assert_eq!(PHYSICAL_OPTIMIZER_CALLS.load(Ordering::Relaxed), 0);
+
+ // `FFI_SessionRef::clone` restores the standard function pointers, so
+ // instrument the clone retained by `ForeignSession` as well.
+ foreign_session.session.query_planner = counting_query_planner;
+ foreign_session.session.physical_optimizers =
counting_physical_optimizers;
+
+ foreign_session.query_planner();
+ foreign_session.query_planner();
+ assert_eq!(QUERY_PLANNER_CALLS.load(Ordering::Relaxed), 1);
+
+ foreign_session.physical_optimizers();
+ foreign_session.physical_optimizers();
+ assert_eq!(PHYSICAL_OPTIMIZER_CALLS.load(Ordering::Relaxed), 1);
+
+ Ok(())
+ }
+
#[tokio::test]
async fn test_ffi_session() -> Result<(), DataFusionError> {
let (ctx, task_ctx_provider) =
crate::util::tests::test_session_and_ctx();
@@ -731,14 +962,16 @@ mod tests {
let logical_plan = LogicalPlan::default();
assert_eq!(foreign_session.optimize(&logical_plan)?, logical_plan);
- assert!(foreign_session.physical_optimizers().is_empty());
+ assert_eq!(
+ foreign_session.physical_optimizers().len(),
+ state.physical_optimizers().len()
+ );
assert!(foreign_session.statistics_registry().is_none());
- let planner_error = foreign_session
+ let planned = foreign_session
.query_planner()
.create_physical_plan(&logical_plan, &foreign_session)
- .await
- .unwrap_err();
- assert!(planner_error.to_string().contains("does not expose"));
+ .await?;
+ assert_eq!(planned.name(), "EmptyExec");
let physical_plan =
foreign_session.create_physical_plan(&logical_plan).await?;
assert_eq!(
diff --git a/datafusion/ffi/src/table_provider.rs
b/datafusion/ffi/src/table_provider.rs
index 5a4b2fa272..ee9377bff0 100644
--- a/datafusion/ffi/src/table_provider.rs
+++ b/datafusion/ffi/src/table_provider.rs
@@ -263,7 +263,7 @@ unsafe extern "C" fn scan_fn_wrapper(
let session = sresult_return!(
session
.as_local()
- .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .map(Ok::<&dyn Session, DataFusionError>)
.unwrap_or_else(|| {
foreign_session =
Some(ForeignSession::try_from(&session)?);
Ok(foreign_session.as_ref().unwrap())
@@ -314,7 +314,7 @@ unsafe extern "C" fn insert_into_fn_wrapper(
let session = sresult_return!(
session
.as_local()
- .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .map(Ok::<&dyn Session, DataFusionError>)
.unwrap_or_else(|| {
foreign_session =
Some(ForeignSession::try_from(&session)?);
Ok(foreign_session.as_ref().unwrap())
diff --git a/datafusion/ffi/src/table_provider_factory.rs
b/datafusion/ffi/src/table_provider_factory.rs
index b70e72f31a..63ebb51bb1 100644
--- a/datafusion/ffi/src/table_provider_factory.rs
+++ b/datafusion/ffi/src/table_provider_factory.rs
@@ -211,7 +211,7 @@ async fn create_fn_wrapper_impl(
let mut foreign_session = None;
let session = session
.as_local()
- .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .map(Ok::<&dyn Session, DataFusionError>)
.unwrap_or_else(|| {
foreign_session = Some(ForeignSession::try_from(&session)?);
Ok(foreign_session.as_ref().unwrap())
diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs
index 59bcc861d0..357953674a 100644
--- a/datafusion/ffi/src/tests/mod.rs
+++ b/datafusion/ffi/src/tests/mod.rs
@@ -43,6 +43,8 @@ use crate::execution_plan::FFI_ExecutionPlan;
use crate::execution_plan::tests::EmptyExec;
use crate::physical_optimizer::FFI_PhysicalOptimizerRule;
use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use crate::query_planner::FFI_QueryPlanner;
use crate::table_provider::FFI_TableProvider;
use crate::table_provider_factory::FFI_TableProviderFactory;
use crate::tests::catalog::create_catalog_provider_list;
@@ -50,11 +52,13 @@ use crate::udaf::FFI_AggregateUDF;
use crate::udf::FFI_ScalarUDF;
use crate::udtf::FFI_TableFunction;
use crate::udwf::FFI_WindowUDF;
+use crate::util::FFI_Option;
mod async_provider;
pub mod catalog;
pub mod config;
mod physical_optimizer;
+mod query_planner;
mod sync_provider;
mod table_provider_factory;
mod udf_udaf_udwf;
@@ -117,6 +121,14 @@ pub struct ForeignLibraryModule {
pub create_context_aware_optimizer_rule: extern "C" fn() ->
FFI_PhysicalOptimizerRule,
+ /// Construct a query planner. When `library_a_planner` is provided the
+ /// planner delegates to it, as library C does after library A swaps
planners.
+ pub create_query_planner: extern "C" fn(
+ logical_codec: FFI_LogicalExtensionCodec,
+ physical_codec: FFI_PhysicalExtensionCodec,
+ library_a_planner: FFI_Option<FFI_QueryPlanner>,
+ ) -> FFI_QueryPlanner,
+
pub version: extern "C" fn() -> u64,
/// Create an aggregate UDAF using first_value
@@ -269,6 +281,7 @@ pub extern "C" fn datafusion_ffi_get_module() ->
ForeignLibraryModule {
physical_optimizer::create_physical_optimizer_rule,
create_context_aware_optimizer_rule:
physical_optimizer::create_context_aware_optimizer_rule,
+ create_query_planner: query_planner::create_query_planner,
version: super::version,
create_first_value_udaf: create_ffi_first_value_func,
}
diff --git a/datafusion/ffi/src/tests/query_planner.rs
b/datafusion/ffi/src/tests/query_planner.rs
new file mode 100644
index 0000000000..90d713f8a6
--- /dev/null
+++ b/datafusion/ffi/src/tests/query_planner.rs
@@ -0,0 +1,172 @@
+// 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::any::Any;
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema};
+use async_trait::async_trait;
+use datafusion_catalog::default_table_source::source_as_provider;
+use datafusion_common::{Result, exec_err};
+use datafusion_expr::LogicalPlan;
+use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_plan::empty::EmptyExec;
+use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::union::UnionExec;
+use datafusion_session::{QueryPlanner, Session};
+
+use crate::execution_plan::ForeignExecutionPlan;
+use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use crate::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner};
+use crate::session::ForeignSession;
+use crate::table_provider::ForeignTableProvider;
+use crate::util::FFI_Option;
+
+#[derive(Debug)]
+struct TestQueryPlanner;
+
+#[async_trait]
+impl QueryPlanner for TestQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if let LogicalPlan::TableScan(scan) = logical_plan {
+ if session.as_any().downcast_ref::<ForeignSession>().is_none() {
+ return exec_err!("library A's session was not foreign to
library C");
+ }
+
+ let provider = source_as_provider(&scan.source)?;
+ if provider.downcast_ref::<ForeignTableProvider>().is_none() {
+ return exec_err!("library B's provider was not foreign to
library C");
+ }
+ let library_b_plan = provider
+ .scan(session, scan.projection.as_ref(), &scan.filters,
scan.fetch)
+ .await?;
+
+ if !library_b_plan.is::<ForeignExecutionPlan>() {
+ return exec_err!("library B's plan unexpectedly downcast as
C-local");
+ }
+
+ let plan = UnionExec::try_new(vec![
+ Arc::clone(&library_b_plan),
+ Arc::clone(&library_b_plan),
+ ])?;
+ if !plan.is::<UnionExec>() {
+ return exec_err!("library C could not downcast its local
UnionExec");
+ }
+ return Ok(plan);
+ }
+
+ let query_planner = session.query_planner();
+ let planner_any: &dyn Any = query_planner.as_ref();
+ if planner_any.downcast_ref::<ForeignQueryPlanner>().is_none() {
+ return exec_err!("query planner did not cross the FFI boundary");
+ }
+ session.optimize(logical_plan)?;
+ if session.physical_optimizers().is_empty() {
+ return exec_err!("physical optimizers did not cross the FFI
boundary");
+ }
+
+ let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::Int32, true)]));
+ Ok(Arc::new(EmptyExec::new(schema)))
+ }
+}
+
+/// Library C's planner for the planner-swap deployment.
+///
+/// It holds the query planner library A exported *before* A swapped this
planner
+/// into its session, so delegating to it cannot re-enter library C.
+#[derive(Debug)]
+struct SwappedQueryPlanner {
+ library_a_planner: Arc<dyn QueryPlanner + Send + Sync>,
+}
+
+#[async_trait]
+impl QueryPlanner for SwappedQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if session.as_any().downcast_ref::<ForeignSession>().is_none() {
+ return exec_err!("library A's session was not foreign to library
C");
+ }
+
+ // After the swap, the planner installed on library A's session is this
+ // planner, so `session.query_planner()` and
`session.create_physical_plan()`
+ // are both self-references. Assert the hazard instead of triggering
it:
+ // calling either would recurse until the stack is exhausted.
+ let installed = session.query_planner();
+ let installed: &dyn Any = installed.as_ref();
+ if installed.downcast_ref::<Self>().is_none() {
+ return exec_err!(
+ "expected the swapped session to report library C's own
planner"
+ );
+ }
+
+ // Delegate to library A. The result crosses the FFI boundary as
+ // serialized bytes, so library C receives nodes carrying its own local
+ // Rust type identities.
+ let plan = self
+ .library_a_planner
+ .create_physical_plan(logical_plan, session)
+ .await?;
+
+ if plan.is::<ForeignExecutionPlan>() {
+ return exec_err!("library A's plan was opaque to library C");
+ }
+ let Some(sort) = plan.downcast_ref::<SortExec>() else {
+ return exec_err!(
+ "library C could not downcast library A's SortExec; got {}",
+ plan.name()
+ );
+ };
+ // Library B's scan is still foreign to library C. Only a codec
boundary
+ // reconstructs it, and library A's codec hands back an A-local node.
+ if !sort.input().is::<ForeignExecutionPlan>() {
+ return exec_err!("library B's scan unexpectedly downcast as
C-local");
+ }
+
+ Ok(UnionExec::try_new(vec![
+ Arc::clone(&plan),
+ Arc::clone(&plan),
+ ])?)
+ }
+}
+
+/// Creates library C's query planner.
+///
+/// `library_a_planner` is the planner library A exported before swapping this
one
+/// onto its session. When it is absent the planner does its own planning
instead
+/// of delegating.
+pub extern "C" fn create_query_planner(
+ logical_codec: FFI_LogicalExtensionCodec,
+ physical_codec: FFI_PhysicalExtensionCodec,
+ library_a_planner: FFI_Option<FFI_QueryPlanner>,
+) -> FFI_QueryPlanner {
+ let planner: Arc<dyn QueryPlanner + Send + Sync> = match
library_a_planner.as_ref() {
+ Some(library_a_planner) => Arc::new(SwappedQueryPlanner {
+ library_a_planner: library_a_planner.into(),
+ }),
+ None => Arc::new(TestQueryPlanner),
+ };
+
+ FFI_QueryPlanner::new_with_ffi_codecs(planner, logical_codec,
physical_codec)
+}
diff --git a/datafusion/ffi/src/tests/utils.rs
b/datafusion/ffi/src/tests/utils.rs
index b6b50cbce8..3119ab96d4 100644
--- a/datafusion/ffi/src/tests/utils.rs
+++ b/datafusion/ffi/src/tests/utils.rs
@@ -62,14 +62,12 @@ fn find_library() -> Result<PathBuf> {
find_cdylib(deps_dir)
}
-pub fn get_module() -> Result<ForeignLibraryModule> {
+fn load_module(lib_path: &Path) -> Result<ForeignLibraryModule> {
let expected_version = crate::version();
- let lib_path = find_library()?;
-
// Load the library using libloading
let lib = unsafe {
- libloading::Library::new(&lib_path)
+ libloading::Library::new(lib_path)
.map_err(|e| DataFusionError::External(Box::new(e)))?
};
@@ -87,3 +85,40 @@ pub fn get_module() -> Result<ForeignLibraryModule> {
Ok(module)
}
+
+pub fn get_module() -> Result<ForeignLibraryModule> {
+ load_module(&find_library()?)
+}
+
+/// Load an independent copy of the integration-test cdylib.
+///
+/// Copying to a unique path makes the dynamic loader create a separate image
+/// with its own library marker and Rust object graph.
+pub fn get_module_copy(name: &str) -> Result<ForeignLibraryModule> {
+ let source = find_library()?;
+ let file_name = source
+ .file_name()
+ .and_then(|name| name.to_str())
+ .ok_or_else(|| DataFusionError::External("Invalid cdylib
filename".into()))?;
+ // Windows cannot remove a loaded DLL, so use a stable name that bounds the
+ // retained test artifacts to one file per library role.
+ #[cfg(target_os = "windows")]
+ let destination = source.with_file_name(format!("{name}_{file_name}"));
+ #[cfg(not(target_os = "windows"))]
+ let destination =
+ source.with_file_name(format!("{}_{}_{}", std::process::id(), name,
file_name));
+
+ std::fs::copy(&source, &destination)
+ .map_err(|e| DataFusionError::External(Box::new(e)))?;
+ match load_module(&destination) {
+ Ok(module) => {
+ #[cfg(not(target_os = "windows"))]
+ let _ = std::fs::remove_file(destination);
+ Ok(module)
+ }
+ Err(error) => {
+ let _ = std::fs::remove_file(destination);
+ Err(error)
+ }
+ }
+}
diff --git a/datafusion/ffi/src/udtf.rs b/datafusion/ffi/src/udtf.rs
index 0a11102879..fa28519d58 100644
--- a/datafusion/ffi/src/udtf.rs
+++ b/datafusion/ffi/src/udtf.rs
@@ -155,7 +155,7 @@ unsafe extern "C" fn call_with_args_wrapper(
let session = sresult_return!(
session
.as_local()
- .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .map(Ok::<&dyn Session, DataFusionError>)
.unwrap_or_else(|| {
foreign_session = Some(ForeignSession::try_from(&session)?);
Ok(foreign_session.as_ref().unwrap())
diff --git a/datafusion/ffi/tests/ffi_query_planner.rs
b/datafusion/ffi/tests/ffi_query_planner.rs
new file mode 100644
index 0000000000..c72b8c3e88
--- /dev/null
+++ b/datafusion/ffi/tests/ffi_query_planner.rs
@@ -0,0 +1,348 @@
+// 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.
+
+mod utils;
+
+#[cfg(feature = "integration-tests")]
+mod tests {
+ use std::sync::{Arc, OnceLock, Weak};
+
+ use arrow::datatypes::SchemaRef;
+ use datafusion::execution::SessionStateBuilder;
+ use datafusion::prelude::SessionContext;
+ use datafusion_catalog::TableProvider;
+ use datafusion_common::{
+ DataFusionError, Result, TableReference, exec_err, not_impl_err,
+ };
+ use datafusion_execution::{TaskContext, TaskContextProvider};
+ use datafusion_expr::logical_plan::Extension;
+ use datafusion_expr::{LogicalPlan, col};
+ use datafusion_ffi::execution::FFI_TaskContextProvider;
+ use datafusion_ffi::execution_plan::ForeignExecutionPlan;
+ 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, ForeignQueryPlanner};
+ use datafusion_ffi::table_provider::ForeignTableProvider;
+ use datafusion_ffi::tests::{
+ create_test_schema,
+ utils::{get_module, get_module_copy},
+ };
+ use datafusion_ffi::util::FFI_Option;
+ use datafusion_physical_plan::ExecutionPlan;
+ use datafusion_physical_plan::empty::EmptyExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_proto::logical_plan::LogicalExtensionCodec;
+ use datafusion_proto::physical_plan::{
+ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec,
+ PhysicalProtoConverterExtension,
+ };
+ use datafusion_session::QueryPlanner;
+
+ #[tokio::test]
+ async fn test_ffi_query_planner() -> Result<(), DataFusionError> {
+ let module = get_module()?;
+ let (ctx, logical_codec) = crate::utils::ctx_and_codec();
+ let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn
TaskContextProvider>;
+ let task_ctx_provider =
FFI_TaskContextProvider::from(&task_ctx_provider);
+ let physical_codec = FFI_PhysicalExtensionCodec::new(
+ Arc::new(DefaultPhysicalExtensionCodec {}),
+ None,
+ task_ctx_provider,
+ );
+
+ let ffi_planner = (module.create_query_planner)(
+ logical_codec,
+ physical_codec,
+ FFI_Option::None,
+ );
+ let planner: Arc<dyn QueryPlanner + Send + Sync> =
(&ffi_planner).into();
+
+ let any_ref: &dyn std::any::Any = planner.as_ref();
+ assert!(any_ref.downcast_ref::<ForeignQueryPlanner>().is_some());
+
+ let logical_plan =
datafusion_expr::LogicalPlanBuilder::empty(false).build()?;
+ let state = ctx.state();
+ let physical_plan = planner.create_physical_plan(&logical_plan,
&state).await?;
+
+ assert_eq!(physical_plan.name(), "EmptyExec");
+ assert!(physical_plan.is::<EmptyExec>());
+
+ Ok(())
+ }
+
+ /// Test-only codec that preserves library B's table provider while the
logical
+ /// plan crosses between library A and library C.
+ ///
+ /// Encoding writes a fixed identifier and stores a weak reference to the
+ /// provider. Decoding validates the identifier and upgrades that
reference.
+ /// This works because all three test libraries run in one process and
library
+ /// A's session continues to own the provider.
+ ///
+ /// This is not a general serialization format for table providers. A
+ /// cross-process deployment must provide its own codec that either
resolves a
+ /// stable identifier through shared state or reconstructs the provider
from a
+ /// portable, provider-specific description. DataFusion passes the table
+ /// reference, schema, and task context separately to the decoder.
+ #[derive(Debug, Default)]
+ struct LibraryALogicalCodec {
+ library_b_provider: OnceLock<Weak<dyn TableProvider>>,
+ }
+
+ impl LogicalExtensionCodec for LibraryALogicalCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[LogicalPlan],
+ _ctx: &TaskContext,
+ ) -> Result<Extension> {
+ not_impl_err!("logical extension nodes are not used in this test")
+ }
+
+ fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) ->
Result<()> {
+ not_impl_err!("logical extension nodes are not used in this test")
+ }
+
+ fn try_decode_table_provider(
+ &self,
+ buf: &[u8],
+ _table_ref: &TableReference,
+ _schema: SchemaRef,
+ _ctx: &TaskContext,
+ ) -> Result<Arc<dyn TableProvider>> {
+ if buf != b"library-b-provider" {
+ return exec_err!("unexpected library B provider payload");
+ }
+ self.library_b_provider
+ .get()
+ .and_then(Weak::upgrade)
+ .ok_or_else(|| DataFusionError::Plan("missing library B
provider".into()))
+ }
+
+ fn try_encode_table_provider(
+ &self,
+ _table_ref: &TableReference,
+ node: Arc<dyn TableProvider>,
+ buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ self.library_b_provider
+ .get_or_init(|| Arc::downgrade(&node));
+ buf.extend_from_slice(b"library-b-provider");
+ Ok(())
+ }
+ }
+
+ /// Library A's physical codec reconstructs B's opaque foreign plan as an
+ /// A-local test plan when the result returns from library C.
+ ///
+ /// Encoding sees B's node in one of two shapes. When A serializes a plan
it
+ /// built itself, B's scan is a [`ForeignExecutionPlan`]. When library C
+ /// serializes a plan containing a node A previously handed it, the FFI
handle
+ /// unwraps back to its home library, so A is asked to encode the very
+ /// [`EmptyExec`] its own `try_decode` produced.
+ #[derive(Debug)]
+ struct LibraryAPhysicalCodec;
+
+ impl PhysicalExtensionCodec for LibraryAPhysicalCodec {
+ fn try_decode(
+ &self,
+ buf: &[u8],
+ inputs: &[Arc<dyn ExecutionPlan>],
+ _ctx: &TaskContext,
+ _proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if buf != b"library-b-empty-exec" || !inputs.is_empty() {
+ return exec_err!("unexpected library B execution plan
payload");
+ }
+ Ok(Arc::new(EmptyExec::new(create_test_schema())))
+ }
+
+ fn try_encode(
+ &self,
+ node: Arc<dyn ExecutionPlan>,
+ buf: &mut Vec<u8>,
+ _proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<()> {
+ if !node.is::<ForeignExecutionPlan>() && !node.is::<EmptyExec>() {
+ return exec_err!(
+ "expected library B's plan to be foreign or A-local; got
{}",
+ node.name()
+ );
+ }
+ buf.extend_from_slice(b"library-b-empty-exec");
+ Ok(())
+ }
+ }
+
+ #[tokio::test]
+ async fn test_three_library_query_planner_restores_type_identity() ->
Result<()> {
+ // Library A: datafusion-python owns the session and codec registry.
+ let state = SessionStateBuilder::new_with_default_features()
+ .with_physical_optimizer_rules(vec![])
+ .build();
+ let ctx = Arc::new(SessionContext::new_with_state(state));
+ let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn
TaskContextProvider>;
+ let ffi_task_ctx_provider =
FFI_TaskContextProvider::from(&task_ctx_provider);
+ let logical_codec = FFI_LogicalExtensionCodec::new(
+ Arc::new(LibraryALogicalCodec::default()),
+ None,
+ ffi_task_ctx_provider.clone(),
+ );
+ let physical_codec = FFI_PhysicalExtensionCodec::new(
+ Arc::new(LibraryAPhysicalCodec),
+ None,
+ ffi_task_ctx_provider,
+ );
+
+ let library_b = get_module_copy("query_planner_library_b")?;
+ let library_c = get_module_copy("query_planner_library_c")?;
+
+ // Library B: reuse the synchronous table provider from the existing
+ // FFI integration-test module.
+ let ffi_provider = (library_b.create_table)(true,
logical_codec.clone());
+ let provider: Arc<dyn TableProvider> = (&ffi_provider).into();
+ assert!(provider.downcast_ref::<ForeignTableProvider>().is_some());
+ ctx.register_table("library_b", provider)?;
+ let logical_plan =
ctx.table("library_b").await?.into_optimized_plan()?;
+
+ // Library C: a foreign query planner sees B's scan result as opaque,
+ // but can downcast its own UnionExec. Its result is serialized rather
+ // than returned as FFI_ExecutionPlan.
+ let ffi_planner = (library_c.create_query_planner)(
+ logical_codec,
+ physical_codec,
+ FFI_Option::None,
+ );
+ let planner: Arc<dyn QueryPlanner + Send + Sync> =
(&ffi_planner).into();
+ let planner_any: &dyn std::any::Any = planner.as_ref();
+ assert!(planner_any.downcast_ref::<ForeignQueryPlanner>().is_some());
+
+ let state = ctx.state();
+ let physical_plan = planner.create_physical_plan(&logical_plan,
&state).await?;
+
+ // Deserialization in A reconstructs the full result as A-local
+ // concrete nodes, including the plans that originated in B.
+ assert!(physical_plan.is::<UnionExec>());
+ assert!(!physical_plan.is::<ForeignExecutionPlan>());
+ let children = physical_plan.children();
+ assert_eq!(children.len(), 2);
+ assert!(children.iter().all(|child| child.is::<EmptyExec>()));
+ assert!(
+ children
+ .iter()
+ .all(|child| !child.is::<ForeignExecutionPlan>())
+ );
+
+ Ok(())
+ }
+
+ /// Exercises the deployment library C actually uses: library A hands its
own
+ /// query planner to C, then installs C's planner on the session it already
+ /// owns. C plans by delegating back to A's captured planner.
+ ///
+ /// This is the case that requires serialized plans in both directions. C
must
+ /// downcast the nodes A produced in order to rewrite them, and A must
downcast
+ /// the nodes C produced in order to run its own passes over the result.
+ #[tokio::test]
+ async fn test_query_planner_swap_round_trips_type_identity() -> Result<()>
{
+ // Library A: datafusion-python owns the session and codec registry.
The
+ // physical optimizer rules are cleared so the assertions below observe
+ // planning alone.
+ let state = SessionStateBuilder::new_with_default_features()
+ .with_physical_optimizer_rules(vec![])
+ .build();
+ let ctx = Arc::new(SessionContext::new_with_state(state));
+ let task_ctx_provider = Arc::clone(&ctx) as Arc<dyn
TaskContextProvider>;
+ let ffi_task_ctx_provider =
FFI_TaskContextProvider::from(&task_ctx_provider);
+ let logical_codec = FFI_LogicalExtensionCodec::new(
+ Arc::new(LibraryALogicalCodec::default()),
+ None,
+ ffi_task_ctx_provider.clone(),
+ );
+ let physical_codec = FFI_PhysicalExtensionCodec::new(
+ Arc::new(LibraryAPhysicalCodec),
+ None,
+ ffi_task_ctx_provider,
+ );
+
+ let library_b = get_module_copy("planner_swap_library_b")?;
+ let library_c = get_module_copy("planner_swap_library_c")?;
+
+ // Library B: a table provider that is foreign to both A and C.
+ let ffi_provider = (library_b.create_table)(true,
logical_codec.clone());
+ let provider: Arc<dyn TableProvider> = (&ffi_provider).into();
+ ctx.register_table("library_b", provider)?;
+
+ // Library A exports its default planner *before* the swap. Fetching it
+ // afterwards through `FFI_SessionRef::query_planner` would hand
library C
+ // its own planner back.
+ let library_a_planner = Arc::clone(ctx.state().query_planner());
+ let ffi_library_a_planner = FFI_QueryPlanner::new_with_ffi_codecs(
+ library_a_planner,
+ logical_codec.clone(),
+ physical_codec.clone(),
+ );
+
+ // Library C: builds its planner around A's planner.
+ let ffi_planner = (library_c.create_query_planner)(
+ logical_codec,
+ physical_codec,
+ FFI_Option::Some(ffi_library_a_planner),
+ );
+ let library_c_planner: Arc<dyn QueryPlanner + Send + Sync> =
+ (&ffi_planner).into();
+ let planner_any: &dyn std::any::Any = library_c_planner.as_ref();
+ assert!(planner_any.downcast_ref::<ForeignQueryPlanner>().is_some());
+
+ // Library A swaps C's planner into the session it already owns.
Mutating
+ // the existing state keeps the `Arc<SessionContext>` identity stable,
so
+ // the task context provider captured by the codecs above stays
current.
+ let state_ref = ctx.state_ref();
+ let swapped =
SessionStateBuilder::new_from_existing(state_ref.read().clone())
+ .with_query_planner(library_c_planner)
+ .build();
+ *state_ref.write() = swapped;
+
+ // A sort keeps a well-known, non-extension node at the root of A's
+ // physical plan. A projection or limit would be pushed into the scan,
+ // leaving only library B's opaque node for C to inspect.
+ let logical_plan = ctx
+ .table("library_b")
+ .await?
+ .sort(vec![col("a").sort(true, true)])?
+ .into_optimized_plan()?;
+
+ // Planning now runs A -> C -> A -> C -> A across three library images.
+ let physical_plan =
ctx.state().create_physical_plan(&logical_plan).await?;
+
+ // Library A reconstructs C's result as A-local concrete nodes,
including
+ // the plan that originated in B.
+ assert!(physical_plan.is::<UnionExec>());
+ assert!(!physical_plan.is::<ForeignExecutionPlan>());
+ let children = physical_plan.children();
+ assert_eq!(children.len(), 2);
+ for child in &children {
+ let sort = child
+ .downcast_ref::<SortExec>()
+ .expect("library A could not downcast the SortExec it
planned");
+ assert!(sort.input().is::<EmptyExec>());
+ assert!(!sort.input().is::<ForeignExecutionPlan>());
+ }
+
+ Ok(())
+ }
+}
diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md
b/docs/source/library-user-guide/upgrading/55.0.0.md
index e8e7f4a68f..2811fb4df2 100644
--- a/docs/source/library-user-guide/upgrading/55.0.0.md
+++ b/docs/source/library-user-guide/upgrading/55.0.0.md
@@ -925,15 +925,28 @@ fn physical_optimizers(&self) -> &[Arc<dyn
PhysicalOptimizerRule + Send + Sync>]
}
```
-`ForeignSession::create_physical_plan` continues to run the complete planning
-pipeline in the library that owns the session. `ForeignSession::query_planner`
-returns `UnsupportedQueryPlanner` until the query planner FFI interface is
-available. FFI wrappers for the individual planner and optimizer interfaces are
-not included in this release.
+`ForeignSession::create_physical_plan` runs the complete planning pipeline in
the
+library that owns the session. `ForeignSession::query_planner`, `optimize`, and
+`physical_optimizers` forward to the owning session across the FFI boundary. A
+foreign query planner can also be installed on a session through the new
+`datafusion_ffi::query_planner::FFI_QueryPlanner`; see that module's
+documentation for how plans and extension codecs cross the boundary.
See [PR #23703](https://github.com/apache/datafusion/pull/23703) for details on
the catalog changes.
+### `FFI_LogicalExtensionCodec::task_ctx_provider` is now private
+
+The `task_ctx_provider` field on
+`datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec` was
+`pub` and is now crate-private, matching `FFI_PhysicalExtensionCodec`.
+
+**Who is affected:**
+
+- Code that read or cloned `FFI_LogicalExtensionCodec::task_ctx_provider`
+ directly. Pass the task context provider to `FFI_LogicalExtensionCodec::new`
+ instead, and keep your own copy if you need it elsewhere.
+
### Unused `async` removed from several public functions
Public functions that were declared `async` but never awaited anything are now
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
