From d22a37dfa3aef4f3c652d7baae1e2194e57c120e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 26 Aug 2026 07:52:26 -0400 Subject: [PATCH] fix(ffi): use the session's LogicalExtensionCodec in create_physical_plan `Session::create_physical_plan` was the one method on `FFI_SessionRef` that never consulted the session's `logical_codec`. Both sides of the boundary serialized the logical plan with the codec-less helpers: `ForeignSession::create_physical_plan` called `logical_plan_to_bytes` and `create_physical_plan_fn_wrapper` called `logical_plan_from_bytes`. As a result, a foreign library holding an `FFI_SessionRef` could not call `create_physical_plan` for any plan referencing a custom table provider or other extension node. Serialization failed with "This feature is not implemented: LogicalExtensionCodec is not provided", even though the neighboring `optimize` path already used the codec-aware helpers. Both sides now use `logical_plan_to_bytes_with_extension_codec` and `logical_plan_from_bytes_with_extension_codec`. In the wrapper the codec is extracted before `session.inner()` shadows the binding. Co-Authored-By: Claude Opus 5 (1M context) --- datafusion/ffi/src/session/mod.rs | 57 ++++++++++++++++++++++++++----- 1 file changed, 49 insertions(+), 8 deletions(-) diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 83f842508ab2c..f215a6ba5a568 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -53,8 +53,8 @@ use datafusion_expr::{ use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; use datafusion_proto::bytes::{ - logical_plan_from_bytes, logical_plan_from_bytes_with_extension_codec, - logical_plan_to_bytes, logical_plan_to_bytes_with_extension_codec, + logical_plan_from_bytes_with_extension_codec, + logical_plan_to_bytes_with_extension_codec, }; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::logical_plan::from_proto::parse_expr; @@ -246,13 +246,17 @@ unsafe extern "C" fn create_physical_plan_fn_wrapper( let runtime = session.runtime().clone(); let session = session.clone(); async move { + let logical_codec: Arc = + (&session.logical_codec).into(); 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; @@ -746,7 +750,10 @@ impl Session for ForeignSession { logical_plan: &LogicalPlan, ) -> datafusion_common::Result> { unsafe { - let logical_plan = logical_plan_to_bytes(logical_plan)?; + let codec: Arc = + (&self.session.logical_codec).into(); + 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, @@ -851,8 +858,9 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; + use arrow::array::record_batch; use arrow_schema::{DataType, Field, Schema}; - use datafusion::catalog::MemoryCatalogProvider; + use datafusion::catalog::{MemTable, MemoryCatalogProvider}; use datafusion::execution::SessionStateBuilder; use datafusion_common::DataFusionError; use datafusion_expr::col; @@ -860,6 +868,7 @@ mod tests { use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; 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); @@ -1004,4 +1013,36 @@ mod tests { Ok(()) } + + /// `create_physical_plan` must serialize with the session's logical codec on + /// both sides of the boundary. A plan that scans a custom table provider is + /// unserializable without it. + #[tokio::test] + async fn test_create_physical_plan_uses_logical_codec() -> Result<(), DataFusionError> + { + let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch = record_batch!(("a", Int32, [1, 2, 3]))?; + let table = MemTable::try_new(schema, vec![vec![batch]])?; + ctx.register_table("test_table", Arc::new(table))?; + + let logical_codec = FFI_LogicalExtensionCodec::new( + Arc::new(TestExtensionCodec), + None, + task_ctx_provider, + ); + + let state = ctx.state(); + let local_session = FFI_SessionRef::new(&state, None, logical_codec); + let foreign_session = ForeignSession::try_from(&local_session)?; + + let logical_plan = ctx.table("test_table").await?.into_optimized_plan()?; + let physical_plan = foreign_session.create_physical_plan(&logical_plan).await?; + + assert_eq!(physical_plan.name(), "DataSourceExec"); + assert_eq!(physical_plan.schema().field(0).name(), "a"); + + Ok(()) + } }