feat: Add GroupsAccumulator and GroupsAccumulatorFlatAdapter, extending API

This is based on upstream GroupsAccumulator and
GroupsAccumulatorAdapter, but extends the API so that existing hash
aggregation works with it.  We basically don't really use the upstream
interface (at this time).

We still use basic Accumulator for types that do not implement
GroupsAccumulator, and hash aggregation code handles this poorly.

Author: srh

datafusion/src/cube_ext/joinagg.rs

@@ -25,7 +25,7 @@ use crate::execution::context::{ExecutionContextState, ExecutionProps};
 use crate::logical_plan::{DFSchemaRef, Expr, LogicalPlan, UserDefinedLogicalNode};
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::utils::from_plan;
-use crate::physical_plan::hash_aggregate::{Accumulators, AggregateMode};
+use crate::physical_plan::hash_aggregate::{create_accumulation_state, AggregateMode};
 use crate::physical_plan::planner::{physical_name, ExtensionPlanner};
 use crate::physical_plan::{hash_aggregate, PhysicalPlanner};
 use crate::physical_plan::{
@@ -245,7 +245,7 @@ impl ExecutionPlan for CrossJoinAggExec {
             &AggregateMode::Full,
             self.group_expr.len(),
         )?;
-        let mut accumulators = Accumulators::new();
+        let mut accumulators = create_accumulation_state(&self.agg_expr)?;
         for partition in 0..self.join.right.output_partitioning().partition_count() {
             let mut batches = self.join.right.execute(partition).await?;
             while let Some(right) = batches.next().await {

datafusion/src/physical_plan/expressions/average.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;
@@ -112,6 +114,23 @@ impl AggregateExpr for Avg {
         )?))
     }
 
+    fn uses_groups_accumulator(&self) -> bool {
+        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>>> {
+        Ok(Some(Box::new(
+            GroupsAccumulatorFlatAdapter::<AvgAccumulator>::new(|| {
+                // avg is f64 (as in create_accumulator)
+                AvgAccumulator::try_new(&DataType::Float64)
+            }),
+        )))
+    }
+
     fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
         vec![self.expr.clone()]
     }

datafusion/src/physical_plan/expressions/sum.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;
@@ -42,7 +44,7 @@ use super::format_state_name;
 use smallvec::smallvec;
 use smallvec::SmallVec;
 
-// SUM aggregate expression
+/// SUM aggregate expression
 #[derive(Debug)]
 pub struct Sum {
     name: String,
@@ -118,6 +120,23 @@ impl AggregateExpr for Sum {
         Ok(Box::new(SumAccumulator::try_new(&self.data_type)?))
     }
 
+    fn uses_groups_accumulator(&self) -> bool {
+        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>>> {
+        let data_type = self.data_type.clone();
+        Ok(Some(Box::new(
+            GroupsAccumulatorFlatAdapter::<SumAccumulator>::new(move || {
+                SumAccumulator::try_new(&data_type)
+            }),
+        )))
+    }
+
     fn name(&self) -> &str {
         &self.name
     }