Elasticsearch Aggregations

Elasticsearch Aggregations.rs

#[macro_use] extern crate serde_derive;
#[macro_use] extern crate serde_json;
extern crate serde;

pub mod aggs {
    use std::slice::Iter;
    use std::borrow::Cow;
    use std::collections::BTreeMap;
    use serde::{Deserialize, Deserializer};
    use serde_json::{Value, Map};

    #[derive(Clone)]
    pub struct Aggregation(AggregationValue);

    #[derive(Clone)]
    enum AggregationValue {
        Anonymous(Value),
        Bucket(Bucket),
        Metric(Metric),
    }

    #[derive(Clone)]
    enum Metric {
        Avg(AvgAggregation),
        Min(MinAggregation),
        Max(MaxAggregation),
    }

    #[derive(Clone, Deserialize)]
    pub struct Bucket {
        buckets: Vec<BucketItem>,
    }

    #[derive(Clone, Deserialize)]
    pub struct BucketItem {
        key: String,
        doc_count: i32,
        #[serde(default)]
        sub_aggs: BTreeMap<String, Aggregation>,
    }

    pub struct BucketIter<'a> {
        inner: Iter<'a, BucketItem>
    }

    impl<'a> Iterator for BucketIter<'a> {
        type Item = &'a BucketItem;

        fn next(&mut self) -> Option<Self::Item> {
            self.inner.next()
        }
    }

    impl Bucket {
        pub fn iter<'a>(&'a self) -> BucketIter<'a> {
            BucketIter {
                inner: self.buckets.iter()
            }
        }
    }

    impl Aggregation {
        pub fn as_bucket<'a>(&'a self) -> Option<Cow<'a, Bucket>> {
            match self.0 {
                AggregationValue::Anonymous(ref agg) => Bucket::deserialize(agg).map(|agg| Cow::Owned(agg)).ok(),
                AggregationValue::Bucket(ref agg) => Some(Cow::Borrowed(agg)),
                _ => None,
            }
        }
    }
    
    impl<'de> Deserialize<'de> for Aggregation {
        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where D: Deserializer<'de>
        {
            let anonymous = Value::deserialize(deserializer)?;

            Ok(Aggregation(AggregationValue::Anonymous(anonymous)))
        }
    }

    macro_rules! simple_metric {
        ($metric_struct:ident, $metric_path:path, $as_conv:ident) => (
            #[derive(Clone, Deserialize)]
            pub struct $metric_struct {
                value: f32,
            }

            impl $metric_struct {
                pub fn value(&self) -> f32 {
                    self.value
                }
            }

            impl Aggregation {
                pub fn $as_conv<'a>(&'a self) -> Option<Cow<'a, $metric_struct>> {
                    match self.0 {
                        AggregationValue::Anonymous(ref agg) => $metric_struct::deserialize(agg).map(|agg| Cow::Owned(agg)).ok(),
                        AggregationValue::Metric($metric_path(ref agg)) => Some(Cow::Borrowed(agg)),
                        _ => None,
                    }
                }
            }
        )
    }

    simple_metric!(AvgAggregation, Metric::Avg, as_avg);
    simple_metric!(MinAggregation, Metric::Min, as_min);
    simple_metric!(MaxAggregation, Metric::Max, as_max);
}

#[cfg(test)]
mod tests {
    use serde_json;
    use aggs::Aggregation;

    #[test]
    fn value_metric_agg() {
        let json = r#"{
            "value": 1.0 
        }"#;

        let agg: Aggregation = serde_json::from_str(json).unwrap();

        let avg = agg.as_avg().unwrap();

        assert_eq!(1.0, avg.value());
    }

    #[test]
    fn value_metric_agg_as_bucket_agg() {
        let json = r#"{
            "value": 1.0 
        }"#;

        let agg: Aggregation = serde_json::from_str(json).unwrap();

        assert!(agg.as_bucket().is_none());
    }

    #[test]
    fn bucket_agg() {
        let json = r#"{
            "buckets": [
                {
                    "key": "1",
                    "doc_count": 2863
                },
                {
                    "key": "2",
                    "doc_count": 9823
                }
            ]
        }"#;

        let agg: Aggregation = serde_json::from_str(json).unwrap();

        let buckets = agg.as_bucket().unwrap().iter();

        let fst = buckets.next().unwrap();

        assert_eq!("1", fst.key());
        assert_eq!(2863, fst.doc_count());
    }
}

The idea is:

• Keep the API for anonymous vs strongly-typed aggregations the same. You need to 'cast' the aggregation to the right type to inspect it
• Strong-typing is an implementation detail that can make parsing aggregations faster

Some things to consider:

• If we attempt to parse an aggregation using typed_keys and it fails, this will behave differently than if we hadn't used typed_keys at all
• The aggregations will mostly mirror what was in the request. This API has you walking the aggregations, converting into the right type each time. That should be straightfoward to do.