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, 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/logical_extension_codec.rs b/datafusion/ffi/src/proto/logical_extension_codec.rs index 97aa5c901a636..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. @@ -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 9e64df82e31b4..95d2ed68a6ea3 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. @@ -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 new file mode 100644 index 0000000000000..6d895d65c5fc1 --- /dev/null +++ b/datafusion/ffi/src/query_planner.rs @@ -0,0 +1,445 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI support for [`QueryPlanner`]. +//! +//! A typical deployment has three libraries. Library A (for example, +//! `datafusion-python`) owns the [`Session`] and codec registry. Library B owns +//! a custom table provider and its extension nodes. Library C (for example, +//! Ballista or `datafusion-distributed`) owns the query planner. A serializes a +//! logical plan and invokes C, while `FFI_SessionRef` lets C call session +//! services in A. C deserializes the logical plan, creates a physical plan, +//! serializes that result, and returns it for A to deserialize. The logical and +//! physical extension codecs preserve nodes supplied by B. +//! +//! The physical result is serialized instead of returned as an +//! [`crate::execution_plan::FFI_ExecutionPlan`]. An FFI execution-plan handle is +//! a foreign trait-object proxy, so even a built-in plan created in C cannot be +//! downcast to its concrete +//! type in A. Serialization reconstructs known plan nodes with A's local Rust +//! type identities, allowing A's optimizers and other consumers to downcast +//! them. Extension codecs control how custom nodes are reconstructed. +//! +//! A node returned by B while C is planning is still foreign to C unless a +//! codec boundary reconstructs it in C. The query-planner boundary guarantees +//! that C-local serializable nodes, and extension nodes understood by the +//! configured codecs, are reconstructed for A when the completed plan returns. +//! +//! # Delegating back to library A +//! +//! C commonly wants A's built-in planning as a starting point, then rewrites the +//! result. A must export its planner *before* installing C's planner on the +//! session, and C must retain that handle: after the swap, +//! [`Session::query_planner`] reports C's own planner, and +//! [`Session::create_physical_plan`] dispatches to it, so either one is a +//! self-call. Delegating to the retained handle is safe, because DataFusion's +//! built-in physical planner never re-dispatches through [`Session`]. +//! +//! Retain the planner rather than the session. [`FFI_QueryPlanner`] owns a +//! reference-counted planner, so it outlives A's original session, whereas +//! `FFI_SessionRef` borrows its session with the lifetime erased. + +use std::ffi::c_void; +use std::sync::Arc; + +use async_ffi::{FfiFuture, FutureExt}; +use async_trait::async_trait; +use datafusion_common::error::{DataFusionError, Result}; +use datafusion_expr::LogicalPlan; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_proto::bytes::{ + logical_plan_from_bytes_with_extension_codec, + logical_plan_to_bytes_with_extension_codec, + physical_plan_from_bytes_with_extension_codec, + physical_plan_to_bytes_with_extension_codec, +}; +use datafusion_proto::logical_plan::LogicalExtensionCodec; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use datafusion_session::{QueryPlanner, Session}; +use stabby::vec::Vec as SVec; +use tokio::runtime::Handle; + +use crate::execution::FFI_TaskContextProvider; +use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::session::{FFI_SessionRef, ForeignSession}; +use crate::util::FFI_Result; +use crate::{df_result, sresult_return}; + +/// An ABI-stable handle to a [`QueryPlanner`] owned by another library. +/// +/// The Rust-facing adapters serialize the input [`LogicalPlan`] and resulting +/// [`ExecutionPlan`]; callers do not invoke the byte-oriented function pointer +/// directly. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_QueryPlanner { + create_physical_plan: unsafe extern "C" fn( + &Self, + logical_plan_serialized: SVec, + 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, + + /// 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, DataFusionError>) + .unwrap_or_else(|| { + foreign_session = Some(ForeignSession::try_from(&session)?); + Ok(foreign_session.as_ref().unwrap()) + }) + ); + + let logical_plan = sresult_return!(logical_plan_from_bytes_with_extension_codec( + logical_plan_serialized.as_slice(), + session.task_ctx().as_ref(), + logical_codec.as_ref(), + )); + + let physical_plan = sresult_return!( + internal_planner + .create_physical_plan(&logical_plan, session) + .await + ); + let physical_plan = sresult_return!(physical_plan_to_bytes_with_extension_codec( + physical_plan, + physical_codec.as_ref(), + )); + + FFI_Result::Ok(SVec::from(physical_plan.as_ref())) + } + .into_ffi() +} + +unsafe extern "C" fn release_fn_wrapper(planner: &mut FFI_QueryPlanner) { + unsafe { + debug_assert!(!planner.private_data.is_null()); + let private_data = + Box::from_raw(planner.private_data as *mut QueryPlannerPrivateData); + drop(private_data); + planner.private_data = std::ptr::null_mut(); + } +} + +unsafe extern "C" fn clone_fn_wrapper(planner: &FFI_QueryPlanner) -> FFI_QueryPlanner { + let old_planner = Arc::clone(planner.inner()); + + let private_data = Box::into_raw(Box::new(QueryPlannerPrivateData { + planner: old_planner, + })) as *mut c_void; + + FFI_QueryPlanner { + create_physical_plan: create_physical_plan_fn_wrapper, + logical_codec: planner.logical_codec.clone(), + physical_codec: planner.physical_codec.clone(), + clone: clone_fn_wrapper, + release: release_fn_wrapper, + version: super::version, + private_data, + library_marker_id: crate::get_library_marker_id, + } +} + +impl Drop for FFI_QueryPlanner { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +impl Clone for FFI_QueryPlanner { + fn clone(&self) -> Self { + unsafe { (self.clone)(self) } + } +} + +impl FFI_QueryPlanner { + /// Creates an [`FFI_QueryPlanner`] with native extension codecs. + /// + /// Both codecs are required so that the caller states which extension nodes + /// survive the boundary. Pass + /// [`DefaultLogicalExtensionCodec`](datafusion_proto::logical_plan::DefaultLogicalExtensionCodec) + /// and + /// [`DefaultPhysicalExtensionCodec`](datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec) + /// when no custom nodes are involved. `runtime` and `task_ctx_provider` + /// support codec callbacks across the FFI boundary. + pub fn new( + planner: Arc, + 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. + /// + /// 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, + physical_codec: FFI_PhysicalExtensionCodec, + ) -> 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; + return planner; + } + + let private_data = Box::new(QueryPlannerPrivateData { planner }); + + Self { + create_physical_plan: create_physical_plan_fn_wrapper, + logical_codec, + physical_codec, + clone: clone_fn_wrapper, + release: release_fn_wrapper, + version: super::version, + private_data: Box::into_raw(private_data) as *mut c_void, + library_marker_id: crate::get_library_marker_id, + } + } + + /// Creates a physical plan through this planner's FFI interface. + /// + /// This serializes `logical_plan`, exports `session` as an + /// `FFI_SessionRef`, invokes the planner's owning library, and + /// deserializes its physical-plan response. `session_runtime` is attached + /// to the exported session for callbacks that need its Tokio runtime. + /// + /// The [`QueryPlanner`] implementation for [`ForeignQueryPlanner`] cannot + /// obtain the session owner's runtime from the trait API, so it calls this + /// method with `None`. Embedders that own the runtime and need session + /// callbacks to enter it must call this method directly with `Some(handle)`. + pub async fn create_physical_plan_with_session_runtime( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + session_runtime: Option, + ) -> 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 = SVec::from(logical_plan.as_ref()); + let task_ctx = session.task_ctx(); + let session = FFI_SessionRef::new_with_ffi_codecs( + session, + session_runtime, + self.logical_codec.clone(), + self.physical_codec.clone(), + ); + + let physical_plan = unsafe { + df_result!((self.create_physical_plan)(self, logical_plan, session).await)? + }; + let physical_codec: Arc = + (&self.physical_codec).into(); + + physical_plan_from_bytes_with_extension_codec( + physical_plan.as_slice(), + task_ctx.as_ref(), + physical_codec.as_ref(), + ) + } +} + +/// Consumer-side [`QueryPlanner`] adapter for an [`FFI_QueryPlanner`]. +/// +/// Calls serialize the logical plan, invoke the producing library, and +/// deserialize its physical-plan response. +#[derive(Debug)] +pub struct ForeignQueryPlanner(pub FFI_QueryPlanner); + +unsafe impl Send for ForeignQueryPlanner {} +unsafe impl Sync for ForeignQueryPlanner {} + +impl From<&FFI_QueryPlanner> for Arc { + 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 datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; + use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; + + use super::*; + + #[derive(Debug)] + struct EmptyQueryPlanner; + + #[async_trait] + impl QueryPlanner for EmptyQueryPlanner { + async fn create_physical_plan( + &self, + _logical_plan: &LogicalPlan, + _session: &dyn Session, + ) -> Result> { + 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, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), + ) + } + + #[test] + fn test_ffi_query_planner_local_bypass() { + let ctx = Arc::new(SessionContext::new()); + let ffi_planner = create_ffi_query_planner(ctx); + let planner: Arc = (&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(()) + } + + #[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(()) + } +} diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 519384379edb8..085a24074b358 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -15,10 +15,25 @@ // specific language governing permissions and limitations // under the License. +//! FFI support for [`Session`]. +//! +//! # Delegating physical planning +//! +//! Consider a session owned by library A that uses a query planner owned by +//! library C. After A installs C's planner, [`ForeignSession::query_planner`] +//! returns C's planner and [`ForeignSession::create_physical_plan`] dispatches +//! to C's planner. C must not call `create_physical_plan`, or invoke the planner +//! returned by `query_planner`, to delegate planning back to A. Repeating either +//! self-call recurses until the stack is exhausted. +//! +//! To delegate safely, A must export its original planner before installing C's +//! planner, and C must retain and invoke that planner directly. See the +//! [`crate::query_planner`] module for details. + use std::any::Any; use std::collections::HashMap; use std::ffi::c_void; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use arrow_schema::SchemaRef; use arrow_schema::ffi::FFI_ArrowSchema; @@ -37,12 +52,18 @@ use datafusion_expr::{ }; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; -use datafusion_proto::bytes::{logical_plan_from_bytes, logical_plan_to_bytes}; +use datafusion_proto::bytes::{ + logical_plan_from_bytes, logical_plan_from_bytes_with_extension_codec, + logical_plan_to_bytes, logical_plan_to_bytes_with_extension_codec, +}; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::logical_plan::from_proto::parse_expr; use datafusion_proto::logical_plan::to_proto::serialize_expr; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; use datafusion_proto::protobuf::LogicalExprNode; -use datafusion_session::{CatalogProviderList, Session}; +use datafusion_session::{ + CatalogProviderList, PhysicalOptimizerRule, QueryPlanner, Session, +}; use prost::Message; use stabby::str::Str as SStr; @@ -55,7 +76,10 @@ use crate::catalog_provider_list::FFI_CatalogProviderList; use crate::execution::FFI_TaskContext; use crate::execution_plan::FFI_ExecutionPlan; use crate::physical_expr::FFI_PhysicalExpr; +use crate::physical_optimizer::FFI_PhysicalOptimizerRule; use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::query_planner::FFI_QueryPlanner; use crate::session::config::FFI_SessionConfig; use crate::udaf::FFI_AggregateUDF; use crate::udf::FFI_ScalarUDF; @@ -86,6 +110,13 @@ pub(crate) struct FFI_SessionRef { catalog_list: unsafe extern "C" fn(&Self) -> FFI_CatalogProviderList, + query_planner: unsafe extern "C" fn(&Self) -> FFI_QueryPlanner, + + optimize: unsafe extern "C" fn( + &Self, + logical_plan_serialized: SVec, + ) -> FFI_Result>, + create_physical_plan: unsafe extern "C" fn( &Self, @@ -110,8 +141,12 @@ pub(crate) struct FFI_SessionRef { task_ctx: unsafe extern "C" fn(&Self) -> FFI_TaskContext, + physical_optimizers: unsafe extern "C" fn(&Self) -> SVec, + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + /// Used to create a clone on the provider of the registry. This should /// only need to be called by the receiver of the plan. clone: unsafe extern "C" fn(plan: &Self) -> Self, @@ -135,12 +170,12 @@ unsafe impl Send for FFI_SessionRef {} unsafe impl Sync for FFI_SessionRef {} struct SessionPrivateData<'a> { - session: &'a (dyn Session + Send + Sync), + session: &'a dyn Session, runtime: Option, } impl FFI_SessionRef { - fn inner(&self) -> &(dyn Session + Send + Sync) { + fn inner(&self) -> &dyn Session { let private_data = self.private_data as *const SessionPrivateData; unsafe { (*private_data).session } } @@ -173,6 +208,36 @@ unsafe extern "C" fn catalog_list_fn_wrapper( ) } +unsafe extern "C" fn query_planner_fn_wrapper( + session: &FFI_SessionRef, +) -> FFI_QueryPlanner { + FFI_QueryPlanner::new_with_ffi_codecs( + session.inner().query_planner(), + session.logical_codec.clone(), + session.physical_codec.clone(), + ) +} + +unsafe extern "C" fn optimize_fn_wrapper( + session: &FFI_SessionRef, + logical_plan_serialized: SVec, +) -> 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(SVec::from(optimized_plan.as_ref())) +} + unsafe extern "C" fn create_physical_plan_fn_wrapper( session: &FFI_SessionRef, logical_plan_serialized: SVec, @@ -303,6 +368,18 @@ unsafe extern "C" fn task_ctx_fn_wrapper(session: &FFI_SessionRef) -> FFI_TaskCo session.inner().task_ctx().into() } +unsafe extern "C" fn physical_optimizers_fn_wrapper( + session: &FFI_SessionRef, +) -> SVec { + let runtime = unsafe { session.runtime().clone() }; + session + .inner() + .physical_optimizers() + .iter() + .map(|rule| FFI_PhysicalOptimizerRule::new(Arc::clone(rule), runtime.clone())) + .collect() +} + unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_SessionRef) { unsafe { let private_data = @@ -324,6 +401,8 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_SessionRef) -> FFI_SessionR session_id: session_id_fn_wrapper, config: config_fn_wrapper, catalog_list: catalog_list_fn_wrapper, + query_planner: query_planner_fn_wrapper, + optimize: optimize_fn_wrapper, create_physical_plan: create_physical_plan_fn_wrapper, create_physical_expr: create_physical_expr_fn_wrapper, scalar_functions: scalar_functions_fn_wrapper, @@ -332,7 +411,9 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_SessionRef) -> FFI_SessionR table_options: table_options_fn_wrapper, default_table_options: default_table_options_fn_wrapper, task_ctx: task_ctx_fn_wrapper, + physical_optimizers: physical_optimizers_fn_wrapper, logical_codec: provider.logical_codec.clone(), + physical_codec: provider.physical_codec.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -350,11 +431,50 @@ impl Drop for FFI_SessionRef { } impl FFI_SessionRef { - /// Creates a new [`FFI_SessionRef`]. + /// Creates a new [`FFI_SessionRef`] with a default physical extension codec. + /// + /// The synthesized [`DefaultPhysicalExtensionCodec`] supports built-in physical + /// nodes only. A query planner obtained through this session reference therefore + /// cannot encode or decode custom physical extension nodes. Use + /// [`Self::new_with_ffi_codecs`] with matching logical and physical codecs when + /// custom physical nodes must cross the FFI boundary. + /// + /// The physical codec wrapper requires a + /// [`FFI_TaskContextProvider`](crate::execution::FFI_TaskContextProvider), but this + /// constructor has only a session reference and a logical codec. It therefore + /// reuses the logical codec's provider. The provider may be owned by another + /// library; this is safe, but it must remain live and return the task context + /// intended for codec callbacks. The default physical codec does not successfully + /// decode extension nodes, so callers that need such callbacks must instead use + /// [`Self::new_with_ffi_codecs`] with an explicitly configured physical codec and + /// task context provider. pub fn new( - session: &(dyn Session + Send + Sync), + session: &dyn Session, runtime: Option, 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, ) -> Self { if let Some(session) = session.as_any().downcast_ref::() { return session.session.clone(); @@ -366,6 +486,8 @@ impl FFI_SessionRef { session_id: session_id_fn_wrapper, config: config_fn_wrapper, catalog_list: catalog_list_fn_wrapper, + query_planner: query_planner_fn_wrapper, + optimize: optimize_fn_wrapper, create_physical_plan: create_physical_plan_fn_wrapper, create_physical_expr: create_physical_expr_fn_wrapper, scalar_functions: scalar_functions_fn_wrapper, @@ -374,7 +496,9 @@ impl FFI_SessionRef { table_options: table_options_fn_wrapper, default_table_options: default_table_options_fn_wrapper, task_ctx: task_ctx_fn_wrapper, + physical_optimizers: physical_optimizers_fn_wrapper, logical_codec, + physical_codec, clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -387,8 +511,17 @@ impl FFI_SessionRef { /// This wrapper struct exists on the receiver side of the FFI interface, so it has /// no guarantees about being able to access the data in `private_data`. Any functions -/// defined on this struct must only use the stable functions provided in -/// FFI_Session to interact with the foreign table provider. +/// defined on this struct must use only the stable function pointers in +/// `FFI_SessionRef` to interact with the foreign session. +/// +/// # Query planner delegation +/// +/// If the session owner installed the current foreign query planner, +/// [`Session::create_physical_plan`] dispatches back to that planner and +/// [`Session::query_planner`] returns that planner. The planner must retain and +/// invoke the session owner's previous planner instead of using either method to +/// delegate back to the session. Otherwise, repeated delegation exhausts the +/// stack. See [`crate::query_planner`] for details. #[derive(Debug)] pub struct ForeignSession { session: FFI_SessionRef, @@ -402,13 +535,15 @@ pub struct ForeignSession { table_options: TableOptions, runtime_env: Arc, props: ExecutionProps, + query_planner: OnceLock>, + physical_optimizers: OnceLock>>, } unsafe impl Send for ForeignSession {} unsafe impl Sync for ForeignSession {} impl FFI_SessionRef { - pub fn as_local(&self) -> Option<&(dyn Session + Send + Sync)> { + pub fn as_local(&self) -> Option<&dyn Session> { if (self.library_marker_id)() == crate::get_library_marker_id() { return Some(self.inner()); } @@ -462,7 +597,6 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { ) }) .collect(); - Ok(Self { session: session.clone(), config, @@ -475,6 +609,8 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { extension_types: Arc::new(MemoryExtensionTypeRegistry::default()), runtime_env: Default::default(), props: Default::default(), + query_planner: OnceLock::new(), + physical_optimizers: OnceLock::new(), }) } } @@ -573,6 +709,31 @@ impl Session for ForeignSession { Arc::clone(&self.catalog_list) } + fn query_planner(&self) -> Arc { + 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 { + 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, + SVec::from(logical_plan.as_ref()), + ))?; + logical_plan_from_bytes_with_extension_codec( + optimized_plan.as_slice(), + self.task_ctx().as_ref(), + codec.as_ref(), + ) + } + } + async fn create_physical_plan( &self, logical_plan: &LogicalPlan, @@ -613,6 +774,15 @@ impl Session for ForeignSession { } } + fn physical_optimizers(&self) -> &[Arc] { + 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> { &self.scalar_functions } @@ -672,6 +842,7 @@ impl Session for ForeignSession { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use arrow_schema::{DataType, Field, Schema}; use datafusion::catalog::MemoryCatalogProvider; @@ -683,6 +854,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(); @@ -731,14 +955,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..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/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 59bcc861d0567..357953674acde 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -43,6 +43,8 @@ use crate::execution_plan::FFI_ExecutionPlan; use crate::execution_plan::tests::EmptyExec; use crate::physical_optimizer::FFI_PhysicalOptimizerRule; use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::query_planner::FFI_QueryPlanner; use crate::table_provider::FFI_TableProvider; use crate::table_provider_factory::FFI_TableProviderFactory; use crate::tests::catalog::create_catalog_provider_list; @@ -50,11 +52,13 @@ use crate::udaf::FFI_AggregateUDF; use crate::udf::FFI_ScalarUDF; use crate::udtf::FFI_TableFunction; use crate::udwf::FFI_WindowUDF; +use crate::util::FFI_Option; mod async_provider; pub mod catalog; pub mod config; mod physical_optimizer; +mod query_planner; mod sync_provider; mod table_provider_factory; mod udf_udaf_udwf; @@ -117,6 +121,14 @@ pub struct ForeignLibraryModule { pub create_context_aware_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, + /// Construct a query planner. When `library_a_planner` is provided the + /// planner delegates to it, as library C does after library A swaps planners. + pub create_query_planner: extern "C" fn( + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + library_a_planner: FFI_Option, + ) -> FFI_QueryPlanner, + pub version: extern "C" fn() -> u64, /// Create an aggregate UDAF using first_value @@ -269,6 +281,7 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { physical_optimizer::create_physical_optimizer_rule, create_context_aware_optimizer_rule: physical_optimizer::create_context_aware_optimizer_rule, + create_query_planner: query_planner::create_query_planner, version: super::version, create_first_value_udaf: create_ffi_first_value_func, } diff --git a/datafusion/ffi/src/tests/query_planner.rs b/datafusion/ffi/src/tests/query_planner.rs new file mode 100644 index 0000000000000..90d713f8a6336 --- /dev/null +++ b/datafusion/ffi/src/tests/query_planner.rs @@ -0,0 +1,172 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::any::Any; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema}; +use async_trait::async_trait; +use datafusion_catalog::default_table_source::source_as_provider; +use datafusion_common::{Result, exec_err}; +use datafusion_expr::LogicalPlan; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::empty::EmptyExec; +use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::union::UnionExec; +use datafusion_session::{QueryPlanner, Session}; + +use crate::execution_plan::ForeignExecutionPlan; +use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner}; +use crate::session::ForeignSession; +use crate::table_provider::ForeignTableProvider; +use crate::util::FFI_Option; + +#[derive(Debug)] +struct TestQueryPlanner; + +#[async_trait] +impl QueryPlanner for TestQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + 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 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))) + } +} + +/// 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 { + 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/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/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()) diff --git a/datafusion/ffi/tests/ffi_query_planner.rs b/datafusion/ffi/tests/ffi_query_planner.rs new file mode 100644 index 0000000000000..c72b8c3e889ae --- /dev/null +++ b/datafusion/ffi/tests/ffi_query_planner.rs @@ -0,0 +1,348 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod utils; + +#[cfg(feature = "integration-tests")] +mod tests { + use std::sync::{Arc, OnceLock, Weak}; + + use arrow::datatypes::SchemaRef; + use datafusion::execution::SessionStateBuilder; + use datafusion::prelude::SessionContext; + use datafusion_catalog::TableProvider; + use datafusion_common::{ + DataFusionError, Result, TableReference, exec_err, not_impl_err, + }; + use datafusion_execution::{TaskContext, TaskContextProvider}; + use datafusion_expr::logical_plan::Extension; + use datafusion_expr::{LogicalPlan, col}; + use datafusion_ffi::execution::FFI_TaskContextProvider; + use datafusion_ffi::execution_plan::ForeignExecutionPlan; + use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; + use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; + use datafusion_ffi::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner}; + use datafusion_ffi::table_provider::ForeignTableProvider; + use datafusion_ffi::tests::{ + create_test_schema, + utils::{get_module, get_module_copy}, + }; + use datafusion_ffi::util::FFI_Option; + use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::empty::EmptyExec; + use datafusion_physical_plan::sorts::sort::SortExec; + use datafusion_physical_plan::union::UnionExec; + use datafusion_proto::logical_plan::LogicalExtensionCodec; + use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, + PhysicalProtoConverterExtension, + }; + use datafusion_session::QueryPlanner; + + #[tokio::test] + async fn test_ffi_query_planner() -> Result<(), DataFusionError> { + let module = get_module()?; + let (ctx, logical_codec) = crate::utils::ctx_and_codec(); + let task_ctx_provider = Arc::clone(&ctx) as Arc; + let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let physical_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(DefaultPhysicalExtensionCodec {}), + None, + task_ctx_provider, + ); + + let ffi_planner = (module.create_query_planner)( + logical_codec, + physical_codec, + FFI_Option::None, + ); + let planner: Arc = (&ffi_planner).into(); + + let any_ref: &dyn std::any::Any = planner.as_ref(); + assert!(any_ref.downcast_ref::().is_some()); + + let logical_plan = datafusion_expr::LogicalPlanBuilder::empty(false).build()?; + let state = ctx.state(); + let physical_plan = planner.create_physical_plan(&logical_plan, &state).await?; + + assert_eq!(physical_plan.name(), "EmptyExec"); + assert!(physical_plan.is::()); + + Ok(()) + } + + /// Test-only codec that preserves library B's table provider while the logical + /// plan crosses between library A and library C. + /// + /// Encoding writes a fixed identifier and stores a weak reference to the + /// provider. Decoding validates the identifier and upgrades that reference. + /// This works because all three test libraries run in one process and library + /// A's session continues to own the provider. + /// + /// This is not a general serialization format for table providers. A + /// cross-process deployment must provide its own codec that either resolves a + /// stable identifier through shared state or reconstructs the provider from a + /// portable, provider-specific description. DataFusion passes the table + /// reference, schema, and task context separately to the decoder. + #[derive(Debug, Default)] + struct LibraryALogicalCodec { + library_b_provider: OnceLock>, + } + + 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. + /// + /// Encoding sees B's node in one of two shapes. When A serializes a plan it + /// built itself, B's scan is a [`ForeignExecutionPlan`]. When library C + /// serializes a plan containing a node A previously handed it, the FFI handle + /// unwraps back to its home library, so A is asked to encode the very + /// [`EmptyExec`] its own `try_decode` produced. + #[derive(Debug)] + struct LibraryAPhysicalCodec; + + impl PhysicalExtensionCodec for LibraryAPhysicalCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + _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::() && !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(()) + } + } + + #[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_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()); + + 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(()) + } + + /// 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(()) + } +} diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 53cba29f9abb6..b98b5d2aa52a1 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -895,15 +895,28 @@ fn physical_optimizers(&self) -> &[Arc] } ``` -`ForeignSession::create_physical_plan` continues to run the complete planning -pipeline in the library that owns the session. `ForeignSession::query_planner` -returns `UnsupportedQueryPlanner` until the query planner FFI interface is -available. FFI wrappers for the individual planner and optimizer interfaces are -not included in this release. +`ForeignSession::create_physical_plan` runs the complete planning pipeline in the +library that owns the session. `ForeignSession::query_planner`, `optimize`, and +`physical_optimizers` forward to the owning session across the FFI boundary. A +foreign query planner can also be installed on a session through the new +`datafusion_ffi::query_planner::FFI_QueryPlanner`; see that module's +documentation for how plans and extension codecs cross the boundary. See [PR #23703](https://github.com/apache/datafusion/pull/23703) for details on the catalog changes. +### `FFI_LogicalExtensionCodec::task_ctx_provider` is now private + +The `task_ctx_provider` field on +`datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec` was +`pub` and is now crate-private, matching `FFI_PhysicalExtensionCodec`. + +**Who is affected:** + +- Code that read or cloned `FFI_LogicalExtensionCodec::task_ctx_provider` + directly. Pass the task context provider to `FFI_LogicalExtensionCodec::new` + instead, and keep your own copy if you need it elsewhere. + ### Unused `async` removed from several public functions Public functions that were declared `async` but never awaited anything are now