Actually accumulating (badly) into a groups accumulator

Author: srh

datafusion/src/cube_ext/joinagg.rs

@@ -25,7 +25,9 @@ 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, AccumulationState, Accumulators, AggregateMode,
+};
 use crate::physical_plan::planner::{physical_name, ExtensionPlanner};
 use crate::physical_plan::{hash_aggregate, PhysicalPlanner};
 use crate::physical_plan::{
@@ -245,7 +247,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,9 @@ 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_adapter::GroupsAccumulatorAdapter;
+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 +115,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,9 @@ 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_adapter::GroupsAccumulatorAdapter;
+use crate::physical_plan::groups_accumulator_flat_adapter::GroupsAccumulatorFlatAdapter;
 use crate::physical_plan::{Accumulator, AggregateExpr, PhysicalExpr};
 use crate::scalar::ScalarValue;
 use arrow::compute;
@@ -118,6 +121,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
     }

datafusion/src/physical_plan/groups_accumulator.rs

@@ -18,7 +18,9 @@
 //! Vectorized [`GroupsAccumulator`]
 
 use crate::error::{DataFusionError, Result};
+use crate::scalar::ScalarValue;
 use arrow::array::{ArrayRef, BooleanArray};
+use smallvec::SmallVec;
 
 /// From upstream:  This replaces a datafusion_common::{not_impl_err} import.
 macro_rules! not_impl_err {
@@ -27,6 +29,18 @@ macro_rules! not_impl_err {
     };
 }
 
+// TODO: Probably drop the #[macro_export] that was copy/pasted in
+// From upstream datafusion.  We don't pass the backtrace:
+/// Exposes a macro to create `DataFusionError::ArrowError` with optional backtrace
+#[macro_export]
+macro_rules! arrow_datafusion_err {
+    ($ERR:expr) => {
+        DataFusionError::ArrowError(
+            $ERR, /* , Some(DataFusionError::get_back_trace() */
+        )
+    };
+}
+
 /// Describes how many rows should be emitted during grouping.
 #[derive(Debug, Clone, Copy)]
 pub enum EmitTo {
@@ -164,6 +178,10 @@ pub trait GroupsAccumulator: Send {
     /// `n`. See [`EmitTo::First`] for more details.
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef>;
 
+    // TODO: Remove this?
+    /// evaluate for a particular group index.
+    fn peek_evaluate(&self, group_index: usize) -> Result<ScalarValue>;
+
     /// Returns the intermediate aggregate state for this accumulator,
     /// used for multi-phase grouping, resetting its internal state.
     ///
@@ -182,6 +200,10 @@ pub trait GroupsAccumulator: Send {
     /// [`Accumulator::state`]: crate::accumulator::Accumulator::state
     fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>>;
 
+    // TODO: Remove this?
+    /// Looks at the state for a particular group index.
+    fn peek_state(&self, group_index: usize) -> Result<SmallVec<[ScalarValue; 2]>>;
+
     /// Merges intermediate state (the output from [`Self::state`])
     /// into this accumulator's current state.
     ///

datafusion/src/physical_plan/groups_accumulator_adapter.rs

@@ -20,6 +20,7 @@
 
 use std::mem::{size_of, size_of_val};
 
+use crate::arrow_datafusion_err;
 use crate::error::{DataFusionError, Result};
 use crate::physical_plan::groups_accumulator::{EmitTo, GroupsAccumulator};
 use crate::physical_plan::Accumulator;
@@ -31,17 +32,7 @@ use arrow::{
     compute,
     datatypes::UInt32Type,
 };
-
-// From upstream datafusion.  We don't pass the backtrace:
-/// Exposes a macro to create `DataFusionError::ArrowError` with optional backtrace
-#[macro_export]
-macro_rules! arrow_datafusion_err {
-    ($ERR:expr) => {
-        DataFusionError::ArrowError(
-            $ERR, /* , Some(DataFusionError::get_back_trace() */
-        )
-    };
-}
+use smallvec::SmallVec;
 
 /// An adapter that implements [`GroupsAccumulator`] for any [`Accumulator`]
 ///
@@ -102,6 +93,7 @@ pub struct GroupsAccumulatorAdapter {
     /// state for each group, stored in group_index order
     states: Vec<AccumulatorState>,
 
+    // TODO: Code maintaining this is commented.
     /// Current memory usage, in bytes.
     ///
     /// Note this is incrementally updated with deltas to avoid the
@@ -209,7 +201,11 @@ impl GroupsAccumulatorAdapter {
     {
         self.make_accumulators_if_needed(total_num_groups)?;
 
-        assert_eq!(values[0].len(), group_indices.len());
+        assert_eq!(
+            values[0].len(),
+            group_indices.len(),
+            "asserting values[0].len() == group_indices.len()"
+        );
 
         // figure out which input rows correspond to which groups.
         // Note that self.state.indices starts empty for all groups
@@ -348,6 +344,10 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
         result
     }
 
+    fn peek_evaluate(&self, group_index: usize) -> Result<ScalarValue> {
+        self.states[group_index].accumulator.evaluate()
+    }
+
     // filtered_null_mask(opt_filter, &values);
     fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
         let vec_size_pre = self.states.allocated_size();
@@ -384,6 +384,10 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
         Ok(arrays)
     }
 
+    fn peek_state(&self, group_index: usize) -> Result<SmallVec<[ScalarValue; 2]>> {
+        self.states[group_index].accumulator.state()
+    }
+
     fn merge_batch(
         &mut self,
         values: &[ArrayRef],