datafusion_functions_aggregate/
covariance.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`CovarianceSample`]: covariance sample aggregations.
19
20use arrow::array::ArrayRef;
21use arrow::datatypes::{DataType, Field, FieldRef};
22use datafusion_common::cast::{as_float64_array, as_uint64_array};
23use datafusion_common::{Result, ScalarValue};
24use datafusion_expr::{
25    Accumulator, AggregateUDFImpl, Documentation, Signature, Volatility,
26    function::{AccumulatorArgs, StateFieldsArgs},
27    utils::format_state_name,
28};
29use datafusion_functions_aggregate_common::stats::StatsType;
30use datafusion_macros::user_doc;
31use std::fmt::Debug;
32use std::mem::size_of_val;
33use std::sync::Arc;
34
35make_udaf_expr_and_func!(
36    CovarianceSample,
37    covar_samp,
38    y x,
39    "Computes the sample covariance.",
40    covar_samp_udaf
41);
42
43make_udaf_expr_and_func!(
44    CovariancePopulation,
45    covar_pop,
46    y x,
47    "Computes the population covariance.",
48    covar_pop_udaf
49);
50
51#[user_doc(
52    doc_section(label = "Statistical Functions"),
53    description = "Returns the sample covariance of a set of number pairs.",
54    syntax_example = "covar_samp(expression1, expression2)",
55    sql_example = r#"```sql
56> SELECT covar_samp(column1, column2) FROM table_name;
57+-----------------------------------+
58| covar_samp(column1, column2)      |
59+-----------------------------------+
60| 8.25                              |
61+-----------------------------------+
62```"#,
63    standard_argument(name = "expression1", prefix = "First"),
64    standard_argument(name = "expression2", prefix = "Second")
65)]
66#[derive(PartialEq, Eq, Hash, Debug)]
67pub struct CovarianceSample {
68    signature: Signature,
69    aliases: Vec<String>,
70}
71
72impl Default for CovarianceSample {
73    fn default() -> Self {
74        Self::new()
75    }
76}
77
78impl CovarianceSample {
79    pub fn new() -> Self {
80        Self {
81            aliases: vec![String::from("covar")],
82            signature: Signature::exact(
83                vec![DataType::Float64, DataType::Float64],
84                Volatility::Immutable,
85            ),
86        }
87    }
88}
89
90impl AggregateUDFImpl for CovarianceSample {
91    fn name(&self) -> &str {
92        "covar_samp"
93    }
94
95    fn signature(&self) -> &Signature {
96        &self.signature
97    }
98
99    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
100        Ok(DataType::Float64)
101    }
102
103    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
104        let name = args.name;
105        Ok(vec![
106            Field::new(format_state_name(name, "count"), DataType::UInt64, true),
107            Field::new(format_state_name(name, "mean1"), DataType::Float64, true),
108            Field::new(format_state_name(name, "mean2"), DataType::Float64, true),
109            Field::new(
110                format_state_name(name, "algo_const"),
111                DataType::Float64,
112                true,
113            ),
114        ]
115        .into_iter()
116        .map(Arc::new)
117        .collect())
118    }
119
120    fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
121        Ok(Box::new(CovarianceAccumulator::try_new(StatsType::Sample)?))
122    }
123
124    fn aliases(&self) -> &[String] {
125        &self.aliases
126    }
127
128    fn documentation(&self) -> Option<&Documentation> {
129        self.doc()
130    }
131}
132
133#[user_doc(
134    doc_section(label = "Statistical Functions"),
135    description = "Returns the sample covariance of a set of number pairs.",
136    syntax_example = "covar_samp(expression1, expression2)",
137    sql_example = r#"```sql
138> SELECT covar_samp(column1, column2) FROM table_name;
139+-----------------------------------+
140| covar_samp(column1, column2)      |
141+-----------------------------------+
142| 8.25                              |
143+-----------------------------------+
144```"#,
145    standard_argument(name = "expression1", prefix = "First"),
146    standard_argument(name = "expression2", prefix = "Second")
147)]
148#[derive(PartialEq, Eq, Hash, Debug)]
149pub struct CovariancePopulation {
150    signature: Signature,
151}
152
153impl Default for CovariancePopulation {
154    fn default() -> Self {
155        Self::new()
156    }
157}
158
159impl CovariancePopulation {
160    pub fn new() -> Self {
161        Self {
162            signature: Signature::exact(
163                vec![DataType::Float64, DataType::Float64],
164                Volatility::Immutable,
165            ),
166        }
167    }
168}
169
170impl AggregateUDFImpl for CovariancePopulation {
171    fn name(&self) -> &str {
172        "covar_pop"
173    }
174
175    fn signature(&self) -> &Signature {
176        &self.signature
177    }
178
179    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
180        Ok(DataType::Float64)
181    }
182
183    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
184        let name = args.name;
185        Ok(vec![
186            Field::new(format_state_name(name, "count"), DataType::UInt64, true),
187            Field::new(format_state_name(name, "mean1"), DataType::Float64, true),
188            Field::new(format_state_name(name, "mean2"), DataType::Float64, true),
189            Field::new(
190                format_state_name(name, "algo_const"),
191                DataType::Float64,
192                true,
193            ),
194        ]
195        .into_iter()
196        .map(Arc::new)
197        .collect())
198    }
199
200    fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
201        Ok(Box::new(CovarianceAccumulator::try_new(
202            StatsType::Population,
203        )?))
204    }
205
206    fn documentation(&self) -> Option<&Documentation> {
207        self.doc()
208    }
209}
210
211/// An accumulator to compute covariance
212/// The algorithm used is an online implementation and numerically stable. It is derived from the following paper
213/// for calculating variance:
214/// Welford, B. P. (1962). "Note on a method for calculating corrected sums of squares and products".
215/// Technometrics. 4 (3): 419–420. doi:10.2307/1266577. JSTOR 1266577.
216///
217/// The algorithm has been analyzed here:
218/// Ling, Robert F. (1974). "Comparison of Several Algorithms for Computing Sample Means and Variances".
219/// Journal of the American Statistical Association. 69 (348): 859–866. doi:10.2307/2286154. JSTOR 2286154.
220///
221/// Though it is not covered in the original paper but is based on the same idea, as a result the algorithm is online,
222/// parallelize and numerically stable.
223
224#[derive(Debug)]
225pub struct CovarianceAccumulator {
226    algo_const: f64,
227    mean1: f64,
228    mean2: f64,
229    count: u64,
230    stats_type: StatsType,
231}
232
233impl CovarianceAccumulator {
234    /// Creates a new `CovarianceAccumulator`
235    pub fn try_new(s_type: StatsType) -> Result<Self> {
236        Ok(Self {
237            algo_const: 0_f64,
238            mean1: 0_f64,
239            mean2: 0_f64,
240            count: 0_u64,
241            stats_type: s_type,
242        })
243    }
244
245    pub fn get_count(&self) -> u64 {
246        self.count
247    }
248
249    pub fn get_mean1(&self) -> f64 {
250        self.mean1
251    }
252
253    pub fn get_mean2(&self) -> f64 {
254        self.mean2
255    }
256
257    pub fn get_algo_const(&self) -> f64 {
258        self.algo_const
259    }
260}
261
262impl Accumulator for CovarianceAccumulator {
263    fn state(&mut self) -> Result<Vec<ScalarValue>> {
264        Ok(vec![
265            ScalarValue::from(self.count),
266            ScalarValue::from(self.mean1),
267            ScalarValue::from(self.mean2),
268            ScalarValue::from(self.algo_const),
269        ])
270    }
271
272    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
273        let values1 = as_float64_array(&values[0])?;
274        let values2 = as_float64_array(&values[1])?;
275
276        for (value1, value2) in values1.iter().zip(values2) {
277            let (value1, value2) = match (value1, value2) {
278                (Some(a), Some(b)) => (a, b),
279                _ => continue,
280            };
281
282            let new_count = self.count + 1;
283            let delta1 = value1 - self.mean1;
284            let new_mean1 = delta1 / new_count as f64 + self.mean1;
285            let delta2 = value2 - self.mean2;
286            let new_mean2 = delta2 / new_count as f64 + self.mean2;
287            let new_c = delta1 * (value2 - new_mean2) + self.algo_const;
288
289            self.count += 1;
290            self.mean1 = new_mean1;
291            self.mean2 = new_mean2;
292            self.algo_const = new_c;
293        }
294
295        Ok(())
296    }
297
298    fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
299        let values1 = as_float64_array(&values[0])?;
300        let values2 = as_float64_array(&values[1])?;
301
302        for (value1, value2) in values1.iter().zip(values2) {
303            let (value1, value2) = match (value1, value2) {
304                (Some(a), Some(b)) => (a, b),
305                _ => continue,
306            };
307
308            let new_count = self.count - 1;
309            let delta1 = self.mean1 - value1;
310            let new_mean1 = delta1 / new_count as f64 + self.mean1;
311            let delta2 = self.mean2 - value2;
312            let new_mean2 = delta2 / new_count as f64 + self.mean2;
313            let new_c = self.algo_const - delta1 * (new_mean2 - value2);
314
315            self.count -= 1;
316            self.mean1 = new_mean1;
317            self.mean2 = new_mean2;
318            self.algo_const = new_c;
319        }
320
321        Ok(())
322    }
323
324    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
325        let counts = as_uint64_array(&states[0])?;
326        let means1 = as_float64_array(&states[1])?;
327        let means2 = as_float64_array(&states[2])?;
328        let cs = as_float64_array(&states[3])?;
329
330        for i in 0..counts.len() {
331            let c = counts.value(i);
332            if c == 0_u64 {
333                continue;
334            }
335            let new_count = self.count + c;
336            let new_mean1 = self.mean1 * self.count as f64 / new_count as f64
337                + means1.value(i) * c as f64 / new_count as f64;
338            let new_mean2 = self.mean2 * self.count as f64 / new_count as f64
339                + means2.value(i) * c as f64 / new_count as f64;
340            let delta1 = self.mean1 - means1.value(i);
341            let delta2 = self.mean2 - means2.value(i);
342            let new_c = self.algo_const
343                + cs.value(i)
344                + delta1 * delta2 * self.count as f64 * c as f64 / new_count as f64;
345
346            self.count = new_count;
347            self.mean1 = new_mean1;
348            self.mean2 = new_mean2;
349            self.algo_const = new_c;
350        }
351        Ok(())
352    }
353
354    fn evaluate(&mut self) -> Result<ScalarValue> {
355        let count = match self.stats_type {
356            StatsType::Population => self.count,
357            StatsType::Sample => {
358                if self.count > 0 {
359                    self.count - 1
360                } else {
361                    self.count
362                }
363            }
364        };
365
366        if count == 0 {
367            Ok(ScalarValue::Float64(None))
368        } else {
369            Ok(ScalarValue::Float64(Some(self.algo_const / count as f64)))
370        }
371    }
372
373    fn size(&self) -> usize {
374        size_of_val(self)
375    }
376}
datafusion_functions_aggregate/covariance.rs

datafusion_functions_aggregate/
covariance.rs
