mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-11-14 06:46:56 +00:00
feat(promql): match evaluated vectors for binary and set operators
Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
Dennis Zhuang
@@ -19,3 +19,4 @@ pub mod grid_series;
|
||||
mod metrics;
|
||||
pub mod range_array;
|
||||
pub mod range_eval;
|
||||
pub mod vector_match;
|
||||
|
||||
@@ -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<Arc<str>>),
|
||||
/// `ignoring(...)`, and the default with no labels: every label except these and
|
||||
/// `__name__`.
|
||||
Ignoring(Vec<Arc<str>>),
|
||||
}
|
||||
|
||||
/// 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<Arc<str>>),
|
||||
/// `group_right(...)`: the left side is the "one" side; the listed labels are copied from it.
|
||||
OneToMany(Vec<Arc<str>>),
|
||||
}
|
||||
|
||||
/// 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<Arc<SeriesSchema>> {
|
||||
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<Vec<Series>> {
|
||||
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<bool> = 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<GroupKey, Vec<usize>> = 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<bool> {
|
||||
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<Vec<Series>> {
|
||||
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<Arc<EvalGrid>> {
|
||||
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<str>, Arc<str>)>;
|
||||
|
||||
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<GroupKey, Vec<usize>> {
|
||||
let mut groups: HashMap<GroupKey, Vec<usize>> = 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<bool>,
|
||||
) -> Result<Option<usize>> {
|
||||
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<str>],
|
||||
drop_name: bool,
|
||||
) -> Result<MetricName> {
|
||||
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::<Vec<_>>())
|
||||
}
|
||||
MatchingLabels::Ignoring(names) => {
|
||||
metric.remove_all(&names.iter().map(|name| &**name).collect::<Vec<_>>())
|
||||
}
|
||||
}
|
||||
}
|
||||
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<usize>,
|
||||
}
|
||||
|
||||
impl<'a> Operand<'a> {
|
||||
fn try_new(schema: &SeriesSchema, series: &'a [Series]) -> Result<Self> {
|
||||
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<usize> {
|
||||
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<Option<f64>> {
|
||||
let samples = &self.series[index].fields.as_slice()[field];
|
||||
sample(samples, point - self.starts[index])
|
||||
}
|
||||
|
||||
fn present(&self, index: usize, point: usize) -> Result<bool> {
|
||||
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<f64>]) -> 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<Vec<bool>> {
|
||||
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<Option<f64>> {
|
||||
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<f64> {
|
||||
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<GroupKey, ResultRows>,
|
||||
}
|
||||
|
||||
/// Rows in arrival order: their points, and their field values laid out row after row.
|
||||
struct ResultRows {
|
||||
metric: MetricName,
|
||||
points: Vec<usize>,
|
||||
values: Vec<Option<f64>>,
|
||||
}
|
||||
|
||||
impl Results {
|
||||
fn insert(&mut self, metric: &MetricName, point: usize, values: &[Option<f64>]) {
|
||||
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<SeriesSchema>,
|
||||
duplicate: impl Fn(&GroupKey, usize) -> DataFusionError,
|
||||
) -> Result<Vec<Series>> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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<String, String>;
|
||||
|
||||
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<EvalGrid>) -> SeriesSchema {
|
||||
SeriesSchema {
|
||||
fields: (0..fields)
|
||||
.map(|index| FieldSpec {
|
||||
name: Arc::from(format!("f{index}")),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.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<Series>,
|
||||
dense: Vec<(Labels, Dense)>,
|
||||
}
|
||||
|
||||
fn random_side(rng: &mut StdRng, grid: &Arc<EvalGrid>, 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<Arc<str>> {
|
||||
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<str>], 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<BTreeMap<Labels, Dense>, BTreeSet<Violation>>;
|
||||
|
||||
fn present(row: &[Option<f64>]) -> 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<Labels, Dense> = 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<Labels, HashSet<Labels>> = 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<Labels, Dense> = 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<Labels> {
|
||||
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<Vec<Series>, DataFusionError>, expected: Expected, seed: u64) {
|
||||
match (actual, expected) {
|
||||
(Ok(actual), Ok(expected)) => {
|
||||
let actual: BTreeMap<Labels, Series> = 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::<Vec<_>>(),
|
||||
expected.keys().collect::<Vec<_>>(),
|
||||
"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<EvalGrid> {
|
||||
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(_))
|
||||
));
|
||||
}
|
||||
Reference in New Issue
Block a user