Skip to main content

datafusion_functions_aggregate/
regr.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//! Defines physical expressions that can evaluated at runtime during query execution
19
20use arrow::datatypes::FieldRef;
21use arrow::{array::ArrayRef, datatypes::DataType, datatypes::Field};
22use datafusion_common::cast::{as_float64_array, as_uint64_array};
23use datafusion_common::{HashMap, Result, ScalarValue};
24use datafusion_doc::aggregate_doc_sections::DOC_SECTION_STATISTICAL;
25use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
26use datafusion_expr::utils::format_state_name;
27use datafusion_expr::{
28    Accumulator, AggregateUDFImpl, Documentation, Signature, Volatility,
29};
30use std::fmt::Debug;
31use std::hash::Hash;
32use std::mem::size_of_val;
33use std::sync::{Arc, LazyLock};
34
35macro_rules! make_regr_udaf_expr_and_func {
36    ($EXPR_FN:ident, $AGGREGATE_UDF_FN:ident, $REGR_TYPE:expr) => {
37        make_udaf_expr!($EXPR_FN, expr_y expr_x, concat!("Compute a linear regression of type [", stringify!($REGR_TYPE), "]"), $AGGREGATE_UDF_FN);
38        create_func!($EXPR_FN, $AGGREGATE_UDF_FN, Regr::new($REGR_TYPE, stringify!($EXPR_FN)));
39    }
40}
41
42make_regr_udaf_expr_and_func!(regr_slope, regr_slope_udaf, RegrType::Slope);
43make_regr_udaf_expr_and_func!(regr_intercept, regr_intercept_udaf, RegrType::Intercept);
44make_regr_udaf_expr_and_func!(regr_count, regr_count_udaf, RegrType::Count);
45make_regr_udaf_expr_and_func!(regr_r2, regr_r2_udaf, RegrType::R2);
46make_regr_udaf_expr_and_func!(regr_avgx, regr_avgx_udaf, RegrType::AvgX);
47make_regr_udaf_expr_and_func!(regr_avgy, regr_avgy_udaf, RegrType::AvgY);
48make_regr_udaf_expr_and_func!(regr_sxx, regr_sxx_udaf, RegrType::SXX);
49make_regr_udaf_expr_and_func!(regr_syy, regr_syy_udaf, RegrType::SYY);
50make_regr_udaf_expr_and_func!(regr_sxy, regr_sxy_udaf, RegrType::SXY);
51
52#[derive(PartialEq, Eq, Hash, Debug)]
53pub struct Regr {
54    signature: Signature,
55    regr_type: RegrType,
56    func_name: &'static str,
57}
58
59impl Regr {
60    pub fn new(regr_type: RegrType, func_name: &'static str) -> Self {
61        Self {
62            signature: Signature::exact(
63                vec![DataType::Float64, DataType::Float64],
64                Volatility::Immutable,
65            ),
66            regr_type,
67            func_name,
68        }
69    }
70}
71
72#[derive(Debug, Clone, PartialEq, Hash, Eq)]
73pub enum RegrType {
74    /// Variant for `regr_slope` aggregate expression
75    /// Returns the slope of the linear regression line for non-null pairs in aggregate columns.
76    /// Given input column Y and X: `regr_slope(Y, X)` returns the slope (k in Y = k*X + b) using minimal
77    /// RSS (Residual Sum of Squares) fitting.
78    Slope,
79    /// Variant for `regr_intercept` aggregate expression
80    /// Returns the intercept of the linear regression line for non-null pairs in aggregate columns.
81    /// Given input column Y and X: `regr_intercept(Y, X)` returns the intercept (b in Y = k*X + b) using minimal
82    /// RSS fitting.
83    Intercept,
84    /// Variant for `regr_count` aggregate expression
85    /// Returns the number of input rows for which both expressions are not null.
86    /// Given input column Y and X: `regr_count(Y, X)` returns the count of non-null pairs.
87    Count,
88    /// Variant for `regr_r2` aggregate expression
89    /// Returns the coefficient of determination (R-squared value) of the linear regression line for non-null pairs in aggregate columns.
90    /// The R-squared value represents the proportion of variance in Y that is predictable from X.
91    R2,
92    /// Variant for `regr_avgx` aggregate expression
93    /// Returns the average of the independent variable for non-null pairs in aggregate columns.
94    /// Given input column X: `regr_avgx(Y, X)` returns the average of X values.
95    AvgX,
96    /// Variant for `regr_avgy` aggregate expression
97    /// Returns the average of the dependent variable for non-null pairs in aggregate columns.
98    /// Given input column Y: `regr_avgy(Y, X)` returns the average of Y values.
99    AvgY,
100    /// Variant for `regr_sxx` aggregate expression
101    /// Returns the sum of squares of the independent variable for non-null pairs in aggregate columns.
102    /// Given input column X: `regr_sxx(Y, X)` returns the sum of squares of deviations of X from its mean.
103    SXX,
104    /// Variant for `regr_syy` aggregate expression
105    /// Returns the sum of squares of the dependent variable for non-null pairs in aggregate columns.
106    /// Given input column Y: `regr_syy(Y, X)` returns the sum of squares of deviations of Y from its mean.
107    SYY,
108    /// Variant for `regr_sxy` aggregate expression
109    /// Returns the sum of products of pairs of numbers for non-null pairs in aggregate columns.
110    /// Given input column Y and X: `regr_sxy(Y, X)` returns the sum of products of the deviations of Y and X from their respective means.
111    SXY,
112}
113
114impl RegrType {
115    /// return the documentation for the `RegrType`
116    fn documentation(&self) -> Option<&Documentation> {
117        get_regr_docs().get(self)
118    }
119}
120
121static DOCUMENTATION: LazyLock<HashMap<RegrType, Documentation>> = LazyLock::new(|| {
122    let mut hash_map = HashMap::new();
123    hash_map.insert(
124            RegrType::Slope,
125            Documentation::builder(
126                DOC_SECTION_STATISTICAL,
127                    "Returns the slope of the linear regression line for non-null pairs in aggregate columns. \
128                    Given input column Y and X: regr_slope(Y, X) returns the slope (k in Y = k*X + b) using minimal RSS fitting.",
129
130                "regr_slope(expression_y, expression_x)")
131                .with_sql_example(
132                    r#"```sql
133create table weekly_performance(day int, user_signups int) as values (1,60), (2,65), (3, 70), (4,75), (5,80);
134select * from weekly_performance;
135+-----+--------------+
136| day | user_signups |
137+-----+--------------+
138| 1   | 60           |
139| 2   | 65           |
140| 3   | 70           |
141| 4   | 75           |
142| 5   | 80           |
143+-----+--------------+
144
145SELECT regr_slope(user_signups, day) AS slope FROM weekly_performance;
146+--------+
147| slope  |
148+--------+
149| 5.0    |
150+--------+
151```
152"#
153                )
154                .with_standard_argument("expression_y", Some("Dependent variable"))
155                .with_standard_argument("expression_x", Some("Independent variable"))
156                .build()
157        );
158
159    hash_map.insert(
160            RegrType::Intercept,
161            Documentation::builder(
162                DOC_SECTION_STATISTICAL,
163                    "Computes the y-intercept of the linear regression line. For the equation (y = kx + b), \
164                    this function returns b.",
165
166                "regr_intercept(expression_y, expression_x)")
167                .with_sql_example(
168                    r#"```sql
169create table weekly_performance(week int, productivity_score int) as values (1,60), (2,65), (3, 70), (4,75), (5,80);
170select * from weekly_performance;
171+------+---------------------+
172| week | productivity_score  |
173| ---- | ------------------- |
174| 1    | 60                  |
175| 2    | 65                  |
176| 3    | 70                  |
177| 4    | 75                  |
178| 5    | 80                  |
179+------+---------------------+
180
181SELECT regr_intercept(productivity_score, week) AS intercept FROM weekly_performance;
182+----------+
183|intercept|
184|intercept |
185+----------+
186|  55      |
187+----------+
188```
189"#
190                )
191                .with_standard_argument("expression_y", Some("Dependent variable"))
192                .with_standard_argument("expression_x", Some("Independent variable"))
193                .build()
194        );
195
196    hash_map.insert(
197        RegrType::Count,
198        Documentation::builder(
199            DOC_SECTION_STATISTICAL,
200            "Counts the number of non-null paired data points.",
201            "regr_count(expression_y, expression_x)",
202        )
203        .with_sql_example(
204            r#"```sql
205create table daily_metrics(day int, user_signups int) as values (1,100), (2,120), (3, NULL), (4,110), (5,NULL);
206select * from daily_metrics;
207+-----+---------------+
208| day | user_signups  |
209| --- | ------------- |
210| 1   | 100           |
211| 2   | 120           |
212| 3   | NULL          |
213| 4   | 110           |
214| 5   | NULL          |
215+-----+---------------+
216
217SELECT regr_count(user_signups, day) AS valid_pairs FROM daily_metrics;
218+-------------+
219| valid_pairs |
220+-------------+
221| 3           |
222+-------------+
223```
224"#
225        )
226        .with_standard_argument("expression_y", Some("Dependent variable"))
227        .with_standard_argument("expression_x", Some("Independent variable"))
228        .build(),
229    );
230
231    hash_map.insert(
232            RegrType::R2,
233            Documentation::builder(
234                DOC_SECTION_STATISTICAL,
235                    "Computes the square of the correlation coefficient between the independent and dependent variables.",
236
237                "regr_r2(expression_y, expression_x)")
238                .with_sql_example(
239                    r#"```sql
240create table weekly_performance(day int ,user_signups int) as values (1,60), (2,65), (3, 70), (4,75), (5,80);
241select * from weekly_performance;
242+-----+--------------+
243| day | user_signups |
244+-----+--------------+
245| 1   | 60           |
246| 2   | 65           |
247| 3   | 70           |
248| 4   | 75           |
249| 5   | 80           |
250+-----+--------------+
251
252SELECT regr_r2(user_signups, day) AS r_squared FROM weekly_performance;
253+---------+
254|r_squared|
255+---------+
256| 1.0     |
257+---------+
258```
259"#
260                )
261                .with_standard_argument("expression_y", Some("Dependent variable"))
262                .with_standard_argument("expression_x", Some("Independent variable"))
263                .build()
264        );
265
266    hash_map.insert(
267            RegrType::AvgX,
268            Documentation::builder(
269                DOC_SECTION_STATISTICAL,
270                    "Computes the average of the independent variable (input) expression_x for the non-null paired data points.",
271
272                "regr_avgx(expression_y, expression_x)")
273                .with_sql_example(
274                    r#"```sql
275create table daily_sales(day int, total_sales int) as values (1,100), (2,150), (3,200), (4,NULL), (5,250);
276select * from daily_sales;
277+-----+-------------+
278| day | total_sales |
279| --- | ----------- |
280| 1   | 100         |
281| 2   | 150         |
282| 3   | 200         |
283| 4   | NULL        |
284| 5   | 250         |
285+-----+-------------+
286
287SELECT regr_avgx(total_sales, day) AS avg_day FROM daily_sales;
288+----------+
289| avg_day  |
290+----------+
291|   2.75   |
292+----------+
293```
294"#
295                )
296                .with_standard_argument("expression_y", Some("Dependent variable"))
297                .with_standard_argument("expression_x", Some("Independent variable"))
298                .build()
299        );
300
301    hash_map.insert(
302            RegrType::AvgY,
303            Documentation::builder(
304                DOC_SECTION_STATISTICAL,
305                    "Computes the average of the dependent variable (output) expression_y for the non-null paired data points.",
306
307                "regr_avgy(expression_y, expression_x)")
308                .with_sql_example(
309                    r#"```sql
310create table daily_temperature(day int, temperature int) as values (1,30), (2,32), (3, NULL), (4,35), (5,36);
311select * from daily_temperature;
312+-----+-------------+
313| day | temperature |
314| --- | ----------- |
315| 1   | 30          |
316| 2   | 32          |
317| 3   | NULL        |
318| 4   | 35          |
319| 5   | 36          |
320+-----+-------------+
321
322-- temperature as Dependent Variable(Y), day as Independent Variable(X)
323SELECT regr_avgy(temperature, day) AS avg_temperature FROM daily_temperature;
324+-----------------+
325| avg_temperature |
326+-----------------+
327| 33.25           |
328+-----------------+
329```
330"#
331                )
332                .with_standard_argument("expression_y", Some("Dependent variable"))
333                .with_standard_argument("expression_x", Some("Independent variable"))
334                .build()
335        );
336
337    hash_map.insert(
338        RegrType::SXX,
339        Documentation::builder(
340            DOC_SECTION_STATISTICAL,
341            "Computes the sum of squares of the independent variable.",
342            "regr_sxx(expression_y, expression_x)",
343        )
344        .with_sql_example(
345            r#"```sql
346create table study_hours(student_id int, hours int, test_score int) as values (1,2,55), (2,4,65), (3,6,75), (4,8,85), (5,10,95);
347select * from study_hours;
348+------------+-------+------------+
349| student_id | hours | test_score |
350+------------+-------+------------+
351| 1          | 2     | 55         |
352| 2          | 4     | 65         |
353| 3          | 6     | 75         |
354| 4          | 8     | 85         |
355| 5          | 10    | 95         |
356+------------+-------+------------+
357
358SELECT regr_sxx(test_score, hours) AS sxx FROM study_hours;
359+------+
360| sxx  |
361+------+
362| 40.0 |
363+------+
364```
365"#
366        )
367        .with_standard_argument("expression_y", Some("Dependent variable"))
368        .with_standard_argument("expression_x", Some("Independent variable"))
369        .build(),
370    );
371
372    hash_map.insert(
373        RegrType::SYY,
374        Documentation::builder(
375            DOC_SECTION_STATISTICAL,
376            "Computes the sum of squares of the dependent variable.",
377            "regr_syy(expression_y, expression_x)",
378        )
379        .with_sql_example(
380            r#"```sql
381create table employee_productivity(week int, productivity_score int) as values (1,60), (2,65), (3,70);
382select * from employee_productivity;
383+------+--------------------+
384| week | productivity_score |
385+------+--------------------+
386| 1    | 60                 |
387| 2    | 65                 |
388| 3    | 70                 |
389+------+--------------------+
390
391SELECT regr_syy(productivity_score, week) AS sum_squares_y FROM employee_productivity;
392+---------------+
393| sum_squares_y |
394+---------------+
395|    50.0       |
396+---------------+
397```
398"#
399        )
400        .with_standard_argument("expression_y", Some("Dependent variable"))
401        .with_standard_argument("expression_x", Some("Independent variable"))
402        .build(),
403    );
404
405    hash_map.insert(
406        RegrType::SXY,
407        Documentation::builder(
408            DOC_SECTION_STATISTICAL,
409            "Computes the sum of products of paired data points.",
410            "regr_sxy(expression_y, expression_x)",
411        )
412        .with_sql_example(
413            r#"```sql
414create table employee_productivity(week int, productivity_score int) as values(1,60), (2,65), (3,70);
415select * from employee_productivity;
416+------+--------------------+
417| week | productivity_score |
418+------+--------------------+
419| 1    | 60                 |
420| 2    | 65                 |
421| 3    | 70                 |
422+------+--------------------+
423
424SELECT regr_sxy(productivity_score, week) AS sum_product_deviations FROM employee_productivity;
425+------------------------+
426| sum_product_deviations |
427+------------------------+
428|       10.0             |
429+------------------------+
430```
431"#
432        )
433        .with_standard_argument("expression_y", Some("Dependent variable"))
434        .with_standard_argument("expression_x", Some("Independent variable"))
435        .build(),
436    );
437    hash_map
438});
439fn get_regr_docs() -> &'static HashMap<RegrType, Documentation> {
440    &DOCUMENTATION
441}
442
443impl AggregateUDFImpl for Regr {
444    fn name(&self) -> &str {
445        self.func_name
446    }
447
448    fn signature(&self) -> &Signature {
449        &self.signature
450    }
451
452    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
453        if self.regr_type == RegrType::Count {
454            Ok(DataType::UInt64)
455        } else {
456            Ok(DataType::Float64)
457        }
458    }
459
460    fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
461        Ok(Box::new(RegrAccumulator::try_new(&self.regr_type)?))
462    }
463
464    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
465        Ok(vec![
466            Field::new(
467                format_state_name(args.name, "count"),
468                DataType::UInt64,
469                true,
470            ),
471            Field::new(
472                format_state_name(args.name, "mean_x"),
473                DataType::Float64,
474                true,
475            ),
476            Field::new(
477                format_state_name(args.name, "mean_y"),
478                DataType::Float64,
479                true,
480            ),
481            Field::new(
482                format_state_name(args.name, "m2_x"),
483                DataType::Float64,
484                true,
485            ),
486            Field::new(
487                format_state_name(args.name, "m2_y"),
488                DataType::Float64,
489                true,
490            ),
491            Field::new(
492                format_state_name(args.name, "algo_const"),
493                DataType::Float64,
494                true,
495            ),
496        ]
497        .into_iter()
498        .map(Arc::new)
499        .collect())
500    }
501
502    fn documentation(&self) -> Option<&Documentation> {
503        self.regr_type.documentation()
504    }
505}
506
507/// `RegrAccumulator` is used to compute linear regression aggregate functions
508/// by maintaining statistics needed to compute them in an online fashion.
509///
510/// This struct uses Welford's online algorithm for calculating variance and covariance:
511/// <https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm>
512///
513/// Given the statistics, the following aggregate functions can be calculated:
514///
515/// - `regr_slope(y, x)`: Slope of the linear regression line, calculated as:
516///   cov_pop(x, y) / var_pop(x).
517///   It represents the expected change in Y for a one-unit change in X.
518///
519/// - `regr_intercept(y, x)`: Intercept of the linear regression line, calculated as:
520///   mean_y - (regr_slope(y, x) * mean_x).
521///   It represents the expected value of Y when X is 0.
522///
523/// - `regr_count(y, x)`: Count of the non-null(both x and y) input rows.
524///
525/// - `regr_r2(y, x)`: R-squared value (coefficient of determination), calculated as:
526///   (cov_pop(x, y) ^ 2) / (var_pop(x) * var_pop(y)).
527///   It provides a measure of how well the model's predictions match the observed data.
528///
529/// - `regr_avgx(y, x)`: Average of the independent variable X, calculated as: mean_x.
530///
531/// - `regr_avgy(y, x)`: Average of the dependent variable Y, calculated as: mean_y.
532///
533/// - `regr_sxx(y, x)`: Sum of squares of the independent variable X, calculated as:
534///   m2_x.
535///
536/// - `regr_syy(y, x)`: Sum of squares of the dependent variable Y, calculated as:
537///   m2_y.
538///
539/// - `regr_sxy(y, x)`: Sum of products of paired values, calculated as:
540///   algo_const.
541///
542/// Here's how the statistics maintained in this struct are calculated:
543/// - `cov_pop(x, y)`: algo_const / count.
544/// - `var_pop(x)`: m2_x / count.
545/// - `var_pop(y)`: m2_y / count.
546#[derive(Debug)]
547pub struct RegrAccumulator {
548    count: u64,
549    mean_x: f64,
550    mean_y: f64,
551    m2_x: f64,
552    m2_y: f64,
553    algo_const: f64,
554    regr_type: RegrType,
555}
556
557impl RegrAccumulator {
558    /// Creates a new `RegrAccumulator`
559    pub fn try_new(regr_type: &RegrType) -> Result<Self> {
560        Ok(Self {
561            count: 0_u64,
562            mean_x: 0_f64,
563            mean_y: 0_f64,
564            m2_x: 0_f64,
565            m2_y: 0_f64,
566            algo_const: 0_f64,
567            regr_type: regr_type.clone(),
568        })
569    }
570}
571
572impl Accumulator for RegrAccumulator {
573    fn state(&mut self) -> Result<Vec<ScalarValue>> {
574        Ok(vec![
575            ScalarValue::from(self.count),
576            ScalarValue::from(self.mean_x),
577            ScalarValue::from(self.mean_y),
578            ScalarValue::from(self.m2_x),
579            ScalarValue::from(self.m2_y),
580            ScalarValue::from(self.algo_const),
581        ])
582    }
583
584    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
585        // regr_slope(Y, X) calculates k in y = k*x + b
586        let values_y = as_float64_array(&values[0])?;
587        let values_x = as_float64_array(&values[1])?;
588
589        for (value_y, value_x) in values_y.iter().zip(values_x) {
590            // skip either x or y is NULL
591            let (value_y, value_x) = match (value_y, value_x) {
592                (Some(y), Some(x)) => (y, x),
593                // skip either x or y is NULL
594                _ => continue,
595            };
596
597            // Update states for regr_slope(y,x) [using cov_pop(x,y)/var_pop(x)]
598            self.count += 1;
599            let delta_x = value_x - self.mean_x;
600            let delta_y = value_y - self.mean_y;
601            self.mean_x += delta_x / self.count as f64;
602            self.mean_y += delta_y / self.count as f64;
603            let delta_x_2 = value_x - self.mean_x;
604            let delta_y_2 = value_y - self.mean_y;
605            self.m2_x += delta_x * delta_x_2;
606            self.m2_y += delta_y * delta_y_2;
607            self.algo_const += delta_x * (value_y - self.mean_y);
608        }
609
610        Ok(())
611    }
612
613    fn supports_retract_batch(&self) -> bool {
614        true
615    }
616
617    fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
618        let values_y = as_float64_array(&values[0])?;
619        let values_x = as_float64_array(&values[1])?;
620
621        for (value_y, value_x) in values_y.iter().zip(values_x) {
622            // skip either x or y is NULL
623            let (value_y, value_x) = match (value_y, value_x) {
624                (Some(y), Some(x)) => (y, x),
625                // skip either x or y is NULL
626                _ => continue,
627            };
628
629            // Update states for regr_slope(y,x) [using cov_pop(x,y)/var_pop(x)]
630            if self.count > 1 {
631                self.count -= 1;
632                let delta_x = value_x - self.mean_x;
633                let delta_y = value_y - self.mean_y;
634                self.mean_x -= delta_x / self.count as f64;
635                self.mean_y -= delta_y / self.count as f64;
636                let delta_x_2 = value_x - self.mean_x;
637                let delta_y_2 = value_y - self.mean_y;
638                self.m2_x -= delta_x * delta_x_2;
639                self.m2_y -= delta_y * delta_y_2;
640                self.algo_const -= delta_x * (value_y - self.mean_y);
641            } else {
642                self.count = 0;
643                self.mean_x = 0.0;
644                self.m2_x = 0.0;
645                self.m2_y = 0.0;
646                self.mean_y = 0.0;
647                self.algo_const = 0.0;
648            }
649        }
650
651        Ok(())
652    }
653
654    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
655        let count_arr = as_uint64_array(&states[0])?;
656        let mean_x_arr = as_float64_array(&states[1])?;
657        let mean_y_arr = as_float64_array(&states[2])?;
658        let m2_x_arr = as_float64_array(&states[3])?;
659        let m2_y_arr = as_float64_array(&states[4])?;
660        let algo_const_arr = as_float64_array(&states[5])?;
661
662        for i in 0..count_arr.len() {
663            let count_b = count_arr.value(i);
664            if count_b == 0_u64 {
665                continue;
666            }
667            let (count_a, mean_x_a, mean_y_a, m2_x_a, m2_y_a, algo_const_a) = (
668                self.count,
669                self.mean_x,
670                self.mean_y,
671                self.m2_x,
672                self.m2_y,
673                self.algo_const,
674            );
675            let (count_b, mean_x_b, mean_y_b, m2_x_b, m2_y_b, algo_const_b) = (
676                count_b,
677                mean_x_arr.value(i),
678                mean_y_arr.value(i),
679                m2_x_arr.value(i),
680                m2_y_arr.value(i),
681                algo_const_arr.value(i),
682            );
683
684            // Assuming two different batches of input have calculated the states:
685            // batch A of Y, X -> {count_a, mean_x_a, mean_y_a, m2_x_a, algo_const_a}
686            // batch B of Y, X -> {count_b, mean_x_b, mean_y_b, m2_x_b, algo_const_b}
687            // The merged states from A and B are {count_ab, mean_x_ab, mean_y_ab, m2_x_ab,
688            // algo_const_ab}
689            //
690            // Reference for the algorithm to merge states:
691            // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
692            let count_ab = count_a + count_b;
693            let (count_a, count_b) = (count_a as f64, count_b as f64);
694            let d_x = mean_x_b - mean_x_a;
695            let d_y = mean_y_b - mean_y_a;
696            let mean_x_ab = mean_x_a + d_x * count_b / count_ab as f64;
697            let mean_y_ab = mean_y_a + d_y * count_b / count_ab as f64;
698            let m2_x_ab =
699                m2_x_a + m2_x_b + d_x * d_x * count_a * count_b / count_ab as f64;
700            let m2_y_ab =
701                m2_y_a + m2_y_b + d_y * d_y * count_a * count_b / count_ab as f64;
702            let algo_const_ab = algo_const_a
703                + algo_const_b
704                + d_x * d_y * count_a * count_b / count_ab as f64;
705
706            self.count = count_ab;
707            self.mean_x = mean_x_ab;
708            self.mean_y = mean_y_ab;
709            self.m2_x = m2_x_ab;
710            self.m2_y = m2_y_ab;
711            self.algo_const = algo_const_ab;
712        }
713        Ok(())
714    }
715
716    fn evaluate(&mut self) -> Result<ScalarValue> {
717        let cov_pop_x_y = self.algo_const / self.count as f64;
718        let var_pop_x = self.m2_x / self.count as f64;
719        let var_pop_y = self.m2_y / self.count as f64;
720
721        let nullif_or_stat = |cond: bool, stat: f64| {
722            if cond {
723                Ok(ScalarValue::Float64(None))
724            } else {
725                Ok(ScalarValue::Float64(Some(stat)))
726            }
727        };
728
729        match self.regr_type {
730            RegrType::Slope => {
731                // Only 0/1 point or slope is infinite
732                let nullif_cond = self.count <= 1 || var_pop_x == 0.0;
733                nullif_or_stat(nullif_cond, cov_pop_x_y / var_pop_x)
734            }
735            RegrType::Intercept => {
736                let slope = cov_pop_x_y / var_pop_x;
737                // Only 0/1 point or slope is infinite
738                let nullif_cond = self.count <= 1 || var_pop_x == 0.0;
739                nullif_or_stat(nullif_cond, self.mean_y - slope * self.mean_x)
740            }
741            RegrType::Count => Ok(ScalarValue::UInt64(Some(self.count))),
742            RegrType::R2 => {
743                // Only 0/1 point or all x(or y) is the same
744                let nullif_cond = self.count <= 1 || var_pop_x == 0.0 || var_pop_y == 0.0;
745                nullif_or_stat(
746                    nullif_cond,
747                    (cov_pop_x_y * cov_pop_x_y) / (var_pop_x * var_pop_y),
748                )
749            }
750            RegrType::AvgX => nullif_or_stat(self.count < 1, self.mean_x),
751            RegrType::AvgY => nullif_or_stat(self.count < 1, self.mean_y),
752            RegrType::SXX => nullif_or_stat(self.count < 1, self.m2_x),
753            RegrType::SYY => nullif_or_stat(self.count < 1, self.m2_y),
754            RegrType::SXY => nullif_or_stat(self.count < 1, self.algo_const),
755        }
756    }
757
758    fn size(&self) -> usize {
759        size_of_val(self)
760    }
761}