feat(cube): Cube Accumulator extensions

Author: srh

datafusion/expr-common/src/accumulator.rs

@@ -18,7 +18,7 @@
 //! Accumulator module contains the trait definition for aggregation function's accumulators.
 
 use arrow::array::ArrayRef;
-use datafusion_common::{internal_err, Result, ScalarValue};
+use datafusion_common::{internal_err, not_impl_err, Result, ScalarValue};
 use std::fmt::Debug;
 
 /// Tracks an aggregate function's state.
@@ -72,6 +72,23 @@ pub trait Accumulator: Send + Sync + Debug {
     /// when possible (for example distinct strings)
     fn evaluate(&mut self) -> Result<ScalarValue>;
 
+    /// Cube: Like evaluate() but doesn't modify the accumulator.
+    fn peek_evaluate(&self) -> Result<ScalarValue> {
+        not_impl_err!("Accumulator::peek_evaluate not implemented for {}", std::any::type_name::<Self>())
+    }
+    /// Cube: Resets the accumulator to its initial (zero-like) state.
+    fn reset(&mut self) -> Result<()> {
+        not_impl_err!("Accumulator::reset not implemented for {}", std::any::type_name::<Self>())
+    }
+    /// Cube: Like state() but doesn't modify the accumulator.
+    fn peek_state(&self) -> Result<Vec<ScalarValue>> {
+        not_impl_err!("Accumulator::peek_state not implemented for {}", std::any::type_name::<Self>())
+    }
+    /// Cube: true if this Accumulator supports these Cube accumulator functions.
+    fn supports_cube_ext(&self) -> bool {
+        false
+    }
+
     /// Returns the allocated size required for this accumulator, in
     /// bytes, including `Self`.
     ///

datafusion/functions-aggregate/src/min_max.rs

@@ -846,6 +846,21 @@ impl Accumulator for MaxAccumulator {
     fn size(&self) -> usize {
         std::mem::size_of_val(self) - std::mem::size_of_val(&self.max) + self.max.size()
     }
+
+    // Cube exts:
+    fn peek_evaluate(&self) -> Result<ScalarValue> {
+        Ok(self.max.clone())
+    }
+    fn reset(&mut self) -> Result<()> {
+        self.max = ScalarValue::try_from(self.max.data_type())?;
+        Ok(())
+    }
+    fn peek_state(&self) -> Result<Vec<ScalarValue>> {
+        Ok(vec![self.peek_evaluate()?])
+    }
+    fn supports_cube_ext(&self) -> bool {
+        true
+    }
 }
 
 #[derive(Debug)]
@@ -1104,6 +1119,21 @@ impl Accumulator for MinAccumulator {
     fn size(&self) -> usize {
         std::mem::size_of_val(self) - std::mem::size_of_val(&self.min) + self.min.size()
     }
+
+    // Cube exts:
+    fn peek_evaluate(&self) -> Result<ScalarValue> {
+        Ok(self.min.clone())
+    }
+    fn reset(&mut self) -> Result<()> {
+        self.min = ScalarValue::try_from(self.min.data_type())?;
+        Ok(())
+    }
+    fn peek_state(&self) -> Result<Vec<ScalarValue>> {
+        Ok(vec![self.peek_evaluate()?])
+    }
+    fn supports_cube_ext(&self) -> bool {
+        true
+    }
 }
 
 #[derive(Debug)]

datafusion/functions-aggregate/src/sum.rs

@@ -281,6 +281,21 @@ impl<T: ArrowNumericType> Accumulator for SumAccumulator<T> {
     fn size(&self) -> usize {
         std::mem::size_of_val(self)
     }
+
+    // Cube exts:
+    fn peek_evaluate(&self) -> Result<ScalarValue> {
+        ScalarValue::new_primitive::<T>(self.sum, &self.data_type)
+    }
+    fn reset(&mut self) -> Result<()> {
+        self.sum = None;
+        Ok(())
+    }
+    fn peek_state(&self) -> Result<Vec<ScalarValue>> {
+        Ok(vec![self.peek_evaluate()?])
+    }
+    fn supports_cube_ext(&self) -> bool {
+        true
+    }
 }
 
 /// This accumulator incrementally computes sums over a sliding window