diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 7851e5705934e..4c9df522c989f 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -673,13 +673,26 @@ 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 output_schema = projector.output_schema(); + // Keep inherited metadata self-contained, and retain empty overrides + // that explicitly clear metadata from the input. + let schema = if *overrides_metadata + || !output_schema.metadata().is_empty() + || output_schema + .fields() + .iter() + .any(|field| !field.metadata().is_empty()) + { + Some(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 +700,7 @@ impl ExecutionPlan for ProjectionExec { input: Some(Box::new(input)), expr, expr_name, + schema, }, )), ), @@ -722,6 +736,7 @@ impl ProjectionExec { input, expr, expr_name, + schema, } = &**projection; let input = ctx.decode_required_child(input.as_deref(), "ProjectionExec", "input")?; @@ -736,7 +751,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 fac5ff27191cd..07d760078596c 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1476,6 +1476,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. + // Present when the projection has field or schema metadata, or explicitly clears + // input metadata. Absent for older plans and metadata-free projections. + 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 e06f5b7011504..7197bc1f66d3e 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -24074,6 +24074,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)?; @@ -24084,6 +24087,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() } } @@ -24098,6 +24104,7 @@ impl<'de> serde::Deserialize<'de> for ProjectionExecNode { "expr", "expr_name", "exprName", + "schema", ]; #[allow(clippy::enum_variant_names)] @@ -24105,6 +24112,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 @@ -24129,6 +24137,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)), } } @@ -24151,6 +24160,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 => { @@ -24171,12 +24181,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 149839dacc967..ae15f1de3e446 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2231,6 +2231,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. + /// Present when the projection has field or schema metadata, or explicitly clears + /// input metadata. Absent for older plans and metadata-free projections. + #[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..456b6522fddb5 100644 --- a/datafusion/proto/tests/cases/plans/exprs.rs +++ b/datafusion/proto/tests/cases/plans/exprs.rs @@ -43,9 +43,275 @@ 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()); + for node in projection_roundtrip_nodes(&node) { + 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_schema() -> 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 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") + }; + // Simulate a plan encoded before the schema field existed. + projection.schema = None; + for node in projection_roundtrip_nodes(&node) { + let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?; + assert_eq!(decoded.schema(), input_schema); + } + Ok(()) +} + +fn projection_roundtrip_nodes(node: &PhysicalPlanNode) -> Vec { + vec![ + PhysicalPlanNode::decode(node.encode_to_vec().as_slice()).unwrap(), + #[cfg(feature = "json")] + serde_json::from_str(&serde_json::to_string(node).unwrap()).unwrap(), + ] +} + +#[test] +fn roundtrip_projection_metadata_without_child_metadata() -> Result<()> { + 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 extension_metadata = HashMap::from([ + ("ARROW:extension:name".to_string(), "arrow.uuid".to_string()), + ("ARROW:extension:metadata".to_string(), String::new()), + ]); + for (data_type, field_metadata, schema_metadata) in [ + ( + DataType::Int32, + field_metadata.clone(), + schema_metadata.clone(), + ), + (DataType::Int32, field_metadata, HashMap::new()), + (DataType::Int32, HashMap::new(), schema_metadata), + ( + DataType::FixedSizeBinary(16), + extension_metadata, + HashMap::new(), + ), + ] { + let input_field = Field::new("value", data_type, false); + let input_without_metadata = Arc::new(Schema::new(vec![input_field.clone()])); + let input_schema = Arc::new(Schema::new_with_metadata( + vec![input_field.with_metadata(field_metadata)], + schema_metadata, + )); + // This constructor derives metadata from the input, without an override. + let plan = Arc::new(ProjectionExec::try_new( + vec![(col("value", &input_schema)?, "output".to_string())], + Arc::new(EmptyExec::new(input_schema)), + )?); + let expected_schema = plan.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") + }; + // The child cannot supply metadata when the projection is reconstructed. + projection.input = Some(Box::new(PhysicalPlanNode::try_from_physical_plan( + Arc::new(EmptyExec::new(input_without_metadata)), + &codec, + )?)); + for node in projection_roundtrip_nodes(&node) { + let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?; + assert_eq!(decoded.schema(), expected_schema); + assert!(decoded.children()[0].schema().metadata().is_empty()); + assert!( + decoded.children()[0] + .schema() + .field(0) + .metadata() + .is_empty() + ); + } + } + Ok(()) +} + +#[test] +fn roundtrip_projection_without_metadata() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + 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") + }; + assert!(projection.schema.is_none()); + for node in projection_roundtrip_nodes(&node) { + 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 expected_schema = Schema::new_with_metadata( + vec![Field::new("output", DataType::Int32, false).with_metadata(field_metadata)], + schema_metadata, + ); + for node in projection_roundtrip_nodes(&node) { + let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?; + assert_eq!(decoded.schema().as_ref(), &expected_schema); + } + for field_count in [0, 2] { + let Some(protobuf::physical_plan_node::PhysicalPlanType::Projection(projection)) = + node.physical_plan_type.as_mut() + else { + unreachable!("expected ProjectionExecNode") + }; + let metadata_schema = + Schema::new(vec![ + Field::new("ignored", DataType::Utf8, true); + field_count + ]); + projection.schema = Some((&metadata_schema).try_into()?); + for node in projection_roundtrip_nodes(&node) { + let error = node + .try_into_physical_plan(&ctx.task_ctx(), &codec) + .unwrap_err(); + assert!(error.strip_backtrace().contains(&format!( + "Projection has 1 output fields but metadata schema has {field_count} fields" + ))); + } + } + Ok(()) +} + #[test] fn roundtrip_date_time_interval() -> Result<()> { let schema = Schema::new(vec![