Use GroupsAccumulator exclusively in grouped hash aggregation (#175)

Makes other AggregateExprs in use GroupsAccumulatorFlatAdapter, and
also uses a GroupsAccumulator implementation that has Box<dyn
Accumulator> inside as a fallback accumulator if some AggregateExpr
implementation does not support that.

This fully removes a batch keys and hash table iteration and brings
that performance benefit from Sum and Avg to other aggregation types.

datafusion/src/cube_ext/joinagg.rs

@@ -259,7 +259,6 @@ impl ExecutionPlan for CrossJoinAggExec {
                         accumulators = hash_aggregate::group_aggregate_batch(
                             &AggregateMode::Full,
                             &group_expr,
-                            &self.agg_expr,
                             joined,
                             std::mem::take(&mut accumulators),
                             &aggs,

datafusion/src/physical_plan/distinct_expressions.rs

@@ -27,6 +27,8 @@ use arrow::datatypes::{DataType, Field};
 
 use crate::error::{DataFusionError, Result};
 use crate::physical_plan::group_scalar::GroupByScalar;
+use crate::physical_plan::groups_accumulator::GroupsAccumulator;
+use crate::physical_plan::groups_accumulator_flat_adapter::GroupsAccumulatorFlatAdapter;
 use crate::physical_plan::{Accumulator, AggregateExpr, PhysicalExpr};
 use crate::scalar::ScalarValue;
 use itertools::Itertools;
@@ -122,6 +124,26 @@ impl AggregateExpr for DistinctCount {
         }))
     }
 
+    fn uses_groups_accumulator(&self) -> bool {
+        return true;
+    }
+
+    fn create_groups_accumulator(
+        &self,
+    ) -> arrow::error::Result<Option<Box<dyn GroupsAccumulator>>> {
+        let state_data_types = self.state_data_types.clone();
+        let count_data_type = self.data_type.clone();
+        Ok(Some(Box::new(GroupsAccumulatorFlatAdapter::<
+            DistinctCountAccumulator,
+        >::new(move || {
+            Ok(DistinctCountAccumulator {
+                values: HashSet::default(),
+                state_data_types: state_data_types.clone(),
+                count_data_type: count_data_type.clone(),
+            })
+        }))))
+    }
+
     fn name(&self) -> &str {
         &self.name
     }

datafusion/src/physical_plan/expressions/average.rs

@@ -118,8 +118,6 @@ impl AggregateExpr for Avg {
         return true;
     }
 
-    /// the groups accumulator used to accumulate values from the expression.  If this returns None,
-    /// create_accumulator must be used.
     fn create_groups_accumulator(
         &self,
     ) -> arrow::error::Result<Option<Box<dyn GroupsAccumulator>>> {

datafusion/src/physical_plan/expressions/count.rs

@@ -21,6 +21,8 @@ use std::any::Any;
 use std::sync::Arc;
 
 use crate::error::Result;
+use crate::physical_plan::groups_accumulator::GroupsAccumulator;
+use crate::physical_plan::groups_accumulator_flat_adapter::GroupsAccumulatorFlatAdapter;
 use crate::physical_plan::{Accumulator, AggregateExpr, PhysicalExpr};
 use crate::scalar::ScalarValue;
 use arrow::compute;
@@ -90,6 +92,20 @@ impl AggregateExpr for Count {
         Ok(Box::new(CountAccumulator::new()))
     }
 
+    fn uses_groups_accumulator(&self) -> bool {
+        return true;
+    }
+
+    fn create_groups_accumulator(
+        &self,
+    ) -> arrow::error::Result<Option<Box<dyn GroupsAccumulator>>> {
+        Ok(Some(Box::new(GroupsAccumulatorFlatAdapter::<
+            CountAccumulator,
+        >::new(move || {
+            Ok(CountAccumulator::new())
+        }))))
+    }
+
     fn name(&self) -> &str {
         &self.name
     }

datafusion/src/physical_plan/expressions/min_max.rs

@@ -22,6 +22,8 @@ use std::convert::TryFrom;
 use std::sync::Arc;
 
 use crate::error::{DataFusionError, Result};
+use crate::physical_plan::groups_accumulator::GroupsAccumulator;
+use crate::physical_plan::groups_accumulator_flat_adapter::GroupsAccumulatorFlatAdapter;
 use crate::physical_plan::{Accumulator, AggregateExpr, PhysicalExpr};
 use crate::scalar::ScalarValue;
 use arrow::compute;
@@ -99,6 +101,21 @@ impl AggregateExpr for Max {
         Ok(Box::new(MaxAccumulator::try_new(&self.data_type)?))
     }
 
+    fn uses_groups_accumulator(&self) -> bool {
+        return true;
+    }
+
+    fn create_groups_accumulator(
+        &self,
+    ) -> arrow::error::Result<Option<Box<dyn GroupsAccumulator>>> {
+        let data_type = self.data_type.clone();
+        Ok(Some(Box::new(
+            GroupsAccumulatorFlatAdapter::<MaxAccumulator>::new(move || {
+                MaxAccumulator::try_new(&data_type)
+            }),
+        )))
+    }
+
     fn name(&self) -> &str {
         &self.name
     }
@@ -523,6 +540,21 @@ impl AggregateExpr for Min {
         Ok(Box::new(MinAccumulator::try_new(&self.data_type)?))
     }
 
+    fn uses_groups_accumulator(&self) -> bool {
+        return true;
+    }
+
+    fn create_groups_accumulator(
+        &self,
+    ) -> arrow::error::Result<Option<Box<dyn GroupsAccumulator>>> {
+        let data_type = self.data_type.clone();
+        Ok(Some(Box::new(
+            GroupsAccumulatorFlatAdapter::<MinAccumulator>::new(move || {
+                MinAccumulator::try_new(&data_type)
+            }),
+        )))
+    }
+
     fn name(&self) -> &str {
         &self.name
     }

datafusion/src/physical_plan/expressions/sum.rs

@@ -124,8 +124,6 @@ impl AggregateExpr for Sum {
         return true;
     }
 
-    /// the groups accumulator used to accumulate values from the expression.  If this returns None,
-    /// create_accumulator must be used.
     fn create_groups_accumulator(
         &self,
     ) -> arrow::error::Result<Option<Box<dyn GroupsAccumulator>>> {