bk_promql_parser/parser/
function.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
15use std::collections::{HashMap, HashSet};
16use std::fmt;
17
18use lazy_static::lazy_static;
19
20use crate::parser::value::ValueType;
21use crate::parser::{Expr, Prettier};
22use crate::util::join_vector;
23
24/// called by func in Call
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct FunctionArgs {
27    pub args: Vec<Box<Expr>>,
28}
29
30impl FunctionArgs {
31    pub fn empty_args() -> Self {
32        Self { args: vec![] }
33    }
34
35    pub fn new_args(expr: Expr) -> Self {
36        Self {
37            args: vec![Box::new(expr)],
38        }
39    }
40
41    pub fn append_args(mut self: FunctionArgs, expr: Expr) -> Self {
42        self.args.push(Box::new(expr));
43        self
44    }
45
46    pub fn is_empty(&self) -> bool {
47        self.args.is_empty()
48    }
49
50    pub fn len(&self) -> usize {
51        self.args.len()
52    }
53
54    pub fn first(&self) -> Option<Box<Expr>> {
55        self.args.first().cloned()
56    }
57
58    pub fn last(&self) -> Option<Box<Expr>> {
59        self.args.last().cloned()
60    }
61}
62
63impl fmt::Display for FunctionArgs {
64    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
65        write!(f, "{}", join_vector(&self.args, ", ", false))
66    }
67}
68
69impl Prettier for FunctionArgs {
70    fn pretty(&self, level: usize, max: usize) -> String {
71        let mut v = vec![];
72        for ex in &self.args {
73            v.push(ex.pretty(level, max));
74        }
75        v.join(",\n")
76    }
77}
78
79/// Functions is a list of all functions supported by PromQL, including their types.
80#[derive(Debug, Clone, PartialEq, Eq)]
81pub struct Function {
82    pub name: &'static str,
83    pub arg_types: Vec<ValueType>,
84    pub variadic: bool,
85    pub return_type: ValueType,
86}
87
88impl Function {
89    pub fn new(
90        name: &'static str,
91        arg_types: Vec<ValueType>,
92        variadic: bool,
93        return_type: ValueType,
94    ) -> Self {
95        Self {
96            name,
97            arg_types,
98            variadic,
99            return_type,
100        }
101    }
102}
103
104macro_rules! map {
105    // if variadic args, then the last is the variadic one
106    ($(($name:literal, $arg:expr, $ret:expr)),*) => (
107        {
108            let mut m: HashMap<&'static str, Function> = HashMap::new();
109            $(
110                let variadic = FUNCTIONS_WITH_VARIADIC_ARGS.contains($name);
111                let func = Function::new($name, $arg, variadic, $ret);
112                m.insert($name, func);
113            )*
114            m
115        }
116    );
117}
118
119lazy_static! {
120    static ref FUNCTIONS_WITH_VARIADIC_ARGS: HashSet<&'static str> = HashSet::from([
121        "days_in_month",
122        "day_of_year",
123        "day_of_month",
124        "day_of_week",
125        "year",
126        "month",
127        "hour",
128        "minute",
129        "label_join",
130        "round",
131    ]);
132    static ref FUNCTIONS: HashMap<&'static str, Function> = map!(
133        ("abs", vec![ValueType::Vector], ValueType::Vector),
134        ("absent", vec![ValueType::Vector], ValueType::Vector),
135        (
136            "absent_over_time",
137            vec![ValueType::Matrix],
138            ValueType::Vector
139        ),
140        ("acos", vec![ValueType::Vector], ValueType::Vector),
141        ("acosh", vec![ValueType::Vector], ValueType::Vector),
142        ("asin", vec![ValueType::Vector], ValueType::Vector),
143        ("asinh", vec![ValueType::Vector], ValueType::Vector),
144        ("atan", vec![ValueType::Vector], ValueType::Vector),
145        ("atanh", vec![ValueType::Vector], ValueType::Vector),
146        ("avg_over_time", vec![ValueType::Matrix], ValueType::Vector),
147        ("ceil", vec![ValueType::Vector], ValueType::Vector),
148        ("changes", vec![ValueType::Matrix], ValueType::Vector),
149        (
150            "clamp",
151            vec![ValueType::Vector, ValueType::Scalar, ValueType::Scalar],
152            ValueType::Vector
153        ),
154        (
155            "clamp_max",
156            vec![ValueType::Vector, ValueType::Scalar],
157            ValueType::Vector
158        ),
159        (
160            "clamp_min",
161            vec![ValueType::Vector, ValueType::Scalar],
162            ValueType::Vector
163        ),
164        ("cos", vec![ValueType::Vector], ValueType::Vector),
165        ("cosh", vec![ValueType::Vector], ValueType::Vector),
166        (
167            "count_over_time",
168            vec![ValueType::Matrix],
169            ValueType::Vector
170        ),
171        ("days_in_month", vec![ValueType::Vector], ValueType::Vector),
172        ("day_of_month", vec![ValueType::Vector], ValueType::Vector),
173        ("day_of_week", vec![ValueType::Vector], ValueType::Vector),
174        ("day_of_year", vec![ValueType::Vector], ValueType::Vector),
175        ("deg", vec![ValueType::Vector], ValueType::Vector),
176        ("delta", vec![ValueType::Matrix], ValueType::Vector),
177        ("deriv", vec![ValueType::Matrix], ValueType::Vector),
178        ("exp", vec![ValueType::Vector], ValueType::Vector),
179        ("floor", vec![ValueType::Vector], ValueType::Vector),
180        (
181            "histogram_count",
182            vec![ValueType::Vector],
183            ValueType::Vector
184        ),
185        ("histogram_sum", vec![ValueType::Vector], ValueType::Vector),
186        (
187            "histogram_fraction",
188            vec![ValueType::Scalar, ValueType::Scalar, ValueType::Vector],
189            ValueType::Vector
190        ),
191        (
192            "histogram_quantile",
193            vec![ValueType::Scalar, ValueType::Vector],
194            ValueType::Vector
195        ),
196        (
197            "holt_winters",
198            vec![ValueType::Matrix, ValueType::Scalar, ValueType::Scalar],
199            ValueType::Vector
200        ),
201        ("hour", vec![ValueType::Vector], ValueType::Vector),
202        ("idelta", vec![ValueType::Matrix], ValueType::Vector),
203        ("increase", vec![ValueType::Matrix], ValueType::Vector),
204        ("irate", vec![ValueType::Matrix], ValueType::Vector),
205        (
206            "label_replace",
207            vec![
208                ValueType::Vector,
209                ValueType::String,
210                ValueType::String,
211                ValueType::String,
212                ValueType::String,
213            ],
214            ValueType::Vector
215        ),
216        (
217            "label_join",
218            vec![
219                ValueType::Vector,
220                ValueType::String,
221                ValueType::String,
222                ValueType::String,
223            ],
224            ValueType::Vector
225        ),
226        ("last_over_time", vec![ValueType::Matrix], ValueType::Vector),
227        ("ln", vec![ValueType::Vector], ValueType::Vector),
228        ("log10", vec![ValueType::Vector], ValueType::Vector),
229        ("log2", vec![ValueType::Vector], ValueType::Vector),
230        ("max_over_time", vec![ValueType::Matrix], ValueType::Vector),
231        ("min_over_time", vec![ValueType::Matrix], ValueType::Vector),
232        ("minute", vec![ValueType::Vector], ValueType::Vector),
233        ("month", vec![ValueType::Vector], ValueType::Vector),
234        ("pi", vec![], ValueType::Scalar),
235        (
236            "predict_linear",
237            vec![ValueType::Matrix, ValueType::Scalar],
238            ValueType::Vector
239        ),
240        (
241            "present_over_time",
242            vec![ValueType::Matrix],
243            ValueType::Vector
244        ),
245        (
246            "quantile_over_time",
247            vec![ValueType::Scalar, ValueType::Matrix],
248            ValueType::Vector
249        ),
250        ("rad", vec![ValueType::Vector], ValueType::Vector),
251        ("rate", vec![ValueType::Matrix], ValueType::Vector),
252        ("resets", vec![ValueType::Matrix], ValueType::Vector),
253        (
254            "round",
255            vec![ValueType::Vector, ValueType::Scalar],
256            ValueType::Vector
257        ),
258        ("scalar", vec![ValueType::Vector], ValueType::Scalar),
259        ("sgn", vec![ValueType::Vector], ValueType::Vector),
260        ("sin", vec![ValueType::Vector], ValueType::Vector),
261        ("sinh", vec![ValueType::Vector], ValueType::Vector),
262        ("sort", vec![ValueType::Vector], ValueType::Vector),
263        ("sort_desc", vec![ValueType::Vector], ValueType::Vector),
264        ("sqrt", vec![ValueType::Vector], ValueType::Vector),
265        (
266            "stddev_over_time",
267            vec![ValueType::Matrix],
268            ValueType::Vector
269        ),
270        (
271            "stdvar_over_time",
272            vec![ValueType::Matrix],
273            ValueType::Vector
274        ),
275        ("sum_over_time", vec![ValueType::Matrix], ValueType::Vector),
276        ("tan", vec![ValueType::Vector], ValueType::Vector),
277        ("tanh", vec![ValueType::Vector], ValueType::Vector),
278        ("time", vec![], ValueType::Scalar),
279        ("timestamp", vec![ValueType::Vector], ValueType::Vector),
280        ("vector", vec![ValueType::Scalar], ValueType::Vector),
281        ("year", vec![ValueType::Vector], ValueType::Vector)
282    );
283}
284
285/// get_function returns a predefined Function object for the given name.
286pub(crate) fn get_function(name: &str) -> Option<Function> {
287    FUNCTIONS.get(name).cloned()
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293    use crate::parser::*;
294
295    #[test]
296    fn test_function_equality() {
297        let func = "month";
298        assert!(get_function(func).is_some());
299        assert_eq!(get_function(func), get_function(func));
300    }
301
302    #[test]
303    fn test_function_args_equality() {
304        assert_eq!(FunctionArgs::empty_args(), FunctionArgs::empty_args());
305
306        let arg1 = Expr::NumberLiteral(NumberLiteral::new(1.0));
307        let arg2 = Expr::StringLiteral(StringLiteral {
308            val: "prometheus".into(),
309        });
310        let args1 = FunctionArgs::new_args(arg1).append_args(arg2);
311
312        let arg1 = Expr::NumberLiteral(NumberLiteral::new(0.5 + 0.5));
313        let arg2 = Expr::StringLiteral(StringLiteral {
314            val: String::from("prometheus"),
315        });
316        let args2 = FunctionArgs::new_args(arg1).append_args(arg2);
317
318        assert_eq!(args1, args2);
319    }
320
321    #[test]
322    fn test_args_display() {
323        let cases = vec![
324            (
325                FunctionArgs::new_args(Expr::from(VectorSelector::from("up"))),
326                "up",
327            ),
328            (
329                FunctionArgs::empty_args()
330                    .append_args(Expr::from("src1"))
331                    .append_args(Expr::from("src2"))
332                    .append_args(Expr::from("src3")),
333                r#""src1", "src2", "src3""#,
334            ),
335        ];
336
337        for (args, expect) in cases {
338            assert_eq!(expect, args.to_string())
339        }
340    }
341}
bk_promql_parser/parser/function.rs

bk_promql_parser/parser/
function.rs
