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