fix: handle ordering by order of inputs

Author: rkrishn7

datafusion/expr/src/logical_plan/plan.rs

@@ -18,7 +18,7 @@
 //! Logical plan types
 
 use std::cmp::Ordering;
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{HashMap, HashSet};
 use std::fmt::{self, Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
 use std::sync::{Arc, LazyLock};
@@ -2741,35 +2741,34 @@ impl Union {
     ) -> Result<DFSchemaRef> {
         type FieldData<'a> =
             (&'a DataType, bool, Vec<&'a HashMap<String, String>>, usize);
-        // Prefer `BTreeMap` as it produces items in order by key when iterated over
-        let mut cols: BTreeMap<&str, FieldData> = BTreeMap::new();
+        let mut cols: Vec<(&str, FieldData)> = Vec::new();
         for input in inputs.iter() {
             for field in input.schema().fields() {
-                match cols.entry(field.name()) {
-                    std::collections::btree_map::Entry::Occupied(mut occupied) => {
-                        let (data_type, is_nullable, metadata, occurrences) =
-                            occupied.get_mut();
-                        if !loose_types && *data_type != field.data_type() {
-                            return plan_err!(
-                                "Found different types for field {}",
-                                field.name()
-                            );
-                        }
-
-                        metadata.push(field.metadata());
-                        // If the field is nullable in any one of the inputs,
-                        // then the field in the final schema is also nullable.
-                        *is_nullable |= field.is_nullable();
-                        *occurrences += 1;
+                if let Some((_, (data_type, is_nullable, metadata, occurrences))) =
+                    cols.iter_mut().find(|(name, _)| name == field.name())
+                {
+                    if !loose_types && *data_type != field.data_type() {
+                        return plan_err!(
+                            "Found different types for field {}",
+                            field.name()
+                        );
                     }
-                    std::collections::btree_map::Entry::Vacant(vacant) => {
-                        vacant.insert((
+
+                    metadata.push(field.metadata());
+                    // If the field is nullable in any one of the inputs,
+                    // then the field in the final schema is also nullable.
+                    *is_nullable |= field.is_nullable();
+                    *occurrences += 1;
+                } else {
+                    cols.push((
+                        field.name(),
+                        (
                             field.data_type(),
                             field.is_nullable(),
                             vec![field.metadata()],
                             1,
-                        ));
-                    }
+                        ),
+                    ));
                 }
             }
         }

datafusion/sql/tests/sql_integration.rs

@@ -1898,10 +1898,10 @@ fn union_by_name_different_columns() {
     let expected = "\
         Distinct:\
         \n  Union\
-        \n    Projection: NULL AS Int64(1), order_id\
+        \n    Projection: order_id, NULL AS Int64(1)\
         \n      Projection: orders.order_id\
         \n        TableScan: orders\
-        \n    Projection: Int64(1), order_id\
+        \n    Projection: order_id, Int64(1)\
         \n      Projection: orders.order_id, Int64(1)\
         \n        TableScan: orders";
     quick_test(sql, expected);
@@ -1937,10 +1937,10 @@ fn union_all_by_name_different_columns() {
         "SELECT order_id from orders UNION ALL BY NAME SELECT order_id, 1 FROM orders";
     let expected = "\
         Union\
-        \n  Projection: NULL AS Int64(1), order_id\
+        \n  Projection: order_id, NULL AS Int64(1)\
         \n    Projection: orders.order_id\
         \n      TableScan: orders\
-        \n  Projection: Int64(1), order_id\
+        \n  Projection: order_id, Int64(1)\
         \n    Projection: orders.order_id, Int64(1)\
         \n      TableScan: orders";
     quick_test(sql, expected);

datafusion/sqllogictest/test_files/union_by_name.slt

@@ -88,38 +88,38 @@ SELECT x FROM t1 UNION ALL BY NAME SELECT x FROM t1 ORDER BY x;
 query II
 (SELECT x FROM t1 UNION ALL SELECT x FROM t1) UNION BY NAME SELECT 5 ORDER BY x;
 ----
-NULL 1
-NULL 3
-5 NULL
+1 NULL
+3 NULL
+NULL 5
 
 query II
 (SELECT x FROM t1 UNION ALL SELECT x FROM t1) UNION ALL BY NAME SELECT 5 ORDER BY x;
 ----
-NULL 1
-NULL 1
-NULL 3
-NULL 3
-NULL 3
-NULL 3
-5 NULL
+1 NULL
+1 NULL
+3 NULL
+3 NULL
+3 NULL
+3 NULL
+NULL 5
 
 query II
 (SELECT x FROM t1 UNION ALL SELECT y FROM t1) UNION BY NAME SELECT 5 ORDER BY x;
 ----
-NULL 1
-NULL 3
-5 NULL
+1 NULL
+3 NULL
+NULL 5
 
 query II
 (SELECT x FROM t1 UNION ALL SELECT y FROM t1) UNION ALL BY NAME SELECT 5 ORDER BY x;
 ----
-NULL 1
-NULL 1
-NULL 3
-NULL 3
-NULL 3
-NULL 3
-5 NULL
+1 NULL
+1 NULL
+3 NULL
+3 NULL
+3 NULL
+3 NULL
+NULL 5
 
 
 # Ambiguous name
@@ -152,22 +152,22 @@ NULL 4
 # Limit
 
 query III
-SELECT 1 UNION BY NAME SELECT * FROM unnest(range(2, 100)) UNION BY NAME SELECT 999 ORDER BY 3, 1 LIMIT 5;
+SELECT 1 UNION BY NAME SELECT * FROM unnest(range(2, 100)) UNION BY NAME SELECT 999 ORDER BY 3, 1, 2 LIMIT 5;
 ----
-NULL NULL 2
-NULL NULL 3
-NULL NULL 4
-NULL NULL 5
-NULL NULL 6
+NULL NULL 999
+1 NULL NULL
+NULL 2 NULL
+NULL 3 NULL
+NULL 4 NULL
 
 query III
 SELECT 1 UNION ALL BY NAME SELECT * FROM unnest(range(2, 100)) UNION ALL BY NAME SELECT 999 ORDER BY 3, 1 LIMIT 5;
 ----
-NULL NULL 2
-NULL NULL 3
-NULL NULL 4
-NULL NULL 5
-NULL NULL 6
+NULL NULL 999
+1 NULL NULL
+NULL 2 NULL
+NULL 3 NULL
+NULL 4 NULL
 
 # Order by