1use 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 Slope,
79 Intercept,
84 Count,
88 R2,
92 AvgX,
96 AvgY,
100 SXX,
104 SYY,
108 SXY,
112}
113
114impl RegrType {
115 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#[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 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 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 let (value_y, value_x) = match (value_y, value_x) {
592 (Some(y), Some(x)) => (y, x),
593 _ => continue,
595 };
596
597 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 let (value_y, value_x) = match (value_y, value_x) {
624 (Some(y), Some(x)) => (y, x),
625 _ => continue,
627 };
628
629 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 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 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 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 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}