From 0a9cccabd73aa2e3ffa0f876a81345dd39812a68 Mon Sep 17 00:00:00 2001 From: Evaldas Buinauskas Date: Mon, 8 Apr 2024 10:24:55 +0300 Subject: [PATCH] feat: add terms aggregation include, exclude --- .../aggregations/bucket/terms_aggregation.rs | 41 ++++++++- src/search/aggregations/params/mod.rs | 4 + .../aggregations/params/terms_exclude.rs | 61 +++++++++++++ .../aggregations/params/terms_include.rs | 90 +++++++++++++++++++ 4 files changed, 194 insertions(+), 2 deletions(-) create mode 100644 src/search/aggregations/params/terms_exclude.rs create mode 100644 src/search/aggregations/params/terms_include.rs diff --git a/src/search/aggregations/bucket/terms_aggregation.rs b/src/search/aggregations/bucket/terms_aggregation.rs index 0792dad..6e8d224 100644 --- a/src/search/aggregations/bucket/terms_aggregation.rs +++ b/src/search/aggregations/bucket/terms_aggregation.rs @@ -31,6 +31,12 @@ struct TermsAggregationInner { #[serde(skip_serializing_if = "ShouldSkip::should_skip")] missing: Option, + + #[serde(skip_serializing_if = "ShouldSkip::should_skip")] + include: Option, + + #[serde(skip_serializing_if = "ShouldSkip::should_skip")] + exclude: Option, } impl Aggregation { @@ -49,6 +55,8 @@ impl Aggregation { order: Default::default(), min_doc_count: None, missing: None, + include: None, + exclude: None, }, aggs: Aggregations::new(), } @@ -115,6 +123,24 @@ impl TermsAggregation { self } + /// The `include` parameter can be set to include only specific terms in the response. + pub fn include(mut self, include: T) -> Self + where + T: Into, + { + self.terms.include = Some(include.into()); + self + } + + /// The `exclude` parameter can be set to exclude specific terms from the response. + pub fn exclude(mut self, exclude: T) -> Self + where + T: Into, + { + self.terms.exclude = Some(exclude.into()); + self + } + add_aggregate!(); } @@ -155,15 +181,22 @@ mod tests { .size(0) .order(TermsOrder::ascending("test_order")) .missing(123) + .include(["mazda", "honda"]) + .exclude("water_.*") .aggregate( "test_sub_agg", - Aggregation::terms("test_field2").size(3).missing(false), + Aggregation::terms("test_field2") + .size(3) + .missing(false) + .include([0, 20]), ), json!({ "terms": { "field": "test_field", "size": 0, "missing": 123, + "include": ["mazda", "honda"], + "exclude": "water_.*", "order": [ { "test_order": "asc" } ] @@ -173,7 +206,11 @@ mod tests { "terms": { "field": "test_field2", "size": 3, - "missing": false + "missing": false, + "include": { + "partition": 0, + "num_partitions": 20 + } } } } diff --git a/src/search/aggregations/params/mod.rs b/src/search/aggregations/params/mod.rs index dc4c4ab..b8d2ff5 100644 --- a/src/search/aggregations/params/mod.rs +++ b/src/search/aggregations/params/mod.rs @@ -3,9 +3,13 @@ mod aggregation_name; mod gap_policy; mod rate_mode; +mod terms_exclude; +mod terms_include; mod terms_order; pub use self::aggregation_name::*; pub use self::gap_policy::*; pub use self::rate_mode::*; +pub use self::terms_exclude::*; +pub use self::terms_include::*; pub use self::terms_order::*; diff --git a/src/search/aggregations/params/terms_exclude.rs b/src/search/aggregations/params/terms_exclude.rs new file mode 100644 index 0000000..d29f540 --- /dev/null +++ b/src/search/aggregations/params/terms_exclude.rs @@ -0,0 +1,61 @@ +use crate::util::ShouldSkip; + +/// Filter the values for which buckets will be created. +#[derive(Debug, Clone, Serialize, PartialEq)] +#[serde(untagged)] +pub enum TermsExclude { + /// Filter buckets by their regular expression pattern. + /// + /// + Regex(String), + + /// Filter buckets by their exact value. + /// + /// + Exact(Vec), +} + +impl ShouldSkip for TermsExclude { + fn should_skip(&self) -> bool { + match self { + Self::Regex(ref s) => s.is_empty(), + Self::Exact(ref v) => v.is_empty(), + } + } +} + +impl From for TermsExclude { + fn from(s: String) -> Self { + Self::Regex(s) + } +} + +impl From<&str> for TermsExclude { + fn from(s: &str) -> Self { + Self::Regex(s.to_string()) + } +} + +impl From> for TermsExclude { + fn from(v: Vec) -> Self { + Self::Exact(v) + } +} + +impl From> for TermsExclude { + fn from(v: Vec<&str>) -> Self { + Self::Exact(v.iter().map(|s| s.to_string()).collect()) + } +} + +impl From<&[&str]> for TermsExclude { + fn from(v: &[&str]) -> Self { + Self::Exact(v.iter().map(|s| s.to_string()).collect()) + } +} + +impl From<[&str; N]> for TermsExclude { + fn from(value: [&str; N]) -> Self { + Self::Exact(value.iter().map(|s| s.to_string()).collect()) + } +} diff --git a/src/search/aggregations/params/terms_include.rs b/src/search/aggregations/params/terms_include.rs new file mode 100644 index 0000000..4473809 --- /dev/null +++ b/src/search/aggregations/params/terms_include.rs @@ -0,0 +1,90 @@ +use crate::util::ShouldSkip; + +/// Filter the values for which buckets will be created. +#[derive(Debug, Clone, Serialize, PartialEq)] +#[serde(untagged)] +pub enum TermsInclude { + /// Filter buckets by their regular expression pattern. + /// + /// + Regex(String), + + /// Filter buckets by their exact value. + /// + /// + Exact(Vec), + + /// A number of partitions at query-time and processing only one partition in each request. + /// + /// + Partitions { + /// The partition number to return. + partition: u32, + /// The number of partitions to create. + num_partitions: u32, + }, +} + +impl ShouldSkip for TermsInclude { + fn should_skip(&self) -> bool { + match self { + Self::Regex(ref s) => s.is_empty(), + Self::Exact(ref v) => v.is_empty(), + Self::Partitions { .. } => false, + } + } +} + +impl From for TermsInclude { + fn from(s: String) -> Self { + Self::Regex(s) + } +} + +impl From<&str> for TermsInclude { + fn from(s: &str) -> Self { + Self::Regex(s.to_string()) + } +} + +impl From> for TermsInclude { + fn from(v: Vec) -> Self { + Self::Exact(v) + } +} + +impl From> for TermsInclude { + fn from(v: Vec<&str>) -> Self { + Self::Exact(v.iter().map(|s| s.to_string()).collect()) + } +} + +impl From<&[&str]> for TermsInclude { + fn from(v: &[&str]) -> Self { + Self::Exact(v.iter().map(|s| s.to_string()).collect()) + } +} + +impl From<[&str; N]> for TermsInclude { + fn from(value: [&str; N]) -> Self { + Self::Exact(value.iter().map(|s| s.to_string()).collect()) + } +} + +impl From<(u32, u32)> for TermsInclude { + fn from(value: (u32, u32)) -> Self { + Self::Partitions { + partition: value.0, + num_partitions: value.1, + } + } +} + +impl From<[u32; 2]> for TermsInclude { + fn from(value: [u32; 2]) -> Self { + Self::Partitions { + partition: value[0], + num_partitions: value[1], + } + } +}