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!("end", vec![], 0, ValueType::Scalar, true),
251        function!("exp", vec![ValueType::Vector], 0, ValueType::Vector, false),
252        function!(
253            "first_over_time",
254            vec![ValueType::Matrix],
255            0,
256            ValueType::Vector,
257            true
258        ),
259        function!(
260            "floor",
261            vec![ValueType::Vector],
262            0,
263            ValueType::Vector,
264            false
265        ),
266        function!(
267            "histogram_avg",
268            vec![ValueType::Vector],
269            0,
270            ValueType::Vector,
271            false
272        ),
273        function!(
274            "histogram_count",
275            vec![ValueType::Vector],
276            0,
277            ValueType::Vector,
278            false
279        ),
280        function!(
281            "histogram_fraction",
282            vec![ValueType::Scalar, ValueType::Scalar, ValueType::Vector],
283            0,
284            ValueType::Vector,
285            false
286        ),
287        function!(
288            "histogram_quantile",
289            vec![ValueType::Scalar, ValueType::Vector],
290            0,
291            ValueType::Vector,
292            false
293        ),
294        function!(
295            "histogram_quantiles",
296            vec![
297                ValueType::Vector,
298                ValueType::String,
299                ValueType::Scalar,
300                ValueType::Scalar
301            ],
302            9,
303            ValueType::Vector,
304            true
305        ),
306        function!(
307            "histogram_stddev",
308            vec![ValueType::Vector],
309            0,
310            ValueType::Vector,
311            false
312        ),
313        function!(
314            "histogram_stdvar",
315            vec![ValueType::Vector],
316            0,
317            ValueType::Vector,
318            false
319        ),
320        function!(
321            "histogram_sum",
322            vec![ValueType::Vector],
323            0,
324            ValueType::Vector,
325            false
326        ),
327        function!(
328            "info",
329            vec![ValueType::Vector, ValueType::Vector],
330            1,
331            ValueType::Vector,
332            true
333        ),
334        function!(
335            "double_exponential_smoothing",
336            vec![ValueType::Matrix, ValueType::Scalar, ValueType::Scalar],
337            0,
338            ValueType::Vector,
339            true
340        ),
341        function!("hour", vec![ValueType::Vector], 1, ValueType::Vector, false),
342        function!(
343            "idelta",
344            vec![ValueType::Matrix],
345            0,
346            ValueType::Vector,
347            false
348        ),
349        function!(
350            "increase",
351            vec![ValueType::Matrix],
352            0,
353            ValueType::Vector,
354            false
355        ),
356        function!(
357            "irate",
358            vec![ValueType::Matrix],
359            0,
360            ValueType::Vector,
361            false
362        ),
363        function!(
364            "label_replace",
365            vec![
366                ValueType::Vector,
367                ValueType::String,
368                ValueType::String,
369                ValueType::String,
370                ValueType::String
371            ],
372            0,
373            ValueType::Vector,
374            false
375        ),
376        function!(
377            "label_join",
378            vec![
379                ValueType::Vector,
380                ValueType::String,
381                ValueType::String,
382                ValueType::String
383            ],
384            -1,
385            ValueType::Vector,
386            false
387        ),
388        function!(
389            "max_of",
390            vec![ValueType::Scalar, ValueType::Scalar],
391            0,
392            ValueType::Scalar,
393            true
394        ),
395        function!(
396            "last_over_time",
397            vec![ValueType::Matrix],
398            0,
399            ValueType::Vector,
400            false
401        ),
402        function!(
403            "min_of",
404            vec![ValueType::Scalar, ValueType::Scalar],
405            0,
406            ValueType::Scalar,
407            true
408        ),
409        function!("ln", vec![ValueType::Vector], 0, ValueType::Vector, false),
410        function!(
411            "log10",
412            vec![ValueType::Vector],
413            0,
414            ValueType::Vector,
415            false
416        ),
417        function!("log2", vec![ValueType::Vector], 0, ValueType::Vector, false),
418        function!(
419            "mad_over_time",
420            vec![ValueType::Matrix],
421            0,
422            ValueType::Vector,
423            true
424        ),
425        function!(
426            "max_over_time",
427            vec![ValueType::Matrix],
428            0,
429            ValueType::Vector,
430            false
431        ),
432        function!(
433            "min_over_time",
434            vec![ValueType::Matrix],
435            0,
436            ValueType::Vector,
437            false
438        ),
439        function!(
440            "ts_of_first_over_time",
441            vec![ValueType::Matrix],
442            0,
443            ValueType::Vector,
444            true
445        ),
446        function!(
447            "ts_of_last_over_time",
448            vec![ValueType::Matrix],
449            0,
450            ValueType::Vector,
451            true
452        ),
453        function!(
454            "ts_of_max_over_time",
455            vec![ValueType::Matrix],
456            0,
457            ValueType::Vector,
458            true
459        ),
460        function!(
461            "ts_of_min_over_time",
462            vec![ValueType::Matrix],
463            0,
464            ValueType::Vector,
465            true
466        ),
467        function!(
468            "minute",
469            vec![ValueType::Vector],
470            1,
471            ValueType::Vector,
472            false
473        ),
474        function!(
475            "month",
476            vec![ValueType::Vector],
477            1,
478            ValueType::Vector,
479            false
480        ),
481        function!("pi", vec![], 0, ValueType::Scalar, false),
482        function!("range", vec![], 0, ValueType::Scalar, true),
483        function!(
484            "predict_linear",
485            vec![ValueType::Matrix, ValueType::Scalar],
486            0,
487            ValueType::Vector,
488            false
489        ),
490        function!(
491            "present_over_time",
492            vec![ValueType::Matrix],
493            0,
494            ValueType::Vector,
495            false
496        ),
497        function!(
498            "quantile_over_time",
499            vec![ValueType::Scalar, ValueType::Matrix],
500            0,
501            ValueType::Vector,
502            false
503        ),
504        function!("rad", vec![ValueType::Vector], 0, ValueType::Vector, false),
505        function!("rate", vec![ValueType::Matrix], 0, ValueType::Vector, false),
506        function!(
507            "resets",
508            vec![ValueType::Matrix],
509            0,
510            ValueType::Vector,
511            false
512        ),
513        function!(
514            "round",
515            vec![ValueType::Vector, ValueType::Scalar],
516            1,
517            ValueType::Vector,
518            false
519        ),
520        function!(
521            "scalar",
522            vec![ValueType::Vector],
523            0,
524            ValueType::Scalar,
525            false
526        ),
527        function!("sgn", vec![ValueType::Vector], 0, ValueType::Vector, false),
528        function!("start", vec![], 0, ValueType::Scalar, true),
529        function!("step", vec![], 0, ValueType::Scalar, true),
530        function!("sin", vec![ValueType::Vector], 0, ValueType::Vector, false),
531        function!("sinh", vec![ValueType::Vector], 0, ValueType::Vector, false),
532        function!("sort", vec![ValueType::Vector], 0, ValueType::Vector, false),
533        function!(
534            "sort_desc",
535            vec![ValueType::Vector],
536            0,
537            ValueType::Vector,
538            false
539        ),
540        function!(
541            "sort_by_label",
542            vec![ValueType::Vector, ValueType::String],
543            -1,
544            ValueType::Vector,
545            true
546        ),
547        function!(
548            "sort_by_label_desc",
549            vec![ValueType::Vector, ValueType::String],
550            -1,
551            ValueType::Vector,
552            true
553        ),
554        function!("sqrt", vec![ValueType::Vector], 0, ValueType::Vector, false),
555        function!(
556            "stddev_over_time",
557            vec![ValueType::Matrix],
558            0,
559            ValueType::Vector,
560            false
561        ),
562        function!(
563            "stdvar_over_time",
564            vec![ValueType::Matrix],
565            0,
566            ValueType::Vector,
567            false
568        ),
569        function!(
570            "sum_over_time",
571            vec![ValueType::Matrix],
572            0,
573            ValueType::Vector,
574            false
575        ),
576        function!("tan", vec![ValueType::Vector], 0, ValueType::Vector, false),
577        function!("tanh", vec![ValueType::Vector], 0, ValueType::Vector, false),
578        function!("time", vec![], 0, ValueType::Scalar, false),
579        function!(
580            "timestamp",
581            vec![ValueType::Vector],
582            0,
583            ValueType::Vector,
584            false
585        ),
586        function!(
587            "vector",
588            vec![ValueType::Scalar],
589            0,
590            ValueType::Vector,
591            false
592        ),
593        function!("year", vec![ValueType::Vector], 1, ValueType::Vector, false),
594    ]);
595}
596
597/// get_function returns a predefined Function object for the given name.
598pub(crate) fn get_function(name: &str) -> Option<Function> {
599    FUNCTIONS.get(name).cloned()
600}
601
602#[cfg(test)]
603mod tests {
604    use super::*;
605    use crate::parser::*;
606
607    #[test]
608    fn test_function_equality() {
609        let func = "month";
610        assert!(get_function(func).is_some());
611        assert_eq!(get_function(func), get_function(func));
612    }
613
614    #[test]
615    fn test_function_args_equality() {
616        assert_eq!(FunctionArgs::empty_args(), FunctionArgs::empty_args());
617
618        let arg1 = Expr::NumberLiteral(NumberLiteral::new(1.0));
619        let arg2 = Expr::StringLiteral(StringLiteral {
620            val: "prometheus".into(),
621        });
622        let args1 = FunctionArgs::new_args(arg1).append_args(arg2);
623
624        let arg1 = Expr::NumberLiteral(NumberLiteral::new(0.5 + 0.5));
625        let arg2 = Expr::StringLiteral(StringLiteral {
626            val: String::from("prometheus"),
627        });
628        let args2 = FunctionArgs::new_args(arg1).append_args(arg2);
629
630        assert_eq!(args1, args2);
631    }
632
633    #[test]
634    fn test_args_display() {
635        let cases = vec![
636            (
637                FunctionArgs::new_args(Expr::from(VectorSelector::from("up"))),
638                "up",
639            ),
640            (
641                FunctionArgs::empty_args()
642                    .append_args(Expr::from("src1"))
643                    .append_args(Expr::from("src2"))
644                    .append_args(Expr::from("src3")),
645                r#""src1", "src2", "src3""#,
646            ),
647        ];
648
649        for (args, expect) in cases {
650            assert_eq!(expect, args.to_string())
651        }
652    }
653
654    #[test]
655    fn test_function_metadata() {
656        let round = get_function("round").unwrap();
657        assert_eq!(round.variadic, 1);
658        assert!(!round.experimental);
659
660        let label_join = get_function("label_join").unwrap();
661        assert_eq!(label_join.variadic, -1);
662        assert!(!label_join.experimental);
663
664        let sort_by_label = get_function("sort_by_label").unwrap();
665        assert_eq!(sort_by_label.variadic, -1);
666        assert!(sort_by_label.experimental);
667
668        let rate = get_function("rate").unwrap();
669        assert_eq!(rate.variadic, 0);
670        assert!(!rate.experimental);
671
672        for func_name in ["max_of", "min_of"] {
673            let func = get_function(func_name).unwrap();
674            assert_eq!(func.arg_types, vec![ValueType::Scalar, ValueType::Scalar]);
675            assert_eq!(func.variadic, 0);
676            assert_eq!(func.return_type, ValueType::Scalar);
677            assert!(func.experimental);
678        }
679    }
680}