From 01d2afaa4379dcdc688c33ef6c1be253d54aca46 Mon Sep 17 00:00:00 2001 From: Zhenchi Date: Wed, 3 Jan 2024 06:13:40 +0000 Subject: [PATCH] feat(inverted_index): get memory usage of appliers Signed-off-by: Zhenchi --- .../src/inverted_index/search/fst_apply.rs | 3 ++ .../search/fst_apply/intersection_apply.rs | 54 +++++++++++++++++++ .../search/fst_apply/keys_apply.rs | 17 ++++++ .../src/inverted_index/search/index_apply.rs | 3 ++ .../search/index_apply/predicates_apply.rs | 27 ++++++++++ 5 files changed, 104 insertions(+) diff --git a/src/index/src/inverted_index/search/fst_apply.rs b/src/index/src/inverted_index/search/fst_apply.rs index 9f54d0d88918..9d09eccda3f7 100644 --- a/src/index/src/inverted_index/search/fst_apply.rs +++ b/src/index/src/inverted_index/search/fst_apply.rs @@ -30,4 +30,7 @@ pub trait FstApplier: Send + Sync { /// /// Returns a `Vec`, with each u64 being a value from the FstMap. fn apply(&self, fst: &FstMap) -> Vec; + + /// Returns the memory usage of the applier. + fn memory_usage(&self) -> usize; } diff --git a/src/index/src/inverted_index/search/fst_apply/intersection_apply.rs b/src/index/src/inverted_index/search/fst_apply/intersection_apply.rs index a608acd0bab5..d76b44fe9d4b 100644 --- a/src/index/src/inverted_index/search/fst_apply/intersection_apply.rs +++ b/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::(); + 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::>>(); + 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::() + 4 + size_of::>>() + 320 + ); + } } diff --git a/src/index/src/inverted_index/search/fst_apply/keys_apply.rs b/src/index/src/inverted_index/search/fst_apply/keys_apply.rs index 4ec5710a3435..118ba1edba05 100644 --- a/src/index/src/inverted_index/search/fst_apply/keys_apply.rs +++ b/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 { self.keys.iter().filter_map(|k| fst.get(k)).collect() } + + fn memory_usage(&self) -> usize { + self.keys.capacity() * size_of::() + + self.keys.iter().map(|k| k.capacity()).sum::() + } } 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::() + 6); + } } diff --git a/src/index/src/inverted_index/search/index_apply.rs b/src/index/src/inverted_index/search/index_apply.rs index 35d8c387a2d6..e87472bdeb0a 100644 --- a/src/index/src/inverted_index/search/index_apply.rs +++ b/src/index/src/inverted_index/search/index_apply.rs @@ -33,6 +33,9 @@ pub trait IndexApplier { context: SearchContext, reader: &mut dyn InvertedIndexReader, ) -> Result>; + + /// Returns the memory usage of the applier. + fn memory_usage(&self) -> usize; } /// A context for searching the inverted index. diff --git a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs index e2bea2756a7f..addf1a07446d 100644 --- a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs +++ b/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)>(); + 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)>() + 5 + 100 + ); + } }