1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179

use serde::Serialize;

use crate::search::*;
use crate::types::Map;
use crate::util::*;

/// Specifies the path to the buckets to filter in a bucket selector aggregation.
///
/// This can either be a single path, referencing a single metric, or multiple paths
/// in case of more complex aggregations.
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(untagged)]
pub enum BucketsPath {
    /// A single path referencing a metric.
    Single(String),
    /// Multiple paths in the form of key-value pairs.
    /// Each key corresponds to an alias, and each value is a path to the metric.
    Multi(Map<String, String>),
}

impl From<&str> for BucketsPath {
    fn from(path: &str) -> Self {
        BucketsPath::Single(path.to_string())
    }
}

impl From<String> for BucketsPath {
    fn from(path: String) -> Self {
        BucketsPath::Single(path)
    }
}

impl From<Vec<(&str, &str)>> for BucketsPath {
    fn from(paths: Vec<(&str, &str)>) -> Self {
        BucketsPath::Multi(
            paths
                .into_iter()
                .map(|(k, v)| (k.to_string(), v.to_string()))
                .collect(),
        )
    }
}

impl From<Vec<(String, String)>> for BucketsPath {
    fn from(paths: Vec<(String, String)>) -> Self {
        BucketsPath::Multi(paths.into_iter().collect())
    }
}

#[derive(Debug, Clone, Serialize, PartialEq)]
/// A parent pipeline aggregation which allows the user to specify a script to run
/// for each bucket on the set of values returned by another aggregation.
pub struct BucketSelectorAggregation {
    bucket_selector: BucketSelectorAggregationInner,

    #[serde(skip_serializing_if = "ShouldSkip::should_skip")]
    aggs: Aggregations,
}

#[derive(Debug, Clone, Serialize, PartialEq)]
struct BucketSelectorAggregationInner {
    buckets_path: BucketsPath,

    script: Script,

    #[serde(skip_serializing_if = "ShouldSkip::should_skip")]
    gap_policy: Option<GapPolicy>,

    #[serde(skip_serializing_if = "ShouldSkip::should_skip")]
    format: Option<String>,
}

impl Aggregation {
    /// Creates an instance of [`BucketSelectorAggregation`]
    ///
    /// - `buckets_path` - the path to the buckets to filter
    /// - `script` - the script to execute for filtering
    pub fn bucket_selector<B, S>(buckets_path: B, script: S) -> BucketSelectorAggregation
    where
        B: Into<BucketsPath>,
        S: Into<Script>,
    {
        BucketSelectorAggregation {
            bucket_selector: BucketSelectorAggregationInner {
                buckets_path: buckets_path.into(),
                script: script.into(),
                gap_policy: None,
                format: None,
            },
            aggs: Aggregations::new(),
        }
    }
}

impl BucketSelectorAggregation {
    /// Sets the gap policy for the bucket selector aggregation.
    ///
    /// The gap policy determines how documents with missing values are treated.
    /// The default policy is to skip gaps.
    pub fn gap_policy(mut self, gap_policy: GapPolicy) -> Self {
        self.bucket_selector.gap_policy = Some(gap_policy);
        self
    }

    /// Sets the format for the bucket selector aggregation.
    ///
    /// The format parameter can be used to specify the format of the output values.
    pub fn format<T>(mut self, format: T) -> Self
    where
        T: ToString,
    {
        self.bucket_selector.format = Some(format.to_string());
        self
    }

    add_aggregate!();
}

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

    #[test]
    fn serialization() {
        assert_serialize_aggregation(
            Aggregation::bucket_selector(
                "the_sum",
                Script::source("params.the_sum > 1000").lang("painless"),
            ),
            json!({
                "bucket_selector": {
                    "buckets_path": "the_sum",
                    "script": {
                        "lang": "painless",
                         "source": "params.the_sum > 1000"
                    }
                }
            }),
        );

        assert_serialize_aggregation(
            Aggregation::bucket_selector(
                "the_sum",
                Script::source("params.the_sum > 1000").lang("painless"),
            )
            .gap_policy(GapPolicy::Skip)
            .format("###.00"),
            json!({
                "bucket_selector": {
                    "buckets_path": "the_sum",
                   "script": {
                        "lang": "painless",
                        "source": "params.the_sum > 1000"
                    },
                    "gap_policy": "skip",
                    "format": "###.00"
                }
            }),
        );

        assert_serialize_aggregation(
            Aggregation::bucket_selector(
                vec![("sum_value", "the_sum")],
                Script::source("params.sum_value > 1000").lang("painless"),
            ),
            json!({
                "bucket_selector": {
                    "buckets_path": {
                        "sum_value": "the_sum"
                    },
                    "script": {
                        "lang": "painless",
                        "source":"params.sum_value > 1000"
                    }
                }
            }),
        );
    }
}