Improve mem consumption estimate

Author: rdettai-sk

src/aggregation/bucket/composite/map.rs

@@ -49,6 +49,13 @@ impl<K: Eq + Hash + Clone + Ord, V, const S: usize> ArrayHeapMap<K, V, S> {
             self.buckets.remove(&highest);
         }
     }
+
+    fn memory_consumption(&self) -> u64 {
+        let key_size = std::mem::size_of::<[K; S]>();
+        let map_size = (key_size + std::mem::size_of::<V>()) * self.buckets.capacity();
+        let heap_size = key_size * self.heap.capacity();
+        (map_size + heap_size) as u64
+    }
 }
 
 impl<K: Copy + Ord + Clone + 'static, V: 'static, const S: usize> ArrayHeapMap<K, V, S> {
@@ -249,6 +256,27 @@ impl<K: Ord + Hash + Clone, V> DynArrayHeapMap<K, V> {
             DynArrayHeapMapInner::Dim16(map) => map.evict_highest(),
         }
     }
+
+    pub(crate) fn memory_consumption(&self) -> u64 {
+        match &self.0 {
+            DynArrayHeapMapInner::Dim1(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim2(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim3(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim4(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim5(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim6(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim7(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim8(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim9(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim10(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim11(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim12(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim13(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim14(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim15(map) => map.memory_consumption(),
+            DynArrayHeapMapInner::Dim16(map) => map.memory_consumption(),
+        }
+    }
 }
 
 impl<K: Ord + Clone + Copy + 'static, V: 'static> DynArrayHeapMap<K, V> {

src/aggregation/bucket/composite/mod.rs

@@ -341,8 +341,9 @@ impl SegmentAggregationCollector for SegmentCompositeCollector {
 
 impl SegmentCompositeCollector {
     fn get_memory_consumption(&self) -> u64 {
-        // TODO get correct estimate (these are just the keys)
-        (self.buckets.size() * (std::mem::size_of::<InternalValueRepr>())) as u64
+        // TODO: the footprint is underestimated because we don't account for the
+        // sub-aggregations which are trait objects
+        self.buckets.memory_consumption()
     }
 
     pub(crate) fn from_req_and_validate(
@@ -601,6 +602,7 @@ fn recursive_key_visitor(
     let sub_level_sources = &sources[1..];
     let mut missing = true;
     for (i, accessor) in current_level_accessor.iter().enumerate() {
+        // TODO: optimize with prefetching using fetch_block
         let values = accessor.column.values_for_doc(doc_id);
         match current_level_source {
             CompositeAggregationSource::Terms(term_source) => {