fix: preserve projection metadata during optimization

Author: gene-bordegaray

datafusion/expr/src/expr_schema.rs

@@ -71,18 +71,31 @@ pub trait ExprSchemable {
     -> Result<(DataType, bool)>;
 }
 
-/// Derives the output field for a cast expression from the source field.
+/// Derives the output field for a cast expression from the source and target
+/// fields. Type-only casts preserve source metadata, while an explicit target
+/// field supplies its own metadata.
+///
 /// 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<Field> {
+    // `Cast::new` and `TryCast::new` use this field when only a target
+    // type is known. In that case, retain the source field's metadata.
+    let type_only_target = target_field.name().is_empty()
+        && target_field.is_nullable()
+        && target_field.metadata().is_empty();
+    let metadata = if type_only_target {
+        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 +636,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)),
@@ -1149,7 +1158,7 @@ mod tests {
     }
 
     #[test]
-    fn test_expr_metadata() {
+    fn test_expr_metadata() -> Result<()> {
         let mut meta = HashMap::new();
         meta.insert("bar".to_string(), "buzz".to_string());
         let meta = FieldMetadata::from(meta);
@@ -1179,13 +1188,38 @@ mod tests {
         // verify to_field method populates metadata
         assert_eq!(meta, expr.metadata(&schema).unwrap());
 
+        // An explicit cast target replaces source metadata. A type-only cast
+        // continues to preserve it.
+        let target_metadata = HashMap::from([(
+            "ARROW:extension:name".to_string(),
+            "arrow.uuid".to_string(),
+        )]);
+        let target_field = Arc::new(
+            Field::new("", DataType::FixedSizeBinary(16), true)
+                .with_metadata(target_metadata.clone()),
+        );
+        let cast = Expr::Cast(Cast::new_from_field(
+            Box::new(expr.clone()),
+            Arc::clone(&target_field),
+        ));
+        assert_eq!(cast.to_field(&schema)?.1.metadata(), &target_metadata);
+
+        let try_cast = Expr::TryCast(TryCast::new_from_field(
+            Box::new(expr.clone()),
+            target_field,
+        ));
+        let try_cast_field = try_cast.to_field(&schema)?.1;
+        assert_eq!(try_cast_field.metadata(), &target_metadata);
+        assert!(try_cast_field.is_nullable());
+
         // outer ref constructed by `out_ref_col_with_metadata` should be metadata-preserving
         let outer_ref = out_ref_col_with_metadata(
             DataType::Int32,
             meta.to_hashmap(),
             Column::from_name("foo"),
         );
         assert_eq!(meta, outer_ref.metadata(&schema).unwrap());
+        Ok(())
     }
 
     #[test]

datafusion/physical-plan/src/projection.rs

@@ -1014,6 +1014,11 @@ pub fn remove_unnecessary_projections(
     plan: Arc<dyn ExecutionPlan>,
 ) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
     let maybe_modified = if let Some(projection) = plan.downcast_ref::<ProjectionExec>() {
+        // Metadata is observable by physical expressions. Moving a projection that
+        // defines metadata can therefore change query results.
+        if projection_overrides_metadata(projection)? {
+            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) {
@@ -1029,10 +1034,12 @@ pub fn remove_unnecessary_projections(
     Ok(maybe_modified.map_or_else(|| Transformed::no(plan), Transformed::yes))
 }
 
-/// Compare the inputs and outputs of the projection. All expressions must be
-/// columns without alias, and projection does not change the order of fields.
-/// 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.
+/// Compare the input and output of a projection. A removable projection contains
+/// only unaliased columns in input order and preserves the exact schema, including
+/// field and schema metadata.
+///
+/// For example, if the input schema is `a, b`, `SELECT a, b` is removable, but
+/// `SELECT b, a`, `SELECT a + 1, b`, and a metadata-only projection are not.
 fn is_projection_removable(projection: &ProjectionExec) -> bool {
     let exprs = projection.expr();
     exprs.iter().enumerate().all(|(idx, proj_expr)| {
@@ -1041,6 +1048,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
@@ -1073,14 +1081,20 @@ pub fn new_projections_for_columns(
         .collect()
 }
 
-/// Creates a new [`ProjectionExec`] instance with the given child plan and
-/// projected expressions.
+/// Creates an equivalent [`ProjectionExec`] with a new child.
+///
+/// The original output metadata is preserved because a parent expression may
+/// observe it; recomputing metadata from `child` could change query results.
 pub fn make_with_child(
     projection: &ProjectionExec,
     child: &Arc<dyn ExecutionPlan>,
 ) -> Result<Arc<dyn ExecutionPlan>> {
-    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.
@@ -1326,17 +1340,44 @@ pub fn update_join_filter(
     })
 }
 
+/// Returns whether a projection defines metadata that its expressions and input
+/// schema cannot reproduce.
+///
+/// Such a projection is an execution boundary: a parent expression such as
+/// `arrow_metadata` can observe its output field metadata.
+fn projection_overrides_metadata(projection: &ProjectionExec) -> Result<bool> {
+    let derived_schema = projection
+        .projector
+        .projection()
+        .project_schema(projection.input().schema().as_ref())?;
+    let output_schema = projection.schema();
+    Ok(derived_schema.metadata() != output_schema.metadata()
+        || derived_schema
+            .fields()
+            .iter()
+            .zip(output_schema.fields())
+            .any(|(derived, output)| derived.metadata() != output.metadata()))
+}
+
 /// Collapse a chain of consecutive [`ProjectionExec`]s into one. Returns
 /// `None` if nothing could be merged.
 fn try_collapse_projection_chain(
     outer: &ProjectionExec,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+    if projection_overrides_metadata(outer)? {
+        return Ok(None);
+    }
+
     let mut current_exprs: Vec<ProjectionExpr> = outer.expr().to_vec();
     let mut current_input: Arc<dyn ExecutionPlan> = Arc::clone(outer.input());
     let mut column_ref_map: HashMap<Column, usize> = HashMap::new();
     let mut collapsed_any = false;
 
     'outer: while let Some(inner_proj) = current_input.downcast_ref::<ProjectionExec>() {
+        if projection_overrides_metadata(inner_proj)? {
+            break;
+        }
+
         // Collect the column references usage in the outer projection.
         column_ref_map.clear();
         for proj_expr in &current_exprs {
@@ -1385,9 +1426,14 @@ fn try_collapse_projection_chain(
         return Ok(None);
     }
 
-    // To unify 3 or more sequential projections:
+    // Expression substitution must not change the outer projection's output
+    // metadata contract.
     let unified: Arc<dyn ExecutionPlan> =
-        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 +1563,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 +1615,91 @@ mod tests {
         Ok(())
     }
 
+    fn identity_projection_with_metadata(
+        input: Arc<dyn ExecutionPlan>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        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::<ProjectionExec>().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::<ProjectionExec>()
+            .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<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new(
+            [ProjectionExpr {
+                expr: Arc::new(arrow_metadata),
+                alias: "metadata".to_string(),
+            }],
+            inner,
+        )?);
+
+        let outer_projection = outer
+            .downcast_ref::<ProjectionExec>()
+            .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::<StringArray>()
+            .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(