Preserve metadata-changing projections

Author: osipovartem

datafusion/physical-plan/src/projection.rs

@@ -968,6 +968,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
@@ -1006,8 +1007,12 @@ 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.
@@ -1441,6 +1446,7 @@ mod tests {
     use crate::statistics::{StatisticsArgs, StatisticsContext};
     use crate::test;
     use crate::test::exec::StatisticsExec;
+    use crate::union::UnionExec;
 
     use arrow::datatypes::{DataType, Field, Schema};
     use datafusion_common::ScalarValue;
@@ -1491,6 +1497,38 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_projection_pushdown_preserves_output_metadata() -> Result<()> {
+        let input_schema = Arc::new(Schema::new(vec![
+            Field::new("input", DataType::Int32, false).with_metadata(HashMap::from([(
+                "source".to_string(),
+                "input".to_string(),
+            )])),
+            Field::new("unused", DataType::Int32, false),
+        ]));
+        let input: Arc<dyn ExecutionPlan> = UnionExec::try_new(vec![
+            Arc::new(EmptyExec::new(Arc::clone(&input_schema))),
+            Arc::new(EmptyExec::new(input_schema)),
+        ])?;
+        let projected_schema =
+            Schema::new(vec![Field::new("input", DataType::Int32, false)]);
+        let projection = ProjectionExec::try_new_with_schema_metadata(
+            [ProjectionExpr {
+                expr: Arc::new(Column::new("input", 0)),
+                alias: "input".to_string(),
+            }],
+            input,
+            &projected_schema,
+        )?;
+
+        assert!(!is_projection_removable(&projection));
+        let plan: Arc<dyn ExecutionPlan> = Arc::new(projection);
+        let optimized = remove_unnecessary_projections(Arc::clone(&plan))?;
+        assert!(optimized.transformed);
+        assert_eq!(optimized.data.schema(), plan.schema());
+        Ok(())
+    }
+
     #[test]
     fn test_collect_column_indices() -> Result<()> {
         let expr = Arc::new(BinaryExpr::new(