Skip to main content

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;
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)]
26#[cfg_attr(feature = "ser", derive(serde::Serialize))]
27pub struct FunctionArgs {
28    pub args: Vec<Box<Expr>>,
29}
30
31impl FunctionArgs {
32    pub fn empty_args() -> Self {
33        Self { args: vec![] }
34    }
35
36    pub fn new_args(expr: Expr) -> Self {
37        Self {
38            args: vec![Box::new(expr)],
39        }
40    }
41
42    pub fn append_args(mut self: FunctionArgs, expr: Expr) -> Self {
43        self.args.push(Box::new(expr));
44        self
45    }
46
47    pub fn is_empty(&self) -> bool {
48        self.args.is_empty()
49    }
50
51    pub fn len(&self) -> usize {
52        self.args.len()
53    }
54
55    pub fn first(&self) -> Option<Box<Expr>> {
56        self.args.first().cloned()
57    }
58
59    pub fn last(&self) -> Option<Box<Expr>> {
60        self.args.last().cloned()
61    }
62}
63
64impl fmt::Display for FunctionArgs {
65    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
66        write!(f, "{}", join_vector(&self.args, ", ", false))
67    }
68}
69
70impl Prettier for FunctionArgs {
71    fn pretty(&self, level: usize, max: usize) -> String {
72        let mut v = vec![];
73        for ex in &self.args {
74            v.push(ex.pretty(level, max));
75        }
76        v.join(",\n")
77    }
78}
79
80/// Functions is a list of all functions supported by PromQL, including their types.
81#[derive(Debug, Clone, PartialEq, Eq)]
82#[cfg_attr(feature = "ser", derive(serde::Serialize))]
83#[cfg_attr(feature = "ser", serde(rename_all = "camelCase"))]
84pub struct Function {
85    pub name: &'static str,
86    pub arg_types: Vec<ValueType>,
87    /// Variadic cardinality follows Prometheus semantics:
88    /// 0 = exact args, >0 = bounded optional args, <0 = unbounded args.
89    pub variadic: i32,
90    pub return_type: ValueType,
91    pub experimental: bool,
92}
93
94impl Function {
95    pub fn new(
96        name: &'static str,
97        arg_types: Vec<ValueType>,
98        variadic: i32,
99        return_type: ValueType,
100        experimental: bool,
101    ) -> Self {
102        Self {
103            name,
104            arg_types,
105            variadic,
106            return_type,
107            experimental,
108        }
109    }
110}
111
112macro_rules! function {
113    ($name:expr, $arg_types:expr, $variadic:expr, $return_type:expr, $experimental:expr) => {
114        (
115            $name,
116            Function::new($name, $arg_types, $variadic, $return_type, $experimental),
117        )
118    };
119}
120
121lazy_static! {
122    static ref FUNCTIONS: HashMap<&'static str, Function> = HashMap::from([
123        function!("abs", vec![ValueType::Vector], 0, ValueType::Vector, false),
124        function!(
125            "absent",
126            vec![ValueType::Vector],
127            0,
128            ValueType::Vector,
129            false
130        ),
131        function!(
132            "absent_over_time",
133            vec![ValueType::Matrix],
134            0,
135            ValueType::Vector,
136            false
137        ),
138        function!("acos", vec![ValueType::Vector], 0, ValueType::Vector, false),
139        function!(
140            "acosh",
141            vec![ValueType::Vector],
142            0,
143            ValueType::Vector,
144            false
145        ),
146        function!("asin", vec![ValueType::Vector], 0, ValueType::Vector, false),
147        function!(
148            "asinh",
149            vec![ValueType::Vector],
150            0,
151            ValueType::Vector,
152            false
153        ),
154        function!("atan", vec![ValueType::Vector], 0, ValueType::Vector, false),
155        function!(
156            "atanh",
157            vec![ValueType::Vector],
158            0,
159            ValueType::Vector,
160            false
161        ),
162        function!(
163            "avg_over_time",
164            vec![ValueType::Matrix],
165            0,
166            ValueType::Vector,
167            false
168        ),
169        function!("ceil", vec![ValueType::Vector], 0, ValueType::Vector, false),
170        function!(
171            "changes",
172            vec![ValueType::Matrix],
173            0,
174            ValueType::Vector,
175            false
176        ),
177        function!(
178            "clamp",
179            vec![ValueType::Vector, ValueType::Scalar, ValueType::Scalar],
180            0,
181            ValueType::Vector,
182            false
183        ),
184        function!(
185            "clamp_max",
186            vec![ValueType::Vector, ValueType::Scalar],
187            0,
188            ValueType::Vector,
189            false
190        ),
191        function!(
192            "clamp_min",
193            vec![ValueType::Vector, ValueType::Scalar],
194            0,
195            ValueType::Vector,
196            false
197        ),
198        function!("cos", vec![ValueType::Vector], 0, ValueType::Vector, false),
199        function!("cosh", vec![ValueType::Vector], 0, ValueType::Vector, false),
200        function!(
201            "count_over_time",
202            vec![ValueType::Matrix],
203            0,
204            ValueType::Vector,
205            false
206        ),
207        function!(
208            "days_in_month",
209            vec![ValueType::Vector],
210            1,
211            ValueType::Vector,
212            false
213        ),
214        function!(
215            "day_of_month",
216            vec![ValueType::Vector],
217            1,
218            ValueType::Vector,
219            false
220        ),
221        function!(
222            "day_of_week",
223            vec![ValueType::Vector],
224            1,
225            ValueType::Vector,
226            false
227        ),
228        function!(
229            "day_of_year",
230            vec![ValueType::Vector],
231            1,
232            ValueType::Vector,
233            false
234        ),
235        function!("deg", vec![ValueType::Vector], 0, ValueType::Vector, false),
236        function!(
237            "delta",
238            vec![ValueType::Matrix],
239            0,
240            ValueType::Vector,
241            false
242        ),
243        function!(
244            "deriv",
245            vec![ValueType::Matrix],
246            0,
247            ValueType::Vector,
248            false
249        ),
250        function!("exp", vec![ValueType::Vector], 0, ValueType::Vector, false),
251        function!(
252            "floor",
253            vec![ValueType::Vector],
254            0,
255            ValueType::Vector,
256            false
257        ),
258        function!(
259            "histogram_count",
260            vec![ValueType::Vector],
261            0,
262            ValueType::Vector,
263            false
264        ),
265        function!(
266            "histogram_sum",
267            vec![ValueType::Vector],
268            0,
269            ValueType::Vector,
270            false
271        ),
272        function!(
273            "histogram_avg",
274            vec![ValueType::Vector],
275            0,
276            ValueType::Vector,
277            false
278        ),
279        function!(
280            "histogram_fraction",
281            vec![ValueType::Scalar, ValueType::Scalar, ValueType::Vector],
282            0,
283            ValueType::Vector,
284            false
285        ),
286        function!(
287            "histogram_quantile",
288            vec![ValueType::Scalar, ValueType::Vector],
289            0,
290            ValueType::Vector,
291            false
292        ),
293        function!(
294            "histogram_stddev",
295            vec![ValueType::Vector],
296            0,
297            ValueType::Vector,
298            false
299        ),
300        function!(
301            "histogram_stdvar",
302            vec![ValueType::Vector],
303            0,
304            ValueType::Vector,
305            false
306        ),
307        function!(
308            "double_exponential_smoothing",
309            vec![ValueType::Matrix, ValueType::Scalar, ValueType::Scalar],
310            0,
311            ValueType::Vector,
312            true
313        ),
314        function!(
315            "holt_winters",
316            vec![ValueType::Matrix, ValueType::Scalar, ValueType::Scalar],
317            0,
318            ValueType::Vector,
319            false
320        ),
321        function!("hour", vec![ValueType::Vector], 1, ValueType::Vector, false),
322        function!(
323            "idelta",
324            vec![ValueType::Matrix],
325            0,
326            ValueType::Vector,
327            false
328        ),
329        function!(
330            "increase",
331            vec![ValueType::Matrix],
332            0,
333            ValueType::Vector,
334            false
335        ),
336        function!(
337            "irate",
338            vec![ValueType::Matrix],
339            0,
340            ValueType::Vector,
341            false
342        ),
343        function!(
344            "label_replace",
345            vec![
346                ValueType::Vector,
347                ValueType::String,
348                ValueType::String,
349                ValueType::String,
350                ValueType::String
351            ],
352            0,
353            ValueType::Vector,
354            false
355        ),
356        function!(
357            "label_join",
358            vec![
359                ValueType::Vector,
360                ValueType::String,
361                ValueType::String,
362                ValueType::String
363            ],
364            -1,
365            ValueType::Vector,
366            false
367        ),
368        function!(
369            "last_over_time",
370            vec![ValueType::Matrix],
371            0,
372            ValueType::Vector,
373            false
374        ),
375        function!("ln", vec![ValueType::Vector], 0, ValueType::Vector, false),
376        function!(
377            "log10",
378            vec![ValueType::Vector],
379            0,
380            ValueType::Vector,
381            false
382        ),
383        function!("log2", vec![ValueType::Vector], 0, ValueType::Vector, false),
384        function!(
385            "max_over_time",
386            vec![ValueType::Matrix],
387            0,
388            ValueType::Vector,
389            false
390        ),
391        function!(
392            "min_over_time",
393            vec![ValueType::Matrix],
394            0,
395            ValueType::Vector,
396            false
397        ),
398        function!(
399            "minute",
400            vec![ValueType::Vector],
401            1,
402            ValueType::Vector,
403            false
404        ),
405        function!(
406            "month",
407            vec![ValueType::Vector],
408            1,
409            ValueType::Vector,
410            false
411        ),
412        function!("pi", vec![], 0, ValueType::Scalar, false),
413        function!(
414            "predict_linear",
415            vec![ValueType::Matrix, ValueType::Scalar],
416            0,
417            ValueType::Vector,
418            false
419        ),
420        function!(
421            "present_over_time",
422            vec![ValueType::Matrix],
423            0,
424            ValueType::Vector,
425            false
426        ),
427        function!(
428            "quantile_over_time",
429            vec![ValueType::Scalar, ValueType::Matrix],
430            0,
431            ValueType::Vector,
432            false
433        ),
434        function!("rad", vec![ValueType::Vector], 0, ValueType::Vector, false),
435        function!("rate", vec![ValueType::Matrix], 0, ValueType::Vector, false),
436        function!(
437            "resets",
438            vec![ValueType::Matrix],
439            0,
440            ValueType::Vector,
441            false
442        ),
443        function!(
444            "round",
445            vec![ValueType::Vector, ValueType::Scalar],
446            1,
447            ValueType::Vector,
448            false
449        ),
450        function!(
451            "scalar",
452            vec![ValueType::Vector],
453            0,
454            ValueType::Scalar,
455            false
456        ),
457        function!("sgn", vec![ValueType::Vector], 0, ValueType::Vector, false),
458        function!("sin", vec![ValueType::Vector], 0, ValueType::Vector, false),
459        function!("sinh", vec![ValueType::Vector], 0, ValueType::Vector, false),
460        function!("sort", vec![ValueType::Vector], 0, ValueType::Vector, false),
461        function!(
462            "sort_desc",
463            vec![ValueType::Vector],
464            0,
465            ValueType::Vector,
466            false
467        ),
468        function!(
469            "sort_by_label",
470            vec![ValueType::Vector, ValueType::String, ValueType::String],
471            -1,
472            ValueType::Vector,
473            true
474        ),
475        function!(
476            "sort_by_label_desc",
477            vec![ValueType::Vector, ValueType::String, ValueType::String],
478            -1,
479            ValueType::Vector,
480            true
481        ),
482        function!("sqrt", vec![ValueType::Vector], 0, ValueType::Vector, false),
483        function!(
484            "stddev_over_time",
485            vec![ValueType::Matrix],
486            0,
487            ValueType::Vector,
488            false
489        ),
490        function!(
491            "stdvar_over_time",
492            vec![ValueType::Matrix],
493            0,
494            ValueType::Vector,
495            false
496        ),
497        function!(
498            "sum_over_time",
499            vec![ValueType::Matrix],
500            0,
501            ValueType::Vector,
502            false
503        ),
504        function!("tan", vec![ValueType::Vector], 0, ValueType::Vector, false),
505        function!("tanh", vec![ValueType::Vector], 0, ValueType::Vector, false),
506        function!("time", vec![], 0, ValueType::Scalar, false),
507        function!(
508            "timestamp",
509            vec![ValueType::Vector],
510            0,
511            ValueType::Vector,
512            false
513        ),
514        function!(
515            "vector",
516            vec![ValueType::Scalar],
517            0,
518            ValueType::Vector,
519            false
520        ),
521        function!("year", vec![ValueType::Vector], 1, ValueType::Vector, false),
522    ]);
523}
524
525/// get_function returns a predefined Function object for the given name.
526pub(crate) fn get_function(name: &str) -> Option<Function> {
527    FUNCTIONS.get(name).cloned()
528}
529
530#[cfg(test)]
531mod tests {
532    use super::*;
533    use crate::parser::*;
534
535    #[test]
536    fn test_function_equality() {
537        let func = "month";
538        assert!(get_function(func).is_some());
539        assert_eq!(get_function(func), get_function(func));
540    }
541
542    #[test]
543    fn test_function_args_equality() {
544        assert_eq!(FunctionArgs::empty_args(), FunctionArgs::empty_args());
545
546        let arg1 = Expr::NumberLiteral(NumberLiteral::new(1.0));
547        let arg2 = Expr::StringLiteral(StringLiteral {
548            val: "prometheus".into(),
549        });
550        let args1 = FunctionArgs::new_args(arg1).append_args(arg2);
551
552        let arg1 = Expr::NumberLiteral(NumberLiteral::new(0.5 + 0.5));
553        let arg2 = Expr::StringLiteral(StringLiteral {
554            val: String::from("prometheus"),
555        });
556        let args2 = FunctionArgs::new_args(arg1).append_args(arg2);
557
558        assert_eq!(args1, args2);
559    }
560
561    #[test]
562    fn test_args_display() {
563        let cases = vec![
564            (
565                FunctionArgs::new_args(Expr::from(VectorSelector::from("up"))),
566                "up",
567            ),
568            (
569                FunctionArgs::empty_args()
570                    .append_args(Expr::from("src1"))
571                    .append_args(Expr::from("src2"))
572                    .append_args(Expr::from("src3")),
573                r#""src1", "src2", "src3""#,
574            ),
575        ];
576
577        for (args, expect) in cases {
578            assert_eq!(expect, args.to_string())
579        }
580    }
581
582    #[test]
583    fn test_function_metadata() {
584        let round = get_function("round").unwrap();
585        assert_eq!(round.variadic, 1);
586        assert!(!round.experimental);
587
588        let label_join = get_function("label_join").unwrap();
589        assert_eq!(label_join.variadic, -1);
590        assert!(!label_join.experimental);
591
592        let sort_by_label = get_function("sort_by_label").unwrap();
593        assert_eq!(sort_by_label.variadic, -1);
594        assert!(sort_by_label.experimental);
595
596        let rate = get_function("rate").unwrap();
597        assert_eq!(rate.variadic, 0);
598        assert!(!rate.experimental);
599    }
600}