feat(promql): add stddev, stdvar, quantile and count_values aggregates

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

src/promql/src/aggregate.rs

@@ -19,6 +19,7 @@
 #[cfg(test)]
 mod tests;
 
+use std::collections::hash_map::Entry;
 use std::collections::{BTreeMap, HashMap};
 use std::sync::Arc;
 
@@ -26,10 +27,11 @@ use datafusion::error::{DataFusionError, Result};
 use datatypes::series::MetricName;
 use datatypes::series::samples::{SampleRef, Series, SeriesSchema};
 
-use crate::grid_series::GridSeriesBuilder;
-use crate::map::{drop_empty_labels, validate_float_series};
+use crate::grid_series::{GridSeriesBuilder, external};
+use crate::map::{drop_empty_labels, validate_float_series, validate_scalar};
+use crate::range_eval::quantile;
 
-/// Basic float aggregations.
+/// Float aggregations without a parameter.
 #[derive(Clone, Copy, Debug)]
 pub enum AggregateOp {
     Sum,
@@ -38,6 +40,8 @@ pub enum AggregateOp {
     Min,
     Max,
     Group,
+    Stddev,
+    Stdvar,
 }
 
 /// Grouping by series labels. No modifier is `By(&[])`, as in Prometheus.
@@ -66,19 +70,150 @@ pub fn aggregate_float(
     grouping: Grouping<'_>,
 ) -> Result<Vec<Series>> {
     validate_input(schema, input, grouping)?;
+    let groups = group_series(input, grouping, |state: &mut FloatState, sample| {
+        state.update(op, sample)
+    })?;
+    build_output(
+        schema,
+        groups,
+        |state| state.seen,
+        |state, _| state.finish(op),
+    )
+}
+
+/// `quantile(q, v)`, with `q` read at each grid point from an unlabelled single-field scalar on
+/// the same grid, as Prometheus evaluates the parameter at every step. A point where `q` has no
+/// value has no result; `q` outside `[0, 1]` or NaN gives what Prometheus' `quantile` gives.
+pub fn aggregate_quantile(
+    schema: &Arc<SeriesSchema>,
+    input: &[Series],
+    q: &Series,
+    grouping: Grouping<'_>,
+) -> Result<Vec<Series>> {
+    validate_input(schema, input, grouping)?;
+    let q_axis = validate_scalar(schema, q)?;
+    let q: Vec<(usize, f64)> = q_axis
+        .positions()
+        .zip(q.fields.as_slice()[0].iter())
+        .filter_map(|((point, _), sample)| match sample {
+            SampleRef::Float(value) => Some((point, value)),
+            _ => None,
+        })
+        .collect();
+    let groups = group_series(input, grouping, |values: &mut Vec<f64>, sample| {
+        values.push(sample)
+    })?;
+    let mut sorted = Vec::new();
+    build_output(
+        schema,
+        groups,
+        |values| !values.is_empty(),
+        |values, point| {
+            let index = q.binary_search_by_key(&point, |&(point, _)| point).ok()?;
+            Some(quantile(q[index].1, values, &mut sorted))
+        },
+    )
+}
+
+/// `count_values(label, v)`: at each grid point, the number of samples in each group of series
+/// that share a value, with the value written to `label` in Prometheus' `FormatFloat(v, 'f', -1)`
+/// form. With `by`, `label` is part of the grouping; with `without`, it is dropped if listed.
+pub fn count_values(
+    schema: &Arc<SeriesSchema>,
+    input: &[Series],
+    label: &str,
+    grouping: Grouping<'_>,
+) -> Result<Vec<Series>> {
+    validate_input(schema, input, grouping)?;
+    if label.is_empty() {
+        return Err(DataFusionError::Execution(format!(
+            "invalid label name {label:?}"
+        )));
+    }
+    let by_labels: Vec<&str>;
+    let grouping = match grouping {
+        Grouping::By(labels) => {
+            by_labels = labels.iter().copied().chain([label]).collect();
+            Grouping::By(&by_labels)
+        }
+        without => without,
+    };
+    let width = schema.fields.len();
+    let mut group_index = HashMap::new();
+    let mut groups: Vec<Group<f64>> = Vec::new();
+    for series in input {
+        // Equal bits format to the same label value, so each distinct value of a series
+        // resolves its group once.
+        let mut group_of_value: HashMap<u64, usize> = HashMap::new();
+        for (field_index, field) in series.fields.iter().enumerate() {
+            for ((point, _), sample) in series.time.positions().zip(field.iter()) {
+                let SampleRef::Float(value) = sample else {
+                    continue;
+                };
+                let group = match group_of_value.entry(value.to_bits()) {
+                    Entry::Occupied(entry) => *entry.get(),
+                    Entry::Vacant(entry) => {
+                        let mut metric = series.metric.clone();
+                        metric
+                            .set(label, Arc::from(format_value(value)))
+                            .map_err(external)?;
+                        let metric = grouping.apply(&metric)?;
+                        let index = *group_index.entry(label_key(&metric)).or_insert_with(|| {
+                            groups.push(Group::new(metric));
+                            groups.len() - 1
+                        });
+                        *entry.insert(index)
+                    }
+                };
+                groups[group].update_point(point, width, field_index, |count| *count += 1.0);
+            }
+        }
+    }
+    build_output(
+        schema,
+        groups,
+        |count| *count > 0.0,
+        |count, _| Some(*count),
+    )
+}
+
+/// Go's `strconv.FormatFloat(value, 'f', -1, 64)`; Rust's shortest round-trip decimal matches it
+/// for finite values.
+fn format_value(value: f64) -> String {
+    match value {
+        f64::INFINITY => "+Inf".to_string(),
+        f64::NEG_INFINITY => "-Inf".to_string(),
+        _ => value.to_string(),
+    }
+}
+
+/// Groups `input` by `grouping` in first-seen order and folds each series into its group.
+fn group_series<S: Clone + Default>(
+    input: &[Series],
+    grouping: Grouping<'_>,
+    mut update: impl FnMut(&mut S, f64),
+) -> Result<Vec<Group<S>>> {
     let mut group_index = HashMap::new();
     let mut groups = Vec::new();
     for series in input {
         let metric = grouping.apply(&series.metric)?;
         let index = *group_index.entry(label_key(&metric)).or_insert_with(|| {
-            groups.push(Group {
-                metric,
-                blocks: BTreeMap::new(),
-            });
+            groups.push(Group::new(metric));
             groups.len() - 1
         });
-        groups[index].accumulate(series, op)?;
+        groups[index].accumulate(series, &mut update);
     }
+    Ok(groups)
+}
+
+/// Writes one series per group, from the points where any field's state has a sample. `finish`
+/// receives the absolute grid index.
+fn build_output<S>(
+    schema: &Arc<SeriesSchema>,
+    groups: Vec<Group<S>>,
+    seen: impl Fn(&S) -> bool,
+    mut finish: impl FnMut(&S, usize) -> Option<f64>,
+) -> Result<Vec<Series>> {
     let mut builder = GridSeriesBuilder::try_new(schema.clone())?;
     let mut output = Vec::with_capacity(groups.len());
     let mut row = Vec::with_capacity(schema.fields.len());
@@ -88,12 +223,19 @@ pub fn aggregate_float(
         for (&block, states) in &group.blocks {
             // The last block can extend past the grid; only points with a sample are pushed.
             for (offset, point) in states.chunks_exact(schema.fields.len()).enumerate() {
-                if !point.iter().any(|state| state.seen) {
+                if !point.iter().any(&seen) {
                     continue;
                 }
+                let index = block * BLOCK_POINTS + offset;
                 row.clear();
-                row.extend(point.iter().map(|state| state.finish(op)));
-                builder.push(block * BLOCK_POINTS + offset, &row)?;
+                row.extend(point.iter().map(|state| {
+                    if seen(state) {
+                        finish(state, index)
+                    } else {
+                        None
+                    }
+                }));
+                builder.push(index, &row)?;
             }
         }
         output.extend(builder.finish(group.metric)?);
@@ -162,15 +304,22 @@ fn label_key(metric: &MetricName) -> Vec<(Arc<str>, Arc<str>)> {
 /// input never expands to the full grid and dense input costs one allocation per block.
 const BLOCK_POINTS: usize = 64;
 
-struct Group {
+struct Group<S> {
     metric: MetricName,
     /// Keyed by absolute grid index / `BLOCK_POINTS`; each block holds one state per field for
     /// each of its points, point by point.
-    blocks: BTreeMap<usize, Box<[FloatState]>>,
+    blocks: BTreeMap<usize, Box<[S]>>,
 }
 
-impl Group {
-    fn accumulate(&mut self, series: &Series, op: AggregateOp) -> Result<()> {
+impl<S: Clone + Default> Group<S> {
+    fn new(metric: MetricName) -> Self {
+        Self {
+            metric,
+            blocks: BTreeMap::new(),
+        }
+    }
+
+    fn accumulate(&mut self, series: &Series, update: &mut impl FnMut(&mut S, f64)) {
         let width = series.fields.len();
         let mut pending = Vec::with_capacity(BLOCK_POINTS);
         for (field_index, field) in series.fields.iter().enumerate() {
@@ -181,16 +330,15 @@ impl Group {
             let mut block = usize::MAX;
             for ((index, _), sample) in series.time.positions().zip(field.iter()) {
                 if index / BLOCK_POINTS != block {
-                    self.update_block(block, width, field_index, op, &mut pending);
+                    self.update_block(block, width, field_index, update, &mut pending);
                     block = index / BLOCK_POINTS;
                 }
                 if let SampleRef::Float(sample) = sample {
                     pending.push((index % BLOCK_POINTS, sample));
                 }
             }
-            self.update_block(block, width, field_index, op, &mut pending);
+            self.update_block(block, width, field_index, update, &mut pending);
         }
-        Ok(())
     }
 
     /// Applies one field's samples within a block, so a block costs one lookup per field run.
@@ -199,19 +347,35 @@ impl Group {
         block: usize,
         width: usize,
         field_index: usize,
-        op: AggregateOp,
+        update: &mut impl FnMut(&mut S, f64),
         pending: &mut Vec<(usize, f64)>,
     ) {
         if pending.is_empty() {
             return;
         }
-        let states = self.blocks.entry(block).or_insert_with(|| {
-            vec![FloatState::default(); BLOCK_POINTS * width].into_boxed_slice()
-        });
+        let states = self.block(block, width);
         for (offset, sample) in pending.drain(..) {
-            states[offset * width + field_index].update(op, sample);
+            update(&mut states[offset * width + field_index], sample);
         }
     }
+
+    /// Updates the state of one field at absolute grid `point`.
+    fn update_point(
+        &mut self,
+        point: usize,
+        width: usize,
+        field_index: usize,
+        update: impl FnOnce(&mut S),
+    ) {
+        let states = self.block(point / BLOCK_POINTS, width);
+        update(&mut states[point % BLOCK_POINTS * width + field_index]);
+    }
+
+    fn block(&mut self, block: usize, width: usize) -> &mut [S] {
+        self.blocks
+            .entry(block)
+            .or_insert_with(|| vec![S::default(); BLOCK_POINTS * width].into_boxed_slice())
+    }
 }
 
 /// The float part of Prometheus' `groupedAggregation`.
@@ -230,10 +394,13 @@ impl FloatState {
         if !self.seen {
             *self = Self {
                 seen: true,
-                value: if matches!(op, AggregateOp::Group) {
-                    1.0
-                } else {
-                    sample
+                value: match op {
+                    AggregateOp::Group => 1.0,
+                    // The running sum of squared deviations; Prometheus starts it at NaN for a
+                    // non-finite first sample.
+                    AggregateOp::Stddev | AggregateOp::Stdvar if !sample.is_finite() => f64::NAN,
+                    AggregateOp::Stddev | AggregateOp::Stdvar => 0.0,
+                    _ => sample,
                 },
                 mean: sample,
                 count: 1.0,
@@ -258,6 +425,13 @@ impl FloatState {
                 }
             }
             AggregateOp::Group => {}
+            // Welford's update, uncompensated as in Prometheus.
+            AggregateOp::Stddev | AggregateOp::Stdvar => {
+                self.count += 1.0;
+                let delta = sample - self.mean;
+                self.mean += delta / self.count;
+                self.value += delta * (sample - self.mean);
+            }
         }
     }
 
@@ -289,6 +463,8 @@ impl FloatState {
             AggregateOp::Avg => self.value / self.count + self.compensation / self.count,
             AggregateOp::Count => self.count,
             AggregateOp::Min | AggregateOp::Max | AggregateOp::Group => self.value,
+            AggregateOp::Stdvar => self.value / self.count,
+            AggregateOp::Stddev => (self.value / self.count).sqrt(),
         })
     }
 }

src/promql/src/aggregate/tests.rs

@@ -199,6 +199,42 @@ fn aggregates_match_prometheus() {
     }
 }
 
+/// `stddev`, `stdvar`, `quantile` and `count_values` evaluated by the Prometheus v3.15.0 engine;
+/// see the header of `stat_cases.txt`.
+#[test]
+fn stat_aggregates_match_prometheus() {
+    let mut matched = BTreeMap::<String, usize>::new();
+    for case in cases(include_str!("stat_cases.txt")) {
+        let h = &case.header;
+        let labels = h[3]
+            .split(',')
+            .filter(|label| *label != "-")
+            .collect::<Vec<_>>();
+        let grouping = grouping(h, &labels);
+        let output = match h[1].as_str() {
+            "stddev" => aggregate_float(&case.schema, &case.series, AggregateOp::Stddev, grouping),
+            "stdvar" => aggregate_float(&case.schema, &case.series, AggregateOp::Stdvar, grouping),
+            "quantile" => {
+                let q = case.scalar.as_ref().expect("quantile cases carry q");
+                aggregate_quantile(&case.schema, &case.series, q, grouping)
+            }
+            "count_values" => count_values(&case.schema, &case.series, &h[4], grouping),
+            other => panic!("unmapped aggregate {other}"),
+        };
+        assert_result(&case, output.unwrap());
+        *matched.entry(format!("{} {}", h[1], h[2])).or_default() += 1;
+    }
+    for op in ["stddev", "stdvar", "quantile", "count_values"] {
+        for grouping in ["none", "by", "without"] {
+            let count = matched
+                .get(&format!("{op} {grouping}"))
+                .copied()
+                .unwrap_or_default();
+            assert!(count >= 12, "{op} {grouping}: {count}");
+        }
+    }
+}
+
 #[test]
 fn compensated_sum_and_overflowing_avg_follow_prometheus() {
     let schema = schema(&["value"]);
@@ -351,18 +387,50 @@ fn fields_are_aggregated_independently() {
             })
             .collect()
     };
+    let grid_len = multi.eval_grid.as_ref().unwrap().timestamps().len();
+    let q_values: Vec<Option<f64>> = (0..grid_len)
+        .map(|i| Some([0.0, 0.25, 0.5, 0.9, 1.0][i % 5]))
+        .collect();
+    let q = input(
+        &schema(&["scalar"]),
+        MetricName::try_new(None, Vec::new()).unwrap(),
+        0,
+        &[q_values],
+    );
+    type Kernel<'a> = Box<dyn Fn(&Arc<SeriesSchema>, &[Series], Grouping<'_>) -> Vec<Series> + 'a>;
+    let mut kernels: Vec<(String, Kernel<'_>)> = OPS
+        .iter()
+        .chain(&[AggregateOp::Stddev, AggregateOp::Stdvar])
+        .map(|&op| -> (String, Kernel<'_>) {
+            (
+                format!("{op:?}"),
+                Box::new(move |schema, input, grouping| {
+                    aggregate_float(schema, input, op, grouping).unwrap()
+                }),
+            )
+        })
+        .collect();
+    kernels.push((
+        "quantile".into(),
+        Box::new(|schema, input, grouping| {
+            aggregate_quantile(schema, input, &q, grouping).unwrap()
+        }),
+    ));
+    kernels.push((
+        "count_values".into(),
+        Box::new(|schema, input, grouping| count_values(schema, input, "value", grouping).unwrap()),
+    ));
     for grouping in [
         Grouping::By(&[]),
         Grouping::By(&["host"]),
         Grouping::Without(&["i"]),
     ] {
-        for op in OPS {
-            let output = aggregate_float(&multi, &series, op, grouping).unwrap();
+        for (op, kernel) in &kernels {
+            let output = kernel(&multi, &series, grouping);
             for field in [0, 1] {
-                let expected = aggregate_float(&single, &project(field), op, grouping).unwrap();
+                let expected = kernel(&single, &project(field), grouping);
                 // The multi-field result is cut into pieces where its values are interrupted, so
                 // the comparison is over the dense values of each label set, not over pieces.
-                let grid_len = multi.eval_grid.as_ref().unwrap().timestamps().len();
                 let grouped = |pieces: &[Series], field: usize| {
                     let mut by_metric: BTreeMap<LabelKey, Positions> = BTreeMap::new();
                     for piece in pieces {
@@ -394,7 +462,7 @@ fn fields_are_aggregated_independently() {
                 };
                 let actual = grouped(&output, field);
                 let expected = grouped(&expected, 0);
-                assert_eq!(actual, expected, "{op:?} {grouping:?} field {field}");
+                assert_eq!(actual, expected, "{op} {grouping:?} field {field}");
             }
         }
     }
@@ -632,6 +700,16 @@ fn unsupported_shapes_and_mutated_schema_fail_without_partial_output() {
         11,
         &[vec![Some(1.0)], vec![Some(2.0)]],
     );
+    // Prometheus rejects an empty `count_values` label name at evaluation.
+    assert!(matches!(
+        count_values(
+            &schema,
+            std::slice::from_ref(&source),
+            "",
+            Grouping::By(&[])
+        ),
+        Err(DataFusionError::Execution(_))
+    ));
     assert!(matches!(
         aggregate_float(
             &schema,

src/promql/src/map.rs

@@ -23,6 +23,7 @@ mod go_math;
 #[cfg(test)]
 mod tests;
 
+pub(crate) use broadcast::validate_scalar;
 pub use broadcast::{ScalarSeriesOperand, map_arithmetic_series, map_compare_series};
 pub use comparison::{Comparison, ComparisonOp, map_compare, scalar_arithmetic, scalar_compare};
 

src/promql/src/map/broadcast.rs

@@ -127,7 +127,10 @@ fn map_grid_scalar(
     builder.finish(metric)
 }
 
-fn validate_scalar<'a>(schema: &SeriesSchema, scalar: &'a Series) -> Result<&'a TimeAxis> {
+pub(crate) fn validate_scalar<'a>(
+    schema: &SeriesSchema,
+    scalar: &'a Series,
+) -> Result<&'a TimeAxis> {
     if scalar.metric.name().is_some() || !scalar.metric.tags().is_empty() {
         return Err(DataFusionError::Execution(
             "grid scalar must be unlabelled".into(),

src/promql/src/range_eval.rs

@@ -39,6 +39,7 @@ mod histogram;
 mod tests;
 
 pub use function::RangeFunction;
+pub(crate) use function::quantile;
 
 /// The samples a range selector reads at evaluation time `T`: source timestamps in
 /// `(E - range_ms, E]`, where `E` is `T - offset_ms`, or `anchor_ms - offset_ms` with `@`.

src/promql/src/range_eval/function.rs

@@ -216,7 +216,7 @@ fn variance(values: &[f64]) -> f64 {
 }
 
 /// Prometheus `quantile`, which sorts NaN before every number.
-fn quantile(q: f64, values: &[f64], sorted: &mut Vec<f64>) -> f64 {
+pub(crate) fn quantile(q: f64, values: &[f64], sorted: &mut Vec<f64>) -> f64 {
     if q.is_nan() {
         return f64::NAN;
     }