promql_parser/parser/

function.rs

// 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::{HashMap, HashSet};
use std::fmt;

use lazy_static::lazy_static;

use crate::parser::value::ValueType;
use crate::parser::{Expr, Prettier};
use crate::util::join_vector;

/// called by func in Call
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "ser", derive(serde::Serialize))]
pub struct FunctionArgs {
    pub args: Vec<Box<Expr>>,
}

impl FunctionArgs {
    pub fn empty_args() -> Self {
        Self { args: vec![] }
    }

    pub fn new_args(expr: Expr) -> Self {
        Self {
            args: vec![Box::new(expr)],
        }
    }

    pub fn append_args(mut self: FunctionArgs, expr: Expr) -> Self {
        self.args.push(Box::new(expr));
        self
    }

    pub fn is_empty(&self) -> bool {
        self.args.is_empty()
    }

    pub fn len(&self) -> usize {
        self.args.len()
    }

    pub fn first(&self) -> Option<Box<Expr>> {
        self.args.first().cloned()
    }

    pub fn last(&self) -> Option<Box<Expr>> {
        self.args.last().cloned()
    }
}

impl fmt::Display for FunctionArgs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", join_vector(&self.args, ", ", false))
    }
}

impl Prettier for FunctionArgs {
    fn pretty(&self, level: usize, max: usize) -> String {
        let mut v = vec![];
        for ex in &self.args {
            v.push(ex.pretty(level, max));
        }
        v.join(",\n")
    }
}

/// Functions is a list of all functions supported by PromQL, including their types.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "ser", derive(serde::Serialize))]
#[cfg_attr(feature = "ser", serde(rename_all = "camelCase"))]
pub struct Function {
    pub name: &'static str,
    pub arg_types: Vec<ValueType>,
    /// Variadic cardinality follows Prometheus semantics:
    /// 0 = exact args, >0 = bounded optional args, <0 = unbounded args.
    pub variadic: i32,
    pub return_type: ValueType,
    pub experimental: bool,
}

impl Function {
    pub fn new(
        name: &'static str,
        arg_types: Vec<ValueType>,
        variadic: i32,
        return_type: ValueType,
        experimental: bool,
    ) -> Self {
        Self {
            name,
            arg_types,
            variadic,
            return_type,
            experimental,
        }
    }
}

macro_rules! map {
    // if variadic args, then the last is the variadic one
    ($(($name:literal, $arg:expr, $ret:expr)),*) => (
        {
            let mut m: HashMap<&'static str, Function> = HashMap::new();
            $(
                let variadic = FUNCTION_VARIADIC_CARDINALITY
                    .get($name)
                    .copied()
                    .unwrap_or_default();
                let experimental = EXPERIMENTAL_FUNCTIONS.contains($name);
                let func = Function::new($name, $arg, variadic, $ret, experimental);
                m.insert($name, func);
            )*
            m
        }
    );
}

lazy_static! {
    // Variadic cardinality follows Prometheus semantics:
    // 0 = exact args, >0 = bounded optional args, <0 = unbounded args.
    // Only non-zero entries are listed here; missing entries default to 0.
    static ref FUNCTION_VARIADIC_CARDINALITY: HashMap<&'static str, i32> = HashMap::from([
        ("days_in_month", 1),
        ("day_of_year", 1),
        ("day_of_month", 1),
        ("day_of_week", 1),
        ("year", 1),
        ("month", 1),
        ("hour", 1),
        ("minute", 1),
        ("label_join", -1),
        ("sort_by_label", -1),
        ("sort_by_label_desc", -1),
        ("round", 1),
    ]);
    // Functions currently supported by this parser that are marked experimental in Prometheus.
    // Keep this set in sync when adding new functions to FUNCTIONS.
    static ref EXPERIMENTAL_FUNCTIONS: HashSet<&'static str> = HashSet::from([
        "double_exponential_smoothing",
        "sort_by_label",
        "sort_by_label_desc",
    ]);
    static ref FUNCTIONS: HashMap<&'static str, Function> = map!(
        ("abs", vec![ValueType::Vector], ValueType::Vector),
        ("absent", vec![ValueType::Vector], ValueType::Vector),
        (
            "absent_over_time",
            vec![ValueType::Matrix],
            ValueType::Vector
        ),
        ("acos", vec![ValueType::Vector], ValueType::Vector),
        ("acosh", vec![ValueType::Vector], ValueType::Vector),
        ("asin", vec![ValueType::Vector], ValueType::Vector),
        ("asinh", vec![ValueType::Vector], ValueType::Vector),
        ("atan", vec![ValueType::Vector], ValueType::Vector),
        ("atanh", vec![ValueType::Vector], ValueType::Vector),
        ("avg_over_time", vec![ValueType::Matrix], ValueType::Vector),
        ("ceil", vec![ValueType::Vector], ValueType::Vector),
        ("changes", vec![ValueType::Matrix], ValueType::Vector),
        (
            "clamp",
            vec![ValueType::Vector, ValueType::Scalar, ValueType::Scalar],
            ValueType::Vector
        ),
        (
            "clamp_max",
            vec![ValueType::Vector, ValueType::Scalar],
            ValueType::Vector
        ),
        (
            "clamp_min",
            vec![ValueType::Vector, ValueType::Scalar],
            ValueType::Vector
        ),
        ("cos", vec![ValueType::Vector], ValueType::Vector),
        ("cosh", vec![ValueType::Vector], ValueType::Vector),
        (
            "count_over_time",
            vec![ValueType::Matrix],
            ValueType::Vector
        ),
        ("days_in_month", vec![ValueType::Vector], ValueType::Vector),
        ("day_of_month", vec![ValueType::Vector], ValueType::Vector),
        ("day_of_week", vec![ValueType::Vector], ValueType::Vector),
        ("day_of_year", vec![ValueType::Vector], ValueType::Vector),
        ("deg", vec![ValueType::Vector], ValueType::Vector),
        ("delta", vec![ValueType::Matrix], ValueType::Vector),
        ("deriv", vec![ValueType::Matrix], ValueType::Vector),
        ("exp", vec![ValueType::Vector], ValueType::Vector),
        ("floor", vec![ValueType::Vector], ValueType::Vector),
        (
            "histogram_count",
            vec![ValueType::Vector],
            ValueType::Vector
        ),
        ("histogram_sum", vec![ValueType::Vector], ValueType::Vector),
        ("histogram_avg", vec![ValueType::Vector], ValueType::Vector),
        (
            "histogram_fraction",
            vec![ValueType::Scalar, ValueType::Scalar, ValueType::Vector],
            ValueType::Vector
        ),
        (
            "histogram_quantile",
            vec![ValueType::Scalar, ValueType::Vector],
            ValueType::Vector
        ),
        ("histogram_stddev", vec![ValueType::Vector], ValueType::Vector),
        ("histogram_stdvar", vec![ValueType::Vector], ValueType::Vector),
        (
            "double_exponential_smoothing",
            vec![ValueType::Matrix, ValueType::Scalar, ValueType::Scalar],
            ValueType::Vector
        ),
        // Kept for backward compatibility; Prometheus 3.x renamed this function.
        (
            "holt_winters",
            vec![ValueType::Matrix, ValueType::Scalar, ValueType::Scalar],
            ValueType::Vector
        ),
        ("hour", vec![ValueType::Vector], ValueType::Vector),
        ("idelta", vec![ValueType::Matrix], ValueType::Vector),
        ("increase", vec![ValueType::Matrix], ValueType::Vector),
        ("irate", vec![ValueType::Matrix], ValueType::Vector),
        (
            "label_replace",
            vec![
                ValueType::Vector,
                ValueType::String,
                ValueType::String,
                ValueType::String,
                ValueType::String,
            ],
            ValueType::Vector
        ),
        (
            "label_join",
            vec![
                ValueType::Vector,
                ValueType::String,
                ValueType::String,
                ValueType::String,
            ],
            ValueType::Vector
        ),
        ("last_over_time", vec![ValueType::Matrix], ValueType::Vector),
        ("ln", vec![ValueType::Vector], ValueType::Vector),
        ("log10", vec![ValueType::Vector], ValueType::Vector),
        ("log2", vec![ValueType::Vector], ValueType::Vector),
        ("max_over_time", vec![ValueType::Matrix], ValueType::Vector),
        ("min_over_time", vec![ValueType::Matrix], ValueType::Vector),
        ("minute", vec![ValueType::Vector], ValueType::Vector),
        ("month", vec![ValueType::Vector], ValueType::Vector),
        ("pi", vec![], ValueType::Scalar),
        (
            "predict_linear",
            vec![ValueType::Matrix, ValueType::Scalar],
            ValueType::Vector
        ),
        (
            "present_over_time",
            vec![ValueType::Matrix],
            ValueType::Vector
        ),
        (
            "quantile_over_time",
            vec![ValueType::Scalar, ValueType::Matrix],
            ValueType::Vector
        ),
        ("rad", vec![ValueType::Vector], ValueType::Vector),
        ("rate", vec![ValueType::Matrix], ValueType::Vector),
        ("resets", vec![ValueType::Matrix], ValueType::Vector),
        (
            "round",
            vec![ValueType::Vector, ValueType::Scalar],
            ValueType::Vector
        ),
        ("scalar", vec![ValueType::Vector], ValueType::Scalar),
        ("sgn", vec![ValueType::Vector], ValueType::Vector),
        ("sin", vec![ValueType::Vector], ValueType::Vector),
        ("sinh", vec![ValueType::Vector], ValueType::Vector),
        ("sort", vec![ValueType::Vector], ValueType::Vector),
        ("sort_desc", vec![ValueType::Vector], ValueType::Vector),
        (
            "sort_by_label",
            vec![ValueType::Vector, ValueType::String, ValueType::String],
            ValueType::Vector
        ),
        (
            "sort_by_label_desc",
            vec![ValueType::Vector, ValueType::String, ValueType::String],
            ValueType::Vector
        ),
        ("sqrt", vec![ValueType::Vector], ValueType::Vector),
        (
            "stddev_over_time",
            vec![ValueType::Matrix],
            ValueType::Vector
        ),
        (
            "stdvar_over_time",
            vec![ValueType::Matrix],
            ValueType::Vector
        ),
        ("sum_over_time", vec![ValueType::Matrix], ValueType::Vector),
        ("tan", vec![ValueType::Vector], ValueType::Vector),
        ("tanh", vec![ValueType::Vector], ValueType::Vector),
        ("time", vec![], ValueType::Scalar),
        ("timestamp", vec![ValueType::Vector], ValueType::Vector),
        ("vector", vec![ValueType::Scalar], ValueType::Vector),
        ("year", vec![ValueType::Vector], ValueType::Vector)
    );
}

/// get_function returns a predefined Function object for the given name.
pub(crate) fn get_function(name: &str) -> Option<Function> {
    FUNCTIONS.get(name).cloned()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::*;

    #[test]
    fn test_function_equality() {
        let func = "month";
        assert!(get_function(func).is_some());
        assert_eq!(get_function(func), get_function(func));
    }

    #[test]
    fn test_function_args_equality() {
        assert_eq!(FunctionArgs::empty_args(), FunctionArgs::empty_args());

        let arg1 = Expr::NumberLiteral(NumberLiteral::new(1.0));
        let arg2 = Expr::StringLiteral(StringLiteral {
            val: "prometheus".into(),
        });
        let args1 = FunctionArgs::new_args(arg1).append_args(arg2);

        let arg1 = Expr::NumberLiteral(NumberLiteral::new(0.5 + 0.5));
        let arg2 = Expr::StringLiteral(StringLiteral {
            val: String::from("prometheus"),
        });
        let args2 = FunctionArgs::new_args(arg1).append_args(arg2);

        assert_eq!(args1, args2);
    }

    #[test]
    fn test_args_display() {
        let cases = vec![
            (
                FunctionArgs::new_args(Expr::from(VectorSelector::from("up"))),
                "up",
            ),
            (
                FunctionArgs::empty_args()
                    .append_args(Expr::from("src1"))
                    .append_args(Expr::from("src2"))
                    .append_args(Expr::from("src3")),
                r#""src1", "src2", "src3""#,
            ),
        ];

        for (args, expect) in cases {
            assert_eq!(expect, args.to_string())
        }
    }

    #[test]
    fn test_function_metadata() {
        let round = get_function("round").unwrap();
        assert_eq!(round.variadic, 1);
        assert!(!round.experimental);

        let label_join = get_function("label_join").unwrap();
        assert_eq!(label_join.variadic, -1);
        assert!(!label_join.experimental);

        let sort_by_label = get_function("sort_by_label").unwrap();
        assert_eq!(sort_by_label.variadic, -1);
        assert!(sort_by_label.experimental);

        let rate = get_function("rate").unwrap();
        assert_eq!(rate.variadic, 0);
        assert!(!rate.experimental);
    }
}