From a7830e69ae8faac84a3510fcabab178a5caba3e7 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Sat, 17 Oct 2026 17:56:41 +0800 Subject: [PATCH] feat(promql): match evaluated vectors for binary and set operators Signed-off-by: Dennis Zhuang --- src/promql/src/lib.rs | 1 + src/promql/src/vector_match.rs | 617 +++++++++++++++++++++++++++ src/promql/src/vector_match/tests.rs | 497 +++++++++++++++++++++ 3 files changed, 1115 insertions(+) create mode 100644 src/promql/src/vector_match.rs create mode 100644 src/promql/src/vector_match/tests.rs diff --git a/src/promql/src/lib.rs b/src/promql/src/lib.rs index f1ac59d792f..8727aafbb2d 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 vector_match; diff --git a/src/promql/src/vector_match.rs b/src/promql/src/vector_match.rs new file mode 100644 index 00000000000..db62feedd50 --- /dev/null +++ b/src/promql/src/vector_match.rs @@ -0,0 +1,617 @@ +// 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. + +//! Vector–vector binary and set operations over evaluated series on one grid. +//! +//! Semantics follow Prometheus `VectorBinop`, `VectorAnd`, `VectorOr` and `VectorUnless` +//! evaluated at every grid point. Series are processed whole, over their own grid ranges, and +//! the cardinality rules are checked per point: series sharing a match group only conflict at +//! points where both are present. + +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::sync::Arc; + +use datafusion::arrow::array::{Array, Float64Array}; +use datafusion::error::{DataFusionError, Result}; +use datatypes::series::MetricName; +use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, Series, SeriesSchema}; +use datatypes::series::time::{EvalGrid, TimeAxis}; + +use crate::grid_series::GridSeriesBuilder; + +#[cfg(test)] +mod tests; + +/// A PromQL binary operator between two instant vectors. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BinaryOp { + Add, + Sub, + Mul, + Div, + Mod, + Pow, + Atan2, + Eq, + Ne, + Gt, + Lt, + Ge, + Le, +} + +impl BinaryOp { + fn is_comparison(self) -> bool { + matches!( + self, + Self::Eq | Self::Ne | Self::Gt | Self::Lt | Self::Ge | Self::Le + ) + } + + /// The value and whether a comparison holds, as Prometheus `vectorElemBinop` computes them. + fn apply(self, lhs: f64, rhs: f64) -> (f64, bool) { + match self { + Self::Add => (lhs + rhs, true), + Self::Sub => (lhs - rhs, true), + Self::Mul => (lhs * rhs, true), + Self::Div => (lhs / rhs, true), + Self::Mod => (lhs % rhs, true), + Self::Pow => (lhs.powf(rhs), true), + Self::Atan2 => (lhs.atan2(rhs), true), + Self::Eq => (lhs, lhs == rhs), + Self::Ne => (lhs, lhs != rhs), + Self::Gt => (lhs, lhs > rhs), + Self::Lt => (lhs, lhs < rhs), + Self::Ge => (lhs, lhs >= rhs), + Self::Le => (lhs, lhs <= rhs), + } + } +} + +/// The labels two series are matched by. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum MatchingLabels { + /// `on(...)`: only these labels, `__name__` included when listed. + On(Vec>), + /// `ignoring(...)`, and the default with no labels: every label except these and + /// `__name__`. + Ignoring(Vec>), +} + +/// How many series of each side may share a match group. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Cardinality { + OneToOne, + /// `group_left(...)`: the right side is the "one" side; the listed labels are copied from it. + ManyToOne(Vec>), + /// `group_right(...)`: the left side is the "one" side; the listed labels are copied from it. + OneToMany(Vec>), +} + +/// A vector–vector binary operation with its modifiers. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BinaryMatching { + pub op: BinaryOp, + /// `bool`: comparisons return 1 or 0 instead of filtering. + pub return_bool: bool, + pub labels: MatchingLabels, + pub cardinality: Cardinality, +} + +/// A PromQL set operator between two instant vectors. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SetOp { + And, + Or, + Unless, +} + +/// The schema of `binary_vectors` output: fields paired by position, named after the left side. +pub fn binary_output_schema(lhs: &SeriesSchema, rhs: &SeriesSchema) -> Result> { + let grid = same_grid(lhs, rhs)?; + let fields = lhs.fields.len().min(rhs.fields.len()); + Ok(Arc::new(SeriesSchema { + fields: lhs.fields[..fields].to_vec().into(), + eval_grid: Some(grid), + })) +} + +/// Applies `matching` to every pair of matched series at every grid point. +/// +/// Fields are paired by position and computed independently; a point is present in a series +/// when any of its fields has a value. Arithmetic and `bool` comparisons drop `__name__`; +/// comparisons without `bool` keep the left value where they hold. +pub fn binary_vectors( + matching: &BinaryMatching, + lhs_schema: &SeriesSchema, + lhs: &[Series], + rhs_schema: &SeriesSchema, + rhs: &[Series], +) -> Result> { + let output = binary_output_schema(lhs_schema, rhs_schema)?; + let grid = output + .eval_grid + .clone() + .expect("set by binary_output_schema"); + let lhs = Operand::try_new(lhs_schema, lhs)?; + let rhs = Operand::try_new(rhs_schema, rhs)?; + let (many, one, include, one_is_left) = match &matching.cardinality { + Cardinality::OneToOne => (&lhs, &rhs, &[][..], false), + Cardinality::ManyToOne(include) => (&lhs, &rhs, &include[..], false), + Cardinality::OneToMany(include) => (&rhs, &lhs, &include[..], true), + }; + let points = grid.timestamps().len(); + // Prometheus short-circuits a point where either side is empty, before any check. + let both_present: Vec = lhs + .presence(points)? + .into_iter() + .zip(rhs.presence(points)?) + .map(|(lhs, rhs)| lhs && rhs) + .collect(); + + let one_groups = group(one, &matching.labels); + for (key, members) in &one_groups { + if let Some(point) = overlap(one, members, |point| Ok(both_present[point]))? { + let side = if one_is_left { "left" } else { "right" }; + return Err(DataFusionError::Execution(format!( + "found duplicate series for the match group {} on the {side} hand-side of the \ + operation at {} ms; many-to-many matching not allowed: matching labels must be \ + unique on one side", + render(key), + grid.timestamps()[point] + ))); + } + } + + let fields = output.fields.len(); + let drop_name = !matching.op.is_comparison() || matching.return_bool; + let mut results = Results::default(); + let mut matched_groups: HashMap> = HashMap::new(); + let mut values = vec![None; fields]; + for (many_index, many_series) in many.series.iter().enumerate() { + let key = group_key(&many_series.metric, &matching.labels); + let Some(candidates) = one_groups.get(&key) else { + continue; + }; + if matches!(matching.cardinality, Cardinality::OneToOne) { + matched_groups.entry(key).or_default().push(many_index); + } + for &one_index in candidates { + let metric = result_metric( + &many_series.metric, + &one.series[one_index].metric, + matching, + include, + drop_name, + )?; + let (start, end) = one.overlap_range(many, one_index, many_index); + for point in start..end { + if !(many.present(many_index, point)? && one.present(one_index, point)?) { + continue; + } + for (field, value) in values.iter_mut().enumerate() { + let (lhs_value, rhs_value) = if one_is_left { + ( + one.value(one_index, field, point)?, + many.value(many_index, field, point)?, + ) + } else { + ( + many.value(many_index, field, point)?, + one.value(one_index, field, point)?, + ) + }; + *value = match (lhs_value, rhs_value) { + (Some(lhs_value), Some(rhs_value)) => { + let (result, holds) = matching.op.apply(lhs_value, rhs_value); + if !matching.op.is_comparison() { + Some(result) + } else if matching.return_bool { + Some(if holds { 1.0 } else { 0.0 }) + } else { + holds.then_some(result) + } + } + _ => None, + }; + } + results.insert(&metric, point, &values); + } + } + } + // One-to-one: a match group may be matched once per point; Prometheus counts a filtered + // comparison as a match too. + for (key, members) in &matched_groups { + let matched = |point: usize| -> Result { + let Some(candidates) = one_groups.get(key) else { + return Ok(false); + }; + for &candidate in candidates { + if one.covers(candidate, point) && one.present(candidate, point)? { + return Ok(true); + } + } + Ok(false) + }; + if let Some(point) = overlap(many, members, matched)? { + return Err(DataFusionError::Execution(format!( + "multiple matches for labels {} at {} ms: many-to-one matching must be explicit \ + (group_left/group_right)", + render(key), + grid.timestamps()[point] + ))); + } + } + results.finish(output, |key, point| { + DataFusionError::Execution(format!( + "multiple matches for labels {} at {} ms: grouping labels must ensure unique matches", + render(key), + grid.timestamps()[point] + )) + }) +} + +/// Applies a set operator at every grid point. Series keep their labels and values. +pub fn set_vectors( + op: SetOp, + labels: &MatchingLabels, + lhs_schema: &SeriesSchema, + lhs: &[Series], + rhs_schema: &SeriesSchema, + rhs: &[Series], +) -> Result> { + let grid = same_grid(lhs_schema, rhs_schema)?; + if op == SetOp::Or && lhs_schema.fields.len() != rhs_schema.fields.len() { + return Err(DataFusionError::NotImplemented( + "`or` between vectors with different field counts".into(), + )); + } + let output = Arc::new(SeriesSchema { + fields: lhs_schema.fields.clone(), + eval_grid: Some(grid.clone()), + }); + let lhs = Operand::try_new(lhs_schema, lhs)?; + let rhs = Operand::try_new(rhs_schema, rhs)?; + let (kept, filter, keep_matched) = match op { + SetOp::And => (&lhs, &rhs, true), + SetOp::Unless => (&lhs, &rhs, false), + SetOp::Or => (&rhs, &lhs, false), + }; + let filter_groups = group(filter, labels); + let mut results = Results::default(); + let fields = output.fields.len(); + let mut values = vec![None; fields]; + if op == SetOp::Or { + for index in 0..lhs.series.len() { + for point in lhs.range(index) { + if lhs.present(index, point)? { + lhs.row(index, point, &mut values)?; + results.insert(&lhs.series[index].metric, point, &values); + } + } + } + } + for (index, series) in kept.series.iter().enumerate() { + let candidates = filter_groups.get(&group_key(&series.metric, labels)); + for point in kept.range(index) { + if !kept.present(index, point)? { + continue; + } + let mut matched = false; + for &candidate in candidates.into_iter().flatten() { + if filter.covers(candidate, point) && filter.present(candidate, point)? { + matched = true; + break; + } + } + if matched == keep_matched { + kept.row(index, point, &mut values)?; + results.insert(&series.metric, point, &values); + } + } + } + results.finish(output, |key, point| { + DataFusionError::Execution(format!( + "vector cannot contain metrics with the same labelset {} at {} ms", + render(key), + grid.timestamps()[point] + )) + }) +} + +fn same_grid(lhs: &SeriesSchema, rhs: &SeriesSchema) -> Result> { + match (&lhs.eval_grid, &rhs.eval_grid) { + (Some(lhs), Some(rhs)) if lhs.spec() == rhs.spec() => Ok(lhs.clone()), + _ => Err(DataFusionError::Execution( + "vector operands must be evaluated on the same grid".into(), + )), + } +} + +/// Label pairs identifying a match group or a result; empty values are absent, as in +/// Prometheus. +type GroupKey = Vec<(Arc, Arc)>; + +fn group_key(metric: &MetricName, labels: &MatchingLabels) -> GroupKey { + let mut key: GroupKey = match labels { + MatchingLabels::On(names) => names + .iter() + .filter_map(|name| { + let value = metric.get(name)?; + (!value.is_empty()).then(|| (name.clone(), Arc::from(value))) + }) + .collect(), + MatchingLabels::Ignoring(names) => metric + .tags() + .iter() + .filter(|tag| !tag.value.is_empty() && !names.contains(&tag.name)) + .map(|tag| (tag.name.clone(), tag.value.clone())) + .collect(), + }; + key.sort(); + key.dedup(); + key +} + +fn key_of(metric: &MetricName) -> GroupKey { + let mut key = group_key(metric, &MatchingLabels::Ignoring(Vec::new())); + if let Some(name) = metric.name() { + key.insert(0, (Arc::from("__name__"), Arc::from(name))); + } + key +} + +fn render(key: &GroupKey) -> String { + let labels: Vec<_> = key + .iter() + .map(|(name, value)| format!("{name}=\"{value}\"")) + .collect(); + format!("{{{}}}", labels.join(", ")) +} + +/// Groups the series of `operand` by their match group. +fn group(operand: &Operand<'_>, labels: &MatchingLabels) -> HashMap> { + let mut groups: HashMap> = HashMap::new(); + for (index, series) in operand.series.iter().enumerate() { + groups + .entry(group_key(&series.metric, labels)) + .or_default() + .push(index); + } + groups +} + +/// The first point where two `members` are present and `relevant` holds. +fn overlap( + operand: &Operand<'_>, + members: &[usize], + mut relevant: impl FnMut(usize) -> Result, +) -> Result> { + if members.len() < 2 { + return Ok(None); + } + let start = members.iter().map(|&m| operand.start(m)).min().unwrap_or(0); + let end = members.iter().map(|&m| operand.end(m)).max().unwrap_or(0); + let mut seen = vec![false; end.saturating_sub(start)]; + for &member in members { + for point in operand.range(member) { + if !operand.present(member, point)? { + continue; + } + if std::mem::replace(&mut seen[point - start], true) && relevant(point)? { + return Ok(Some(point)); + } + } + } + Ok(None) +} + +fn result_metric( + many: &MetricName, + one: &MetricName, + matching: &BinaryMatching, + include: &[Arc], + drop_name: bool, +) -> Result { + let mut metric = many.clone(); + if drop_name { + metric.remove("__name__"); + } + if matching.cardinality == Cardinality::OneToOne { + match &matching.labels { + MatchingLabels::On(names) => { + metric.keep_only(&names.iter().map(|name| &**name).collect::>()) + } + MatchingLabels::Ignoring(names) => { + metric.remove_all(&names.iter().map(|name| &**name).collect::>()) + } + } + } + for name in include { + match one.get(name) { + Some(value) if !value.is_empty() => metric + .set(name, Arc::from(value)) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + _ => metric.remove(name), + } + } + Ok(metric) +} + +/// One side's series with checked grid ranges. +struct Operand<'a> { + series: &'a [Series], + starts: Vec, +} + +impl<'a> Operand<'a> { + fn try_new(schema: &SeriesSchema, series: &'a [Series]) -> Result { + let grid = schema.eval_grid.as_ref().map(|grid| grid.spec()); + let mut starts = Vec::with_capacity(series.len()); + for item in series { + let TimeAxis::Grid(range) = &item.time else { + return Err(DataFusionError::Execution( + "vector operand series must be on the evaluation grid".into(), + )); + }; + if Some(range.grid().spec()) != grid + || item.fields.len() != schema.fields.len() + || item.fields.iter().any(|field| field.len() != range.len()) + { + return Err(DataFusionError::Execution( + "vector operand series does not match its schema".into(), + )); + } + starts.push(range.start()); + } + Ok(Self { series, starts }) + } + + fn start(&self, index: usize) -> usize { + self.starts[index] + } + + fn end(&self, index: usize) -> usize { + self.starts[index] + self.series[index].time.len() + } + + fn range(&self, index: usize) -> std::ops::Range { + self.start(index)..self.end(index) + } + + fn covers(&self, index: usize, point: usize) -> bool { + self.range(index).contains(&point) + } + + /// The shared points of `self[own]` and `other[theirs]`. + fn overlap_range(&self, other: &Operand<'_>, own: usize, theirs: usize) -> (usize, usize) { + let start = self.start(own).max(other.start(theirs)); + let end = self.end(own).min(other.end(theirs)); + (start, end.max(start)) + } + + fn value(&self, index: usize, field: usize, point: usize) -> Result> { + let samples = &self.series[index].fields.as_slice()[field]; + sample(samples, point - self.starts[index]) + } + + fn present(&self, index: usize, point: usize) -> Result { + for field in 0..self.series[index].fields.len() { + if self.value(index, field, point)?.is_some() { + return Ok(true); + } + } + Ok(false) + } + + fn row(&self, index: usize, point: usize, values: &mut [Option]) -> Result<()> { + for (field, value) in values.iter_mut().enumerate() { + *value = self.value(index, field, point)?; + } + Ok(()) + } + + /// Whether any series is present at each grid point. + fn presence(&self, points: usize) -> Result> { + let mut presence = vec![false; points]; + for index in 0..self.series.len() { + for point in self.range(index) { + if !presence[point] && self.present(index, point)? { + presence[point] = true; + } + } + } + Ok(presence) + } +} + +fn sample(samples: &Samples, row: usize) -> Result> { + if let Some(DenseSamplesRef::Float(values)) = samples.dense() { + return Ok(float_at(values, row)); + } + match samples.sample(row) { + SampleRef::Missing => Ok(None), + SampleRef::Float(value) => Ok(Some(value)), + SampleRef::Histogram(_) => Err(DataFusionError::NotImplemented( + "vector matching of histogram samples".into(), + )), + } +} + +fn float_at(values: &Float64Array, row: usize) -> Option { + values.is_valid(row).then(|| values.value(row)) +} + +/// Result rows by result label set, assembled into trimmed grid series. +#[derive(Default)] +struct Results { + series: HashMap, +} + +/// Rows in arrival order: their points, and their field values laid out row after row. +struct ResultRows { + metric: MetricName, + points: Vec, + values: Vec>, +} + +impl Results { + fn insert(&mut self, metric: &MetricName, point: usize, values: &[Option]) { + let rows = match self.series.entry(key_of(metric)) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => entry.insert(ResultRows { + metric: metric.clone(), + points: Vec::new(), + values: Vec::new(), + }), + }; + rows.points.push(point); + rows.values.extend_from_slice(values); + } + + /// Builds the result series; two rows of one label set at the same point fail with + /// `duplicate`. + fn finish( + self, + schema: Arc, + duplicate: impl Fn(&GroupKey, usize) -> DataFusionError, + ) -> Result> { + let fields = schema.fields.len(); + let mut builder = GridSeriesBuilder::try_new(schema)?; + let mut output = Vec::with_capacity(self.series.len()); + let mut order = Vec::new(); + for (key, rows) in self.series { + order.clear(); + order.extend(0..rows.points.len()); + order.sort_unstable_by_key(|&row| rows.points[row]); + if let Some(pair) = order + .windows(2) + .find(|pair| rows.points[pair[0]] == rows.points[pair[1]]) + { + return Err(duplicate(&key, rows.points[pair[0]])); + } + for &row in &order { + builder.push( + rows.points[row], + &rows.values[row * fields..(row + 1) * fields], + )?; + } + if let Some(series) = builder.finish(rows.metric)? { + output.push(series); + } + } + Ok(output) + } +} diff --git a/src/promql/src/vector_match/tests.rs b/src/promql/src/vector_match/tests.rs new file mode 100644 index 00000000000..3a9fd2f3771 --- /dev/null +++ b/src/promql/src/vector_match/tests.rs @@ -0,0 +1,497 @@ +// 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. + +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::sync::Arc; + +use datafusion::error::DataFusionError; +use datatypes::series::samples::{FieldSpec, Series, SeriesSchema}; +use datatypes::series::time::{EvalGrid, GridSpec}; +use datatypes::series::{Label, MetricName}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +use super::{ + BinaryMatching, BinaryOp, Cardinality, MatchingLabels, SetOp, binary_vectors, set_vectors, +}; +use crate::grid_series::GridSeriesBuilder; +use crate::grid_series::reference::{Dense, assert_series_eq, random_dense}; + +/// Labels with `__name__` as an ordinary entry, the way Prometheus holds them. +type Labels = BTreeMap; + +const OPS: [BinaryOp; 13] = [ + BinaryOp::Add, + BinaryOp::Sub, + BinaryOp::Mul, + BinaryOp::Div, + BinaryOp::Mod, + BinaryOp::Pow, + BinaryOp::Atan2, + BinaryOp::Eq, + BinaryOp::Ne, + BinaryOp::Gt, + BinaryOp::Lt, + BinaryOp::Ge, + BinaryOp::Le, +]; + +const LABEL_NAMES: [&str; 4] = ["__name__", "host", "job", "zone"]; + +fn schema(fields: usize, grid: &Arc) -> SeriesSchema { + SeriesSchema { + fields: (0..fields) + .map(|index| FieldSpec { + name: Arc::from(format!("f{index}")), + }) + .collect::>() + .into(), + eval_grid: Some(grid.clone()), + } +} + +fn metric(labels: &Labels) -> MetricName { + let mut labels = labels.clone(); + let name = labels.remove("__name__").map(Arc::from); + MetricName::try_new( + name, + labels + .into_iter() + .map(|(name, value)| Label { + name: Arc::from(name), + value: Arc::from(value), + }) + .collect(), + ) + .unwrap() +} + +fn labels_of(metric: &MetricName) -> Labels { + let mut labels: Labels = metric + .tags() + .iter() + .map(|tag| (tag.name.to_string(), tag.value.to_string())) + .collect(); + if let Some(name) = metric.name() { + labels.insert("__name__".to_string(), name.to_string()); + } + labels +} + +/// One operand: distinct label sets with their dense values. +struct Side { + schema: SeriesSchema, + series: Vec, + dense: Vec<(Labels, Dense)>, +} + +fn random_side(rng: &mut StdRng, grid: &Arc, fields: usize) -> Side { + let schema = schema(fields, grid); + let mut builder = GridSeriesBuilder::try_new(Arc::new(schema.clone())).unwrap(); + let mut seen = HashSet::new(); + let mut series = Vec::new(); + let mut dense = Vec::new(); + for _ in 0..rng.random_range(0..6) { + let mut labels = Labels::new(); + for (name, values) in [ + ("__name__", ["m1", "m2", ""]), + ("host", ["a", "b", ""]), + ("job", ["x", "y", "y"]), + ("zone", ["z1", "", ""]), + ] { + let value = values[rng.random_range(0..values.len())]; + if !value.is_empty() { + labels.insert(name.to_string(), value.to_string()); + } + } + if !seen.insert(labels.clone()) { + continue; + } + let values = random_dense(rng, grid.timestamps().len(), fields); + for (index, row) in values.iter().enumerate() { + builder.push(index, row).unwrap(); + } + if let Some(built) = builder.finish(metric(&labels)).unwrap() { + series.push(built); + dense.push((labels, values)); + } + } + Side { + schema, + series, + dense, + } +} + +fn random_names(rng: &mut StdRng, candidates: &[&str]) -> Vec> { + candidates + .iter() + .filter(|_| rng.random_bool(0.4)) + .map(|name| Arc::from(*name)) + .collect() +} + +fn random_labels(rng: &mut StdRng) -> MatchingLabels { + if rng.random_bool(0.5) { + MatchingLabels::On(random_names(rng, &LABEL_NAMES)) + } else { + MatchingLabels::Ignoring(random_names(rng, &LABEL_NAMES)) + } +} + +fn listed(names: &[Arc], name: &str) -> bool { + names.iter().any(|listed| &**listed == name) +} + +/// Prometheus `signatureFunc`: `on` keeps the listed labels, `ignoring` every label except the +/// listed ones and `__name__`. +fn signature(labels: &Labels, matching: &MatchingLabels) -> Labels { + labels + .iter() + .filter(|(name, _)| match matching { + MatchingLabels::On(names) => listed(names, name), + MatchingLabels::Ignoring(names) => name.as_str() != "__name__" && !listed(names, name), + }) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +/// Prometheus `resultMetric`. +fn result_labels(many: &Labels, one: &Labels, matching: &BinaryMatching) -> Labels { + let mut labels = many.clone(); + if matching.return_bool || !matching.op.is_comparison() { + labels.remove("__name__"); + } + let include = match &matching.cardinality { + Cardinality::OneToOne => { + match &matching.labels { + MatchingLabels::On(names) => labels.retain(|name, _| listed(names, name)), + MatchingLabels::Ignoring(names) => labels.retain(|name, _| !listed(names, name)), + } + &[][..] + } + Cardinality::ManyToOne(include) | Cardinality::OneToMany(include) => &include[..], + }; + for name in include { + match one.get(&**name) { + Some(value) => labels.insert(name.to_string(), value.clone()), + None => labels.remove(&**name), + }; + } + labels +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum Violation { + DuplicateOneSide, + ImplicitManyToOne, + Grouping, + DuplicateLabelSet, +} + +fn violation_of(error: &DataFusionError) -> Violation { + let message = error.to_string(); + if message.contains("duplicate series for the match group") { + Violation::DuplicateOneSide + } else if message.contains("many-to-one matching must be explicit") { + Violation::ImplicitManyToOne + } else if message.contains("grouping labels must ensure unique matches") { + Violation::Grouping + } else if message.contains("same labelset") { + Violation::DuplicateLabelSet + } else { + panic!("unexpected error {message}") + } +} + +type Expected = Result, BTreeSet>; + +fn present(row: &[Option]) -> bool { + row.iter().any(Option::is_some) +} + +fn present_at(side: &Side, point: usize) -> Vec<&(Labels, Dense)> { + side.dense + .iter() + .filter(|(_, dense)| present(&dense[point])) + .collect() +} + +/// Prometheus `VectorBinop` evaluated on its own at every point, collecting every violation +/// instead of stopping at the first. +fn reference_binary( + matching: &BinaryMatching, + lhs: &Side, + rhs: &Side, + points: usize, + fields: usize, +) -> Expected { + let mut output: BTreeMap = BTreeMap::new(); + let mut violations = BTreeSet::new(); + let swapped = matches!(matching.cardinality, Cardinality::OneToMany(_)); + for point in 0..points { + let left = present_at(lhs, point); + let right = present_at(rhs, point); + if left.is_empty() || right.is_empty() { + continue; + } + let (many, one) = if swapped { + (&right, &left) + } else { + (&left, &right) + }; + let mut one_by_signature = HashMap::new(); + let mut duplicate = false; + for (labels, dense) in one { + duplicate |= one_by_signature + .insert(signature(labels, &matching.labels), (labels, dense)) + .is_some(); + } + if duplicate { + violations.insert(Violation::DuplicateOneSide); + continue; + } + let mut matched = HashSet::new(); + let mut inserted: HashMap> = HashMap::new(); + for (many_labels, many_dense) in many { + let signature = signature(many_labels, &matching.labels); + let Some((one_labels, one_dense)) = one_by_signature.get(&signature) else { + continue; + }; + let result = result_labels(many_labels, one_labels, matching); + if matching.cardinality == Cardinality::OneToOne { + if !matched.insert(signature.clone()) { + violations.insert(Violation::ImplicitManyToOne); + continue; + } + } else if !inserted + .entry(signature.clone()) + .or_default() + .insert(result.clone()) + { + violations.insert(Violation::Grouping); + continue; + } + let (lhs_row, rhs_row) = if swapped { + (&one_dense[point], &many_dense[point]) + } else { + (&many_dense[point], &one_dense[point]) + }; + let row: Vec<_> = (0..fields) + .map(|field| { + let (value, holds) = matching.op.apply(lhs_row[field]?, rhs_row[field]?); + if !matching.op.is_comparison() { + Some(value) + } else if matching.return_bool { + Some(if holds { 1.0 } else { 0.0 }) + } else { + holds.then_some(value) + } + }) + .collect(); + let dense = output + .entry(result) + .or_insert_with(|| vec![vec![None; fields]; points]); + assert!(!present(&dense[point]), "result label sets collide"); + dense[point] = row; + } + } + if violations.is_empty() { + Ok(output) + } else { + Err(violations) + } +} + +/// Prometheus `VectorAnd`, `VectorOr` and `VectorUnless` at every point. +fn reference_set( + op: SetOp, + labels: &MatchingLabels, + lhs: &Side, + rhs: &Side, + points: usize, +) -> Expected { + let fields = lhs.schema.fields.len(); + let mut output: BTreeMap = BTreeMap::new(); + let mut violations = BTreeSet::new(); + for point in 0..points { + let left = present_at(lhs, point); + let right = present_at(rhs, point); + let signatures = |side: &[&(Labels, Dense)]| -> HashSet { + side.iter().map(|(l, _)| signature(l, labels)).collect() + }; + let (left_signatures, right_signatures) = (signatures(&left), signatures(&right)); + let kept: Vec<_> = match op { + SetOp::And => left + .into_iter() + .filter(|(l, _)| right_signatures.contains(&signature(l, labels))) + .collect(), + SetOp::Unless => left + .into_iter() + .filter(|(l, _)| !right_signatures.contains(&signature(l, labels))) + .collect(), + SetOp::Or => left + .into_iter() + .chain( + right + .into_iter() + .filter(|(l, _)| !left_signatures.contains(&signature(l, labels))), + ) + .collect(), + }; + for (labels, dense) in kept { + let target = output + .entry(labels.clone()) + .or_insert_with(|| vec![vec![None; fields]; points]); + if present(&target[point]) { + violations.insert(Violation::DuplicateLabelSet); + } + target[point] = dense[point].clone(); + } + } + if violations.is_empty() { + Ok(output) + } else { + Err(violations) + } +} + +fn assert_output(actual: Result, DataFusionError>, expected: Expected, seed: u64) { + match (actual, expected) { + (Ok(actual), Ok(expected)) => { + let actual: BTreeMap = actual + .into_iter() + .map(|series| (labels_of(&series.metric), series)) + .collect(); + let expected: BTreeMap<_, _> = expected + .into_iter() + .filter(|(_, dense)| dense.iter().any(|row| present(row))) + .collect(); + assert_eq!( + actual.keys().collect::>(), + expected.keys().collect::>(), + "seed {seed}" + ); + for (labels, dense) in &expected { + assert_series_eq(actual.get(labels), dense, &metric(labels)); + } + } + (Err(error), Err(violations)) => assert!( + violations.contains(&violation_of(&error)), + "seed {seed}: {error} not in {violations:?}" + ), + (actual, expected) => panic!("seed {seed}: got {actual:?}, expected {expected:?}"), + } +} + +fn random_grid(rng: &mut StdRng) -> Arc { + let spec = GridSpec::try_new(rng.random_range(-1_000..1_000), 10, rng.random_range(1..40)); + Arc::new(EvalGrid::try_new(spec.unwrap()).unwrap()) +} + +#[test] +fn binary_vectors_match_prometheus_reference() { + let (mut matched, mut failed) = (0, 0); + for seed in 0..3_000 { + let mut rng = StdRng::seed_from_u64(seed); + let grid = random_grid(&mut rng); + let lhs_fields = rng.random_range(1..=2); + let lhs = random_side(&mut rng, &grid, lhs_fields); + let rhs_fields = rng.random_range(1..=2); + let rhs = random_side(&mut rng, &grid, rhs_fields); + let op = OPS[rng.random_range(0..OPS.len())]; + let return_bool = op.is_comparison() && rng.random_bool(0.5); + let labels = random_labels(&mut rng); + let include = random_names(&mut rng, &["host", "job", "zone"]); + let cardinality = match rng.random_range(0..3) { + 0 => Cardinality::OneToOne, + 1 => Cardinality::ManyToOne(include), + _ => Cardinality::OneToMany(include), + }; + let matching = BinaryMatching { + op, + return_bool, + labels, + cardinality, + }; + let fields = lhs_fields.min(rhs_fields); + let expected = reference_binary(&matching, &lhs, &rhs, grid.timestamps().len(), fields); + if expected.is_ok() { + matched += 1; + } else { + failed += 1; + } + let actual = binary_vectors( + &matching, + &lhs.schema, + &lhs.series, + &rhs.schema, + &rhs.series, + ); + assert_output(actual, expected, seed); + } + assert!( + matched > 1_000 && failed > 100, + "{matched} matched, {failed} failed" + ); +} + +#[test] +fn set_vectors_match_prometheus_reference() { + for seed in 0..3_000 { + let mut rng = StdRng::seed_from_u64(seed); + let grid = random_grid(&mut rng); + let fields = rng.random_range(1..=2); + let lhs = random_side(&mut rng, &grid, fields); + let rhs = random_side(&mut rng, &grid, fields); + let op = [SetOp::And, SetOp::Or, SetOp::Unless][rng.random_range(0..3)]; + let labels = random_labels(&mut rng); + let expected = reference_set(op, &labels, &lhs, &rhs, grid.timestamps().len()); + let actual = set_vectors( + op, + &labels, + &lhs.schema, + &lhs.series, + &rhs.schema, + &rhs.series, + ); + assert_output(actual, expected, seed); + } +} + +#[test] +fn rejects_operands_on_other_grids_and_unpaired_or_fields() { + let mut rng = StdRng::seed_from_u64(0); + let grid = random_grid(&mut rng); + let other = Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 7, 3).unwrap()).unwrap()); + let lhs = random_side(&mut rng, &grid, 1); + let matching = BinaryMatching { + op: BinaryOp::Add, + return_bool: false, + labels: MatchingLabels::Ignoring(Vec::new()), + cardinality: Cardinality::OneToOne, + }; + assert!(binary_vectors(&matching, &lhs.schema, &lhs.series, &schema(1, &other), &[]).is_err()); + assert!(matches!( + set_vectors( + SetOp::Or, + &matching.labels, + &lhs.schema, + &lhs.series, + &schema(2, &grid), + &[] + ), + Err(DataFusionError::NotImplemented(_)) + )); +}