fix(ffi): use the session's LogicalExtensionCodec in create_physical_plan

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<dyn LogicalExtensionCodec> =
+                (&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<Arc<dyn ExecutionPlan>> {
         unsafe {
-            let logical_plan = logical_plan_to_bytes(logical_plan)?;
+            let codec: Arc<dyn LogicalExtensionCodec> =
+                (&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,15 +858,17 @@ 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;
     use datafusion_expr::registry::FunctionRegistry;
     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(())
+    }
 }