From 15eb10b6d9aeef0c65aa67dd40528e2c00211a10 Mon Sep 17 00:00:00 2001 From: Vadim Piven Date: Mon, 13 Jul 2026 13:19:04 +0200 Subject: [PATCH 1/3] fix: keep heterogeneous constants in required orderings (add_sort_above) `add_sort_above` and `add_sort_above_with_distribution` pruned every expression that equivalence properties reported as constant, via `is_expr_constant(..).is_none()`. That treats `AcrossPartitions::Heterogeneous` the same as `AcrossPartitions::Uniform`, which is wrong. A `Heterogeneous` constant is constant *within* each partition but varies *across* partitions (e.g. `sample` in `SELECT 1 AS sample, c FROM t UNION ALL SELECT 2 AS sample, c FROM t`). Both helpers build a partition-preserving `SortExec` whose output is merged across partitions downstream, where such a column is still the leading ordering discriminator. Dropping it silently discards the global ordering, producing non-deterministic output order for `ORDER BY sample, c` once the plan is rooted at a sink / output requirement. Only `AcrossPartitions::Uniform` is globally constant and safe to drop. Introduce `is_uniform_constant` and use it at both prune sites. Adds a regression test that builds a union with a heterogeneous-constant column and asserts the key survives `add_sort_above`. --- datafusion/physical-optimizer/src/utils.rs | 107 +++++++++++++++++++-- 1 file changed, 97 insertions(+), 10 deletions(-) diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index 1fbf8c6fb78cd..27115a2dc3a54 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -19,8 +19,8 @@ use std::sync::Arc; use datafusion_common::Result; use datafusion_physical_expr::{ - Distribution, EquivalenceProperties, LexOrdering, LexRequirement, Partitioning, - PhysicalExpr, physical_exprs_equal, + AcrossPartitions, Distribution, EquivalenceProperties, LexOrdering, LexRequirement, + Partitioning, PhysicalExpr, physical_exprs_equal, }; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; @@ -32,6 +32,25 @@ use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +/// Returns `true` when `expr` is a *globally* constant column in `eq_properties` +/// and can therefore be safely dropped from a required ordering. +/// +/// Only [`AcrossPartitions::Uniform`] qualifies. An +/// [`AcrossPartitions::Heterogeneous`] value is constant *within* each partition +/// but may differ *across* partitions, so it must be kept in the sort key: the +/// `SortExec` these helpers build preserves partitioning, and its output is +/// merged across partitions downstream, where the column is still an ordering +/// discriminator. Dropping it there silently loses the global ordering. +fn is_uniform_constant( + eq_properties: &EquivalenceProperties, + expr: &Arc, +) -> bool { + matches!( + eq_properties.is_expr_constant(expr), + Some(AcrossPartitions::Uniform(_)) + ) +} + /// This utility function adds a `SortExec` above an operator according to the /// given ordering requirements while preserving the original partitioning. /// @@ -45,10 +64,7 @@ pub fn add_sort_above( ) -> PlanContext { let mut sort_reqs: Vec<_> = sort_requirements.into(); sort_reqs.retain(|sort_expr| { - node.plan - .equivalence_properties() - .is_expr_constant(&sort_expr.expr) - .is_none() + !is_uniform_constant(node.plan.equivalence_properties(), &sort_expr.expr) }); let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::>(); let Some(ordering) = LexOrdering::new(sort_exprs) else { @@ -73,10 +89,7 @@ pub fn add_sort_above_with_distribution( ) -> PlanContext { let mut sort_reqs: Vec<_> = sort_requirements.into(); sort_reqs.retain(|sort_expr| { - node.plan - .equivalence_properties() - .is_expr_constant(&sort_expr.expr) - .is_none() + !is_uniform_constant(node.plan.equivalence_properties(), &sort_expr.expr) }); let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::>(); let Some(ordering) = LexOrdering::new(sort_exprs) else { @@ -220,3 +233,77 @@ pub(crate) fn range_partitioning_satisfies_key_partitioning( pub fn is_limit(plan: &Arc) -> bool { plan.is::() || plan.is::() } + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr::PhysicalSortRequirement; + use datafusion_physical_expr::expressions::{col, lit}; + use datafusion_physical_expr::projection::ProjectionExpr; + use datafusion_physical_plan::displayable; + use datafusion_physical_plan::empty::EmptyExec; + use datafusion_physical_plan::projection::ProjectionExec; + + /// A `UnionExec` whose `sample` column is 1 on the left branch and 2 on the + /// right makes `sample` a `Heterogeneous` constant: constant within each + /// partition, but different across them. A required ordering of + /// `[sample, other]` must retain `sample`; the `SortExec` built here + /// preserves partitioning and its output is merged across partitions + /// downstream, where `sample` is still the leading ordering discriminator. + /// Regression test: `add_sort_above` used to drop every constant, uniform + /// or not, and silently discarded the global ordering. + #[test] + fn add_sort_above_keeps_heterogeneous_constant_key() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)])); + + let branch = |value: i32| -> Result> { + let source = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let projection = ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: lit(value), + alias: "sample".to_string(), + }, + ProjectionExpr { + expr: col("c", &schema)?, + alias: "other".to_string(), + }, + ], + source, + )?; + Ok(Arc::new(projection)) + }; + + let union: Arc = + UnionExec::try_new(vec![branch(1)?, branch(2)?])?; + let union_schema = union.schema(); + + // Precondition: the union reports `sample` as a heterogeneous constant. + let sample = col("sample", &union_schema)?; + assert!(matches!( + union.equivalence_properties().is_expr_constant(&sample), + Some(AcrossPartitions::Heterogeneous) + )); + + let asc = SortOptions::default(); + let requirement: LexRequirement = [ + PhysicalSortRequirement::new(sample, Some(asc)), + PhysicalSortRequirement::new(col("other", &union_schema)?, Some(asc)), + ] + .into(); + + let node = PlanContext::<()>::new_default(union); + let result = add_sort_above(node, requirement, None); + + let plan = displayable(result.plan.as_ref()).indent(true).to_string(); + assert!( + plan.contains("sample@0 ASC"), + "add_sort_above dropped the heterogeneous-constant sort key:\n{plan}" + ); + + Ok(()) + } +} From b301d9e8ac1fb68446ef7d49b954c429be9c3bff Mon Sep 17 00:00:00 2001 From: Vadim Piven Date: Thu, 16 Jul 2026 13:55:02 +0200 Subject: [PATCH 2/3] fix: keep heterogeneous constants in normalized required orderings The previous fix touched `add_sort_above`/`add_sort_above_with_distribution` in physical-optimizer, but that path is not where real SQL loses the key. A global sort feeding a sink (`CopyTo`) still dropped its leading ordering key when that key is an `AcrossPartitions::Heterogeneous` constant, producing non-deterministically ordered sink output. The actual prune happens in `EquivalenceGroup::normalize_sort_exprs` and `normalize_sort_requirements`, which filtered out *every* constant via `is_expr_constant(..).is_none()`. That drops `Heterogeneous` constants too, even though such a column is constant only *within* a partition and still discriminates the global order once partitions are merged (e.g. a partition-preserving `SortExec` feeding a `SortPreservingMergeExec` under a sink). Only `AcrossPartitions::Uniform` is globally constant and safe to drop. Revert the physical-optimizer change and instead keep heterogeneous constants at both normalize sites via a new `is_uniform_constant` helper. Replace the Rust unit test with an `.slt` regression test that plans `EXPLAIN COPY ( ORDER BY a, b) TO parquet` and asserts the sink-rooted plan keeps `[a@0 ASC, b@1 ASC]` in the merge above the union. This exercises the real planning path end to end rather than calling `add_sort_above` directly. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../physical-expr/src/equivalence/class.rs | 16 ++- datafusion/physical-optimizer/src/utils.rs | 107 ++---------------- datafusion/sqllogictest/test_files/order.slt | 38 +++++++ 3 files changed, 62 insertions(+), 99 deletions(-) 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/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index 27115a2dc3a54..1fbf8c6fb78cd 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -19,8 +19,8 @@ use std::sync::Arc; use datafusion_common::Result; use datafusion_physical_expr::{ - AcrossPartitions, Distribution, EquivalenceProperties, LexOrdering, LexRequirement, - Partitioning, PhysicalExpr, physical_exprs_equal, + Distribution, EquivalenceProperties, LexOrdering, LexRequirement, Partitioning, + PhysicalExpr, physical_exprs_equal, }; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; @@ -32,25 +32,6 @@ use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; -/// Returns `true` when `expr` is a *globally* constant column in `eq_properties` -/// and can therefore be safely dropped from a required ordering. -/// -/// Only [`AcrossPartitions::Uniform`] qualifies. An -/// [`AcrossPartitions::Heterogeneous`] value is constant *within* each partition -/// but may differ *across* partitions, so it must be kept in the sort key: the -/// `SortExec` these helpers build preserves partitioning, and its output is -/// merged across partitions downstream, where the column is still an ordering -/// discriminator. Dropping it there silently loses the global ordering. -fn is_uniform_constant( - eq_properties: &EquivalenceProperties, - expr: &Arc, -) -> bool { - matches!( - eq_properties.is_expr_constant(expr), - Some(AcrossPartitions::Uniform(_)) - ) -} - /// This utility function adds a `SortExec` above an operator according to the /// given ordering requirements while preserving the original partitioning. /// @@ -64,7 +45,10 @@ pub fn add_sort_above( ) -> PlanContext { let mut sort_reqs: Vec<_> = sort_requirements.into(); sort_reqs.retain(|sort_expr| { - !is_uniform_constant(node.plan.equivalence_properties(), &sort_expr.expr) + node.plan + .equivalence_properties() + .is_expr_constant(&sort_expr.expr) + .is_none() }); let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::>(); let Some(ordering) = LexOrdering::new(sort_exprs) else { @@ -89,7 +73,10 @@ pub fn add_sort_above_with_distribution( ) -> PlanContext { let mut sort_reqs: Vec<_> = sort_requirements.into(); sort_reqs.retain(|sort_expr| { - !is_uniform_constant(node.plan.equivalence_properties(), &sort_expr.expr) + node.plan + .equivalence_properties() + .is_expr_constant(&sort_expr.expr) + .is_none() }); let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::>(); let Some(ordering) = LexOrdering::new(sort_exprs) else { @@ -233,77 +220,3 @@ pub(crate) fn range_partitioning_satisfies_key_partitioning( pub fn is_limit(plan: &Arc) -> bool { plan.is::() || plan.is::() } - -#[cfg(test)] -mod tests { - use super::*; - - use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_physical_expr::PhysicalSortRequirement; - use datafusion_physical_expr::expressions::{col, lit}; - use datafusion_physical_expr::projection::ProjectionExpr; - use datafusion_physical_plan::displayable; - use datafusion_physical_plan::empty::EmptyExec; - use datafusion_physical_plan::projection::ProjectionExec; - - /// A `UnionExec` whose `sample` column is 1 on the left branch and 2 on the - /// right makes `sample` a `Heterogeneous` constant: constant within each - /// partition, but different across them. A required ordering of - /// `[sample, other]` must retain `sample`; the `SortExec` built here - /// preserves partitioning and its output is merged across partitions - /// downstream, where `sample` is still the leading ordering discriminator. - /// Regression test: `add_sort_above` used to drop every constant, uniform - /// or not, and silently discarded the global ordering. - #[test] - fn add_sort_above_keeps_heterogeneous_constant_key() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)])); - - let branch = |value: i32| -> Result> { - let source = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let projection = ProjectionExec::try_new( - vec![ - ProjectionExpr { - expr: lit(value), - alias: "sample".to_string(), - }, - ProjectionExpr { - expr: col("c", &schema)?, - alias: "other".to_string(), - }, - ], - source, - )?; - Ok(Arc::new(projection)) - }; - - let union: Arc = - UnionExec::try_new(vec![branch(1)?, branch(2)?])?; - let union_schema = union.schema(); - - // Precondition: the union reports `sample` as a heterogeneous constant. - let sample = col("sample", &union_schema)?; - assert!(matches!( - union.equivalence_properties().is_expr_constant(&sample), - Some(AcrossPartitions::Heterogeneous) - )); - - let asc = SortOptions::default(); - let requirement: LexRequirement = [ - PhysicalSortRequirement::new(sample, Some(asc)), - PhysicalSortRequirement::new(col("other", &union_schema)?, Some(asc)), - ] - .into(); - - let node = PlanContext::<()>::new_default(union); - let result = add_sort_above(node, requirement, None); - - let plan = displayable(result.plan.as_ref()).indent(true).to_string(); - assert!( - plan.contains("sample@0 ASC"), - "add_sort_above dropped the heterogeneous-constant sort key:\n{plan}" - ); - - Ok(()) - } -} diff --git a/datafusion/sqllogictest/test_files/order.slt b/datafusion/sqllogictest/test_files/order.slt index 79fb676f4b410..f4d130daa90b0 100644 --- a/datafusion/sqllogictest/test_files/order.slt +++ b/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; From 53719c485618f1e3d257d637551acf67615e4cc4 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 16 Jul 2026 13:21:55 -0400 Subject: [PATCH 3/3] add output validation to slt --- datafusion/sqllogictest/test_files/order.slt | 32 +++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/order.slt b/datafusion/sqllogictest/test_files/order.slt index f4d130daa90b0..978fcc197c0de 100644 --- a/datafusion/sqllogictest/test_files/order.slt +++ b/datafusion/sqllogictest/test_files/order.slt @@ -1777,7 +1777,7 @@ reset datafusion.sql_parser.dialect; # 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); +CREATE TABLE t2(b INT) AS VALUES (10), (20); query TT EXPLAIN COPY ( @@ -1806,5 +1806,35 @@ physical_plan 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;