diff --git a/src/promql/src/lib.rs b/src/promql/src/lib.rs index f1ac59d792f..ab2a53f79f2 100644 --- a/src/promql/src/lib.rs +++ b/src/promql/src/lib.rs @@ -19,3 +19,4 @@ pub mod grid_series; mod metrics; pub mod range_array; pub mod range_eval; +pub mod select; diff --git a/src/promql/src/select.rs b/src/promql/src/select.rs new file mode 100644 index 00000000000..a9d4035d98e --- /dev/null +++ b/src/promql/src/select.rs @@ -0,0 +1,450 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Operators that depend on the whole set of evaluated series at each grid point: `topk`, +//! `bottomk`, `scalar`, `absent` and the instant-result sorts. +//! +//! Semantics follow Prometheus (`aggregationK`, `funcScalar`, `funcAbsent`, +//! `createLabelsForAbsentFunction`, `funcSort*`). A scalar is a grid series without labels that +//! has a value at every point. + +use std::cmp::Ordering; +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::error::{DataFusionError, Result}; +use datatypes::series::samples::{SampleRef, Series, SeriesSchema}; +use datatypes::series::time::TimeAxis; +use datatypes::series::{Label, MetricName}; + +use crate::grid_series::GridSeriesBuilder; + +#[cfg(test)] +mod tests; + +/// The labels that form an aggregation group. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Grouping { + /// `by(...)`: only these labels, `__name__` included when listed. + By(Vec>), + /// `without(...)`: every label except these and `__name__`. + Without(Vec>), +} + +/// The order `sort_instant` arranges an instant result in. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum SortOrder { + /// `sort`: ascending values, NaN last. + Asc, + /// `sort_desc`: descending values, NaN last. + Desc, + /// `sort_by_label`: natural order of the listed label values, then the label sets. + ByLabel(Vec>), + /// `sort_by_label_desc`. + ByLabelDesc(Vec>), +} + +/// `topk(k, ...)` (`top`) or `bottomk(k, ...)` with a constant `k`. +/// +/// At every point each group keeps the `k` largest (or smallest) values, NaN ranking last for +/// both; the selected samples keep their series' labels and values. Prometheus breaks ties by +/// the order its storage returns series in, which is not a property of the series; here ties +/// go to the series whose label set sorts first. +pub fn topk( + k: f64, + top: bool, + grouping: &Grouping, + schema: &SeriesSchema, + series: &[Series], +) -> Result> { + single_field(schema, "topk/bottomk")?; + let points = grid_points(schema, series)?; + // Prometheus converts k with int64() and caps it at the number of input series. + let k = if k.is_nan() || k < 1.0 { + 0 + } else { + (k.trunc().min(series.len() as f64)) as usize + }; + if k == 0 { + return Ok(Vec::new()); + } + + let mut order: Vec = (0..series.len()).collect(); + order.sort_by(|&a, &b| compare_labels(&series[a].metric, &series[b].metric)); + let mut groups: HashMap> = HashMap::new(); + for index in order { + groups + .entry(group_key(&series[index].metric, grouping)) + .or_default() + .push(index); + } + + let mut selected: Vec> = vec![Vec::new(); series.len()]; + let mut candidates: Vec<(usize, f64)> = Vec::new(); + for members in groups.values() { + for point in 0..points { + candidates.clear(); + for &member in members { + if let Some(value) = value_at(&series[member], point)? { + candidates.push((member, value)); + } + } + // Members are in label order, so a stable sort hands ties to the earlier label set. + candidates.sort_by(|a, b| rank(a.1, b.1, top)); + for &(member, _) in candidates.iter().take(k) { + selected[member].push(point); + } + } + } + + let output = Arc::new(schema.clone()); + let mut builder = GridSeriesBuilder::try_new(output)?; + let mut result = Vec::new(); + for (index, points) in selected.into_iter().enumerate() { + if points.is_empty() { + continue; + } + let mut sorted = points; + sorted.sort_unstable(); + for point in sorted { + builder.push(point, &[value_at(&series[index], point)?])?; + } + if let Some(built) = builder.finish(series[index].metric.clone())? { + result.push(built); + } + } + Ok(result) +} + +/// `scalar(v)`: the value of the only series present at each point, NaN otherwise. +pub fn scalar(schema: &SeriesSchema, series: &[Series]) -> Result { + single_field(schema, "scalar")?; + let points = grid_points(schema, series)?; + let mut values: Vec> = vec![None; points]; + let mut counts = vec![0u8; points]; + for item in series { + let TimeAxis::Grid(range) = &item.time else { + unreachable!("checked by grid_points"); + }; + for point in range.start()..range.start() + range.len() { + if let Some(value) = value_at(item, point)? { + counts[point] = counts[point].saturating_add(1); + values[point] = Some(value); + } + } + } + let mut builder = GridSeriesBuilder::try_new(Arc::new(schema.clone()))?; + for (point, (value, count)) in values.into_iter().zip(counts).enumerate() { + let value = if count == 1 { value } else { Some(GO_NAN) }; + builder.push(point, &[value])?; + } + let empty = MetricName::try_new(None, Vec::new()).map_err(external)?; + Ok(builder + .finish(empty)? + .expect("every point of a scalar has a value")) +} + +/// `absent(v)`: 1 with `labels` at every point where no series is present. +pub fn absent( + schema: &SeriesSchema, + series: &[Series], + labels: MetricName, +) -> Result> { + let points = grid_points(schema, series)?; + let mut present = vec![false; points]; + for item in series { + let TimeAxis::Grid(range) = &item.time else { + unreachable!("checked by grid_points"); + }; + let points = &mut present[range.start()..range.start() + range.len()]; + for (row, present) in points.iter_mut().enumerate() { + if item + .fields + .iter() + .any(|field| !matches!(field.sample(row), SampleRef::Missing)) + { + *present = true; + } + } + } + let ones = vec![Some(1.0); schema.fields.len()]; + let mut builder = GridSeriesBuilder::try_new(Arc::new(schema.clone()))?; + for (point, present) in present.into_iter().enumerate() { + if !present { + builder.push(point, &ones)?; + } + } + builder.finish(labels) +} + +/// A label matcher of the selector `absent` is applied to. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LabelMatcher { + pub name: Arc, + pub op: MatchOp, + pub value: Arc, + /// Whether the matcher accepts an empty value, i.e. a missing label; the caller evaluates + /// regular expressions. + pub matches_empty: bool, +} + +/// The operator of a [`LabelMatcher`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MatchOp { + Equal, + NotEqual, + Regex, + NotRegex, +} + +/// The labels of an `absent` result, as Prometheus `createLabelsForAbsentFunction` derives them: +/// the first equality matcher of a name sets that label, any other matcher of the name removes +/// it, and `__name__` is skipped. +/// +/// Prometheus reads the matchers after TSDB's `PostingsForMatchers` has reordered them in place, +/// moving the matchers that narrow the result before those that only subtract from it; that is +/// the order users see whenever the storage holds data, and the order used here. It changes the +/// result only when a name has both an equality and another matcher. +pub fn absent_labels(matchers: &[LabelMatcher]) -> Result { + let must_be_set: Vec<&str> = matchers + .iter() + .filter(|matcher| !matcher.matches_empty) + .map(|matcher| &*matcher.name) + .collect(); + let subtracting = |matcher: &LabelMatcher| { + !must_be_set.contains(&&*matcher.name) + || (matches!(matcher.op, MatchOp::NotEqual | MatchOp::NotRegex) + && matcher.matches_empty) + }; + let ordered = matchers + .iter() + .filter(|matcher| !subtracting(matcher)) + .chain(matchers.iter().filter(|matcher| subtracting(matcher))); + + let mut labels: Vec