kosiew commented on code in PR #24953:
URL: https://github.com/apache/datafusion/pull/24953#discussion_r3978513685
##########
datafusion/functions-aggregate/src/correlation.rs:
##########
@@ -625,33 +632,134 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator {
fn size(&self) -> usize {
self.count.capacity() * size_of::<u64>()
- + self.sum_x.capacity() * size_of::<f64>()
- + self.sum_y.capacity() * size_of::<f64>()
- + self.sum_xy.capacity() * size_of::<f64>()
- + self.sum_xx.capacity() * size_of::<f64>()
- + self.sum_yy.capacity() * size_of::<f64>()
+ + self.mean_x.capacity() * size_of::<f64>()
+ + self.m2_x.capacity() * size_of::<f64>()
+ + self.mean_y.capacity() * size_of::<f64>()
+ + self.m2_y.capacity() * size_of::<f64>()
+ + self.co_moment.capacity() * size_of::<f64>()
}
}
#[cfg(test)]
mod tests {
use super::*;
+ #[test]
+ fn correlation_groups_large_offsets() -> Result<()> {
+ let values: Vec<ArrayRef> = vec![
+ Arc::new(Float64Array::from(vec![
+ 1e9,
+ 1e9,
+ 1e9 + 7.0,
+ 1e9 + 7.0,
+ 1e9 + 15.0,
+ 1e9 + 15.0,
+ ])),
+ Arc::new(Float64Array::from(vec![
+ 2e9,
+ -2e9,
+ 2e9 + 14.0,
+ -2e9 - 14.0,
+ 2e9 + 30.0,
+ -2e9 - 30.0,
+ ])),
+ ];
+ let group_indices = [0, 1, 0, 1, 0, 1];
+ for batch_size in 1..=6 {
+ let mut direct = CorrelationGroupsAccumulator::new();
+ let mut merged = CorrelationGroupsAccumulator::new();
+ let mut converted = CorrelationGroupsAccumulator::new();
+ for start in (0..6).step_by(batch_size) {
+ let len = batch_size.min(6 - start);
+ let batch: Vec<_> = values.iter().map(|a| a.slice(start,
len)).collect();
+ let groups = &group_indices[start..start + len];
+ direct.update_batch(&batch, groups, None, 2)?;
+
+ let mut partial = CorrelationGroupsAccumulator::new();
+ partial.update_batch(&batch, groups, None, 2)?;
+ merged.merge_batch(&partial.state(EmitTo::All)?, &[0, 1], 2)?;
+ converted.merge_batch(
+ &partial.convert_to_state(&batch, None)?,
+ groups,
+ 2,
+ )?;
+ }
+ for mut accumulator in [direct, merged, converted] {
+ let result = accumulator.evaluate(EmitTo::All)?;
+ let result = result.as_primitive::<Float64Type>();
+ assert_eq!(result.null_count(), 0);
+ for (actual, expected) in result.values().iter().zip([1.0,
-1.0]) {
+ assert!((actual - expected).abs() < 1e-12, "{result:?}");
+ }
+ }
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn correlation_groups_merge_large_range() -> Result<()> {
+ let values: Vec<ArrayRef> = vec![
+ Arc::new(Float64Array::from(vec![-8e153, 8e153])),
+ Arc::new(Float64Array::from(vec![-0.5, 0.5])),
+ ];
+ let mut accumulator = CorrelationGroupsAccumulator::new();
+ let states = accumulator.convert_to_state(&values, None)?;
+ accumulator.merge_batch(&states, &[0, 0], 1)?;
+ let result = accumulator.evaluate(EmitTo::All)?;
+ let result = result.as_primitive::<Float64Type>();
+ assert_eq!(result.null_count(), 0);
+ assert!((result.value(0) - 1.0).abs() < 1e-12, "{result:?}");
+ Ok(())
+ }
+
+ #[test]
+ fn correlation_scalar_and_grouped_states_are_compatible() -> Result<()> {
+ let values: Vec<ArrayRef> = vec![
+ Arc::new(Float64Array::from(vec![1e9, 1e9 + 7.0, 1e9 + 15.0])),
Review Comment:
Nice to have coverage for scalar/grouped state interchange here. One
thought: in the second direction, the grouped accumulator gets the same linear
input added again before its state is merged into the scalar accumulator. Since
duplicating that input does not change its correlation, a state layout or
merge-ordering issue could potentially be hidden.
Would it make sense to use a non-linear, multi-partition input and test
scalar to grouped and grouped to scalar parity separately? I think that would
give us a stronger regression test for the newly aligned six-field state layout
and its ordering.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
