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