From 78c543618542ca2bc33853f95aae74d5be4e919a Mon Sep 17 00:00:00 2001 From: Kristofferjoha Date: Sun, 25 Oct 2026 15:43:09 +0100 Subject: [PATCH 1/2] retain valid orderings --- .../physical-expr/src/equivalence/class.rs | 10 +++++- .../src/equivalence/properties/mod.rs | 32 ++++++++++++++++++ .../physical-plan/src/repartition/mod.rs | 12 ++++--- .../test_files/equivalence_properties.slt | 33 +++++++++++++++++++ .../test_files/monotonic_projection_test.slt | 8 +++++ datafusion/sqllogictest/test_files/order.slt | 24 ++++++++++++++ .../push_down_filter_regression.slt | 4 +-- datafusion/sqllogictest/test_files/window.slt | 2 +- 8 files changed, 116 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index ff34190bb1c1b..bc185459378a9 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/datafusion/physical-expr/src/equivalence/class.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use std::vec::IntoIter; use super::ProjectionMapping; -use crate::expressions::Literal; +use crate::expressions::{CastExpr, Literal}; use crate::physical_expr::add_offset_to_expr; use crate::projection::ProjectionTargets; use crate::{PhysicalExpr, PhysicalExprRef, PhysicalSortExpr, PhysicalSortRequirement}; @@ -741,6 +741,7 @@ impl EquivalenceGroup { let mut new_constants = vec![]; let mut new_classes = IndexMap::<_, EquivalenceClass>::with_capacity(mapping.len()); + let mut augmented = None; for (source, targets) in mapping.iter() { // We need to find equivalent projected expressions. For example, // consider a table with columns `[a, b, c]` with `a` == `b`, and @@ -752,6 +753,13 @@ impl EquivalenceGroup { for (target, _) in targets.iter() { cls.push(Arc::clone(target)); } + if source.is::() { + let augmented = augmented + .get_or_insert_with(|| self.augment_projection_mapping(mapping)); + if let Some(projected) = Self::project_expr_indirect(augmented, source) { + cls.push(projected); + } + } // Save new constants arising from the projection: if let Some(across) = self.is_expr_constant(source) { for (target, _) in targets.iter() { diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index ed824e5712f96..6214b838231b6 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -702,7 +702,17 @@ impl EquivalenceProperties { /// sorted within each such group. fn add_satisfied_key_constants(&mut self, expr: Arc) -> Result<()> { let mut stack = vec![expr]; + let mut visited = IndexSet::new(); while let Some(expr) = stack.pop() { + if let Some(class) = self.eq_group.get_equivalence_class(&expr) { + for equivalent in class.iter() { + if !equivalent.children().is_empty() + && visited.insert(Arc::clone(equivalent)) + { + stack.push(Arc::clone(equivalent)); + } + } + } let properties = self.get_expr_properties(Arc::clone(&expr)); if properties.strictly_order_preserving { for child in expr.children() { @@ -1402,6 +1412,28 @@ impl EquivalenceProperties { /// Returns an [`ExprProperties`] object containing the ordering and range /// information for the given expression. pub fn get_expr_properties(&self, expr: Arc) -> ExprProperties { + let class = self.eq_group.get_equivalence_class(&expr); + let mut properties = self.get_expr_properties_from_tree(expr); + if properties.sort_properties != SortProperties::Unordered { + return properties; + } + if let Some(class) = class { + for equivalent in class.iter() { + let equivalent_properties = + self.get_expr_properties_from_tree(Arc::clone(equivalent)); + if equivalent_properties.sort_properties != SortProperties::Unordered { + properties.sort_properties = equivalent_properties.sort_properties; + return properties; + } + } + } + properties + } + + fn get_expr_properties_from_tree( + &self, + expr: Arc, + ) -> ExprProperties { ExprPropertiesNode::new_unknown(expr) .transform_up(|expr| update_properties(expr, self)) .data() diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index b4984ab5b234f..3c0d9b7873d36 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -2200,15 +2200,17 @@ impl RepartitionExec { input: &Arc, preserve_order: bool, ) -> EquivalenceProperties { - // Equivalence Properties let mut eq_properties = input.equivalence_properties().clone(); - // If the ordering is lost, reset the ordering equivalence class: + let input_partitions = input.output_partitioning().partition_count(); if !Self::maintains_input_order_helper(input, preserve_order)[0] { eq_properties.clear_orderings(); + } else if input_partitions > 1 { + eq_properties.clear_orderings(); + if let Some(ordering) = input.output_ordering() { + eq_properties.add_ordering(ordering.clone()); + } } - // When there are more than one input partitions, they will be fused at the output. - // Therefore, remove per partition constants. - if input.output_partitioning().partition_count() > 1 { + if input_partitions > 1 { eq_properties.clear_per_partition_constants(); } eq_properties diff --git a/datafusion/sqllogictest/test_files/equivalence_properties.slt b/datafusion/sqllogictest/test_files/equivalence_properties.slt index a1484231e499b..ecc78fb41b9e3 100644 --- a/datafusion/sqllogictest/test_files/equivalence_properties.slt +++ b/datafusion/sqllogictest/test_files/equivalence_properties.slt @@ -147,6 +147,39 @@ SELECT l.c, s.x FROM (SELECT c FROM t) l CROSS JOIN (SELECT x FROM small WHERE x 1 1 2 1 +### +### Equal expressions need not be injective +### + +# c = abs(a) transfers ordering to abs(a), but equal c values do not pin +# down a. The sort on b must remain for the two rows with c = 1. +statement ok +set datafusion.execution.target_partitions = 1; + +query I +COPY (VALUES (-1, 2, 1), (1, 1, 1)) +TO 'test_files/scratch/equivalence_properties/abs.csv' +OPTIONS ('format.has_header' 'true'); +---- +2 + +statement ok +CREATE EXTERNAL TABLE abs_equivalence (a INT, b INT, c INT) +STORED AS CSV +WITH ORDER (a, b) +WITH ORDER (c) +LOCATION 'test_files/scratch/equivalence_properties/abs.csv' +OPTIONS ('format.has_header' 'true'); + +query III +SELECT a, b, c FROM abs_equivalence WHERE c = abs(a) ORDER BY c, b; +---- +1 1 1 +-1 2 1 + +statement ok +DROP TABLE abs_equivalence; + ### ### Order-preserving repartition of an unbounded input ### diff --git a/datafusion/sqllogictest/test_files/monotonic_projection_test.slt b/datafusion/sqllogictest/test_files/monotonic_projection_test.slt index 4ecdbca24a116..3e665fe754baa 100644 --- a/datafusion/sqllogictest/test_files/monotonic_projection_test.slt +++ b/datafusion/sqllogictest/test_files/monotonic_projection_test.slt @@ -289,6 +289,14 @@ SELECT k, cast(v as float) v_ FROM t2_int_float ORDER BY v_ DESC, k DESC; 2 16777216 1 16777216 +# Retaining v lets projection record v_ = CAST(v AS FLOAT). Equal v_ values +# still do not imply equal v values: the suffix key must be sorted again. +query IIR +SELECT k, v, cast(v as float) v_ FROM t1_int_float ORDER BY v_ DESC, k DESC; +---- +2 16777216 16777216 +1 16777217 16777216 + # Widening cast (Int32 -> Int64) is strictly 1-to-1, so the optimizer CAN # legally reuse the pre-existing sort order and omit a SortExec. statement ok diff --git a/datafusion/sqllogictest/test_files/order.slt b/datafusion/sqllogictest/test_files/order.slt index 3514c30e3cce0..106e50e088050 100644 --- a/datafusion/sqllogictest/test_files/order.slt +++ b/datafusion/sqllogictest/test_files/order.slt @@ -2328,6 +2328,30 @@ physical_plan 02)--SortPreservingMergeExec: [a@0 ASC NULLS LAST], fetch=10 03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/order/spm_orderings/1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/order/spm_orderings/2.csv]]}, projection=[a, b], limit=10, output_orderings=[[a@0 ASC NULLS LAST], [b@1 ASC NULLS LAST]], file_type=csv, has_header=true +# An order-preserving RepartitionExec also only guarantees its merge key +# ordering. Grouping by both columns merges on [a, b], so ORDER BY b still +# needs a sort even though each input file is individually sorted by b. +# https://github.com/apache/datafusion/issues/26745 +statement ok +SET datafusion.optimizer.prefer_existing_sort = true; + +statement ok +SET datafusion.optimizer.repartition_file_scans = false; + +query III +SELECT a, b, count(*) FROM spm_orderings GROUP BY a, b ORDER BY b; +---- +1 1 1 +3 2 1 +1 10 1 +2 20 1 + +statement ok +RESET datafusion.optimizer.prefer_existing_sort; + +statement ok +RESET datafusion.optimizer.repartition_file_scans; + statement ok DROP TABLE spm_orderings; diff --git a/datafusion/sqllogictest/test_files/push_down_filter_regression.slt b/datafusion/sqllogictest/test_files/push_down_filter_regression.slt index f85d568efc789..445e5c81cd51b 100644 --- a/datafusion/sqllogictest/test_files/push_down_filter_regression.slt +++ b/datafusion/sqllogictest/test_files/push_down_filter_regression.slt @@ -598,9 +598,9 @@ EXPLAIN SELECT a, b, cnt FROM ( physical_plan 01)FilterExec: cnt@2 = 2 02)--ProjectionExec: expr=[a@0 as a, b@1 as b, count(agg_filter_pushdown.b)@2 as cnt] -03)----AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[count(agg_filter_pushdown.b)] +03)----AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[count(agg_filter_pushdown.b)], group_clustering_mode=Partial([1]) 04)------RepartitionExec: partitioning=Hash([a@0, b@1], 4), input_partitions=4 -05)--------AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[count(agg_filter_pushdown.b)] +05)--------AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[count(agg_filter_pushdown.b)], group_clustering_mode=Partial([1]) 06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_filter_pushdown.parquet]]}, projection=[a, b], file_type=parquet, predicate=b@1 = CAST(foo AS Utf8View), pruning_predicate=b_null_count@2 != row_count@3 AND b_min@0 <= foo AND foo <= b_max@1, required_guarantees=[] diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 1422875db3594..dc4d3c9bc5564 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -3349,7 +3349,7 @@ logical_plan physical_plan 01)ProjectionExec: expr=[sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.a, annotated_data_infinite2.b] ORDER BY [annotated_data_infinite2.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as sum1, sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.b, annotated_data_infinite2.a] ORDER BY [annotated_data_infinite2.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 as sum2, sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.a, annotated_data_infinite2.d] ORDER BY [annotated_data_infinite2.b ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as sum3, sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.d] ORDER BY [annotated_data_infinite2.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@6 as sum4] 02)--BoundedWindowAggExec: wdw=[sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.d] ORDER BY [annotated_data_infinite2.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.d] ORDER BY [annotated_data_infinite2.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Linear] -03)----RepartitionExec: partitioning=Hash([d@2], 2), input_partitions=2, preserve_order=true, sort_exprs=__common_expr_1@0 ASC NULLS LAST, a@1 ASC NULLS LAST +03)----RepartitionExec: partitioning=Hash([d@2], 2), input_partitions=2, preserve_order=true, sort_exprs=a@1 ASC NULLS LAST 04)------ProjectionExec: expr=[__common_expr_1@0 as __common_expr_1, a@1 as a, d@4 as d, sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.a, annotated_data_infinite2.b] ORDER BY [annotated_data_infinite2.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 as sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.a, annotated_data_infinite2.b] ORDER BY [annotated_data_infinite2.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.a, annotated_data_infinite2.d] ORDER BY [annotated_data_infinite2.b ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@6 as sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.a, annotated_data_infinite2.d] ORDER BY [annotated_data_infinite2.b ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.b, annotated_data_infinite2.a] ORDER BY [annotated_data_infinite2.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@7 as sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.b, annotated_data_infinite2.a] ORDER BY [annotated_data_infinite2.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] 05)--------BoundedWindowAggExec: wdw=[sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.b, annotated_data_infinite2.a] ORDER BY [annotated_data_infinite2.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(annotated_data_infinite2.a) PARTITION BY [annotated_data_infinite2.b, annotated_data_infinite2.a] ORDER BY [annotated_data_infinite2.c ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] 06)----------RepartitionExec: partitioning=Hash([b@2, a@1], 2), input_partitions=2, preserve_order=true, sort_exprs=a@1 ASC NULLS LAST, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST, __common_expr_1@0 ASC NULLS LAST From e12ea26f7b6d757e7dd8a6dc54f29c43f17cfe94 Mon Sep 17 00:00:00 2001 From: Kristofferjoha Date: Sun, 25 Oct 2026 15:50:29 +0100 Subject: [PATCH 2/2] fixed whitespace --- datafusion/physical-expr/src/equivalence/properties/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 6214b838231b6..11b581ee747bb 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -1417,7 +1417,7 @@ impl EquivalenceProperties { if properties.sort_properties != SortProperties::Unordered { return properties; } - if let Some(class) = class { + if let Some(class) = class { for equivalent in class.iter() { let equivalent_properties = self.get_expr_properties_from_tree(Arc::clone(equivalent));