Temporary cherry pick of apache#16064

Author: LiaCastaneda

datafusion/core/src/physical_planner.rs

@@ -62,7 +62,9 @@ use arrow::array::{builder::StringBuilder, RecordBatch};
 use arrow::compute::SortOptions;
 use arrow::datatypes::{Schema, SchemaRef};
 use datafusion_common::display::ToStringifiedPlan;
-use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
+use datafusion_common::tree_node::{
+    Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeVisitor,
+};
 use datafusion_common::{
     exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema,
     ScalarValue,
@@ -2075,29 +2077,37 @@ fn maybe_fix_physical_column_name(
     expr: Result<Arc<dyn PhysicalExpr>>,
     input_physical_schema: &SchemaRef,
 ) -> Result<Arc<dyn PhysicalExpr>> {
-    if let Ok(e) = &expr {
-        if let Some(column) = e.as_any().downcast_ref::<Column>() {
-            let physical_field = input_physical_schema.field(column.index());
-            let expr_col_name = column.name();
-            let physical_name = physical_field.name();
-
-            if physical_name != expr_col_name {
-                // handle edge cases where the physical_name contains ':'.
-                let colon_count = physical_name.matches(':').count();
-                let mut splits = expr_col_name.match_indices(':');
-                let split_pos = splits.nth(colon_count);
-
-                if let Some((idx, _)) = split_pos {
-                    let base_name = &expr_col_name[..idx];
-                    if base_name == physical_name {
-                        let updated_column = Column::new(physical_name, column.index());
-                        return Ok(Arc::new(updated_column));
+    expr.and_then(|e| {
+        e.transform_down(|node| {
+            if let Some(column) = node.as_any().downcast_ref::<Column>() {
+                let idx = column.index();
+                let physical_field = input_physical_schema.field(idx);
+                let expr_col_name = column.name();
+                let physical_name = physical_field.name();
+
+                if expr_col_name != physical_name {
+                    // handle edge cases where the physical_name contains ':'.
+                    let colon_count = physical_name.matches(':').count();
+                    let mut splits = expr_col_name.match_indices(':');
+                    let split_pos = splits.nth(colon_count);
+
+                    if let Some((i, _)) = split_pos {
+                        let base_name = &expr_col_name[..i];
+                        if base_name == physical_name {
+                            let updated_column = Column::new(physical_name, idx);
+                            return Ok(Transformed::yes(Arc::new(updated_column)));
+                        }
                     }
                 }
+
+                // If names already match or fix is not possible, just leave it as it is
+                Ok(Transformed::no(node))
+            } else {
+                Ok(Transformed::no(node))
             }
-        }
-    }
-    expr
+        })
+        .data()
+    })
 }
 
 struct OptimizationInvariantChecker<'a> {

datafusion/physical-plan/src/union.rs

@@ -513,7 +513,12 @@ fn union_schema(inputs: &[Arc<dyn ExecutionPlan>]) -> SchemaRef {
 
     let fields = (0..first_schema.fields().len())
         .map(|i| {
-            inputs
+            // We take the name from the left side of the union to match how names are coerced during logical planning,
+            // which also uses the left side names.
+            let base_field = first_schema.field(i).clone();
+
+            // Coerce metadata and nullability across all inputs
+            let merged_field = inputs
                 .iter()
                 .enumerate()
                 .map(|(input_idx, input)| {
@@ -535,6 +540,9 @@ fn union_schema(inputs: &[Arc<dyn ExecutionPlan>]) -> SchemaRef {
                 // We can unwrap this because if inputs was empty, this would've already panic'ed when we
                 // indexed into inputs[0].
                 .unwrap()
+                .with_name(base_field.name());
+
+            merged_field
         })
         .collect::<Vec<_>>();
 

datafusion/substrait/tests/cases/consumer_integration.rs

@@ -560,4 +560,26 @@ mod tests {
                         );
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_multiple_unions() -> Result<()> {
+        let plan_str = test_plan_to_string("multiple_unions.json").await?;
+        assert_snapshot!(
+        plan_str,
+        @"Projection: Utf8(\"people\") AS product_category, Utf8(\"people\")__temp__0 AS product_type, product_key\
+        \n  Union\
+        \n    Projection: Utf8(\"people\"), Utf8(\"people\") AS Utf8(\"people\")__temp__0, sales.product_key\
+        \n      Left Join: sales.product_key = food.@food_id\
+        \n        TableScan: sales\
+        \n        TableScan: food\
+        \n    Union\
+        \n      Projection: people.$f3, people.$f5, people.product_key0\
+        \n        Left Join: people.product_key0 = food.@food_id\
+        \n          TableScan: people\
+        \n          TableScan: food\
+        \n      TableScan: more_products"
+        );
+
+        Ok(())
+    }
 }