Dandandan commented on code in PR #13581:
URL: https://github.com/apache/datafusion/pull/13581#discussion_r1863046233
##########
datafusion/functions-aggregate/src/correlation.rs:
##########
@@ -263,3 +283,307 @@ impl Accumulator for CorrelationAccumulator {
Ok(())
}
}
+
+#[derive(Default)]
+pub struct CorrelationGroupsAccumulator {
+ // Number of elements for each group
+ // This is also used to track nulls: if a group has 0 valid values
accumulated,
+ // final aggregation result will be null.
+ count: Vec<u64>,
+ // Sum of x values for each group
+ sum_x: Vec<f64>,
+ // Sum of y
+ sum_y: Vec<f64>,
+ // Sum of x*y
+ sum_xy: Vec<f64>,
+ // Sum of x^2
+ sum_xx: Vec<f64>,
+ // Sum of y^2
+ sum_yy: Vec<f64>,
+}
+
+impl CorrelationGroupsAccumulator {
+ pub fn new() -> Self {
+ Default::default()
+ }
+}
+
+/// Specialized version of `accumulate_multiple` for correlation's merge_batch
+///
+/// Note: Arrays in `state_arrays` should not have null values, because they
are all
+/// intermediate states created within the accumulator, instead of inputs from
+/// outside.
+fn accumulate_correlation_states(
+ group_indices: &[usize],
+ state_arrays: (
+ &UInt64Array, // count
+ &Float64Array, // sum_x
+ &Float64Array, // sum_y
+ &Float64Array, // sum_xy
+ &Float64Array, // sum_xx
+ &Float64Array, // sum_yy
+ ),
+ mut value_fn: impl FnMut(usize, u64, &[f64]),
+) {
+ let (counts, sum_x, sum_y, sum_xy, sum_xx, sum_yy) = state_arrays;
+
+ assert_eq!(counts.null_count(), 0);
+ assert_eq!(sum_x.null_count(), 0);
+ assert_eq!(sum_y.null_count(), 0);
+ assert_eq!(sum_xy.null_count(), 0);
+ assert_eq!(sum_xx.null_count(), 0);
+ assert_eq!(sum_yy.null_count(), 0);
+
+ let counts_values = counts.values().as_ref();
+ let sum_x_values = sum_x.values().as_ref();
+ let sum_y_values = sum_y.values().as_ref();
+ let sum_xy_values = sum_xy.values().as_ref();
+ let sum_xx_values = sum_xx.values().as_ref();
+ let sum_yy_values = sum_yy.values().as_ref();
+
+ let mut row = [0.0; 5];
Review Comment:
This doesn't have to be a mutable slice if we initialize it in the loop
based on the values (small win, but might generate better code now or in
future?).
--
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
