vrl/stdlib/
tally.rs

1use crate::compiler::prelude::*;
2use std::collections::{BTreeMap, HashMap};
3
4fn tally(value: Value) -> Resolved {
5    let value = value.try_array()?;
6    #[allow(clippy::mutable_key_type)] // false positive due to bytes::Bytes
7    let mut map: HashMap<Bytes, usize> = HashMap::new();
8    for value in value {
9        if let Value::Bytes(value) = value {
10            *map.entry(value).or_insert(0) += 1;
11        } else {
12            return Err(format!("all values must be strings, found: {value:?}").into());
13        }
14    }
15    let map: BTreeMap<_, _> = map
16        .into_iter()
17        .map(|(k, v)| {
18            (
19                String::from_utf8_lossy(&k).into_owned().into(),
20                Value::from(v),
21            )
22        })
23        .collect();
24    Ok(map.into())
25}
26
27#[derive(Clone, Copy, Debug)]
28pub struct Tally;
29
30impl Function for Tally {
31    fn identifier(&self) -> &'static str {
32        "tally"
33    }
34
35    fn usage(&self) -> &'static str {
36        "Counts the occurrences of each string value in the provided array and returns an object with the counts."
37    }
38
39    fn category(&self) -> &'static str {
40        Category::Enumerate.as_ref()
41    }
42
43    fn return_kind(&self) -> u16 {
44        kind::OBJECT
45    }
46
47    fn examples(&self) -> &'static [Example] {
48        &[example! {
49            title: "tally",
50            source: r#"tally!(["foo", "bar", "foo", "baz"])"#,
51            result: Ok(r#"{"foo": 2, "bar": 1, "baz": 1}"#),
52        }]
53    }
54
55    fn compile(
56        &self,
57        _state: &state::TypeState,
58        _ctx: &mut FunctionCompileContext,
59        arguments: ArgumentList,
60    ) -> Compiled {
61        let value = arguments.required("value");
62
63        Ok(TallyFn { value }.as_expr())
64    }
65
66    fn parameters(&self) -> &'static [Parameter] {
67        const PARAMETERS: &[Parameter] = &[Parameter::required(
68            "value",
69            kind::ARRAY,
70            "The array of strings to count occurrences for.",
71        )];
72        PARAMETERS
73    }
74}
75
76#[derive(Debug, Clone)]
77pub(crate) struct TallyFn {
78    value: Box<dyn Expression>,
79}
80
81impl FunctionExpression for TallyFn {
82    fn resolve(&self, ctx: &mut Context) -> Resolved {
83        let value = self.value.resolve(ctx)?;
84        tally(value)
85    }
86
87    fn type_def(&self, _: &state::TypeState) -> TypeDef {
88        TypeDef::object(Collection::from_unknown(Kind::integer())).fallible()
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use crate::value;
96
97    test_function![
98        tally => Tally;
99
100        default {
101            args: func_args![
102                value: value!(["bar", "foo", "baz", "foo"]),
103            ],
104            want: Ok(value!({"bar": 1, "foo": 2, "baz": 1})),
105            tdef: TypeDef::object(Collection::from_unknown(Kind::integer())).fallible(),
106        }
107
108        non_string_values {
109            args: func_args![
110                value: value!(["foo", [1,2,3], "123abc", 1, true, [1,2,3], "foo", true, 1]),
111            ],
112            want: Err("all values must be strings, found: Array([Integer(1), Integer(2), Integer(3)])"),
113            tdef: TypeDef::object(Collection::from_unknown(Kind::integer())).fallible(),
114        }
115    ];
116}
vrl/stdlib/tally.rs

vrl/stdlib/
tally.rs
