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(()) + } }