alamb commented on a change in pull request #1547:
URL: https://github.com/apache/arrow-datafusion/pull/1547#discussion_r783085457
##########
File path: datafusion/src/physical_plan/expressions/variance.rs
##########
@@ -230,93 +235,186 @@ impl VarianceAccumulator {
self.count
}
- pub fn get_mean(&self) -> ScalarValue {
- self.mean.clone()
+ pub fn get_mean(&self) -> f64 {
+ self.mean
}
- pub fn get_m2(&self) -> ScalarValue {
- self.m2.clone()
+ pub fn get_m2(&self) -> f64 {
+ self.m2
}
}
impl Accumulator for VarianceAccumulator {
fn state(&self) -> Result<Vec<ScalarValue>> {
Ok(vec![
ScalarValue::from(self.count),
- self.mean.clone(),
- self.m2.clone(),
+ ScalarValue::from(self.mean),
+ ScalarValue::from(self.m2),
])
}
+ fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+ let values = &cast(&values[0], &DataType::Float64)?;
+ let arr = values.as_any().downcast_ref::<Float64Array>().unwrap();
+
+ for i in 0..arr.len() {
+ let value = arr.value(i);
+
+ if value == 0_f64 && values.is_null(i) {
+ continue;
+ }
+ let new_count = self.count + 1;
+ let delta1 = value - self.mean;
+ let new_mean = delta1 / new_count as f64 + self.mean;
+ let delta2 = value - new_mean;
+ let new_m2 = self.m2 + delta1 * delta2;
+
+ self.count += 1;
+ self.mean = new_mean;
+ self.m2 = new_m2;
+ }
+
+ Ok(())
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
+ let counts = states[0].as_any().downcast_ref::<UInt64Array>().unwrap();
+ let means = states[1].as_any().downcast_ref::<Float64Array>().unwrap();
+ let m2s = states[2].as_any().downcast_ref::<Float64Array>().unwrap();
+
+ for i in 0..counts.len() {
+ let c = counts.value(i);
+ if c == 0_u64 {
+ continue;
+ }
+ let new_count = self.count + c;
Review comment:
seems reasonable to me
--
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...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
