From 6a05f617a40ae28fb8b1d22c42c40d0ce7b17ac6 Mon Sep 17 00:00:00 2001
From: Eugene Tolbakov <ev.tolbakov@gmail.com>
Date: Sat, 1 Apr 2023 10:16:51 +0100
Subject: [PATCH] feat(stddev_over_time): add initial implementation (#1289)

* feat(stddev_over_time): add initial implementation

* feat(stddev_over_time): address code review remarks, add compensated summation

* feat(stddev_over_time): fix fmt issues

* feat(stddev_over_time): add docs, minor renamings
---
 src/promql/src/functions.rs                | 18 ++++-
 src/promql/src/functions/aggr_over_time.rs | 85 +++++++++++++++++++++-
 2 files changed, 100 insertions(+), 3 deletions(-)

diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs
index ceba5c6181c..f83c4e5d071 100644
--- a/src/promql/src/functions.rs
+++ b/src/promql/src/functions.rs
@@ -23,7 +23,7 @@ mod test_util;
 
 pub use aggr_over_time::{
     AbsentOverTime, AvgOverTime, CountOverTime, LastOverTime, MaxOverTime, MinOverTime,
-    PresentOverTime, SumOverTime,
+    PresentOverTime, StddevOverTime, SumOverTime,
 };
 use datafusion::arrow::array::ArrayRef;
 use datafusion::error::DataFusionError;
@@ -40,3 +40,19 @@ pub(crate) fn extract_array(columnar_value: &ColumnarValue) -> Result<ArrayRef,
         ))
     }
 }
+
+/// compensation(Kahan) summation algorithm - a technique for reducing the numerical error
+/// in floating-point arithmetic. The algorithm also includes the modification ("Neumaier improvement")
+/// that reduces the numerical error further in cases
+/// where the numbers being summed have a large difference in magnitude
+/// Prometheus's implementation:
+/// https://github.com/prometheus/prometheus/blob/f55ab2217984770aa1eecd0f2d5f54580029b1c0/promql/functions.go#L782)
+pub(crate) fn compensated_sum_inc(inc: f64, sum: f64, mut compensation: f64) -> (f64, f64) {
+    let new_sum = sum + inc;
+    if sum.abs() >= inc.abs() {
+        compensation += (sum - new_sum) + inc;
+    } else {
+        compensation += (inc - new_sum) + sum;
+    }
+    (new_sum, compensation)
+}
diff --git a/src/promql/src/functions/aggr_over_time.rs b/src/promql/src/functions/aggr_over_time.rs
index 451008d4208..be164a60f0c 100644
--- a/src/promql/src/functions/aggr_over_time.rs
+++ b/src/promql/src/functions/aggr_over_time.rs
@@ -24,7 +24,7 @@ use datatypes::arrow::array::Array;
 use datatypes::arrow::compute;
 use datatypes::arrow::datatypes::DataType;
 
-use crate::functions::extract_array;
+use crate::functions::{compensated_sum_inc, extract_array};
 use crate::range_array::RangeArray;
 
 /// The average value of all points in the specified interval.
@@ -117,7 +117,42 @@ pub fn present_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -
     }
 }
 
-// TODO(ruihang): support quantile_over_time, stddev_over_time, and stdvar_over_time
+// TODO(ruihang): support quantile_over_time, and stdvar_over_time
+
+/// the population standard deviation of the values in the specified interval.
+/// Prometheus's implementation: https://github.com/prometheus/prometheus/blob/f55ab2217984770aa1eecd0f2d5f54580029b1c0/promql/functions.go#L556-L569
+#[range_fn(
+    name = "StddevOverTime",
+    ret = "Float64Array",
+    display_name = "prom_stddev_over_time"
+)]
+pub fn stddev_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option<f64> {
+    if values.is_empty() {
+        None
+    } else {
+        let mut count = 0.0;
+        let mut mean = 0.0;
+        let mut comp_mean = 0.0;
+        let mut deviations_sum_sq = 0.0;
+        let mut comp_deviations_sum_sq = 0.0;
+        for v in values {
+            count += 1.0;
+            let current_value = v.unwrap();
+            let delta = current_value - (mean + comp_mean);
+            let (new_mean, new_comp_mean) = compensated_sum_inc(delta / count, mean, comp_mean);
+            mean = new_mean;
+            comp_mean = new_comp_mean;
+            let (new_deviations_sum_sq, new_comp_deviations_sum_sq) = compensated_sum_inc(
+                delta * (current_value - (mean + comp_mean)),
+                deviations_sum_sq,
+                comp_deviations_sum_sq,
+            );
+            deviations_sum_sq = new_deviations_sum_sq;
+            comp_deviations_sum_sq = new_comp_deviations_sum_sq;
+        }
+        Some(((deviations_sum_sq + comp_deviations_sum_sq) / count).sqrt())
+    }
+}
 
 #[cfg(test)]
 mod test {
@@ -332,4 +367,50 @@ mod test {
             ],
         );
     }
+
+    #[test]
+    fn calculate_std_dev_over_time() {
+        let (ts_array, value_array) = build_test_range_arrays();
+        simple_range_udf_runner(
+            StddevOverTime::scalar_udf(),
+            ts_array,
+            value_array,
+            vec![
+                Some(37.6543215),
+                Some(28.442923895289123),
+                Some(0.0),
+                None,
+                None,
+                Some(18.12081352042062),
+                Some(11.983172291869804),
+                Some(11.441953741554055),
+                Some(0.0),
+                None,
+            ],
+        );
+
+        // add more assertions
+        let ts_array = Arc::new(TimestampMillisecondArray::from_iter(
+            [1000i64, 3000, 5000, 7000, 9000, 11000, 13000, 15000]
+                .into_iter()
+                .map(Some),
+        ));
+        let values_array = Arc::new(Float64Array::from_iter([
+            1.5990505637277868,
+            1.5990505637277868,
+            1.5990505637277868,
+            0.0,
+            8.0,
+            8.0,
+            2.0,
+            3.0,
+        ]));
+        let ranges = [(0, 3), (3, 5)];
+        simple_range_udf_runner(
+            StddevOverTime::scalar_udf(),
+            RangeArray::from_ranges(ts_array, ranges).unwrap(),
+            RangeArray::from_ranges(values_array, ranges).unwrap(),
+            vec![Some(0.0), Some(3.249615361854384)],
+        );
+    }
 }