fix: preserve ProjectionExec metadata during serialization

datafusion/physical-plan/src/projection.rs

@@ -673,20 +673,25 @@ impl ExecutionPlan for ProjectionExec {
             metrics: _,
             // Derived plan properties, recomputed on decode.
             cache: _,
-            // Derived metadata comparison, recomputed with the projector.
-            overrides_metadata: _,
+            overrides_metadata,
         } = self;
         let projection_exprs = projector.projection().as_ref();
         let input = ctx.encode_child(input)?;
         let expr = ctx.encode_expressions(projection_exprs.iter().map(|p| &p.expr))?;
         let expr_name = projection_exprs.iter().map(|p| p.alias.clone()).collect();
+        let schema = if *overrides_metadata {
+            Some(projector.output_schema().as_ref().try_into()?)
+        } else {
+            None
+        };
         Ok(Some(protobuf::PhysicalPlanNode {
             physical_plan_type: Some(
                 protobuf::physical_plan_node::PhysicalPlanType::Projection(Box::new(
                     protobuf::ProjectionExecNode {
                         input: Some(Box::new(input)),
                         expr,
                         expr_name,
+                        schema,
                     },
                 )),
             ),
@@ -722,6 +727,7 @@ impl ProjectionExec {
             input,
             expr,
             expr_name,
+            schema,
         } = &**projection;
         let input =
             ctx.decode_required_child(input.as_deref(), "ProjectionExec", "input")?;
@@ -736,7 +742,15 @@ impl ProjectionExec {
                 })
             })
             .collect::<Result<Vec<ProjectionExpr>>>()?;
-        Ok(Arc::new(ProjectionExec::try_new(exprs, input)?))
+        let projection = match schema {
+            Some(schema) => ProjectionExec::try_new_with_schema_metadata(
+                exprs,
+                input,
+                &Schema::try_from(schema)?,
+            )?,
+            None => ProjectionExec::try_new(exprs, input)?,
+        };
+        Ok(Arc::new(projection))
     }
 }
 

datafusion/proto-models/proto/datafusion.proto

@@ -1452,6 +1452,10 @@ message ProjectionExecNode {
   PhysicalPlanNode input = 1;
   repeated PhysicalExprNode expr = 2;
   repeated string expr_name = 3;
+  // Only field and schema metadata are used; output types are derived from expr.
+  // Absent when metadata can be derived from the expressions and input, including
+  // plans encoded before this field existed.
+  datafusion_common.Schema schema = 4;
 }
 
 enum AggregateMode {

datafusion/proto-models/src/generated/pbjson.rs

@@ -23527,6 +23527,9 @@ impl serde::Serialize for ProjectionExecNode {
         if !self.expr_name.is_empty() {
             len += 1;
         }
+        if self.schema.is_some() {
+            len += 1;
+        }
         let mut struct_ser = serializer.serialize_struct("datafusion.ProjectionExecNode", len)?;
         if let Some(v) = self.input.as_ref() {
             struct_ser.serialize_field("input", v)?;
@@ -23537,6 +23540,9 @@ impl serde::Serialize for ProjectionExecNode {
         if !self.expr_name.is_empty() {
             struct_ser.serialize_field("exprName", &self.expr_name)?;
         }
+        if let Some(v) = self.schema.as_ref() {
+            struct_ser.serialize_field("schema", v)?;
+        }
         struct_ser.end()
     }
 }
@@ -23551,13 +23557,15 @@ impl<'de> serde::Deserialize<'de> for ProjectionExecNode {
             "expr",
             "expr_name",
             "exprName",
+            "schema",
         ];
 
         #[allow(clippy::enum_variant_names)]
         enum GeneratedField {
             Input,
             Expr,
             ExprName,
+            Schema,
         }
         impl<'de> serde::Deserialize<'de> for GeneratedField {
             fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error>
@@ -23582,6 +23590,7 @@ impl<'de> serde::Deserialize<'de> for ProjectionExecNode {
                             "input" => Ok(GeneratedField::Input),
                             "expr" => Ok(GeneratedField::Expr),
                             "exprName" | "expr_name" => Ok(GeneratedField::ExprName),
+                            "schema" => Ok(GeneratedField::Schema),
                             _ => Err(serde::de::Error::unknown_field(value, FIELDS)),
                         }
                     }
@@ -23604,6 +23613,7 @@ impl<'de> serde::Deserialize<'de> for ProjectionExecNode {
                 let mut input__ = None;
                 let mut expr__ = None;
                 let mut expr_name__ = None;
+                let mut schema__ = None;
                 while let Some(k) = map_.next_key()? {
                     match k {
                         GeneratedField::Input => {
@@ -23624,12 +23634,19 @@ impl<'de> serde::Deserialize<'de> for ProjectionExecNode {
                             }
                             expr_name__ = Some(map_.next_value()?);
                         }
+                        GeneratedField::Schema => {
+                            if schema__.is_some() {
+                                return Err(serde::de::Error::duplicate_field("schema"));
+                            }
+                            schema__ = map_.next_value()?;
+                        }
                     }
                 }
                 Ok(ProjectionExecNode {
                     input: input__,
                     expr: expr__.unwrap_or_default(),
                     expr_name: expr_name__.unwrap_or_default(),
+                    schema: schema__,
                 })
             }
         }

datafusion/proto-models/src/generated/prost.rs

@@ -2200,6 +2200,11 @@ pub struct ProjectionExecNode {
     pub expr: ::prost::alloc::vec::Vec<PhysicalExprNode>,
     #[prost(string, repeated, tag = "3")]
     pub expr_name: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
+    /// Only field and schema metadata are used; output types are derived from expr.
+    /// Absent when metadata can be derived from the expressions and input, including
+    /// plans encoded before this field existed.
+    #[prost(message, optional, tag = "4")]
+    pub schema: ::core::option::Option<super::datafusion_common::Schema>,
 }
 #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
 pub struct PartiallySortedInputOrderMode {

datafusion/proto/tests/cases/plans/exprs.rs

@@ -43,9 +43,176 @@ use datafusion_proto::physical_plan::{
 };
 use datafusion_proto::protobuf;
 use datafusion_proto::protobuf::PhysicalPlanNode;
+use prost::Message;
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::vec;
 
+#[test]
+fn roundtrip_projection_metadata() -> Result<()> {
+    let input_schema = Arc::new(Schema::new(vec![Field::new(
+        "value",
+        DataType::Int32,
+        false,
+    )]));
+    let projected_schema = Schema::new_with_metadata(
+        vec![Field::new("value", DataType::Int32, false).with_metadata(
+            [("field-key".to_string(), "field-value".to_string())].into(),
+        )],
+        [("schema-key".to_string(), "schema-value".to_string())].into(),
+    );
+    let plan = Arc::new(ProjectionExec::try_new_with_schema_metadata(
+        vec![(col("value", &input_schema)?, "value".to_string())],
+        Arc::new(EmptyExec::new(input_schema)),
+        &projected_schema,
+    )?);
+    let ctx = SessionContext::new();
+    let codec = DefaultPhysicalExtensionCodec {};
+    let converter = DefaultPhysicalProtoConverter {};
+    let decoded = roundtrip_test_and_return(plan, &ctx, &codec, &converter)?;
+    assert_eq!(decoded.schema().as_ref(), &projected_schema);
+    Ok(())
+}
+
+#[test]
+fn roundtrip_projection_metadata_overrides() -> Result<()> {
+    let field_metadata = [("field-key".to_string(), "field-value".to_string())].into();
+    let extension_metadata = [
+        ("ARROW:extension:name".to_string(), "arrow.uuid".to_string()),
+        ("ARROW:extension:metadata".to_string(), String::new()),
+    ]
+    .into();
+    for (input_field, output_field, input_metadata) in [
+        (
+            Field::new("value", DataType::Int32, false).with_metadata(field_metadata),
+            Field::new("value", DataType::Int32, false),
+            [("input-schema".to_string(), "input-value".to_string())].into(),
+        ),
+        (
+            Field::new("value", DataType::FixedSizeBinary(16), true),
+            Field::new("value", DataType::FixedSizeBinary(16), true)
+                .with_metadata(extension_metadata),
+            HashMap::new(),
+        ),
+    ] {
+        let input_schema =
+            Arc::new(Schema::new_with_metadata(vec![input_field], input_metadata));
+        let projected_schema = Schema::new(vec![output_field]);
+        let plan = Arc::new(ProjectionExec::try_new_with_schema_metadata(
+            vec![(col("value", &input_schema)?, "value".to_string())],
+            Arc::new(EmptyExec::new(input_schema)),
+            &projected_schema,
+        )?);
+        let codec = DefaultPhysicalExtensionCodec {};
+        let ctx = SessionContext::new();
+        let node = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?;
+        let Some(protobuf::physical_plan_node::PhysicalPlanType::Projection(projection)) =
+            node.physical_plan_type.as_ref()
+        else {
+            unreachable!("expected ProjectionExecNode")
+        };
+        assert!(projection.schema.is_some());
+        let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()).unwrap();
+        #[cfg(feature = "json")]
+        let node: PhysicalPlanNode =
+            serde_json::from_str(&serde_json::to_string(&node).unwrap()).unwrap();
+        let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?;
+        assert_eq!(decoded.schema().as_ref(), &projected_schema);
+    }
+    Ok(())
+}
+
+#[test]
+fn roundtrip_projection_without_metadata_override() -> Result<()> {
+    let input_schema = Arc::new(Schema::new_with_metadata(
+        vec![Field::new("value", DataType::Int32, false).with_metadata(
+            [("field-key".to_string(), "field-value".to_string())].into(),
+        )],
+        [("schema-key".to_string(), "schema-value".to_string())].into(),
+    ));
+    let plan = Arc::new(ProjectionExec::try_new(
+        vec![(col("value", &input_schema)?, "value".to_string())],
+        Arc::new(EmptyExec::new(Arc::clone(&input_schema))),
+    )?);
+    let codec = DefaultPhysicalExtensionCodec {};
+    let ctx = SessionContext::new();
+    let node = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?;
+    let Some(protobuf::physical_plan_node::PhysicalPlanType::Projection(projection)) =
+        node.physical_plan_type.as_ref()
+    else {
+        unreachable!("expected ProjectionExecNode")
+    };
+    // The payload also represents plans encoded before the schema field existed.
+    assert!(projection.schema.is_none());
+    let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()).unwrap();
+    #[cfg(feature = "json")]
+    let node: PhysicalPlanNode =
+        serde_json::from_str(&serde_json::to_string(&node).unwrap()).unwrap();
+    let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?;
+    assert_eq!(decoded.schema(), input_schema);
+    Ok(())
+}
+
+#[test]
+fn decode_projection_schema_only_replaces_metadata() -> Result<()> {
+    let input_schema = Arc::new(Schema::new(vec![Field::new(
+        "input",
+        DataType::Int32,
+        false,
+    )]));
+    let plan = Arc::new(ProjectionExec::try_new(
+        vec![(col("input", &input_schema)?, "output".to_string())],
+        Arc::new(EmptyExec::new(input_schema)),
+    )?);
+    let codec = DefaultPhysicalExtensionCodec {};
+    let ctx = SessionContext::new();
+    let mut node = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?;
+    let Some(protobuf::physical_plan_node::PhysicalPlanType::Projection(projection)) =
+        node.physical_plan_type.as_mut()
+    else {
+        unreachable!("expected ProjectionExecNode")
+    };
+    let field_metadata =
+        HashMap::from([("field-key".to_string(), "field-value".to_string())]);
+    let schema_metadata =
+        HashMap::from([("schema-key".to_string(), "schema-value".to_string())]);
+    let metadata_schema = Schema::new_with_metadata(
+        vec![
+            Field::new("ignored", DataType::Utf8, true)
+                .with_metadata(field_metadata.clone()),
+        ],
+        schema_metadata.clone(),
+    );
+    projection.schema = Some((&metadata_schema).try_into()?);
+    let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?;
+    assert_eq!(
+        decoded.schema().as_ref(),
+        &Schema::new_with_metadata(
+            vec![
+                Field::new("output", DataType::Int32, false)
+                    .with_metadata(field_metadata)
+            ],
+            schema_metadata,
+        ),
+    );
+
+    let Some(protobuf::physical_plan_node::PhysicalPlanType::Projection(projection)) =
+        node.physical_plan_type.as_mut()
+    else {
+        unreachable!("expected ProjectionExecNode")
+    };
+    projection.schema.as_mut().unwrap().columns.clear();
+    let error = node
+        .try_into_physical_plan(&ctx.task_ctx(), &codec)
+        .unwrap_err();
+    assert!(
+        error
+            .strip_backtrace()
+            .contains("Projection has 1 output fields but metadata schema has 0 fields")
+    );
+    Ok(())
+}
+
 #[test]
 fn roundtrip_date_time_interval() -> Result<()> {
     let schema = Schema::new(vec![