promql/functions/
vector_matching.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Cardinality checks for PromQL vector matching.
16
17use datafusion::arrow::array::{Array, Int64Array};
18use datafusion::arrow::util::display::array_value_to_string;
19use datafusion::common::{DataFusionError, Result as DfResult};
20use datafusion::logical_expr::{ScalarUDF, Volatility};
21use datafusion::physical_plan::ColumnarValue;
22use datafusion_common::ScalarValue;
23use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
24use datatypes::arrow::datatypes::DataType;
25
26use crate::functions::extract_array;
27
28/// The rule a repeated match group violates.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
30pub enum MatchGroupViolation {
31    /// Several series share a match group on the side that must hold one series per group.
32    DuplicateOnOneSide { one_side_is_left: bool },
33    /// One-to-one matching found several matches for the same group.
34    ImplicitManyToOne,
35    /// A group modifier left several matches with the same result label set.
36    AmbiguousGroupLabels,
37    /// A function rewrote labels so that several series share the same label set.
38    DuplicateLabelSet,
39}
40
41impl MatchGroupViolation {
42    fn message(&self, group: &str) -> String {
43        match self {
44            Self::DuplicateOnOneSide { one_side_is_left } => {
45                let side = if *one_side_is_left { "left" } else { "right" };
46                format!(
47                    "found duplicate series for the match group {group} on the {side} hand-side \
48                     of the operation; many-to-many matching not allowed: matching labels must \
49                     be unique on one side"
50                )
51            }
52            Self::ImplicitManyToOne => format!(
53                "multiple matches for labels {group}: many-to-one matching must be explicit \
54                 (group_left/group_right)"
55            ),
56            Self::AmbiguousGroupLabels => format!(
57                "multiple matches for labels {group}: grouping labels must ensure unique matches"
58            ),
59            Self::DuplicateLabelSet => {
60                "vector cannot contain metrics with the same labelset".to_string()
61            }
62        }
63    }
64}
65
66/// Rejects a vector matching whose match groups are not unique.
67///
68/// Takes the per-group row count as first argument and the label columns that form the group
69/// as the remaining ones. Returns `true` when every group holds a single row, and fails the
70/// query otherwise; PromQL has no way to express the duplicate series it would produce.
71pub struct UniqueMatchGroup;
72
73impl UniqueMatchGroup {
74    pub const fn name() -> &'static str {
75        "prom_assert_unique_match_group"
76    }
77
78    pub fn scalar_udf(labels: Vec<String>, violation: MatchGroupViolation) -> ScalarUDF {
79        ScalarUDF::new_from_impl(AssertUniqueMatchGroup {
80            signature: Signature::variadic_any(Volatility::Volatile),
81            labels,
82            violation,
83        })
84    }
85}
86
87#[derive(Debug, Clone, PartialEq, Eq, Hash)]
88struct AssertUniqueMatchGroup {
89    signature: Signature,
90    /// Label names of the group, in the same order as the label arguments.
91    labels: Vec<String>,
92    violation: MatchGroupViolation,
93}
94
95impl AssertUniqueMatchGroup {
96    fn render_group(&self, args: &[ColumnarValue], row: usize) -> DfResult<String> {
97        let mut rendered = Vec::with_capacity(self.labels.len());
98        for (label, arg) in self.labels.iter().zip(args) {
99            let array = extract_array(arg)?;
100            // A scalar argument was expanded to a single row above.
101            let row = if array.len() == 1 { 0 } else { row };
102            if row >= array.len() || array.is_null(row) {
103                continue;
104            }
105            let value = array_value_to_string(&array, row)
106                .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
107            if value.is_empty() {
108                continue;
109            }
110            rendered.push(format!("{label}=\"{value}\""));
111        }
112        Ok(format!("{{{}}}", rendered.join(", ")))
113    }
114}
115
116impl ScalarUDFImpl for AssertUniqueMatchGroup {
117    fn name(&self) -> &str {
118        UniqueMatchGroup::name()
119    }
120
121    fn signature(&self) -> &Signature {
122        &self.signature
123    }
124
125    fn return_type(&self, _arg_types: &[DataType]) -> DfResult<DataType> {
126        Ok(DataType::Boolean)
127    }
128
129    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult<ColumnarValue> {
130        let Some((counts, labels)) = args.args.split_first() else {
131            return Err(DataFusionError::Execution(format!(
132                "{} expects the match group row count as first argument",
133                UniqueMatchGroup::name()
134            )));
135        };
136        let counts = extract_array(counts)?;
137        let counts = counts
138            .as_any()
139            .downcast_ref::<Int64Array>()
140            .ok_or_else(|| {
141                DataFusionError::Execution(format!(
142                    "{} expects an Int64 match group row count, found {}",
143                    UniqueMatchGroup::name(),
144                    counts.data_type()
145                ))
146            })?;
147
148        if let Some(row) =
149            (0..counts.len()).find(|row| !counts.is_null(*row) && counts.value(*row) > 1)
150        {
151            let group = self.render_group(labels, row)?;
152            return Err(DataFusionError::Execution(self.violation.message(&group)));
153        }
154
155        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))))
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use std::sync::Arc;
162
163    use datafusion::arrow::array::StringArray;
164    use datafusion::arrow::datatypes::Field;
165
166    use super::*;
167
168    fn invoke(counts: Vec<i64>, hosts: Vec<Option<&str>>) -> DfResult<ColumnarValue> {
169        let udf = UniqueMatchGroup::scalar_udf(
170            vec!["host".to_string()],
171            MatchGroupViolation::ImplicitManyToOne,
172        );
173        let number_rows = counts.len();
174        udf.invoke_with_args(ScalarFunctionArgs {
175            args: vec![
176                ColumnarValue::Array(Arc::new(Int64Array::from(counts))),
177                ColumnarValue::Array(Arc::new(StringArray::from(hosts))),
178            ],
179            arg_fields: vec![
180                Arc::new(Field::new("count", DataType::Int64, true)),
181                Arc::new(Field::new("host", DataType::Utf8, true)),
182            ],
183            number_rows,
184            return_field: Arc::new(Field::new("assert", DataType::Boolean, false)),
185            config_options: Arc::new(Default::default()),
186        })
187    }
188
189    #[test]
190    fn unique_groups_pass() {
191        let result = invoke(vec![1, 1], vec![Some("a"), Some("b")]).unwrap();
192        assert!(matches!(
193            result,
194            ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)))
195        ));
196    }
197
198    #[test]
199    fn duplicate_group_reports_its_labels() {
200        let err = invoke(vec![1, 2], vec![Some("a"), Some("b")]).unwrap_err();
201        assert!(
202            err.to_string()
203                .contains("multiple matches for labels {host=\"b\"}"),
204            "{err}"
205        );
206    }
207
208    #[test]
209    fn null_label_is_omitted_from_the_group() {
210        let err = invoke(vec![2], vec![None]).unwrap_err();
211        assert!(err.to_string().contains("labels {}"), "{err}");
212    }
213}
promql/functions/vector_matching.rs

promql/functions/
vector_matching.rs
