feat(inverted_index): get memory usage of appliers

Signed-off-by: Zhenchi <[email protected]>

src/index/src/inverted_index/search/fst_apply.rs

@@ -30,4 +30,7 @@ pub trait FstApplier: Send + Sync {
     ///
     /// Returns a `Vec<u64>`, with each u64 being a value from the FstMap.
     fn apply(&self, fst: &FstMap) -> Vec<u64>;
+
+    /// Returns the memory usage of the applier.
+    fn memory_usage(&self) -> usize;
 }

src/index/src/inverted_index/search/fst_apply/intersection_apply.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::mem::size_of;
+
 use fst::map::OpBuilder;
 use fst::{IntoStreamer, Streamer};
 use regex_automata::dfa::dense::DFA;
@@ -68,6 +70,26 @@ impl FstApplier for IntersectionFstApplier {
         }
         values
     }
+
+    fn memory_usage(&self) -> usize {
+        let mut size = self.ranges.capacity() * size_of::<Range>();
+        for range in &self.ranges {
+            size += range
+                .lower
+                .as_ref()
+                .map_or(0, |bound| bound.value.capacity());
+            size += range
+                .upper
+                .as_ref()
+                .map_or(0, |bound| bound.value.capacity());
+        }
+
+        size += self.dfas.capacity() * size_of::<DFA<Vec<u32>>>();
+        for dfa in &self.dfas {
+            size += dfa.memory_usage();
+        }
+        size
+    }
 }
 
 impl IntersectionFstApplier {
@@ -340,4 +362,36 @@ mod tests {
             Err(Error::IntersectionApplierWithInList { .. })
         ));
     }
+
+    #[test]
+    fn test_intersection_fst_applier_memory_usage() {
+        let applier = IntersectionFstApplier {
+            ranges: vec![],
+            dfas: vec![],
+        };
+
+        assert_eq!(applier.memory_usage(), 0);
+
+        let dfa = DFA::new("^abc$").unwrap();
+        assert_eq!(dfa.memory_usage(), 320);
+
+        let applier = IntersectionFstApplier {
+            ranges: vec![Range {
+                lower: Some(Bound {
+                    value: b"aa".to_vec(),
+                    inclusive: true,
+                }),
+                upper: Some(Bound {
+                    value: b"cc".to_vec(),
+                    inclusive: true,
+                }),
+            }],
+            dfas: vec![dfa],
+        };
+
+        assert_eq!(
+            applier.memory_usage(),
+            size_of::<Range>() + 4 + size_of::<DFA<Vec<u32>>>() + 320
+        );
+    }
 }

src/index/src/inverted_index/search/fst_apply/keys_apply.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::collections::HashSet;
+use std::mem::size_of;
 
 use snafu::{ensure, ResultExt};
 
@@ -35,6 +36,11 @@ impl FstApplier for KeysFstApplier {
     fn apply(&self, fst: &FstMap) -> Vec<u64> {
         self.keys.iter().filter_map(|k| fst.get(k)).collect()
     }
+
+    fn memory_usage(&self) -> usize {
+        self.keys.capacity() * size_of::<Bytes>()
+            + self.keys.iter().map(|k| k.capacity()).sum::<usize>()
+    }
 }
 
 impl KeysFstApplier {
@@ -302,4 +308,15 @@ mod tests {
         let result = KeysFstApplier::try_from(predicates);
         assert!(matches!(result, Err(Error::ParseRegex { .. })));
     }
+
+    #[test]
+    fn test_keys_fst_applier_memory_usage() {
+        let applier = KeysFstApplier { keys: vec![] };
+        assert_eq!(applier.memory_usage(), 0);
+
+        let applier = KeysFstApplier {
+            keys: vec![b("foo"), b("bar")],
+        };
+        assert_eq!(applier.memory_usage(), 2 * size_of::<Bytes>() + 6);
+    }
 }

src/index/src/inverted_index/search/index_apply.rs

@@ -33,6 +33,9 @@ pub trait IndexApplier {
         context: SearchContext,
         reader: &mut dyn InvertedIndexReader,
     ) -> Result<Vec<usize>>;
+
+    /// Returns the memory usage of the applier.
+    fn memory_usage(&self) -> usize;
 }
 
 /// A context for searching the inverted index.

src/index/src/inverted_index/search/index_apply/predicates_apply.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::mem::size_of;
+
 use async_trait::async_trait;
 use common_base::BitVec;
 use greptime_proto::v1::index::InvertedIndexMetas;
@@ -80,6 +82,16 @@ impl IndexApplier for PredicatesIndexApplier {
 
         Ok(bitmap.iter_ones().collect())
     }
+
+    /// Returns the memory usage of the applier.
+    fn memory_usage(&self) -> usize {
+        let mut size = self.fst_appliers.capacity() * size_of::<(IndexName, Box<dyn FstApplier>)>();
+        for (name, fst_applier) in &self.fst_appliers {
+            size += name.capacity();
+            size += fst_applier.memory_usage();
+        }
+        size
+    }
 }
 
 impl PredicatesIndexApplier {
@@ -343,4 +355,19 @@ mod tests {
             .unwrap();
         assert_eq!(indices, vec![0, 1, 2, 3, 4, 5, 6, 7]);
     }
+
+    #[test]
+    fn test_index_applier_memory_usage() {
+        let mut mock_fst_applier = MockFstApplier::new();
+        mock_fst_applier.expect_memory_usage().returning(|| 100);
+
+        let applier = PredicatesIndexApplier {
+            fst_appliers: vec![(s("tag-0"), Box::new(mock_fst_applier))],
+        };
+
+        assert_eq!(
+            applier.memory_usage(),
+            size_of::<(IndexName, Box<dyn FstApplier>)>() + 5 + 100
+        );
+    }
 }