fix: retain only valid orderings after order-preserving repartition

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::<CastExpr>() {
+                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() {

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<dyn PhysicalExpr>) -> 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<dyn PhysicalExpr>) -> 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<dyn PhysicalExpr>,
+    ) -> ExprProperties {
         ExprPropertiesNode::new_unknown(expr)
             .transform_up(|expr| update_properties(expr, self))
             .data()

datafusion/physical-plan/src/repartition/mod.rs

@@ -2200,15 +2200,17 @@ impl RepartitionExec {
         input: &Arc<dyn ExecutionPlan>,
         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

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
 ###

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

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;
 

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=[]
 

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