OptimizerHints having more fields

Author: srh

datafusion/src/physical_plan/coalesce_partitions.rs

@@ -139,10 +139,8 @@ impl ExecutionPlan for CoalescePartitionsExec {
         } else {
             sort_order = None
         }
-        OptimizerHints {
-            sort_order,
-            single_value_columns: input_hints.single_value_columns,
-        }
+        // TODO: This could do approximate sort order, no?
+        OptimizerHints::new_sorted(sort_order, input_hints.single_value_columns)
     }
 
     fn fmt_as(

datafusion/src/physical_plan/filter.rs

@@ -128,10 +128,7 @@ impl ExecutionPlan for FilterExec {
         single_value_columns.sort_unstable();
         single_value_columns.dedup();
 
-        OptimizerHints {
-            sort_order: inputs_hints.sort_order,
-            single_value_columns,
-        }
+        OptimizerHints::new_sorted(inputs_hints.sort_order, single_value_columns)
     }
 
     async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {

datafusion/src/physical_plan/hash_aggregate.rs

@@ -188,15 +188,15 @@ impl HashAggregateExec {
         match strategy {
             AggregateStrategy::Hash => assert!(output_sort_order.is_none()),
             AggregateStrategy::InplaceSorted => {
-                assert!(output_sort_order.is_some());
-                assert!(
-                    output_sort_order
-                        .as_ref()
-                        .unwrap()
-                        .iter()
-                        .all(|i| *i < group_expr.len()),
-                    "sort_order mentions value columns"
-                );
+                // TODO: If we make output hints contain partial sort orders, this will need to assert it's .is_some() and such.
+                if let Some(output_sort_order) = &output_sort_order {
+                    assert!(
+                        output_sort_order
+                            .iter()
+                            .all(|i| *i < group_expr.len()),
+                        "sort_order mentions value columns"
+                    );
+                }
             }
         }
 
@@ -325,10 +325,10 @@ impl ExecutionPlan for HashAggregateExec {
             AggregateStrategy::Hash => None,
             AggregateStrategy::InplaceSorted => self.output_sort_order.clone(),
         };
-        OptimizerHints {
+        OptimizerHints::new_sorted(
             sort_order,
-            single_value_columns: Vec::new(),
-        }
+            Vec::new(),
+        )
     }
 
     fn metrics(&self) -> HashMap<String, SQLMetric> {

datafusion/src/physical_plan/merge.rs

@@ -164,10 +164,10 @@ impl ExecutionPlan for MergeExec {
         } else {
             sort_order = None
         }
-        OptimizerHints {
+        OptimizerHints::new_sorted(
             sort_order,
-            single_value_columns: input_hints.single_value_columns,
-        }
+            input_hints.single_value_columns,
+        )
     }
 }
 

datafusion/src/physical_plan/merge_sort.rs

@@ -101,10 +101,10 @@ impl ExecutionPlan for MergeSortExec {
     }
 
     fn output_hints(&self) -> OptimizerHints {
-        OptimizerHints {
-            single_value_columns: self.input.output_hints().single_value_columns,
-            sort_order: Some(self.columns.iter().map(|c| c.index()).collect()),
-        }
+        OptimizerHints::new_sorted(
+            Some(self.columns.iter().map(|c| c.index()).collect()),
+            self.input.output_hints().single_value_columns,
+        )
     }
 
     async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
@@ -115,6 +115,8 @@ impl ExecutionPlan for MergeSortExec {
             )));
         }
 
+        log::error!("MergeSortExec executing across {} partitions", self.input.output_partitioning().partition_count());
+
         let inputs = join_all(
             (0..self.input.output_partitioning().partition_count())
                 .map(|i| self.input.execute(i))
@@ -614,10 +616,10 @@ impl ExecutionPlan for LastRowByUniqueKeyExec {
     }
 
     fn output_hints(&self) -> OptimizerHints {
-        OptimizerHints {
-            single_value_columns: self.input.output_hints().single_value_columns,
-            sort_order: self.input.output_hints().sort_order,
-        }
+        OptimizerHints::new_sorted(
+            self.input.output_hints().sort_order,
+            self.input.output_hints().single_value_columns,
+        )
     }
 
     async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {

datafusion/src/physical_plan/mod.rs

@@ -129,14 +129,56 @@ use smallvec::SmallVec;
 pub struct OptimizerHints {
     /// If the output is sorted, contains indices of the sort key columns in the output schema.
     /// Each partition should meet this sort order, but order between partitions is unspecified.
-    /// Note that this does **not** guarantee the exact ordering inside each of the columns, e.g.
+    /// Note that this does **not** specify the exact ordering inside each of the columns, e.g.
     /// the values may end up in ascending or descending order, nulls can go first or last.
     pub sort_order: Option<Vec<usize>>,
+
+    // Describes the sawtoothing runs of the stream that is partially sorted.  If sort_order is
+    // present, the first element of this should be sort_order.unwrap().  If we take a sorted stream
+    // and add a projection that removes a column in the middle of sort_order, and it isn't a single
+    // value column, approximate_sort_order.len() would be 2, and it would be the input's sort order
+    // split on the missing column.
+    //
+    // However, this is free to have jumps outside of the sort order.  We might have a MergeNode
+    // which retains the approximate_sort_order optimizer hint despite merging stuff out of order.
+    // The approximate sort order is more "statistical" in nature.
+    pub approximate_sort_order: Vec<Vec<usize>>,
+    /// True if the sort order has no jumps other than those permitted by approximate_sort_order.
+    /// This means that the ordering represents a truly sorted order with some columns missing.
+    pub approximate_sort_order_is_strict: bool,
+    /// True there are no missing columns in front of the approximate sort order.  If and only if
+    /// this and approximate_sort_order_is_strict are true, that implies sort_order should equal
+    /// Some(approximate_sort_order[0]).
+    pub approximate_sort_order_is_prefix: bool,
+
     /// Indices of columns that will always have the same value in each row. No information about
     /// the value is provided.
     pub single_value_columns: Vec<usize>,
 }
 
+impl OptimizerHints {
+    /// Use with None for sort_order is arguably deprecated.  Used to adapt code that preceded
+    /// approximate_sort_order information.
+    fn new_sorted(sort_order: Option<Vec<usize>>, single_value_columns: Vec<usize>) -> OptimizerHints {
+        let mut approximate_sort_order = Vec::new();
+        let mut approximate_sort_order_is_strict = false;
+        let mut approximate_sort_order_is_prefix = false;
+        if let Some(order) = &sort_order {
+            approximate_sort_order.push(order.clone());
+            approximate_sort_order_is_strict = true;
+            approximate_sort_order_is_prefix = true;
+        }
+        let hints = OptimizerHints {
+            sort_order,
+            approximate_sort_order,
+            approximate_sort_order_is_prefix,
+            approximate_sort_order_is_strict,
+            single_value_columns,
+        };
+        hints
+    }
+}
+
 /// `ExecutionPlan` represent nodes in the DataFusion Physical Plan.
 ///
 /// Each `ExecutionPlan` is Partition-aware and is responsible for

datafusion/src/physical_plan/planner.rs

@@ -513,7 +513,7 @@ impl DefaultPhysicalPlanner {
                 //single-value columns
                 let input_sortedness =
                     input_sortedness_by_group_key(input_exec.as_ref(), &groups);
-                let (strategy, order): (AggregateStrategy, Option<Vec<usize>>) =
+                let (strategy, partial_strategy, order): (AggregateStrategy, AggregateStrategy, Option<Vec<usize>>) =
                     match input_sortedness.sawtooth_levels() {
                         Some(0) => {
                             log::error!("DefaultPhysicalExpr: Perfect match for inplace aggregation");
@@ -523,15 +523,17 @@ impl DefaultPhysicalPlanner {
                                     *group_key_offset
                                 })
                                 .collect_vec();
-                            (AggregateStrategy::InplaceSorted, Some(order))
+                            (AggregateStrategy::InplaceSorted, AggregateStrategy::InplaceSorted, Some(order))
                         }
                         Some(n) => {
                             log::error!("DefaultPhysicalExpr: Non-perfect match for inplace aggregation: {} clumps", n);
-                            (AggregateStrategy::Hash, None)
+                            // TODO: Note that this is very oversimplified
+                            (AggregateStrategy::Hash, AggregateStrategy::InplaceSorted, None)
+                            // (AggregateStrategy::Hash, AggregateStrategy::Hash, None)
                         },
                         _ => {
                             log::error!("DefaultPhysicalExpr: No match for inplace aggregation");
-                            (AggregateStrategy::Hash, None)
+                            (AggregateStrategy::Hash, AggregateStrategy::Hash, None)
                         },
                     };
 
@@ -551,7 +553,7 @@ impl DefaultPhysicalPlanner {
 
                 let mut initial_aggr: Arc<dyn ExecutionPlan> =
                     Arc::new(HashAggregateExec::try_new(
-                        strategy,
+                        partial_strategy,
                         order.clone(),
                         AggregateMode::Partial,
                         groups.clone(),
@@ -590,7 +592,8 @@ impl DefaultPhysicalPlanner {
                     && ctx_state.config.concurrency > 1
                     && ctx_state.config.repartition_aggregations
                     && !contains_dict
-                    && strategy == AggregateStrategy::Hash;
+                    && strategy == AggregateStrategy::Hash
+                    && partial_strategy == AggregateStrategy::Hash;
 
                 let (initial_aggr, next_partition_mode): (
                     Arc<dyn ExecutionPlan>,
@@ -1764,6 +1767,8 @@ pub fn input_sortedness_by_group_key(
     }
 
     let hints = input.output_hints();
+    // log::error!("input_sortedness_by_group_key OptimizerHints is: {:?}", hints);
+    // log::error!("input_sortedness_by_group_key input is: {:#?}", input);
     // We check the group key is a prefix of the sort key.
     let sort_key = hints.sort_order;
     if sort_key.is_none() {

datafusion/src/physical_plan/projection.rs

@@ -170,14 +170,14 @@ impl ExecutionPlan for ProjectionExec {
             }
         };
 
-        OptimizerHints {
-            single_value_columns,
-            sort_order: if sort_order.is_empty() {
+        OptimizerHints::new_sorted(
+            if sort_order.is_empty() {
                 None
             } else {
                 Some(sort_order)
             },
-        }
+            single_value_columns,
+        )
     }
 
     fn fmt_as(

datafusion/src/physical_plan/sort.rs

@@ -209,10 +209,10 @@ impl ExecutionPlan for SortExec {
         let input_hints = self.input.output_hints();
         // TODO: If sort_order_truncated is false, we can combine input_hints.sort_order.  Do this.
 
-        OptimizerHints {
-            sort_order: Some(order),
-            single_value_columns: input_hints.single_value_columns.clone(),
-        }
+        OptimizerHints::new_sorted(
+            Some(order),
+            input_hints.single_value_columns,
+        )
     }
 }