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/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/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/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/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 97aa5c901a636..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; @@ -99,7 +100,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. @@ -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, )) } @@ -295,7 +305,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 { @@ -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) @@ -712,14 +728,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/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/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..2a607c334d39c --- /dev/null +++ b/datafusion/ffi/src/query_planner.rs @@ -0,0 +1,413 @@ +// 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_session::{QueryPlanner, Session}; +use stabby::vec::Vec as SVec; +use tokio::runtime::Handle; + +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}; + +/// 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>>, + + /// 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, + + /// 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 } + } + + /// 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( + planner: &FFI_QueryPlanner, + logical_plan_serialized: SVec, + session: FFI_SessionRef, +) -> FfiFuture>> { + let internal_planner = Arc::clone(planner.inner()); + let logical_codec = planner.codecs.to_logical_codec(); + let physical_codec = planner.codecs.to_physical_codec(); + + 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, + codecs: planner.codecs.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`]. + /// + /// `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, 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, + 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.codecs = codecs; + return planner; + } + + let private_data = Box::new(QueryPlannerPrivateData { planner }); + + Self { + create_physical_plan: create_physical_plan_fn_wrapper, + codecs, + 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 = 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(session, session_runtime, self.codecs.clone()); + + let physical_plan = unsafe { + df_result!((self.create_physical_plan)(self, logical_plan, session).await)? + }; + let physical_codec = self.codecs.to_physical_codec(); + + 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; + let codecs = FFI_ExtensionCodecBundle::new( + &task_ctx_provider, + None, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), + ); + FFI_QueryPlanner::new(Arc::new(EmptyQueryPlanner), codecs) + } + + #[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/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 519384379edb8..b4ddeb1f4518f 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -15,10 +15,37 @@ // specific language governing permissions and limitations // under the License. +//! 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 +//! 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 +64,16 @@ 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::logical_plan::LogicalExtensionCodec; +use datafusion_proto::bytes::{ + logical_plan_from_bytes_with_extension_codec, + logical_plan_to_bytes_with_extension_codec, +}; use datafusion_proto::logical_plan::from_proto::parse_expr; use datafusion_proto::logical_plan::to_proto::serialize_expr; 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 +86,9 @@ 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::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::physical_optimizer::FFI_PhysicalOptimizerRule; +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; use crate::udf::FFI_ScalarUDF; @@ -86,6 +119,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,7 +150,10 @@ pub(crate) struct FFI_SessionRef { task_ctx: unsafe extern "C" fn(&Self) -> FFI_TaskContext, - logical_codec: FFI_LogicalExtensionCodec, + physical_optimizers: unsafe extern "C" fn(&Self) -> SVec, + + /// 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. @@ -135,12 +178,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 } } @@ -166,13 +209,39 @@ 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(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 = 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(), + 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, @@ -181,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; @@ -202,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(); @@ -303,6 +375,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 +408,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 +418,8 @@ 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, - logical_codec: provider.logical_codec.clone(), + physical_optimizers: physical_optimizers_fn_wrapper, + codecs: provider.codecs.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -351,10 +438,23 @@ impl Drop for FFI_SessionRef { impl FFI_SessionRef { /// Creates a new [`FFI_SessionRef`]. + /// + /// 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. + /// + /// 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 + Send + Sync), + session: &dyn Session, runtime: Option, - logical_codec: FFI_LogicalExtensionCodec, + codecs: FFI_ExtensionCodecBundle, ) -> Self { if let Some(session) = session.as_any().downcast_ref::() { return session.session.clone(); @@ -366,6 +466,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 +476,8 @@ impl FFI_SessionRef { table_options: table_options_fn_wrapper, default_table_options: default_table_options_fn_wrapper, task_ctx: task_ctx_fn_wrapper, - logical_codec, + physical_optimizers: physical_optimizers_fn_wrapper, + codecs, clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -387,8 +490,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 +514,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 +576,6 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { ) }) .collect(); - Ok(Self { session: session.clone(), config, @@ -475,6 +588,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,12 +688,38 @@ 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 = 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)( + &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, ) -> 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, @@ -598,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())?); @@ -613,6 +753,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,16 +821,94 @@ 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; 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); + + 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 codecs = FFI_ExtensionCodecBundle::new_default(task_ctx_provider, None); + let state = ctx.state(); + let mut local_session = FFI_SessionRef::new(&state, None, codecs); + 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(()) + } + + /// 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> { @@ -699,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)])); @@ -731,14 +954,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..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 { @@ -263,7 +259,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 +310,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()) @@ -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 b70e72f31aa4d..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,25 +188,25 @@ 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)?; 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()) })?; 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 59bcc861d0567..f04701e362f95 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -42,7 +42,8 @@ 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::extension_codec_bundle::FFI_ExtensionCodecBundle; +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 +51,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; @@ -67,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, @@ -93,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, @@ -111,12 +114,30 @@ 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, 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( + 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 @@ -142,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 { @@ -233,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); @@ -242,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. @@ -269,6 +290,10 @@ 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, + 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 new file mode 100644 index 0000000000000..e4d303065b978 --- /dev/null +++ b/datafusion/ffi/src/tests/query_planner.rs @@ -0,0 +1,259 @@ +// 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, 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::{Expr, LogicalPlan, LogicalPlanBuilder, TableType}; +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::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::{FFI_TableProvider, 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), + ])?) + } +} + +/// 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( + codecs: FFI_ExtensionCodecBundle, + 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(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/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..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())); @@ -155,7 +150,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()) @@ -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 new file mode 100644 index 0000000000000..03d7bdb127c0b --- /dev/null +++ b/datafusion/ffi/tests/ffi_query_planner.rs @@ -0,0 +1,440 @@ +// 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_plan::ForeignExecutionPlan; + 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::{ + 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::{ + DefaultLogicalExtensionCodec, LogicalExtensionCodec, + }; + use datafusion_proto::physical_plan::{ + PhysicalExtensionCodec, PhysicalProtoConverterExtension, + }; + use datafusion_session::QueryPlanner; + + #[tokio::test] + async fn test_ffi_query_planner() -> Result<(), DataFusionError> { + let module = get_module()?; + let (ctx, codecs) = crate::utils::ctx_and_codecs(); + + 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(); + 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 codecs = FFI_ExtensionCodecBundle::new( + &task_ctx_provider, + None, + Arc::new(LibraryALogicalCodec::default()), + Arc::new(LibraryAPhysicalCodec), + ); + + 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, codecs.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)(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()); + + 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 codecs = FFI_ExtensionCodecBundle::new( + &task_ctx_provider, + None, + Arc::new(LibraryALogicalCodec::default()), + Arc::new(LibraryAPhysicalCodec), + ); + + 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, codecs.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(library_a_planner, codecs.clone()); + + // Library C: builds its planner around A's planner. + let ffi_planner = (library_c.create_query_planner)( + codecs, + 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(()) + } + + /// 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