Fix order hints of inplace aggregations

They were incorrect when the order of grouping columns differed
from the input columns.

Author: ilya-biryukov

rust/datafusion/src/physical_plan/hash_aggregate.rs

@@ -111,6 +111,7 @@ pub enum AggregateStrategy {
 #[derive(Debug)]
 pub struct HashAggregateExec {
     strategy: AggregateStrategy,
+    output_sort_order: Option<Vec<usize>>,
     /// Aggregation mode (full, partial)
     mode: AggregateMode,
     /// Grouping expressions
@@ -166,6 +167,7 @@ impl HashAggregateExec {
     /// Create a new hash aggregate execution plan
     pub fn try_new(
         strategy: AggregateStrategy,
+        output_sort_order: Option<Vec<usize>>,
         mode: AggregateMode,
         group_expr: Vec<(Arc<dyn PhysicalExpr>, String)>,
         aggr_expr: Vec<Arc<dyn AggregateExpr>>,
@@ -178,8 +180,24 @@ impl HashAggregateExec {
 
         let output_rows = SQLMetric::counter("outputRows");
 
+        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"
+                );
+            }
+        }
+
         Ok(HashAggregateExec {
             strategy,
+            output_sort_order,
             mode,
             group_expr,
             aggr_expr,
@@ -281,6 +299,7 @@ impl ExecutionPlan for HashAggregateExec {
         match children.len() {
             1 => Ok(Arc::new(HashAggregateExec::try_new(
                 self.strategy,
+                self.output_sort_order.clone(),
                 self.mode,
                 self.group_expr.clone(),
                 self.aggr_expr.clone(),
@@ -296,9 +315,7 @@ impl ExecutionPlan for HashAggregateExec {
     fn output_hints(&self) -> OptimizerHints {
         let sort_order = match self.strategy {
             AggregateStrategy::Hash => None,
-            AggregateStrategy::InplaceSorted => {
-                Some((0..self.group_expr.len()).collect_vec())
-            }
+            AggregateStrategy::InplaceSorted => self.output_sort_order.clone(),
         };
         OptimizerHints {
             sort_order,
@@ -1814,6 +1831,7 @@ mod tests {
         let input_schema = input.schema();
         let partial_aggregate = Arc::new(HashAggregateExec::try_new(
             AggregateStrategy::Hash,
+            None,
             AggregateMode::Partial,
             groups.clone(),
             aggregates.clone(),
@@ -1841,6 +1859,7 @@ mod tests {
 
         let merged_aggregate = Arc::new(HashAggregateExec::try_new(
             AggregateStrategy::Hash,
+            None,
             AggregateMode::Final,
             final_group
                 .iter()

rust/datafusion/src/physical_plan/planner.rs

@@ -197,12 +197,14 @@ impl DefaultPhysicalPlanner {
                     })
                     .collect::<Result<Vec<_>>>()?;
 
-                let strategy = compute_aggregation_strategy(input_exec.as_ref(), &groups);
+                let (strategy, order) =
+                    compute_aggregation_strategy(input_exec.as_ref(), &groups);
                 // TODO: fix cubestore planning and re-enable.
                 if false && input_exec.output_partitioning().partition_count() == 1 {
                     // A single pass is enough for 1 partition.
                     return Ok(Arc::new(HashAggregateExec::try_new(
                         strategy,
+                        order,
                         AggregateMode::Full,
                         groups,
                         aggregates,
@@ -214,6 +216,7 @@ impl DefaultPhysicalPlanner {
                 let mut initial_aggr: Arc<dyn ExecutionPlan> =
                     Arc::new(HashAggregateExec::try_new(
                         strategy,
+                        order.clone(),
                         AggregateMode::Partial,
                         groups.clone(),
                         aggregates.clone(),
@@ -238,6 +241,7 @@ impl DefaultPhysicalPlanner {
                 // and the expressions corresponding to the respective aggregate
                 Ok(Arc::new(HashAggregateExec::try_new(
                     strategy,
+                    order,
                     AggregateMode::Final,
                     final_group
                         .iter()
@@ -957,19 +961,25 @@ pub fn evaluate_const(expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExp
 pub fn compute_aggregation_strategy(
     input: &dyn ExecutionPlan,
     group_key: &[(Arc<dyn PhysicalExpr>, String)],
-) -> AggregateStrategy {
-    if !group_key.is_empty() && input_sorted_by_group_key(input, &group_key) {
-        AggregateStrategy::InplaceSorted
+) -> (AggregateStrategy, /*sort_order*/ Option<Vec<usize>>) {
+    let mut sort_order = Vec::new();
+    if !group_key.is_empty()
+        && input_sorted_by_group_key(input, &group_key, &mut sort_order)
+    {
+        (AggregateStrategy::InplaceSorted, Some(sort_order))
     } else {
-        AggregateStrategy::Hash
+        (AggregateStrategy::Hash, None)
     }
 }
 
 fn input_sorted_by_group_key(
     input: &dyn ExecutionPlan,
     group_key: &[(Arc<dyn PhysicalExpr>, String)],
+    sort_order: &mut Vec<usize>,
 ) -> bool {
     assert!(!group_key.is_empty());
+    sort_order.clear();
+
     let hints = input.output_hints();
     // We check the group key is a prefix of the sort key.
     let sort_key = hints.sort_order;
@@ -979,7 +989,8 @@ fn input_sorted_by_group_key(
     let sort_key = sort_key.unwrap();
     // Tracks which elements of sort key are used in the group key or have a single value.
     let mut sort_key_hit = vec![false; sort_key.len()];
-    for (g, _) in group_key {
+    let mut sort_to_group = vec![usize::MAX; sort_key.len()];
+    for (group_i, (g, _)) in group_key.iter().enumerate() {
         let col = g.as_any().downcast_ref::<Column>();
         if col.is_none() {
             return false;
@@ -989,11 +1000,15 @@ fn input_sorted_by_group_key(
             return false;
         }
         let input_col = input_col.unwrap();
-        let sort_key_pos = sort_key.iter().find_position(|i| **i == input_col);
-        if sort_key_pos.is_none() {
-            return false;
+        let sort_key_pos = match sort_key.iter().find_position(|i| **i == input_col) {
+            None => return false,
+            Some((p, _)) => p,
+        };
+        sort_key_hit[sort_key_pos] = true;
+        if sort_to_group[sort_key_pos] != usize::MAX {
+            return false; // Bail out to simplify code a bit. This should not happen in practice.
         }
-        sort_key_hit[sort_key_pos.unwrap().0] = true;
+        sort_to_group[sort_key_pos] = group_i;
     }
     for i in 0..sort_key.len() {
         if hints.single_value_columns.contains(&sort_key[i]) {
@@ -1003,12 +1018,21 @@ fn input_sorted_by_group_key(
 
     // At this point all elements of the group key mapped into some column of the sort key.
     // This checks the group key is mapped into a prefix of the sort key.
-    sort_key_hit
-        .iter()
-        .skip_while(|present| **present)
-        .skip_while(|present| !**present)
-        .next()
-        .is_none()
+    let pref_len = sort_key_hit.iter().take_while(|present| **present).count();
+    if sort_key_hit[pref_len..].iter().any(|present| *present) {
+        return false;
+    }
+
+    assert!(sort_order.is_empty()); // Cleared at the beginning of the function.
+
+    // Note that single-value columns might not have a mapping to the grouping key.
+    sort_order.extend(
+        sort_to_group
+            .iter()
+            .take(pref_len)
+            .filter(|i| **i != usize::MAX),
+    );
+    true
 }
 
 fn tuple_err<T, R>(value: (Result<T>, Result<R>)) -> Result<(T, R)> {