Encapsulate PhysicalSortRequrement and add more doc comment

alamb · alamb · commit fc774091f362 · 2023-04-04T11:51:15.000-04:00
diff --git a/datafusion/core/src/physical_optimizer/repartition.rs b/datafusion/core/src/physical_optimizer/repartition.rs
@@ -339,9 +339,7 @@ mod tests {
     use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
     use crate::physical_plan::union::UnionExec;
     use crate::physical_plan::{displayable, DisplayFormatType, Statistics};
-    use datafusion_physical_expr::{
-        make_sort_requirements_from_exprs, PhysicalSortRequirement,
-    };
+    use datafusion_physical_expr::PhysicalSortRequirement;
 
     fn schema() -> SchemaRef {
         Arc::new(Schema::new(vec![Field::new("c1", DataType::Boolean, true)]))
@@ -1162,7 +1160,7 @@ mod tests {
         fn required_input_ordering(&self) -> Vec<Option<Vec<PhysicalSortRequirement>>> {
             vec![self
                 .output_ordering()
-                .map(make_sort_requirements_from_exprs)]
+                .map(PhysicalSortRequirement::from_sort_exprs)]
         }
 
         fn with_new_children(
diff --git a/datafusion/core/src/physical_optimizer/sort_enforcement.rs b/datafusion/core/src/physical_optimizer/sort_enforcement.rs
@@ -50,10 +50,9 @@ use arrow::datatypes::SchemaRef;
 use datafusion_common::tree_node::{Transformed, TreeNode, VisitRecursion};
 use datafusion_common::{reverse_sort_options, DataFusionError};
 use datafusion_physical_expr::utils::{
-    make_sort_exprs_from_requirements, ordering_satisfy,
-    ordering_satisfy_requirement_concrete,
+    ordering_satisfy, ordering_satisfy_requirement_concrete,
 };
-use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr};
+use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr, PhysicalSortRequirement};
 use itertools::{concat, izip};
 use std::iter::zip;
 use std::sync::Arc;
@@ -468,7 +467,8 @@ fn ensure_sorting(
                 ) {
                     // Make sure we preserve the ordering requirements:
                     update_child_to_remove_unnecessary_sort(child, sort_onwards, &plan)?;
-                    let sort_expr = make_sort_exprs_from_requirements(&required_ordering);
+                    let sort_expr =
+                        PhysicalSortRequirement::to_sort_exprs(required_ordering);
                     add_sort_above(child, sort_expr)?;
                     if is_sort(child) {
                         *sort_onwards = Some(ExecTree::new(child.clone(), idx, vec![]));
@@ -479,7 +479,7 @@ fn ensure_sorting(
             }
             (Some(required), None) => {
                 // Ordering requirement is not met, we should add a `SortExec` to the plan.
-                let sort_expr = make_sort_exprs_from_requirements(&required);
+                let sort_expr = PhysicalSortRequirement::to_sort_exprs(required);
                 add_sort_above(child, sort_expr)?;
                 *sort_onwards = Some(ExecTree::new(child.clone(), idx, vec![]));
             }
diff --git a/datafusion/core/src/physical_optimizer/sort_pushdown.rs b/datafusion/core/src/physical_optimizer/sort_pushdown.rs
@@ -27,12 +27,9 @@ use datafusion_common::{DataFusionError, Result};
 use datafusion_expr::JoinType;
 use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr::utils::{
-    make_sort_exprs_from_requirements, ordering_satisfy_requirement,
-    requirements_compatible,
-};
-use datafusion_physical_expr::{
-    make_sort_requirements_from_exprs, PhysicalSortExpr, PhysicalSortRequirement,
+    ordering_satisfy_requirement, requirements_compatible,
 };
+use datafusion_physical_expr::{PhysicalSortExpr, PhysicalSortRequirement};
 use itertools::izip;
 use std::ops::Deref;
 use std::sync::Arc;
@@ -130,14 +127,15 @@ pub(crate) fn pushdown_sorts(
             plan.equivalence_properties()
         }) {
             // If the current plan is a SortExec, modify it to satisfy parent requirements:
-            let parent_required_expr =
-                make_sort_exprs_from_requirements(parent_required.ok_or_else(err)?);
+            let parent_required_expr = PhysicalSortRequirement::to_sort_exprs(
+                parent_required.ok_or_else(err)?.iter().cloned(),
+            );
             new_plan = sort_exec.input.clone();
             add_sort_above(&mut new_plan, parent_required_expr)?;
         };
         let required_ordering = new_plan
             .output_ordering()
-            .map(make_sort_requirements_from_exprs);
+            .map(PhysicalSortRequirement::from_sort_exprs);
         // Since new_plan is a SortExec, we can safely get the 0th index.
         let child = &new_plan.children()[0];
         if let Some(adjusted) =
@@ -173,8 +171,9 @@ pub(crate) fn pushdown_sorts(
             }))
         } else {
             // Can not push down requirements, add new SortExec:
-            let parent_required_expr =
-                make_sort_exprs_from_requirements(parent_required.ok_or_else(err)?);
+            let parent_required_expr = PhysicalSortRequirement::to_sort_exprs(
+                parent_required.ok_or_else(err)?.iter().cloned(),
+            );
             let mut new_plan = plan.clone();
             add_sort_above(&mut new_plan, parent_required_expr)?;
             Ok(Transformed::Yes(SortPushDown::init(new_plan)))
@@ -210,8 +209,9 @@ fn pushdown_requirement_to_children(
     } else if let Some(smj) = plan.as_any().downcast_ref::<SortMergeJoinExec>() {
         // If the current plan is SortMergeJoinExec
         let left_columns_len = smj.left.schema().fields().len();
-        let parent_required_expr =
-            make_sort_exprs_from_requirements(parent_required.ok_or_else(err)?);
+        let parent_required_expr = PhysicalSortRequirement::to_sort_exprs(
+            parent_required.ok_or_else(err)?.iter().cloned(),
+        );
         let expr_source_side =
             expr_source_sides(&parent_required_expr, smj.join_type, left_columns_len);
         match expr_source_side {
@@ -383,15 +383,17 @@ fn shift_right_required(
     let new_right_required: Vec<PhysicalSortRequirement> = parent_required
         .iter()
         .filter_map(|r| {
-            r.expr.as_any().downcast_ref::<Column>().and_then(|col| {
-                (col.index() >= left_columns_len).then_some(PhysicalSortRequirement {
-                    expr: Arc::new(Column::new(
-                        col.name(),
-                        col.index() - left_columns_len,
-                    )) as _,
-                    options: r.options,
-                })
-            })
+            let Some(col) = r.expr().as_any().downcast_ref::<Column>() else {
+                return None;
+            };
+
+            if col.index() < left_columns_len {
+                return None;
+            }
+
+            let new_col =
+                Arc::new(Column::new(col.name(), col.index() - left_columns_len));
+            Some(r.clone().with_expr(new_col))
         })
         .collect::<Vec<_>>();
     if new_right_required.len() == parent_required.len() {
diff --git a/datafusion/core/src/physical_plan/joins/sort_merge_join.rs b/datafusion/core/src/physical_plan/joins/sort_merge_join.rs
@@ -35,9 +35,7 @@ use arrow::compute::{concat_batches, take, SortOptions};
 use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
 use arrow::error::ArrowError;
 use arrow::record_batch::RecordBatch;
-use datafusion_physical_expr::{
-    make_sort_requirements_from_exprs, PhysicalSortRequirement,
-};
+use datafusion_physical_expr::PhysicalSortRequirement;
 use futures::{Stream, StreamExt};
 
 use crate::error::DataFusionError;
@@ -230,8 +228,12 @@ impl ExecutionPlan for SortMergeJoinExec {
 
     fn required_input_ordering(&self) -> Vec<Option<Vec<PhysicalSortRequirement>>> {
         vec![
-            Some(make_sort_requirements_from_exprs(&self.left_sort_exprs)),
-            Some(make_sort_requirements_from_exprs(&self.right_sort_exprs)),
+            Some(PhysicalSortRequirement::from_sort_exprs(
+                &self.left_sort_exprs,
+            )),
+            Some(PhysicalSortRequirement::from_sort_exprs(
+                &self.right_sort_exprs,
+            )),
         ]
     }
 
diff --git a/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs b/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs
@@ -46,9 +46,7 @@ use crate::physical_plan::{
     Distribution, ExecutionPlan, Partitioning, PhysicalExpr, RecordBatchStream,
     SendableRecordBatchStream, Statistics,
 };
-use datafusion_physical_expr::{
-    make_sort_requirements_from_exprs, EquivalenceProperties, PhysicalSortRequirement,
-};
+use datafusion_physical_expr::{EquivalenceProperties, PhysicalSortRequirement};
 
 /// Sort preserving merge execution plan
 ///
@@ -128,7 +126,7 @@ impl ExecutionPlan for SortPreservingMergeExec {
     }
 
     fn required_input_ordering(&self) -> Vec<Option<Vec<PhysicalSortRequirement>>> {
-        vec![Some(make_sort_requirements_from_exprs(&self.expr))]
+        vec![Some(PhysicalSortRequirement::from_sort_exprs(&self.expr))]
     }
 
     fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
diff --git a/datafusion/core/src/physical_plan/windows/mod.rs b/datafusion/core/src/physical_plan/windows/mod.rs
@@ -194,18 +194,12 @@ pub(crate) fn calc_requirements(
 ) -> Option<Vec<PhysicalSortRequirement>> {
     let mut sort_reqs = vec![];
     for partition_by in partition_by_exprs {
-        sort_reqs.push(PhysicalSortRequirement {
-            expr: partition_by.clone(),
-            options: None,
-        });
+        sort_reqs.push(PhysicalSortRequirement::new_expr_only(partition_by.clone()))
     }
-    for PhysicalSortExpr { expr, options } in orderby_sort_exprs {
-        let contains = sort_reqs.iter().any(|e| expr.eq(&e.expr));
+    for sort_expr in orderby_sort_exprs {
+        let contains = sort_reqs.iter().any(|e| sort_expr.expr.eq(e.expr()));
         if !contains {
-            sort_reqs.push(PhysicalSortRequirement {
-                expr: expr.clone(),
-                options: Some(*options),
-            });
+            sort_reqs.push(PhysicalSortRequirement::from(sort_expr));
         }
     }
     // Convert empty result to None. Otherwise wrap result inside Some()
@@ -297,7 +291,11 @@ mod tests {
                     nulls_first,
                 });
                 let expr = col(col_name, &schema)?;
-                let res = PhysicalSortRequirement { expr, options };
+                let res = if let Some(options) = options {
+                    PhysicalSortRequirement::new_exact(expr, options)
+                } else {
+                    PhysicalSortRequirement::new_expr_only(expr)
+                };
                 if let Some(expected) = &mut expected {
                     expected.push(res);
                 } else {
diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs
@@ -53,9 +53,7 @@ pub use equivalence::EquivalentClass;
 pub use physical_expr::{AnalysisContext, ExprBoundaries, PhysicalExpr, PhysicalExprRef};
 pub use planner::create_physical_expr;
 pub use scalar_function::ScalarFunctionExpr;
-pub use sort_expr::{
-    make_sort_requirements_from_exprs, PhysicalSortExpr, PhysicalSortRequirement,
-};
+pub use sort_expr::{PhysicalSortExpr, PhysicalSortRequirement};
 pub use utils::{
     expr_list_eq_any_order, expr_list_eq_strict_order,
     normalize_expr_with_equivalence_properties, normalize_out_expr_with_alias_schema,
diff --git a/datafusion/physical-expr/src/sort_expr.rs b/datafusion/physical-expr/src/sort_expr.rs
diff --git a/datafusion/physical-expr/src/utils.rs b/datafusion/physical-expr/src/utils.rs
