1pub(crate) mod cache_size;
18
19pub(crate) mod file_cache;
20pub(crate) mod index;
21pub(crate) mod manifest_cache;
22pub(crate) mod range_cache;
23#[cfg(test)]
24pub(crate) mod test_util;
25pub(crate) mod write_cache;
26
27use std::collections::{BTreeMap, HashMap};
28use std::mem;
29use std::ops::Range;
30use std::sync::{Arc, RwLock, Weak};
31
32use bytes::Bytes;
33use common_base::readable_size::ReadableSize;
34use common_datasource::compression::CompressionType;
35use common_runtime::Runtime;
36use common_runtime::runtime::RuntimeTrait;
37use common_telemetry::warn;
38use datatypes::arrow::buffer::BooleanBuffer;
39use datatypes::arrow::record_batch::RecordBatch;
40use datatypes::types::json_type::JsonNativeType;
41use datatypes::value::Value;
42use datatypes::vectors::VectorRef;
43use index::bloom_filter_index::{BloomFilterIndexCache, BloomFilterIndexCacheRef};
44use index::result_cache::IndexResultCache;
45use moka::notification::RemovalCause;
46use moka::sync::Cache;
47use object_store::ObjectStore;
48use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
49use parquet::file::metadata::{
50 FileMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, ParquetMetaDataWriter,
51};
52use puffin::puffin_manager::cache::{PuffinMetadataCache, PuffinMetadataCacheRef};
53use smallvec::SmallVec;
54use snafu::{OptionExt, ResultExt};
55use store_api::metadata::{RegionMetadata, RegionMetadataRef};
56use store_api::storage::{ColumnId, ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector};
57pub use write_cache::{WriteCacheUploadStoreWrapper, WriteCacheUploadStoreWrapperRef};
58
59use crate::cache::cache_size::parquet_meta_size;
60use crate::cache::file_cache::{FileType, IndexKey};
61use crate::cache::index::inverted_index::{InvertedIndexCache, InvertedIndexCacheRef};
62use crate::cache::write_cache::WriteCacheRef;
63use crate::error::{
64 CompressObjectSnafu, DecompressObjectSnafu, InvalidMetadataSnafu, InvalidParquetSnafu,
65 JoinSnafu, ReadParquetSnafu, Result, UnexpectedSnafu, WriteParquetSnafu,
66};
67use crate::memtable::record_batch_estimated_size;
68use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS};
69use crate::read::Batch;
70use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue};
71use crate::read::read_columns::JsonTargetTypes;
72use crate::sst::file::{RegionFileId, RegionIndexId};
73use crate::sst::parquet::PARQUET_METADATA_KEY;
74use crate::sst::parquet::read_columns::ParquetReadColumns;
75use crate::sst::parquet::reader::MetadataCacheMetrics;
76
77const SST_META_TYPE: &str = "sst_meta";
79const SST_META_DECODED_TYPE: &str = "sst_meta_decoded";
81const VECTOR_TYPE: &str = "vector";
83const PAGE_TYPE: &str = "page";
85const FILE_TYPE: &str = "file";
87const INDEX_TYPE: &str = "index";
89const SELECTOR_RESULT_TYPE: &str = "selector_result";
91const RANGE_RESULT_TYPE: &str = "range_result";
93const PREFILTER_RESULT_TYPE: &str = "prefilter_result";
95const RANGE_RESULT_CONCAT_MEMORY_LIMIT: ReadableSize = ReadableSize::mb(512);
96const RANGE_RESULT_CONCAT_MEMORY_PERMIT: ReadableSize = ReadableSize::kb(1);
97
98#[derive(Debug)]
99pub(crate) struct RangeResultMemoryLimiter {
100 semaphore: Arc<tokio::sync::Semaphore>,
101 permit_bytes: usize,
102 total_permits: usize,
103 #[cfg(test)]
106 waited_acquires: std::sync::atomic::AtomicUsize,
107}
108
109impl Default for RangeResultMemoryLimiter {
110 fn default() -> Self {
111 Self::new(
112 RANGE_RESULT_CONCAT_MEMORY_LIMIT.as_bytes() as usize,
113 RANGE_RESULT_CONCAT_MEMORY_PERMIT.as_bytes() as usize,
114 )
115 }
116}
117
118impl RangeResultMemoryLimiter {
119 pub(crate) fn new(limit_bytes: usize, permit_bytes: usize) -> Self {
120 let permit_bytes = permit_bytes.max(1);
121 let total_permits = limit_bytes
122 .div_ceil(permit_bytes)
123 .clamp(1, tokio::sync::Semaphore::MAX_PERMITS);
124 Self {
125 semaphore: Arc::new(tokio::sync::Semaphore::new(total_permits)),
126 permit_bytes,
127 total_permits,
128 #[cfg(test)]
129 waited_acquires: std::sync::atomic::AtomicUsize::new(0),
130 }
131 }
132
133 #[cfg(test)]
134 pub(crate) fn permit_bytes(&self) -> usize {
135 self.permit_bytes
136 }
137
138 #[cfg(test)]
139 pub(crate) fn waited_acquires(&self) -> usize {
140 self.waited_acquires
141 .load(std::sync::atomic::Ordering::Acquire)
142 }
143
144 #[cfg(test)]
145 pub(crate) fn available_permits(&self) -> usize {
146 self.semaphore.available_permits()
147 }
148
149 pub(crate) fn try_acquire(&self, bytes: usize) -> Option<tokio::sync::OwnedSemaphorePermit> {
152 let permits = bytes.div_ceil(self.permit_bytes).max(1);
153 if permits > self.total_permits {
154 return None;
155 }
156 self.semaphore
157 .clone()
158 .try_acquire_many_owned(permits.try_into().ok()?)
159 .ok()
160 }
161
162 pub(crate) fn try_resize(
166 &self,
167 permit: Option<tokio::sync::OwnedSemaphorePermit>,
168 bytes: usize,
169 ) -> Option<tokio::sync::OwnedSemaphorePermit> {
170 let Some(mut permit) = permit else {
171 return self.try_acquire(bytes);
172 };
173 let required = bytes.div_ceil(self.permit_bytes).max(1);
174 let reserved = permit.num_permits();
175 if required > reserved {
176 permit.merge(self.try_acquire((required - reserved).checked_mul(self.permit_bytes)?)?);
177 } else {
178 drop(permit.split(reserved - required));
179 }
180 Some(permit)
181 }
182
183 pub(crate) async fn acquire(&self, bytes: usize) -> Result<tokio::sync::SemaphorePermit<'_>> {
184 let permits = bytes.div_ceil(self.permit_bytes).max(1);
185 if permits > self.total_permits {
186 return UnexpectedSnafu {
187 reason: format!(
188 "range result memory request of {bytes} bytes exceeds limiter capacity of {} bytes",
189 self.total_permits.saturating_mul(self.permit_bytes)
190 ),
191 }
192 .fail();
193 }
194 #[cfg(test)]
197 if self.semaphore.available_permits() < permits {
198 self.waited_acquires
199 .fetch_add(1, std::sync::atomic::Ordering::Release);
200 }
201
202 self.semaphore
203 .acquire_many(permits as u32)
204 .await
205 .map_err(|_| {
206 UnexpectedSnafu {
207 reason: "range result memory limiter is unexpectedly closed",
208 }
209 .build()
210 })
211 }
212}
213
214#[derive(Debug)]
219pub(crate) struct CachedSstMeta {
220 parquet_metadata: Arc<ParquetMetaData>,
221 parquet_metadata_size: usize,
222 region_metadata: RegionMetadataRef,
223 region_metadata_weight: usize,
224 page_index_policy: PageIndexPolicy,
225}
226
227#[derive(Debug)]
233pub(crate) struct CompactSstMeta {
234 encoded_metadata: Bytes,
235 decoded_size: usize,
236 region_metadata: Weak<RegionMetadata>,
237 page_index_policy: PageIndexPolicy,
238 decoded: tokio::sync::Mutex<Weak<CachedSstMeta>>,
239}
240
241#[derive(Debug)]
243pub(crate) struct PreparedSstMeta {
244 compact: Arc<CompactSstMeta>,
245 decoded: Arc<CachedSstMeta>,
246}
247
248impl PreparedSstMeta {
249 pub(crate) fn decoded(&self) -> Arc<CachedSstMeta> {
250 self.decoded.clone()
251 }
252}
253
254#[derive(Debug)]
256pub(crate) enum SstMetaPreparation {
257 Prepared(PreparedSstMeta),
259 DecodedOnly {
261 decoded: Arc<CachedSstMeta>,
262 encoding_error: crate::error::Error,
263 },
264}
265
266impl SstMetaPreparation {
267 pub(crate) fn decoded(&self) -> Arc<CachedSstMeta> {
268 match self {
269 Self::Prepared(metadata) => metadata.decoded(),
270 Self::DecodedOnly { decoded, .. } => decoded.clone(),
271 }
272 }
273}
274
275impl CompactSstMeta {
276 async fn decode(&self, runtime: &Runtime) -> Result<Arc<CachedSstMeta>> {
277 let mut decoded_guard = self.decoded.lock().await;
278 if let Some(decoded) = decoded_guard.upgrade() {
279 return Ok(decoded);
280 }
281
282 let encoded_metadata = self.encoded_metadata.clone();
283 let decoded_size = self.decoded_size;
284 let region_metadata = self.region_metadata.upgrade();
285 let page_index_policy = self.page_index_policy;
286 let decoded = runtime
287 .spawn_blocking(move || {
288 let bytes = zstd::bulk::decompress(&encoded_metadata, decoded_size).context(
289 DecompressObjectSnafu {
290 compress_type: CompressionType::Zstd,
291 path: "cached SST metadata",
292 },
293 )?;
294 let bytes = Bytes::from(bytes);
295 let mut reader = ParquetMetaDataReader::new()
296 .with_column_index_policy(PageIndexPolicy::Skip)
297 .with_offset_index_policy(page_index_policy);
298 reader.try_parse(&bytes).context(ReadParquetSnafu {
299 path: "cached SST metadata",
300 })?;
301 let metadata = reader.finish().context(ReadParquetSnafu {
302 path: "cached SST metadata",
303 })?;
304 CachedSstMeta::try_new_with_page_index_policy(
305 "cached SST metadata",
306 metadata,
307 region_metadata,
308 page_index_policy,
309 )
310 .map(Arc::new)
311 })
312 .await
313 .context(JoinSnafu)??;
314
315 *decoded_guard = Arc::downgrade(&decoded);
316 Ok(decoded)
317 }
318
319 fn satisfies_page_index_policy(&self, requested: PageIndexPolicy) -> bool {
320 satisfies_page_index_policy(self.page_index_policy, requested)
321 }
322}
323
324pub(crate) async fn decode_sst_meta(
326 file_path: &str,
327 parquet_metadata: ParquetMetaData,
328 region_metadata: Option<RegionMetadataRef>,
329 page_index_policy: PageIndexPolicy,
330 runtime: &Runtime,
331) -> Result<Arc<CachedSstMeta>> {
332 let file_path = file_path.to_string();
333 runtime
334 .spawn_blocking(move || {
335 let parquet_metadata = strip_column_indexes(parquet_metadata);
336 CachedSstMeta::try_new_with_page_index_policy(
337 &file_path,
338 parquet_metadata,
339 region_metadata,
340 page_index_policy,
341 )
342 .map(Arc::new)
343 })
344 .await
345 .context(JoinSnafu)?
346}
347
348pub(crate) async fn prepare_sst_meta(
350 file_path: &str,
351 parquet_metadata: ParquetMetaData,
352 region_metadata: Option<RegionMetadataRef>,
353 page_index_policy: PageIndexPolicy,
354 runtime: &Runtime,
355) -> Result<SstMetaPreparation> {
356 let file_path = file_path.to_string();
357 runtime
358 .spawn_blocking(move || {
359 prepare_sst_meta_sync(
360 &file_path,
361 parquet_metadata,
362 region_metadata,
363 page_index_policy,
364 )
365 })
366 .await
367 .context(JoinSnafu)?
368}
369
370pub(crate) fn prepare_sst_meta_sync(
372 file_path: &str,
373 parquet_metadata: ParquetMetaData,
374 region_metadata: Option<RegionMetadataRef>,
375 page_index_policy: PageIndexPolicy,
376) -> Result<SstMetaPreparation> {
377 let parquet_metadata = strip_column_indexes(parquet_metadata);
378 let cache_encoding = encode_compact_sst_meta(file_path, &parquet_metadata);
379 finish_sst_meta_preparation(
380 file_path,
381 parquet_metadata,
382 region_metadata,
383 page_index_policy,
384 cache_encoding,
385 )
386}
387
388fn strip_column_indexes(parquet_metadata: ParquetMetaData) -> ParquetMetaData {
389 let mut builder = parquet_metadata.into_builder();
392 builder.take_column_index();
393 builder.build()
394}
395
396fn encode_compact_sst_meta(
397 file_path: &str,
398 parquet_metadata: &ParquetMetaData,
399) -> Result<(Bytes, usize)> {
400 let mut encoded = Vec::new();
401 ParquetMetaDataWriter::new(&mut encoded, parquet_metadata)
402 .finish()
403 .context(WriteParquetSnafu)?;
404 let decoded_size = encoded.len();
405 let encoded_metadata = zstd::bulk::compress(&encoded, 3).context(CompressObjectSnafu {
406 compress_type: CompressionType::Zstd,
407 path: file_path,
408 })?;
409
410 Ok((
414 Bytes::from(encoded_metadata.into_boxed_slice()),
415 decoded_size,
416 ))
417}
418
419fn finish_sst_meta_preparation(
420 file_path: &str,
421 parquet_metadata: ParquetMetaData,
422 region_metadata: Option<RegionMetadataRef>,
423 page_index_policy: PageIndexPolicy,
424 cache_encoding: Result<(Bytes, usize)>,
425) -> Result<SstMetaPreparation> {
426 let decoded = Arc::new(CachedSstMeta::try_new_with_page_index_policy(
427 file_path,
428 parquet_metadata,
429 region_metadata,
430 page_index_policy,
431 )?);
432 let (encoded_metadata, decoded_size) = match cache_encoding {
433 Ok(encoded) => encoded,
434 Err(encoding_error) => {
435 return Ok(SstMetaPreparation::DecodedOnly {
436 decoded,
437 encoding_error,
438 });
439 }
440 };
441 let compact = Arc::new(CompactSstMeta {
442 encoded_metadata,
443 decoded_size,
444 region_metadata: Arc::downgrade(&decoded.region_metadata),
445 page_index_policy,
446 decoded: tokio::sync::Mutex::new(Arc::downgrade(&decoded)),
447 });
448
449 Ok(SstMetaPreparation::Prepared(PreparedSstMeta {
450 compact,
451 decoded,
452 }))
453}
454
455impl CachedSstMeta {
456 #[cfg(test)]
457 pub(crate) fn try_new(file_path: &str, parquet_metadata: ParquetMetaData) -> Result<Self> {
458 let page_index_policy = infer_loaded_page_index_policy(&parquet_metadata);
459 Self::try_new_with_page_index_policy(file_path, parquet_metadata, None, page_index_policy)
460 }
461
462 pub(crate) fn try_new_with_region_metadata(
463 file_path: &str,
464 parquet_metadata: ParquetMetaData,
465 region_metadata: Option<RegionMetadataRef>,
466 ) -> Result<Self> {
467 let page_index_policy = infer_loaded_page_index_policy(&parquet_metadata);
468 Self::try_new_with_page_index_policy(
469 file_path,
470 parquet_metadata,
471 region_metadata,
472 page_index_policy,
473 )
474 }
475
476 pub(crate) fn try_new_with_page_index_policy(
477 file_path: &str,
478 parquet_metadata: ParquetMetaData,
479 region_metadata: Option<RegionMetadataRef>,
480 page_index_policy: PageIndexPolicy,
481 ) -> Result<Self> {
482 let file_metadata = parquet_metadata.file_metadata();
483 let key_values = file_metadata
484 .key_value_metadata()
485 .context(InvalidParquetSnafu {
486 file: file_path,
487 reason: "missing key value meta",
488 })?;
489 let meta_value = key_values
490 .iter()
491 .find(|kv| kv.key == PARQUET_METADATA_KEY)
492 .with_context(|| InvalidParquetSnafu {
493 file: file_path,
494 reason: format!("key {} not found", PARQUET_METADATA_KEY),
495 })?;
496 let json = meta_value
497 .value
498 .as_ref()
499 .with_context(|| InvalidParquetSnafu {
500 file: file_path,
501 reason: format!("No value for key {}", PARQUET_METADATA_KEY),
502 })?;
503 let region_metadata = match region_metadata {
504 Some(region_metadata) => region_metadata,
505 None => Arc::new(
506 store_api::metadata::RegionMetadata::from_json(json)
507 .context(InvalidMetadataSnafu)?,
508 ),
509 };
510 let region_metadata_weight = region_metadata.estimated_size().max(json.len());
512 let parquet_metadata = Arc::new(strip_region_metadata_from_parquet(parquet_metadata));
513 let parquet_metadata_size = parquet_meta_size(&parquet_metadata);
514
515 Ok(Self {
516 parquet_metadata,
517 parquet_metadata_size,
518 region_metadata,
519 region_metadata_weight,
520 page_index_policy,
521 })
522 }
523
524 pub(crate) fn parquet_metadata(&self) -> Arc<ParquetMetaData> {
525 self.parquet_metadata.clone()
526 }
527
528 pub(crate) fn parquet_metadata_size(&self) -> usize {
530 self.parquet_metadata_size
531 }
532
533 pub(crate) fn region_metadata(&self) -> RegionMetadataRef {
534 self.region_metadata.clone()
535 }
536
537 fn satisfies_page_index_policy(&self, requested: PageIndexPolicy) -> bool {
538 satisfies_page_index_policy(self.page_index_policy, requested)
539 }
540}
541
542fn satisfies_page_index_policy(cached: PageIndexPolicy, requested: PageIndexPolicy) -> bool {
543 match requested {
544 PageIndexPolicy::Skip => true,
545 PageIndexPolicy::Optional => cached != PageIndexPolicy::Skip,
546 PageIndexPolicy::Required => cached == PageIndexPolicy::Required,
547 }
548}
549
550fn infer_loaded_page_index_policy(parquet_metadata: &ParquetMetaData) -> PageIndexPolicy {
551 if parquet_metadata.offset_index().is_some() {
552 PageIndexPolicy::Optional
553 } else {
554 PageIndexPolicy::Skip
555 }
556}
557
558fn strip_region_metadata_from_parquet(parquet_metadata: ParquetMetaData) -> ParquetMetaData {
559 let file_metadata = parquet_metadata.file_metadata();
560 let filtered_key_values = file_metadata.key_value_metadata().and_then(|key_values| {
561 let filtered = key_values
562 .iter()
563 .filter(|kv| kv.key != PARQUET_METADATA_KEY)
564 .cloned()
565 .collect::<Vec<_>>();
566 (!filtered.is_empty()).then_some(filtered)
567 });
568 let stripped_file_metadata = FileMetaData::new(
569 file_metadata.version(),
570 file_metadata.num_rows(),
571 file_metadata.created_by().map(ToString::to_string),
572 filtered_key_values,
573 file_metadata.schema_descr_ptr(),
574 file_metadata.column_orders().cloned(),
575 );
576
577 let mut builder = parquet_metadata.into_builder();
578 let row_groups = builder.take_row_groups();
579 let offset_index = builder.take_offset_index();
580
581 parquet::file::metadata::ParquetMetaDataBuilder::new(stripped_file_metadata)
582 .set_row_groups(row_groups)
583 .set_offset_index(offset_index)
584 .build()
585}
586
587fn removal_cause_str(cause: RemovalCause) -> &'static str {
588 match cause {
589 RemovalCause::Expired => "expired",
590 RemovalCause::Explicit => "explicit",
591 RemovalCause::Replaced => "replaced",
592 RemovalCause::Size => "size",
593 }
594}
595
596#[derive(Debug, Clone, PartialEq, Eq, Hash)]
597pub(crate) struct PrefilterRowSelector {
598 row_count: usize,
599 skip: bool,
600}
601
602impl From<&RowSelector> for PrefilterRowSelector {
607 fn from(selector: &RowSelector) -> Self {
608 Self {
609 row_count: selector.row_count,
610 skip: selector.skip,
611 }
612 }
613}
614
615#[derive(Debug, Clone, PartialEq, Eq, Hash)]
617pub(crate) struct PrefilterKey {
618 file_id: FileId,
619 row_group_idx: u32,
620 row_selection: Option<Arc<Vec<PrefilterRowSelector>>>,
621 schema_version: u64,
622 filter_exprs: SmallVec<[String; 1]>,
623 mem_usage: usize,
624}
625
626impl PrefilterKey {
627 pub(crate) fn row_selection_snapshot(
628 row_selection: Option<&RowSelection>,
629 ) -> Option<Arc<Vec<PrefilterRowSelector>>> {
630 row_selection.map(|selection| {
631 Arc::new(
632 selection
633 .iter()
634 .map(PrefilterRowSelector::from)
635 .collect::<Vec<_>>(),
636 )
637 })
638 }
639
640 pub(crate) fn new(
641 file_id: FileId,
642 row_group_idx: u32,
643 row_selection: Option<Arc<Vec<PrefilterRowSelector>>>,
644 schema_version: u64,
645 filter_exprs: SmallVec<[String; 1]>,
646 ) -> Self {
647 let row_selection_bytes = row_selection
648 .as_ref()
649 .map(|selection| selection.len() * mem::size_of::<PrefilterRowSelector>())
650 .unwrap_or(0);
651 let spilled_expr_bytes = if filter_exprs.spilled() {
652 filter_exprs.capacity() * mem::size_of::<String>()
653 } else {
654 0
655 };
656 let expr_bytes = filter_exprs.iter().map(|s| s.capacity()).sum::<usize>();
657
658 Self {
659 file_id,
660 row_group_idx,
661 row_selection,
662 schema_version,
663 filter_exprs,
664 mem_usage: mem::size_of::<Self>()
665 + row_selection_bytes
666 + spilled_expr_bytes
667 + expr_bytes,
668 }
669 }
670
671 fn mem_usage(&self) -> usize {
672 self.mem_usage
673 }
674}
675
676type PrefilterResultCache = Cache<PrefilterKey, Arc<BooleanBuffer>>;
677
678fn new_prefilter_result_cache(capacity: u64) -> PrefilterResultCache {
679 Cache::builder()
680 .max_capacity(capacity)
681 .weigher(prefilter_result_cache_weight)
682 .eviction_listener(|k, v, cause| {
683 let size = prefilter_result_cache_weight(&k, &v);
684 CACHE_BYTES
685 .with_label_values(&[PREFILTER_RESULT_TYPE])
686 .sub(size.into());
687 CACHE_EVICTION
688 .with_label_values(&[PREFILTER_RESULT_TYPE, removal_cause_str(cause)])
689 .inc();
690 })
691 .build()
692}
693
694fn prefilter_result_cache_weight(k: &PrefilterKey, v: &Arc<BooleanBuffer>) -> u32 {
695 (k.mem_usage() + mem::size_of::<BooleanBuffer>() + v.values().len()) as u32
696}
697
698#[derive(Clone)]
700pub enum CacheStrategy {
701 EnableAll(CacheManagerRef),
704 Compaction(CacheManagerRef),
709 Disabled,
711}
712
713impl CacheStrategy {
714 pub(crate) fn sst_meta_runtime(&self) -> Runtime {
716 match self {
717 CacheStrategy::Compaction(_) => common_runtime::compact_runtime(),
718 CacheStrategy::EnableAll(_) | CacheStrategy::Disabled => {
719 common_runtime::global_runtime()
720 }
721 }
722 }
723
724 pub(crate) fn sst_meta_cache_enabled(&self) -> bool {
726 match self {
727 CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => {
728 cache_manager.sst_meta_cache_enabled()
729 }
730 CacheStrategy::Disabled => false,
731 }
732 }
733
734 pub(crate) async fn get_sst_meta_data(
736 &self,
737 file_id: RegionFileId,
738 metrics: &mut MetadataCacheMetrics,
739 page_index_policy: PageIndexPolicy,
740 ) -> Option<Arc<CachedSstMeta>> {
741 match self {
742 CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => {
743 cache_manager
744 .get_sst_meta_data(
745 file_id,
746 metrics,
747 page_index_policy,
748 &self.sst_meta_runtime(),
749 )
750 .await
751 }
752 CacheStrategy::Disabled => {
753 metrics.cache_miss += 1;
754 None
755 }
756 }
757 }
758
759 pub(crate) fn get_sst_meta_data_from_mem_cache(
761 &self,
762 file_id: RegionFileId,
763 page_index_policy: PageIndexPolicy,
764 ) -> Option<Arc<CachedSstMeta>> {
765 match self {
766 CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => {
767 cache_manager.get_sst_meta_data_from_mem_cache(file_id, page_index_policy)
768 }
769 CacheStrategy::Disabled => None,
770 }
771 }
772
773 pub fn get_parquet_meta_data_from_mem_cache(
775 &self,
776 file_id: RegionFileId,
777 ) -> Option<Arc<ParquetMetaData>> {
778 self.get_sst_meta_data_from_mem_cache(file_id, PageIndexPolicy::Skip)
779 .map(|metadata| metadata.parquet_metadata())
780 }
781
782 pub(crate) fn put_prepared_sst_meta(
784 &self,
785 file_id: RegionFileId,
786 metadata: PreparedSstMeta,
787 retain_decoded: bool,
788 ) {
789 match self {
790 CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => {
791 cache_manager.put_prepared_sst_meta(file_id, metadata, retain_decoded);
792 }
793 CacheStrategy::Disabled => {}
794 }
795 }
796
797 pub fn put_parquet_meta_data(
799 &self,
800 file_id: RegionFileId,
801 metadata: Arc<ParquetMetaData>,
802 region_metadata: Option<RegionMetadataRef>,
803 ) {
804 match self {
805 CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => {
806 cache_manager.put_parquet_meta_data(file_id, metadata, region_metadata);
807 }
808 CacheStrategy::Disabled => {}
809 }
810 }
811
812 pub(crate) fn get_prefilter_result(&self, key: &PrefilterKey) -> Option<Arc<BooleanBuffer>> {
815 match self {
816 CacheStrategy::EnableAll(cache_manager) => cache_manager.get_prefilter_result(key),
817 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
818 }
819 }
820
821 pub(crate) fn put_prefilter_result(&self, key: PrefilterKey, result: Arc<BooleanBuffer>) {
824 if let CacheStrategy::EnableAll(cache_manager) = self {
825 cache_manager.put_prefilter_result(key, result);
826 }
827 }
828
829 pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) {
831 match self {
832 CacheStrategy::EnableAll(cache_manager) => {
833 cache_manager.remove_parquet_meta_data(file_id);
834 }
835 CacheStrategy::Compaction(cache_manager) => {
836 cache_manager.remove_parquet_meta_data(file_id);
837 }
838 CacheStrategy::Disabled => {}
839 }
840 }
841
842 pub fn get_repeated_vector(
845 &self,
846 data_type: &ConcreteDataType,
847 value: &Value,
848 ) -> Option<VectorRef> {
849 match self {
850 CacheStrategy::EnableAll(cache_manager) => {
851 cache_manager.get_repeated_vector(data_type, value)
852 }
853 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
854 }
855 }
856
857 pub fn put_repeated_vector(&self, value: Value, vector: VectorRef) {
860 if let CacheStrategy::EnableAll(cache_manager) = self {
861 cache_manager.put_repeated_vector(value, vector);
862 }
863 }
864
865 pub fn get_page_ranges(
868 &self,
869 file_id: FileId,
870 row_group_idx: usize,
871 ranges: &[Range<u64>],
872 ) -> Option<PageRangeLookup> {
873 match self {
874 CacheStrategy::EnableAll(cache_manager) => {
875 cache_manager.get_page_ranges(file_id, row_group_idx, ranges)
876 }
877 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
878 }
879 }
880
881 pub fn put_page_ranges(
884 &self,
885 file_id: FileId,
886 row_group_idx: usize,
887 ranges: &[Range<u64>],
888 pages: &[Bytes],
889 ) {
890 if let CacheStrategy::EnableAll(cache_manager) = self {
891 cache_manager.put_page_ranges(file_id, row_group_idx, ranges, pages);
892 }
893 }
894
895 pub async fn evict_puffin_cache(&self, file_id: RegionIndexId) {
897 match self {
898 CacheStrategy::EnableAll(cache_manager) => {
899 cache_manager.evict_puffin_cache(file_id).await
900 }
901 CacheStrategy::Compaction(cache_manager) => {
902 cache_manager.evict_puffin_cache(file_id).await
903 }
904 CacheStrategy::Disabled => {}
905 }
906 }
907
908 pub fn get_selector_result(
911 &self,
912 selector_key: &SelectorResultKey,
913 ) -> Option<Arc<SelectorResultValue>> {
914 match self {
915 CacheStrategy::EnableAll(cache_manager) => {
916 cache_manager.get_selector_result(selector_key)
917 }
918 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
919 }
920 }
921
922 pub fn put_selector_result(
925 &self,
926 selector_key: SelectorResultKey,
927 result: Arc<SelectorResultValue>,
928 ) {
929 if let CacheStrategy::EnableAll(cache_manager) = self {
930 cache_manager.put_selector_result(selector_key, result);
931 }
932 }
933
934 #[cfg(test)]
937 pub(crate) fn get_range_result(
938 &self,
939 key: &RangeScanCacheKey,
940 ) -> Option<Arc<RangeScanCacheValue>> {
941 match self {
942 CacheStrategy::EnableAll(cache_manager) => cache_manager.get_range_result(key),
943 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
944 }
945 }
946
947 #[cfg(any(test, feature = "test"))]
950 pub(crate) fn put_range_result(
951 &self,
952 key: RangeScanCacheKey,
953 result: Arc<RangeScanCacheValue>,
954 ) {
955 if let CacheStrategy::EnableAll(cache_manager) = self {
956 cache_manager.put_range_result(key, result);
957 }
958 }
959
960 pub(crate) fn range_cache_worker(&self) -> Option<&Arc<range_cache::RangeCacheWorker>> {
961 match self {
962 Self::EnableAll(manager) => manager.range_result_cache.as_ref(),
963 Self::Compaction(_) | Self::Disabled => None,
964 }
965 }
966
967 pub(crate) fn has_range_result_cache(&self) -> bool {
969 match self {
970 CacheStrategy::EnableAll(cache_manager) => cache_manager.has_range_result_cache(),
971 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => false,
972 }
973 }
974
975 pub(crate) fn range_result_memory_limiter(&self) -> Option<&Arc<RangeResultMemoryLimiter>> {
976 match self {
977 CacheStrategy::EnableAll(cache_manager) => {
978 Some(cache_manager.range_result_memory_limiter())
979 }
980 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
981 }
982 }
983
984 pub(crate) fn range_result_cache_size(&self) -> Option<usize> {
985 match self {
986 CacheStrategy::EnableAll(cache_manager) => {
987 Some(cache_manager.range_result_cache_size())
988 }
989 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
990 }
991 }
992
993 pub fn write_cache(&self) -> Option<&WriteCacheRef> {
996 match self {
997 CacheStrategy::EnableAll(cache_manager) => cache_manager.write_cache(),
998 CacheStrategy::Compaction(cache_manager) => cache_manager.write_cache(),
999 CacheStrategy::Disabled => None,
1000 }
1001 }
1002
1003 pub fn inverted_index_cache(&self) -> Option<&InvertedIndexCacheRef> {
1006 match self {
1007 CacheStrategy::EnableAll(cache_manager) => cache_manager.inverted_index_cache(),
1008 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
1009 }
1010 }
1011
1012 pub fn bloom_filter_index_cache(&self) -> Option<&BloomFilterIndexCacheRef> {
1015 match self {
1016 CacheStrategy::EnableAll(cache_manager) => cache_manager.bloom_filter_index_cache(),
1017 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
1018 }
1019 }
1020
1021 pub fn puffin_metadata_cache(&self) -> Option<&PuffinMetadataCacheRef> {
1024 match self {
1025 CacheStrategy::EnableAll(cache_manager) => cache_manager.puffin_metadata_cache(),
1026 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
1027 }
1028 }
1029
1030 pub fn index_result_cache(&self) -> Option<&IndexResultCache> {
1033 match self {
1034 CacheStrategy::EnableAll(cache_manager) => cache_manager.index_result_cache(),
1035 CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None,
1036 }
1037 }
1038
1039 pub fn maybe_download_background(
1041 &self,
1042 index_key: IndexKey,
1043 remote_path: String,
1044 remote_store: ObjectStore,
1045 file_size: u64,
1046 ) {
1047 if let CacheStrategy::EnableAll(cache_manager) = self
1048 && let Some(write_cache) = cache_manager.write_cache()
1049 {
1050 write_cache.file_cache().maybe_download_background(
1051 index_key,
1052 remote_path,
1053 remote_store,
1054 file_size,
1055 );
1056 }
1057 }
1058}
1059
1060#[derive(Default)]
1064pub struct CacheManager {
1065 sst_meta_cache: Option<SstMetaCache>,
1067 sst_decoded_meta_cache: Option<SstDecodedMetaCache>,
1069 vector_cache: Option<VectorCache>,
1071 page_cache: Option<Arc<PageRangeCache>>,
1073 write_cache: Option<WriteCacheRef>,
1075 inverted_index_cache: Option<InvertedIndexCacheRef>,
1077 bloom_filter_index_cache: Option<BloomFilterIndexCacheRef>,
1079 puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
1081 selector_result_cache: Option<SelectorResultCache>,
1083 range_result_cache: Option<RangeResultCache>,
1085 range_result_cache_size: u64,
1087 range_result_memory_limiter: Arc<RangeResultMemoryLimiter>,
1089 index_result_cache: Option<IndexResultCache>,
1091 prefilter_result_cache: Option<PrefilterResultCache>,
1093}
1094
1095pub type CacheManagerRef = Arc<CacheManager>;
1096
1097impl CacheManager {
1098 pub fn builder() -> CacheManagerBuilder {
1100 CacheManagerBuilder::default()
1101 }
1102
1103 pub(crate) async fn get_sst_meta_data(
1106 &self,
1107 file_id: RegionFileId,
1108 metrics: &mut MetadataCacheMetrics,
1109 page_index_policy: PageIndexPolicy,
1110 runtime: &Runtime,
1111 ) -> Option<Arc<CachedSstMeta>> {
1112 let cache_key = SstMetaKey(file_id.region_id(), file_id.file_id());
1113 let compact = self
1114 .get_compact_sst_meta(&cache_key)
1115 .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy));
1116 let decoded = self
1117 .get_decoded_sst_meta(&cache_key)
1118 .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy));
1119
1120 if let Some(compact) = compact {
1121 CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc();
1122 metrics.mem_cache_hit += 1;
1123 if let Some(decoded) = decoded {
1124 CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc();
1125 return Some(decoded);
1126 }
1127
1128 CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc();
1129 match compact.decode(runtime).await {
1130 Ok(decoded) => {
1131 self.put_sst_meta_data(file_id, decoded.clone());
1132 return Some(decoded);
1133 }
1134 Err(err) => {
1135 warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id());
1136 self.remove_parquet_meta_data(file_id);
1137 }
1138 }
1139 } else if let Some(decoded) = decoded {
1140 CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc();
1143 CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc();
1144 metrics.mem_cache_hit += 1;
1145 return Some(decoded);
1146 } else {
1147 CACHE_MISS.with_label_values(&[SST_META_TYPE]).inc();
1148 }
1149
1150 let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet);
1151 if let Some(write_cache) = &self.write_cache {
1152 let file_cache = write_cache.file_cache();
1153 if self.sst_meta_cache_enabled() {
1154 if let Some(metadata) = file_cache
1155 .get_sst_meta_data(key, metrics, page_index_policy, runtime)
1156 .await
1157 {
1158 metrics.file_cache_hit += 1;
1159 let decoded = metadata.decoded();
1160 match metadata {
1161 SstMetaPreparation::Prepared(metadata) => {
1162 self.put_prepared_sst_meta(file_id, metadata, true);
1163 }
1164 SstMetaPreparation::DecodedOnly { encoding_error, .. } => {
1165 warn!(
1166 encoding_error;
1167 "Failed to encode file-cached SST metadata for memory cache, region_id: {}, file_id: {}",
1168 file_id.region_id(),
1169 file_id.file_id()
1170 );
1171 }
1172 }
1173 return Some(decoded);
1174 }
1175 } else if let Some(decoded) = file_cache
1176 .get_decoded_sst_meta_data(key, metrics, page_index_policy, runtime)
1177 .await
1178 {
1179 metrics.file_cache_hit += 1;
1180 return Some(decoded);
1181 }
1182 }
1183
1184 metrics.cache_miss += 1;
1185 None
1186 }
1187
1188 pub(crate) fn get_sst_meta_data_from_mem_cache(
1191 &self,
1192 file_id: RegionFileId,
1193 page_index_policy: PageIndexPolicy,
1194 ) -> Option<Arc<CachedSstMeta>> {
1195 let key = SstMetaKey(file_id.region_id(), file_id.file_id());
1196 let value = self
1197 .get_decoded_sst_meta(&key)
1198 .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy));
1199 update_hit_miss(value, SST_META_DECODED_TYPE)
1200 }
1201
1202 pub fn get_parquet_meta_data_from_mem_cache(
1205 &self,
1206 file_id: RegionFileId,
1207 ) -> Option<Arc<ParquetMetaData>> {
1208 self.get_sst_meta_data_from_mem_cache(file_id, PageIndexPolicy::Skip)
1209 .map(|metadata| metadata.parquet_metadata())
1210 }
1211
1212 pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc<CachedSstMeta>) {
1214 if let Some(cache) = &self.sst_decoded_meta_cache {
1215 let key = SstMetaKey(file_id.region_id(), file_id.file_id());
1216 CACHE_BYTES
1217 .with_label_values(&[SST_META_DECODED_TYPE])
1218 .add(decoded_meta_cache_weight(&key, &metadata).into());
1219 cache.insert(key, metadata);
1220 }
1221 }
1222
1223 pub(crate) fn put_prepared_sst_meta(
1225 &self,
1226 file_id: RegionFileId,
1227 metadata: PreparedSstMeta,
1228 retain_decoded: bool,
1229 ) {
1230 let key = SstMetaKey(file_id.region_id(), file_id.file_id());
1231 if let Some(cache) = &self.sst_meta_cache {
1232 CACHE_BYTES
1233 .with_label_values(&[SST_META_TYPE])
1234 .add(meta_cache_weight(&key, &metadata.compact).into());
1235 cache.insert(key.clone(), metadata.compact);
1236 }
1237 if retain_decoded {
1238 self.put_sst_meta_data(file_id, metadata.decoded);
1239 }
1240 }
1241
1242 fn get_compact_sst_meta(&self, key: &SstMetaKey) -> Option<Arc<CompactSstMeta>> {
1243 self.sst_meta_cache.as_ref()?.get(key)
1244 }
1245
1246 fn get_decoded_sst_meta(&self, key: &SstMetaKey) -> Option<Arc<CachedSstMeta>> {
1247 self.sst_decoded_meta_cache.as_ref()?.get(key)
1248 }
1249
1250 pub(crate) async fn get_compact_sst_meta_data(
1255 &self,
1256 file_id: RegionFileId,
1257 page_index_policy: PageIndexPolicy,
1258 ) -> Option<Arc<CachedSstMeta>> {
1259 let key = SstMetaKey(file_id.region_id(), file_id.file_id());
1260 let compact = self
1261 .get_compact_sst_meta(&key)
1262 .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy))?;
1263 match compact.decode(&common_runtime::global_runtime()).await {
1264 Ok(metadata) => Some(metadata),
1265 Err(err) => {
1266 warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id());
1267 self.remove_parquet_meta_data(file_id);
1268 None
1269 }
1270 }
1271 }
1272
1273 pub fn put_parquet_meta_data(
1275 &self,
1276 file_id: RegionFileId,
1277 metadata: Arc<ParquetMetaData>,
1278 region_metadata: Option<RegionMetadataRef>,
1279 ) {
1280 if self.sst_decoded_meta_cache.is_some() {
1281 let file_path = format!(
1282 "region_id={}, file_id={}",
1283 file_id.region_id(),
1284 file_id.file_id()
1285 );
1286 match CachedSstMeta::try_new_with_region_metadata(
1287 &file_path,
1288 Arc::unwrap_or_clone(metadata),
1289 region_metadata,
1290 ) {
1291 Ok(metadata) => self.put_sst_meta_data(file_id, Arc::new(metadata)),
1292 Err(err) => warn!(
1293 err; "Failed to decode region metadata while caching parquet metadata, region_id: {}, file_id: {}",
1294 file_id.region_id(),
1295 file_id.file_id()
1296 ),
1297 }
1298 }
1299 }
1300
1301 pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) {
1303 let key = SstMetaKey(file_id.region_id(), file_id.file_id());
1304 if let Some(cache) = &self.sst_meta_cache {
1305 cache.remove(&key);
1306 }
1307 if let Some(cache) = &self.sst_decoded_meta_cache {
1308 cache.remove(&key);
1309 }
1310 }
1311
1312 pub(crate) fn sst_meta_cache_is_full(&self) -> bool {
1314 let Some(cache) = &self.sst_meta_cache else {
1315 return true;
1316 };
1317 cache
1318 .policy()
1319 .max_capacity()
1320 .is_some_and(|capacity| cache.weighted_size() >= capacity)
1321 }
1322
1323 pub(crate) fn sst_meta_cache_enabled(&self) -> bool {
1325 self.sst_meta_cache.is_some()
1326 }
1327
1328 pub fn get_repeated_vector(
1330 &self,
1331 data_type: &ConcreteDataType,
1332 value: &Value,
1333 ) -> Option<VectorRef> {
1334 self.vector_cache.as_ref().and_then(|vector_cache| {
1335 let value = vector_cache.get(&(data_type.clone(), value.clone()));
1336 update_hit_miss(value, VECTOR_TYPE)
1337 })
1338 }
1339
1340 pub fn put_repeated_vector(&self, value: Value, vector: VectorRef) {
1342 if let Some(cache) = &self.vector_cache {
1343 let key = (vector.data_type(), value);
1344 CACHE_BYTES
1345 .with_label_values(&[VECTOR_TYPE])
1346 .add(vector_cache_weight(&key, &vector).into());
1347 cache.insert(key, vector);
1348 }
1349 }
1350
1351 pub fn get_page_ranges(
1353 &self,
1354 file_id: FileId,
1355 row_group_idx: usize,
1356 ranges: &[Range<u64>],
1357 ) -> Option<PageRangeLookup> {
1358 self.page_cache.as_ref().map(|page_cache| {
1359 let lookup = page_cache.lookup(file_id, row_group_idx, ranges);
1360 if lookup.cached_bytes > 0 {
1361 CACHE_HIT.with_label_values(&[PAGE_TYPE]).inc();
1362 }
1363 if !lookup.missing_ranges.is_empty() {
1364 CACHE_MISS.with_label_values(&[PAGE_TYPE]).inc();
1365 }
1366 lookup
1367 })
1368 }
1369
1370 pub fn put_page_ranges(
1372 &self,
1373 file_id: FileId,
1374 row_group_idx: usize,
1375 ranges: &[Range<u64>],
1376 pages: &[Bytes],
1377 ) {
1378 if let Some(cache) = &self.page_cache {
1379 cache.insert_ranges(file_id, row_group_idx, ranges, pages);
1380 }
1381 }
1382
1383 pub async fn evict_puffin_cache(&self, file_id: RegionIndexId) {
1385 if let Some(cache) = &self.bloom_filter_index_cache {
1386 cache.invalidate_file(file_id.file_id());
1387 }
1388
1389 if let Some(cache) = &self.inverted_index_cache {
1390 cache.invalidate_file(file_id.file_id());
1391 }
1392
1393 if let Some(cache) = &self.index_result_cache {
1394 cache.invalidate_file(file_id.file_id());
1395 }
1396
1397 if let Some(cache) = &self.puffin_metadata_cache {
1398 cache.remove(&file_id.to_string());
1399 }
1400
1401 if let Some(write_cache) = &self.write_cache {
1402 write_cache
1403 .remove(IndexKey::new(
1404 file_id.region_id(),
1405 file_id.file_id(),
1406 FileType::Puffin(file_id.version),
1407 ))
1408 .await;
1409 }
1410 }
1411
1412 pub fn get_selector_result(
1414 &self,
1415 selector_key: &SelectorResultKey,
1416 ) -> Option<Arc<SelectorResultValue>> {
1417 self.selector_result_cache
1418 .as_ref()
1419 .and_then(|selector_result_cache| selector_result_cache.get(selector_key))
1420 }
1421
1422 pub fn put_selector_result(
1424 &self,
1425 selector_key: SelectorResultKey,
1426 result: Arc<SelectorResultValue>,
1427 ) {
1428 if let Some(cache) = &self.selector_result_cache {
1429 CACHE_BYTES
1430 .with_label_values(&[SELECTOR_RESULT_TYPE])
1431 .add(selector_result_cache_weight(&selector_key, &result).into());
1432 cache.insert(selector_key, result);
1433 }
1434 }
1435
1436 pub(crate) fn invalidate_range_results(&self, file_id: FileId) {
1438 if let Some(worker) = &self.range_result_cache {
1439 worker.invalidate_file(file_id);
1440 }
1441 }
1442
1443 #[cfg(test)]
1445 pub(crate) fn get_range_result(
1446 &self,
1447 key: &RangeScanCacheKey,
1448 ) -> Option<Arc<RangeScanCacheValue>> {
1449 self.range_result_cache
1450 .as_ref()
1451 .and_then(|cache| update_hit_miss(cache.get(key), RANGE_RESULT_TYPE))
1452 }
1453
1454 #[cfg(any(test, feature = "test"))]
1456 pub(crate) fn put_range_result(
1457 &self,
1458 key: RangeScanCacheKey,
1459 result: Arc<RangeScanCacheValue>,
1460 ) {
1461 if let Some(cache) = &self.range_result_cache {
1462 cache.insert(key, result);
1463 }
1464 }
1465
1466 pub(crate) fn has_range_result_cache(&self) -> bool {
1468 self.range_result_cache.is_some()
1469 }
1470
1471 pub(crate) fn range_result_memory_limiter(&self) -> &Arc<RangeResultMemoryLimiter> {
1472 &self.range_result_memory_limiter
1473 }
1474
1475 pub(crate) fn range_result_cache_size(&self) -> usize {
1476 self.range_result_cache_size as usize
1477 }
1478
1479 pub(crate) fn write_cache(&self) -> Option<&WriteCacheRef> {
1481 self.write_cache.as_ref()
1482 }
1483
1484 pub(crate) fn inverted_index_cache(&self) -> Option<&InvertedIndexCacheRef> {
1485 self.inverted_index_cache.as_ref()
1486 }
1487
1488 pub(crate) fn bloom_filter_index_cache(&self) -> Option<&BloomFilterIndexCacheRef> {
1489 self.bloom_filter_index_cache.as_ref()
1490 }
1491
1492 pub(crate) fn puffin_metadata_cache(&self) -> Option<&PuffinMetadataCacheRef> {
1493 self.puffin_metadata_cache.as_ref()
1494 }
1495
1496 pub(crate) fn index_result_cache(&self) -> Option<&IndexResultCache> {
1497 self.index_result_cache.as_ref()
1498 }
1499
1500 pub(crate) fn get_prefilter_result(&self, key: &PrefilterKey) -> Option<Arc<BooleanBuffer>> {
1501 self.prefilter_result_cache
1502 .as_ref()
1503 .and_then(|cache| update_hit_miss(cache.get(key), PREFILTER_RESULT_TYPE))
1504 }
1505
1506 pub(crate) fn put_prefilter_result(&self, key: PrefilterKey, result: Arc<BooleanBuffer>) {
1507 if let Some(cache) = &self.prefilter_result_cache {
1508 CACHE_BYTES
1509 .with_label_values(&[PREFILTER_RESULT_TYPE])
1510 .add(prefilter_result_cache_weight(&key, &result).into());
1511 cache.insert(key, result);
1512 }
1513 }
1514}
1515
1516pub fn selector_result_cache_miss() {
1518 CACHE_MISS.with_label_values(&[SELECTOR_RESULT_TYPE]).inc()
1519}
1520
1521pub fn selector_result_cache_hit() {
1523 CACHE_HIT.with_label_values(&[SELECTOR_RESULT_TYPE]).inc()
1524}
1525
1526#[derive(Default)]
1528pub struct CacheManagerBuilder {
1529 sst_meta_cache_size: u64,
1530 vector_cache_size: u64,
1531 page_cache_size: u64,
1532 index_metadata_size: u64,
1533 index_content_size: u64,
1534 index_content_page_size: u64,
1535 index_result_cache_size: u64,
1536 prefilter_result_cache_size: u64,
1537 puffin_metadata_size: u64,
1538 write_cache: Option<WriteCacheRef>,
1539 selector_result_cache_size: u64,
1540 range_result_cache_size: u64,
1541}
1542
1543impl CacheManagerBuilder {
1544 pub fn sst_meta_cache_size(mut self, bytes: u64) -> Self {
1546 self.sst_meta_cache_size = bytes;
1547 self
1548 }
1549
1550 pub fn vector_cache_size(mut self, bytes: u64) -> Self {
1552 self.vector_cache_size = bytes;
1553 self
1554 }
1555
1556 pub fn page_cache_size(mut self, bytes: u64) -> Self {
1558 self.page_cache_size = bytes;
1559 self
1560 }
1561
1562 pub fn write_cache(mut self, cache: Option<WriteCacheRef>) -> Self {
1564 self.write_cache = cache;
1565 self
1566 }
1567
1568 pub fn index_metadata_size(mut self, bytes: u64) -> Self {
1570 self.index_metadata_size = bytes;
1571 self
1572 }
1573
1574 pub fn index_content_size(mut self, bytes: u64) -> Self {
1576 self.index_content_size = bytes;
1577 self
1578 }
1579
1580 pub fn index_content_page_size(mut self, bytes: u64) -> Self {
1582 self.index_content_page_size = bytes;
1583 self
1584 }
1585
1586 pub fn index_result_cache_size(mut self, bytes: u64) -> Self {
1588 self.index_result_cache_size = bytes;
1589 self
1590 }
1591
1592 pub fn prefilter_result_cache_size(mut self, bytes: u64) -> Self {
1594 self.prefilter_result_cache_size = bytes;
1595 self
1596 }
1597
1598 pub fn puffin_metadata_size(mut self, bytes: u64) -> Self {
1600 self.puffin_metadata_size = bytes;
1601 self
1602 }
1603
1604 pub fn selector_result_cache_size(mut self, bytes: u64) -> Self {
1606 self.selector_result_cache_size = bytes;
1607 self
1608 }
1609
1610 pub fn range_result_cache_size(mut self, bytes: u64) -> Self {
1612 self.range_result_cache_size = bytes;
1613 self
1614 }
1615
1616 pub fn build(self) -> CacheManager {
1618 let compact_meta_capacity = self.sst_meta_cache_size.div_ceil(2);
1622 let decoded_meta_capacity = self.sst_meta_cache_size / 2;
1623 let sst_meta_cache = (compact_meta_capacity != 0).then(|| {
1624 Cache::builder()
1625 .max_capacity(compact_meta_capacity)
1626 .weigher(meta_cache_weight)
1627 .eviction_listener(|k, v, cause| {
1628 let size = meta_cache_weight(&k, &v);
1629 CACHE_BYTES
1630 .with_label_values(&[SST_META_TYPE])
1631 .sub(size.into());
1632 CACHE_EVICTION
1633 .with_label_values(&[SST_META_TYPE, removal_cause_str(cause)])
1634 .inc();
1635 })
1636 .build()
1637 });
1638 let sst_decoded_meta_cache = (decoded_meta_capacity != 0).then(|| {
1639 Cache::builder()
1640 .max_capacity(decoded_meta_capacity)
1641 .weigher(decoded_meta_cache_weight)
1642 .eviction_listener(|k, v, cause| {
1643 let size = decoded_meta_cache_weight(&k, &v);
1644 CACHE_BYTES
1645 .with_label_values(&[SST_META_DECODED_TYPE])
1646 .sub(size.into());
1647 CACHE_EVICTION
1648 .with_label_values(&[SST_META_DECODED_TYPE, removal_cause_str(cause)])
1649 .inc();
1650 })
1651 .build()
1652 });
1653 let vector_cache = (self.vector_cache_size != 0).then(|| {
1654 Cache::builder()
1655 .max_capacity(self.vector_cache_size)
1656 .weigher(vector_cache_weight)
1657 .eviction_listener(|k, v, cause| {
1658 let size = vector_cache_weight(&k, &v);
1659 CACHE_BYTES
1660 .with_label_values(&[VECTOR_TYPE])
1661 .sub(size.into());
1662 CACHE_EVICTION
1663 .with_label_values(&[VECTOR_TYPE, removal_cause_str(cause)])
1664 .inc();
1665 })
1666 .build()
1667 });
1668 let page_cache =
1669 (self.page_cache_size != 0).then(|| PageRangeCache::new(self.page_cache_size));
1670 let inverted_index_cache = InvertedIndexCache::new(
1671 self.index_metadata_size,
1672 self.index_content_size,
1673 self.index_content_page_size,
1674 );
1675 let bloom_filter_index_cache = BloomFilterIndexCache::new(
1677 self.index_metadata_size,
1678 self.index_content_size,
1679 self.index_content_page_size,
1680 );
1681 let index_result_cache = (self.index_result_cache_size != 0)
1682 .then(|| IndexResultCache::new(self.index_result_cache_size));
1683 let prefilter_result_cache = (self.prefilter_result_cache_size != 0)
1684 .then(|| new_prefilter_result_cache(self.prefilter_result_cache_size));
1685 let puffin_metadata_cache =
1686 PuffinMetadataCache::new(self.puffin_metadata_size, &CACHE_BYTES);
1687 let selector_result_cache = (self.selector_result_cache_size != 0).then(|| {
1688 Cache::builder()
1689 .max_capacity(self.selector_result_cache_size)
1690 .weigher(selector_result_cache_weight)
1691 .eviction_listener(|k, v, cause| {
1692 let size = selector_result_cache_weight(&k, &v);
1693 CACHE_BYTES
1694 .with_label_values(&[SELECTOR_RESULT_TYPE])
1695 .sub(size.into());
1696 CACHE_EVICTION
1697 .with_label_values(&[SELECTOR_RESULT_TYPE, removal_cause_str(cause)])
1698 .inc();
1699 })
1700 .build()
1701 });
1702 let range_result_cache = (self.range_result_cache_size != 0)
1703 .then(|| range_cache::RangeCacheWorker::new(self.range_result_cache_size));
1704 CacheManager {
1705 sst_meta_cache,
1706 sst_decoded_meta_cache,
1707 vector_cache,
1708 page_cache,
1709 write_cache: self.write_cache,
1710 inverted_index_cache: Some(Arc::new(inverted_index_cache)),
1711 bloom_filter_index_cache: Some(Arc::new(bloom_filter_index_cache)),
1712 puffin_metadata_cache: Some(Arc::new(puffin_metadata_cache)),
1713 selector_result_cache,
1714 range_result_cache,
1715 range_result_cache_size: self.range_result_cache_size,
1716 range_result_memory_limiter: Arc::new(RangeResultMemoryLimiter::new(
1717 self.range_result_cache_size as usize,
1718 RANGE_RESULT_CONCAT_MEMORY_PERMIT.as_bytes() as usize,
1719 )),
1720 index_result_cache,
1721 prefilter_result_cache,
1722 }
1723 }
1724}
1725
1726fn meta_cache_weight(k: &SstMetaKey, v: &Arc<CompactSstMeta>) -> u32 {
1727 let size = k.estimated_size() + mem::size_of::<CompactSstMeta>() + v.encoded_metadata.len();
1730 u32::try_from(size).unwrap_or(u32::MAX)
1731}
1732
1733fn decoded_meta_cache_weight(k: &SstMetaKey, v: &Arc<CachedSstMeta>) -> u32 {
1734 let size = k.estimated_size() + v.parquet_metadata_size + v.region_metadata_weight;
1735 u32::try_from(size).unwrap_or(u32::MAX)
1736}
1737
1738fn vector_cache_weight(_k: &(ConcreteDataType, Value), v: &VectorRef) -> u32 {
1739 (mem::size_of::<ConcreteDataType>() + mem::size_of::<Value>() + v.memory_size()) as u32
1741}
1742
1743fn page_cache_weight(k: &PageFragmentKey, v: &Bytes) -> u32 {
1744 (k.estimated_size() + mem::size_of::<Bytes>() + v.len()) as u32
1745}
1746
1747fn selector_result_cache_weight(k: &SelectorResultKey, v: &Arc<SelectorResultValue>) -> u32 {
1748 (mem::size_of_val(k) + v.estimated_size()) as u32
1749}
1750
1751fn range_result_cache_weight(k: &RangeScanCacheKey, v: &Arc<RangeScanCacheValue>) -> u32 {
1752 (k.estimated_size() + v.estimated_size()) as u32
1753}
1754
1755fn update_hit_miss<T>(value: Option<T>, cache_type: &str) -> Option<T> {
1757 if value.is_some() {
1758 CACHE_HIT.with_label_values(&[cache_type]).inc();
1759 } else {
1760 CACHE_MISS.with_label_values(&[cache_type]).inc();
1761 }
1762 value
1763}
1764
1765#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1767struct SstMetaKey(RegionId, FileId);
1768
1769impl SstMetaKey {
1770 fn estimated_size(&self) -> usize {
1772 mem::size_of::<Self>()
1773 }
1774}
1775
1776#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1777struct PageFragmentGroupKey {
1778 file_id: FileId,
1779 row_group_idx: usize,
1780}
1781
1782#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1784pub struct PageFragmentKey {
1785 file_id: FileId,
1787 row_group_idx: usize,
1789 start: u64,
1791 end: u64,
1793}
1794
1795impl PageFragmentKey {
1796 fn new(file_id: FileId, row_group_idx: usize, range: &Range<u64>) -> PageFragmentKey {
1797 PageFragmentKey {
1798 file_id,
1799 row_group_idx,
1800 start: range.start,
1801 end: range.end,
1802 }
1803 }
1804
1805 fn group_key(&self) -> PageFragmentGroupKey {
1806 PageFragmentGroupKey {
1807 file_id: self.file_id,
1808 row_group_idx: self.row_group_idx,
1809 }
1810 }
1811
1812 fn estimated_size(&self) -> usize {
1814 mem::size_of::<Self>()
1815 }
1816}
1817
1818#[derive(Clone)]
1820pub struct PageRangePart {
1821 pub range: Range<u64>,
1823 pub bytes: Bytes,
1825}
1826
1827pub struct PageRangeLookup {
1829 pub cached_parts: Vec<Vec<PageRangePart>>,
1831 pub missing_ranges: Vec<Range<u64>>,
1833 pub cached_range_count: usize,
1835 pub cached_bytes: u64,
1837}
1838
1839impl PageRangeLookup {
1840 pub fn is_fully_cached(&self) -> bool {
1841 self.missing_ranges.is_empty()
1842 }
1843}
1844
1845type PageFragmentRangeIndex = BTreeMap<(u64, u64), PageFragmentKey>;
1846type PageFragmentIndex = HashMap<PageFragmentGroupKey, PageFragmentRangeIndex>;
1847
1848pub struct PageRangeCache {
1850 cache: Cache<PageFragmentKey, Bytes>,
1851 index: RwLock<PageFragmentIndex>,
1852}
1853
1854impl PageRangeCache {
1855 fn new(capacity: u64) -> Arc<PageRangeCache> {
1856 Arc::new_cyclic(|weak_cache: &std::sync::Weak<PageRangeCache>| {
1857 let cache = Cache::builder()
1858 .max_capacity(capacity)
1859 .weigher(page_cache_weight)
1860 .eviction_listener({
1861 let weak_cache = weak_cache.clone();
1862 move |k, v, cause| {
1863 let size = page_cache_weight(&k, &v);
1864 CACHE_BYTES.with_label_values(&[PAGE_TYPE]).sub(size.into());
1865 CACHE_EVICTION
1866 .with_label_values(&[PAGE_TYPE, removal_cause_str(cause)])
1867 .inc();
1868
1869 if let Some(cache) = weak_cache.upgrade()
1870 && !matches!(cause, RemovalCause::Replaced)
1871 {
1872 cache.remove_index_entry(*k);
1873 }
1874 }
1875 })
1876 .build();
1877
1878 PageRangeCache {
1879 cache,
1880 index: RwLock::new(HashMap::new()),
1881 }
1882 })
1883 }
1884
1885 fn lookup(
1886 &self,
1887 file_id: FileId,
1888 row_group_idx: usize,
1889 ranges: &[Range<u64>],
1890 ) -> PageRangeLookup {
1891 let mut cached_parts = Vec::with_capacity(ranges.len());
1892 let mut missing_ranges = Vec::new();
1893 let mut cached_range_count = 0;
1894 let mut cached_bytes = 0;
1895
1896 for range in ranges {
1897 if range.start >= range.end {
1898 cached_parts.push(Vec::new());
1899 continue;
1900 }
1901
1902 let mut parts = Vec::new();
1903 let candidates = self.find_index_candidates(file_id, row_group_idx, range);
1904 let mut stale_keys = Vec::new();
1905
1906 for fragment_key in candidates {
1907 if let Some(bytes) = self.cache.get(&fragment_key) {
1908 let part_start = range.start.max(fragment_key.start);
1909 let part_end = range.end.min(fragment_key.end);
1910 let slice_start = (part_start - fragment_key.start) as usize;
1911 let slice_end = (part_end - fragment_key.start) as usize;
1912 parts.push(PageRangePart {
1913 range: part_start..part_end,
1914 bytes: bytes.slice(slice_start..slice_end),
1915 });
1916 } else {
1917 stale_keys.push(fragment_key);
1918 }
1919 }
1920 for key in stale_keys {
1921 self.remove_uncached_index_entry(key);
1922 }
1923
1924 let mut cursor = range.start;
1925 let mut compacted_parts: Vec<PageRangePart> = Vec::with_capacity(parts.len());
1926 for part in parts {
1927 if part.range.end <= cursor {
1928 continue;
1929 }
1930
1931 let part = if part.range.start < cursor {
1932 let offset = (cursor - part.range.start) as usize;
1933 PageRangePart {
1934 range: cursor..part.range.end,
1935 bytes: part.bytes.slice(offset..),
1936 }
1937 } else {
1938 part
1939 };
1940
1941 if cursor < part.range.start {
1942 missing_ranges.push(cursor..part.range.start);
1943 }
1944 cached_bytes += part.range.end - part.range.start;
1945 cached_range_count += 1;
1946 cursor = part.range.end;
1947 compacted_parts.push(part);
1948
1949 if cursor >= range.end {
1950 break;
1951 }
1952 }
1953
1954 if cursor < range.end {
1955 missing_ranges.push(cursor..range.end);
1956 }
1957 cached_parts.push(compacted_parts);
1958 }
1959
1960 PageRangeLookup {
1961 cached_parts,
1962 missing_ranges,
1963 cached_range_count,
1964 cached_bytes,
1965 }
1966 }
1967
1968 fn insert_ranges(
1969 &self,
1970 file_id: FileId,
1971 row_group_idx: usize,
1972 ranges: &[Range<u64>],
1973 pages: &[Bytes],
1974 ) {
1975 for (range, bytes) in ranges.iter().zip(pages) {
1976 if range.start >= range.end || bytes.len() as u64 != range.end - range.start {
1977 continue;
1978 }
1979
1980 let key = PageFragmentKey::new(file_id, row_group_idx, range);
1981 let bytes = Bytes::copy_from_slice(bytes.as_ref());
1982 let size = page_cache_weight(&key, &bytes);
1983 CACHE_BYTES.with_label_values(&[PAGE_TYPE]).add(size.into());
1984 self.cache.insert(key, bytes);
1985 let mut index = self.index.write().unwrap();
1986 index
1987 .entry(key.group_key())
1988 .or_default()
1989 .insert((key.start, key.end), key);
1990 }
1991 }
1992
1993 fn find_index_candidates(
1994 &self,
1995 file_id: FileId,
1996 row_group_idx: usize,
1997 range: &Range<u64>,
1998 ) -> Vec<PageFragmentKey> {
1999 let group_key = PageFragmentGroupKey {
2000 file_id,
2001 row_group_idx,
2002 };
2003 let index = self.index.read().unwrap();
2004 index
2005 .get(&group_key)
2006 .map(|ranges| {
2007 ranges
2008 .range(..(range.end, 0))
2009 .filter_map(|(_, fragment_key)| {
2010 (fragment_key.end > range.start).then_some(*fragment_key)
2011 })
2012 .collect()
2013 })
2014 .unwrap_or_default()
2015 }
2016
2017 fn remove_uncached_index_entry(&self, key: PageFragmentKey) {
2018 let group_key = key.group_key();
2019 let mut index = self.index.write().unwrap();
2020 if self.cache.contains_key(&key) {
2021 return;
2022 }
2023
2024 Self::remove_index_entry_locked(&mut index, group_key, key);
2025 }
2026
2027 fn remove_index_entry(&self, key: PageFragmentKey) {
2028 let group_key = key.group_key();
2029 let mut index = self.index.write().unwrap();
2030 Self::remove_index_entry_locked(&mut index, group_key, key);
2031 }
2032
2033 fn remove_index_entry_locked(
2034 index: &mut PageFragmentIndex,
2035 group_key: PageFragmentGroupKey,
2036 key: PageFragmentKey,
2037 ) {
2038 let Some(ranges) = index.get_mut(&group_key) else {
2039 return;
2040 };
2041
2042 let removed = ranges
2043 .get(&(key.start, key.end))
2044 .is_some_and(|current| current == &key);
2045 if removed {
2046 ranges.remove(&(key.start, key.end));
2047 }
2048 if ranges.is_empty() {
2049 index.remove(&group_key);
2050 }
2051 }
2052}
2053
2054#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2056pub struct SelectorResultKey {
2057 pub file_id: FileId,
2059 pub row_group_idx: usize,
2061 pub selector: TimeSeriesRowSelector,
2063}
2064
2065pub enum SelectorResult {
2067 PrimaryKey(Vec<Batch>),
2069 Flat(Vec<RecordBatch>),
2071}
2072
2073pub struct SelectorResultValue {
2075 pub result: SelectorResult,
2077 pub read_cols: ParquetReadColumns,
2079 pub json_target_types: JsonTargetTypes,
2084}
2085
2086impl SelectorResultValue {
2087 pub fn new(result: Vec<Batch>, read_cols: ParquetReadColumns) -> SelectorResultValue {
2089 SelectorResultValue {
2090 result: SelectorResult::PrimaryKey(result),
2091 read_cols,
2092 json_target_types: Arc::default(),
2093 }
2094 }
2095
2096 pub fn new_flat(
2098 result: Vec<RecordBatch>,
2099 read_cols: ParquetReadColumns,
2100 json_target_types: JsonTargetTypes,
2101 ) -> SelectorResultValue {
2102 SelectorResultValue {
2103 result: SelectorResult::Flat(result),
2104 read_cols,
2105 json_target_types,
2106 }
2107 }
2108
2109 fn estimated_size(&self) -> usize {
2111 let result_size: usize = match &self.result {
2112 SelectorResult::PrimaryKey(batches) => {
2113 batches.iter().map(|batch| batch.memory_size()).sum()
2114 }
2115 SelectorResult::Flat(batches) => batches.iter().map(record_batch_estimated_size).sum(),
2116 };
2117 result_size
2118 + self.json_target_types.len() * (size_of::<ColumnId>() + size_of::<JsonNativeType>())
2119 }
2120}
2121
2122type SstMetaCache = Cache<SstMetaKey, Arc<CompactSstMeta>>;
2124type SstDecodedMetaCache = Cache<SstMetaKey, Arc<CachedSstMeta>>;
2126type VectorCache = Cache<(ConcreteDataType, Value), VectorRef>;
2130type SelectorResultCache = Cache<SelectorResultKey, Arc<SelectorResultValue>>;
2132type RangeResultCache = Arc<range_cache::RangeCacheWorker>;
2134
2135#[cfg(test)]
2136mod tests {
2137 use std::sync::Arc;
2138
2139 use api::v1::SemanticType;
2140 use api::v1::index::{BloomFilterMeta, InvertedIndexMetas};
2141 use datatypes::schema::ColumnSchema;
2142 use datatypes::vectors::Int64Vector;
2143 use parquet::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation};
2144 use puffin::file_metadata::FileMetadata;
2145 use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
2146 use store_api::storage::ColumnId;
2147
2148 use super::*;
2149 use crate::cache::index::bloom_filter_index::Tag;
2150 use crate::cache::index::result_cache::PredicateKey;
2151 use crate::cache::test_util::{
2152 parquet_meta, sst_parquet_meta, sst_parquet_meta_with_region_metadata,
2153 };
2154 use crate::read::range_cache::{
2155 RangeScanCacheKey, RangeScanCacheValue, ScanRequestFingerprintBuilder,
2156 };
2157 use crate::read::read_columns::ReadColumns;
2158 use crate::sst::parquet::row_selection::RowGroupSelection;
2159
2160 #[tokio::test]
2161 async fn test_disable_cache() {
2162 let cache = CacheManager::default();
2163 assert!(cache.sst_meta_cache.is_none());
2164 assert!(cache.sst_decoded_meta_cache.is_none());
2165 assert!(cache.vector_cache.is_none());
2166 assert!(cache.page_cache.is_none());
2167
2168 let region_id = RegionId::new(1, 1);
2169 let file_id = RegionFileId::new(region_id, FileId::random());
2170 let metadata = parquet_meta();
2171 let mut metrics = MetadataCacheMetrics::default();
2172 cache.put_parquet_meta_data(file_id, metadata, None);
2173 assert!(
2174 cache
2175 .get_sst_meta_data(
2176 file_id,
2177 &mut metrics,
2178 Default::default(),
2179 &common_runtime::global_runtime()
2180 )
2181 .await
2182 .is_none()
2183 );
2184
2185 let value = Value::Int64(10);
2186 let vector: VectorRef = Arc::new(Int64Vector::from_slice([10, 10, 10, 10]));
2187 cache.put_repeated_vector(value.clone(), vector.clone());
2188 assert!(
2189 cache
2190 .get_repeated_vector(&ConcreteDataType::int64_datatype(), &value)
2191 .is_none()
2192 );
2193
2194 cache.put_page_ranges(
2195 file_id.file_id(),
2196 1,
2197 &[Range { start: 0, end: 5 }],
2198 &[Bytes::from_static(b"abcde")],
2199 );
2200 assert!(
2201 cache
2202 .get_page_ranges(file_id.file_id(), 1, &[Range { start: 0, end: 5 }])
2203 .is_none()
2204 );
2205
2206 assert!(cache.write_cache().is_none());
2207 }
2208
2209 #[test]
2210 fn test_sst_meta_cache_splits_capacity() {
2211 let cache = CacheManager::builder().sst_meta_cache_size(101).build();
2212
2213 assert_eq!(
2214 Some(51),
2215 cache
2216 .sst_meta_cache
2217 .as_ref()
2218 .unwrap()
2219 .policy()
2220 .max_capacity()
2221 );
2222 assert_eq!(
2223 Some(50),
2224 cache
2225 .sst_decoded_meta_cache
2226 .as_ref()
2227 .unwrap()
2228 .policy()
2229 .max_capacity()
2230 );
2231 }
2232
2233 #[tokio::test]
2234 async fn test_parquet_meta_cache() {
2235 let cache = CacheManager::builder().sst_meta_cache_size(2000).build();
2236 let mut metrics = MetadataCacheMetrics::default();
2237 let region_id = RegionId::new(1, 1);
2238 let file_id = RegionFileId::new(region_id, FileId::random());
2239 assert!(
2240 cache
2241 .get_sst_meta_data(
2242 file_id,
2243 &mut metrics,
2244 Default::default(),
2245 &common_runtime::global_runtime()
2246 )
2247 .await
2248 .is_none()
2249 );
2250 let (metadata, region_metadata) = sst_parquet_meta();
2251 cache.put_parquet_meta_data(file_id, metadata, None);
2252 let cached = cache
2253 .get_sst_meta_data(
2254 file_id,
2255 &mut metrics,
2256 Default::default(),
2257 &common_runtime::global_runtime(),
2258 )
2259 .await
2260 .unwrap();
2261 assert_eq!(region_metadata, cached.region_metadata());
2262 assert!(
2263 cached
2264 .parquet_metadata()
2265 .file_metadata()
2266 .key_value_metadata()
2267 .is_none_or(|key_values| {
2268 key_values
2269 .iter()
2270 .all(|key_value| key_value.key != PARQUET_METADATA_KEY)
2271 })
2272 );
2273 cache.remove_parquet_meta_data(file_id);
2274 assert!(
2275 cache
2276 .get_sst_meta_data(
2277 file_id,
2278 &mut metrics,
2279 Default::default(),
2280 &common_runtime::global_runtime()
2281 )
2282 .await
2283 .is_none()
2284 );
2285 }
2286
2287 #[tokio::test]
2288 async fn test_parquet_meta_cache_with_provided_region_metadata() {
2289 let cache = CacheManager::builder().sst_meta_cache_size(2000).build();
2290 let mut metrics = MetadataCacheMetrics::default();
2291 let region_id = RegionId::new(1, 1);
2292 let file_id = RegionFileId::new(region_id, FileId::random());
2293 let (metadata, region_metadata) = sst_parquet_meta();
2294
2295 cache.put_parquet_meta_data(file_id, metadata, Some(region_metadata.clone()));
2296
2297 let cached = cache
2298 .get_sst_meta_data(
2299 file_id,
2300 &mut metrics,
2301 Default::default(),
2302 &common_runtime::global_runtime(),
2303 )
2304 .await
2305 .unwrap();
2306 assert!(Arc::ptr_eq(®ion_metadata, &cached.region_metadata()));
2307 }
2308
2309 #[tokio::test]
2310 async fn test_compact_sst_meta_round_trip() {
2311 let cache = CacheManager::builder()
2312 .sst_meta_cache_size(1024 * 1024)
2313 .build();
2314 let region_id = RegionId::new(1, 1);
2315 let file_id = RegionFileId::new(region_id, FileId::random());
2316 let (metadata, expected_region_metadata) = sst_parquet_meta();
2317 let metadata = Arc::unwrap_or_clone(metadata);
2318 let offset_indexes = metadata
2319 .row_groups()
2320 .iter()
2321 .map(|row_group| {
2322 row_group
2323 .columns()
2324 .iter()
2325 .map(|_| OffsetIndexMetaData {
2326 page_locations: vec![PageLocation {
2327 offset: 42,
2328 compressed_page_size: 128,
2329 first_row_index: 0,
2330 }],
2331 unencoded_byte_array_data_bytes: None,
2332 })
2333 .collect()
2334 })
2335 .collect();
2336 let metadata = metadata
2337 .into_builder()
2338 .set_offset_index(Some(offset_indexes))
2339 .build();
2340 let expected_rows = metadata.file_metadata().num_rows();
2341 let prepared = prepare_sst_meta(
2342 "test.parquet",
2343 metadata,
2344 None,
2345 PageIndexPolicy::Required,
2346 &common_runtime::global_runtime(),
2347 )
2348 .await
2349 .unwrap();
2350 let SstMetaPreparation::Prepared(prepared) = prepared else {
2351 panic!("valid metadata should produce a compact cache entry");
2352 };
2353
2354 cache.put_prepared_sst_meta(file_id, prepared, false);
2356 let mut metrics = MetadataCacheMetrics::default();
2357 let cached = cache
2358 .get_sst_meta_data(
2359 file_id,
2360 &mut metrics,
2361 PageIndexPolicy::Required,
2362 &common_runtime::global_runtime(),
2363 )
2364 .await
2365 .unwrap();
2366
2367 assert_eq!(1, metrics.mem_cache_hit);
2368 assert_eq!(0, metrics.cache_miss);
2369 assert_eq!(
2370 expected_rows,
2371 cached.parquet_metadata.file_metadata().num_rows()
2372 );
2373 assert!(cached.parquet_metadata.offset_index().is_some());
2374 assert_eq!(expected_region_metadata, cached.region_metadata());
2375 assert!(
2376 cached
2377 .parquet_metadata
2378 .file_metadata()
2379 .key_value_metadata()
2380 .is_none_or(|key_values| key_values
2381 .iter()
2382 .all(|key_value| key_value.key != PARQUET_METADATA_KEY))
2383 );
2384 }
2385
2386 #[test]
2387 fn test_cache_encoding_failure_preserves_decoded_metadata() {
2388 let (metadata, expected_region_metadata) = sst_parquet_meta();
2389 let expected_rows = metadata.file_metadata().num_rows();
2390 let cache_encoding: Result<(Bytes, usize)> = Err(UnexpectedSnafu {
2391 reason: "injected compact metadata encoding failure",
2392 }
2393 .build());
2394
2395 let preparation = finish_sst_meta_preparation(
2396 "test.parquet",
2397 Arc::unwrap_or_clone(metadata),
2398 None,
2399 PageIndexPolicy::Skip,
2400 cache_encoding,
2401 )
2402 .unwrap();
2403 let SstMetaPreparation::DecodedOnly {
2404 decoded,
2405 encoding_error,
2406 } = preparation
2407 else {
2408 panic!("cache encoding failure should return decoded-only metadata");
2409 };
2410
2411 assert_eq!(
2412 expected_rows,
2413 decoded.parquet_metadata.file_metadata().num_rows()
2414 );
2415 assert_eq!(expected_region_metadata, decoded.region_metadata());
2416 assert!(
2417 encoding_error
2418 .to_string()
2419 .contains("injected compact metadata encoding failure")
2420 );
2421 }
2422
2423 #[tokio::test]
2424 async fn test_parquet_meta_cache_respects_page_index_policy() {
2425 let cache = CacheManager::builder().sst_meta_cache_size(2000).build();
2426 let region_id = RegionId::new(1, 1);
2427 let file_id = RegionFileId::new(region_id, FileId::random());
2428 let (metadata, _) = sst_parquet_meta();
2429
2430 let skip_metadata = Arc::new(
2431 CachedSstMeta::try_new_with_page_index_policy(
2432 "test.parquet",
2433 Arc::unwrap_or_clone(metadata.clone()),
2434 None,
2435 PageIndexPolicy::Skip,
2436 )
2437 .unwrap(),
2438 );
2439 cache.put_sst_meta_data(file_id, skip_metadata);
2440
2441 let mut metrics = MetadataCacheMetrics::default();
2442 assert!(
2443 cache
2444 .get_sst_meta_data(
2445 file_id,
2446 &mut metrics,
2447 PageIndexPolicy::Optional,
2448 &common_runtime::global_runtime()
2449 )
2450 .await
2451 .is_none()
2452 );
2453 assert_eq!(1, metrics.cache_miss);
2454
2455 let optional_metadata = Arc::new(
2456 CachedSstMeta::try_new_with_page_index_policy(
2457 "test.parquet",
2458 Arc::unwrap_or_clone(metadata),
2459 None,
2460 PageIndexPolicy::Optional,
2461 )
2462 .unwrap(),
2463 );
2464 cache.put_sst_meta_data(file_id, optional_metadata);
2465
2466 let mut metrics = MetadataCacheMetrics::default();
2467 assert!(
2468 cache
2469 .get_sst_meta_data(
2470 file_id,
2471 &mut metrics,
2472 PageIndexPolicy::Optional,
2473 &common_runtime::global_runtime()
2474 )
2475 .await
2476 .is_some()
2477 );
2478 assert_eq!(1, metrics.mem_cache_hit);
2479
2480 let mut metrics = MetadataCacheMetrics::default();
2481 assert!(
2482 cache
2483 .get_sst_meta_data(
2484 file_id,
2485 &mut metrics,
2486 PageIndexPolicy::Skip,
2487 &common_runtime::global_runtime()
2488 )
2489 .await
2490 .is_some()
2491 );
2492 assert_eq!(1, metrics.mem_cache_hit);
2493 }
2494
2495 #[test]
2496 fn test_decoded_meta_cache_weight_accounts_for_region_metadata() {
2497 let region_metadata = Arc::new(wide_region_metadata(128));
2498 let json_len = region_metadata.to_json().unwrap().len();
2499 let metadata = sst_parquet_meta_with_region_metadata(region_metadata.clone());
2500 let cached = Arc::new(
2501 CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(),
2502 );
2503 let key = SstMetaKey(region_metadata.region_id, FileId::random());
2504
2505 assert!(cached.region_metadata_weight > json_len);
2506 assert_eq!(
2507 decoded_meta_cache_weight(&key, &cached) as usize,
2508 key.estimated_size() + cached.parquet_metadata_size + cached.region_metadata_weight
2509 );
2510 assert_eq!(
2511 cached.parquet_metadata_size,
2512 parquet_meta_size(&cached.parquet_metadata)
2513 );
2514 }
2515
2516 #[test]
2517 fn test_decoded_meta_cache_weight_saturates_on_overflow() {
2518 let region_metadata = Arc::new(wide_region_metadata(1));
2519 let metadata = sst_parquet_meta_with_region_metadata(region_metadata.clone());
2520 let mut cached =
2521 CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap();
2522 cached.region_metadata_weight = u32::MAX as usize + 1;
2523 let cached = Arc::new(cached);
2524 let key = SstMetaKey(region_metadata.region_id, FileId::random());
2525
2526 assert_eq!(u32::MAX, decoded_meta_cache_weight(&key, &cached));
2527 }
2528
2529 #[test]
2530 fn test_repeated_vector_cache() {
2531 let cache = CacheManager::builder().vector_cache_size(4096).build();
2532 let value = Value::Int64(10);
2533 assert!(
2534 cache
2535 .get_repeated_vector(&ConcreteDataType::int64_datatype(), &value)
2536 .is_none()
2537 );
2538 let vector: VectorRef = Arc::new(Int64Vector::from_slice([10, 10, 10, 10]));
2539 cache.put_repeated_vector(value.clone(), vector.clone());
2540 let cached = cache
2541 .get_repeated_vector(&ConcreteDataType::int64_datatype(), &value)
2542 .unwrap();
2543 assert_eq!(vector, cached);
2544 }
2545
2546 #[test]
2547 fn test_page_cache() {
2548 let cache = CacheManager::builder().page_cache_size(1000).build();
2549 let file_id = FileId::random();
2550 let uncached = 0..10;
2551 assert_eq!(
2552 vec![0..10],
2553 cache
2554 .get_page_ranges(file_id, 0, std::slice::from_ref(&uncached))
2555 .unwrap()
2556 .missing_ranges
2557 );
2558
2559 let cached = 100..500;
2560 cache.put_page_ranges(
2561 file_id,
2562 0,
2563 std::slice::from_ref(&cached),
2564 &[Bytes::from(vec![7; 400])],
2565 );
2566
2567 let subrange = 200..300;
2568 let lookup = cache
2569 .get_page_ranges(file_id, 0, std::slice::from_ref(&subrange))
2570 .unwrap();
2571 assert!(lookup.is_fully_cached());
2572 assert_eq!(100, lookup.cached_bytes);
2573 assert_eq!(1, lookup.cached_parts.len());
2574 assert_eq!(200..300, lookup.cached_parts[0][0].range);
2575 assert_eq!(100, lookup.cached_parts[0][0].bytes.len());
2576
2577 let overlapping = 400..600;
2578 let lookup = cache
2579 .get_page_ranges(file_id, 0, std::slice::from_ref(&overlapping))
2580 .unwrap();
2581 assert!(!lookup.is_fully_cached());
2582 assert_eq!(100, lookup.cached_bytes);
2583 assert_eq!(vec![500..600], lookup.missing_ranges);
2584 assert_eq!(400..500, lookup.cached_parts[0][0].range);
2585 }
2586
2587 #[test]
2588 fn test_page_cache_detaches_fragment_bytes() {
2589 let cache = PageRangeCache::new(1000);
2590 let file_id = FileId::random();
2591 let backing = Bytes::from(vec![1; 1024]);
2592 let page = backing.slice(512..522);
2593 let page_ptr = page.as_ptr();
2594 let range = 0..10;
2595
2596 cache.insert_ranges(
2597 file_id,
2598 0,
2599 std::slice::from_ref(&range),
2600 std::slice::from_ref(&page),
2601 );
2602
2603 let lookup = cache.lookup(file_id, 0, std::slice::from_ref(&range));
2604 assert!(lookup.is_fully_cached());
2605 assert_eq!(1, lookup.cached_parts[0].len());
2606 assert_eq!(&page[..], &lookup.cached_parts[0][0].bytes[..]);
2607 assert_ne!(page_ptr, lookup.cached_parts[0][0].bytes.as_ptr());
2608 }
2609
2610 #[test]
2611 fn test_page_cache_replaces_fragment() {
2612 let cache = PageRangeCache::new(1000);
2613 let file_id = FileId::random();
2614 let range = 0..10;
2615
2616 cache.insert_ranges(
2617 file_id,
2618 0,
2619 std::slice::from_ref(&range),
2620 &[Bytes::from(vec![1; 10])],
2621 );
2622 cache.insert_ranges(
2623 file_id,
2624 0,
2625 std::slice::from_ref(&range),
2626 &[Bytes::from(vec![2; 10])],
2627 );
2628 cache.cache.run_pending_tasks();
2629 assert_eq!(
2630 vec![PageFragmentKey::new(file_id, 0, &range)],
2631 cache.find_index_candidates(file_id, 0, &range)
2632 );
2633
2634 let lookup = cache.lookup(file_id, 0, std::slice::from_ref(&range));
2635 assert!(lookup.is_fully_cached());
2636 assert_eq!(&vec![2; 10][..], &lookup.cached_parts[0][0].bytes[..]);
2637 }
2638
2639 #[test]
2640 fn test_page_cache_retains_disjoint_inserts_for_same_row_group() {
2641 let cache = PageRangeCache::new(1000);
2642 let file_id = FileId::random();
2643 let range1 = 0..10;
2644 let range2 = 20..30;
2645
2646 cache.insert_ranges(
2647 file_id,
2648 0,
2649 std::slice::from_ref(&range1),
2650 &[Bytes::from(vec![1; 10])],
2651 );
2652 cache.insert_ranges(
2653 file_id,
2654 0,
2655 std::slice::from_ref(&range2),
2656 &[Bytes::from(vec![2; 10])],
2657 );
2658
2659 let lookup = cache.lookup(file_id, 0, &[range1, range2]);
2660 assert!(lookup.is_fully_cached());
2661 assert_eq!(2, lookup.cached_range_count);
2662 assert_eq!(&vec![1; 10][..], &lookup.cached_parts[0][0].bytes[..]);
2663 assert_eq!(&vec![2; 10][..], &lookup.cached_parts[1][0].bytes[..]);
2664 }
2665
2666 #[test]
2667 fn test_page_cache_fragment_eviction() {
2668 let file_id = FileId::random();
2669 let range = 0..10;
2670 let key = PageFragmentKey::new(file_id, 0, &range);
2671 let page = Bytes::from(vec![1; 10]);
2672 let cache = PageRangeCache::new(page_cache_weight(&key, &page) as u64);
2673
2674 cache.insert_ranges(
2675 file_id,
2676 0,
2677 std::slice::from_ref(&range),
2678 &[Bytes::from(vec![1; 10])],
2679 );
2680 assert!(
2681 cache
2682 .lookup(file_id, 0, std::slice::from_ref(&range))
2683 .is_fully_cached()
2684 );
2685
2686 cache.cache.invalidate(&key);
2687 cache.cache.run_pending_tasks();
2688 assert!(cache.find_index_candidates(file_id, 0, &range).is_empty());
2689
2690 let lookup = cache.lookup(file_id, 0, std::slice::from_ref(&range));
2691 assert!(!lookup.is_fully_cached());
2692 assert_eq!(vec![0..10], lookup.missing_ranges);
2693 }
2694
2695 #[test]
2696 fn test_page_cache_rejects_oversized_fragment() {
2697 let cache = PageRangeCache::new(1);
2698 let file_id = FileId::random();
2699 let range = 0..10;
2700
2701 cache.insert_ranges(
2702 file_id,
2703 0,
2704 std::slice::from_ref(&range),
2705 &[Bytes::from(vec![1; 10])],
2706 );
2707 cache.cache.run_pending_tasks();
2708
2709 let lookup = cache.lookup(file_id, 0, std::slice::from_ref(&range));
2710 assert!(!lookup.is_fully_cached());
2711 assert_eq!(vec![0..10], lookup.missing_ranges);
2712 }
2713
2714 #[test]
2715 fn test_selector_result_cache() {
2716 let cache = CacheManager::builder()
2717 .selector_result_cache_size(1000)
2718 .build();
2719 let file_id = FileId::random();
2720 let key = SelectorResultKey {
2721 file_id,
2722 row_group_idx: 0,
2723 selector: TimeSeriesRowSelector::LastRow { after_merge: false },
2724 };
2725 assert!(cache.get_selector_result(&key).is_none());
2726 let result = Arc::new(SelectorResultValue::new(
2727 Vec::new(),
2728 ParquetReadColumns::from_deduped(Vec::new()),
2729 ));
2730 cache.put_selector_result(key, result);
2731 assert!(cache.get_selector_result(&key).is_some());
2732 }
2733
2734 #[test]
2735 fn test_prefilter_result_cache() {
2736 let disabled = CacheManager::builder().build();
2737 let file_id = FileId::random();
2738 let key = PrefilterKey::new(
2739 file_id,
2740 0,
2741 None,
2742 1,
2743 SmallVec::from_vec(vec!["tag_0 IN ([a])".to_string()]),
2744 );
2745 let selection = Arc::new(BooleanBuffer::new_set(3));
2746
2747 disabled.put_prefilter_result(key.clone(), selection.clone());
2748 assert!(disabled.get_prefilter_result(&key).is_none());
2749
2750 let cache = Arc::new(
2751 CacheManager::builder()
2752 .prefilter_result_cache_size(1000)
2753 .build(),
2754 );
2755 assert!(cache.get_prefilter_result(&key).is_none());
2756 cache.put_prefilter_result(key.clone(), selection.clone());
2757 assert_eq!(
2758 cache.get_prefilter_result(&key).unwrap().as_ref(),
2759 selection.as_ref()
2760 );
2761
2762 let enable_all = CacheStrategy::EnableAll(cache.clone());
2763 assert!(enable_all.get_prefilter_result(&key).is_some());
2764
2765 let compaction = CacheStrategy::Compaction(cache.clone());
2766 assert!(compaction.get_prefilter_result(&key).is_none());
2767 compaction.put_prefilter_result(key.clone(), selection.clone());
2768 assert!(cache.get_prefilter_result(&key).is_some());
2769
2770 let disabled_strategy = CacheStrategy::Disabled;
2771 assert!(disabled_strategy.get_prefilter_result(&key).is_none());
2772 disabled_strategy.put_prefilter_result(key.clone(), selection);
2773 assert!(cache.get_prefilter_result(&key).is_some());
2774 }
2775
2776 #[test]
2777 fn test_prefilter_key_distinguishes_dimensions() {
2778 let file_id = FileId::random();
2779 let row_selection = RowSelection::from(vec![RowSelector::skip(1), RowSelector::select(3)]);
2780 let other_row_selection =
2781 RowSelection::from(vec![RowSelector::skip(2), RowSelector::select(2)]);
2782 let row_selection = PrefilterKey::row_selection_snapshot(Some(&row_selection));
2783 let other_row_selection = PrefilterKey::row_selection_snapshot(Some(&other_row_selection));
2784 let base = PrefilterKey::new(
2785 file_id,
2786 0,
2787 row_selection.clone(),
2788 1,
2789 SmallVec::from_vec(vec!["tag_0 IN ([a])".to_string()]),
2790 );
2791
2792 assert_ne!(
2793 base,
2794 PrefilterKey::new(
2795 FileId::random(),
2796 0,
2797 row_selection.clone(),
2798 1,
2799 SmallVec::from_vec(vec!["tag_0 IN ([a])".to_string()])
2800 )
2801 );
2802 assert_ne!(
2803 base,
2804 PrefilterKey::new(
2805 file_id,
2806 1,
2807 row_selection.clone(),
2808 1,
2809 SmallVec::from_vec(vec!["tag_0 IN ([a])".to_string()])
2810 )
2811 );
2812 assert_ne!(
2813 base,
2814 PrefilterKey::new(
2815 file_id,
2816 0,
2817 other_row_selection,
2818 1,
2819 SmallVec::from_vec(vec!["tag_0 IN ([a])".to_string()])
2820 )
2821 );
2822 assert_ne!(
2823 base,
2824 PrefilterKey::new(
2825 file_id,
2826 0,
2827 row_selection.clone(),
2828 1,
2829 SmallVec::from_vec(vec!["tag_0 IN ([b])".to_string()])
2830 )
2831 );
2832 assert_ne!(
2833 base,
2834 PrefilterKey::new(
2835 file_id,
2836 0,
2837 row_selection.clone(),
2838 2,
2839 SmallVec::from_vec(vec!["tag_0 IN ([a])".to_string()])
2840 )
2841 );
2842 let pk_group = PrefilterKey::new(
2843 file_id,
2844 0,
2845 row_selection,
2846 1,
2847 SmallVec::from_vec(vec![
2848 "tag_0 IN ([a])".to_string(),
2849 "tag_1 IN ([x])".to_string(),
2850 ]),
2851 );
2852 assert_ne!(base, pk_group);
2853 }
2854
2855 #[test]
2856 fn test_range_result_cache() {
2857 let cache = Arc::new(
2858 CacheManager::builder()
2859 .range_result_cache_size(1024 * 1024)
2860 .build(),
2861 );
2862
2863 let key = RangeScanCacheKey {
2864 region_id: RegionId::new(1, 1),
2865 row_groups: vec![(FileId::random(), 0)],
2866 scan: ScanRequestFingerprintBuilder {
2867 read_columns: ReadColumns::new(std::iter::empty()),
2868 read_column_types: vec![],
2869 filters: vec!["tag_0 = 1".to_string()],
2870 time_filters: vec![],
2871 series_row_selector: None,
2872 append_mode: false,
2873 filter_deleted: true,
2874 merge_mode: crate::region::options::MergeMode::LastRow,
2875 sequence_range: None,
2876 partition_expr_version: 0,
2877 }
2878 .build(),
2879 };
2880 let value = Arc::new(RangeScanCacheValue::new(Vec::new(), 0));
2881
2882 assert!(cache.get_range_result(&key).is_none());
2883 cache.put_range_result(key.clone(), value.clone());
2884 assert!(cache.get_range_result(&key).is_some());
2885
2886 let enable_all = CacheStrategy::EnableAll(cache.clone());
2887 assert!(enable_all.get_range_result(&key).is_some());
2888
2889 let compaction = CacheStrategy::Compaction(cache.clone());
2890 assert!(compaction.get_range_result(&key).is_none());
2891 compaction.put_range_result(key.clone(), value.clone());
2892 assert!(cache.get_range_result(&key).is_some());
2893
2894 let disabled = CacheStrategy::Disabled;
2895 assert!(disabled.get_range_result(&key).is_none());
2896 disabled.put_range_result(key.clone(), value);
2897 assert!(cache.get_range_result(&key).is_some());
2898 }
2899
2900 #[test]
2901 fn test_range_result_cache_size_configures_limiter() {
2902 let cache_size = 3 * 1024_u64;
2903 let cache = CacheManager::builder()
2904 .range_result_cache_size(cache_size)
2905 .build();
2906
2907 assert_eq!(cache.range_result_cache_size(), cache_size as usize);
2908 assert_eq!(
2909 cache.range_result_memory_limiter().permit_bytes(),
2910 RANGE_RESULT_CONCAT_MEMORY_PERMIT.as_bytes() as usize
2911 );
2912 assert_eq!(
2913 cache.range_result_memory_limiter().available_permits(),
2914 (cache_size as usize).div_ceil(RANGE_RESULT_CONCAT_MEMORY_PERMIT.as_bytes() as usize)
2915 );
2916 }
2917
2918 #[tokio::test]
2919 async fn range_result_memory_limiter_rejects_oversized_request() {
2920 let limiter = RangeResultMemoryLimiter::new(2 * 1024, 1024);
2921 assert_eq!(limiter.available_permits(), 2);
2922
2923 let err = limiter.acquire(10 * 1024).await.unwrap_err();
2924 assert!(
2925 err.to_string().contains("exceeds limiter capacity"),
2926 "unexpected error: {err}"
2927 );
2928 assert_eq!(limiter.available_permits(), 2);
2929 }
2930
2931 #[tokio::test]
2932 async fn range_result_memory_limiter_allows_request_up_to_capacity() {
2933 let limiter = RangeResultMemoryLimiter::new(2 * 1024, 1024);
2934 let permit = limiter.acquire(2 * 1024).await.unwrap();
2935 assert_eq!(limiter.available_permits(), 0);
2936 drop(permit);
2937 assert_eq!(limiter.available_permits(), 2);
2938 }
2939
2940 #[rstest::rstest]
2941 #[case::empty_result(None, 0, Some(1))]
2942 #[case::new_result(None, 1025, Some(2))]
2943 #[case::unchanged(Some(1024), 1024, Some(1))]
2944 #[case::grow(Some(1024), 2048, Some(2))]
2945 #[case::shrink(Some(3072), 1024, Some(1))]
2946 #[case::exhausted(Some(1024), 4096, None)]
2947 #[case::oversized(Some(1024), usize::MAX, None)]
2948 #[test]
2949 fn range_result_memory_reservation_resizes_without_waiting(
2950 #[case] input_size: Option<usize>,
2951 #[case] result_size: usize,
2952 #[case] expected_permits: Option<usize>,
2953 ) {
2954 let limiter = RangeResultMemoryLimiter::new(4096, 1024);
2955 let blocker = limiter.try_acquire(1024).unwrap();
2956 let input_permit = input_size.map(|size| limiter.try_acquire(size).unwrap());
2957 let result_permit = limiter.try_resize(input_permit, result_size);
2958 assert_eq!(
2959 result_permit.as_ref().map(|permit| permit.num_permits()),
2960 expected_permits
2961 );
2962 assert_eq!(
2963 limiter.available_permits(),
2964 3 - expected_permits.unwrap_or(0)
2965 );
2966 drop(result_permit);
2967 assert_eq!(limiter.available_permits(), 3);
2968 drop(blocker);
2969 assert_eq!(limiter.available_permits(), 4);
2970 }
2971
2972 #[tokio::test]
2973 async fn test_evict_puffin_cache_clears_all_entries() {
2974 use std::collections::{BTreeMap, HashMap};
2975
2976 let cache = CacheManager::builder()
2977 .index_metadata_size(128)
2978 .index_content_size(128)
2979 .index_content_page_size(64)
2980 .index_result_cache_size(128)
2981 .puffin_metadata_size(128)
2982 .build();
2983 let cache = Arc::new(cache);
2984
2985 let region_id = RegionId::new(1, 1);
2986 let index_id = RegionIndexId::new(RegionFileId::new(region_id, FileId::random()), 0);
2987 let column_id: ColumnId = 1;
2988
2989 let bloom_cache = cache.bloom_filter_index_cache().unwrap().clone();
2990 let inverted_cache = cache.inverted_index_cache().unwrap().clone();
2991 let result_cache = cache.index_result_cache().unwrap();
2992 let puffin_metadata_cache = cache.puffin_metadata_cache().unwrap().clone();
2993
2994 let bloom_key = (
2995 index_id.file_id(),
2996 index_id.version,
2997 column_id,
2998 Tag::Skipping,
2999 );
3000 bloom_cache.put_metadata(bloom_key, Arc::new(BloomFilterMeta::default()));
3001 inverted_cache.put_metadata(
3002 (index_id.file_id(), index_id.version),
3003 Arc::new(InvertedIndexMetas::default()),
3004 );
3005 let predicate = PredicateKey::new_bloom(Arc::new(BTreeMap::new()));
3006 let selection = Arc::new(RowGroupSelection::default());
3007 result_cache.put(predicate.clone(), index_id.file_id(), selection);
3008 let file_id_str = index_id.to_string();
3009 let metadata = Arc::new(FileMetadata {
3010 blobs: Vec::new(),
3011 properties: HashMap::new(),
3012 });
3013 puffin_metadata_cache.put_metadata(file_id_str.clone(), metadata);
3014
3015 assert!(bloom_cache.get_metadata(bloom_key).is_some());
3016 assert!(
3017 inverted_cache
3018 .get_metadata((index_id.file_id(), index_id.version))
3019 .is_some()
3020 );
3021 assert!(result_cache.get(&predicate, index_id.file_id()).is_some());
3022 assert!(puffin_metadata_cache.get_metadata(&file_id_str).is_some());
3023
3024 cache.evict_puffin_cache(index_id).await;
3025
3026 assert!(bloom_cache.get_metadata(bloom_key).is_none());
3027 assert!(
3028 inverted_cache
3029 .get_metadata((index_id.file_id(), index_id.version))
3030 .is_none()
3031 );
3032 assert!(result_cache.get(&predicate, index_id.file_id()).is_none());
3033 assert!(puffin_metadata_cache.get_metadata(&file_id_str).is_none());
3034
3035 bloom_cache.put_metadata(bloom_key, Arc::new(BloomFilterMeta::default()));
3037 inverted_cache.put_metadata(
3038 (index_id.file_id(), index_id.version),
3039 Arc::new(InvertedIndexMetas::default()),
3040 );
3041 result_cache.put(
3042 predicate.clone(),
3043 index_id.file_id(),
3044 Arc::new(RowGroupSelection::default()),
3045 );
3046 puffin_metadata_cache.put_metadata(
3047 file_id_str.clone(),
3048 Arc::new(FileMetadata {
3049 blobs: Vec::new(),
3050 properties: HashMap::new(),
3051 }),
3052 );
3053
3054 let strategy = CacheStrategy::EnableAll(cache.clone());
3055 strategy.evict_puffin_cache(index_id).await;
3056
3057 assert!(bloom_cache.get_metadata(bloom_key).is_none());
3058 assert!(
3059 inverted_cache
3060 .get_metadata((index_id.file_id(), index_id.version))
3061 .is_none()
3062 );
3063 assert!(result_cache.get(&predicate, index_id.file_id()).is_none());
3064 assert!(puffin_metadata_cache.get_metadata(&file_id_str).is_none());
3065 }
3066
3067 fn wide_region_metadata(column_count: u32) -> RegionMetadata {
3068 let region_id = RegionId::new(1024, 7);
3069 let mut builder = RegionMetadataBuilder::new(region_id);
3070 let mut primary_key = Vec::new();
3071
3072 for column_id in 0..column_count {
3073 let semantic_type = if column_id < 32 {
3074 primary_key.push(column_id);
3075 SemanticType::Tag
3076 } else {
3077 SemanticType::Field
3078 };
3079 let mut column_schema = ColumnSchema::new(
3080 format!("wide_column_{column_id}"),
3081 ConcreteDataType::string_datatype(),
3082 true,
3083 );
3084 column_schema
3085 .mut_metadata()
3086 .insert(format!("cache_key_{column_id}"), "cache_value".repeat(4));
3087 builder.push_column_metadata(ColumnMetadata {
3088 column_schema,
3089 semantic_type,
3090 column_id,
3091 });
3092 }
3093
3094 builder.push_column_metadata(ColumnMetadata {
3095 column_schema: ColumnSchema::new(
3096 "ts",
3097 ConcreteDataType::timestamp_millisecond_datatype(),
3098 false,
3099 ),
3100 semantic_type: SemanticType::Timestamp,
3101 column_id: column_count,
3102 });
3103 builder.primary_key(primary_key);
3104
3105 builder.build().unwrap()
3106 }
3107}