Fix ordering for UNION ALL over heterogeneous constants

datafusion/physical-expr/src/equivalence/class.rs

@@ -551,7 +551,19 @@ impl EquivalenceGroup {
         sort_exprs
             .into_iter()
             .map(|sort_expr| self.normalize_sort_expr(sort_expr))
-            .filter(|sort_expr| self.is_expr_constant(&sort_expr.expr).is_none())
+            .filter(|sort_expr| !self.is_uniform_constant(&sort_expr.expr))
+    }
+
+    /// Returns `true` when `expr` is a *globally* constant column, safe to drop
+    /// from a required ordering. Only [`AcrossPartitions::Uniform`] qualifies; a
+    /// [`AcrossPartitions::Heterogeneous`] value is constant within a partition
+    /// but varies across partitions, so it still discriminates the order once
+    /// partitions are merged and must be kept.
+    fn is_uniform_constant(&self, expr: &Arc<dyn PhysicalExpr>) -> bool {
+        matches!(
+            self.is_expr_constant(expr),
+            Some(AcrossPartitions::Uniform(_))
+        )
     }
 
     /// Normalizes the given sort requirement according to this group. The
@@ -582,7 +594,7 @@ impl EquivalenceGroup {
         sort_reqs
             .into_iter()
             .map(|req| self.normalize_sort_requirement(req))
-            .filter(|req| self.is_expr_constant(&req.expr).is_none())
+            .filter(|req| !self.is_uniform_constant(&req.expr))
     }
 
     /// Perform an indirect projection of `expr` by consulting the equivalence

datafusion/sqllogictest/test_files/order.slt

@@ -1770,3 +1770,41 @@ reset datafusion.sql_parser.default_null_ordering;
 
 statement ok
 reset datafusion.sql_parser.dialect;
+
+# A global sort feeding a sink (CopyTo) must keep a leading key that is constant
+# within each partition but differs across them ("a" is 2 on one union branch,
+# 1 on the other). The merge above the union has to reorder rows across branches,
+# so the physical plan must keep "a" in its ordering; dropping it (leaving only
+# [b@1 ASC]) silently loses the global order under the sink.
+statement ok
+CREATE TABLE t2(b INT) AS VALUES (10);
+
+query TT
+EXPLAIN COPY (
+  SELECT 2 AS a, b FROM t2
+  UNION ALL
+  SELECT 1 AS a, b FROM t2
+  ORDER BY a, b
+) TO 'test_files/scratch/order/sort_key_sink.parquet';
+----
+logical_plan
+01)CopyTo: format=parquet output_url=test_files/scratch/order/sort_key_sink.parquet options: ()
+02)--Sort: a ASC NULLS LAST, b ASC NULLS LAST
+03)----Union
+04)------Projection: Int64(2) AS a, t2.b
+05)--------TableScan: t2 projection=[b]
+06)------Projection: Int64(1) AS a, t2.b
+07)--------TableScan: t2 projection=[b]
+physical_plan
+01)DataSinkExec: sink=ParquetSink(file_groups=[])
+02)--SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST]
+03)----UnionExec
+04)------SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false]
+05)--------ProjectionExec: expr=[2 as a, b@0 as b]
+06)----------DataSourceExec: partitions=1, partition_sizes=[1]
+07)------SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false]
+08)--------ProjectionExec: expr=[1 as a, b@0 as b]
+09)----------DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+DROP TABLE t2;