diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index d00a4a32278f0..1f9a6a583cc44 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/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) -> 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 diff --git a/datafusion/sqllogictest/test_files/order.slt b/datafusion/sqllogictest/test_files/order.slt index 79fb676f4b410..978fcc197c0de 100644 --- a/datafusion/sqllogictest/test_files/order.slt +++ b/datafusion/sqllogictest/test_files/order.slt @@ -1770,3 +1770,71 @@ 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), (20); + +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] + +# Actually execute the COPY and verify the rows written to the file are in +# global (a, b) order, interleaving the two union branches. If "a" were +# dropped from the sort, the file would instead contain rows ordered only +# by "b" within each branch. +query I +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'; +---- +4 + +statement ok +CREATE EXTERNAL TABLE sort_key_sink STORED AS PARQUET +LOCATION 'test_files/scratch/order/sort_key_sink.parquet'; + +# Note: no ORDER BY here, so this checks the order rows were written in +query II +SELECT * FROM sort_key_sink; +---- +1 10 +1 20 +2 10 +2 20 + +statement ok +DROP TABLE sort_key_sink; + +statement ok +DROP TABLE t2;