From ba0a3dd07d3d7b21f5cce6eea0e0a81c00929a33 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 31 Jul 2026 08:31:39 -0400 Subject: [PATCH 01/17] Add clippy exception so that no warnings are generated when parquet_encryption feature is not enabled --- datafusion/datasource-parquet/src/file_format.rs | 1 + datafusion/datasource-parquet/src/opener/encryption.rs | 1 + datafusion/datasource-parquet/src/sink.rs | 1 + 3 files changed, 3 insertions(+) diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index 29083ebfb2e72..8d19bedeb7155 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 b725198237bbf..498fe8acf7530 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 df2f17c6be22d..ceac70938c642 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -239,6 +239,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, From 74842d9655b1cc7821d12298b32a81bb6408750a Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 31 Jul 2026 08:40:48 -0400 Subject: [PATCH 02/17] feat: add FFI query planner support Add FFI_QueryPlanner and ForeignQueryPlanner with logical and physical plan codec support. Forward query planners, logical optimization, and physical optimizer rules through foreign sessions, with unit and cross-library coverage.\n\nAI Disclosure: This code was written in part by an AI agent. --- datafusion/ffi/src/lib.rs | 1 + .../ffi/src/proto/physical_extension_codec.rs | 2 +- datafusion/ffi/src/query_planner.rs | 358 ++++++++++++++++++ datafusion/ffi/src/session/mod.rs | 131 ++++++- datafusion/ffi/src/table_provider.rs | 6 +- datafusion/ffi/src/table_provider_factory.rs | 3 +- datafusion/ffi/src/tests/mod.rs | 9 + datafusion/ffi/src/tests/query_planner.rs | 66 ++++ datafusion/ffi/src/udtf.rs | 1 + datafusion/ffi/tests/ffi_query_planner.rs | 61 +++ 10 files changed, 626 insertions(+), 12 deletions(-) create mode 100644 datafusion/ffi/src/query_planner.rs create mode 100644 datafusion/ffi/src/tests/query_planner.rs create mode 100644 datafusion/ffi/tests/ffi_query_planner.rs diff --git a/datafusion/ffi/src/lib.rs b/datafusion/ffi/src/lib.rs index fd2ac58576b09..de8f8cba9ca9b 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/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index 9e64df82e31b4..bf45013e12645 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>, /// 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. diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs new file mode 100644 index 0000000000000..cb1f3e2f9d8d4 --- /dev/null +++ b/datafusion/ffi/src/query_planner.rs @@ -0,0 +1,358 @@ +// 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::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::{ + DefaultLogicalExtensionCodec, LogicalExtensionCodec, +}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, 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}; + +/// A stable struct for sharing [`QueryPlanner`] across FFI boundaries. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_QueryPlanner { + create_physical_plan: unsafe extern "C" fn( + &Self, + logical_plan_serialized: SVec, + session: FFI_SessionRef, + ) -> FfiFuture>>, + + pub logical_codec: FFI_LogicalExtensionCodec, + + pub 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, +} + +impl FFI_QueryPlanner { + fn inner(&self) -> &Arc { + 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, + session: FFI_SessionRef, +) -> FfiFuture>> { + let internal_planner = Arc::clone(planner.inner()); + let logical_codec: Arc = (&planner.logical_codec).into(); + let physical_codec: Arc = + (&planner.physical_codec).into(); + + async move { + let mut foreign_session = None; + let session = sresult_return!( + session + .as_local() + .map(Ok::<&(dyn Session + Send + Sync), 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(physical_plan.iter().copied().collect()) + } + .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 a new [`FFI_QueryPlanner`]. + pub fn new( + planner: Arc, + runtime: Option, + task_ctx_provider: impl Into, + logical_codec: Option>, + physical_codec: Option>, + ) -> Self { + let task_ctx_provider = task_ctx_provider.into(); + let logical_codec = + logical_codec.unwrap_or_else(|| Arc::new(DefaultLogicalExtensionCodec {})); + let logical_codec = FFI_LogicalExtensionCodec::new( + logical_codec, + runtime.clone(), + task_ctx_provider.clone(), + ); + let physical_codec = + physical_codec.unwrap_or_else(|| Arc::new(DefaultPhysicalExtensionCodec {})); + let physical_codec = + FFI_PhysicalExtensionCodec::new(physical_codec, runtime, task_ctx_provider); + Self::new_with_ffi_codecs(planner, logical_codec, physical_codec) + } + + pub fn new_with_ffi_codecs( + planner: Arc, + 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::() { + return planner.0.clone(); + } + + 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, + } + } + + /// Calls this query planner with a [`Session`] exported over FFI using + /// `session_runtime` as that session provider's local Tokio runtime. + pub async fn create_physical_plan_with_session_runtime( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + session_runtime: Option, + ) -> Result> { + let codec: Arc = (&self.logical_codec).into(); + let logical_plan = + logical_plan_to_bytes_with_extension_codec(logical_plan, codec.as_ref())?; + let logical_plan = logical_plan.iter().copied().collect(); + let task_ctx = session.task_ctx(); + let session = FFI_SessionRef::new( + session, + session_runtime, + self.logical_codec.clone(), + Some(self.physical_codec.clone()), + ); + + let physical_plan = unsafe { + df_result!((self.create_physical_plan)(self, logical_plan, session).await)? + }; + let physical_codec: Arc = + (&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 wrapper for a foreign [`QueryPlanner`]. +#[derive(Debug)] +pub struct ForeignQueryPlanner(pub FFI_QueryPlanner); + +unsafe impl Send for ForeignQueryPlanner {} +unsafe impl Sync for ForeignQueryPlanner {} + +impl From<&FFI_QueryPlanner> for Arc { + 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> { + 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 super::*; + + #[derive(Debug)] + struct EmptyQueryPlanner; + + #[async_trait] + impl QueryPlanner for EmptyQueryPlanner { + async fn create_physical_plan( + &self, + _logical_plan: &LogicalPlan, + _session: &dyn Session, + ) -> Result> { + 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) -> FFI_QueryPlanner { + let task_ctx_provider = Arc::clone(&ctx) as Arc; + FFI_QueryPlanner::new( + Arc::new(EmptyQueryPlanner), + None, + &task_ctx_provider, + None, + None, + ) + } + + #[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 = (&ffi_planner).into(); + let any_ref: &dyn std::any::Any = planner.as_ref(); + assert!(any_ref.downcast_ref::().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 = (&ffi_planner).into(); + let any_ref: &dyn std::any::Any = planner.as_ref(); + assert!(any_ref.downcast_ref::().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::()); + + Ok(()) + } +} diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 519384379edb8..a97ef97dca131 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -37,12 +37,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 +61,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 +95,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, + ) -> FFI_Result>, + create_physical_plan: unsafe extern "C" fn( &Self, @@ -110,8 +126,12 @@ pub(crate) struct FFI_SessionRef { task_ctx: unsafe extern "C" fn(&Self) -> FFI_TaskContext, + physical_optimizers: unsafe extern "C" fn(&Self) -> SVec, + 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, @@ -173,6 +193,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, +) -> FFI_Result> { + let logical_codec: Arc = (&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(optimized_plan.iter().copied().collect()) +} + unsafe extern "C" fn create_physical_plan_fn_wrapper( session: &FFI_SessionRef, logical_plan_serialized: SVec, @@ -303,6 +353,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 { + 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 +386,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 +396,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, @@ -355,17 +421,27 @@ impl FFI_SessionRef { session: &(dyn Session + Send + Sync), runtime: Option, logical_codec: FFI_LogicalExtensionCodec, + physical_codec: Option, ) -> Self { if let Some(session) = session.as_any().downcast_ref::() { return session.session.clone(); } + let physical_codec = physical_codec.unwrap_or_else(|| { + FFI_PhysicalExtensionCodec::new( + Arc::new(DefaultPhysicalExtensionCodec {}), + runtime.clone(), + logical_codec.task_ctx_provider.clone(), + ) + }); let private_data = Box::new(SessionPrivateData { session, runtime }); Self { 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 +450,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, @@ -402,6 +480,8 @@ pub struct ForeignSession { table_options: TableOptions, runtime_env: Arc, props: ExecutionProps, + query_planner: Arc, + physical_optimizers: Vec>, } unsafe impl Send for ForeignSession {} @@ -462,6 +542,11 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { ) }) .collect(); + let query_planner = (&(session.query_planner)(session)).into(); + let physical_optimizers = (session.physical_optimizers)(session) + .into_iter() + .map(|rule| (&rule).into()) + .collect(); Ok(Self { session: session.clone(), @@ -475,6 +560,8 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { extension_types: Arc::new(MemoryExtensionTypeRegistry::default()), runtime_env: Default::default(), props: Default::default(), + query_planner, + physical_optimizers, }) } } @@ -573,6 +660,28 @@ impl Session for ForeignSession { Arc::clone(&self.catalog_list) } + fn query_planner(&self) -> Arc { + Arc::clone(&self.query_planner) + } + + fn optimize(&self, plan: &LogicalPlan) -> datafusion_common::Result { + unsafe { + let codec: Arc = + (&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, + logical_plan.iter().copied().collect(), + ))?; + 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 +722,10 @@ impl Session for ForeignSession { } } + fn physical_optimizers(&self) -> &[Arc] { + &self.physical_optimizers + } + fn scalar_functions(&self) -> &HashMap> { &self.scalar_functions } @@ -705,7 +818,7 @@ mod tests { task_ctx_provider, ); - let local_session = FFI_SessionRef::new(&state, None, logical_codec); + let local_session = FFI_SessionRef::new(&state, None, logical_codec, None); let foreign_session = ForeignSession::try_from(&local_session)?; let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); @@ -731,14 +844,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 5a4b2fa27256f..3bb8409c3672c 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -492,7 +492,8 @@ impl TableProvider for ForeignTableProvider { filters: &[Expr], limit: Option, ) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + let session = + FFI_SessionRef::new(session, None, self.0.logical_codec.clone(), None); let projections: FFI_Option> = projection .map(|p| p.iter().map(|v| v.to_owned()).collect()) @@ -562,7 +563,8 @@ impl TableProvider for ForeignTableProvider { input: Arc, insert_op: InsertOp, ) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + let session = + FFI_SessionRef::new(session, None, self.0.logical_codec.clone(), None); let rc = Handle::try_current().ok(); let input = FFI_ExecutionPlan::new(input, rc); diff --git a/datafusion/ffi/src/table_provider_factory.rs b/datafusion/ffi/src/table_provider_factory.rs index b70e72f31aa4d..b1a5036685035 100644 --- a/datafusion/ffi/src/table_provider_factory.rs +++ b/datafusion/ffi/src/table_provider_factory.rs @@ -293,7 +293,8 @@ impl TableProviderFactory for ForeignTableProviderFactory { session: &dyn Session, cmd: &CreateExternalTable, ) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + let session = + FFI_SessionRef::new(session, None, self.0.logical_codec.clone(), None); let cmd = self.serialize_cmd(cmd.clone())?; let provider = unsafe { diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index d372dcf9177e6..5c3d71adc91ae 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -42,6 +42,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; @@ -54,6 +56,7 @@ 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; @@ -116,6 +119,11 @@ pub struct ForeignLibraryModule { pub create_context_aware_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, + pub create_query_planner: extern "C" fn( + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + ) -> FFI_QueryPlanner, + pub version: extern "C" fn() -> u64, } @@ -265,6 +273,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, } } diff --git a/datafusion/ffi/src/tests/query_planner.rs b/datafusion/ffi/src/tests/query_planner.rs new file mode 100644 index 0000000000000..f2eda6d8ff620 --- /dev/null +++ b/datafusion/ffi/src/tests/query_planner.rs @@ -0,0 +1,66 @@ +// 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 arrow::datatypes::{DataType, Field, Schema}; +use async_trait::async_trait; +use datafusion_common::{Result, exec_err}; +use datafusion_expr::LogicalPlan; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::empty::EmptyExec; +use datafusion_session::{QueryPlanner, Session}; + +use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner}; + +#[derive(Debug)] +struct EmptyQueryPlanner; + +#[async_trait] +impl QueryPlanner for EmptyQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + let query_planner = session.query_planner(); + let planner_any: &dyn std::any::Any = query_planner.as_ref(); + if planner_any.downcast_ref::().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))) + } +} + +pub extern "C" fn create_query_planner( + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, +) -> FFI_QueryPlanner { + FFI_QueryPlanner::new_with_ffi_codecs( + Arc::new(EmptyQueryPlanner), + logical_codec, + physical_codec, + ) +} diff --git a/datafusion/ffi/src/udtf.rs b/datafusion/ffi/src/udtf.rs index 0a111028798d1..a6b71e85c10ac 100644 --- a/datafusion/ffi/src/udtf.rs +++ b/datafusion/ffi/src/udtf.rs @@ -278,6 +278,7 @@ impl TableFunctionImpl for ForeignTableFunction { args.session(), self.0.runtime(), self.0.logical_codec.clone(), + None, ); let codec: Arc = (&self.0.logical_codec).into(); let expr_list = LogicalExprList { diff --git a/datafusion/ffi/tests/ffi_query_planner.rs b/datafusion/ffi/tests/ffi_query_planner.rs new file mode 100644 index 0000000000000..3c8eae9bdec6d --- /dev/null +++ b/datafusion/ffi/tests/ffi_query_planner.rs @@ -0,0 +1,61 @@ +// 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; + + use datafusion_common::DataFusionError; + use datafusion_execution::TaskContextProvider; + use datafusion_expr::LogicalPlanBuilder; + use datafusion_ffi::execution::FFI_TaskContextProvider; + use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; + use datafusion_ffi::query_planner::ForeignQueryPlanner; + use datafusion_ffi::tests::utils::get_module; + use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; + 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; + 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); + let planner: Arc = (&ffi_planner).into(); + + let any_ref: &dyn std::any::Any = planner.as_ref(); + assert!(any_ref.downcast_ref::().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::()); + + Ok(()) + } +} From 57be2c2dcbe39f139f0f9a06c0c945daa46f4c37 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 31 Jul 2026 09:05:52 -0400 Subject: [PATCH 03/17] refactor: split FFI session codec constructors Keep the standard FFI_SessionRef constructor focused on the required logical codec and derive a default physical codec. Add an explicit constructor for callers that already own matching logical and physical codecs.\n\nAI Disclosure: This code was written in part by an AI agent. --- datafusion/ffi/src/query_planner.rs | 4 ++-- datafusion/ffi/src/session/mod.rs | 25 +++++++++++++------- datafusion/ffi/src/table_provider.rs | 6 ++--- datafusion/ffi/src/table_provider_factory.rs | 3 +-- datafusion/ffi/src/udtf.rs | 1 - 5 files changed, 21 insertions(+), 18 deletions(-) diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index cb1f3e2f9d8d4..24ee0a5e0b6be 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -237,11 +237,11 @@ impl FFI_QueryPlanner { logical_plan_to_bytes_with_extension_codec(logical_plan, codec.as_ref())?; let logical_plan = logical_plan.iter().copied().collect(); let task_ctx = session.task_ctx(); - let session = FFI_SessionRef::new( + let session = FFI_SessionRef::new_with_ffi_codecs( session, session_runtime, self.logical_codec.clone(), - Some(self.physical_codec.clone()), + self.physical_codec.clone(), ); let physical_plan = unsafe { diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index a97ef97dca131..f4b97649dd1aa 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -421,19 +421,26 @@ impl FFI_SessionRef { session: &(dyn Session + Send + Sync), runtime: Option, logical_codec: FFI_LogicalExtensionCodec, - physical_codec: Option, + ) -> Self { + 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. + pub fn new_with_ffi_codecs( + session: &(dyn Session + Send + Sync), + runtime: Option, + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, ) -> Self { if let Some(session) = session.as_any().downcast_ref::() { return session.session.clone(); } - let physical_codec = physical_codec.unwrap_or_else(|| { - FFI_PhysicalExtensionCodec::new( - Arc::new(DefaultPhysicalExtensionCodec {}), - runtime.clone(), - logical_codec.task_ctx_provider.clone(), - ) - }); let private_data = Box::new(SessionPrivateData { session, runtime }); Self { @@ -818,7 +825,7 @@ mod tests { task_ctx_provider, ); - let local_session = FFI_SessionRef::new(&state, None, logical_codec, None); + let local_session = FFI_SessionRef::new(&state, None, logical_codec); let foreign_session = ForeignSession::try_from(&local_session)?; let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index 3bb8409c3672c..5a4b2fa27256f 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -492,8 +492,7 @@ impl TableProvider for ForeignTableProvider { filters: &[Expr], limit: Option, ) -> Result> { - let session = - FFI_SessionRef::new(session, None, self.0.logical_codec.clone(), None); + let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); let projections: FFI_Option> = projection .map(|p| p.iter().map(|v| v.to_owned()).collect()) @@ -563,8 +562,7 @@ impl TableProvider for ForeignTableProvider { input: Arc, insert_op: InsertOp, ) -> Result> { - let session = - FFI_SessionRef::new(session, None, self.0.logical_codec.clone(), None); + let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); let rc = Handle::try_current().ok(); let input = FFI_ExecutionPlan::new(input, rc); diff --git a/datafusion/ffi/src/table_provider_factory.rs b/datafusion/ffi/src/table_provider_factory.rs index b1a5036685035..b70e72f31aa4d 100644 --- a/datafusion/ffi/src/table_provider_factory.rs +++ b/datafusion/ffi/src/table_provider_factory.rs @@ -293,8 +293,7 @@ impl TableProviderFactory for ForeignTableProviderFactory { session: &dyn Session, cmd: &CreateExternalTable, ) -> Result> { - let session = - FFI_SessionRef::new(session, None, self.0.logical_codec.clone(), None); + let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); let cmd = self.serialize_cmd(cmd.clone())?; let provider = unsafe { diff --git a/datafusion/ffi/src/udtf.rs b/datafusion/ffi/src/udtf.rs index a6b71e85c10ac..0a111028798d1 100644 --- a/datafusion/ffi/src/udtf.rs +++ b/datafusion/ffi/src/udtf.rs @@ -278,7 +278,6 @@ impl TableFunctionImpl for ForeignTableFunction { args.session(), self.0.runtime(), self.0.logical_codec.clone(), - None, ); let codec: Arc = (&self.0.logical_codec).into(); let expr_list = LogicalExprList { From 45631bfb5f29c3b7acd0a963e145b169e4d92f5e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 31 Jul 2026 10:23:41 -0400 Subject: [PATCH 04/17] test: cover three-library FFI query planning Document the query planner serialization boundary and exercise an A/B/C ownership model with independently loaded cdylib images. Reuse the existing FFI table provider and verify foreign plans are reconstructed as local nodes through A's codecs. AI Disclosure: This code was written in part by an AI agent. --- datafusion/ffi/src/query_planner.rs | 54 ++++++- datafusion/ffi/src/tests/mod.rs | 6 + datafusion/ffi/src/tests/query_planner.rs | 56 +++++++ datafusion/ffi/src/tests/utils.rs | 43 +++++- datafusion/ffi/tests/ffi_query_planner.rs | 179 +++++++++++++++++++++- 5 files changed, 321 insertions(+), 17 deletions(-) diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index 24ee0a5e0b6be..b4fba21675c2b 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -15,6 +15,30 @@ // 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. + use std::ffi::c_void; use std::sync::Arc; @@ -46,7 +70,11 @@ use crate::session::{FFI_SessionRef, ForeignSession}; use crate::util::FFI_Result; use crate::{df_result, sresult_return}; -/// A stable struct for sharing [`QueryPlanner`] across FFI boundaries. +/// 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 { @@ -56,8 +84,10 @@ pub struct FFI_QueryPlanner { session: FFI_SessionRef, ) -> FfiFuture>>, + /// Codec used to encode and decode logical plans and extension nodes. pub logical_codec: FFI_LogicalExtensionCodec, + /// Codec used to encode and decode physical plans and extension nodes. pub physical_codec: FFI_PhysicalExtensionCodec, /// Used to create a clone of the query planner. @@ -177,7 +207,10 @@ impl Clone for FFI_QueryPlanner { } impl FFI_QueryPlanner { - /// Creates a new [`FFI_QueryPlanner`]. + /// Creates an [`FFI_QueryPlanner`] with native extension codecs. + /// + /// Missing codecs use DataFusion's defaults. `runtime` and + /// `task_ctx_provider` support codec callbacks across the FFI boundary. pub fn new( planner: Arc, runtime: Option, @@ -200,6 +233,10 @@ impl FFI_QueryPlanner { 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 returns its original FFI handle + /// rather than adding another wrapper layer. pub fn new_with_ffi_codecs( planner: Arc, logical_codec: FFI_LogicalExtensionCodec, @@ -224,8 +261,12 @@ impl FFI_QueryPlanner { } } - /// Calls this query planner with a [`Session`] exported over FFI using - /// `session_runtime` as that session provider's local Tokio runtime. + /// 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. pub async fn create_physical_plan_with_session_runtime( &self, logical_plan: &LogicalPlan, @@ -258,7 +299,10 @@ impl FFI_QueryPlanner { } } -/// Consumer-side wrapper for a foreign [`QueryPlanner`]. +/// 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); diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 5c3d71adc91ae..6d56c9117cba1 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -124,6 +124,11 @@ pub struct ForeignLibraryModule { physical_codec: FFI_PhysicalExtensionCodec, ) -> FFI_QueryPlanner, + pub create_library_c_query_planner: extern "C" fn( + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + ) -> FFI_QueryPlanner, + pub version: extern "C" fn() -> u64, } @@ -274,6 +279,7 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_context_aware_optimizer_rule: physical_optimizer::create_context_aware_optimizer_rule, create_query_planner: query_planner::create_query_planner, + create_library_c_query_planner: query_planner::create_library_c_query_planner, version: super::version, } } diff --git a/datafusion/ffi/src/tests/query_planner.rs b/datafusion/ffi/src/tests/query_planner.rs index f2eda6d8ff620..e57b28c34bbc9 100644 --- a/datafusion/ffi/src/tests/query_planner.rs +++ b/datafusion/ffi/src/tests/query_planner.rs @@ -19,15 +19,20 @@ 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::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; #[derive(Debug)] struct EmptyQueryPlanner; @@ -64,3 +69,54 @@ pub extern "C" fn create_query_planner( physical_codec, ) } + +#[derive(Debug)] +struct LibraryCQueryPlanner; + +#[async_trait] +impl QueryPlanner for LibraryCQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + if session.as_any().downcast_ref::().is_none() { + return exec_err!("library A's session was not foreign to library C"); + } + + let LogicalPlan::TableScan(scan) = logical_plan else { + return exec_err!("library C expected a table scan"); + }; + let provider = source_as_provider(&scan.source)?; + if provider.downcast_ref::().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::() { + 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::() { + return exec_err!("library C could not downcast its local UnionExec"); + } + Ok(plan) + } +} + +pub extern "C" fn create_library_c_query_planner( + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, +) -> FFI_QueryPlanner { + FFI_QueryPlanner::new_with_ffi_codecs( + Arc::new(LibraryCQueryPlanner), + logical_codec, + physical_codec, + ) +} diff --git a/datafusion/ffi/src/tests/utils.rs b/datafusion/ffi/src/tests/utils.rs index b6b50cbce875c..3119ab96d4032 100644 --- a/datafusion/ffi/src/tests/utils.rs +++ b/datafusion/ffi/src/tests/utils.rs @@ -62,14 +62,12 @@ fn find_library() -> Result { find_cdylib(deps_dir) } -pub fn get_module() -> Result { +fn load_module(lib_path: &Path) -> Result { 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 { Ok(module) } + +pub fn get_module() -> Result { + 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 { + 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/tests/ffi_query_planner.rs b/datafusion/ffi/tests/ffi_query_planner.rs index 3c8eae9bdec6d..09e459bab2098 100644 --- a/datafusion/ffi/tests/ffi_query_planner.rs +++ b/datafusion/ffi/tests/ffi_query_planner.rs @@ -19,16 +19,36 @@ mod utils; #[cfg(feature = "integration-tests")] mod tests { - use std::sync::Arc; + use std::sync::{Arc, OnceLock, Weak}; - use datafusion_common::DataFusionError; - use datafusion_execution::TaskContextProvider; - use datafusion_expr::LogicalPlanBuilder; + 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::LogicalPlan; + use datafusion_expr::logical_plan::Extension; 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::ForeignQueryPlanner; - use datafusion_ffi::tests::utils::get_module; - use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; + use datafusion_ffi::table_provider::ForeignTableProvider; + use datafusion_ffi::tests::{ + create_test_schema, + utils::{get_module, get_module_copy}, + }; + use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::empty::EmptyExec; + 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] @@ -49,12 +69,155 @@ mod tests { let any_ref: &dyn std::any::Any = planner.as_ref(); assert!(any_ref.downcast_ref::().is_some()); - let logical_plan = LogicalPlanBuilder::empty(false).build()?; + 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::()); + assert!(physical_plan.is::()); + + Ok(()) + } + + /// Library A's logical codec stores library B's provider while the logical + /// plan crosses into library C. A real application would encode enough + /// metadata to reconstruct or locate the provider instead. + #[derive(Debug, Default)] + struct LibraryALogicalCodec { + library_b_provider: OnceLock>, + } + + impl LogicalExtensionCodec for LibraryALogicalCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[LogicalPlan], + _ctx: &TaskContext, + ) -> Result { + not_impl_err!("logical extension nodes are not used in this test") + } + + fn try_encode(&self, _node: &Extension, _buf: &mut Vec) -> 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> { + 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, + buf: &mut Vec, + ) -> 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. + #[derive(Debug)] + struct LibraryAPhysicalCodec; + + impl PhysicalExtensionCodec for LibraryAPhysicalCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + 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, + buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + if !node.is::() { + return exec_err!("expected library B's plan to be foreign"); + } + 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; + 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 = (&ffi_provider).into(); + assert!(provider.downcast_ref::().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_library_c_query_planner)(logical_codec, physical_codec); + let planner: Arc = (&ffi_planner).into(); + let planner_any: &dyn std::any::Any = planner.as_ref(); + assert!(planner_any.downcast_ref::().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::()); + assert!(!physical_plan.is::()); + let children = physical_plan.children(); + assert_eq!(children.len(), 2); + assert!(children.iter().all(|child| child.is::())); + assert!( + children + .iter() + .all(|child| !child.is::()) + ); Ok(()) } From 9ca40b30a77db90921281a93dd9149ec0f211250 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 31 Jul 2026 10:35:45 -0400 Subject: [PATCH 05/17] test: consolidate FFI query planner fixtures Use one test query planner for both the basic round trip and the three-library table scan scenario. Remove the redundant library C constructor from the integration-test module. AI Disclosure: This code was written in part by an AI agent. --- datafusion/ffi/src/tests/mod.rs | 6 -- datafusion/ffi/src/tests/query_planner.rs | 84 ++++++++--------------- datafusion/ffi/tests/ffi_query_planner.rs | 3 +- 3 files changed, 31 insertions(+), 62 deletions(-) diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 6d56c9117cba1..5c3d71adc91ae 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -124,11 +124,6 @@ pub struct ForeignLibraryModule { physical_codec: FFI_PhysicalExtensionCodec, ) -> FFI_QueryPlanner, - pub create_library_c_query_planner: extern "C" fn( - logical_codec: FFI_LogicalExtensionCodec, - physical_codec: FFI_PhysicalExtensionCodec, - ) -> FFI_QueryPlanner, - pub version: extern "C" fn() -> u64, } @@ -279,7 +274,6 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_context_aware_optimizer_rule: physical_optimizer::create_context_aware_optimizer_rule, create_query_planner: query_planner::create_query_planner, - create_library_c_query_planner: query_planner::create_library_c_query_planner, version: super::version, } } diff --git a/datafusion/ffi/src/tests/query_planner.rs b/datafusion/ffi/src/tests/query_planner.rs index e57b28c34bbc9..c0c375c03f5fd 100644 --- a/datafusion/ffi/src/tests/query_planner.rs +++ b/datafusion/ffi/src/tests/query_planner.rs @@ -35,15 +35,42 @@ use crate::session::ForeignSession; use crate::table_provider::ForeignTableProvider; #[derive(Debug)] -struct EmptyQueryPlanner; +struct TestQueryPlanner; #[async_trait] -impl QueryPlanner for EmptyQueryPlanner { +impl QueryPlanner for TestQueryPlanner { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, session: &dyn Session, ) -> Result> { + if let LogicalPlan::TableScan(scan) = logical_plan { + if session.as_any().downcast_ref::().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::().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::() { + 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::() { + return exec_err!("library C could not downcast its local UnionExec"); + } + return Ok(plan); + } + let query_planner = session.query_planner(); let planner_any: &dyn std::any::Any = query_planner.as_ref(); if planner_any.downcast_ref::().is_none() { @@ -64,58 +91,7 @@ pub extern "C" fn create_query_planner( physical_codec: FFI_PhysicalExtensionCodec, ) -> FFI_QueryPlanner { FFI_QueryPlanner::new_with_ffi_codecs( - Arc::new(EmptyQueryPlanner), - logical_codec, - physical_codec, - ) -} - -#[derive(Debug)] -struct LibraryCQueryPlanner; - -#[async_trait] -impl QueryPlanner for LibraryCQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session: &dyn Session, - ) -> Result> { - if session.as_any().downcast_ref::().is_none() { - return exec_err!("library A's session was not foreign to library C"); - } - - let LogicalPlan::TableScan(scan) = logical_plan else { - return exec_err!("library C expected a table scan"); - }; - let provider = source_as_provider(&scan.source)?; - if provider.downcast_ref::().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::() { - 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::() { - return exec_err!("library C could not downcast its local UnionExec"); - } - Ok(plan) - } -} - -pub extern "C" fn create_library_c_query_planner( - logical_codec: FFI_LogicalExtensionCodec, - physical_codec: FFI_PhysicalExtensionCodec, -) -> FFI_QueryPlanner { - FFI_QueryPlanner::new_with_ffi_codecs( - Arc::new(LibraryCQueryPlanner), + Arc::new(TestQueryPlanner), logical_codec, physical_codec, ) diff --git a/datafusion/ffi/tests/ffi_query_planner.rs b/datafusion/ffi/tests/ffi_query_planner.rs index 09e459bab2098..60fbd3ad40a4a 100644 --- a/datafusion/ffi/tests/ffi_query_planner.rs +++ b/datafusion/ffi/tests/ffi_query_planner.rs @@ -197,8 +197,7 @@ mod tests { // 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_library_c_query_planner)(logical_codec, physical_codec); + let ffi_planner = (library_c.create_query_planner)(logical_codec, physical_codec); let planner: Arc = (&ffi_planner).into(); let planner_any: &dyn std::any::Any = planner.as_ref(); assert!(planner_any.downcast_ref::().is_some()); From b18d41d81089ab95d824b658cd5232f9c1a695ed Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 31 Jul 2026 10:41:57 -0400 Subject: [PATCH 06/17] docs: fix private query planner links Render the private FFI_SessionRef type as code so public query planner documentation passes rustdoc's private intra-doc link checks. AI Disclosure: This code was written in part by an AI agent. --- datafusion/ffi/src/query_planner.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index b4fba21675c2b..806bbc2a13efa 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -21,7 +21,7 @@ //! `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 +//! 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. @@ -264,7 +264,7 @@ impl FFI_QueryPlanner { /// 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 + /// `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. pub async fn create_physical_plan_with_session_runtime( From 05be26558a135ba5cfc3c4d72ceb9ca939abd756 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 3 Aug 2026 09:16:04 -0400 Subject: [PATCH 07/17] test: cover FFI query planner swap deployment Library C typically captures library A's query planner, then A installs C's planner on its session. C plans by delegating back to the captured handle. This is the deployment 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 to run its own passes over the result. Add an integration test for that topology. Replacing either serialization step with an FFI_ExecutionPlan handoff makes it fail, which the prior tests could not detect on the inbound leg. The test also asserts that after the swap the session reports C's own planner, documenting why C must delegate to the captured handle rather than call Session::query_planner or Session::create_physical_plan, both of which are self-references at that point. Widen the test physical codec to accept an A-local node during encode. FFI_ExecutionPlan::new unwraps a ForeignExecutionPlan back to its origin handle, so when C serializes a node A gave it, A is asked to encode the very plan its own try_decode produced. Take the delegate planner as FFI_Option on the existing create_query_planner module entry instead of adding a second entry point. Co-Authored-By: Claude Opus 5 (1M context) --- datafusion/ffi/src/query_planner.rs | 15 +++ datafusion/ffi/src/tests/mod.rs | 4 + datafusion/ffi/src/tests/query_planner.rs | 86 ++++++++++++++- datafusion/ffi/tests/ffi_query_planner.rs | 127 +++++++++++++++++++++- 4 files changed, 220 insertions(+), 12 deletions(-) diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index 806bbc2a13efa..5f75f77275e26 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -38,6 +38,21 @@ //! 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 +//! [`crate::session::FFI_SessionRef`] borrows its session with the lifetime +//! erased. use std::ffi::c_void; use std::sync::Arc; diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 5c3d71adc91ae..72eb9f5728b67 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -51,6 +51,7 @@ 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; @@ -119,9 +120,12 @@ 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, pub version: extern "C" fn() -> u64, diff --git a/datafusion/ffi/src/tests/query_planner.rs b/datafusion/ffi/src/tests/query_planner.rs index c0c375c03f5fd..90d713f8a6336 100644 --- a/datafusion/ffi/src/tests/query_planner.rs +++ b/datafusion/ffi/src/tests/query_planner.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::any::Any; use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema}; @@ -24,6 +25,7 @@ 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}; @@ -33,6 +35,7 @@ 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; @@ -72,7 +75,7 @@ impl QueryPlanner for TestQueryPlanner { } let query_planner = session.query_planner(); - let planner_any: &dyn std::any::Any = query_planner.as_ref(); + let planner_any: &dyn Any = query_planner.as_ref(); if planner_any.downcast_ref::().is_none() { return exec_err!("query planner did not cross the FFI boundary"); } @@ -86,13 +89,84 @@ impl QueryPlanner for TestQueryPlanner { } } +/// 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, +} + +#[async_trait] +impl QueryPlanner for SwappedQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + if session.as_any().downcast_ref::().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::().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::() { + return exec_err!("library A's plan was opaque to library C"); + } + let Some(sort) = plan.downcast_ref::() 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::() { + 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::new_with_ffi_codecs( - Arc::new(TestQueryPlanner), - logical_codec, - physical_codec, - ) + let planner: Arc = 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/tests/ffi_query_planner.rs b/datafusion/ffi/tests/ffi_query_planner.rs index 60fbd3ad40a4a..e1bf9478dc88f 100644 --- a/datafusion/ffi/tests/ffi_query_planner.rs +++ b/datafusion/ffi/tests/ffi_query_planner.rs @@ -29,20 +29,22 @@ mod tests { DataFusionError, Result, TableReference, exec_err, not_impl_err, }; use datafusion_execution::{TaskContext, TaskContextProvider}; - use datafusion_expr::LogicalPlan; 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::ForeignQueryPlanner; + 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::{ @@ -63,7 +65,11 @@ mod tests { task_ctx_provider, ); - let ffi_planner = (module.create_query_planner)(logical_codec, physical_codec); + let ffi_planner = (module.create_query_planner)( + logical_codec, + physical_codec, + FFI_Option::None, + ); let planner: Arc = (&ffi_planner).into(); let any_ref: &dyn std::any::Any = planner.as_ref(); @@ -132,6 +138,12 @@ mod tests { /// 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; @@ -155,8 +167,11 @@ mod tests { buf: &mut Vec, _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { - if !node.is::() { - return exec_err!("expected library B's plan to be foreign"); + if !node.is::() && !node.is::() { + 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(()) @@ -197,7 +212,11 @@ mod tests { // 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); + let ffi_planner = (library_c.create_query_planner)( + logical_codec, + physical_codec, + FFI_Option::None, + ); let planner: Arc = (&ffi_planner).into(); let planner_any: &dyn std::any::Any = planner.as_ref(); assert!(planner_any.downcast_ref::().is_some()); @@ -220,4 +239,100 @@ mod tests { 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; + 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 = (&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 = + (&ffi_planner).into(); + let planner_any: &dyn std::any::Any = library_c_planner.as_ref(); + assert!(planner_any.downcast_ref::().is_some()); + + // Library A swaps C's planner into the session it already owns. Mutating + // the existing state keeps the `Arc` 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::()); + assert!(!physical_plan.is::()); + let children = physical_plan.children(); + assert_eq!(children.len(), 2); + for child in &children { + let sort = child + .downcast_ref::() + .expect("library A could not downcast the SortExec it planned"); + assert!(sort.input().is::()); + assert!(!sort.input().is::()); + } + + Ok(()) + } } From a72ce75b333f3b88f62a93cdce76743784f20854 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 3 Aug 2026 14:29:40 -0400 Subject: [PATCH 08/17] Fix cargo doc --- datafusion/ffi/src/query_planner.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index 5f75f77275e26..dcb177abbbca5 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -51,8 +51,7 @@ //! //! Retain the planner rather than the session. [`FFI_QueryPlanner`] owns a //! reference-counted planner, so it outlives A's original session, whereas -//! [`crate::session::FFI_SessionRef`] borrows its session with the lifetime -//! erased. +//! `FFI_SessionRef` borrows its session with the lifetime erased. use std::ffi::c_void; use std::sync::Arc; From bb15fc9dd72b2b2ccd805ee15475c933865e460c Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 4 Aug 2026 08:26:17 -0400 Subject: [PATCH 09/17] refactor: address FFI query planner review feedback Drop redundant `Send`/`Sync` bounds and remove two codec footguns in the FFI query planner API, per review on #24028. - `Session` already requires `Send + Sync`, so `&(dyn Session + Send + Sync)` was noise. Narrowing to `&dyn Session` also widens what callers can pass. - `LogicalExtensionCodec` and `PhysicalExtensionCodec` already require `Send`. The `+ Send` on the codec constructor parameters bought nothing and blocked callers holding an existing `Arc`, since Rust will not coerce that to `Arc`. - `FFI_QueryPlanner::new_with_ffi_codecs` silently dropped the supplied codecs when re-exporting an already-foreign planner. It now adopts them while keeping the original planner identity. - `FFI_QueryPlanner::new` no longer takes `Option` codecs. Passing `None` used to install the default codecs, clobbering extension-node handling; requiring an explicit codec makes that unrepresentable. Co-Authored-By: Claude Opus 5 (1M context) --- .../ffi/src/proto/logical_extension_codec.rs | 6 +-- .../ffi/src/proto/physical_extension_codec.rs | 6 +-- datafusion/ffi/src/query_planner.rs | 43 ++++++++++--------- datafusion/ffi/src/session/mod.rs | 10 ++--- datafusion/ffi/src/table_provider.rs | 4 +- datafusion/ffi/src/table_provider_factory.rs | 2 +- datafusion/ffi/src/udtf.rs | 2 +- 7 files changed, 36 insertions(+), 37 deletions(-) diff --git a/datafusion/ffi/src/proto/logical_extension_codec.rs b/datafusion/ffi/src/proto/logical_extension_codec.rs index 97aa5c901a636..0ffdccd4e4830 100644 --- a/datafusion/ffi/src/proto/logical_extension_codec.rs +++ b/datafusion/ffi/src/proto/logical_extension_codec.rs @@ -295,7 +295,7 @@ impl Drop for FFI_LogicalExtensionCodec { impl FFI_LogicalExtensionCodec { /// Creates a new [`FFI_LogicalExtensionCodec`]. pub fn new( - codec: Arc, + codec: Arc, runtime: Option, task_ctx_provider: impl Into, ) -> Self { @@ -712,14 +712,12 @@ mod tests { #[test] fn ffi_logical_extension_codec_local_bypass() { - let codec = - Arc::new(TestExtensionCodec {}) as Arc; + let codec = Arc::new(TestExtensionCodec {}) as Arc; 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; // Verify local libraries can be downcast to their original let foreign_codec: Arc = (&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 bf45013e12645..95d2ed68a6ea3 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/datafusion/ffi/src/proto/physical_extension_codec.rs @@ -281,7 +281,7 @@ impl Drop for FFI_PhysicalExtensionCodec { impl FFI_PhysicalExtensionCodec { /// Creates a new [`FFI_PhysicalExtensionCodec`]. pub fn new( - codec: Arc, + codec: Arc, runtime: Option, task_ctx_provider: impl Into, ) -> Self { @@ -695,14 +695,12 @@ pub(crate) mod tests { #[test] fn ffi_physical_extension_codec_local_bypass() { - let codec = - Arc::new(TestExtensionCodec {}) as Arc; + let codec = Arc::new(TestExtensionCodec {}) as Arc; 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; // Verify local libraries can be downcast to their original let foreign_codec: Arc = (&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 index dcb177abbbca5..395c4cd342787 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -67,12 +67,8 @@ use datafusion_proto::bytes::{ physical_plan_from_bytes_with_extension_codec, physical_plan_to_bytes_with_extension_codec, }; -use datafusion_proto::logical_plan::{ - DefaultLogicalExtensionCodec, LogicalExtensionCodec, -}; -use datafusion_proto::physical_plan::{ - DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, -}; +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; @@ -151,7 +147,7 @@ unsafe extern "C" fn create_physical_plan_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()) @@ -223,25 +219,26 @@ impl Clone for FFI_QueryPlanner { impl FFI_QueryPlanner { /// Creates an [`FFI_QueryPlanner`] with native extension codecs. /// - /// Missing codecs use DataFusion's defaults. `runtime` and - /// `task_ctx_provider` support codec callbacks across the FFI boundary. + /// 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, runtime: Option, task_ctx_provider: impl Into, - logical_codec: Option>, - physical_codec: Option>, + logical_codec: Arc, + physical_codec: Arc, ) -> Self { let task_ctx_provider = task_ctx_provider.into(); - let logical_codec = - logical_codec.unwrap_or_else(|| Arc::new(DefaultLogicalExtensionCodec {})); let logical_codec = FFI_LogicalExtensionCodec::new( logical_codec, runtime.clone(), task_ctx_provider.clone(), ); - let physical_codec = - physical_codec.unwrap_or_else(|| Arc::new(DefaultPhysicalExtensionCodec {})); let physical_codec = FFI_PhysicalExtensionCodec::new(physical_codec, runtime, task_ctx_provider); Self::new_with_ffi_codecs(planner, logical_codec, physical_codec) @@ -249,8 +246,9 @@ impl FFI_QueryPlanner { /// Creates an [`FFI_QueryPlanner`] using prebuilt FFI extension codecs. /// - /// If `planner` is already foreign, this returns its original FFI handle - /// rather than adding another wrapper layer. + /// 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, logical_codec: FFI_LogicalExtensionCodec, @@ -258,7 +256,10 @@ impl FFI_QueryPlanner { ) -> Self { let any_ref: &dyn std::any::Any = planner.as_ref(); if let Some(planner) = any_ref.downcast_ref::() { - return planner.0.clone(); + 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 }); @@ -356,6 +357,8 @@ mod tests { 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::*; @@ -381,8 +384,8 @@ mod tests { Arc::new(EmptyQueryPlanner), None, &task_ctx_provider, - None, - None, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), ) } diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index f4b97649dd1aa..5815d2af85db6 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -155,12 +155,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, } 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 } } @@ -418,7 +418,7 @@ impl Drop for FFI_SessionRef { impl FFI_SessionRef { /// Creates a new [`FFI_SessionRef`]. pub fn new( - session: &(dyn Session + Send + Sync), + session: &dyn Session, runtime: Option, logical_codec: FFI_LogicalExtensionCodec, ) -> Self { @@ -432,7 +432,7 @@ impl FFI_SessionRef { /// Creates a new [`FFI_SessionRef`] using existing FFI codecs. pub fn new_with_ffi_codecs( - session: &(dyn Session + Send + Sync), + session: &dyn Session, runtime: Option, logical_codec: FFI_LogicalExtensionCodec, physical_codec: FFI_PhysicalExtensionCodec, @@ -495,7 +495,7 @@ 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()); } diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index 5a4b2fa27256f..ee9377bff064e 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 b70e72f31aa4d..63ebb51bb1db8 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/udtf.rs b/datafusion/ffi/src/udtf.rs index 0a111028798d1..fa28519d58de5 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()) From 01966adff9d6f96718be93d7a43f2eee44d4b8ea Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 5 Aug 2026 08:05:31 -0400 Subject: [PATCH 10/17] perf: lazily export FFI session planning state AI Disclosure: This code was written in part by an AI agent.: --- datafusion/ffi/src/session/mod.rs | 82 ++++++++++++++++++++++++++----- 1 file changed, 69 insertions(+), 13 deletions(-) diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 5815d2af85db6..c3126e21248ee 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -18,7 +18,7 @@ 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; @@ -487,8 +487,8 @@ pub struct ForeignSession { table_options: TableOptions, runtime_env: Arc, props: ExecutionProps, - query_planner: Arc, - physical_optimizers: Vec>, + query_planner: OnceLock>, + physical_optimizers: OnceLock>>, } unsafe impl Send for ForeignSession {} @@ -549,12 +549,6 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { ) }) .collect(); - let query_planner = (&(session.query_planner)(session)).into(); - let physical_optimizers = (session.physical_optimizers)(session) - .into_iter() - .map(|rule| (&rule).into()) - .collect(); - Ok(Self { session: session.clone(), config, @@ -567,8 +561,8 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { extension_types: Arc::new(MemoryExtensionTypeRegistry::default()), runtime_env: Default::default(), props: Default::default(), - query_planner, - physical_optimizers, + query_planner: OnceLock::new(), + physical_optimizers: OnceLock::new(), }) } } @@ -668,7 +662,10 @@ impl Session for ForeignSession { } fn query_planner(&self) -> Arc { - Arc::clone(&self.query_planner) + 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 { @@ -730,7 +727,12 @@ impl Session for ForeignSession { } fn physical_optimizers(&self) -> &[Arc] { - &self.physical_optimizers + 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> { @@ -792,6 +794,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; @@ -803,6 +806,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 { + 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(); From df6ed3d5c55273fad83ccb4f135466023cda46ca Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 5 Aug 2026 08:26:50 -0400 Subject: [PATCH 11/17] docs: clarify FFI session planner delegation AI Disclosure: This code was written in part by an AI agent.: --- datafusion/ffi/src/session/mod.rs | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index c3126e21248ee..8d8e3bbe469d7 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -15,6 +15,21 @@ // 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; @@ -472,8 +487,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, From 7dc6a8ea06d17c0febc829405cced55259e54e9e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 5 Aug 2026 09:14:35 -0400 Subject: [PATCH 12/17] docs: clarify FFI session codec limitations AI Disclosure: This code was written in part by an AI agent.: AI Disclosure: This code was written in part by an AI agent.: --- datafusion/ffi/src/session/mod.rs | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 8d8e3bbe469d7..1378661cdd88f 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -431,12 +431,31 @@ 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, runtime: Option, 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(), @@ -446,6 +465,11 @@ impl FFI_SessionRef { } /// 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. pub fn new_with_ffi_codecs( session: &dyn Session, runtime: Option, From 5fd78ecb4378d8dd0ba539f081aba96292dea5cb Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 5 Aug 2026 09:39:27 -0400 Subject: [PATCH 13/17] refactor: restrict FFI codec field visibility AI Disclosure: This code was written in part by an AI agent.: --- datafusion/ffi/src/proto/logical_extension_codec.rs | 2 +- datafusion/ffi/src/query_planner.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/ffi/src/proto/logical_extension_codec.rs b/datafusion/ffi/src/proto/logical_extension_codec.rs index 0ffdccd4e4830..ed2c594f1bc02 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>, - 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. diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index 395c4cd342787..e36a60764f9e9 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -95,10 +95,10 @@ pub struct FFI_QueryPlanner { ) -> FfiFuture>>, /// Codec used to encode and decode logical plans and extension nodes. - pub logical_codec: FFI_LogicalExtensionCodec, + logical_codec: FFI_LogicalExtensionCodec, /// Codec used to encode and decode physical plans and extension nodes. - pub physical_codec: FFI_PhysicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, /// Used to create a clone of the query planner. clone: unsafe extern "C" fn(planner: &Self) -> Self, From 15a4dfef4b59a9be32651c0aac8c684a9736e802 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 5 Aug 2026 09:57:04 -0400 Subject: [PATCH 14/17] docs: clarify FFI query planner runtime handling AI Disclosure: This code was written in part by an AI agent.: --- datafusion/ffi/src/query_planner.rs | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index e36a60764f9e9..bfe3bf61840a3 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -282,6 +282,11 @@ impl FFI_QueryPlanner { /// `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, @@ -416,4 +421,25 @@ mod tests { 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::()); + + Ok(()) + } } From 704df62c1cde6c8d38a5e4ab0dc9c8da42ecf7eb Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 5 Aug 2026 10:21:38 -0400 Subject: [PATCH 15/17] perf: bulk copy serialized FFI plans AI Disclosure: This code was written in part by an AI agent.: --- datafusion/ffi/src/query_planner.rs | 4 ++-- datafusion/ffi/src/session/mod.rs | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index bfe3bf61840a3..6d895d65c5fc1 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -170,7 +170,7 @@ unsafe extern "C" fn create_physical_plan_fn_wrapper( physical_codec.as_ref(), )); - FFI_Result::Ok(physical_plan.iter().copied().collect()) + FFI_Result::Ok(SVec::from(physical_plan.as_ref())) } .into_ffi() } @@ -296,7 +296,7 @@ impl FFI_QueryPlanner { let codec: Arc = (&self.logical_codec).into(); let logical_plan = logical_plan_to_bytes_with_extension_codec(logical_plan, codec.as_ref())?; - let logical_plan = logical_plan.iter().copied().collect(); + let logical_plan = SVec::from(logical_plan.as_ref()); let task_ctx = session.task_ctx(); let session = FFI_SessionRef::new_with_ffi_codecs( session, diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 1378661cdd88f..085a24074b358 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -235,7 +235,7 @@ unsafe extern "C" fn optimize_fn_wrapper( logical_codec.as_ref(), )); - FFI_Result::Ok(optimized_plan.iter().copied().collect()) + FFI_Result::Ok(SVec::from(optimized_plan.as_ref())) } unsafe extern "C" fn create_physical_plan_fn_wrapper( @@ -724,7 +724,7 @@ impl Session for ForeignSession { logical_plan_to_bytes_with_extension_codec(plan, codec.as_ref())?; let optimized_plan = df_result!((self.session.optimize)( &self.session, - logical_plan.iter().copied().collect(), + SVec::from(logical_plan.as_ref()), ))?; logical_plan_from_bytes_with_extension_codec( optimized_plan.as_slice(), From 57797fea446d876280261873368e1401d47631d4 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 5 Aug 2026 10:45:54 -0400 Subject: [PATCH 16/17] docs: clarify FFI planner test codec AI Disclosure: This code was written in part by an AI agent.: --- datafusion/ffi/tests/ffi_query_planner.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/datafusion/ffi/tests/ffi_query_planner.rs b/datafusion/ffi/tests/ffi_query_planner.rs index e1bf9478dc88f..c72b8c3e889ae 100644 --- a/datafusion/ffi/tests/ffi_query_planner.rs +++ b/datafusion/ffi/tests/ffi_query_planner.rs @@ -85,9 +85,19 @@ mod tests { Ok(()) } - /// Library A's logical codec stores library B's provider while the logical - /// plan crosses into library C. A real application would encode enough - /// metadata to reconstruct or locate the provider instead. + /// 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>, From 9f970e38742796c5d40cd1421d5eb10cf8e24c91 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 5 Aug 2026 12:16:08 -0400 Subject: [PATCH 17/17] Update upgrade guide to reflect current status --- .../library-user-guide/upgrading/55.0.0.md | 23 +++++++++++++++---- 1 file changed, 18 insertions(+), 5 deletions(-) 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 53cba29f9abb6..b98b5d2aa52a1 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -895,15 +895,28 @@ fn physical_optimizers(&self) -> &[Arc] } ``` -`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