Add composite aggregation

Uses the latest composite PR in Tantivy

Author: rdettai-sk

quickwit/Cargo.lock

@@ -346,7 +346,7 @@ quickwit-serve = { path = "quickwit-serve" }
 quickwit-storage = { path = "quickwit-storage" }
 quickwit-telemetry = { path = "quickwit-telemetry" }
 
-tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "dabcaa5", default-features = false, features = [
+tantivy = { git = "https://github.com/SekoiaLab/tantivy/", rev = "bbdf83e", default-features = false, features = [
   "lz4-compression",
   "mmap",
   "quickwit",

quickwit/Cargo.toml

@@ -18,7 +18,8 @@ use tantivy::aggregation::Key as TantivyKey;
 use tantivy::aggregation::agg_result::{
     AggregationResult as TantivyAggregationResult, AggregationResults as TantivyAggregationResults,
     BucketEntries as TantivyBucketEntries, BucketEntry as TantivyBucketEntry,
-    BucketResult as TantivyBucketResult, MetricResult as TantivyMetricResult,
+    BucketResult as TantivyBucketResult, CompositeBucketEntry as TantivyCompositeBucketEntry,
+    CompositeKey as TantivyCompositeKey, MetricResult as TantivyMetricResult,
     RangeBucketEntry as TantivyRangeBucketEntry,
 };
 use tantivy::aggregation::metric::{
@@ -169,6 +170,13 @@ pub enum BucketResult {
         /// The upper bound error for the doc count of each term.
         doc_count_error_upper_bound: Option<u64>,
     },
+    /// This is the composite aggregation result
+    Composite {
+        /// The buckets
+        buckets: Vec<CompositeBucketEntry>,
+        /// The key to start after when paginating
+        after_key: FxHashMap<String, CompositeKey>,
+    },
 }
 
 impl From<TantivyBucketResult> for BucketResult {
@@ -189,6 +197,10 @@ impl From<TantivyBucketResult> for BucketResult {
                 sum_other_doc_count,
                 doc_count_error_upper_bound,
             },
+            TantivyBucketResult::Composite { buckets, after_key } => BucketResult::Composite {
+                buckets: buckets.into_iter().map(Into::into).collect(),
+                after_key: after_key.into_iter().map(|(k, v)| (k, v.into())).collect(),
+            },
         }
     }
 }
@@ -211,6 +223,10 @@ impl From<BucketResult> for TantivyBucketResult {
                 sum_other_doc_count,
                 doc_count_error_upper_bound,
             },
+            BucketResult::Composite { buckets, after_key } => TantivyBucketResult::Composite {
+                buckets: buckets.into_iter().map(Into::into).collect(),
+                after_key: after_key.into_iter().map(|(k, v)| (k, v.into())).collect(),
+            },
         }
     }
 }
@@ -410,3 +426,75 @@ impl From<PercentilesMetricResult> for TantivyPercentilesMetricResult {
         TantivyPercentilesMetricResult { values }
     }
 }
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub enum CompositeKey {
+    /// Boolean key
+    Bool(bool),
+    /// String key
+    Str(String),
+    /// `i64` key
+    I64(i64),
+    /// `u64` key
+    U64(u64),
+    /// `f64` key
+    F64(f64),
+    /// Null key
+    Null,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct CompositeBucketEntry {
+    /// The identifier of the bucket.
+    pub key: FxHashMap<String, CompositeKey>,
+    /// Number of documents in the bucket.
+    pub doc_count: u64,
+    /// Sub-aggregations in this bucket.
+    pub sub_aggregation: AggregationResults,
+}
+
+impl From<TantivyCompositeKey> for CompositeKey {
+    fn from(value: TantivyCompositeKey) -> CompositeKey {
+        match value {
+            TantivyCompositeKey::Bool(b) => CompositeKey::Bool(b),
+            TantivyCompositeKey::Str(s) => CompositeKey::Str(s),
+            TantivyCompositeKey::I64(i) => CompositeKey::I64(i),
+            TantivyCompositeKey::U64(u) => CompositeKey::U64(u),
+            TantivyCompositeKey::F64(f) => CompositeKey::F64(f),
+            TantivyCompositeKey::Null => CompositeKey::Null,
+        }
+    }
+}
+
+impl From<CompositeKey> for TantivyCompositeKey {
+    fn from(value: CompositeKey) -> TantivyCompositeKey {
+        match value {
+            CompositeKey::Bool(b) => TantivyCompositeKey::Bool(b),
+            CompositeKey::Str(s) => TantivyCompositeKey::Str(s),
+            CompositeKey::I64(i) => TantivyCompositeKey::I64(i),
+            CompositeKey::U64(u) => TantivyCompositeKey::U64(u),
+            CompositeKey::F64(f) => TantivyCompositeKey::F64(f),
+            CompositeKey::Null => TantivyCompositeKey::Null,
+        }
+    }
+}
+
+impl From<TantivyCompositeBucketEntry> for CompositeBucketEntry {
+    fn from(value: TantivyCompositeBucketEntry) -> CompositeBucketEntry {
+        CompositeBucketEntry {
+            key: value.key.into_iter().map(|(k, v)| (k, v.into())).collect(),
+            doc_count: value.doc_count,
+            sub_aggregation: value.sub_aggregation.into(),
+        }
+    }
+}
+
+impl From<CompositeBucketEntry> for TantivyCompositeBucketEntry {
+    fn from(value: CompositeBucketEntry) -> TantivyCompositeBucketEntry {
+        TantivyCompositeBucketEntry {
+            key: value.key.into_iter().map(|(k, v)| (k, v.into())).collect(),
+            doc_count: value.doc_count,
+            sub_aggregation: value.sub_aggregation.into(),
+        }
+    }
+}

quickwit/quickwit-query/src/aggregations.rs

@@ -375,6 +375,7 @@ expected:
   aggregations:
     response_stats:
       sum_of_squares: 55300.0
+---
 # Test term aggs number precision
 method: [GET]
 engines:
@@ -393,3 +394,86 @@ expected:
       buckets:
       - doc_count: 1 
         key: 1769070189829214200
+--- 
+# Test composite aggregation
+method: [GET]
+engines:
+  - quickwit
+endpoint: _elastic/aggregations/_search
+json:
+  size: 0
+  aggs:
+    host_name_composite:
+      composite:
+        size: 5
+        sources:
+          - host:
+              terms:
+                field: "host"
+                missing_bucket: true
+          - name:
+              terms:
+                field: "name"
+          - response:
+              histogram:
+                field: "response"
+                interval: 50
+expected:
+  aggregations:
+    host_name_composite:
+      buckets:
+        - key: { "host": null, "name": "Bernhard", "response": 100.0 }
+          doc_count: 1
+        - key: { "host": null, "name": "Fritz", "response": 0.0 }
+          doc_count: 2
+        - key: { "host": "192.168.0.1", "name": "Fred", "response": 100.0 }
+          doc_count: 1
+        - key: { "host": "192.168.0.1", "name": "Fritz", "response": 0.0 }
+          doc_count: 1
+        - key: { "host": "192.168.0.10", "name": "Albert", "response": 100.0 }
+          doc_count: 1
+      after_key:
+        host: "192.168.0.10"
+        name: "Albert"
+        response: 100.0
+
+---
+# Test composite aggregation paging
+method: [GET]
+engines:
+  - quickwit
+endpoint: _elastic/aggregations/_search
+json:
+  size: 0
+  aggs:
+    host_name_composite:
+      composite:
+        size: 5
+        sources:
+          - host:
+              terms:
+                field: "host"
+                missing_bucket: true
+          - name:
+              terms:
+                field: "name"
+          - response:
+              histogram:
+                field: "response"
+                interval: 50
+        after:
+          host: "192.168.0.10"
+          name: "Albert"
+          response: 100.0
+expected:
+  aggregations:
+    host_name_composite:
+      buckets:
+        - key: { "host": "192.168.0.10", "name": "Holger", "response": 0.0 }
+          doc_count: 1
+        # Horst is missing because his response field is missing
+        - key: { "host": "192.168.0.10", "name": "Werner", "response": 0.0 }
+          doc_count: 1
+        - key: { "host": "192.168.0.11", "name": "Manfred", "response": 100.0 }
+          doc_count: 1
+---