This is an automated email from the ASF dual-hosted git repository.
jayzhan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new d32747d09a Convert Correlation to UDAF (#11064)
d32747d09a is described below
commit d32747d09add5e1c670aa32fbe3294ecee15e3b7
Author: Kevin Su <[email protected]>
AuthorDate: Sun Jun 23 19:30:27 2024 +0800
Convert Correlation to UDAF (#11064)
* init
Signed-off-by: Kevin Su <[email protected]>
* test
Signed-off-by: Kevin Su <[email protected]>
* test
Signed-off-by: Kevin Su <[email protected]>
* test
Signed-off-by: Kevin Su <[email protected]>
* remove files
Signed-off-by: Kevin Su <[email protected]>
---------
Signed-off-by: Kevin Su <[email protected]>
Co-authored-by: Andrew Lamb <[email protected]>
---
datafusion/expr/src/aggregate_function.rs | 11 -
datafusion/expr/src/type_coercion/aggregates.rs | 11 -
datafusion/functions-aggregate/src/correlation.rs | 225 +++++++++
datafusion/functions-aggregate/src/lib.rs | 8 +-
datafusion/physical-expr/src/aggregate/build_in.rs | 11 -
.../physical-expr/src/aggregate/correlation.rs | 524 ---------------------
.../physical-expr/src/aggregate/covariance.rs | 227 ---------
datafusion/physical-expr/src/aggregate/mod.rs | 4 -
datafusion/physical-expr/src/aggregate/stddev.rs | 87 ----
datafusion/physical-expr/src/aggregate/variance.rs | 176 -------
datafusion/physical-expr/src/expressions/mod.rs | 1 -
datafusion/proto/proto/datafusion.proto | 2 +-
datafusion/proto/src/generated/pbjson.rs | 3 -
datafusion/proto/src/generated/prost.rs | 4 +-
datafusion/proto/src/logical_plan/from_proto.rs | 1 -
datafusion/proto/src/logical_plan/to_proto.rs | 4 -
datafusion/proto/src/physical_plan/to_proto.rs | 10 +-
.../proto/tests/cases/roundtrip_logical_plan.rs | 4 +-
18 files changed, 240 insertions(+), 1073 deletions(-)
diff --git a/datafusion/expr/src/aggregate_function.rs
b/datafusion/expr/src/aggregate_function.rs
index 7663dcd610..2997aef58e 100644
--- a/datafusion/expr/src/aggregate_function.rs
+++ b/datafusion/expr/src/aggregate_function.rs
@@ -41,8 +41,6 @@ pub enum AggregateFunction {
ArrayAgg,
/// N'th value in a group according to some ordering
NthValue,
- /// Correlation
- Correlation,
/// Grouping
Grouping,
}
@@ -55,7 +53,6 @@ impl AggregateFunction {
Max => "MAX",
ArrayAgg => "ARRAY_AGG",
NthValue => "NTH_VALUE",
- Correlation => "CORR",
Grouping => "GROUPING",
}
}
@@ -76,8 +73,6 @@ impl FromStr for AggregateFunction {
"min" => AggregateFunction::Min,
"array_agg" => AggregateFunction::ArrayAgg,
"nth_value" => AggregateFunction::NthValue,
- // statistical
- "corr" => AggregateFunction::Correlation,
// other
"grouping" => AggregateFunction::Grouping,
_ => {
@@ -115,9 +110,6 @@ impl AggregateFunction {
// The coerced_data_types is same with input_types.
Ok(coerced_data_types[0].clone())
}
- AggregateFunction::Correlation => {
- correlation_return_type(&coerced_data_types[0])
- }
AggregateFunction::ArrayAgg =>
Ok(DataType::List(Arc::new(Field::new(
"item",
coerced_data_types[0].clone(),
@@ -150,9 +142,6 @@ impl AggregateFunction {
Signature::uniform(1, valid, Volatility::Immutable)
}
AggregateFunction::NthValue => Signature::any(2,
Volatility::Immutable),
- AggregateFunction::Correlation => {
- Signature::uniform(2, NUMERICS.to_vec(), Volatility::Immutable)
- }
}
}
}
diff --git a/datafusion/expr/src/type_coercion/aggregates.rs
b/datafusion/expr/src/type_coercion/aggregates.rs
index 7448af2502..36a789d5b0 100644
--- a/datafusion/expr/src/type_coercion/aggregates.rs
+++ b/datafusion/expr/src/type_coercion/aggregates.rs
@@ -91,7 +91,6 @@ pub fn coerce_types(
input_types: &[DataType],
signature: &Signature,
) -> Result<Vec<DataType>> {
- use DataType::*;
// Validate input_types matches (at least one of) the func signature.
check_arg_count(agg_fun.name(), input_types, &signature.type_signature)?;
@@ -102,16 +101,6 @@ pub fn coerce_types(
// unpack the dictionary to get the value
get_min_max_result_type(input_types)
}
- AggregateFunction::Correlation => {
- if !is_correlation_support_arg_type(&input_types[0]) {
- return plan_err!(
- "The function {:?} does not support inputs of type {:?}.",
- agg_fun,
- input_types[0]
- );
- }
- Ok(vec![Float64, Float64])
- }
AggregateFunction::NthValue => Ok(input_types.to_vec()),
AggregateFunction::Grouping => Ok(vec![input_types[0].clone()]),
}
diff --git a/datafusion/functions-aggregate/src/correlation.rs
b/datafusion/functions-aggregate/src/correlation.rs
new file mode 100644
index 0000000000..10d5563086
--- /dev/null
+++ b/datafusion/functions-aggregate/src/correlation.rs
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`Correlation`]: correlation sample aggregations.
+
+use std::any::Any;
+use std::fmt::Debug;
+
+use arrow::compute::{and, filter, is_not_null};
+use arrow::{
+ array::ArrayRef,
+ datatypes::{DataType, Field},
+};
+
+use crate::covariance::CovarianceAccumulator;
+use crate::stddev::StddevAccumulator;
+use datafusion_common::{plan_err, Result, ScalarValue};
+use datafusion_expr::{
+ function::{AccumulatorArgs, StateFieldsArgs},
+ type_coercion::aggregates::NUMERICS,
+ utils::format_state_name,
+ Accumulator, AggregateUDFImpl, Signature, Volatility,
+};
+use datafusion_physical_expr_common::aggregate::stats::StatsType;
+
+make_udaf_expr_and_func!(
+ Correlation,
+ corr,
+ y x,
+ "Correlation between two numeric values.",
+ corr_udaf
+);
+
+#[derive(Debug)]
+pub struct Correlation {
+ signature: Signature,
+}
+
+impl Default for Correlation {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl Correlation {
+ /// Create a new COVAR_POP aggregate function
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::uniform(2, NUMERICS.to_vec(),
Volatility::Immutable),
+ }
+ }
+}
+
+impl AggregateUDFImpl for Correlation {
+ /// Return a reference to Any that can be used for downcasting
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "corr"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if !arg_types[0].is_numeric() {
+ return plan_err!("Correlation requires numeric input types");
+ }
+
+ Ok(DataType::Float64)
+ }
+
+ fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn
Accumulator>> {
+ Ok(Box::new(CorrelationAccumulator::try_new()?))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
+ let name = args.name;
+ Ok(vec![
+ Field::new(format_state_name(name, "count"), DataType::UInt64,
true),
+ Field::new(format_state_name(name, "mean1"), DataType::Float64,
true),
+ Field::new(format_state_name(name, "m2_1"), DataType::Float64,
true),
+ Field::new(format_state_name(name, "mean2"), DataType::Float64,
true),
+ Field::new(format_state_name(name, "m2_2"), DataType::Float64,
true),
+ Field::new(
+ format_state_name(name, "algo_const"),
+ DataType::Float64,
+ true,
+ ),
+ ])
+ }
+}
+
+/// An accumulator to compute correlation
+#[derive(Debug)]
+pub struct CorrelationAccumulator {
+ covar: CovarianceAccumulator,
+ stddev1: StddevAccumulator,
+ stddev2: StddevAccumulator,
+}
+
+impl CorrelationAccumulator {
+ /// Creates a new `CorrelationAccumulator`
+ pub fn try_new() -> Result<Self> {
+ Ok(Self {
+ covar: CovarianceAccumulator::try_new(StatsType::Population)?,
+ stddev1: StddevAccumulator::try_new(StatsType::Population)?,
+ stddev2: StddevAccumulator::try_new(StatsType::Population)?,
+ })
+ }
+}
+
+impl Accumulator for CorrelationAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+ // TODO: null input skipping logic duplicated across Correlation
+ // and its children accumulators.
+ // This could be simplified by splitting up input filtering and
+ // calculation logic in children accumulators, and calling only
+ // calculation part from Correlation
+ let values = if values[0].null_count() != 0 || values[1].null_count()
!= 0 {
+ let mask = and(&is_not_null(&values[0])?,
&is_not_null(&values[1])?)?;
+ let values1 = filter(&values[0], &mask)?;
+ let values2 = filter(&values[1], &mask)?;
+
+ vec![values1, values2]
+ } else {
+ values.to_vec()
+ };
+
+ self.covar.update_batch(&values)?;
+ self.stddev1.update_batch(&values[0..1])?;
+ self.stddev2.update_batch(&values[1..2])?;
+ Ok(())
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ let covar = self.covar.evaluate()?;
+ let stddev1 = self.stddev1.evaluate()?;
+ let stddev2 = self.stddev2.evaluate()?;
+
+ if let ScalarValue::Float64(Some(c)) = covar {
+ if let ScalarValue::Float64(Some(s1)) = stddev1 {
+ if let ScalarValue::Float64(Some(s2)) = stddev2 {
+ if s1 == 0_f64 || s2 == 0_f64 {
+ return Ok(ScalarValue::Float64(Some(0_f64)));
+ } else {
+ return Ok(ScalarValue::Float64(Some(c / s1 / s2)));
+ }
+ }
+ }
+ }
+
+ Ok(ScalarValue::Float64(None))
+ }
+
+ fn size(&self) -> usize {
+ std::mem::size_of_val(self) - std::mem::size_of_val(&self.covar)
+ + self.covar.size()
+ - std::mem::size_of_val(&self.stddev1)
+ + self.stddev1.size()
+ - std::mem::size_of_val(&self.stddev2)
+ + self.stddev2.size()
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ Ok(vec![
+ ScalarValue::from(self.covar.get_count()),
+ ScalarValue::from(self.covar.get_mean1()),
+ ScalarValue::from(self.stddev1.get_m2()),
+ ScalarValue::from(self.covar.get_mean2()),
+ ScalarValue::from(self.stddev2.get_m2()),
+ ScalarValue::from(self.covar.get_algo_const()),
+ ])
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
+ let states_c = [
+ states[0].clone(),
+ states[1].clone(),
+ states[3].clone(),
+ states[5].clone(),
+ ];
+ let states_s1 = [states[0].clone(), states[1].clone(),
states[2].clone()];
+ let states_s2 = [states[0].clone(), states[3].clone(),
states[4].clone()];
+
+ self.covar.merge_batch(&states_c)?;
+ self.stddev1.merge_batch(&states_s1)?;
+ self.stddev2.merge_batch(&states_s2)?;
+ Ok(())
+ }
+
+ fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+ let values = if values[0].null_count() != 0 || values[1].null_count()
!= 0 {
+ let mask = and(&is_not_null(&values[0])?,
&is_not_null(&values[1])?)?;
+ let values1 = filter(&values[0], &mask)?;
+ let values2 = filter(&values[1], &mask)?;
+
+ vec![values1, values2]
+ } else {
+ values.to_vec()
+ };
+
+ self.covar.retract_batch(&values)?;
+ self.stddev1.retract_batch(&values[0..1])?;
+ self.stddev2.retract_batch(&values[1..2])?;
+ Ok(())
+ }
+}
diff --git a/datafusion/functions-aggregate/src/lib.rs
b/datafusion/functions-aggregate/src/lib.rs
index 6b38c4ed61..063e6000b4 100644
--- a/datafusion/functions-aggregate/src/lib.rs
+++ b/datafusion/functions-aggregate/src/lib.rs
@@ -56,6 +56,7 @@
pub mod macros;
pub mod approx_distinct;
+pub mod correlation;
pub mod count;
pub mod covariance;
pub mod first_last;
@@ -73,6 +74,7 @@ pub mod average;
pub mod bit_and_or_xor;
pub mod bool_and_or;
pub mod string_agg;
+
use crate::approx_percentile_cont::approx_percentile_cont_udaf;
use
crate::approx_percentile_cont_with_weight::approx_percentile_cont_with_weight_udaf;
use datafusion_common::Result;
@@ -93,6 +95,7 @@ pub mod expr_fn {
pub use super::bit_and_or_xor::bit_xor;
pub use super::bool_and_or::bool_and;
pub use super::bool_and_or::bool_or;
+ pub use super::correlation::corr;
pub use super::count::count;
pub use super::count::count_distinct;
pub use super::covariance::covar_pop;
@@ -122,8 +125,9 @@ pub fn all_default_aggregate_functions() ->
Vec<Arc<AggregateUDF>> {
first_last::first_value_udaf(),
first_last::last_value_udaf(),
covariance::covar_samp_udaf(),
- sum::sum_udaf(),
covariance::covar_pop_udaf(),
+ correlation::corr_udaf(),
+ sum::sum_udaf(),
median::median_udaf(),
count::count_udaf(),
regr::regr_slope_udaf(),
@@ -179,7 +183,7 @@ mod tests {
let mut names = HashSet::new();
for func in all_default_aggregate_functions() {
// TODO: remove this
- // These functions are in intermidiate migration state, skip them
+ // These functions are in intermediate migration state, skip them
if func.name().to_lowercase() == "count" {
continue;
}
diff --git a/datafusion/physical-expr/src/aggregate/build_in.rs
b/datafusion/physical-expr/src/aggregate/build_in.rs
index 2b80af9bf3..19b54daa2e 100644
--- a/datafusion/physical-expr/src/aggregate/build_in.rs
+++ b/datafusion/physical-expr/src/aggregate/build_in.rs
@@ -107,17 +107,6 @@ pub fn create_aggregate_expr(
name,
data_type,
)),
- (AggregateFunction::Correlation, false) => {
- Arc::new(expressions::Correlation::new(
- input_phy_exprs[0].clone(),
- input_phy_exprs[1].clone(),
- name,
- data_type,
- ))
- }
- (AggregateFunction::Correlation, true) => {
- return not_impl_err!("CORR(DISTINCT) aggregations are not
available");
- }
(AggregateFunction::NthValue, _) => {
let expr = &input_phy_exprs[0];
let Some(n) = input_phy_exprs[1]
diff --git a/datafusion/physical-expr/src/aggregate/correlation.rs
b/datafusion/physical-expr/src/aggregate/correlation.rs
deleted file mode 100644
index a47d350532..0000000000
--- a/datafusion/physical-expr/src/aggregate/correlation.rs
+++ /dev/null
@@ -1,524 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines physical expressions that can evaluated at runtime during query
execution
-
-use crate::aggregate::covariance::CovarianceAccumulator;
-use crate::aggregate::stats::StatsType;
-use crate::aggregate::stddev::StddevAccumulator;
-use crate::aggregate::utils::down_cast_any_ref;
-use crate::expressions::format_state_name;
-use crate::{AggregateExpr, PhysicalExpr};
-use arrow::{
- array::ArrayRef,
- compute::{and, filter, is_not_null},
- datatypes::{DataType, Field},
-};
-use datafusion_common::Result;
-use datafusion_common::ScalarValue;
-use datafusion_expr::Accumulator;
-use std::any::Any;
-use std::sync::Arc;
-
-/// CORR aggregate expression
-#[derive(Debug)]
-pub struct Correlation {
- name: String,
- expr1: Arc<dyn PhysicalExpr>,
- expr2: Arc<dyn PhysicalExpr>,
-}
-
-impl Correlation {
- /// Create a new COVAR_POP aggregate function
- pub fn new(
- expr1: Arc<dyn PhysicalExpr>,
- expr2: Arc<dyn PhysicalExpr>,
- name: impl Into<String>,
- data_type: DataType,
- ) -> Self {
- // the result of correlation just support FLOAT64 data type.
- assert!(matches!(data_type, DataType::Float64));
- Self {
- name: name.into(),
- expr1,
- expr2,
- }
- }
-}
-
-impl AggregateExpr for Correlation {
- /// Return a reference to Any that can be used for downcasting
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- fn field(&self) -> Result<Field> {
- Ok(Field::new(&self.name, DataType::Float64, true))
- }
-
- fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
- Ok(Box::new(CorrelationAccumulator::try_new()?))
- }
-
- fn state_fields(&self) -> Result<Vec<Field>> {
- Ok(vec![
- Field::new(
- format_state_name(&self.name, "count"),
- DataType::UInt64,
- true,
- ),
- Field::new(
- format_state_name(&self.name, "mean1"),
- DataType::Float64,
- true,
- ),
- Field::new(
- format_state_name(&self.name, "m2_1"),
- DataType::Float64,
- true,
- ),
- Field::new(
- format_state_name(&self.name, "mean2"),
- DataType::Float64,
- true,
- ),
- Field::new(
- format_state_name(&self.name, "m2_2"),
- DataType::Float64,
- true,
- ),
- Field::new(
- format_state_name(&self.name, "algo_const"),
- DataType::Float64,
- true,
- ),
- ])
- }
-
- fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
- vec![self.expr1.clone(), self.expr2.clone()]
- }
-
- fn name(&self) -> &str {
- &self.name
- }
-}
-
-impl PartialEq<dyn Any> for Correlation {
- fn eq(&self, other: &dyn Any) -> bool {
- down_cast_any_ref(other)
- .downcast_ref::<Self>()
- .map(|x| {
- self.name == x.name && self.expr1.eq(&x.expr1) &&
self.expr2.eq(&x.expr2)
- })
- .unwrap_or(false)
- }
-}
-
-/// An accumulator to compute correlation
-#[derive(Debug)]
-pub struct CorrelationAccumulator {
- covar: CovarianceAccumulator,
- stddev1: StddevAccumulator,
- stddev2: StddevAccumulator,
-}
-
-impl CorrelationAccumulator {
- /// Creates a new `CorrelationAccumulator`
- pub fn try_new() -> Result<Self> {
- Ok(Self {
- covar: CovarianceAccumulator::try_new(StatsType::Population)?,
- stddev1: StddevAccumulator::try_new(StatsType::Population)?,
- stddev2: StddevAccumulator::try_new(StatsType::Population)?,
- })
- }
-}
-
-impl Accumulator for CorrelationAccumulator {
- fn state(&mut self) -> Result<Vec<ScalarValue>> {
- Ok(vec![
- ScalarValue::from(self.covar.get_count()),
- ScalarValue::from(self.covar.get_mean1()),
- ScalarValue::from(self.stddev1.get_m2()),
- ScalarValue::from(self.covar.get_mean2()),
- ScalarValue::from(self.stddev2.get_m2()),
- ScalarValue::from(self.covar.get_algo_const()),
- ])
- }
-
- fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- // TODO: null input skipping logic duplicated across Correlation
- // and its children accumulators.
- // This could be simplified by splitting up input filtering and
- // calculation logic in children accumulators, and calling only
- // calculation part from Correlation
- let values = if values[0].null_count() != 0 || values[1].null_count()
!= 0 {
- let mask = and(&is_not_null(&values[0])?,
&is_not_null(&values[1])?)?;
- let values1 = filter(&values[0], &mask)?;
- let values2 = filter(&values[1], &mask)?;
-
- vec![values1, values2]
- } else {
- values.to_vec()
- };
-
- self.covar.update_batch(&values)?;
- self.stddev1.update_batch(&values[0..1])?;
- self.stddev2.update_batch(&values[1..2])?;
- Ok(())
- }
-
- fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- let values = if values[0].null_count() != 0 || values[1].null_count()
!= 0 {
- let mask = and(&is_not_null(&values[0])?,
&is_not_null(&values[1])?)?;
- let values1 = filter(&values[0], &mask)?;
- let values2 = filter(&values[1], &mask)?;
-
- vec![values1, values2]
- } else {
- values.to_vec()
- };
-
- self.covar.retract_batch(&values)?;
- self.stddev1.retract_batch(&values[0..1])?;
- self.stddev2.retract_batch(&values[1..2])?;
- Ok(())
- }
-
- fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
- let states_c = [
- states[0].clone(),
- states[1].clone(),
- states[3].clone(),
- states[5].clone(),
- ];
- let states_s1 = [states[0].clone(), states[1].clone(),
states[2].clone()];
- let states_s2 = [states[0].clone(), states[3].clone(),
states[4].clone()];
-
- self.covar.merge_batch(&states_c)?;
- self.stddev1.merge_batch(&states_s1)?;
- self.stddev2.merge_batch(&states_s2)?;
- Ok(())
- }
-
- fn evaluate(&mut self) -> Result<ScalarValue> {
- let covar = self.covar.evaluate()?;
- let stddev1 = self.stddev1.evaluate()?;
- let stddev2 = self.stddev2.evaluate()?;
-
- if let ScalarValue::Float64(Some(c)) = covar {
- if let ScalarValue::Float64(Some(s1)) = stddev1 {
- if let ScalarValue::Float64(Some(s2)) = stddev2 {
- if s1 == 0_f64 || s2 == 0_f64 {
- return Ok(ScalarValue::Float64(Some(0_f64)));
- } else {
- return Ok(ScalarValue::Float64(Some(c / s1 / s2)));
- }
- }
- }
- }
-
- Ok(ScalarValue::Float64(None))
- }
-
- fn size(&self) -> usize {
- std::mem::size_of_val(self) - std::mem::size_of_val(&self.covar)
- + self.covar.size()
- - std::mem::size_of_val(&self.stddev1)
- + self.stddev1.size()
- - std::mem::size_of_val(&self.stddev2)
- + self.stddev2.size()
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::aggregate::utils::get_accum_scalar_values_as_arrays;
- use crate::expressions::col;
- use crate::expressions::tests::aggregate;
- use crate::generic_test_op2;
- use arrow::{array::*, datatypes::*};
-
- #[test]
- fn correlation_f64_1() -> Result<()> {
- let a: ArrayRef = Arc::new(Float64Array::from(vec![1_f64, 2_f64,
3_f64]));
- let b: ArrayRef = Arc::new(Float64Array::from(vec![4_f64, 5_f64,
7_f64]));
-
- generic_test_op2!(
- a,
- b,
- DataType::Float64,
- DataType::Float64,
- Correlation,
- ScalarValue::from(0.9819805060619659_f64)
- )
- }
-
- #[test]
- fn correlation_f64_2() -> Result<()> {
- let a: ArrayRef = Arc::new(Float64Array::from(vec![1_f64, 2_f64,
3_f64]));
- let b: ArrayRef = Arc::new(Float64Array::from(vec![4_f64, -5_f64,
6_f64]));
-
- generic_test_op2!(
- a,
- b,
- DataType::Float64,
- DataType::Float64,
- Correlation,
- ScalarValue::from(0.17066403719657236_f64)
- )
- }
-
- #[test]
- fn correlation_f64_4() -> Result<()> {
- let a: ArrayRef = Arc::new(Float64Array::from(vec![1.1_f64, 2_f64,
3_f64]));
- let b: ArrayRef = Arc::new(Float64Array::from(vec![4.1_f64, 5_f64,
6_f64]));
-
- generic_test_op2!(
- a,
- b,
- DataType::Float64,
- DataType::Float64,
- Correlation,
- ScalarValue::from(1_f64)
- )
- }
-
- #[test]
- fn correlation_f64_6() -> Result<()> {
- let a = Arc::new(Float64Array::from(vec![
- 1_f64, 2_f64, 3_f64, 1.1_f64, 2.2_f64, 3.3_f64,
- ]));
- let b = Arc::new(Float64Array::from(vec![
- 4_f64, 5_f64, 6_f64, 4.4_f64, 5.5_f64, 6.6_f64,
- ]));
-
- generic_test_op2!(
- a,
- b,
- DataType::Float64,
- DataType::Float64,
- Correlation,
- ScalarValue::from(0.9860135594710389_f64)
- )
- }
-
- #[test]
- fn correlation_i32() -> Result<()> {
- let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
- let b: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6]));
-
- generic_test_op2!(
- a,
- b,
- DataType::Int32,
- DataType::Int32,
- Correlation,
- ScalarValue::from(1_f64)
- )
- }
-
- #[test]
- fn correlation_u32() -> Result<()> {
- let a: ArrayRef = Arc::new(UInt32Array::from(vec![1_u32, 2_u32,
3_u32]));
- let b: ArrayRef = Arc::new(UInt32Array::from(vec![4_u32, 5_u32,
6_u32]));
- generic_test_op2!(
- a,
- b,
- DataType::UInt32,
- DataType::UInt32,
- Correlation,
- ScalarValue::from(1_f64)
- )
- }
-
- #[test]
- fn correlation_f32() -> Result<()> {
- let a: ArrayRef = Arc::new(Float32Array::from(vec![1_f32, 2_f32,
3_f32]));
- let b: ArrayRef = Arc::new(Float32Array::from(vec![4_f32, 5_f32,
6_f32]));
- generic_test_op2!(
- a,
- b,
- DataType::Float32,
- DataType::Float32,
- Correlation,
- ScalarValue::from(1_f64)
- )
- }
-
- #[test]
- fn correlation_i32_with_nulls_1() -> Result<()> {
- let a: ArrayRef =
- Arc::new(Int32Array::from(vec![Some(1), None, Some(3), Some(3)]));
- let b: ArrayRef =
- Arc::new(Int32Array::from(vec![Some(4), None, Some(6), Some(3)]));
-
- generic_test_op2!(
- a,
- b,
- DataType::Int32,
- DataType::Int32,
- Correlation,
- ScalarValue::from(0.1889822365046137_f64)
- )
- }
-
- #[test]
- fn correlation_i32_with_nulls_2() -> Result<()> {
- let a: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- None,
- Some(2),
- Some(9),
- Some(3),
- ]));
- let b: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(4),
- Some(5),
- Some(5),
- None,
- Some(6),
- ]));
-
- generic_test_op2!(
- a,
- b,
- DataType::Int32,
- DataType::Int32,
- Correlation,
- ScalarValue::from(1_f64)
- )
- }
-
- #[test]
- fn correlation_i32_all_nulls() -> Result<()> {
- let a: ArrayRef = Arc::new(Int32Array::from(vec![None, None]));
- let b: ArrayRef = Arc::new(Int32Array::from(vec![None, None]));
-
- generic_test_op2!(
- a,
- b,
- DataType::Int32,
- DataType::Int32,
- Correlation,
- ScalarValue::Float64(None)
- )
- }
-
- #[test]
- fn correlation_f64_merge_1() -> Result<()> {
- let a = Arc::new(Float64Array::from(vec![1_f64, 2_f64, 3_f64]));
- let b = Arc::new(Float64Array::from(vec![4_f64, 5_f64, 6_f64]));
- let c = Arc::new(Float64Array::from(vec![1.1_f64, 2.2_f64, 3.3_f64]));
- let d = Arc::new(Float64Array::from(vec![4.4_f64, 5.5_f64, 9.9_f64]));
-
- let schema = Schema::new(vec![
- Field::new("a", DataType::Float64, true),
- Field::new("b", DataType::Float64, true),
- ]);
-
- let batch1 = RecordBatch::try_new(Arc::new(schema.clone()), vec![a,
b])?;
- let batch2 = RecordBatch::try_new(Arc::new(schema.clone()), vec![c,
d])?;
-
- let agg1 = Arc::new(Correlation::new(
- col("a", &schema)?,
- col("b", &schema)?,
- "bla".to_string(),
- DataType::Float64,
- ));
-
- let agg2 = Arc::new(Correlation::new(
- col("a", &schema)?,
- col("b", &schema)?,
- "bla".to_string(),
- DataType::Float64,
- ));
-
- let actual = merge(&batch1, &batch2, agg1, agg2)?;
- assert!(actual == ScalarValue::from(0.8443707186481967));
-
- Ok(())
- }
-
- #[test]
- fn correlation_f64_merge_2() -> Result<()> {
- let a = Arc::new(Float64Array::from(vec![1_f64, 2_f64, 3_f64]));
- let b = Arc::new(Float64Array::from(vec![4_f64, 5_f64, 6_f64]));
- let c = Arc::new(Float64Array::from(vec![None]));
- let d = Arc::new(Float64Array::from(vec![None]));
-
- let schema = Schema::new(vec![
- Field::new("a", DataType::Float64, true),
- Field::new("b", DataType::Float64, true),
- ]);
-
- let batch1 = RecordBatch::try_new(Arc::new(schema.clone()), vec![a,
b])?;
- let batch2 = RecordBatch::try_new(Arc::new(schema.clone()), vec![c,
d])?;
-
- let agg1 = Arc::new(Correlation::new(
- col("a", &schema)?,
- col("b", &schema)?,
- "bla".to_string(),
- DataType::Float64,
- ));
-
- let agg2 = Arc::new(Correlation::new(
- col("a", &schema)?,
- col("b", &schema)?,
- "bla".to_string(),
- DataType::Float64,
- ));
-
- let actual = merge(&batch1, &batch2, agg1, agg2)?;
- assert!(actual == ScalarValue::from(1_f64));
-
- Ok(())
- }
-
- fn merge(
- batch1: &RecordBatch,
- batch2: &RecordBatch,
- agg1: Arc<dyn AggregateExpr>,
- agg2: Arc<dyn AggregateExpr>,
- ) -> Result<ScalarValue> {
- let mut accum1 = agg1.create_accumulator()?;
- let mut accum2 = agg2.create_accumulator()?;
- let expr1 = agg1.expressions();
- let expr2 = agg2.expressions();
-
- let values1 = expr1
- .iter()
- .map(|e| {
- e.evaluate(batch1)
- .and_then(|v| v.into_array(batch1.num_rows()))
- })
- .collect::<Result<Vec<_>>>()?;
- let values2 = expr2
- .iter()
- .map(|e| {
- e.evaluate(batch2)
- .and_then(|v| v.into_array(batch2.num_rows()))
- })
- .collect::<Result<Vec<_>>>()?;
- accum1.update_batch(&values1)?;
- accum2.update_batch(&values2)?;
- let state2 = get_accum_scalar_values_as_arrays(accum2.as_mut())?;
- accum1.merge_batch(&state2)?;
- accum1.evaluate()
- }
-}
diff --git a/datafusion/physical-expr/src/aggregate/covariance.rs
b/datafusion/physical-expr/src/aggregate/covariance.rs
deleted file mode 100644
index 639d8a098c..0000000000
--- a/datafusion/physical-expr/src/aggregate/covariance.rs
+++ /dev/null
@@ -1,227 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines physical expressions that can evaluated at runtime during query
execution
-
-use arrow::array::Float64Array;
-use arrow::{
- array::{ArrayRef, UInt64Array},
- compute::cast,
- datatypes::DataType,
-};
-use datafusion_common::{downcast_value, unwrap_or_internal_err, ScalarValue};
-use datafusion_common::{DataFusionError, Result};
-use datafusion_expr::Accumulator;
-
-use crate::aggregate::stats::StatsType;
-
-/// An accumulator to compute covariance
-/// The algrithm used is an online implementation and numerically stable. It
is derived from the following paper
-/// for calculating variance:
-/// Welford, B. P. (1962). "Note on a method for calculating corrected sums of
squares and products".
-/// Technometrics. 4 (3): 419–420. doi:10.2307/1266577. JSTOR 1266577.
-///
-/// The algorithm has been analyzed here:
-/// Ling, Robert F. (1974). "Comparison of Several Algorithms for Computing
Sample Means and Variances".
-/// Journal of the American Statistical Association. 69 (348): 859–866.
doi:10.2307/2286154. JSTOR 2286154.
-///
-/// Though it is not covered in the original paper but is based on the same
idea, as a result the algorithm is online,
-/// parallelizable and numerically stable.
-
-#[derive(Debug)]
-pub struct CovarianceAccumulator {
- algo_const: f64,
- mean1: f64,
- mean2: f64,
- count: u64,
- stats_type: StatsType,
-}
-
-impl CovarianceAccumulator {
- /// Creates a new `CovarianceAccumulator`
- pub fn try_new(s_type: StatsType) -> Result<Self> {
- Ok(Self {
- algo_const: 0_f64,
- mean1: 0_f64,
- mean2: 0_f64,
- count: 0_u64,
- stats_type: s_type,
- })
- }
-
- pub fn get_count(&self) -> u64 {
- self.count
- }
-
- pub fn get_mean1(&self) -> f64 {
- self.mean1
- }
-
- pub fn get_mean2(&self) -> f64 {
- self.mean2
- }
-
- pub fn get_algo_const(&self) -> f64 {
- self.algo_const
- }
-}
-
-impl Accumulator for CovarianceAccumulator {
- fn state(&mut self) -> Result<Vec<ScalarValue>> {
- Ok(vec![
- ScalarValue::from(self.count),
- ScalarValue::from(self.mean1),
- ScalarValue::from(self.mean2),
- ScalarValue::from(self.algo_const),
- ])
- }
-
- fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- let values1 = &cast(&values[0], &DataType::Float64)?;
- let values2 = &cast(&values[1], &DataType::Float64)?;
-
- let mut arr1 = downcast_value!(values1, Float64Array).iter().flatten();
- let mut arr2 = downcast_value!(values2, Float64Array).iter().flatten();
-
- for i in 0..values1.len() {
- let value1 = if values1.is_valid(i) {
- arr1.next()
- } else {
- None
- };
- let value2 = if values2.is_valid(i) {
- arr2.next()
- } else {
- None
- };
-
- if value1.is_none() || value2.is_none() {
- continue;
- }
-
- let value1 = unwrap_or_internal_err!(value1);
- let value2 = unwrap_or_internal_err!(value2);
- let new_count = self.count + 1;
- let delta1 = value1 - self.mean1;
- let new_mean1 = delta1 / new_count as f64 + self.mean1;
- let delta2 = value2 - self.mean2;
- let new_mean2 = delta2 / new_count as f64 + self.mean2;
- let new_c = delta1 * (value2 - new_mean2) + self.algo_const;
-
- self.count += 1;
- self.mean1 = new_mean1;
- self.mean2 = new_mean2;
- self.algo_const = new_c;
- }
-
- Ok(())
- }
-
- fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- let values1 = &cast(&values[0], &DataType::Float64)?;
- let values2 = &cast(&values[1], &DataType::Float64)?;
- let mut arr1 = downcast_value!(values1, Float64Array).iter().flatten();
- let mut arr2 = downcast_value!(values2, Float64Array).iter().flatten();
-
- for i in 0..values1.len() {
- let value1 = if values1.is_valid(i) {
- arr1.next()
- } else {
- None
- };
- let value2 = if values2.is_valid(i) {
- arr2.next()
- } else {
- None
- };
-
- if value1.is_none() || value2.is_none() {
- continue;
- }
-
- let value1 = unwrap_or_internal_err!(value1);
- let value2 = unwrap_or_internal_err!(value2);
-
- let new_count = self.count - 1;
- let delta1 = self.mean1 - value1;
- let new_mean1 = delta1 / new_count as f64 + self.mean1;
- let delta2 = self.mean2 - value2;
- let new_mean2 = delta2 / new_count as f64 + self.mean2;
- let new_c = self.algo_const - delta1 * (new_mean2 - value2);
-
- self.count -= 1;
- self.mean1 = new_mean1;
- self.mean2 = new_mean2;
- self.algo_const = new_c;
- }
-
- Ok(())
- }
-
- fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
- let counts = downcast_value!(states[0], UInt64Array);
- let means1 = downcast_value!(states[1], Float64Array);
- let means2 = downcast_value!(states[2], Float64Array);
- let cs = downcast_value!(states[3], Float64Array);
-
- for i in 0..counts.len() {
- let c = counts.value(i);
- if c == 0_u64 {
- continue;
- }
- let new_count = self.count + c;
- let new_mean1 = self.mean1 * self.count as f64 / new_count as f64
- + means1.value(i) * c as f64 / new_count as f64;
- let new_mean2 = self.mean2 * self.count as f64 / new_count as f64
- + means2.value(i) * c as f64 / new_count as f64;
- let delta1 = self.mean1 - means1.value(i);
- let delta2 = self.mean2 - means2.value(i);
- let new_c = self.algo_const
- + cs.value(i)
- + delta1 * delta2 * self.count as f64 * c as f64 / new_count
as f64;
-
- self.count = new_count;
- self.mean1 = new_mean1;
- self.mean2 = new_mean2;
- self.algo_const = new_c;
- }
- Ok(())
- }
-
- fn evaluate(&mut self) -> Result<ScalarValue> {
- let count = match self.stats_type {
- StatsType::Population => self.count,
- StatsType::Sample => {
- if self.count > 0 {
- self.count - 1
- } else {
- self.count
- }
- }
- };
-
- if count == 0 {
- Ok(ScalarValue::Float64(None))
- } else {
- Ok(ScalarValue::Float64(Some(self.algo_const / count as f64)))
- }
- }
-
- fn size(&self) -> usize {
- std::mem::size_of_val(self)
- }
-}
diff --git a/datafusion/physical-expr/src/aggregate/mod.rs
b/datafusion/physical-expr/src/aggregate/mod.rs
index 7cb50fc252..ca5bf32934 100644
--- a/datafusion/physical-expr/src/aggregate/mod.rs
+++ b/datafusion/physical-expr/src/aggregate/mod.rs
@@ -20,16 +20,12 @@ pub use
datafusion_physical_expr_common::aggregate::AggregateExpr;
pub(crate) mod array_agg;
pub(crate) mod array_agg_distinct;
pub(crate) mod array_agg_ordered;
-pub(crate) mod correlation;
-pub(crate) mod covariance;
pub(crate) mod grouping;
pub(crate) mod nth_value;
#[macro_use]
pub(crate) mod min_max;
pub(crate) mod groups_accumulator;
pub(crate) mod stats;
-pub(crate) mod stddev;
-pub(crate) mod variance;
pub mod build_in;
pub mod moving_min_max;
diff --git a/datafusion/physical-expr/src/aggregate/stddev.rs
b/datafusion/physical-expr/src/aggregate/stddev.rs
deleted file mode 100644
index 3ade67b519..0000000000
--- a/datafusion/physical-expr/src/aggregate/stddev.rs
+++ /dev/null
@@ -1,87 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines physical expressions that can evaluated at runtime during query
execution
-
-use arrow::array::ArrayRef;
-
-use datafusion_common::ScalarValue;
-use datafusion_common::{internal_err, Result};
-use datafusion_expr::Accumulator;
-
-use crate::aggregate::stats::StatsType;
-use crate::aggregate::variance::VarianceAccumulator;
-
-/// An accumulator to compute the average
-#[derive(Debug)]
-pub struct StddevAccumulator {
- variance: VarianceAccumulator,
-}
-
-impl StddevAccumulator {
- /// Creates a new `StddevAccumulator`
- pub fn try_new(s_type: StatsType) -> Result<Self> {
- Ok(Self {
- variance: VarianceAccumulator::try_new(s_type)?,
- })
- }
-
- pub fn get_m2(&self) -> f64 {
- self.variance.get_m2()
- }
-}
-
-impl Accumulator for StddevAccumulator {
- fn state(&mut self) -> Result<Vec<ScalarValue>> {
- Ok(vec![
- ScalarValue::from(self.variance.get_count()),
- ScalarValue::from(self.variance.get_mean()),
- ScalarValue::from(self.variance.get_m2()),
- ])
- }
-
- fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- self.variance.update_batch(values)
- }
-
- fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- self.variance.retract_batch(values)
- }
-
- fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
- self.variance.merge_batch(states)
- }
-
- fn evaluate(&mut self) -> Result<ScalarValue> {
- let variance = self.variance.evaluate()?;
- match variance {
- ScalarValue::Float64(e) => {
- if e.is_none() {
- Ok(ScalarValue::Float64(None))
- } else {
- Ok(ScalarValue::Float64(e.map(|f| f.sqrt())))
- }
- }
- _ => internal_err!("Variance should be f64"),
- }
- }
-
- fn size(&self) -> usize {
- std::mem::align_of_val(self) - std::mem::align_of_val(&self.variance)
- + self.variance.size()
- }
-}
diff --git a/datafusion/physical-expr/src/aggregate/variance.rs
b/datafusion/physical-expr/src/aggregate/variance.rs
deleted file mode 100644
index 27c67a2f9c..0000000000
--- a/datafusion/physical-expr/src/aggregate/variance.rs
+++ /dev/null
@@ -1,176 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! Defines physical expressions that can evaluated at runtime during query
execution
-
-use crate::aggregate::stats::StatsType;
-use arrow::array::Float64Array;
-use arrow::{
- array::{ArrayRef, UInt64Array},
- compute::cast,
- datatypes::DataType,
-};
-use datafusion_common::downcast_value;
-use datafusion_common::{DataFusionError, Result, ScalarValue};
-use datafusion_expr::Accumulator;
-
-// TODO only holds the definition of `VarianceAccumulator` for use by
`StddevAccumulator` in `physical-expr`,
-// which in turn only has it there for legacy `CorrelationAccumulator`, but
this whole file should go away
-// once the latter is moved to `functions-aggregate`.
-
-/// An accumulator to compute variance
-/// The algrithm used is an online implementation and numerically stable. It
is based on this paper:
-/// Welford, B. P. (1962). "Note on a method for calculating corrected sums of
squares and products".
-/// Technometrics. 4 (3): 419–420. doi:10.2307/1266577. JSTOR 1266577.
-///
-/// The algorithm has been analyzed here:
-/// Ling, Robert F. (1974). "Comparison of Several Algorithms for Computing
Sample Means and Variances".
-/// Journal of the American Statistical Association. 69 (348): 859–866.
doi:10.2307/2286154. JSTOR 2286154.
-
-#[derive(Debug)]
-pub struct VarianceAccumulator {
- m2: f64,
- mean: f64,
- count: u64,
- stats_type: StatsType,
-}
-
-impl VarianceAccumulator {
- /// Creates a new `VarianceAccumulator`
- pub fn try_new(s_type: StatsType) -> Result<Self> {
- Ok(Self {
- m2: 0_f64,
- mean: 0_f64,
- count: 0_u64,
- stats_type: s_type,
- })
- }
-
- pub fn get_count(&self) -> u64 {
- self.count
- }
-
- pub fn get_mean(&self) -> f64 {
- self.mean
- }
-
- pub fn get_m2(&self) -> f64 {
- self.m2
- }
-}
-
-impl Accumulator for VarianceAccumulator {
- fn state(&mut self) -> Result<Vec<ScalarValue>> {
- Ok(vec![
- ScalarValue::from(self.count),
- 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 = downcast_value!(values, Float64Array).iter().flatten();
-
- for value in arr {
- 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 retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- let values = &cast(&values[0], &DataType::Float64)?;
- let arr = downcast_value!(values, Float64Array).iter().flatten();
-
- for value in arr {
- let new_count = self.count - 1;
- let delta1 = self.mean - value;
- let new_mean = delta1 / new_count as f64 + self.mean;
- let delta2 = new_mean - value;
- 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 = downcast_value!(states[0], UInt64Array);
- let means = downcast_value!(states[1], Float64Array);
- let m2s = downcast_value!(states[2], Float64Array);
-
- for i in 0..counts.len() {
- let c = counts.value(i);
- if c == 0_u64 {
- continue;
- }
- let new_count = self.count + c;
- let new_mean = self.mean * self.count as f64 / new_count as f64
- + means.value(i) * c as f64 / new_count as f64;
- let delta = self.mean - means.value(i);
- let new_m2 = self.m2
- + m2s.value(i)
- + delta * delta * self.count as f64 * c as f64 / new_count as
f64;
-
- self.count = new_count;
- self.mean = new_mean;
- self.m2 = new_m2;
- }
- Ok(())
- }
-
- fn evaluate(&mut self) -> Result<ScalarValue> {
- let count = match self.stats_type {
- StatsType::Population => self.count,
- StatsType::Sample => {
- if self.count > 0 {
- self.count - 1
- } else {
- self.count
- }
- }
- };
-
- Ok(ScalarValue::Float64(match self.count {
- 0 => None,
- 1 => {
- if let StatsType::Population = self.stats_type {
- Some(0.0)
- } else {
- None
- }
- }
- _ => Some(self.m2 / count as f64),
- }))
- }
-
- fn size(&self) -> usize {
- std::mem::size_of_val(self)
- }
-}
diff --git a/datafusion/physical-expr/src/expressions/mod.rs
b/datafusion/physical-expr/src/expressions/mod.rs
index 59de4d0874..c98bcc56ad 100644
--- a/datafusion/physical-expr/src/expressions/mod.rs
+++ b/datafusion/physical-expr/src/expressions/mod.rs
@@ -39,7 +39,6 @@ pub use crate::aggregate::array_agg::ArrayAgg;
pub use crate::aggregate::array_agg_distinct::DistinctArrayAgg;
pub use crate::aggregate::array_agg_ordered::OrderSensitiveArrayAgg;
pub use crate::aggregate::build_in::create_aggregate_expr;
-pub use crate::aggregate::correlation::Correlation;
pub use crate::aggregate::grouping::Grouping;
pub use crate::aggregate::min_max::{Max, MaxAccumulator, Min, MinAccumulator};
pub use crate::aggregate::nth_value::NthValueAgg;
diff --git a/datafusion/proto/proto/datafusion.proto
b/datafusion/proto/proto/datafusion.proto
index 8d73643f9c..7710acf137 100644
--- a/datafusion/proto/proto/datafusion.proto
+++ b/datafusion/proto/proto/datafusion.proto
@@ -485,7 +485,7 @@ enum AggregateFunction {
// COVARIANCE_POP = 10;
// STDDEV = 11;
// STDDEV_POP = 12;
- CORRELATION = 13;
+ // CORRELATION = 13;
// APPROX_PERCENTILE_CONT = 14;
// APPROX_MEDIAN = 15;
// APPROX_PERCENTILE_CONT_WITH_WEIGHT = 16;
diff --git a/datafusion/proto/src/generated/pbjson.rs
b/datafusion/proto/src/generated/pbjson.rs
index 10d3f54dc8..48a0c1ad61 100644
--- a/datafusion/proto/src/generated/pbjson.rs
+++ b/datafusion/proto/src/generated/pbjson.rs
@@ -535,7 +535,6 @@ impl serde::Serialize for AggregateFunction {
Self::Min => "MIN",
Self::Max => "MAX",
Self::ArrayAgg => "ARRAY_AGG",
- Self::Correlation => "CORRELATION",
Self::Grouping => "GROUPING",
Self::NthValueAgg => "NTH_VALUE_AGG",
};
@@ -552,7 +551,6 @@ impl<'de> serde::Deserialize<'de> for AggregateFunction {
"MIN",
"MAX",
"ARRAY_AGG",
- "CORRELATION",
"GROUPING",
"NTH_VALUE_AGG",
];
@@ -598,7 +596,6 @@ impl<'de> serde::Deserialize<'de> for AggregateFunction {
"MIN" => Ok(AggregateFunction::Min),
"MAX" => Ok(AggregateFunction::Max),
"ARRAY_AGG" => Ok(AggregateFunction::ArrayAgg),
- "CORRELATION" => Ok(AggregateFunction::Correlation),
"GROUPING" => Ok(AggregateFunction::Grouping),
"NTH_VALUE_AGG" => Ok(AggregateFunction::NthValueAgg),
_ => Err(serde::de::Error::unknown_variant(value, FIELDS)),
diff --git a/datafusion/proto/src/generated/prost.rs
b/datafusion/proto/src/generated/prost.rs
index 9504d26d44..5959b6fba1 100644
--- a/datafusion/proto/src/generated/prost.rs
+++ b/datafusion/proto/src/generated/prost.rs
@@ -1939,7 +1939,7 @@ pub enum AggregateFunction {
/// COVARIANCE_POP = 10;
/// STDDEV = 11;
/// STDDEV_POP = 12;
- Correlation = 13,
+ /// CORRELATION = 13;
/// APPROX_PERCENTILE_CONT = 14;
/// APPROX_MEDIAN = 15;
/// APPROX_PERCENTILE_CONT_WITH_WEIGHT = 16;
@@ -1972,7 +1972,6 @@ impl AggregateFunction {
AggregateFunction::Min => "MIN",
AggregateFunction::Max => "MAX",
AggregateFunction::ArrayAgg => "ARRAY_AGG",
- AggregateFunction::Correlation => "CORRELATION",
AggregateFunction::Grouping => "GROUPING",
AggregateFunction::NthValueAgg => "NTH_VALUE_AGG",
}
@@ -1983,7 +1982,6 @@ impl AggregateFunction {
"MIN" => Some(Self::Min),
"MAX" => Some(Self::Max),
"ARRAY_AGG" => Some(Self::ArrayAgg),
- "CORRELATION" => Some(Self::Correlation),
"GROUPING" => Some(Self::Grouping),
"NTH_VALUE_AGG" => Some(Self::NthValueAgg),
_ => None,
diff --git a/datafusion/proto/src/logical_plan/from_proto.rs
b/datafusion/proto/src/logical_plan/from_proto.rs
index 420e59688d..2a65bc1033 100644
--- a/datafusion/proto/src/logical_plan/from_proto.rs
+++ b/datafusion/proto/src/logical_plan/from_proto.rs
@@ -140,7 +140,6 @@ impl From<protobuf::AggregateFunction> for
AggregateFunction {
protobuf::AggregateFunction::Min => Self::Min,
protobuf::AggregateFunction::Max => Self::Max,
protobuf::AggregateFunction::ArrayAgg => Self::ArrayAgg,
- protobuf::AggregateFunction::Correlation => Self::Correlation,
protobuf::AggregateFunction::Grouping => Self::Grouping,
protobuf::AggregateFunction::NthValueAgg => Self::NthValue,
}
diff --git a/datafusion/proto/src/logical_plan/to_proto.rs
b/datafusion/proto/src/logical_plan/to_proto.rs
index d12bfd9ffd..d22052d5a7 100644
--- a/datafusion/proto/src/logical_plan/to_proto.rs
+++ b/datafusion/proto/src/logical_plan/to_proto.rs
@@ -111,7 +111,6 @@ impl From<&AggregateFunction> for
protobuf::AggregateFunction {
AggregateFunction::Min => Self::Min,
AggregateFunction::Max => Self::Max,
AggregateFunction::ArrayAgg => Self::ArrayAgg,
- AggregateFunction::Correlation => Self::Correlation,
AggregateFunction::Grouping => Self::Grouping,
AggregateFunction::NthValue => Self::NthValueAgg,
}
@@ -373,9 +372,6 @@ pub fn serialize_expr(
AggregateFunction::ArrayAgg =>
protobuf::AggregateFunction::ArrayAgg,
AggregateFunction::Min => protobuf::AggregateFunction::Min,
AggregateFunction::Max => protobuf::AggregateFunction::Max,
- AggregateFunction::Correlation => {
- protobuf::AggregateFunction::Correlation
- }
AggregateFunction::Grouping =>
protobuf::AggregateFunction::Grouping,
AggregateFunction::NthValue => {
protobuf::AggregateFunction::NthValueAgg
diff --git a/datafusion/proto/src/physical_plan/to_proto.rs
b/datafusion/proto/src/physical_plan/to_proto.rs
index 4554e529c3..c02b59d062 100644
--- a/datafusion/proto/src/physical_plan/to_proto.rs
+++ b/datafusion/proto/src/physical_plan/to_proto.rs
@@ -23,10 +23,10 @@ use
datafusion::datasource::file_format::parquet::ParquetSink;
use datafusion::physical_expr::window::{NthValueKind,
SlidingAggregateWindowExpr};
use datafusion::physical_expr::{PhysicalSortExpr, ScalarFunctionExpr};
use datafusion::physical_plan::expressions::{
- ArrayAgg, BinaryExpr, CaseExpr, CastExpr, Column, Correlation, CumeDist,
- DistinctArrayAgg, Grouping, InListExpr, IsNotNullExpr, IsNullExpr,
Literal, Max, Min,
- NegativeExpr, NotExpr, NthValue, NthValueAgg, Ntile,
OrderSensitiveArrayAgg, Rank,
- RankType, RowNumber, TryCastExpr, WindowShift,
+ ArrayAgg, BinaryExpr, CaseExpr, CastExpr, Column, CumeDist,
DistinctArrayAgg,
+ Grouping, InListExpr, IsNotNullExpr, IsNullExpr, Literal, Max, Min,
NegativeExpr,
+ NotExpr, NthValue, NthValueAgg, Ntile, OrderSensitiveArrayAgg, Rank,
RankType,
+ RowNumber, TryCastExpr, WindowShift,
};
use datafusion::physical_plan::udaf::AggregateFunctionExpr;
use datafusion::physical_plan::windows::{BuiltInWindowExpr,
PlainAggregateWindowExpr};
@@ -258,8 +258,6 @@ fn aggr_expr_to_aggr_fn(expr: &dyn AggregateExpr) ->
Result<AggrFn> {
protobuf::AggregateFunction::Min
} else if aggr_expr.downcast_ref::<Max>().is_some() {
protobuf::AggregateFunction::Max
- } else if aggr_expr.downcast_ref::<Correlation>().is_some() {
- protobuf::AggregateFunction::Correlation
} else if aggr_expr.downcast_ref::<NthValueAgg>().is_some() {
protobuf::AggregateFunction::NthValueAgg
} else {
diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
index 1d61a65f0b..486a911f92 100644
--- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
@@ -61,7 +61,7 @@ use datafusion_expr::{
};
use datafusion_functions_aggregate::average::avg_udaf;
use datafusion_functions_aggregate::expr_fn::{
- bit_and, bit_or, bit_xor, bool_and, bool_or,
+ avg, bit_and, bit_or, bit_xor, bool_and, bool_or, corr,
};
use datafusion_functions_aggregate::string_agg::string_agg;
use datafusion_proto::bytes::{
@@ -659,8 +659,10 @@ async fn roundtrip_expr_api() -> Result<()> {
count_distinct(lit(1)),
first_value(lit(1), None),
first_value(lit(1), Some(vec![lit(2).sort(true, true)])),
+ avg(lit(1.5)),
covar_samp(lit(1.5), lit(2.2)),
covar_pop(lit(1.5), lit(2.2)),
+ corr(lit(1.5), lit(2.2)),
sum(lit(1)),
median(lit(2)),
var_sample(lit(2.2)),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
