From d58dcb752c8c5a18765ca39ca0595d415568375b Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sun, 6 Sep 2026 23:40:37 -0700 Subject: [PATCH] fix: preserve ProjectionExec metadata during serialization --- datafusion/physical-plan/src/projection.rs | 20 ++- .../proto-models/proto/datafusion.proto | 4 + .../proto-models/src/generated/pbjson.rs | 17 ++ .../proto-models/src/generated/prost.rs | 5 + datafusion/proto/tests/cases/plans/exprs.rs | 167 ++++++++++++++++++ 5 files changed, 210 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 7851e5705934e..221fd9753dba2 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -673,13 +673,17 @@ 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( @@ -687,6 +691,7 @@ impl ExecutionPlan for ProjectionExec { 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::>>()?; - 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)) } } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index e4607a9b28b55..76a869312ae62 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/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 { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 0df2a4bd077fa..21dd89955a994 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/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,6 +23557,7 @@ impl<'de> serde::Deserialize<'de> for ProjectionExecNode { "expr", "expr_name", "exprName", + "schema", ]; #[allow(clippy::enum_variant_names)] @@ -23558,6 +23565,7 @@ impl<'de> serde::Deserialize<'de> for ProjectionExecNode { Input, Expr, ExprName, + Schema, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -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__, }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 1af60e27a00d8..646f45dfae820 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2200,6 +2200,11 @@ pub struct ProjectionExecNode { pub expr: ::prost::alloc::vec::Vec, #[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, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct PartiallySortedInputOrderMode { diff --git a/datafusion/proto/tests/cases/plans/exprs.rs b/datafusion/proto/tests/cases/plans/exprs.rs index f2b14b043959f..9981260eb32b9 100644 --- a/datafusion/proto/tests/cases/plans/exprs.rs +++ b/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![