eejbyfeldt commented on code in PR #12095:
URL: https://github.com/apache/datafusion/pull/12095#discussion_r1730273749
##########
datafusion/functions-aggregate/src/variance.rs:
##########
@@ -344,3 +391,181 @@ impl Accumulator for VarianceAccumulator {
true
}
}
+
+#[derive(Debug)]
+pub struct VarianceGroupsAccumulator {
+ m2s: Vec<f64>,
+ means: Vec<f64>,
+ counts: Vec<u64>,
+ stats_type: StatsType,
+}
+
+impl VarianceGroupsAccumulator {
+ pub fn new(s_type: StatsType) -> Self {
+ Self {
+ m2s: Vec::new(),
+ means: Vec::new(),
+ counts: Vec::new(),
+ stats_type: s_type,
+ }
+ }
+
+ fn resize(&mut self, total_num_groups: usize) {
+ self.m2s.resize(total_num_groups, 0.0);
+ self.means.resize(total_num_groups, 0.0);
+ self.counts.resize(total_num_groups, 0);
+ }
+
+ fn merge<F>(
+ group_indices: &[usize],
+ counts: &UInt64Array,
+ means: &Float64Array,
+ m2s: &Float64Array,
+ opt_filter: Option<&BooleanArray>,
+ mut value_fn: F,
+ ) where
+ F: FnMut(usize, u64, f64, f64) + Send,
+ {
+ assert_eq!(counts.null_count(), 0);
+ assert_eq!(means.null_count(), 0);
+ assert_eq!(m2s.null_count(), 0);
+
+ match opt_filter {
+ None => {
+ group_indices
+ .iter()
+ .zip(counts.values().iter())
+ .zip(means.values().iter())
+ .zip(m2s.values().iter())
+ .for_each(|(((&group_index, &count), &mean), &m2)| {
+ value_fn(group_index, count, mean, m2);
+ });
+ }
+ Some(filter) => {
+ group_indices
+ .iter()
+ .zip(counts.values().iter())
+ .zip(means.values().iter())
+ .zip(m2s.values().iter())
+ .zip(filter.iter())
+ .for_each(
+ |((((&group_index, &count), &mean), &m2),
filter_value)| {
+ if let Some(true) = filter_value {
+ value_fn(group_index, count, mean, m2);
+ }
+ },
+ );
+ }
+ }
+ }
+
+ pub fn variance(
+ &mut self,
+ emit_to: datafusion_expr::EmitTo,
+ ) -> (Vec<f64>, NullBuffer) {
+ let mut counts = emit_to.take_needed(&mut self.counts);
+ let _ = emit_to.take_needed(&mut self.means);
+ let m2s = emit_to.take_needed(&mut self.m2s);
+
+ if let StatsType::Sample = self.stats_type {
+ counts.iter_mut().for_each(|count| {
+ *count -= 1;
+ });
+ }
+ let nulls = NullBuffer::from_iter(counts.iter().map(|&count| count !=
0));
+ let variance = m2s
+ .iter()
+ .zip(counts)
+ .map(|(m2, count)| m2 / count as f64)
+ .collect();
+ (variance, nulls)
+ }
+}
+
+impl GroupsAccumulator for VarianceGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&arrow::array::BooleanArray>,
+ total_num_groups: usize,
+ ) -> Result<()> {
+ assert_eq!(values.len(), 1, "single argument to update_batch");
+ let values = &cast(&values[0], &DataType::Float64)?;
+ let values = downcast_value!(values, Float64Array);
+
+ self.resize(total_num_groups);
+ accumulate(group_indices, values, opt_filter, |group_index, value| {
+ let (new_count, new_mean, new_m2) = update(
+ self.counts[group_index],
+ self.means[group_index],
+ self.m2s[group_index],
+ value,
+ );
+ self.counts[group_index] = new_count;
+ self.means[group_index] = new_mean;
+ self.m2s[group_index] = new_m2;
+ });
+ Ok(())
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&arrow::array::BooleanArray>,
+ total_num_groups: usize,
+ ) -> Result<()> {
+ assert_eq!(values.len(), 3, "two arguments to merge_batch");
+ // first batch is counts, second is partial sums
Review Comment:
Fixed/improved in
https://github.com/apache/datafusion/pull/12095/commits/0f1b64ca75898cb0f179b70e1b55f953b306a858
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
