Review

Author: ozankabak

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -27,36 +27,34 @@ use crate::physical_optimizer::test_utils::{
     union_exec, RequirementsTestExec,
 };
 
-use datafusion_physical_plan::{displayable, InputOrderMode};
 use arrow::compute::SortOptions;
-use arrow::datatypes::SchemaRef;
-use arrow_schema::DataType;
+use arrow::datatypes::{DataType, SchemaRef};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{TreeNode, TransformedResult};
 use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::{JoinType, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition};
+use datafusion_execution::object_store::ObjectStoreUrl;
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
 use datafusion_physical_expr::expressions::{col, Column, NotExpr};
-use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_expr::Partitioning;
-use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
-use datafusion_physical_optimizer::enforce_sorting::{EnforceSorting,PlanWithCorrespondingCoalescePartitions,PlanWithCorrespondingSort,parallelize_sorts,ensure_sorting};
-use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{replace_with_order_preserving_variants,OrderPreservationContext};
-use datafusion_physical_optimizer::enforce_sorting::sort_pushdown::{SortPushDown, assign_initial_requirements, pushdown_sorts};
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
-use datafusion_physical_plan::{get_plan_string, ExecutionPlan};
-use datafusion_common::config::ConfigOptions;
-use datafusion_common::tree_node::{TreeNode, TransformedResult};
+use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::windows::{create_window_expr, BoundedWindowAggExec, WindowAggExec};
+use datafusion_physical_plan::{displayable, get_plan_string, ExecutionPlan, InputOrderMode};
 use datafusion::datasource::physical_plan::{CsvSource, FileScanConfig, ParquetSource};
-use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion::datasource::listing::PartitionedFile;
+use datafusion_physical_optimizer::enforce_sorting::{EnforceSorting, PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, parallelize_sorts, ensure_sorting};
+use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{replace_with_order_preserving_variants, OrderPreservationContext};
+use datafusion_physical_optimizer::enforce_sorting::sort_pushdown::{SortPushDown, assign_initial_requirements, pushdown_sorts};
 use datafusion_physical_optimizer::enforce_distribution::EnforceDistribution;
-use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
-use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_functions_aggregate::average::avg_udaf;
 use datafusion_functions_aggregate::count::count_udaf;
 use datafusion_functions_aggregate::min_max::{max_udaf, min_udaf};
-use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
-use datafusion_physical_plan::windows::{create_window_expr, BoundedWindowAggExec, WindowAggExec};
 
 use rstest::rstest;
 

datafusion/core/tests/physical_optimizer/test_utils.rs

@@ -40,8 +40,9 @@ use datafusion_functions_aggregate::count::count_udaf;
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
 use datafusion_physical_expr::expressions::col;
 use datafusion_physical_expr::{expressions, PhysicalExpr};
-use datafusion_physical_expr_common::sort_expr::LexRequirement;
-use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
+use datafusion_physical_expr_common::sort_expr::{
+    LexOrdering, LexRequirement, PhysicalSortExpr,
+};
 use datafusion_physical_optimizer::limited_distinct_aggregation::LimitedDistinctAggregation;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::aggregates::{
@@ -60,11 +61,10 @@ use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec};
 use datafusion_physical_plan::tree_node::PlanContext;
 use datafusion_physical_plan::union::UnionExec;
 use datafusion_physical_plan::windows::{create_window_expr, BoundedWindowAggExec};
-use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::{
-    displayable, DisplayAs, DisplayFormatType, PlanProperties,
+    displayable, DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode,
+    Partitioning, PlanProperties,
 };
-use datafusion_physical_plan::{InputOrderMode, Partitioning};
 
 /// Create a non sorted parquet exec
 pub fn parquet_exec(schema: &SchemaRef) -> Arc<DataSourceExec> {

datafusion/physical-expr-common/src/sort_expr.rs

@@ -361,6 +361,11 @@ impl LexOrdering {
         self.inner.clear()
     }
 
+    /// Takes ownership of the actual vector of `PhysicalSortExpr`s in the LexOrdering.
+    pub fn take_exprs(self) -> Vec<PhysicalSortExpr> {
+        self.inner
+    }
+
     /// Returns `true` if the LexOrdering contains `expr`
     pub fn contains(&self, expr: &PhysicalSortExpr) -> bool {
         self.inner.contains(expr)

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

@@ -456,17 +456,19 @@ impl EquivalenceGroup {
     /// The expression is replaced with the first expression in the equivalence
     /// class it matches with (if any).
     pub fn normalize_expr(&self, expr: Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> {
-        Arc::clone(&expr)
-            .transform(|expr| {
-                for cls in self.iter() {
-                    if cls.contains(&expr) {
-                        return Ok(Transformed::yes(cls.canonical_expr().unwrap()));
-                    }
+        expr.transform(|expr| {
+            for cls in self.iter() {
+                if cls.contains(&expr) {
+                    // The unwrap below is safe because the guard above ensures
+                    // that the class is not empty.
+                    return Ok(Transformed::yes(cls.canonical_expr().unwrap()));
                 }
-                Ok(Transformed::no(expr))
-            })
-            .data()
-            .unwrap_or(expr)
+            }
+            Ok(Transformed::no(expr))
+        })
+        .data()
+        .unwrap()
+        // The unwrap above is safe because the closure always returns `Ok`.
     }
 
     /// Normalizes the given sort expression according to this group.
@@ -585,20 +587,21 @@ impl EquivalenceGroup {
             (new_class.len() > 1).then_some(EquivalenceClass::new(new_class))
         });
 
-        // the key is the source expression and the value is the EquivalenceClass that contains the target expression of the source expression.
-        let mut new_classes: IndexMap<Arc<dyn PhysicalExpr>, EquivalenceClass> =
-            IndexMap::new();
-        mapping.iter().for_each(|(source, target)| {
-            // We need to find equivalent projected expressions.
-            // e.g. table with columns [a,b,c] and a == b, projection: [a+c, b+c].
-            // To conclude that a + c == b + c we firsty normalize all source expressions
-            // in the mapping, then merge all equivalent expressions into the classes.
+        // The key is the source expression, and the value is the equivalence
+        // class that contains the corresponding target expression.
+        let mut new_classes: IndexMap<_, _> = IndexMap::new();
+        for (source, target) in mapping.iter() {
+            // We need to find equivalent projected expressions. For example,
+            // consider a table with columns `[a, b, c]` with `a` == `b`, and
+            // projection `[a + c, b + c]`. To conclude that `a + c == b + c`,
+            // we first normalize all source expressions in the mapping, then
+            // merge all equivalent expressions into the classes.
             let normalized_expr = self.normalize_expr(Arc::clone(source));
             new_classes
                 .entry(normalized_expr)
                 .or_insert_with(EquivalenceClass::new_empty)
                 .push(Arc::clone(target));
-        });
+        }
         // Only add equivalence classes with at least two members as singleton
         // equivalence classes are meaningless.
         let new_classes = new_classes
@@ -642,7 +645,7 @@ impl EquivalenceGroup {
                 // are equal in the resulting table.
                 if join_type == &JoinType::Inner {
                     for (lhs, rhs) in on.iter() {
-                        let new_lhs = Arc::clone(lhs) as _;
+                        let new_lhs = Arc::clone(lhs);
                         // Rewrite rhs to point to the right side of the join:
                         let new_rhs = Arc::clone(rhs)
                             .transform(|expr| {

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

@@ -237,40 +237,58 @@ impl OrderingEquivalenceClass {
         None
     }
 
-    /// Checks whether the given expression is constant according to `self`'s ordering equivalence class.
+    /// Checks whether the given expression is partially constant according to
+    /// this ordering equivalence class.
     ///
-    /// This function determines whether `expr` appears in at least one combination of `descending`
-    /// and `nulls_first` options that indicate constantness. Specifically, an expression is
-    /// considered constant if it satisfies either of the following conditions:
-    /// - `descending & nulls_first` and `ascending & nulls_last`
-    /// - `descending & nulls_last` and `ascending & nulls_first`
+    /// This function determines whether `expr` appears in at least one combination
+    /// of `descending` and `nulls_first` options that indicate partial constantness
+    /// in a lexicographical ordering. Specifically, an expression is considered
+    /// a partial constant in this context if its `SortOptions` satisfies either
+    /// of the following conditions:
+    /// - It is `descending` with `nulls_first` and _also_ `ascending` with
+    ///   `nulls_last`, OR
+    /// - It is `descending` with `nulls_last` and _also_ `ascending` with
+    ///   `nulls_first`.
     ///
-    /// We primarily use `ConstExpr` to represent globally constant expressions. However, some expressions
-    /// may only be constant within secondary ordering constraints. This function helps identify such cases.
-    /// If an expression is constant within a prefix ordering, it is added as a constant during
-    /// `ordering_satisfy_requirement()` iterations after the corresponding prefix requirement is satisfied.
+    /// The equivalence mechanism primarily uses `ConstExpr`s to represent globally
+    /// constant expressions. However, some expressions may only be partially
+    /// constant within a lexicographical ordering. This function helps identify
+    /// such cases. If an expression is constant within a prefix ordering, it is
+    /// added as a constant during `ordering_satisfy_requirement()` iterations
+    /// after the corresponding prefix requirement is satisfied.
     ///
-    /// ### Example Scenarios (Assuming All Expressions Share the Same Sort Properties)
+    /// ### Example Scenarios
+    ///
+    /// In these scenarios, we assume that all expressions share the same sort
+    /// properties.
     ///
     /// #### Case 1: Sort Requirement `[a, c]`
-    /// - **Existing Orderings:** `[[a, b, c], [a, d]]`, **Constants:** `[]`
-    /// 1. `ordering_satisfy_single()` returns `true` because the requirement `[a]` is satisfied by `[a, b, c].first()`
-    /// 2. `[a]` is added as a constant to the existing orderings
-    /// 3. The normalized orderings become `[[b, c], [d]]`
-    /// 4. `ordering_satisfy_single()` returns `false` for `[c]`, as neither `[b, c]` nor `[d]` satisfies `[c]`
+    ///
+    /// **Existing Orderings:** `[[a, b, c], [a, d]]`, **Constants:** `[]`
+    /// 1. `ordering_satisfy_single()` returns `true` because the requirement
+    ///    `a` is satisfied by `[a, b, c].first()`.
+    /// 2. `a` is added as a constant for the next iteration.
+    /// 3. The normalized orderings become `[[b, c], [d]]`.
+    /// 4. `ordering_satisfy_single()` returns `false` for `c`, as neither
+    ///    `[b, c]` nor `[d]` satisfies `c`.
     ///
     /// #### Case 2: Sort Requirement `[a, d]`
-    /// - **Existing Orderings:** `[[a, b, c], [a, d]]`, **Constants:** `[]`
-    /// 1. `ordering_satisfy_single()` returns `true` because `[a]` is satisfied by `[a, b, c].first()`
-    /// 2. `[a]` is added as a constant to the existing orderings
-    /// 3. The normalized orderings become `[[b, c], [d]]`
-    /// 4. `ordering_satisfy_single()` returns `true` for `[d]`, as `[d]` satisfies `[d]`
+    ///
+    /// **Existing Orderings:** `[[a, b, c], [a, d]]`, **Constants:** `[]`
+    /// 1. `ordering_satisfy_single()` returns `true` because the requirement
+    ///    `a` is satisfied by `[a, b, c].first()`.
+    /// 2. `a` is added as a constant for the next iteration.
+    /// 3. The normalized orderings become `[[b, c], [d]]`.
+    /// 4. `ordering_satisfy_single()` returns `true` for `d`, as `[d]` satisfies
+    ///    `d`.
     ///
     /// ### Future Improvements
-    /// This function may become unnecessary if any of the following improvements are implemented:
+    ///
+    /// This function may become unnecessary if any of the following improvements
+    /// are implemented:
     /// 1. `SortOptions` supports encoding constantness information.
-    /// 2. `EquivalenceProperties` gains `FunctionalDependency` awareness, eliminating the need for
-    ///    `Constant` and `Constraints`.
+    /// 2. `EquivalenceProperties` gains `FunctionalDependency` awareness, eliminating
+    ///    the need for `Constant` and `Constraints`.
     pub fn is_expr_partial_const(&self, expr: &Arc<dyn PhysicalExpr>) -> bool {
         let mut constantness_defining_pairs = [
             HashSet::from([(false, false), (true, true)]),

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

@@ -1464,16 +1464,12 @@ fn update_properties(
     let normalized_expr = eq_properties
         .eq_group
         .normalize_expr(Arc::clone(&node.expr));
+    let oeq_class = eq_properties.normalized_oeq_class();
     if eq_properties.is_expr_constant(&normalized_expr)
-        || eq_properties
-            .normalized_oeq_class()
-            .is_expr_partial_const(&normalized_expr)
+        || oeq_class.is_expr_partial_const(&normalized_expr)
     {
         node.data.sort_properties = SortProperties::Singleton;
-    } else if let Some(options) = eq_properties
-        .normalized_oeq_class()
-        .get_options(&normalized_expr)
-    {
+    } else if let Some(options) = oeq_class.get_options(&normalized_expr) {
         node.data.sort_properties = SortProperties::Ordered(options);
     }
     Ok(Transformed::yes(node))

datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs

@@ -192,13 +192,14 @@ fn pushdown_requirement_to_children(
                 Ok(Some(vec![req]))
             }
             RequirementsCompatibility::Compatible(adjusted) => {
-                // If parent requirements are more specific than output ordering of the window plan,
-                // then we can deduce that parent expects an ordering from the columns constructed
-                // by the window functions. If that's the case, we block the pushdown of sort operation.
-                let avoid_pushdown = !plan
+                // If parent requirements are more specific than output ordering
+                // of the window plan, then we can deduce that the parent expects
+                // an ordering from the columns created by window functions. If
+                // that's the case, we block the pushdown of sort operation.
+                if !plan
                     .equivalence_properties()
-                    .ordering_satisfy_requirement(parent_required);
-                if avoid_pushdown {
+                    .ordering_satisfy_requirement(parent_required)
+                {
                     return Ok(None);
                 }