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 815acd95e323161b12c6b756331b65c2850040c4 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 5 Aug 2026 11:59:54 -0400 Subject: [PATCH 17/17] feat(ffi): add an FFI extension codec bundle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Serializing plans across an FFI boundary needs three values that must agree with one another: a task context provider, a logical extension codec, and a physical extension codec. The wrappers passed them separately, and the ones that carried only a logical codec synthesized a DefaultPhysicalExtensionCodec when they had to export a session. A query planner reached through such a session could not round-trip custom physical extension nodes; it failed later with "PhysicalExtensionCodec is not provided". The synthesized codec also borrowed the logical codec's task context provider, which is not guaranteed to represent the exported session. Add FFI_ExtensionCodecBundle, which carries all three as one unit with private fields so the constructors are the only way to pair them. The bundle owns no private data of its own; each member already carries its own lifecycle pointers and marker. The dependency direction is bundle to codecs to task context provider — a bundle inside a codec would make cloning recurse forever. Propagate it through every wrapper that exports a session or builds a nested provider: FFI_TableProvider, FFI_TableProviderFactory, FFI_TableFunction, FFI_CatalogProvider, FFI_CatalogProviderList, FFI_SchemaProvider, FFI_SessionRef, and FFI_QueryPlanner. Each keeps one constructor taking the bundle; the new_with_ffi_codec(s) variants and the Option> argument whose None meant "default" are gone, so choosing the defaults is now explicit. The two paths inside FFI_LogicalExtensionCodec that rebuild a table provider receive only the codec, so they pair it with an explicit default physical codec and document what that costs a consumer. Session::create_physical_plan serialized its plan with no extension codec on either side of the boundary; it now uses the bundle's logical codec, matching optimize and create_physical_expr. Tests: a three-library integration test where library A owns the session and a custom physical codec, installs library C's planner, and queries library B's provider, which reaches that planner through the session A handed it and returns a custom physical extension node. Restoring the old default-codec behaviour makes it fail with "PhysicalExtensionCodec is not provided". Plus unit coverage that nested catalog/schema/table construction preserves the bundle and that an expired task context provider reports a clear error. BREAKING CHANGE: these FFI struct layouts changed. Rebuild FFI providers and consumers against this release. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/datafusion-ffi/SKILL.md | 21 +- .../ffi/ffi_example_table_provider/src/lib.rs | 6 +- .../ffi/ffi_module_interface/src/lib.rs | 4 +- .../ffi/ffi_module_loader/src/main.rs | 10 +- datafusion/ffi/README.md | 42 ++ datafusion/ffi/src/catalog_provider.rs | 66 +-- datafusion/ffi/src/catalog_provider_list.rs | 121 ++++-- .../ffi/src/proto/extension_codec_bundle.rs | 380 ++++++++++++++++++ .../ffi/src/proto/logical_extension_codec.rs | 24 +- datafusion/ffi/src/proto/mod.rs | 1 + datafusion/ffi/src/query_planner.rs | 98 ++--- datafusion/ffi/src/schema_provider.rs | 68 +--- datafusion/ffi/src/session/mod.rs | 163 ++++---- datafusion/ffi/src/table_provider.rs | 138 +++---- datafusion/ffi/src/table_provider_factory.rs | 67 ++- datafusion/ffi/src/tests/async_provider.rs | 11 +- datafusion/ffi/src/tests/catalog.rs | 10 +- datafusion/ffi/src/tests/mod.rs | 46 ++- datafusion/ffi/src/tests/query_planner.rs | 103 ++++- datafusion/ffi/src/tests/sync_provider.rs | 6 +- .../ffi/src/tests/table_provider_factory.rs | 6 +- datafusion/ffi/src/tests/udf_udaf_udwf.rs | 6 +- datafusion/ffi/src/udtf.rs | 95 ++--- datafusion/ffi/tests/ffi_catalog.rs | 8 +- datafusion/ffi/tests/ffi_integration.rs | 13 +- datafusion/ffi/tests/ffi_query_planner.rs | 190 ++++++--- datafusion/ffi/tests/ffi_udtf.rs | 4 +- datafusion/ffi/tests/utils/mod.rs | 19 +- .../library-user-guide/upgrading/55.0.0.md | 73 ++++ 29 files changed, 1196 insertions(+), 603 deletions(-) create mode 100644 datafusion/ffi/src/proto/extension_codec_bundle.rs diff --git a/.ai/skills/datafusion-ffi/SKILL.md b/.ai/skills/datafusion-ffi/SKILL.md index c105da653641c..1a7c73e83a09e 100644 --- a/.ai/skills/datafusion-ffi/SKILL.md +++ b/.ai/skills/datafusion-ffi/SKILL.md @@ -42,7 +42,7 @@ A new `FFI_X` for trait `X` must follow this template. Use `FFI_CatalogProvider` pub struct FFI_X { some_method: unsafe extern "C" fn(this: &Self, ...) -> FFI_Result<...>, optional_method: Option FFI_Result<...>>, - pub logical_codec: FFI_LogicalExtensionCodec, + pub codecs: FFI_ExtensionCodecBundle, clone: unsafe extern "C" fn(&Self) -> Self, release: unsafe extern "C" fn(&mut Self), @@ -60,7 +60,7 @@ Field rules: - **One `unsafe extern "C" fn` per trait method.** Always populate — `Arc` dispatch picks override-or-default at call time, so the producer side gets the right answer without the consumer needing to know. See § "Method coverage". - **`Option` is the capability-flag exception**, not a template. Crate uses it exactly once: `FFI_TableProvider::supports_filters_pushdown`. See § "Method coverage". -- **Codec field** (`FFI_LogicalExtensionCodec` / `FFI_PhysicalExtensionCodec`) only if the trait moves `Expr`s / `LogicalPlan`s / `ExecutionPlan`s across the boundary. +- **Codec field** — carry an `FFI_ExtensionCodecBundle` (`src/proto/extension_codec_bundle.rs`), not a bare `FFI_LogicalExtensionCodec` / `FFI_PhysicalExtensionCodec`, whenever the trait moves `Expr`s / `LogicalPlan`s / `ExecutionPlan`s across the boundary **or** exports an `FFI_SessionRef`. The bundle pairs one `FFI_TaskContextProvider` with both codecs so the three cannot drift apart, and a wrapper that only serializes logical data still needs the physical codec for any planner a consumer reaches through the session it exports. Forward the bundle unchanged to every nested `FFI_X` the wrapper creates, including in `clone_fn_wrapper`. Never store a bundle inside a codec: cloning a bundle clones its logical codec, so that would recurse forever — the dependency direction is bundle → codecs → task context provider. - **Method function pointers are private by default.** Mark `pub` only if a downstream library needs to invoke them directly (rare — typically only `version`, `library_marker_id`, embedded codecs are `pub`). - **`version: super::version` is mandatory.** Consumers gate compatibility on it. - **`library_marker_id: crate::get_library_marker_id` is mandatory *when the wrapper uses the standard `ForeignX` adapter pattern*.** Two flavors exist: @@ -126,19 +126,14 @@ impl Clone for FFI_X { fn clone(&self) -> Self { unsafe { (self.clone)(self) } } `release` must null `private_data` so a double-free debug-asserts loudly. -### 5. Constructor split +### 5. Constructor + +One constructor per wrapper, taking the bundle: ```rust impl FFI_X { pub fn new(inner: Arc, runtime: Option, - task_ctx_provider: impl Into, - logical_codec: Option>) -> Self { - // build FFI_LogicalExtensionCodec from defaults, then forward - Self::new_with_ffi_codec(inner, runtime, ffi_codec) - } - - pub fn new_with_ffi_codec(inner: Arc, runtime: Option, - logical_codec: FFI_LogicalExtensionCodec) -> Self { + codecs: FFI_ExtensionCodecBundle) -> Self { // Round-trip downcast: if inner is already a ForeignX, return its FFI directly. if let Some(foreign) = inner.downcast_ref::() { return foreign.0.clone(); @@ -150,6 +145,8 @@ impl FFI_X { The round-trip downcast is **mandatory** — without it, repeated FFI hops nest `ForeignX(FFI_X(ForeignX(...)))` and you pay the boundary cost every layer. +Do **not** add a second constructor that builds the bundle from native codecs plus an `Option>`. That shape makes "use the default physical codec" implicit, which is the failure `FFI_ExtensionCodecBundle` exists to prevent. Callers with native codecs build the bundle themselves via `FFI_ExtensionCodecBundle::new`, or state the defaults explicitly via `FFI_ExtensionCodecBundle::new_default`. + ### 6. The `Foreign` consumer ```rust @@ -356,5 +353,5 @@ When reviewing a PR that touches `datafusion/ffi/`: - Canonical wrapper to model after: `src/catalog_provider.rs`. Async + capability-flag variants: `src/table_provider.rs`. - Mutable-trait variant: `src/udaf/accumulator.rs` (`Box`). - Optional-method pattern: `FFI_TableProvider::supports_filters_pushdown`. -- Codec wiring: `src/proto/logical_extension_codec.rs`, `src/proto/physical_extension_codec.rs`. +- Codec wiring: `src/proto/extension_codec_bundle.rs` (the pairing every wrapper carries), `src/proto/logical_extension_codec.rs`, `src/proto/physical_extension_codec.rs`. - Examples crate: `datafusion-examples/examples/ffi` (end-to-end producer + consumer). diff --git a/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs b/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs index 29b04d0042547..f7bd101ccdde2 100644 --- a/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs +++ b/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs @@ -21,7 +21,7 @@ use arrow::array::{RecordBatch, record_batch}; use arrow::datatypes as arrow_schema; use arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::MemTable; -use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use datafusion_ffi::table_provider::FFI_TableProvider; use ffi_module_interface::TableProviderModule; @@ -36,7 +36,7 @@ fn create_record_batch(start_value: i32, num_values: usize) -> RecordBatch { /// Here we only wish to create a simple table provider as an example. /// We create an in-memory table and convert it to it's FFI counterpart. extern "C" fn construct_simple_table_provider( - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_TableProvider { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::Int32, true), @@ -53,7 +53,7 @@ extern "C" fn construct_simple_table_provider( let table_provider = MemTable::try_new(schema, vec![batches]).unwrap(); - FFI_TableProvider::new_with_ffi_codec(Arc::new(table_provider), true, None, codec) + FFI_TableProvider::new(Arc::new(table_provider), true, None, codecs) } #[unsafe(no_mangle)] diff --git a/datafusion-examples/examples/ffi/ffi_module_interface/src/lib.rs b/datafusion-examples/examples/ffi/ffi_module_interface/src/lib.rs index 54a59c9e5d073..c6b30dedc5c15 100644 --- a/datafusion-examples/examples/ffi/ffi_module_interface/src/lib.rs +++ b/datafusion-examples/examples/ffi/ffi_module_interface/src/lib.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use datafusion_ffi::table_provider::FFI_TableProvider; /// This struct defines the module interfaces. It is to be shared by @@ -27,5 +27,5 @@ use datafusion_ffi::table_provider::FFI_TableProvider; pub struct TableProviderModule { /// Constructs the table provider pub create_table: - extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProvider, + extern "C" fn(codecs: FFI_ExtensionCodecBundle) -> FFI_TableProvider, } diff --git a/datafusion-examples/examples/ffi/ffi_module_loader/src/main.rs b/datafusion-examples/examples/ffi/ffi_module_loader/src/main.rs index 0657c4a08fa86..02a01319105f8 100644 --- a/datafusion-examples/examples/ffi/ffi_module_loader/src/main.rs +++ b/datafusion-examples/examples/ffi/ffi_module_loader/src/main.rs @@ -23,7 +23,7 @@ use datafusion::{ execution::TaskContextProvider, prelude::SessionContext, }; -use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use ffi_module_interface::TableProviderModule; #[tokio::main] @@ -67,13 +67,17 @@ async fn main() -> Result<()> { let table_provider_module = get_module(); let ctx = Arc::new(SessionContext::new()); - let codec = FFI_LogicalExtensionCodec::new_default( + // The bundle pairs the task context provider with the logical and physical + // extension codecs the module should serialize with. This example moves no + // custom extension nodes, so both codecs are the defaults. + let codecs = FFI_ExtensionCodecBundle::new_default( &(Arc::clone(&ctx) as Arc), + None, ); // By calling the code below, the table provided will be created within // the module's code. - let ffi_table_provider = (table_provider_module.create_table)(codec); + let ffi_table_provider = (table_provider_module.create_table)(codecs); // In order to access the table provider within this executable, we need to // turn it into a `TableProvider`. diff --git a/datafusion/ffi/README.md b/datafusion/ffi/README.md index ded54b0a88d09..87262b30b0eeb 100644 --- a/datafusion/ffi/README.md +++ b/datafusion/ffi/README.md @@ -218,6 +218,48 @@ these methods that your provider remains valid for the lifetime of the calls. The `FFI_TaskContextProvider` is implemented on `SessionContext` and it is easy to implement on any struct that implements `Session`. +## Extension Codec Bundle + +Serializing plans across the boundary needs three values that have to agree +with one another: the `FFI_TaskContextProvider` above, a logical extension +codec, and a physical extension codec. `FFI_ExtensionCodecBundle` carries +them as a single unit, and every wrapper that serializes a plan or expression +takes one: + +```rust,ignore +let codecs = FFI_ExtensionCodecBundle::new( + &task_ctx_provider, + None, // Option + Arc::new(MyLogicalCodec), + Arc::new(MyPhysicalCodec), +); + +let ffi_provider = FFI_TableProvider::new(provider, true, None, codecs.clone()); +let ffi_catalog = FFI_CatalogProvider::new(catalog, None, codecs); +``` + +Use `FFI_ExtensionCodecBundle::new_default` when no custom extension nodes +cross the boundary; it selects the default logical and physical codecs +explicitly. + +Both codecs are needed even by a wrapper that only serializes logical data, +because such a wrapper still exports a `Session` and a consumer can reach +`Session::query_planner` through it. That planner serializes physical plans, +so it needs the physical codec belonging to the same environment. + +The bundle is propagated unchanged through nested construction — a catalog +provider list hands it to each catalog, which hands it to each schema, which +hands it to each table provider — so a table found by walking the hierarchy +serializes exactly like the list it came from. + +One case does not get a full bundle: a table provider decoded out of a +serialized logical plan is reconstructed by `FFI_LogicalExtensionCodec`, which +holds no physical codec, so it is paired with the default one. Such a provider +cannot carry custom physical extension nodes if the decoding library scans it +with a session local to that library. Scanning it with the exporting library's +own session handle is unaffected, because exporting an already-foreign session +returns the original handle and its complete bundle. + [apache datafusion]: https://datafusion.apache.org/ [api docs]: http://docs.rs/datafusion-ffi/latest [rust abi]: https://doc.rust-lang.org/reference/abi.html diff --git a/datafusion/ffi/src/catalog_provider.rs b/datafusion/ffi/src/catalog_provider.rs index ff37e8e7e0462..a24847b7788e9 100644 --- a/datafusion/ffi/src/catalog_provider.rs +++ b/datafusion/ffi/src/catalog_provider.rs @@ -20,15 +20,11 @@ use std::sync::Arc; use datafusion_catalog::{CatalogProvider, SchemaProvider}; use datafusion_common::error::Result; -use datafusion_proto::logical_plan::{ - DefaultLogicalExtensionCodec, LogicalExtensionCodec, -}; use stabby::string::String as SString; 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::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::schema_provider::{FFI_SchemaProvider, ForeignSchemaProvider}; use crate::util::{FFI_Option, FFI_Result}; use crate::{df_result, sresult_return}; @@ -58,7 +54,9 @@ pub struct FFI_CatalogProvider { cascade: bool, ) -> FFI_Result>, - pub logical_codec: FFI_LogicalExtensionCodec, + /// The serialization environment propagated to every schema reached through + /// this catalog. + pub codecs: FFI_ExtensionCodecBundle, /// 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. @@ -121,10 +119,10 @@ unsafe extern "C" fn schema_fn_wrapper( let maybe_schema = provider.inner().schema(name.as_str()); maybe_schema .map(|schema| { - FFI_SchemaProvider::new_with_ffi_codec( + FFI_SchemaProvider::new( schema, provider.runtime(), - provider.logical_codec.clone(), + provider.codecs.clone(), ) }) .into() @@ -144,11 +142,7 @@ unsafe extern "C" fn register_schema_fn_wrapper( let returned_schema = sresult_return!(inner_provider.register_schema(name.as_str(), schema)) .map(|schema| { - FFI_SchemaProvider::new_with_ffi_codec( - schema, - runtime, - provider.logical_codec.clone(), - ) + FFI_SchemaProvider::new(schema, runtime, provider.codecs.clone()) }) .into(); @@ -171,11 +165,7 @@ unsafe extern "C" fn deregister_schema_fn_wrapper( FFI_Result::Ok( maybe_schema .map(|schema| { - FFI_SchemaProvider::new_with_ffi_codec( - schema, - runtime, - provider.logical_codec.clone(), - ) + FFI_SchemaProvider::new(schema, runtime, provider.codecs.clone()) }) .into(), ) @@ -209,7 +199,7 @@ unsafe extern "C" fn clone_fn_wrapper( schema: schema_fn_wrapper, register_schema: register_schema_fn_wrapper, deregister_schema: deregister_schema_fn_wrapper, - logical_codec: provider.logical_codec.clone(), + codecs: provider.codecs.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -227,27 +217,13 @@ impl Drop for FFI_CatalogProvider { impl FFI_CatalogProvider { /// Creates a new [`FFI_CatalogProvider`]. + /// + /// `codecs` must describe the extension nodes used by every table below this + /// catalog, since schemas and tables reached through it inherit it. pub fn new( provider: Arc, runtime: Option, - task_ctx_provider: impl Into, - logical_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(), - ); - Self::new_with_ffi_codec(provider, runtime, logical_codec) - } - - pub fn new_with_ffi_codec( - provider: Arc, - runtime: Option, - logical_codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> Self { if let Some(provider) = provider.downcast_ref::() { return provider.0.clone(); @@ -260,7 +236,7 @@ impl FFI_CatalogProvider { schema: schema_fn_wrapper, register_schema: register_schema_fn_wrapper, deregister_schema: deregister_schema_fn_wrapper, - logical_codec, + codecs, clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -325,11 +301,7 @@ impl CatalogProvider for ForeignCatalogProvider { unsafe { let schema = match schema.downcast_ref::() { Some(s) => &s.0, - None => &FFI_SchemaProvider::new_with_ffi_codec( - schema, - None, - self.0.logical_codec.clone(), - ), + None => &FFI_SchemaProvider::new(schema, None, self.0.codecs.clone()), }; let returned_schema: Option = df_result!((self.0.register_schema)(&self.0, name.into(), schema))? @@ -375,9 +347,9 @@ mod tests { .is_none() ); let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let codecs = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); - let mut ffi_catalog = - FFI_CatalogProvider::new(catalog, None, task_ctx_provider, None); + let mut ffi_catalog = FFI_CatalogProvider::new(catalog, None, codecs); ffi_catalog.library_marker_id = crate::mock_foreign_marker_id; let foreign_catalog: Arc = (&ffi_catalog).into(); @@ -421,8 +393,8 @@ mod tests { let catalog = Arc::new(MemoryCatalogProvider::new()); let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); - let mut ffi_catalog = - FFI_CatalogProvider::new(catalog, None, task_ctx_provider, None); + let codecs = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); + let mut ffi_catalog = FFI_CatalogProvider::new(catalog, None, codecs); // Verify local libraries can be downcast to their original let foreign_catalog: Arc = (&ffi_catalog).into(); diff --git a/datafusion/ffi/src/catalog_provider_list.rs b/datafusion/ffi/src/catalog_provider_list.rs index af7cfd0f870cd..2127df5dc9b71 100644 --- a/datafusion/ffi/src/catalog_provider_list.rs +++ b/datafusion/ffi/src/catalog_provider_list.rs @@ -19,16 +19,12 @@ use std::ffi::c_void; use std::sync::Arc; use datafusion_catalog::{CatalogProvider, CatalogProviderList}; -use datafusion_proto::logical_plan::{ - DefaultLogicalExtensionCodec, LogicalExtensionCodec, -}; use stabby::string::String as SString; use stabby::vec::Vec as SVec; use tokio::runtime::Handle; use crate::catalog_provider::{FFI_CatalogProvider, ForeignCatalogProvider}; -use crate::execution::FFI_TaskContextProvider; -use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::util::FFI_Option; /// A stable struct for sharing [`CatalogProviderList`] across FFI boundaries. @@ -49,7 +45,9 @@ pub struct FFI_CatalogProviderList { pub catalog: unsafe extern "C" fn(&Self, name: SString) -> FFI_Option, - pub logical_codec: FFI_LogicalExtensionCodec, + /// The serialization environment propagated to every catalog reached through + /// this list. + pub codecs: FFI_ExtensionCodecBundle, /// 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. @@ -117,11 +115,7 @@ unsafe extern "C" fn register_catalog_fn_wrapper( inner_provider .register_catalog(name.into(), catalog) .map(|catalog| { - FFI_CatalogProvider::new_with_ffi_codec( - catalog, - runtime, - provider.logical_codec.clone(), - ) + FFI_CatalogProvider::new(catalog, runtime, provider.codecs.clone()) }) .into() } @@ -137,11 +131,7 @@ unsafe extern "C" fn catalog_fn_wrapper( inner_provider .catalog(name.as_str()) .map(|catalog| { - FFI_CatalogProvider::new_with_ffi_codec( - catalog, - runtime, - provider.logical_codec.clone(), - ) + FFI_CatalogProvider::new(catalog, runtime, provider.codecs.clone()) }) .into() } @@ -173,7 +163,7 @@ unsafe extern "C" fn clone_fn_wrapper( register_catalog: register_catalog_fn_wrapper, catalog_names: catalog_names_fn_wrapper, catalog: catalog_fn_wrapper, - logical_codec: provider.logical_codec.clone(), + codecs: provider.codecs.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -191,26 +181,13 @@ impl Drop for FFI_CatalogProviderList { impl FFI_CatalogProviderList { /// Creates a new [`FFI_CatalogProviderList`]. + /// + /// `codecs` must describe the extension nodes used by every table below this + /// list, since catalogs, schemas, and tables reached through it inherit it. pub fn new( provider: Arc, runtime: Option, - task_ctx_provider: impl Into, - logical_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(), - ); - Self::new_with_ffi_codec(provider, runtime, logical_codec) - } - pub fn new_with_ffi_codec( - provider: Arc, - runtime: Option, - logical_codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> Self { if let Some(provider) = provider.downcast_ref::() { return provider.0.clone(); @@ -222,7 +199,7 @@ impl FFI_CatalogProviderList { register_catalog: register_catalog_fn_wrapper, catalog_names: catalog_names_fn_wrapper, catalog: catalog_fn_wrapper, - logical_codec, + codecs, clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -268,11 +245,7 @@ impl CatalogProviderList for ForeignCatalogProviderList { unsafe { let catalog = match catalog.downcast_ref::() { Some(s) => &s.0, - None => &FFI_CatalogProvider::new_with_ffi_codec( - catalog, - None, - self.0.logical_codec.clone(), - ), + None => &FFI_CatalogProvider::new(catalog, None, self.0.codecs.clone()), }; (self.0.register_catalog)(&self.0, name.into(), catalog) @@ -303,9 +276,71 @@ impl CatalogProviderList for ForeignCatalogProviderList { #[cfg(test)] mod tests { + use arrow::array::record_batch; use datafusion::catalog::{MemoryCatalogProvider, MemoryCatalogProviderList}; + use datafusion_catalog::{MemTable, MemorySchemaProvider, SchemaProvider}; + use datafusion_common::Result; + use datafusion_expr::ptr_eq::arc_ptr_eq; + use datafusion_proto::physical_plan::PhysicalExtensionCodec; use super::*; + use crate::proto::physical_extension_codec::tests::TestExtensionCodec; + use crate::schema_provider::FFI_SchemaProvider; + use crate::table_provider::FFI_TableProvider; + + /// Walking a catalog hierarchy must hand the same serialization environment to + /// every wrapper it creates. Otherwise a table found this way would export + /// sessions that cannot serialize the nodes its owner's codecs describe. + #[tokio::test] + async fn test_nested_construction_preserves_codecs() -> Result<()> { + let batch = record_batch!(("a", Int32, [1, 2, 3]))?; + let table = Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]])?); + let schema = Arc::new(MemorySchemaProvider::new()); + schema.register_table("t".to_owned(), table)?; + let catalog = Arc::new(MemoryCatalogProvider::new()); + catalog.register_schema("s", schema)?; + let catalog_list = Arc::new(MemoryCatalogProviderList::new()); + catalog_list.register_catalog("c".to_owned(), catalog); + + let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let physical_codec = + Arc::new(TestExtensionCodec {}) as Arc; + let codecs = FFI_ExtensionCodecBundle::new( + task_ctx_provider, + None, + Arc::new(datafusion_proto::logical_plan::DefaultLogicalExtensionCodec {}), + Arc::clone(&physical_codec), + ); + + let assert_carries_codec = |codecs: &FFI_ExtensionCodecBundle, label: &str| { + assert!( + arc_ptr_eq(&codecs.to_physical_codec(), &physical_codec), + "{label} lost the physical codec" + ); + }; + + let ffi_catalog_list = FFI_CatalogProviderList::new(catalog_list, None, codecs); + assert_carries_codec(&ffi_catalog_list.codecs, "catalog list"); + + let ffi_catalog: Option = + unsafe { (ffi_catalog_list.catalog)(&ffi_catalog_list, "c".into()) }.into(); + let ffi_catalog = ffi_catalog.expect("catalog \"c\" should exist"); + assert_carries_codec(&ffi_catalog.codecs, "catalog"); + + let ffi_schema: Option = + unsafe { (ffi_catalog.schema)(&ffi_catalog, "s".into()) }.into(); + let ffi_schema = ffi_schema.expect("schema \"s\" should exist"); + assert_carries_codec(&ffi_schema.codecs, "schema"); + + let ffi_table = crate::df_result!(unsafe { + (ffi_schema.table)(&ffi_schema, "t".into()).await + })?; + let ffi_table: Option = ffi_table.into(); + let ffi_table = ffi_table.expect("table \"t\" should exist"); + assert_carries_codec(&ffi_table.codecs, "table provider"); + + Ok(()) + } #[test] fn test_round_trip_ffi_catalog_provider_list() { @@ -320,8 +355,9 @@ mod tests { ); let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let codecs = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); let mut ffi_catalog_list = - FFI_CatalogProviderList::new(catalog_list, None, task_ctx_provider, None); + FFI_CatalogProviderList::new(catalog_list, None, codecs); ffi_catalog_list.library_marker_id = crate::mock_foreign_marker_id; let foreign_catalog_list: Arc = @@ -361,8 +397,9 @@ mod tests { let catalog_list = Arc::new(MemoryCatalogProviderList::new()); let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let codecs = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); let mut ffi_catalog_list = - FFI_CatalogProviderList::new(catalog_list, None, task_ctx_provider, None); + FFI_CatalogProviderList::new(catalog_list, None, codecs); // Verify local libraries can be downcast to their original let foreign_catalog_list: Arc = diff --git a/datafusion/ffi/src/proto/extension_codec_bundle.rs b/datafusion/ffi/src/proto/extension_codec_bundle.rs new file mode 100644 index 0000000000000..40875d27ef668 --- /dev/null +++ b/datafusion/ffi/src/proto/extension_codec_bundle.rs @@ -0,0 +1,380 @@ +// 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. + +//! The complete serialization environment used at an FFI boundary. +//! +//! Serializing DataFusion plans across a library boundary requires three values +//! that must agree with one another: a [`FFI_TaskContextProvider`] that resolves +//! the current [`TaskContext`], a logical extension codec, and a physical +//! extension codec. [`FFI_ExtensionCodecBundle`] carries them as one unit so +//! that a wrapper cannot be built from a codec and a provider that were never +//! configured together, and so that a wrapper which only needs to serialize +//! logical plans still has a physical codec to hand to any session it exports. +//! +//! # Structure +//! +//! The bundle is exactly its three members. Each one already carries its own +//! `clone` / `release` function pointers, `version` extern, and +//! `library_marker_id`, so the bundle holds no private data and needs no +//! foreign adapter of its own; `Clone` and `Drop` come from the members. +//! +//! The dependency direction is bundle to codecs to task context provider. A +//! bundle must never be stored inside a codec: cloning a bundle clones its +//! logical codec, so a bundle field on a codec would make cloning recurse until +//! the stack is exhausted. + +use std::sync::Arc; + +use datafusion_common::error::Result; +use datafusion_execution::TaskContext; +use datafusion_proto::logical_plan::{ + DefaultLogicalExtensionCodec, LogicalExtensionCodec, +}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, +}; +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; + +/// A stable struct describing one complete serialization environment. +/// +/// The fields are private so that the constructors are the only way to pair a +/// task context provider with the codecs that use it. Read access is through +/// [`Self::task_ctx_provider`], [`Self::logical_codec`], and +/// [`Self::physical_codec`]. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_ExtensionCodecBundle { + task_ctx_provider: FFI_TaskContextProvider, + + logical_codec: FFI_LogicalExtensionCodec, + + physical_codec: FFI_PhysicalExtensionCodec, +} + +impl FFI_ExtensionCodecBundle { + /// Creates a bundle from one task context provider and native codecs. + /// + /// Both codecs are exported using `task_ctx_provider`. `runtime` is attached + /// to both codecs for callbacks that must enter the exporting library's Tokio + /// runtime. + /// + /// The provider is held weakly. It must outlive every wrapper built from + /// this bundle; otherwise codec callbacks fail with a clear error rather + /// than resolving a stale [`TaskContext`]. See [`Self::task_ctx`]. + pub fn new( + task_ctx_provider: impl Into, + runtime: Option, + logical_codec: Arc, + physical_codec: Arc, + ) -> Self { + let task_ctx_provider = task_ctx_provider.into(); + let logical_codec = FFI_LogicalExtensionCodec::new( + logical_codec, + runtime.clone(), + task_ctx_provider.clone(), + ); + let physical_codec = FFI_PhysicalExtensionCodec::new( + physical_codec, + runtime, + task_ctx_provider.clone(), + ); + + Self { + task_ctx_provider, + logical_codec, + physical_codec, + } + } + + /// Creates a bundle whose codecs support built-in nodes only. + /// + /// Use this when no custom logical or physical extension nodes cross the + /// boundary. A wrapper built from this bundle cannot round-trip custom + /// extension nodes; attempting it fails during encoding or decoding rather + /// than silently dropping the node. + pub fn new_default( + task_ctx_provider: impl Into, + runtime: Option, + ) -> Self { + Self::new( + task_ctx_provider, + runtime, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), + ) + } + + /// Creates a bundle from codecs that already crossed an FFI boundary. + /// + /// Cloning or re-exporting FFI codecs does not nest foreign wrappers: each + /// codec's `clone` returns a handle owned by its original library. + /// + /// This constructor cannot verify that the three arguments belong together, + /// so the caller asserts that both codecs were exported with + /// `task_ctx_provider` and that they can round-trip every extension node + /// the resulting wrappers expose. Prefer [`Self::new`] whenever the native + /// codecs are available. + pub fn new_with_ffi_codecs( + task_ctx_provider: FFI_TaskContextProvider, + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + ) -> Self { + Self { + task_ctx_provider, + logical_codec, + physical_codec, + } + } + + /// Pairs an existing logical codec with an explicit default physical codec. + /// + /// Used by the paths inside [`FFI_LogicalExtensionCodec`] that build a nested + /// [`FFI_TableProvider`](crate::table_provider::FFI_TableProvider) from a + /// function pointer receiving the codec alone, where a codec cannot carry a + /// bundle (see the [module documentation](self)). The task context provider is + /// cloned out of `logical_codec`. + /// + /// A table provider built from this bundle cannot carry custom physical + /// extension nodes through a session callback. If the library that decoded + /// the provider scans it with a session local to that library, and the + /// provider reaches back through [`Session::query_planner`] for a custom + /// physical node, the call fails with `PhysicalExtensionCodec is not + /// provided`. When the provider is instead scanned with the exporting + /// library's own session handle, exporting that session short-circuits to + /// the original handle and its complete bundle, so that topology is + /// unaffected. + /// + /// [`Session::query_planner`]: datafusion_session::Session::query_planner + pub(crate) fn new_logical_with_default_physical( + logical_codec: FFI_LogicalExtensionCodec, + runtime: Option, + ) -> Self { + let task_ctx_provider = logical_codec.task_ctx_provider.clone(); + let physical_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(DefaultPhysicalExtensionCodec {}), + runtime, + task_ctx_provider.clone(), + ); + + Self { + task_ctx_provider, + logical_codec, + physical_codec, + } + } + + /// The task context provider shared by both codecs. + pub fn task_ctx_provider(&self) -> &FFI_TaskContextProvider { + &self.task_ctx_provider + } + + /// The logical extension codec. + pub fn logical_codec(&self) -> &FFI_LogicalExtensionCodec { + &self.logical_codec + } + + /// The physical extension codec. + pub fn physical_codec(&self) -> &FFI_PhysicalExtensionCodec { + &self.physical_codec + } + + /// Resolves the current [`TaskContext`]. + /// + /// Returns an error if the [`TaskContextProvider`] this bundle was built + /// from has been dropped. The context is resolved on every call rather than + /// captured at construction, so functions and other session state + /// registered after the bundle was created are visible to codec callbacks. + /// + /// [`TaskContextProvider`]: datafusion_execution::TaskContextProvider + pub fn task_ctx(&self) -> Result> { + (&self.task_ctx_provider).try_into() + } + + /// The logical codec as a native trait object. + /// + /// Returns the underlying codec directly when it is owned by this library, + /// and a foreign adapter otherwise. + pub fn to_logical_codec(&self) -> Arc { + (&self.logical_codec).into() + } + + /// The physical codec as a native trait object. + /// + /// Returns the underlying codec directly when it is owned by this library, + /// and a foreign adapter otherwise. + pub fn to_physical_codec(&self) -> Arc { + (&self.physical_codec).into() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::prelude::SessionContext; + use datafusion_common::{DataFusionError, Result}; + use datafusion_execution::TaskContextProvider; + use datafusion_expr::ptr_eq::arc_ptr_eq; + use datafusion_proto::logical_plan::{ + DefaultLogicalExtensionCodec, LogicalExtensionCodec, + }; + use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, + }; + + use crate::execution::FFI_TaskContextProvider; + use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; + use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; + use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; + use crate::proto::physical_extension_codec::tests::TestExtensionCodec; + + /// Both codec traits require [`Any`](std::any::Any), so a codec trait object + /// can be re-borrowed as one to check the concrete type behind it. + fn logical_is(codec: &Arc) -> bool { + let any_ref: &dyn std::any::Any = codec.as_ref(); + any_ref.is::() + } + + fn physical_is(codec: &Arc) -> bool { + let any_ref: &dyn std::any::Any = codec.as_ref(); + any_ref.is::() + } + + #[test] + fn bundle_new_retains_both_codecs() -> Result<()> { + let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let logical = Arc::new(TestExtensionCodec {}) as Arc; + let physical = Arc::new(TestExtensionCodec {}) as Arc; + + let bundle = FFI_ExtensionCodecBundle::new( + task_ctx_provider, + None, + Arc::clone(&logical), + Arc::clone(&physical), + ); + + // Both codecs are local, so the conversions hand back the originals. + assert!(arc_ptr_eq(&bundle.to_logical_codec(), &logical)); + assert!(arc_ptr_eq(&bundle.to_physical_codec(), &physical)); + bundle.task_ctx()?; + + Ok(()) + } + + #[test] + fn bundle_new_default_uses_default_codecs() { + let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let bundle = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); + + assert!(logical_is::( + &bundle.to_logical_codec() + )); + assert!(physical_is::( + &bundle.to_physical_codec() + )); + } + + #[test] + fn bundle_clone_preserves_codec_identity() -> Result<()> { + let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let logical = Arc::new(TestExtensionCodec {}) as Arc; + let physical = Arc::new(TestExtensionCodec {}) as Arc; + + let bundle = FFI_ExtensionCodecBundle::new( + task_ctx_provider, + None, + Arc::clone(&logical), + Arc::clone(&physical), + ); + let cloned = bundle.clone(); + + // Cloning must not wrap the codecs in another foreign layer. + assert!(arc_ptr_eq(&cloned.to_logical_codec(), &logical)); + assert!(arc_ptr_eq(&cloned.to_physical_codec(), &physical)); + cloned.task_ctx()?; + + Ok(()) + } + + #[test] + fn bundle_new_with_ffi_codecs_retains_supplied_codecs() { + let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let logical = Arc::new(TestExtensionCodec {}) as Arc; + let physical = Arc::new(TestExtensionCodec {}) as Arc; + let ffi_logical = FFI_LogicalExtensionCodec::new( + Arc::clone(&logical), + None, + task_ctx_provider.clone(), + ); + let ffi_physical = FFI_PhysicalExtensionCodec::new( + Arc::clone(&physical), + None, + task_ctx_provider.clone(), + ); + + let bundle = FFI_ExtensionCodecBundle::new_with_ffi_codecs( + task_ctx_provider, + ffi_logical, + ffi_physical, + ); + + assert!(arc_ptr_eq(&bundle.to_logical_codec(), &logical)); + assert!(arc_ptr_eq(&bundle.to_physical_codec(), &physical)); + } + + #[test] + fn bundle_logical_with_default_physical_is_explicit() { + let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let logical = Arc::new(TestExtensionCodec {}) as Arc; + let ffi_logical = + FFI_LogicalExtensionCodec::new(Arc::clone(&logical), None, task_ctx_provider); + + let bundle = FFI_ExtensionCodecBundle::new_logical_with_default_physical( + ffi_logical, + None, + ); + + assert!(arc_ptr_eq(&bundle.to_logical_codec(), &logical)); + assert!(physical_is::( + &bundle.to_physical_codec() + )); + // The provider is inherited from the logical codec, so it still resolves. + assert!(bundle.task_ctx().is_ok()); + } + + #[test] + fn bundle_reports_expired_task_context_provider() { + fn bundle_with_dropped_provider() -> FFI_ExtensionCodecBundle { + let ctx = Arc::new(SessionContext::new()); + let task_ctx_provider = Arc::clone(&ctx) as Arc; + FFI_ExtensionCodecBundle::new_default( + FFI_TaskContextProvider::from(&task_ctx_provider), + None, + ) + } + + let bundle = bundle_with_dropped_provider(); + let Err(DataFusionError::Ffi(message)) = bundle.task_ctx() else { + panic!("expected an out of scope error from an expired provider") + }; + assert!(message.contains("went out of scope"), "{message}"); + } +} diff --git a/datafusion/ffi/src/proto/logical_extension_codec.rs b/datafusion/ffi/src/proto/logical_extension_codec.rs index ed2c594f1bc02..1d3fdc4bda13f 100644 --- a/datafusion/ffi/src/proto/logical_extension_codec.rs +++ b/datafusion/ffi/src/proto/logical_extension_codec.rs @@ -40,6 +40,7 @@ use tokio::runtime::Handle; use crate::arrow_wrappers::WrappedSchema; use crate::execution::FFI_TaskContextProvider; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::table_provider::FFI_TableProvider; use crate::udaf::FFI_AggregateUDF; use crate::udf::FFI_ScalarUDF; @@ -163,11 +164,20 @@ unsafe extern "C" fn try_decode_table_provider_fn_wrapper( ctx.as_ref() )); - FFI_Result::Ok(FFI_TableProvider::new_with_ffi_codec( + // The function pointer receives only the codec, so there is no bundle to + // forward here. Pair this codec with an explicit default physical codec; see + // `FFI_ExtensionCodecBundle::new_logical_with_default_physical` for what that + // costs a consumer of the returned provider. + let codecs = FFI_ExtensionCodecBundle::new_logical_with_default_physical( + codec.clone(), + runtime.clone(), + ); + + FFI_Result::Ok(FFI_TableProvider::new( table_provider, true, runtime, - codec.clone(), + codecs, )) } @@ -404,8 +414,14 @@ impl LogicalExtensionCodec for ForeignLogicalExtensionCodec { buf: &mut Vec, ) -> Result<()> { let table_ref = table_ref.to_string(); - let node = - FFI_TableProvider::new_with_ffi_codec(node, true, None, self.0.clone()); + // As in `try_decode_table_provider_fn_wrapper`, only the codec is in scope. + // This handle exists solely so the owning library can encode `node`, so the + // default physical codec is never exercised. + let codecs = FFI_ExtensionCodecBundle::new_logical_with_default_physical( + self.0.clone(), + None, + ); + let node = FFI_TableProvider::new(node, true, None, codecs); let bytes = df_result!(unsafe { (self.0.try_encode_table_provider)(&self.0, table_ref.as_str().into(), node) diff --git a/datafusion/ffi/src/proto/mod.rs b/datafusion/ffi/src/proto/mod.rs index ae76027ecb64e..59703ce450571 100644 --- a/datafusion/ffi/src/proto/mod.rs +++ b/datafusion/ffi/src/proto/mod.rs @@ -15,5 +15,6 @@ // specific language governing permissions and limitations // under the License. +pub mod extension_codec_bundle; pub mod logical_extension_codec; pub mod physical_extension_codec; diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index 6d895d65c5fc1..2a607c334d39c 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -67,15 +67,11 @@ use datafusion_proto::bytes::{ physical_plan_from_bytes_with_extension_codec, physical_plan_to_bytes_with_extension_codec, }; -use datafusion_proto::logical_plan::LogicalExtensionCodec; -use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_session::{QueryPlanner, Session}; use stabby::vec::Vec as SVec; use tokio::runtime::Handle; -use crate::execution::FFI_TaskContextProvider; -use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; -use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::session::{FFI_SessionRef, ForeignSession}; use crate::util::FFI_Result; use crate::{df_result, sresult_return}; @@ -94,11 +90,9 @@ pub struct FFI_QueryPlanner { session: FFI_SessionRef, ) -> FfiFuture>>, - /// Codec used to encode and decode logical plans and extension nodes. - logical_codec: FFI_LogicalExtensionCodec, - - /// Codec used to encode and decode physical plans and extension nodes. - physical_codec: FFI_PhysicalExtensionCodec, + /// The serialization environment used for the logical plan travelling in and + /// the physical plan travelling out, including their extension nodes. + codecs: FFI_ExtensionCodecBundle, /// Used to create a clone of the query planner. clone: unsafe extern "C" fn(planner: &Self) -> Self, @@ -130,6 +124,11 @@ impl FFI_QueryPlanner { let private_data = self.private_data as *const QueryPlannerPrivateData; unsafe { &(*private_data).planner } } + + /// The serialization environment this planner encodes plans with. + pub fn codecs(&self) -> &FFI_ExtensionCodecBundle { + &self.codecs + } } unsafe extern "C" fn create_physical_plan_fn_wrapper( @@ -138,9 +137,8 @@ unsafe extern "C" fn create_physical_plan_fn_wrapper( 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(); + let logical_codec = planner.codecs.to_logical_codec(); + let physical_codec = planner.codecs.to_physical_codec(); async move { let mut foreign_session = None; @@ -194,8 +192,7 @@ unsafe extern "C" fn clone_fn_wrapper(planner: &FFI_QueryPlanner) -> FFI_QueryPl FFI_QueryPlanner { create_physical_plan: create_physical_plan_fn_wrapper, - logical_codec: planner.logical_codec.clone(), - physical_codec: planner.physical_codec.clone(), + codecs: planner.codecs.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -217,48 +214,26 @@ impl Clone for FFI_QueryPlanner { } impl FFI_QueryPlanner { - /// Creates an [`FFI_QueryPlanner`] with native extension codecs. + /// Creates an [`FFI_QueryPlanner`]. /// - /// 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: Arc, - physical_codec: Arc, - ) -> Self { - let task_ctx_provider = task_ctx_provider.into(); - let logical_codec = FFI_LogicalExtensionCodec::new( - logical_codec, - runtime.clone(), - task_ctx_provider.clone(), - ); - let physical_codec = - FFI_PhysicalExtensionCodec::new(physical_codec, runtime, task_ctx_provider); - Self::new_with_ffi_codecs(planner, logical_codec, physical_codec) - } - - /// Creates an [`FFI_QueryPlanner`] using prebuilt FFI extension codecs. + /// `codecs` decides which extension nodes survive the boundary: its logical + /// codec decodes the incoming plan and its physical codec encodes the planner's + /// result. Use + /// [`FFI_ExtensionCodecBundle::new_default`](crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle::new_default) + /// when no custom nodes are involved. /// - /// 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( + /// If `planner` is already foreign, its original FFI handle is re-exported + /// instead of being wrapped again, and that handle adopts `codecs`: a planner + /// serializes in the environment of the session that invokes it, not the one it + /// came from. + pub fn new( planner: Arc, - logical_codec: FFI_LogicalExtensionCodec, - physical_codec: FFI_PhysicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> Self { let any_ref: &dyn std::any::Any = planner.as_ref(); if let Some(planner) = any_ref.downcast_ref::() { let mut planner = planner.0.clone(); - planner.logical_codec = logical_codec; - planner.physical_codec = physical_codec; + planner.codecs = codecs; return planner; } @@ -266,8 +241,7 @@ impl FFI_QueryPlanner { Self { create_physical_plan: create_physical_plan_fn_wrapper, - logical_codec, - physical_codec, + codecs, clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -293,23 +267,17 @@ impl FFI_QueryPlanner { session: &dyn Session, session_runtime: Option, ) -> Result> { - let codec: Arc = (&self.logical_codec).into(); + let codec = self.codecs.to_logical_codec(); let logical_plan = logical_plan_to_bytes_with_extension_codec(logical_plan, codec.as_ref())?; let logical_plan = SVec::from(logical_plan.as_ref()); let task_ctx = session.task_ctx(); - let session = FFI_SessionRef::new_with_ffi_codecs( - session, - session_runtime, - self.logical_codec.clone(), - self.physical_codec.clone(), - ); + let session = FFI_SessionRef::new(session, session_runtime, self.codecs.clone()); let physical_plan = unsafe { df_result!((self.create_physical_plan)(self, logical_plan, session).await)? }; - let physical_codec: Arc = - (&self.physical_codec).into(); + let physical_codec = self.codecs.to_physical_codec(); physical_plan_from_bytes_with_extension_codec( physical_plan.as_slice(), @@ -385,13 +353,13 @@ mod tests { 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, + let codecs = FFI_ExtensionCodecBundle::new( &task_ctx_provider, + None, Arc::new(DefaultLogicalExtensionCodec {}), Arc::new(DefaultPhysicalExtensionCodec {}), - ) + ); + FFI_QueryPlanner::new(Arc::new(EmptyQueryPlanner), codecs) } #[test] diff --git a/datafusion/ffi/src/schema_provider.rs b/datafusion/ffi/src/schema_provider.rs index 8441b84b48697..1a35932921b27 100644 --- a/datafusion/ffi/src/schema_provider.rs +++ b/datafusion/ffi/src/schema_provider.rs @@ -22,15 +22,11 @@ use async_ffi::{FfiFuture, FutureExt}; use async_trait::async_trait; use datafusion_catalog::{SchemaProvider, TableProvider}; use datafusion_common::error::{DataFusionError, Result}; -use datafusion_proto::logical_plan::{ - DefaultLogicalExtensionCodec, LogicalExtensionCodec, -}; use stabby::string::String as SString; 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::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::table_provider::{FFI_TableProvider, ForeignTableProvider}; use crate::util::{FFI_Option, FFI_Result}; use crate::{df_result, sresult_return}; @@ -65,7 +61,9 @@ pub struct FFI_SchemaProvider { pub table_exist: unsafe extern "C" fn(provider: &Self, name: SString) -> bool, - pub logical_codec: FFI_LogicalExtensionCodec, + /// The serialization environment propagated to every table provider reached + /// through this schema. + pub codecs: FFI_ExtensionCodecBundle, /// 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. @@ -128,14 +126,12 @@ unsafe extern "C" fn table_fn_wrapper( ) -> FfiFuture>> { unsafe { let runtime = provider.runtime(); - let logical_codec = provider.logical_codec.clone(); + let codecs = provider.codecs.clone(); let provider = Arc::clone(provider.inner()); async move { let table = sresult_return!(provider.table(name.as_str()).await) - .map(|t| { - FFI_TableProvider::new_with_ffi_codec(t, true, runtime, logical_codec) - }) + .map(|t| FFI_TableProvider::new(t, true, runtime, codecs)) .into(); FFI_Result::Ok(table) @@ -151,15 +147,13 @@ unsafe extern "C" fn register_table_fn_wrapper( ) -> FFI_Result> { unsafe { let runtime = provider.runtime(); - let logical_codec = provider.logical_codec.clone(); + let codecs = provider.codecs.clone(); let provider = provider.inner(); let table = Arc::new(ForeignTableProvider(table)); let returned_table = sresult_return!(provider.register_table(name.into(), table)) - .map(|t| { - FFI_TableProvider::new_with_ffi_codec(t, true, runtime, logical_codec) - }); + .map(|t| FFI_TableProvider::new(t, true, runtime, codecs)); FFI_Result::Ok(returned_table.into()) } @@ -171,13 +165,11 @@ unsafe extern "C" fn deregister_table_fn_wrapper( ) -> FFI_Result> { unsafe { let runtime = provider.runtime(); - let logical_codec = provider.logical_codec.clone(); + let codecs = provider.codecs.clone(); let provider = provider.inner(); let returned_table = sresult_return!(provider.deregister_table(name.as_str())) - .map(|t| { - FFI_TableProvider::new_with_ffi_codec(t, true, runtime, logical_codec) - }); + .map(|t| FFI_TableProvider::new(t, true, runtime, codecs)); FFI_Result::Ok(returned_table.into()) } @@ -219,7 +211,7 @@ unsafe extern "C" fn clone_fn_wrapper( register_table: register_table_fn_wrapper, deregister_table: deregister_table_fn_wrapper, table_exist: table_exist_fn_wrapper, - logical_codec: provider.logical_codec.clone(), + codecs: provider.codecs.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -237,27 +229,13 @@ impl Drop for FFI_SchemaProvider { impl FFI_SchemaProvider { /// Creates a new [`FFI_SchemaProvider`]. + /// + /// `codecs` must describe the extension nodes used by every table in this + /// schema, since tables reached through it inherit it. pub fn new( provider: Arc, runtime: Option, - task_ctx_provider: impl Into, - logical_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(), - ); - Self::new_with_ffi_codec(provider, runtime, logical_codec) - } - - pub fn new_with_ffi_codec( - provider: Arc, - runtime: Option, - logical_codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> Self { if let Some(provider) = provider.downcast_ref::() { return provider.0.clone(); @@ -273,7 +251,7 @@ impl FFI_SchemaProvider { register_table: register_table_fn_wrapper, deregister_table: deregister_table_fn_wrapper, table_exist: table_exist_fn_wrapper, - logical_codec, + codecs, clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -347,12 +325,7 @@ impl SchemaProvider for ForeignSchemaProvider { unsafe { let ffi_table = match table.downcast_ref::() { Some(t) => t.0.clone(), - None => FFI_TableProvider::new_with_ffi_codec( - table, - true, - None, - self.0.logical_codec.clone(), - ), + None => FFI_TableProvider::new(table, true, None, self.0.codecs.clone()), }; let returned_provider: Option = @@ -403,9 +376,10 @@ mod tests { ); let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let codecs = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); let mut ffi_schema_provider = - FFI_SchemaProvider::new(schema_provider, None, task_ctx_provider, None); + FFI_SchemaProvider::new(schema_provider, None, codecs); ffi_schema_provider.library_marker_id = crate::mock_foreign_marker_id; let foreign_schema_provider: Arc = @@ -457,8 +431,8 @@ mod tests { let schema_provider = Arc::new(MemorySchemaProvider::new()); let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); - let mut ffi_schema = - FFI_SchemaProvider::new(schema_provider, None, task_ctx_provider, None); + let codecs = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); + let mut ffi_schema = FFI_SchemaProvider::new(schema_provider, None, codecs); // Verify local libraries can be downcast to their original let foreign_schema: Arc = (&ffi_schema).into(); diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 085a24074b358..b4ddeb1f4518f 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -17,6 +17,18 @@ //! FFI support for [`Session`]. //! +//! # Serialization environment +//! +//! An exported session carries one +//! [`crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle`]. +//! Every callback that moves a plan or expression across the boundary uses it: +//! the logical codec for [`Session::optimize`], [`Session::create_physical_expr`], +//! and [`Session::create_physical_plan`], and the physical codec for the planner +//! returned by [`Session::query_planner`]. Because the bundle pairs both codecs +//! with the task context provider that exported them, a planner reached through a +//! session can round-trip the same custom extension nodes as the wrapper that +//! exported the session. +//! //! # Delegating physical planning //! //! Consider a session owned by library A that uses a query planner owned by @@ -53,13 +65,11 @@ use datafusion_expr::{ use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; 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, + logical_plan_from_bytes_with_extension_codec, + 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, PhysicalOptimizerRule, QueryPlanner, Session, @@ -77,8 +87,7 @@ 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::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::query_planner::FFI_QueryPlanner; use crate::session::config::FFI_SessionConfig; use crate::udaf::FFI_AggregateUDF; @@ -143,9 +152,8 @@ pub(crate) struct FFI_SessionRef { physical_optimizers: unsafe extern "C" fn(&Self) -> SVec, - logical_codec: FFI_LogicalExtensionCodec, - - physical_codec: FFI_PhysicalExtensionCodec, + /// The serialization environment used by every callback on this session. + codecs: FFI_ExtensionCodecBundle, /// Used to create a clone on the provider of the registry. This should /// only need to be called by the receiver of the plan. @@ -201,28 +209,24 @@ unsafe extern "C" fn config_fn_wrapper(session: &FFI_SessionRef) -> FFI_SessionC unsafe extern "C" fn catalog_list_fn_wrapper( session: &FFI_SessionRef, ) -> FFI_CatalogProviderList { - FFI_CatalogProviderList::new_with_ffi_codec( + FFI_CatalogProviderList::new( session.inner().catalog_list(), unsafe { session.runtime() }.clone(), - session.logical_codec.clone(), + session.codecs.clone(), ) } 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(), - ) + FFI_QueryPlanner::new(session.inner().query_planner(), session.codecs.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 logical_codec = session.codecs.to_logical_codec(); let inner = session.inner(); let logical_plan = sresult_return!(logical_plan_from_bytes_with_extension_codec( logical_plan_serialized.as_slice(), @@ -246,13 +250,16 @@ unsafe extern "C" fn create_physical_plan_fn_wrapper( let runtime = session.runtime().clone(); let session = session.clone(); async move { + let logical_codec = session.codecs.to_logical_codec(); let session = session.inner(); let task_ctx = session.task_ctx(); - let logical_plan = sresult_return!(logical_plan_from_bytes( - logical_plan_serialized.as_slice(), - task_ctx.as_ref(), - )); + let logical_plan = + sresult_return!(logical_plan_from_bytes_with_extension_codec( + logical_plan_serialized.as_slice(), + task_ctx.as_ref(), + logical_codec.as_ref(), + )); let physical_plan = session.create_physical_plan(&logical_plan).await; @@ -267,7 +274,7 @@ unsafe extern "C" fn create_physical_expr_fn_wrapper( expr_serialized: SVec, schema: WrappedSchema, ) -> FFI_Result { - let codec: Arc = (&session.logical_codec).into(); + let codec = session.codecs.to_logical_codec(); let session = session.inner(); let logical_expr = LogicalExprNode::decode(expr_serialized.as_slice()).unwrap(); @@ -412,8 +419,7 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_SessionRef) -> FFI_SessionR 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(), + codecs: provider.codecs.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -431,50 +437,24 @@ impl Drop for FFI_SessionRef { } impl FFI_SessionRef { - /// Creates a new [`FFI_SessionRef`] with a default physical extension codec. + /// Creates a new [`FFI_SessionRef`]. /// - /// 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. + /// Both codecs in `codecs` are used by this session's callbacks: the logical + /// codec for [`Session::optimize`] and [`Session::create_physical_expr`], and + /// the physical codec for the query planner returned by + /// [`Session::query_planner`]. Pass a bundle whose codecs can round-trip every + /// extension node exposed through the session; the bundle's task context + /// provider must also remain live for the lifetime of the exported session. /// - /// 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. + /// When `session` is already a [`ForeignSession`], this re-exports its original + /// handle and discards `codecs`: the codecs a session was exported with are the + /// ones its callbacks need. Note that + /// [`FFI_QueryPlanner::new`](crate::query_planner::FFI_QueryPlanner::new) does + /// the reverse and adopts the supplied bundle. 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(), - logical_codec.task_ctx_provider.clone(), - ); - Self::new_with_ffi_codecs(session, runtime, logical_codec, physical_codec) - } - - /// Creates a new [`FFI_SessionRef`] using existing FFI codecs. - /// - /// The codecs must form a matching pair that can round-trip every logical and - /// physical extension node exposed through the session. Their task context - /// providers must remain live and return contexts appropriate for their decode - /// callbacks. - pub fn new_with_ffi_codecs( - session: &dyn Session, - runtime: Option, - logical_codec: FFI_LogicalExtensionCodec, - physical_codec: FFI_PhysicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> Self { if let Some(session) = session.as_any().downcast_ref::() { return session.session.clone(); @@ -497,8 +477,7 @@ impl FFI_SessionRef { default_table_options: default_table_options_fn_wrapper, task_ctx: task_ctx_fn_wrapper, physical_optimizers: physical_optimizers_fn_wrapper, - logical_codec, - physical_codec, + codecs, clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -718,8 +697,7 @@ impl Session for ForeignSession { fn optimize(&self, plan: &LogicalPlan) -> datafusion_common::Result { unsafe { - let codec: Arc = - (&self.session.logical_codec).into(); + let codec = self.session.codecs.to_logical_codec(); let logical_plan = logical_plan_to_bytes_with_extension_codec(plan, codec.as_ref())?; let optimized_plan = df_result!((self.session.optimize)( @@ -739,7 +717,9 @@ impl Session for ForeignSession { logical_plan: &LogicalPlan, ) -> datafusion_common::Result> { unsafe { - let logical_plan = logical_plan_to_bytes(logical_plan)?; + let codec = self.session.codecs.to_logical_codec(); + let logical_plan = + logical_plan_to_bytes_with_extension_codec(logical_plan, codec.as_ref())?; let physical_plan = df_result!( (self.session.create_physical_plan)( &self.session, @@ -759,8 +739,7 @@ impl Session for ForeignSession { df_schema: &DFSchema, ) -> datafusion_common::Result> { unsafe { - let codec: Arc = - (&self.session.logical_codec).into(); + let codec = self.session.codecs.to_logical_codec(); let logical_expr = serialize_expr(&expr, codec.as_ref())?.encode_to_vec(); let schema = WrappedSchema(FFI_ArrowSchema::try_from(df_schema.as_arrow())?); @@ -849,10 +828,13 @@ mod tests { use datafusion::execution::SessionStateBuilder; use datafusion_common::DataFusionError; use datafusion_expr::col; + use datafusion_expr::ptr_eq::arc_ptr_eq; use datafusion_expr::registry::FunctionRegistry; use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; + use datafusion_proto::physical_plan::PhysicalExtensionCodec; use super::*; + use crate::proto::physical_extension_codec::tests::TestExtensionCodec; static QUERY_PLANNER_CALLS: AtomicUsize = AtomicUsize::new(0); static PHYSICAL_OPTIMIZER_CALLS: AtomicUsize = AtomicUsize::new(0); @@ -877,13 +859,9 @@ mod tests { 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 codecs = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); let state = ctx.state(); - let mut local_session = FFI_SessionRef::new(&state, None, logical_codec); + let mut local_session = FFI_SessionRef::new(&state, None, codecs); local_session.query_planner = counting_query_planner; local_session.physical_optimizers = counting_physical_optimizers; @@ -907,6 +885,31 @@ mod tests { Ok(()) } + /// An exported session must hand its own physical codec to the query planner + /// it exports. A default codec there stops custom physical extension nodes + /// from crossing a session callback. + #[test] + fn test_exported_query_planner_keeps_session_physical_codec() { + let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let physical_codec = + Arc::new(TestExtensionCodec {}) as Arc; + let codecs = FFI_ExtensionCodecBundle::new( + task_ctx_provider, + None, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::clone(&physical_codec), + ); + + let state = ctx.state(); + let local_session = FFI_SessionRef::new(&state, None, codecs); + let planner = unsafe { (local_session.query_planner)(&local_session) }; + + assert!(arc_ptr_eq( + &planner.codecs().to_physical_codec(), + &physical_codec + )); + } + #[tokio::test] async fn test_ffi_session() -> Result<(), DataFusionError> { let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); @@ -923,13 +926,9 @@ mod tests { .with_table_options(table_options) .build(); - let logical_codec = FFI_LogicalExtensionCodec::new( - Arc::new(DefaultLogicalExtensionCodec {}), - None, - task_ctx_provider, - ); + let codecs = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); - let local_session = FFI_SessionRef::new(&state, None, logical_codec); + let local_session = FFI_SessionRef::new(&state, None, codecs); 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 ee9377bff064e..70b8237c77147 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -28,11 +28,9 @@ use datafusion_execution::TaskContext; use datafusion_expr::dml::InsertOp; use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; use datafusion_physical_plan::ExecutionPlan; +use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::logical_plan::from_proto::parse_exprs; use datafusion_proto::logical_plan::to_proto::serialize_exprs; -use datafusion_proto::logical_plan::{ - DefaultLogicalExtensionCodec, LogicalExtensionCodec, -}; use datafusion_proto::protobuf::LogicalExprList; use prost::Message; @@ -42,8 +40,7 @@ use tokio::runtime::Handle; use super::execution_plan::FFI_ExecutionPlan; use super::insert_op::FFI_InsertOp; use crate::arrow_wrappers::WrappedSchema; -use crate::execution::FFI_TaskContextProvider; -use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::session::{FFI_SessionRef, ForeignSession}; use crate::statistics::{deserialize_statistics, serialize_statistics}; use crate::table_source::{FFI_TableProviderFilterPushDown, FFI_TableType}; @@ -140,7 +137,9 @@ pub struct FFI_TableProvider { /// `Some(bytes)` is a prost-encoded `datafusion_proto_common::Statistics`. pub statistics: unsafe extern "C" fn(provider: &Self) -> FFI_Option>, - pub logical_codec: FFI_LogicalExtensionCodec, + /// The serialization environment used by this provider's callbacks and by the + /// sessions it exports for `scan` and `insert_into`. + pub codecs: FFI_ExtensionCodecBundle, /// 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. @@ -232,10 +231,8 @@ unsafe extern "C" fn supports_filters_pushdown_fn_wrapper( provider: &FFI_TableProvider, filters_serialized: SVec, ) -> FFI_Result> { - let logical_codec: Arc = (&provider.logical_codec).into(); - let task_ctx = sresult_return!(>::try_from( - &provider.logical_codec.task_ctx_provider - )); + let logical_codec = provider.codecs.to_logical_codec(); + let task_ctx = sresult_return!(provider.codecs.task_ctx()); supports_filters_pushdown_internal( provider.inner(), &filters_serialized, @@ -252,10 +249,9 @@ unsafe extern "C" fn scan_fn_wrapper( filters_serialized: SVec, limit: FFI_Option, ) -> FfiFuture> { - let task_ctx: Result, DataFusionError> = - (&provider.logical_codec.task_ctx_provider).try_into(); + let task_ctx = provider.codecs.task_ctx(); let runtime = provider.runtime().clone(); - let logical_codec: Arc = (&provider.logical_codec).into(); + let logical_codec = provider.codecs.to_logical_codec(); let internal_provider = Arc::clone(provider.inner()); async move { @@ -362,7 +358,7 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_TableProvider) -> FFI_Table supports_filters_pushdown: provider.supports_filters_pushdown, insert_into: provider.insert_into, statistics: statistics_fn_wrapper, - logical_codec: provider.logical_codec.clone(), + codecs: provider.codecs.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -379,34 +375,19 @@ impl Drop for FFI_TableProvider { impl FFI_TableProvider { /// Creates a new [`FFI_TableProvider`]. + /// + /// `codecs`'s logical codec serializes the filter expressions handed to + /// `scan` and `supports_filters_pushdown`, and the whole bundle is attached to + /// the session this provider exports, so a consumer that reaches + /// [`Session::query_planner`] through that session gets a planner configured + /// with the same codecs. Use + /// [`FFI_ExtensionCodecBundle::new_default`] when no custom extension nodes are + /// involved. pub fn new( provider: Arc, can_support_pushdown_filters: bool, runtime: Option, - task_ctx_provider: impl Into, - logical_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(), - ); - Self::new_with_ffi_codec( - provider, - can_support_pushdown_filters, - runtime, - logical_codec, - ) - } - - pub fn new_with_ffi_codec( - provider: Arc, - can_support_pushdown_filters: bool, - runtime: Option, - logical_codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> Self { if let Some(provider) = provider.downcast_ref::() { return provider.0.clone(); @@ -423,7 +404,7 @@ impl FFI_TableProvider { }, insert_into: insert_into_fn_wrapper, statistics: statistics_fn_wrapper, - logical_codec, + codecs, clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -492,13 +473,13 @@ 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.codecs.clone()); let projections: FFI_Option> = projection .map(|p| p.iter().map(|v| v.to_owned()).collect()) .into(); - let codec: Arc = (&self.0.logical_codec).into(); + let codec = self.0.codecs.to_logical_codec(); let filter_list = LogicalExprList { expr: serialize_exprs(filters, codec.as_ref())?, }; @@ -537,7 +518,7 @@ impl TableProvider for ForeignTableProvider { } }; - let codec: Arc = (&self.0.logical_codec).into(); + let codec = self.0.codecs.to_logical_codec(); let expr_list = LogicalExprList { expr: serialize_exprs( @@ -562,7 +543,7 @@ 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.codecs.clone()); let rc = Handle::try_current().ok(); let input = FFI_ExecutionPlan::new(input, rc); @@ -613,15 +594,46 @@ mod tests { )?)) } + /// A bundle holds its task context provider weakly. When the provider is gone, + /// a wrapper carrying that bundle must report it rather than panic or resolve a + /// stale context. + #[test] + fn test_expired_task_context_provider_reports_clear_error() -> Result<()> { + fn codecs_with_dropped_provider() -> FFI_ExtensionCodecBundle { + let ctx = Arc::new(SessionContext::new()); + let task_ctx_provider = Arc::clone(&ctx) as Arc; + FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None) + } + + let mut ffi_provider = FFI_TableProvider::new( + create_test_table_provider()?, + true, + None, + codecs_with_dropped_provider(), + ); + ffi_provider.library_marker_id = crate::mock_foreign_marker_id; + let foreign: Arc = (&ffi_provider).into(); + + let filter = col("a").gt(lit(3.0)); + let error = foreign + .supports_filters_pushdown(&[&filter]) + .expect_err("an expired task context provider should surface an error"); + assert!( + error.to_string().contains("went out of scope"), + "unexpected error: {error}" + ); + + Ok(()) + } + #[tokio::test] async fn test_round_trip_ffi_table_provider_scan() -> Result<()> { let provider = create_test_table_provider()?; let ctx = Arc::new(SessionContext::new()); let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); - let mut ffi_provider = - FFI_TableProvider::new(provider, true, None, task_ctx_provider, None); + let mut ffi_provider = FFI_TableProvider::new(provider, true, None, codecs); ffi_provider.library_marker_id = crate::mock_foreign_marker_id; let foreign_table_provider: Arc = (&ffi_provider).into(); @@ -643,10 +655,9 @@ mod tests { let provider = create_test_table_provider()?; let ctx = Arc::new(SessionContext::new()); let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); - let mut ffi_provider = - FFI_TableProvider::new(provider, true, None, task_ctx_provider, None); + let mut ffi_provider = FFI_TableProvider::new(provider, true, None, codecs); ffi_provider.library_marker_id = crate::mock_foreign_marker_id; let foreign_table_provider: Arc = (&ffi_provider).into(); @@ -691,12 +702,11 @@ mod tests { let ctx = Arc::new(SessionContext::new()); let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); let provider = Arc::new(MemTable::try_new(schema, vec![vec![batch1]])?); - let mut ffi_provider = - FFI_TableProvider::new(provider, true, None, task_ctx_provider, None); + let mut ffi_provider = FFI_TableProvider::new(provider, true, None, codecs); ffi_provider.library_marker_id = crate::mock_foreign_marker_id; let foreign_table_provider: Arc = (&ffi_provider).into(); @@ -725,9 +735,8 @@ mod tests { let table_provider = create_test_table_provider()?; let ctx = Arc::new(SessionContext::new()) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&ctx); - let mut ffi_table = - FFI_TableProvider::new(table_provider, false, None, task_ctx_provider, None); + let codecs = FFI_ExtensionCodecBundle::new_default(&ctx, None); + let mut ffi_table = FFI_TableProvider::new(table_provider, false, None, codecs); // Verify local libraries can be downcast to their original let foreign_table: Arc = (&ffi_table).into(); @@ -778,11 +787,10 @@ mod tests { let ctx = Arc::new(SessionContext::new()); let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); // Wrap in FFI and force the foreign path (not local bypass) - let mut ffi_provider = - FFI_TableProvider::new(provider, true, None, task_ctx_provider, None); + let mut ffi_provider = FFI_TableProvider::new(provider, true, None, codecs); ffi_provider.library_marker_id = crate::mock_foreign_marker_id; let foreign_table_provider: Arc = (&ffi_provider).into(); @@ -854,7 +862,7 @@ mod tests { let ctx = Arc::new(SessionContext::new()); let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); // Provider without statistics should cross the boundary as None. let no_stats_inner = Arc::new(MemTable::try_new( @@ -865,13 +873,8 @@ mod tests { inner: no_stats_inner, stats: None, }); - let mut ffi_provider = FFI_TableProvider::new( - no_stats_provider, - true, - None, - task_ctx_provider.clone(), - None, - ); + let mut ffi_provider = + FFI_TableProvider::new(no_stats_provider, true, None, codecs.clone()); ffi_provider.library_marker_id = crate::mock_foreign_marker_id; let foreign: Arc = (&ffi_provider).into(); assert!(foreign.statistics().is_none()); @@ -895,8 +898,7 @@ mod tests { inner: stats_inner, stats: Some(original_stats.clone()), }); - let mut ffi_provider = - FFI_TableProvider::new(stats_provider, true, None, task_ctx_provider, None); + let mut ffi_provider = FFI_TableProvider::new(stats_provider, true, None, codecs); ffi_provider.library_marker_id = crate::mock_foreign_marker_id; let foreign: Arc = (&ffi_provider).into(); assert_eq!(foreign.statistics().as_ref(), Some(&original_stats)); diff --git a/datafusion/ffi/src/table_provider_factory.rs b/datafusion/ffi/src/table_provider_factory.rs index 63ebb51bb1db8..20d5ba87eec4b 100644 --- a/datafusion/ffi/src/table_provider_factory.rs +++ b/datafusion/ffi/src/table_provider_factory.rs @@ -21,19 +21,15 @@ use async_ffi::{FfiFuture, FutureExt}; use async_trait::async_trait; use datafusion_catalog::{Session, TableProvider, TableProviderFactory}; use datafusion_common::error::{DataFusionError, Result}; -use datafusion_execution::TaskContext; use datafusion_expr::{CreateExternalTable, DdlStatement, LogicalPlan}; -use datafusion_proto::logical_plan::{ - AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec, -}; +use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::protobuf::LogicalPlanNode; use prost::Message; 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::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::session::{FFI_SessionRef, ForeignSession}; use crate::table_provider::{FFI_TableProvider, ForeignTableProvider}; use crate::util::FFI_Result; @@ -64,7 +60,9 @@ pub struct FFI_TableProviderFactory { cmd_serialized: SVec, ) -> FfiFuture>, - logical_codec: FFI_LogicalExtensionCodec, + /// The serialization environment used to decode the `CreateExternalTable` + /// command and propagated to the table providers this factory creates. + codecs: FFI_ExtensionCodecBundle, /// Used to create a clone of the factory. This should only need to be called /// by the receiver of the factory. @@ -95,34 +93,21 @@ struct FactoryPrivateData { } impl FFI_TableProviderFactory { - /// Creates a new [`FFI_TableProvider`]. + /// Creates a new [`FFI_TableProviderFactory`]. + /// + /// `codecs`'s logical codec encodes and decodes the `CreateExternalTable` + /// command, and the whole bundle is handed to every + /// [`FFI_TableProvider`] this factory creates. pub fn new( factory: Arc, runtime: Option, - task_ctx_provider: impl Into, - logical_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(), - ); - Self::new_with_ffi_codec(factory, runtime, logical_codec) - } - - pub fn new_with_ffi_codec( - factory: Arc, - runtime: Option, - logical_codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> Self { let private_data = Box::new(FactoryPrivateData { factory, runtime }); Self { create: create_fn_wrapper, - logical_codec, + codecs, clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -145,9 +130,8 @@ impl FFI_TableProviderFactory { &self, cmd_serialized: &SVec, ) -> Result { - let task_ctx: Arc = - (&self.logical_codec.task_ctx_provider).try_into()?; - let logical_codec: Arc = (&self.logical_codec).into(); + let task_ctx = self.codecs.task_ctx()?; + let logical_codec = self.codecs.to_logical_codec(); let plan = LogicalPlanNode::decode(cmd_serialized.as_ref()) .map_err(|e| DataFusionError::Internal(format!("{e:?}")))?; @@ -204,7 +188,7 @@ async fn create_fn_wrapper_impl( cmd_serialized: SVec, ) -> Result { let runtime = factory.runtime().clone(); - let ffi_logical_codec = factory.logical_codec.clone(); + let ffi_codecs = factory.codecs.clone(); let internal_factory = Arc::clone(factory.inner()); let cmd = factory.deserialize_cmd(&cmd_serialized)?; @@ -218,11 +202,11 @@ async fn create_fn_wrapper_impl( })?; let provider = internal_factory.create(session, &cmd).await?; - Ok(FFI_TableProvider::new_with_ffi_codec( + Ok(FFI_TableProvider::new( provider, true, runtime.clone(), - ffi_logical_codec, + ffi_codecs, )) } @@ -239,7 +223,7 @@ unsafe extern "C" fn clone_fn_wrapper( FFI_TableProviderFactory { create: create_fn_wrapper, - logical_codec: factory.logical_codec.clone(), + codecs: factory.codecs.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, version: super::version, @@ -269,8 +253,7 @@ impl ForeignTableProviderFactory { &self, cmd: CreateExternalTable, ) -> Result, DataFusionError> { - let logical_codec: Arc = - (&self.0.logical_codec).into(); + let logical_codec = self.0.codecs.to_logical_codec(); let plan = LogicalPlan::Ddl(DdlStatement::CreateExternalTable(Box::new(cmd))); let plan: LogicalPlanNode = @@ -293,7 +276,7 @@ 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.codecs.clone()); let cmd = self.serialize_cmd(cmd.clone())?; let provider = unsafe { @@ -356,11 +339,10 @@ mod tests { async fn test_round_trip_ffi_table_provider_factory() -> Result<()> { let ctx = Arc::new(SessionContext::new()); let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); let factory = Arc::new(TestTableProviderFactory {}); - let mut ffi_factory = - FFI_TableProviderFactory::new(factory, None, task_ctx_provider, None); + let mut ffi_factory = FFI_TableProviderFactory::new(factory, None, codecs); ffi_factory.library_marker_id = crate::mock_foreign_marker_id; let factory: Arc = (&ffi_factory).into(); @@ -393,11 +375,10 @@ mod tests { async fn test_ffi_table_provider_factory_clone() -> Result<()> { let ctx = Arc::new(SessionContext::new()); let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); let factory = Arc::new(TestTableProviderFactory {}); - let ffi_factory = - FFI_TableProviderFactory::new(factory, None, task_ctx_provider, None); + let ffi_factory = FFI_TableProviderFactory::new(factory, None, codecs); // Test that we can clone the factory let cloned_factory = ffi_factory.clone(); diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs index 9821c3e501f67..03742a3c42f9f 100644 --- a/datafusion/ffi/src/tests/async_provider.rs +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -43,7 +43,7 @@ use tokio::runtime::Handle; use tokio::sync::{broadcast, mpsc}; use super::create_record_batch; -use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::table_provider::FFI_TableProvider; #[derive(Debug)] @@ -285,13 +285,8 @@ impl Stream for AsyncTestRecordBatchStream { } pub(crate) fn create_async_table_provider( - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_TableProvider { let (table_provider, tokio_rt) = start_async_provider(); - FFI_TableProvider::new_with_ffi_codec( - Arc::new(table_provider), - true, - Some(tokio_rt), - codec, - ) + FFI_TableProvider::new(Arc::new(table_provider), true, Some(tokio_rt), codecs) } diff --git a/datafusion/ffi/src/tests/catalog.rs b/datafusion/ffi/src/tests/catalog.rs index b0b0858a8a3d7..67d425b072f23 100644 --- a/datafusion/ffi/src/tests/catalog.rs +++ b/datafusion/ffi/src/tests/catalog.rs @@ -38,7 +38,7 @@ use datafusion_common::{Result, exec_err}; use crate::catalog_provider::FFI_CatalogProvider; use crate::catalog_provider_list::FFI_CatalogProviderList; -use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; /// This schema provider is intended only for unit tests. It prepopulates with one /// table and only allows for tables named sales and purchases. @@ -169,10 +169,10 @@ impl CatalogProvider for FixedCatalogProvider { } pub(crate) extern "C" fn create_catalog_provider( - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_CatalogProvider { let catalog_provider = Arc::new(FixedCatalogProvider::default()); - FFI_CatalogProvider::new_with_ffi_codec(catalog_provider, None, codec) + FFI_CatalogProvider::new(catalog_provider, None, codecs) } /// This catalog provider list is intended only for unit tests. It prepopulates with one @@ -221,8 +221,8 @@ impl CatalogProviderList for FixedCatalogProviderList { } pub(crate) extern "C" fn create_catalog_provider_list( - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_CatalogProviderList { let catalog_provider_list = Arc::new(FixedCatalogProviderList::default()); - FFI_CatalogProviderList::new_with_ffi_codec(catalog_provider_list, None, codec) + FFI_CatalogProviderList::new(catalog_provider_list, None, codecs) } diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 357953674acde..f04701e362f95 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -42,8 +42,7 @@ use crate::config::extension_options::FFI_ExtensionOptions; 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::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::query_planner::FFI_QueryPlanner; use crate::table_provider::FFI_TableProvider; use crate::table_provider_factory::FFI_TableProviderFactory; @@ -71,21 +70,21 @@ pub mod utils; pub struct ForeignLibraryModule { /// Construct an opinionated catalog provider pub create_catalog: - extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_CatalogProvider, + extern "C" fn(codecs: FFI_ExtensionCodecBundle) -> FFI_CatalogProvider, /// Construct an opinionated catalog provider list pub create_catalog_list: - extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_CatalogProviderList, + extern "C" fn(codecs: FFI_ExtensionCodecBundle) -> FFI_CatalogProviderList, /// Constructs the table provider pub create_table: extern "C" fn( synchronous: bool, - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_TableProvider, /// Constructs the table provider factory pub create_table_factory: - extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProviderFactory, + extern "C" fn(codecs: FFI_ExtensionCodecBundle) -> FFI_TableProviderFactory, /// Create a scalar UDF pub create_scalar_udf: extern "C" fn() -> FFI_ScalarUDF, @@ -97,7 +96,7 @@ pub struct ForeignLibraryModule { pub create_placement_udf: extern "C" fn() -> FFI_ScalarUDF, pub create_table_function: - extern "C" fn(FFI_LogicalExtensionCodec) -> FFI_TableFunction, + extern "C" fn(FFI_ExtensionCodecBundle) -> FFI_TableFunction, /// Create an aggregate UDAF using sum pub create_sum_udaf: extern "C" fn() -> FFI_AggregateUDF, @@ -115,7 +114,7 @@ pub struct ForeignLibraryModule { pub create_exec_with_statistics: extern "C" fn() -> FFI_ExecutionPlan, pub create_table_with_statistics: - extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProvider, + extern "C" fn(codecs: FFI_ExtensionCodecBundle) -> FFI_TableProvider, pub create_physical_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, @@ -124,11 +123,21 @@ pub struct ForeignLibraryModule { /// 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, + codecs: FFI_ExtensionCodecBundle, library_a_planner: FFI_Option, ) -> FFI_QueryPlanner, + /// Construct a query planner that returns a custom physical extension node. + /// Used to check that a node produced behind a session callback survives the + /// boundary through the session's own physical codec. + pub create_extension_node_query_planner: + extern "C" fn(codecs: FFI_ExtensionCodecBundle) -> FFI_QueryPlanner, + + /// Construct a table provider whose `scan` plans through the query planner it + /// reaches on the session it is given, rather than planning locally. + pub create_session_planning_table: + extern "C" fn(codecs: FFI_ExtensionCodecBundle) -> FFI_TableProvider, + pub version: extern "C" fn() -> u64, /// Create an aggregate UDAF using first_value @@ -154,20 +163,20 @@ pub fn create_record_batch(start_value: i32, num_values: usize) -> RecordBatch { /// We create an in-memory table and convert it to it's FFI counterpart. extern "C" fn construct_table_provider( synchronous: bool, - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_TableProvider { match synchronous { - true => create_sync_table_provider(codec), - false => create_async_table_provider(codec), + true => create_sync_table_provider(codecs), + false => create_async_table_provider(codecs), } } /// Here we only wish to create a simple table provider as an example. /// We create an in-memory table and convert it to it's FFI counterpart. extern "C" fn construct_table_provider_factory( - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_TableProviderFactory { - table_provider_factory::create(codec) + table_provider_factory::create(codecs) } pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan { @@ -245,7 +254,7 @@ impl TableProvider for TableWithStats { } pub(crate) extern "C" fn create_table_with_statistics( - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_TableProvider { let schema = create_test_schema(); let batch = create_record_batch(1, 5); @@ -254,7 +263,7 @@ pub(crate) extern "C" fn create_table_with_statistics( inner, stats: make_test_statistics(), }); - FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec) + FFI_TableProvider::new(provider, true, None, codecs) } /// This defines the entry point for using the module. @@ -282,6 +291,9 @@ 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_extension_node_query_planner: + query_planner::create_extension_node_query_planner, + create_session_planning_table: query_planner::create_session_planning_table, version: super::version, create_first_value_udaf: create_ffi_first_value_func, } diff --git a/datafusion/ffi/src/tests/query_planner.rs b/datafusion/ffi/src/tests/query_planner.rs index 90d713f8a6336..e4d303065b978 100644 --- a/datafusion/ffi/src/tests/query_planner.rs +++ b/datafusion/ffi/src/tests/query_planner.rs @@ -18,11 +18,12 @@ use std::any::Any; use std::sync::Arc; -use arrow::datatypes::{DataType, Field, Schema}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use async_trait::async_trait; +use datafusion_catalog::TableProvider; use datafusion_catalog::default_table_source::source_as_provider; use datafusion_common::{Result, exec_err}; -use datafusion_expr::LogicalPlan; +use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder, TableType}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::sorts::sort::SortExec; @@ -30,11 +31,11 @@ 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::execution_plan::tests::EmptyExec as TestExtensionExec; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner}; use crate::session::ForeignSession; -use crate::table_provider::ForeignTableProvider; +use crate::table_provider::{FFI_TableProvider, ForeignTableProvider}; use crate::util::FFI_Option; #[derive(Debug)] @@ -151,14 +152,86 @@ impl QueryPlanner for SwappedQueryPlanner { } } +/// Library C's planner for the extension-node deployment. +/// +/// It returns a node that has no built-in protobuf representation, so the node can +/// only reach another library through a physical extension codec. +#[derive(Debug)] +struct ExtensionNodeQueryPlanner; + +#[async_trait] +impl QueryPlanner for ExtensionNodeQueryPlanner { + async fn create_physical_plan( + &self, + _logical_plan: &LogicalPlan, + _session: &dyn Session, + ) -> Result> { + Ok(Arc::new( + TestExtensionExec::new(super::create_test_schema()), + )) + } +} + +/// Library B's table provider for the extension-node deployment. +/// +/// Instead of planning its own scan, it plans through the query planner reachable +/// on the session library A handed it. That planner is library C's, and the node it +/// returns must travel back through library A's physical codec. +#[derive(Debug)] +struct SessionPlanningTableProvider; + +#[async_trait] +impl TableProvider for SessionPlanningTableProvider { + fn schema(&self) -> SchemaRef { + super::create_test_schema() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + session: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + if session.as_any().downcast_ref::().is_none() { + return exec_err!("library A's session was not foreign to library B"); + } + + let planner = session.query_planner(); + let planner_any: &dyn Any = planner.as_ref(); + if planner_any.downcast_ref::().is_none() { + return exec_err!("library C's planner did not cross the FFI boundary"); + } + + // The planner ignores the plan; it exists so the call has an argument that + // library A's logical codec can serialize. + let logical_plan = LogicalPlanBuilder::empty(false).build()?; + let plan = planner.create_physical_plan(&logical_plan, session).await?; + + // The node was reconstructed by library A's physical codec, so it is + // A-local and therefore foreign here. + if !plan.is::() { + return exec_err!( + "expected library A to reconstruct the extension node; got {}", + plan.name() + ); + } + + Ok(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, + codecs: FFI_ExtensionCodecBundle, library_a_planner: FFI_Option, ) -> FFI_QueryPlanner { let planner: Arc = match library_a_planner.as_ref() { @@ -168,5 +241,19 @@ pub extern "C" fn create_query_planner( None => Arc::new(TestQueryPlanner), }; - FFI_QueryPlanner::new_with_ffi_codecs(planner, logical_codec, physical_codec) + FFI_QueryPlanner::new(planner, codecs) +} + +/// Creates library C's planner for the extension-node deployment. +pub extern "C" fn create_extension_node_query_planner( + codecs: FFI_ExtensionCodecBundle, +) -> FFI_QueryPlanner { + FFI_QueryPlanner::new(Arc::new(ExtensionNodeQueryPlanner), codecs) +} + +/// Creates library B's table provider for the extension-node deployment. +pub extern "C" fn create_session_planning_table( + codecs: FFI_ExtensionCodecBundle, +) -> FFI_TableProvider { + FFI_TableProvider::new(Arc::new(SessionPlanningTableProvider), false, None, codecs) } diff --git a/datafusion/ffi/src/tests/sync_provider.rs b/datafusion/ffi/src/tests/sync_provider.rs index e3cb54fff90eb..01383fbca2de1 100644 --- a/datafusion/ffi/src/tests/sync_provider.rs +++ b/datafusion/ffi/src/tests/sync_provider.rs @@ -20,11 +20,11 @@ use std::sync::Arc; use datafusion_catalog::MemTable; use super::{create_record_batch, create_test_schema}; -use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::table_provider::FFI_TableProvider; pub(crate) fn create_sync_table_provider( - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_TableProvider { let schema = create_test_schema(); @@ -38,5 +38,5 @@ pub(crate) fn create_sync_table_provider( let table_provider = MemTable::try_new(schema, vec![batches]).unwrap(); - FFI_TableProvider::new_with_ffi_codec(Arc::new(table_provider), true, None, codec) + FFI_TableProvider::new(Arc::new(table_provider), true, None, codecs) } diff --git a/datafusion/ffi/src/tests/table_provider_factory.rs b/datafusion/ffi/src/tests/table_provider_factory.rs index 29af6aacf6484..c417d817c19cd 100644 --- a/datafusion/ffi/src/tests/table_provider_factory.rs +++ b/datafusion/ffi/src/tests/table_provider_factory.rs @@ -23,7 +23,7 @@ use datafusion_common::Result; use datafusion_expr::CreateExternalTable; use super::{create_record_batch, create_test_schema}; -use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::table_provider_factory::FFI_TableProviderFactory; #[derive(Debug)] @@ -52,7 +52,7 @@ impl TableProviderFactory for TestTableProviderFactory { } } -pub(crate) fn create(codec: FFI_LogicalExtensionCodec) -> FFI_TableProviderFactory { +pub(crate) fn create(codecs: FFI_ExtensionCodecBundle) -> FFI_TableProviderFactory { let factory = TestTableProviderFactory {}; - FFI_TableProviderFactory::new_with_ffi_codec(Arc::new(factory), None, codec) + FFI_TableProviderFactory::new(Arc::new(factory), None, codecs) } diff --git a/datafusion/ffi/src/tests/udf_udaf_udwf.rs b/datafusion/ffi/src/tests/udf_udaf_udwf.rs index 04d6fb26c1bc3..44dafa0fd4a1c 100644 --- a/datafusion/ffi/src/tests/udf_udaf_udwf.rs +++ b/datafusion/ffi/src/tests/udf_udaf_udwf.rs @@ -33,7 +33,7 @@ use datafusion_functions_aggregate::sum::Sum; use datafusion_functions_table::generate_series::RangeFunc; use datafusion_functions_window::rank::Rank; -use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::udaf::FFI_AggregateUDF; use crate::udf::FFI_ScalarUDF; use crate::udtf::FFI_TableFunction; @@ -172,11 +172,11 @@ pub(crate) extern "C" fn create_placement_func() -> FFI_ScalarUDF { } pub(crate) extern "C" fn create_ffi_table_func( - codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> FFI_TableFunction { let udtf: Arc = Arc::new(RangeFunc {}); - FFI_TableFunction::new_with_ffi_codec(udtf, None, codec) + FFI_TableFunction::new(udtf, None, codecs) } pub(crate) extern "C" fn create_ffi_sum_func() -> FFI_AggregateUDF { diff --git a/datafusion/ffi/src/udtf.rs b/datafusion/ffi/src/udtf.rs index fa28519d58de5..d03793a6312ff 100644 --- a/datafusion/ffi/src/udtf.rs +++ b/datafusion/ffi/src/udtf.rs @@ -22,20 +22,15 @@ use std::sync::Arc; use datafusion_catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider}; use datafusion_common::DataFusionError; use datafusion_common::error::Result; -use datafusion_execution::TaskContext; use datafusion_proto::logical_plan::from_proto::parse_exprs; use datafusion_proto::logical_plan::to_proto::serialize_exprs; -use datafusion_proto::logical_plan::{ - DefaultLogicalExtensionCodec, LogicalExtensionCodec, -}; use datafusion_proto::protobuf::LogicalExprList; use datafusion_session::Session; use prost::Message; 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::extension_codec_bundle::FFI_ExtensionCodecBundle; use crate::session::{FFI_SessionRef, ForeignSession}; use crate::table_provider::FFI_TableProvider; use crate::util::FFI_Result; @@ -63,7 +58,9 @@ pub struct FFI_TableFunction { session: FFI_SessionRef, ) -> FFI_Result, - pub logical_codec: FFI_LogicalExtensionCodec, + /// The serialization environment used for the argument expressions and + /// propagated to the table providers this function returns. + pub codecs: FFI_ExtensionCodecBundle, /// Used to create a clone on the provider of the udtf. This should /// only need to be called by the receiver of the udtf. @@ -109,9 +106,8 @@ unsafe extern "C" fn call_fn_wrapper( let runtime = udtf.runtime(); let udtf_inner = udtf.inner(); - let ctx: Arc = - sresult_return!((&udtf.logical_codec.task_ctx_provider).try_into()); - let codec: Arc = (&udtf.logical_codec).into(); + let ctx = sresult_return!(udtf.codecs.task_ctx()); + let codec = udtf.codecs.to_logical_codec(); let proto_filters = sresult_return!(LogicalExprList::decode(args.as_ref())); @@ -123,11 +119,11 @@ unsafe extern "C" fn call_fn_wrapper( #[expect(deprecated)] let table_provider = sresult_return!(udtf_inner.call(&args)); - FFI_Result::Ok(FFI_TableProvider::new_with_ffi_codec( + FFI_Result::Ok(FFI_TableProvider::new( table_provider, false, runtime, - udtf.logical_codec.clone(), + udtf.codecs.clone(), )) } @@ -139,9 +135,8 @@ unsafe extern "C" fn call_with_args_wrapper( let runtime = udtf.runtime(); let udtf_inner = udtf.inner(); - let ctx: Arc = - sresult_return!((&udtf.logical_codec.task_ctx_provider).try_into()); - let codec: Arc = (&udtf.logical_codec).into(); + let ctx = sresult_return!(udtf.codecs.task_ctx()); + let codec = udtf.codecs.to_logical_codec(); let proto_filters = sresult_return!(LogicalExprList::decode(args.as_ref())); @@ -164,11 +159,11 @@ unsafe extern "C" fn call_with_args_wrapper( let table_provider = sresult_return!( udtf_inner.call_with_args(TableFunctionArgs::new(&args, session)) ); - FFI_Result::Ok(FFI_TableProvider::new_with_ffi_codec( + FFI_Result::Ok(FFI_TableProvider::new( table_provider, false, runtime, - udtf.logical_codec.clone(), + udtf.codecs.clone(), )) } @@ -186,11 +181,7 @@ unsafe extern "C" fn clone_fn_wrapper(udtf: &FFI_TableFunction) -> FFI_TableFunc let runtime = udtf.runtime(); let udtf_inner = udtf.inner(); - FFI_TableFunction::new_with_ffi_codec( - Arc::clone(udtf_inner), - runtime, - udtf.logical_codec.clone(), - ) + FFI_TableFunction::new(Arc::clone(udtf_inner), runtime, udtf.codecs.clone()) } impl Clone for FFI_TableFunction { @@ -200,28 +191,15 @@ impl Clone for FFI_TableFunction { } impl FFI_TableFunction { + /// Creates a new [`FFI_TableFunction`]. + /// + /// `codecs`'s logical codec encodes and decodes the argument expressions, and + /// the whole bundle is attached both to the session this function exports and + /// to the table providers it returns. pub fn new( udtf: Arc, runtime: Option, - task_ctx_provider: impl Into, - logical_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(), - ); - - Self::new_with_ffi_codec(udtf, runtime, logical_codec) - } - - pub fn new_with_ffi_codec( - udtf: Arc, - runtime: Option, - logical_codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> Self { if let Some(udtf) = (Arc::clone(&udtf) as Arc).downcast_ref::() @@ -235,7 +213,7 @@ impl FFI_TableFunction { #[expect(deprecated)] call: call_fn_wrapper, call_with_args: call_with_args_wrapper, - logical_codec, + codecs, clone: clone_fn_wrapper, release: release_fn_wrapper, private_data: Box::into_raw(private_data) as *mut c_void, @@ -274,12 +252,9 @@ impl From for Arc { impl TableFunctionImpl for ForeignTableFunction { fn call_with_args(&self, args: TableFunctionArgs) -> Result> { - let session = FFI_SessionRef::new( - args.session(), - self.0.runtime(), - self.0.logical_codec.clone(), - ); - let codec: Arc = (&self.0.logical_codec).into(); + let session = + FFI_SessionRef::new(args.session(), self.0.runtime(), self.0.codecs.clone()); + let codec = self.0.codecs.to_logical_codec(); let expr_list = LogicalExprList { expr: serialize_exprs(args.exprs(), codec.as_ref())?, }; @@ -295,7 +270,7 @@ impl TableFunctionImpl for ForeignTableFunction { } fn call(&self, args: &[datafusion_expr::Expr]) -> Result> { - let codec: Arc = (&self.0.logical_codec).into(); + let codec = self.0.codecs.to_logical_codec(); let expr_list = LogicalExprList { expr: serialize_exprs(args, codec.as_ref())?, }; @@ -415,14 +390,10 @@ mod tests { let original_udtf = Arc::new(TestUDTF {}) as Arc; let ctx = Arc::new(SessionContext::default()); let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); - - let mut local_udtf: FFI_TableFunction = FFI_TableFunction::new( - Arc::clone(&original_udtf), - None, - task_ctx_provider, - None, - ); + let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); + + let mut local_udtf: FFI_TableFunction = + FFI_TableFunction::new(Arc::clone(&original_udtf), None, codecs); local_udtf.library_marker_id = crate::mock_foreign_marker_id; let foreign_udf: Arc = local_udtf.into(); @@ -459,13 +430,9 @@ mod tests { let original_udtf = Arc::new(TestUDTF {}) as Arc; let ctx = Arc::new(SessionContext::default()) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&ctx); - let mut ffi_udtf = FFI_TableFunction::new( - Arc::clone(&original_udtf), - None, - task_ctx_provider, - None, - ); + let codecs = FFI_ExtensionCodecBundle::new_default(&ctx, None); + let mut ffi_udtf = + FFI_TableFunction::new(Arc::clone(&original_udtf), None, codecs); // Verify local libraries can be downcast to their original let foreign_udtf: Arc = ffi_udtf.clone().into(); diff --git a/datafusion/ffi/tests/ffi_catalog.rs b/datafusion/ffi/tests/ffi_catalog.rs index 440a435f75c7d..ea55b7e713607 100644 --- a/datafusion/ffi/tests/ffi_catalog.rs +++ b/datafusion/ffi/tests/ffi_catalog.rs @@ -29,9 +29,9 @@ mod tests { #[tokio::test] async fn test_catalog() -> datafusion_common::Result<()> { let module = get_module()?; - let (ctx, codec) = super::utils::ctx_and_codec(); + let (ctx, codecs) = super::utils::ctx_and_codecs(); - let ffi_catalog = (module.create_catalog)(codec); + let ffi_catalog = (module.create_catalog)(codecs); let foreign_catalog: Arc = (&ffi_catalog).into(); let _ = ctx.register_catalog("fruit", foreign_catalog); @@ -50,9 +50,9 @@ mod tests { #[tokio::test] async fn test_catalog_list() -> datafusion_common::Result<()> { let module = get_module()?; - let (ctx, codec) = super::utils::ctx_and_codec(); + let (ctx, codecs) = super::utils::ctx_and_codecs(); - let ffi_catalog_list = (module.create_catalog_list)(codec); + let ffi_catalog_list = (module.create_catalog_list)(codecs); let foreign_catalog_list: Arc = (&ffi_catalog_list).into(); diff --git a/datafusion/ffi/tests/ffi_integration.rs b/datafusion/ffi/tests/ffi_integration.rs index 86f953e262ead..6df92d3ffaf1a 100644 --- a/datafusion/ffi/tests/ffi_integration.rs +++ b/datafusion/ffi/tests/ffi_integration.rs @@ -38,11 +38,12 @@ mod tests { /// testing it via a different executable. async fn test_table_provider(synchronous: bool) -> Result<()> { let table_provider_module = get_module()?; - let (ctx, codec) = super::utils::ctx_and_codec(); + let (ctx, codecs) = super::utils::ctx_and_codecs(); // By calling the code below, the table provided will be created within // the module's code. - let ffi_table_provider = (table_provider_module.create_table)(synchronous, codec); + let ffi_table_provider = + (table_provider_module.create_table)(synchronous, codecs); // In order to access the table provider within this executable, we need to // turn it into a `TableProvider`. @@ -83,11 +84,11 @@ mod tests { #[test] fn test_ffi_table_provider_statistics_cross_library() -> Result<()> { let module = get_module()?; - let (_, codec) = super::utils::ctx_and_codec(); + let (_, codecs) = super::utils::ctx_and_codecs(); let expected = datafusion_ffi::tests::make_test_statistics(); - let ffi_provider = (module.create_table_with_statistics)(codec); + let ffi_provider = (module.create_table_with_statistics)(codecs); let foreign: Arc = (&ffi_provider).into(); assert_eq!(foreign.statistics().as_ref(), Some(&expected)); @@ -98,10 +99,10 @@ mod tests { #[tokio::test] async fn test_table_provider_factory() -> Result<()> { let table_provider_module = get_module()?; - let (ctx, codec) = super::utils::ctx_and_codec(); + let (ctx, codecs) = super::utils::ctx_and_codecs(); let ffi_table_provider_factory = - (table_provider_module.create_table_factory)(codec); + (table_provider_module.create_table_factory)(codecs); let foreign_table_provider_factory: Arc = (&ffi_table_provider_factory).into(); diff --git a/datafusion/ffi/tests/ffi_query_planner.rs b/datafusion/ffi/tests/ffi_query_planner.rs index c72b8c3e889ae..03d7bdb127c0b 100644 --- a/datafusion/ffi/tests/ffi_query_planner.rs +++ b/datafusion/ffi/tests/ffi_query_planner.rs @@ -31,10 +31,9 @@ mod tests { use datafusion_execution::{TaskContext, TaskContextProvider}; use datafusion_expr::logical_plan::Extension; use datafusion_expr::{LogicalPlan, col}; - use datafusion_ffi::execution::FFI_TaskContextProvider; use datafusion_ffi::execution_plan::ForeignExecutionPlan; - use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; - use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; + use datafusion_ffi::execution_plan::tests::EmptyExec as TestExtensionExec; + use datafusion_ffi::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; use datafusion_ffi::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner}; use datafusion_ffi::table_provider::ForeignTableProvider; use datafusion_ffi::tests::{ @@ -46,30 +45,20 @@ mod tests { 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::logical_plan::{ + DefaultLogicalExtensionCodec, LogicalExtensionCodec, + }; use datafusion_proto::physical_plan::{ - DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, - PhysicalProtoConverterExtension, + PhysicalExtensionCodec, PhysicalProtoConverterExtension, }; use datafusion_session::QueryPlanner; #[tokio::test] async fn test_ffi_query_planner() -> Result<(), DataFusionError> { let module = get_module()?; - let (ctx, logical_codec) = crate::utils::ctx_and_codec(); - let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); - let physical_codec = FFI_PhysicalExtensionCodec::new( - Arc::new(DefaultPhysicalExtensionCodec {}), - None, - task_ctx_provider, - ); + let (ctx, codecs) = crate::utils::ctx_and_codecs(); - let ffi_planner = (module.create_query_planner)( - logical_codec, - physical_codec, - FFI_Option::None, - ); + let ffi_planner = (module.create_query_planner)(codecs, FFI_Option::None); let planner: Arc = (&ffi_planner).into(); let any_ref: &dyn std::any::Any = planner.as_ref(); @@ -196,16 +185,11 @@ mod tests { .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()), + let codecs = FFI_ExtensionCodecBundle::new( + &task_ctx_provider, None, - ffi_task_ctx_provider.clone(), - ); - let physical_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(LibraryALogicalCodec::default()), Arc::new(LibraryAPhysicalCodec), - None, - ffi_task_ctx_provider, ); let library_b = get_module_copy("query_planner_library_b")?; @@ -213,7 +197,7 @@ mod tests { // 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 ffi_provider = (library_b.create_table)(true, codecs.clone()); let provider: Arc = (&ffi_provider).into(); assert!(provider.downcast_ref::().is_some()); ctx.register_table("library_b", provider)?; @@ -222,11 +206,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_query_planner)( - logical_codec, - physical_codec, - FFI_Option::None, - ); + let ffi_planner = (library_c.create_query_planner)(codecs, 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()); @@ -267,23 +247,18 @@ mod tests { .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()), + let codecs = FFI_ExtensionCodecBundle::new( + &task_ctx_provider, None, - ffi_task_ctx_provider.clone(), - ); - let physical_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(LibraryALogicalCodec::default()), 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 ffi_provider = (library_b.create_table)(true, codecs.clone()); let provider: Arc = (&ffi_provider).into(); ctx.register_table("library_b", provider)?; @@ -291,16 +266,12 @@ mod tests { // 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(), - ); + let ffi_library_a_planner = + FFI_QueryPlanner::new(library_a_planner, codecs.clone()); // Library C: builds its planner around A's planner. let ffi_planner = (library_c.create_query_planner)( - logical_codec, - physical_codec, + codecs, FFI_Option::Some(ffi_library_a_planner), ); let library_c_planner: Arc = @@ -345,4 +316,125 @@ mod tests { Ok(()) } + + /// Library A's physical codec for the extension-node deployment. + /// + /// The node it decodes has no built-in protobuf representation, so it can only + /// cross a boundary through this codec. Encoding accepts the foreign handle + /// library C hands over; decoding rebuilds the node with library A's own type + /// identity. + #[derive(Debug)] + struct LibraryAExtensionNodeCodec; + + impl PhysicalExtensionCodec for LibraryAExtensionNodeCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + if buf != b"library-c-extension-node" || !inputs.is_empty() { + return exec_err!("unexpected library C extension node payload"); + } + Ok(Arc::new(TestExtensionExec::new(create_test_schema()))) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + // Library C's node arrives as a foreign handle. If library A ever + // re-encodes a node its own `try_decode` produced, it sees the local + // type instead. + if !node.is::() && !node.is::() { + return exec_err!( + "expected library C's node to be foreign or A-local; got {}", + node.name() + ); + } + buf.extend_from_slice(b"library-c-extension-node"); + Ok(()) + } + } + + /// The scenario the extension codec bundle exists for. + /// + /// Library A owns the session, the task context provider, and a custom physical + /// codec. It installs library C's query planner and queries library B's table + /// provider. B reaches C's planner *through the session A handed it*, and the + /// node C returns is a custom physical extension node. Only A's physical codec + /// can move that node, so the session A exported must be carrying it. + /// + /// Before the bundle, exporting a session synthesized a + /// `DefaultPhysicalExtensionCodec`, and this path failed with + /// `PhysicalExtensionCodec is not provided`. + #[tokio::test] + async fn test_session_planner_round_trips_custom_physical_node() -> Result<()> { + // Library A: owns the session and the 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 codecs = FFI_ExtensionCodecBundle::new( + &task_ctx_provider, + None, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(LibraryAExtensionNodeCodec), + ); + + let library_b = get_module_copy("session_planner_library_b")?; + let library_c = get_module_copy("session_planner_library_c")?; + + // Library C: a planner that returns a custom physical extension node. + let ffi_planner = (library_c.create_extension_node_query_planner)(codecs.clone()); + 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 installs C's planner on the session it already owns. Mutating + // the existing state keeps the `Arc` identity stable, so the + // task context provider captured by the bundle 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; + + // Library B: a table provider that plans through the planner it finds on the + // session, rather than planning locally. + let ffi_provider = (library_b.create_session_planning_table)(codecs); + let provider: Arc = (&ffi_provider).into(); + assert!(provider.downcast_ref::().is_some()); + ctx.register_table("library_b", provider)?; + + // Scanning runs A -> B -> A -> C -> A -> B -> A. The extension node is built + // in C, encoded and decoded by A's physical codec, and handed back to A. + let plan = provider_scan(&ctx).await?; + + // The node was reconstructed inside library A, so A can downcast it. + assert!(!plan.is::()); + assert!( + plan.is::(), + "library A could not downcast the node its codec rebuilt; got {}", + plan.name() + ); + + Ok(()) + } + + /// Scans the registered `library_b` table directly, so the assertions above see + /// the plan library B returned rather than a wrapper library A added. + async fn provider_scan(ctx: &SessionContext) -> Result> { + let provider = ctx + .table_provider("library_b") + .await + .map_err(|e| DataFusionError::Plan(e.to_string()))?; + let state = ctx.state(); + provider.scan(&state, None, &[], None).await + } } diff --git a/datafusion/ffi/tests/ffi_udtf.rs b/datafusion/ffi/tests/ffi_udtf.rs index 69e5de90e1364..9e1560be2564d 100644 --- a/datafusion/ffi/tests/ffi_udtf.rs +++ b/datafusion/ffi/tests/ffi_udtf.rs @@ -35,9 +35,9 @@ mod tests { #[tokio::test] async fn test_user_defined_table_function() -> Result<()> { let module = get_module()?; - let (ctx, codec) = super::utils::ctx_and_codec(); + let (ctx, codecs) = super::utils::ctx_and_codecs(); - let ffi_table_func = (module.create_table_function)(codec); + let ffi_table_func = (module.create_table_function)(codecs); let foreign_table_func: Arc = ffi_table_func.into(); ctx.register_udtf("my_range", foreign_table_func); diff --git a/datafusion/ffi/tests/utils/mod.rs b/datafusion/ffi/tests/utils/mod.rs index acf59de7f3464..a4dbc2229cbda 100644 --- a/datafusion/ffi/tests/utils/mod.rs +++ b/datafusion/ffi/tests/utils/mod.rs @@ -19,25 +19,18 @@ use std::sync::Arc; use datafusion::prelude::SessionContext; use datafusion_execution::TaskContextProvider; -use datafusion_ffi::execution::FFI_TaskContextProvider; -use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; -use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; +use datafusion_ffi::proto::extension_codec_bundle::FFI_ExtensionCodecBundle; -// Creates a default SessionContext and FFI Logical Extension Codec -// for use in FFI integration tests. +// Creates a default SessionContext and an extension codec bundle carrying the +// default logical and physical codecs, for use in FFI integration tests. // // This helper centralizes setup logic and is kept intentionally // for upcoming FFI test expansions. #[cfg_attr(not(feature = "integration-tests"), expect(dead_code))] -pub fn ctx_and_codec() -> (Arc, FFI_LogicalExtensionCodec) { +pub fn ctx_and_codecs() -> (Arc, FFI_ExtensionCodecBundle) { let ctx = Arc::new(SessionContext::default()); let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); - let codec = FFI_LogicalExtensionCodec::new( - Arc::new(DefaultLogicalExtensionCodec {}), - None, - task_ctx_provider, - ); + let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); - (ctx, codec) + (ctx, codecs) } 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..8161097ae64cf 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -904,6 +904,79 @@ not included in this release. See [PR #23703](https://github.com/apache/datafusion/pull/23703) for details on the catalog changes. +### FFI wrappers now take an `FFI_ExtensionCodecBundle` + +Serializing plans across an FFI boundary requires three values that must agree +with one another: an `FFI_TaskContextProvider`, a logical extension codec, and a +physical extension codec. The `datafusion-ffi` wrappers previously took the +provider and the logical codec separately and, when they needed a physical codec, +synthesized a `DefaultPhysicalExtensionCodec`. A query planner reached through an +exported session therefore could not round-trip custom physical extension nodes; +it failed later with `PhysicalExtensionCodec is not provided`. + +The new `datafusion_ffi::proto::extension_codec_bundle::FFI_ExtensionCodecBundle` +carries all three together. Its fields are private, so the constructors are the +only way to pair a provider with the codecs that use it. + +**Who is affected:** + +- Anyone constructing `FFI_TableProvider`, `FFI_TableProviderFactory`, + `FFI_TableFunction`, `FFI_CatalogProvider`, `FFI_CatalogProviderList`, + `FFI_SchemaProvider`, or `FFI_QueryPlanner`. +- FFI providers and consumers generally: these struct layouts changed, so both + sides must be rebuilt against DataFusion 55. + +**Migration guide:** + +Each wrapper now has a single constructor taking the bundle. The `new_with_ffi_codec` +and `new_with_ffi_codecs` variants are gone, as is the `Option>` +argument whose `None` meant "use the default": + +```rust,ignore +// Before +let ffi_table = FFI_TableProvider::new(table, true, None, task_ctx_provider, None); +let ffi_table = FFI_TableProvider::new_with_ffi_codec(table, true, None, ffi_logical_codec); + +// After: state both codecs once, then reuse the bundle +let codecs = FFI_ExtensionCodecBundle::new( + &task_ctx_provider, + None, // Option + Arc::new(MyLogicalCodec), + Arc::new(MyPhysicalCodec), +); +let ffi_table = FFI_TableProvider::new(table, true, None, codecs.clone()); +let ffi_catalog = FFI_CatalogProvider::new(catalog, None, codecs); +``` + +When no custom extension nodes cross the boundary, select the defaults +explicitly: + +```rust,ignore +let codecs = FFI_ExtensionCodecBundle::new_default(&task_ctx_provider, None); +``` + +`FFI_QueryPlanner::new` no longer takes a runtime, a provider, and two codecs +separately: + +```rust,ignore +// Before +let ffi_planner = FFI_QueryPlanner::new( + planner, + None, + &task_ctx_provider, + Arc::new(MyLogicalCodec), + Arc::new(MyPhysicalCodec), +); + +// After +let ffi_planner = FFI_QueryPlanner::new(planner, codecs); +``` + +The bundle is propagated unchanged through nested construction, so a table +provider reached by walking a catalog list serializes exactly like the list it +came from, and the session such a provider exports hands the same physical codec +to any planner a consumer reaches through it. + ### Unused `async` removed from several public functions Public functions that were declared `async` but never awaited anything are now