bvolpato commented on code in PR #24953:
URL: https://github.com/apache/datafusion/pull/24953#discussion_r3998888369
##########
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:
updated to use nonlinear input split across two unequal partitions. both
merge directions are checked independently against direct scalar evaluation,
and the test checks all six state fields.
--
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]
