WIP: try_switch_to_inplace_aggregates now coarsely supporting inplace partial aggregates

Author: srh

rust/cubestore/cubestore/src/queryplanner/optimizations/prefer_inplace_aggregates.rs

@@ -6,7 +6,7 @@ use datafusion::physical_plan::filter::FilterExec;
 use datafusion::physical_plan::hash_aggregate::{AggregateStrategy, HashAggregateExec};
 use datafusion::physical_plan::merge::MergeExec;
 use datafusion::physical_plan::merge_sort::MergeSortExec;
-use datafusion::physical_plan::planner::input_sortedness_by_group_key;
+use datafusion::physical_plan::planner::{input_sortedness_by_group_key, input_sortedness_by_group_key_using_approximate};
 use datafusion::physical_plan::projection::ProjectionExec;
 use datafusion::physical_plan::union::UnionExec;
 use datafusion::physical_plan::ExecutionPlan;
@@ -29,7 +29,10 @@ pub fn try_switch_to_inplace_aggregates(
     // Try to cheaply rearrange the plan so that it produces sorted inputs.
     let new_input = try_regroup_columns(agg.input().clone())?;
 
-    let input_sortedness = input_sortedness_by_group_key(new_input.as_ref(), agg.group_expr());
+    let input_sortedness = input_sortedness_by_group_key_using_approximate(new_input.as_ref(), agg.group_expr());
+
+    log::error!("try_switch_to_inplace_aggregates: input_sortedness is {:?}\n and... output_hints is {:?}\n and... new_input is {}",
+        input_sortedness, new_input.output_hints(), crate::queryplanner::pretty_printers::pp_phys_plan(new_input.as_ref()));
 
 
     let (strategy, order): (AggregateStrategy, Option<Vec<usize>>) =

rust/cubestore/cubestore/src/queryplanner/query_executor.rs

@@ -838,35 +838,57 @@ impl ExecutionPlan for CubeTableExec {
     }
 
     fn output_hints(&self) -> OptimizerHints {
-        let sort_order;
         if let Some(snapshot_sort_on) = self.index_snapshot.sort_on() {
             // Note that this returns `None` if any of the columns were not found.
             // This only happens on programming errors.
-            sort_order = snapshot_sort_on
+            let sort_order = snapshot_sort_on
                 .iter()
                 .map(|c| self.schema.index_of(&c).ok())
-                .collect()
+                .collect();
+            log::error!("CubeTableExec output_hints returning sort_order {:?}", sort_order);
+            OptimizerHints::new_sorted(sort_order, Vec::new())
         } else {
             let index = self.index_snapshot.index().get_row();
+            let mut approximate_sort_order = Vec::<Vec<usize>>::new();
+            let mut chunk = Vec::<usize>::new();
             let sort_cols = index
                 .get_columns()
                 .iter()
                 .take(index.sort_key_size() as usize)
-                .map(|sort_col| self.schema.index_of(&sort_col.get_name()).ok())
-                .take_while(|i| i.is_some())
-                .map(|i| i.unwrap())
-                .collect_vec();
-            if !sort_cols.is_empty() {
-                sort_order = Some(sort_cols)
+                .map(|sort_col| self.schema.index_of(&sort_col.get_name()).ok());
+            let mut prefix = true;
+            for sort_col in sort_cols {
+                if let Some(col) = sort_col {
+                    chunk.push(col);
+                } else {
+                    if chunk.is_empty() {
+                        prefix &= !approximate_sort_order.is_empty();
+                    } else {
+                        approximate_sort_order.push(chunk);
+                        chunk = Vec::new();
+                    }
+                }
+            }
+
+            if !chunk.is_empty() {
+                approximate_sort_order.push(chunk);
+            }
+
+            let sort_order = if prefix && !approximate_sort_order.is_empty() {
+                Some(approximate_sort_order[0].clone())
             } else {
-                sort_order = None
+                None
+            };
+            log::error!("CubeTableExec output_hints returning approximate sort order {:?} and sort_order {:?}", approximate_sort_order, sort_order);
+
+            OptimizerHints {
+                sort_order,
+                approximate_sort_order,
+                approximate_sort_order_is_strict: true,
+                approximate_sort_order_is_prefix: prefix,
+                single_value_columns: Vec::new(),
             }
         }
-
-        OptimizerHints::new_sorted(
-            sort_order,
-            Vec::new(),
-        )
     }
 
     #[tracing::instrument(level = "trace", skip(self))]