Support re-use of accumulator states and multiple evaluations

Author: mwylde

datafusion/expr-common/src/accumulator.rs

@@ -58,7 +58,11 @@ pub trait Accumulator: Send + Sync + Debug {
     /// running sum.
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()>;
 
-    /// Returns the final aggregate value, consuming the internal state.
+    /// Returns the final aggregate value, without modifying the internal state.
+    /// This method is safe to call multiple times.
+    fn evaluate(&self) -> Result<ScalarValue>;
+
+    /// Returns the final aggregate value, possibly consuming the internal state.
     ///
     /// For example, the `SUM` accumulator maintains a running sum,
     /// and `evaluate` will produce that running sum as its output.
@@ -69,7 +73,9 @@ pub trait Accumulator: Send + Sync + Debug {
     /// This function gets `&mut self` to allow for the accumulator to build
     /// arrow-compatible internal state that can be returned without copying
     /// when possible (for example distinct strings)
-    fn evaluate(&mut self) -> Result<ScalarValue>;
+    fn evaluate_mut(&mut self) -> Result<ScalarValue> {
+        self.evaluate()
+    }
 
     /// Returns the allocated size required for this accumulator, in
     /// bytes, including `Self`.
@@ -248,7 +254,13 @@ pub trait Accumulator: Send + Sync + Debug {
     ///           group values           group values
     ///           in partition 0         in partition 1
     /// ```
-    fn state(&mut self) -> Result<Vec<ScalarValue>>;
+    fn state_mut(&mut self) -> Result<Vec<ScalarValue>> {
+        self.state()
+    }
+
+    /// Returns the internal state without consuming it; this method is safe to
+    /// call multiple times
+    fn state(&self) -> Result<Vec<ScalarValue>>;
 
     /// Updates the accumulator's state from an `Array` containing one
     /// or more intermediate values.

datafusion/ffi/src/udaf/accumulator.rs

@@ -45,12 +45,12 @@ pub struct FFI_Accumulator {
 
     // Evaluate and return a ScalarValues as protobuf bytes
     pub evaluate:
-        unsafe extern "C" fn(accumulator: &mut Self) -> RResult<RVec<u8>, RString>,
+        unsafe extern "C" fn(accumulator: &Self) -> RResult<RVec<u8>, RString>,
 
     pub size: unsafe extern "C" fn(accumulator: &Self) -> usize,
 
     pub state:
-        unsafe extern "C" fn(accumulator: &mut Self) -> RResult<RVec<RVec<u8>>, RString>,
+        unsafe extern "C" fn(accumulator: &Self) -> RResult<RVec<RVec<u8>>, RString>,
 
     pub merge_batch: unsafe extern "C" fn(
         accumulator: &mut Self,
@@ -109,9 +109,9 @@ unsafe extern "C" fn update_batch_fn_wrapper(
 }
 
 unsafe extern "C" fn evaluate_fn_wrapper(
-    accumulator: &mut FFI_Accumulator,
+    accumulator: &FFI_Accumulator,
 ) -> RResult<RVec<u8>, RString> {
-    let accumulator = accumulator.inner_mut();
+    let accumulator = accumulator.inner();
 
     let scalar_result = rresult_return!(accumulator.evaluate());
     let proto_result: datafusion_proto::protobuf::ScalarValue =
@@ -125,9 +125,9 @@ unsafe extern "C" fn size_fn_wrapper(accumulator: &FFI_Accumulator) -> usize {
 }
 
 unsafe extern "C" fn state_fn_wrapper(
-    accumulator: &mut FFI_Accumulator,
+    accumulator: &FFI_Accumulator,
 ) -> RResult<RVec<RVec<u8>>, RString> {
-    let accumulator = accumulator.inner_mut();
+    let accumulator = accumulator.inner();
 
     let state = rresult_return!(accumulator.state());
     let state = state
@@ -237,10 +237,10 @@ impl Accumulator for ForeignAccumulator {
         }
     }
 
-    fn evaluate(&mut self) -> Result<ScalarValue> {
+    fn evaluate(&self) -> Result<ScalarValue> {
         unsafe {
             let scalar_bytes =
-                df_result!((self.accumulator.evaluate)(&mut self.accumulator))?;
+                df_result!((self.accumulator.evaluate)(&self.accumulator))?;
 
             let proto_scalar =
                 datafusion_proto::protobuf::ScalarValue::decode(scalar_bytes.as_ref())
@@ -254,10 +254,10 @@ impl Accumulator for ForeignAccumulator {
         unsafe { (self.accumulator.size)(&self.accumulator) }
     }
 
-    fn state(&mut self) -> Result<Vec<ScalarValue>> {
+    fn state(&self) -> Result<Vec<ScalarValue>> {
         unsafe {
             let state_protos =
-                df_result!((self.accumulator.state)(&mut self.accumulator))?;
+                df_result!((self.accumulator.state)(&self.accumulator))?;
 
             state_protos
                 .into_iter()

datafusion/ffi/src/udaf/mod.rs

@@ -274,6 +274,7 @@ unsafe extern "C" fn state_fields_fn_wrapper(
         return_field,
         ordering_fields,
         is_distinct,
+        for_sliding: false,
     };
 
     let state_fields = rresult_return!(udaf.state_fields(args));
@@ -682,6 +683,7 @@ mod tests {
             return_field: Field::new("f", DataType::Float64, true).into(),
             ordering_fields: &[Arc::clone(&a_field)],
             is_distinct: false,
+            for_sliding: false,
         })?;
 
         assert_eq!(state_fields.len(), 3);

datafusion/functions-aggregate-common/src/accumulator.rs

@@ -99,6 +99,9 @@ pub struct StateFieldsArgs<'a> {
 
     /// Whether the aggregate function is distinct.
     pub is_distinct: bool,
+
+    /// Returns fields for the sliding accumulator, if one exists
+    pub for_sliding: bool,
 }
 
 impl StateFieldsArgs<'_> {

datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs

@@ -20,7 +20,7 @@
 use arrow::array::{ArrayRef, OffsetSizeTrait};
 use datafusion_common::cast::as_list_array;
 use datafusion_common::utils::SingleRowListArrayBuilder;
-use datafusion_common::ScalarValue;
+use datafusion_common::{exec_err, ScalarValue};
 use datafusion_expr_common::accumulator::Accumulator;
 use datafusion_physical_expr_common::binary_map::{ArrowBytesSet, OutputType};
 use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewSet;
@@ -45,12 +45,16 @@ impl<O: OffsetSizeTrait> BytesDistinctCountAccumulator<O> {
 }
 
 impl<O: OffsetSizeTrait> Accumulator for BytesDistinctCountAccumulator<O> {
-    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+    fn state_mut(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
         let set = self.0.take();
         let arr = set.into_state();
         Ok(vec![SingleRowListArrayBuilder::new(arr).build_list_scalar()])
     }
 
+    fn state(&self) -> datafusion_common::Result<Vec<ScalarValue>> {
+        exec_err!("immutable state not supported for BytesDistinctCount")
+    }
+
     fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
         if values.is_empty() {
             return Ok(());
@@ -80,7 +84,7 @@ impl<O: OffsetSizeTrait> Accumulator for BytesDistinctCountAccumulator<O> {
         })
     }
 
-    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+    fn evaluate(&self) -> datafusion_common::Result<ScalarValue> {
         Ok(ScalarValue::Int64(Some(self.0.non_null_len() as i64)))
     }
 
@@ -104,12 +108,16 @@ impl BytesViewDistinctCountAccumulator {
 }
 
 impl Accumulator for BytesViewDistinctCountAccumulator {
-    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+    fn state_mut(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
         let set = self.0.take();
         let arr = set.into_state();
         Ok(vec![SingleRowListArrayBuilder::new(arr).build_list_scalar()])
     }
 
+    fn state(&self) -> datafusion_common::Result<Vec<ScalarValue>> {
+        exec_err!("immutable state not suported for BytesViewDistinctCount")
+    }
+
     fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
         if values.is_empty() {
             return Ok(());
@@ -139,7 +147,7 @@ impl Accumulator for BytesViewDistinctCountAccumulator {
         })
     }
 
-    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+    fn evaluate(&self) -> datafusion_common::Result<ScalarValue> {
         Ok(ScalarValue::Int64(Some(self.0.non_null_len() as i64)))
     }
 

datafusion/functions-aggregate-common/src/aggregate/count_distinct/dict.rs

@@ -52,15 +52,15 @@ impl Accumulator for DictionaryCountAccumulator {
         self.inner.update_batch(values.as_slice())
     }
 
-    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+    fn evaluate(&self) -> datafusion_common::Result<ScalarValue> {
         self.inner.evaluate()
     }
 
     fn size(&self) -> usize {
         self.inner.size()
     }
 
-    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+    fn state(&self) -> datafusion_common::Result<Vec<ScalarValue>> {
         self.inner.state()
     }
 

datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs

@@ -68,7 +68,7 @@ where
     T: ArrowPrimitiveType + Send + Debug,
     T::Native: Eq + Hash,
 {
-    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+    fn state(&self) -> datafusion_common::Result<Vec<ScalarValue>> {
         let arr = Arc::new(
             PrimitiveArray::<T>::from_iter_values(self.values.iter().cloned())
                 .with_data_type(self.data_type.clone()),
@@ -111,7 +111,7 @@ where
         })
     }
 
-    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+    fn evaluate(&self) -> datafusion_common::Result<ScalarValue> {
         Ok(ScalarValue::Int64(Some(self.values.len() as i64)))
     }
 
@@ -155,7 +155,7 @@ impl<T> Accumulator for FloatDistinctCountAccumulator<T>
 where
     T: ArrowPrimitiveType + Send + Debug,
 {
-    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+    fn state(&self) -> datafusion_common::Result<Vec<ScalarValue>> {
         let arr = Arc::new(PrimitiveArray::<T>::from_iter_values(
             self.values.iter().map(|v| v.0),
         )) as ArrayRef;
@@ -198,7 +198,7 @@ where
         })
     }
 
-    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+    fn evaluate(&self) -> datafusion_common::Result<ScalarValue> {
         Ok(ScalarValue::Int64(Some(self.values.len() as i64)))
     }
 

datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs

@@ -325,7 +325,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
 
         let results: Vec<ScalarValue> = states
             .into_iter()
-            .map(|mut state| {
+            .map(|state| {
                 self.free_allocation(state.size());
                 state.accumulator.evaluate()
             })
@@ -347,7 +347,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
         // which we need to form into columns
         let mut results: Vec<Vec<ScalarValue>> = vec![];
 
-        for mut state in states {
+        for state in states {
             self.free_allocation(state.size());
             let accumulator_state = state.accumulator.state()?;
             results.resize_with(accumulator_state.len(), Vec::new);

datafusion/functions-aggregate/src/approx_distinct.rs

@@ -186,12 +186,12 @@ macro_rules! default_accumulator_impl {
             Ok(())
         }
 
-        fn state(&mut self) -> Result<Vec<ScalarValue>> {
+        fn state(&self) -> Result<Vec<ScalarValue>> {
             let value = ScalarValue::from(&self.hll);
             Ok(vec![value])
         }
 
-        fn evaluate(&mut self) -> Result<ScalarValue> {
+        fn evaluate(&self) -> Result<ScalarValue> {
             Ok(ScalarValue::UInt64(Some(self.hll.count() as u64)))
         }
 

datafusion/functions-aggregate/src/approx_percentile_cont.rs

@@ -457,7 +457,7 @@ impl ApproxPercentileAccumulator {
 }
 
 impl Accumulator for ApproxPercentileAccumulator {
-    fn state(&mut self) -> Result<Vec<ScalarValue>> {
+    fn state(&self) -> Result<Vec<ScalarValue>> {
         Ok(self.digest.to_scalar_state().into_iter().collect())
     }
 
@@ -473,7 +473,7 @@ impl Accumulator for ApproxPercentileAccumulator {
         Ok(())
     }
 
-    fn evaluate(&mut self) -> Result<ScalarValue> {
+    fn evaluate(&self) -> Result<ScalarValue> {
         if self.digest.count() == 0 {
             return ScalarValue::try_from(self.return_type.clone());
         }