mito2/cache/index/
result_cache.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeMap;
16use std::sync::Arc;
17
18use index::bloom_filter::applier::InListPredicate;
19use index::inverted_index::search::predicate::{Predicate, RangePredicate};
20use index::target::IndexTarget;
21use moka::notification::RemovalCause;
22use moka::sync::Cache;
23use store_api::storage::{ColumnId, FileId};
24
25use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS};
26use crate::sst::index::fulltext_index::applier::builder::{
27    FulltextLike, FulltextQuery, FulltextRequest, FulltextTerm,
28};
29use crate::sst::parquet::row_selection::RowGroupSelection;
30
31const INDEX_RESULT_TYPE: &str = "index_result";
32
33/// Cache for storing index query results.
34///
35/// The `RowGroupSelection` is a collection of row groups that match the predicate.
36///
37/// Row groups can be partially searched. Row groups that not contained in `RowGroupSelection` are not searched.
38/// User can retrieve the partial results and handle uncontained row groups required by the predicate subsequently.
39pub struct IndexResultCache {
40    cache: Cache<(PredicateKey, FileId), Arc<RowGroupSelection>>,
41}
42
43impl IndexResultCache {
44    /// Creates a new cache with the given capacity.
45    pub fn new(capacity: u64) -> Self {
46        fn to_str(cause: RemovalCause) -> &'static str {
47            match cause {
48                RemovalCause::Expired => "expired",
49                RemovalCause::Explicit => "explicit",
50                RemovalCause::Replaced => "replaced",
51                RemovalCause::Size => "size",
52            }
53        }
54
55        let cache = Cache::builder()
56            .max_capacity(capacity)
57            .weigher(Self::index_result_cache_weight)
58            .eviction_listener(|k, v, cause| {
59                let size = Self::index_result_cache_weight(&k, &v);
60                CACHE_BYTES
61                    .with_label_values(&[INDEX_RESULT_TYPE])
62                    .sub(size.into());
63                CACHE_EVICTION
64                    .with_label_values(&[INDEX_RESULT_TYPE, to_str(cause)])
65                    .inc();
66            })
67            .support_invalidation_closures()
68            .build();
69        Self { cache }
70    }
71
72    /// Puts a query result into the cache.
73    ///
74    /// Allow user to put a partial result (not containing all row groups) into the cache.
75    pub fn put(&self, key: PredicateKey, file_id: FileId, result: Arc<RowGroupSelection>) {
76        let key = (key, file_id);
77        let size = Self::index_result_cache_weight(&key, &result);
78        CACHE_BYTES
79            .with_label_values(&[INDEX_RESULT_TYPE])
80            .add(size.into());
81        self.cache.insert(key, result);
82    }
83
84    /// Gets a query result from the cache.
85    ///
86    /// Note: the returned `RowGroupSelection` only contains the row groups that are searched.
87    ///       Caller should handle the uncontained row groups required by the predicate subsequently.
88    pub fn get(&self, key: &PredicateKey, file_id: FileId) -> Option<Arc<RowGroupSelection>> {
89        let res = self.cache.get(&(key.clone(), file_id));
90        if res.is_some() {
91            CACHE_HIT.with_label_values(&[INDEX_RESULT_TYPE]).inc();
92        } else {
93            CACHE_MISS.with_label_values(&[INDEX_RESULT_TYPE]).inc()
94        }
95        res
96    }
97
98    /// Calculates the memory usage of a cache entry.
99    fn index_result_cache_weight(k: &(PredicateKey, FileId), v: &Arc<RowGroupSelection>) -> u32 {
100        k.0.mem_usage() as u32 + v.mem_usage() as u32
101    }
102
103    /// Removes cached results for the given file.
104    pub fn invalidate_file(&self, file_id: FileId) {
105        self.cache
106            .invalidate_entries_if(move |(_, cached_file_id), _| *cached_file_id == file_id)
107            .expect("cache should support invalidation closures");
108    }
109}
110
111/// Key for different types of index predicates.
112#[derive(Debug, Clone, PartialEq, Eq, Hash)]
113pub enum PredicateKey {
114    /// Fulltext index predicate.
115    Fulltext(FulltextIndexKey),
116    /// Bloom filter predicate.
117    Bloom(BloomFilterKey),
118    /// Inverted index predicate.
119    Inverted(InvertedIndexKey),
120    /// Min-max pruning predicate.
121    MinMax(MinMaxKey),
122}
123
124impl PredicateKey {
125    /// Creates a new fulltext index key.
126    pub fn new_fulltext(predicates: Arc<BTreeMap<ColumnId, FulltextRequest>>) -> Self {
127        Self::Fulltext(FulltextIndexKey::new(predicates))
128    }
129
130    /// Creates a new bloom filter key.
131    pub fn new_bloom(predicates: Arc<BTreeMap<ColumnId, Vec<InListPredicate>>>) -> Self {
132        Self::Bloom(BloomFilterKey::new(predicates))
133    }
134
135    /// Creates a new inverted index key.
136    pub fn new_inverted(predicates: Arc<BTreeMap<IndexTarget, Vec<Predicate>>>) -> Self {
137        Self::Inverted(InvertedIndexKey::new(predicates))
138    }
139
140    /// Creates a new min-max pruning key.
141    pub fn new_minmax(exprs: Arc<Vec<String>>, schema_version: u64, skip_fields: bool) -> Self {
142        Self::MinMax(MinMaxKey::new(exprs, schema_version, skip_fields))
143    }
144
145    /// Returns the memory usage of this key.
146    pub fn mem_usage(&self) -> usize {
147        match self {
148            Self::Fulltext(key) => key.mem_usage,
149            Self::Bloom(key) => key.mem_usage,
150            Self::Inverted(key) => key.mem_usage,
151            Self::MinMax(key) => key.mem_usage,
152        }
153    }
154}
155
156/// Key for fulltext index queries.
157#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
158pub struct FulltextIndexKey {
159    predicates: Arc<BTreeMap<ColumnId, FulltextRequest>>,
160    mem_usage: usize,
161}
162
163impl FulltextIndexKey {
164    /// Creates a new fulltext index key with the given predicates.
165    /// Calculates memory usage based on the size of queries and terms.
166    pub fn new(predicates: Arc<BTreeMap<ColumnId, FulltextRequest>>) -> Self {
167        let mem_usage = predicates
168            .values()
169            .map(|request| {
170                let query_size = request
171                    .queries
172                    .iter()
173                    .map(|query| query.0.len() + size_of::<FulltextQuery>())
174                    .sum::<usize>();
175                let term_size = request
176                    .terms
177                    .iter()
178                    .map(|term| term.term.len() + size_of::<FulltextTerm>())
179                    .sum::<usize>();
180                let like_size = request
181                    .likes
182                    .iter()
183                    .map(|like| like.pattern.len() + size_of::<FulltextLike>())
184                    .sum::<usize>();
185                query_size + term_size + like_size
186            })
187            .sum();
188        Self {
189            predicates,
190            mem_usage,
191        }
192    }
193}
194
195/// Key for bloom filter queries.
196#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
197pub struct BloomFilterKey {
198    predicates: Arc<BTreeMap<ColumnId, Vec<InListPredicate>>>,
199    mem_usage: usize,
200}
201
202impl BloomFilterKey {
203    /// Creates a new bloom filter key with the given predicates.
204    /// Calculates memory usage based on the size of predicate lists.
205    pub fn new(predicates: Arc<BTreeMap<ColumnId, Vec<InListPredicate>>>) -> Self {
206        let mem_usage = predicates
207            .values()
208            .map(|predicates| {
209                predicates
210                    .iter()
211                    .map(|predicate| predicate.list.iter().map(|list| list.len()).sum::<usize>())
212                    .sum::<usize>()
213            })
214            .sum();
215        Self {
216            predicates,
217            mem_usage,
218        }
219    }
220}
221
222/// Key for inverted index queries.
223#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
224pub struct InvertedIndexKey {
225    predicates: Arc<BTreeMap<IndexTarget, Vec<Predicate>>>,
226    mem_usage: usize,
227}
228
229impl InvertedIndexKey {
230    /// Creates a new inverted index key with the given predicates.
231    /// Calculates memory usage based on the type and size of predicates.
232    pub fn new(predicates: Arc<BTreeMap<IndexTarget, Vec<Predicate>>>) -> Self {
233        let mem_usage = predicates
234            .iter()
235            .map(|(target, predicates)| {
236                let target_size = size_of::<IndexTarget>()
237                    + match target {
238                        IndexTarget::ColumnId(_) => 0,
239                        IndexTarget::JsonPath(target) => target
240                            .path()
241                            .iter()
242                            .map(|part| size_of::<String>() + part.len())
243                            .sum(),
244                    };
245                target_size
246                    + predicates
247                        .iter()
248                        .map(|predicate| match predicate {
249                            Predicate::InList(predicate) => {
250                                predicate.list.iter().map(|list| list.len()).sum::<usize>()
251                            }
252                            Predicate::Range(_) => size_of::<RangePredicate>(),
253                            Predicate::RegexMatch(predicate) => predicate.pattern.len(),
254                        })
255                        .sum::<usize>()
256            })
257            .sum();
258
259        Self {
260            predicates,
261            mem_usage,
262        }
263    }
264}
265
266/// Key for min-max pruning.
267#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
268pub struct MinMaxKey {
269    exprs: Arc<Vec<String>>,
270    schema_version: u64,
271    skip_fields: bool,
272    mem_usage: usize,
273}
274
275impl MinMaxKey {
276    pub fn new(exprs: Arc<Vec<String>>, schema_version: u64, skip_fields: bool) -> Self {
277        let mem_usage = size_of::<Self>()
278            + size_of::<Vec<String>>()
279            + exprs.len() * size_of::<String>()
280            + exprs.iter().map(|s| s.len()).sum::<usize>();
281        Self {
282            exprs,
283            schema_version,
284            skip_fields,
285            mem_usage,
286        }
287    }
288}
289
290#[cfg(test)]
291#[allow(clippy::single_range_in_vec_init)]
292mod tests {
293    use std::collections::{BTreeMap, BTreeSet};
294    use std::sync::Arc;
295
296    use index::bloom_filter::applier::InListPredicate as BloomInListPredicate;
297    use index::inverted_index::search::predicate::{Predicate, Range, RangePredicate};
298
299    use super::*;
300    use crate::sst::index::fulltext_index::applier::builder::{
301        FulltextQuery, FulltextRequest, FulltextTerm,
302    };
303    use crate::sst::parquet::row_selection::RowGroupSelection;
304
305    #[test]
306    fn test_inverted_cache_separates_json_paths() {
307        let key = |path: &str| {
308            PredicateKey::new_inverted(Arc::new(BTreeMap::from([(
309                IndexTarget::new_json_path(1, vec![path.into()]).unwrap(),
310                vec![Predicate::Range(RangePredicate {
311                    range: Range {
312                        lower: None,
313                        upper: None,
314                    },
315                })],
316            )])))
317        };
318        let a = key("a");
319        let b = key("b");
320        let cache = IndexResultCache::new(10000);
321        let file = FileId::random();
322        cache.put(a.clone(), file, Arc::new(RowGroupSelection::new(10, 10)));
323        assert!(cache.get(&a, file).is_some());
324        assert!(cache.get(&b, file).is_none());
325    }
326
327    #[test]
328    fn test_cache_basic_operations() {
329        let cache = IndexResultCache::new(1000);
330        let file_id = FileId::random();
331
332        // Create a test key and value
333        let predicates = BTreeMap::new();
334        let key = PredicateKey::new_fulltext(Arc::new(predicates));
335        let selection = Arc::new(RowGroupSelection::from_row_ids(
336            [1, 2, 3].into_iter().collect(),
337            1,
338            10,
339        ));
340
341        // Test put and get
342        cache.put(key.clone(), file_id, selection.clone());
343        let retrieved = cache.get(&key, file_id);
344        assert!(retrieved.is_some());
345        assert_eq!(
346            retrieved.unwrap().as_ref().row_count(),
347            selection.as_ref().row_count()
348        );
349
350        // Test get non-existent key
351        let non_existent_file_id = FileId::random();
352        assert!(cache.get(&key, non_existent_file_id).is_none());
353    }
354
355    #[test]
356    fn test_minmax_key_should_distinguish_schema_version_and_skip_fields() {
357        let exprs = Arc::new(vec!["col > 1".to_string()]);
358
359        let key1 = PredicateKey::new_minmax(exprs.clone(), 1, false);
360        let key2 = PredicateKey::new_minmax(exprs.clone(), 2, false);
361        assert_ne!(key1, key2);
362
363        let key3 = PredicateKey::new_minmax(exprs, 1, true);
364        assert_ne!(key1, key3);
365    }
366
367    #[test]
368    fn test_cache_capacity_limit() {
369        // Create a cache with small capacity (100 bytes)
370        let cache = IndexResultCache::new(100);
371        let file_id1 = FileId::random();
372        let file_id2 = FileId::random();
373
374        // Create two large keys that will exceed capacity
375        let mut predicates1 = BTreeMap::new();
376        let request1 = FulltextRequest {
377            queries: vec![
378                FulltextQuery(
379                    "test query 1 with a very long string to ensure large weight".to_string(),
380                ),
381                FulltextQuery("another long query string".to_string()),
382            ],
383            terms: vec![],
384            likes: vec![],
385        };
386        predicates1.insert(1, request1);
387        let key1 = PredicateKey::new_fulltext(Arc::new(predicates1));
388        let selection1 = Arc::new(RowGroupSelection::default());
389
390        let mut predicates2 = BTreeMap::new();
391        let request2 = FulltextRequest {
392            queries: vec![
393                FulltextQuery(
394                    "test query 2 with a very long string to ensure large weight".to_string(),
395                ),
396                FulltextQuery("another long query string".to_string()),
397            ],
398            terms: vec![],
399            likes: vec![],
400        };
401        predicates2.insert(1, request2);
402        let key2 = PredicateKey::new_fulltext(Arc::new(predicates2));
403        let selection2 = Arc::new(RowGroupSelection::default());
404
405        // Calculate weights
406        let weight1 =
407            IndexResultCache::index_result_cache_weight(&(key1.clone(), file_id1), &selection1);
408        let weight2 =
409            IndexResultCache::index_result_cache_weight(&(key2.clone(), file_id2), &selection2);
410        assert!(weight1 > 100);
411        assert!(weight2 > 100);
412
413        // Put first key-value pair
414        cache.put(key1.clone(), file_id1, selection1.clone());
415
416        // Verify first key is in cache
417        let retrieved1 = cache.get(&key1, file_id1);
418        assert!(retrieved1.is_some());
419        assert_eq!(
420            retrieved1.unwrap().as_ref().row_count(),
421            selection1.as_ref().row_count()
422        );
423
424        // Put second key-value pair, which should trigger eviction
425        cache.put(key2.clone(), file_id2, selection2.clone());
426
427        // Verify second key is in cache
428        let retrieved2 = cache.get(&key2, file_id2);
429        assert!(retrieved2.is_some());
430        assert_eq!(
431            retrieved2.unwrap().as_ref().row_count(),
432            selection2.as_ref().row_count()
433        );
434
435        // Verify first key was evicted
436        cache.cache.run_pending_tasks();
437        let retrieved1_after_eviction = cache.get(&key1, file_id1);
438        assert!(
439            retrieved1_after_eviction.is_none(),
440            "First key should have been evicted"
441        );
442    }
443
444    #[test]
445    fn test_index_result_cache_weight() {
446        let file_id = FileId::random();
447
448        // Test empty values
449        let empty_predicates = BTreeMap::new();
450        let empty_key = PredicateKey::new_fulltext(Arc::new(empty_predicates));
451        let empty_selection = Arc::new(RowGroupSelection::default());
452        let empty_weight = IndexResultCache::index_result_cache_weight(
453            &(empty_key.clone(), file_id),
454            &empty_selection,
455        );
456        assert_eq!(empty_weight, 0);
457        assert_eq!(
458            empty_weight,
459            empty_key.mem_usage() as u32 + empty_selection.mem_usage() as u32
460        );
461
462        // Test 1: FulltextIndexKey
463        let mut predicates1 = BTreeMap::new();
464        let request1 = FulltextRequest {
465            queries: vec![FulltextQuery("test query".to_string())],
466            terms: vec![FulltextTerm {
467                col_lowered: false,
468                term: "test term".to_string(),
469            }],
470            likes: vec![],
471        };
472        predicates1.insert(1, request1);
473        let key1 = PredicateKey::new_fulltext(Arc::new(predicates1));
474        let selection1 = Arc::new(RowGroupSelection::new(100, 250));
475        let weight1 =
476            IndexResultCache::index_result_cache_weight(&(key1.clone(), file_id), &selection1);
477        assert!(weight1 > 0);
478        assert_eq!(
479            weight1,
480            key1.mem_usage() as u32 + selection1.mem_usage() as u32
481        );
482
483        // Test 2: BloomFilterKey
484        let mut predicates2 = BTreeMap::new();
485        let predicate2 = BloomInListPredicate {
486            list: BTreeSet::from([b"test1".to_vec(), b"test2".to_vec()]),
487        };
488        predicates2.insert(1, vec![predicate2]);
489        let key2 = PredicateKey::new_bloom(Arc::new(predicates2));
490        let selection2 = Arc::new(RowGroupSelection::from_row_ids(
491            [1, 2, 3].into_iter().collect(),
492            100,
493            1,
494        ));
495        let weight2 =
496            IndexResultCache::index_result_cache_weight(&(key2.clone(), file_id), &selection2);
497        assert!(weight2 > 0);
498        assert_eq!(
499            weight2,
500            key2.mem_usage() as u32 + selection2.mem_usage() as u32
501        );
502
503        // Test 3: InvertedIndexKey
504        let mut predicates3 = BTreeMap::new();
505        let predicate3 = Predicate::Range(RangePredicate {
506            range: Range {
507                lower: None,
508                upper: None,
509            },
510        });
511        predicates3.insert(IndexTarget::ColumnId(1), vec![predicate3]);
512        let key3 = PredicateKey::new_inverted(Arc::new(predicates3));
513        let selection3 = Arc::new(RowGroupSelection::from_row_ranges(
514            vec![(0, vec![5..15])],
515            20,
516        ));
517        let weight3 =
518            IndexResultCache::index_result_cache_weight(&(key3.clone(), file_id), &selection3);
519        assert!(weight3 > 0);
520        assert_eq!(
521            weight3,
522            key3.mem_usage() as u32 + selection3.mem_usage() as u32
523        );
524    }
525}
mito2/cache/index/result_cache.rs

mito2/cache/index/
result_cache.rs
