From 69f9e1be671e6cf9804eb84aa81ac5e680b8f88f Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Sun, 6 Sep 2026 16:11:44 +0000 Subject: [PATCH] fix: preserve projection metadata during optimization (#24670) - Closes https://github.com/apache/datafusion/issues/24721 There were four ways metadata could disappear. Consider: ```text ProjectionExec: i@0 AS i output field metadata = {"event_field": "true"} DataSourceExec: i field metadata = {} ``` The check to remove the projection asked: - Is every expression a column? - Does column 0 remain column 0? - Does the alias match the column name? - Does the number of columns match? All answers yes so optimizer removed projection: ```text DataSourceExec: i metadata = {} ``` Metadata lost. Consider: ```text ProjectionExec: arrow_metadata(i, 'event_field') AS metadata ProjectionExec: i@0 AS i output metadata = {"event_field": "true"} DataSourceExec: i metadata = {} ``` The correct result is `true`. The previous projection colapse logic would substitute the outer expression through the inner projection: ```text ProjectionExec: arrow_metadata(i, 'event_field') AS metadata DataSourceExec: i metadata = {} ``` Now the func sees the scan field instead of the inner projection field giving use result as `NULL` now. ### 3 Rebuilding a projection with a new child Some optimizer paths replace the child of a projection: ```text Old: ProjectionExec(metadata={"key": "value"}) OldChild --- New: ProjectionExec(...?) NewChild ``` The previous `make_with_child` implementation did this: ```rust ProjectionExec::try_new(projection.expr().to_vec(), new_child) ``` where try_new derives the output schema from the expressions and new child so it woudlnt retain metadata from the original projection. ### 4 Cast target metadata lost before optimization This one was a little confusing because main passed the UUID metadata test, but the first version of this PR did not (@gabotechs this is what you called out) Basically a cast can have an explicit target field with metadata. For example, the UUID type planner produces: ```text FixedSizeBinary(16) metadata = {"ARROW:extension:name": "arrow.uuid"} ``` But logical cast schema only used the target data type when deriving and kept th source metadata: ```text Source field: raw: FixedSizeBinary(16) metadata: {} Cast target: FixedSizeBinary(16) metadata: {"ARROW:extension:name": "arrow.uuid"} Derived cast output: FixedSizeBinary(16) metadata: {} ``` This appeared in CI when common sub-expr elimination extracts a repeated cast into its own projection: ```text ProjectionExec: arrow_metadata(__common_expr_1, 'ARROW:extension:name') ProjectionExec: CAST(raw AS UUID) AS __common_expr_1 ``` The inner projection was initially created with incorrect empty metadata, so the physical optimizer rebuilt that projection and rederived its schema so isthe was accidentally repairing the logical schema bug. Once this PR started preserving projection metadata correctly had this pop up this other bug. --- So then I solve the optimizer bugs in this PR --------- Co-authored-by: Tim Saucer Co-authored-by: Claude Opus 5 (1M context) (cherry picked from commit 50cab1410cbf1a440a8c84f35159f57551ee0031) --- datafusion/physical-plan/src/projection.rs | 333 ++++++++++++++++++++- 1 file changed, 319 insertions(+), 14 deletions(-) diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index cf362cdee55d3..1e672bb8b98e7 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -82,6 +82,9 @@ pub struct ProjectionExec { metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. cache: Arc, + /// Whether the output metadata differs from the metadata derived from the + /// projection expressions and input schema. + overrides_metadata: bool, } impl ProjectionExec { @@ -143,7 +146,7 @@ impl ProjectionExec { let expr_arc = expr.into_iter().map(Into::into).collect::>(); let projection = ProjectionExprs::from_expressions(expr_arc); let projector = projection.make_projector(&input_schema)?; - Self::try_from_projector(projector, input) + Self::try_from_projector(projector, input, false) } /// Create a projection using field and schema metadata from @@ -171,12 +174,15 @@ impl ProjectionExec { let projection = ProjectionExprs::from_expressions(expr_arc); let projector = projection .make_projector_with_schema_metadata(&input_schema, projected_schema)?; - Self::try_from_projector(projector, input) + let overrides_metadata = + Self::compute_overrides_metadata(&projector, &input_schema)?; + Self::try_from_projector(projector, input, overrides_metadata) } fn try_from_projector( projector: Projector, input: Arc, + overrides_metadata: bool, ) -> Result { // Construct a map from the input expressions to the output expression of the Projection let projection_mapping = @@ -191,6 +197,7 @@ impl ProjectionExec { input, metrics: ExecutionPlanMetricsSet::new(), cache: Arc::new(cache), + overrides_metadata, }) } @@ -231,6 +238,33 @@ impl ProjectionExec { )) } + /// Returns whether `projector`'s output metadata differs from the metadata + /// derived from its expressions and `input_schema`. + fn compute_overrides_metadata( + projector: &Projector, + input_schema: &Schema, + ) -> Result { + let output_schema = projector.output_schema(); + if input_schema.metadata() != output_schema.metadata() { + return Ok(true); + } + for (projection, output_field) in + projector.projection().iter().zip(output_schema.fields()) + { + let derived_field = projection.expr.return_field(input_schema)?; + if derived_field.metadata() != output_field.metadata() { + return Ok(true); + } + } + Ok(false) + } + + /// Returns whether this projection's output metadata differs from the + /// metadata derived when the projection was constructed. + fn overrides_metadata(&self) -> bool { + self.overrides_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( @@ -352,11 +386,16 @@ impl ExecutionPlan for ProjectionExec { metrics: ExecutionPlanMetricsSet::new(), ..Self::clone(&*self) })), - ChildrenPropertiesMode::Recompute => ProjectionExec::try_from_projector( - self.projector.clone(), - children.swap_remove(0), - ) - .map(|p| Arc::new(p) as _), + ChildrenPropertiesMode::Recompute => { + let input = children.swap_remove(0); + let projector = self.projector.clone(); + let overrides_metadata = ProjectionExec::compute_overrides_metadata( + &projector, + input.schema().as_ref(), + )?; + ProjectionExec::try_from_projector(projector, input, overrides_metadata) + .map(|p| Arc::new(p) as _) + } } } @@ -946,7 +985,12 @@ pub fn remove_unnecessary_projections( if is_projection_removable(projection) { return Ok(Transformed::yes(Arc::clone(projection.input()))); } - // If it does, check if we can push it under its child(ren): + // Swapping a projection with observable metadata can change query results + // by changing the metadata visible to its child expressions. + if projection.overrides_metadata() { + return Ok(Transformed::no(plan)); + } + // Otherwise, check if we can push it under its child(ren): projection .input() .try_swapping_with_projection(projection)? @@ -958,6 +1002,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 { @@ -968,6 +1013,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 @@ -1001,13 +1047,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. @@ -1255,6 +1305,11 @@ pub fn update_join_filter( /// Collapse a chain of consecutive [`ProjectionExec`]s into one. Returns /// `None` if nothing could be merged. +/// +/// The projection-removal optimizer checks `outer.overrides_metadata()` before +/// reaching this helper. The unified projection also keeps `outer`'s schema, so +/// collapsing cannot lose its output metadata. Inner projections still need the +/// check below because outer expressions may observe their metadata. fn try_collapse_projection_chain( outer: &ProjectionExec, ) -> Result>> { @@ -1264,6 +1319,10 @@ fn try_collapse_projection_chain( 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 { @@ -1313,8 +1372,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) } @@ -1436,19 +1500,23 @@ mod tests { use crate::common::collect; use crate::empty::EmptyExec; + use crate::filter::FilterExec; use crate::filter_pushdown::PushedDown; use crate::statistics::{StatisticsArgs, StatisticsContext}; 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, + BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, is_null, lit, }; #[test] @@ -1491,6 +1559,243 @@ mod tests { Ok(()) } + fn identity_projection_with_metadata( + input: Arc, + field_metadata: HashMap, + schema_metadata: HashMap, + ) -> Result> { + let metadata_schema = Schema::new_with_metadata( + vec![Field::new("i", DataType::Int32, true).with_metadata(field_metadata)], + schema_metadata, + ); + 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_field_metadata_projection_is_not_removable() -> Result<()> { + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::new(), + )?; + 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_schema_metadata_projection_is_not_removable() -> Result<()> { + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::new(), + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + )?; + 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_replace_children_recomputes_metadata_override() -> Result<()> { + let field_metadata = + HashMap::from([("event_field".to_string(), "true".to_string())]); + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + field_metadata.clone(), + HashMap::new(), + )?; + assert!( + projection + .downcast_ref::() + .expect("test plan should be a ProjectionExec") + .overrides_metadata() + ); + + let replacement_schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, true).with_metadata(field_metadata), + ])); + let replacement: Arc = + Arc::new(EmptyExec::new(replacement_schema)); + let replaced = projection.replace_children( + vec![replacement], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + assert!( + !replaced + .downcast_ref::() + .expect("replaced plan should be a ProjectionExec") + .overrides_metadata() + ); + Ok(()) + } + + #[test] + fn test_make_with_child_preserves_output_metadata() -> Result<()> { + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + )?; + 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), + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::new(), + )?; + 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(()) + } + + #[tokio::test] + async fn test_metadata_observing_filter_blocks_projection_pushdown() -> Result<()> { + let widened: Arc = Arc::new(ProjectionExec::try_new( + [ + ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "i".to_string(), + }, + ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "j".to_string(), + }, + ], + test::scan_partitioned(1), + )?); + let arrow_metadata = Arc::new(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 filter: Arc = + Arc::new(FilterExec::try_new(is_null(arrow_metadata)?, widened)?); + let projection = identity_projection_with_metadata( + filter, + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::new(), + )?; + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + assert_eq!(optimized.schema(), expected_schema); + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 100 + ); + Ok(()) + } + + // A schema-only metadata override must block projection embedding. The filter + // rebuilds the schema from expressions and would otherwise drop this metadata. + #[tokio::test] + async fn test_schema_level_metadata_blocks_projection_embedding() -> Result<()> { + let scan = test::scan_partitioned(1); + let predicate = binary( + col("i", &scan.schema())?, + Operator::Gt, + lit(ScalarValue::Int32(Some(-1))), + &scan.schema(), + )?; + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, scan)?); + let projection = identity_projection_with_metadata( + filter, + HashMap::new(), + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + )?; + // Field metadata matches, so this checks the schema-level comparison. + let projection_exec = projection + .downcast_ref::() + .expect("test plan should be a ProjectionExec"); + assert!(projection_exec.overrides_metadata()); + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + + assert_eq!(optimized.schema(), expected_schema); + assert_eq!( + optimized.schema().metadata(), + &HashMap::from([("schema-key".to_string(), "schema-value".to_string())]) + ); + + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 100 + ); + Ok(()) + } + #[test] fn test_collect_column_indices() -> Result<()> { let expr = Arc::new(BinaryExpr::new(