From e4b20faba507962557acb7be147ab45451d33553 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:21:01 -0500 Subject: [PATCH 1/2] fix: preserve projection field metadata during physical optimization Field metadata on a ProjectionExec's output schema could silently disappear when the physical optimizer removed or rewrote projections: 1. A metadata-only identity projection was treated as removable, because the check only compared column indices, aliases, and counts. 2. Collapsing a projection across a metadata boundary substituted the outer expression through the inner projection, so metadata-reading expressions saw the scan field instead of the projected field. 3. `make_with_child` rebuilt the projection with `try_new`, rederiving the output schema and dropping the original metadata. This commit is taken verbatim from @gene-bordegaray's work in https://github.com/apache/datafusion/pull/24670. Co-Authored-By: Gene Bordegaray --- datafusion/physical-plan/src/projection.rs | 137 ++++++++++++++++++++- 1 file changed, 132 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 2a24eb60e6fbc..4b7236a9ddcc5 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -270,6 +270,22 @@ impl ProjectionExec { )) } + /// Returns whether this projection's output metadata differs from the + /// metadata derived from its expressions and input schema. + fn overrides_metadata(&self) -> Result { + let derived_schema = self + .projector + .projection() + .project_schema(self.input.schema().as_ref())?; + let output_schema = self.schema(); + Ok(derived_schema.metadata() != output_schema.metadata() + || derived_schema + .fields() + .iter() + .zip(output_schema.fields()) + .any(|(derived, output)| derived.metadata() != output.metadata())) + } + /// Collect reverse alias mapping from projection expressions. /// The result hash map is a map from aliased Column in parent to original expr. fn collect_reverse_alias( @@ -1014,6 +1030,10 @@ pub fn remove_unnecessary_projections( plan: Arc, ) -> Result>> { let maybe_modified = if let Some(projection) = plan.downcast_ref::() { + // Removing a projection with observable metadata can change query results. + if projection.overrides_metadata()? { + return Ok(Transformed::no(plan)); + } // If the projection does not cause any change on the input, we can // safely remove it: if is_projection_removable(projection) { @@ -1031,6 +1051,7 @@ pub fn remove_unnecessary_projections( /// Compare the inputs and outputs of the projection. All expressions must be /// columns without alias, and projection does not change the order of fields. +/// The input and output schemas must also match exactly to preserve metadata. /// For example, if the input schema is `a, b`, `SELECT a, b` is removable, /// but `SELECT b, a` and `SELECT a+1, b` and `SELECT a AS c, b` are not. fn is_projection_removable(projection: &ProjectionExec) -> bool { @@ -1041,6 +1062,7 @@ fn is_projection_removable(projection: &ProjectionExec) -> bool { }; col.name() == proj_expr.alias && col.index() == idx }) && exprs.len() == projection.input().schema().fields().len() + && projection.schema() == projection.input().schema() } /// Given the expression set of a projection, checks if the projection causes @@ -1074,13 +1096,17 @@ pub fn new_projections_for_columns( } /// Creates a new [`ProjectionExec`] instance with the given child plan and -/// projected expressions. +/// projected expressions, preserving the original output metadata. pub fn make_with_child( projection: &ProjectionExec, child: &Arc, ) -> Result> { - ProjectionExec::try_new(projection.expr().to_vec(), Arc::clone(child)) - .map(|e| Arc::new(e) as _) + ProjectionExec::try_new_with_schema_metadata( + projection.expr().to_vec(), + Arc::clone(child), + projection.schema().as_ref(), + ) + .map(|e| Arc::new(e) as _) } /// Returns `true` if all the expressions in the argument are `Column`s. @@ -1331,12 +1357,20 @@ pub fn update_join_filter( fn try_collapse_projection_chain( outer: &ProjectionExec, ) -> Result>> { + if outer.overrides_metadata()? { + return Ok(None); + } + let mut current_exprs: Vec = outer.expr().to_vec(); let mut current_input: Arc = Arc::clone(outer.input()); let mut column_ref_map: HashMap = HashMap::new(); let mut collapsed_any = false; 'outer: while let Some(inner_proj) = current_input.downcast_ref::() { + if inner_proj.overrides_metadata()? { + break; + } + // Collect the column references usage in the outer projection. column_ref_map.clear(); for proj_expr in ¤t_exprs { @@ -1386,8 +1420,13 @@ fn try_collapse_projection_chain( } // To unify 3 or more sequential projections: + // Preserve the outer projection's output metadata. let unified: Arc = - Arc::new(ProjectionExec::try_new(current_exprs, current_input)?); + Arc::new(ProjectionExec::try_new_with_schema_metadata( + current_exprs, + current_input, + outer.schema().as_ref(), + )?); remove_unnecessary_projections(unified).data().map(Some) } @@ -1517,11 +1556,14 @@ mod tests { use crate::test; use crate::test::exec::StatisticsExec; + use arrow::array::StringArray; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::ScalarValue; use datafusion_common::stats::{ColumnStatistics, Precision, Statistics}; - use datafusion_expr::Operator; + use datafusion_expr::{Operator, ScalarUDF}; + use datafusion_functions::core::arrow_metadata::ArrowMetadataFunc; + use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::expressions::{ BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, lit, }; @@ -1566,6 +1608,91 @@ mod tests { Ok(()) } + fn identity_projection_with_metadata( + input: Arc, + ) -> Result> { + let metadata_schema = + Schema::new_with_metadata( + vec![Field::new("i", DataType::Int32, true).with_metadata( + HashMap::from([("event_field".to_string(), "true".to_string())]), + )], + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + ); + Ok(Arc::new(ProjectionExec::try_new_with_schema_metadata( + [ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "i".to_string(), + }], + input, + &metadata_schema, + )?)) + } + + #[test] + fn test_metadata_projection_is_not_removable() -> Result<()> { + let projection = identity_projection_with_metadata(test::scan_partitioned(1))?; + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + + assert!(optimized.downcast_ref::().is_some()); + assert_eq!(optimized.schema(), expected_schema); + Ok(()) + } + + #[test] + fn test_make_with_child_preserves_output_metadata() -> Result<()> { + let projection = identity_projection_with_metadata(test::scan_partitioned(1))?; + let projection = projection + .downcast_ref::() + .expect("test plan should be a ProjectionExec"); + + let rebuilt = make_with_child(projection, &test::scan_partitioned(1))?; + + assert_eq!(rebuilt.schema(), projection.schema()); + Ok(()) + } + + #[tokio::test] + async fn test_metadata_observing_parent_blocks_projection_collapse() -> Result<()> { + let inner = identity_projection_with_metadata(test::scan_partitioned(1))?; + let arrow_metadata = ScalarFunctionExpr::new( + "arrow_metadata", + Arc::new(ScalarUDF::new_from_impl(ArrowMetadataFunc::new())), + vec![ + Arc::new(Column::new("i", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "event_field".to_string(), + )))), + ], + Arc::new(Field::new("arrow_metadata", DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + ); + let outer: Arc = Arc::new(ProjectionExec::try_new( + [ProjectionExpr { + expr: Arc::new(arrow_metadata), + alias: "metadata".to_string(), + }], + inner, + )?); + + let outer_projection = outer + .downcast_ref::() + .expect("test plan should be a ProjectionExec"); + assert!(try_collapse_projection_chain(outer_projection)?.is_none()); + + let optimized = remove_unnecessary_projections(outer)?.data; + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + let values = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("metadata expression should return Utf8"); + assert_eq!(values.value(0), "true"); + Ok(()) + } + #[test] fn test_collect_column_indices() -> Result<()> { let expr = Arc::new(BinaryExpr::new( From 7d650aaca1e732b328d2dd184988edea464b5e39 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:22:32 -0500 Subject: [PATCH 2/2] fix: use the cast target's metadata when it carries any The logical `Expr::Cast`/`Expr::TryCast` carry a `FieldRef` target so a cast can express a destination that is more than a `DataType` (for example an extension type produced by a `TypePlanner`). `cast_output_field` ignored that field's metadata entirely and always inherited the source's, so `Expr::to_field()` disagreed with the physical `CastExpr`, which already treats a non-synthesized target field as authoritative. The divergence was masked because the physical optimizer rederives a projection's schema from its expressions, repairing the logical schema on the way through. Once projections preserve their metadata faithfully (previous commit) the underlying bug surfaces, and a cast to an extension type loses it: SELECT arrow_metadata(CAST(raw AS UUID), 'ARROW:extension:name') -- 'arrow.uuid' before, NULL after Take the target's metadata when it carries any, and otherwise inherit the source's. A plain `CAST(expr AS type)` synthesizes a target with no metadata, so its long-standing behaviour is unchanged. --- datafusion/expr/src/expr_schema.rs | 86 +++++++++++++++++++++++++----- 1 file changed, 73 insertions(+), 13 deletions(-) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 36b76f076d26a..7f695747a2b3b 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -72,17 +72,28 @@ pub trait ExprSchemable { } /// Derives the output field for a cast expression from the source field. +/// +/// The cast target's metadata is authoritative when it carries any; otherwise the +/// source's metadata is inherited. This mirrors the physical `CastExpr`, whose +/// target field is already authoritative when it is not the synthesized type-only +/// field. +/// /// For `TryCast`, `force_nullable` is `true` since a failed cast returns NULL. fn cast_output_field( source_field: &FieldRef, - target_type: &DataType, + target_field: &FieldRef, force_nullable: bool, ) -> Arc { + let metadata = if target_field.metadata().is_empty() { + source_field.metadata().clone() + } else { + target_field.metadata().clone() + }; let mut f = source_field .as_ref() .clone() - .with_data_type(target_type.clone()) - .with_metadata(source_field.metadata().clone()); + .with_data_type(target_field.data_type().clone()) + .with_metadata(metadata); if force_nullable { f = f.with_nullable(true); } @@ -623,20 +634,16 @@ impl ExprSchemable for Expr { func.return_field_from_args(args) } // _ => Ok((self.get_type(schema)?, self.nullable(schema)?)), - Expr::Cast(Cast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), false) - }) - } + Expr::Cast(Cast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, false)), Expr::Placeholder(Placeholder { id: _, field: Some(field), }) => Ok(Arc::clone(field).renamed(&schema_name)), - Expr::TryCast(TryCast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), true) - }) - } + Expr::TryCast(TryCast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, true)), Expr::LambdaVariable(LambdaVariable { field: Some(field), .. }) => Ok(Arc::clone(field).renamed(&schema_name)), @@ -1427,4 +1434,57 @@ mod tests { assert_eq!(meta, expr.metadata(&schema).unwrap()); } + + #[test] + fn test_cast_output_field_metadata() { + use crate::expr::{Cast, TryCast}; + + let source_meta = + HashMap::from([("source_key".to_string(), "source_value".to_string())]); + let schema = MockExprSchema::new() + .with_data_type(DataType::FixedSizeBinary(16)) + .with_metadata(FieldMetadata::from(source_meta.clone())); + + // A target field carrying metadata is authoritative: the source metadata + // does not leak into the output. + let target_meta = + HashMap::from([("target_key".to_string(), "target_value".to_string())]); + let target = + Arc::new(Field::new("", DataType::Utf8, true).with_metadata(target_meta)); + + for expr in [ + Expr::Cast(Cast::new_from_field( + Box::new(col("foo")), + Arc::clone(&target), + )), + Expr::TryCast(TryCast::new_from_field( + Box::new(col("foo")), + Arc::clone(&target), + )), + ] { + let field = expr.to_field(&schema).unwrap().1; + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()) + ); + assert!(field.metadata().get("source_key").is_none()); + } + + // A target field with no metadata inherits the source's, preserving the + // long-standing behaviour of a plain `CAST(expr AS type)`. + let bare = Arc::new(Field::new("", DataType::Utf8, true)); + for expr in [ + Expr::Cast(Cast::new_from_field( + Box::new(col("foo")), + Arc::clone(&bare), + )), + Expr::TryCast(TryCast::new_from_field( + Box::new(col("foo")), + Arc::clone(&bare), + )), + ] { + let field = expr.to_field(&schema).unwrap().1; + assert_eq!(field.metadata(), &source_meta); + } + } }