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