Skip to main content

mito2/
cache.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Cache for the engine.
16
17pub(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
79/// Finds a typed shared error through the range-reader I/O wrappers.
80pub(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
96/// Metrics type key for sst meta.
97const SST_META_TYPE: &str = "sst_meta";
98/// Metrics type key for the optional decoded SST metadata acceleration tier.
99const SST_META_DECODED_TYPE: &str = "sst_meta_decoded";
100/// Metrics type key for vector.
101const VECTOR_TYPE: &str = "vector";
102/// Metrics type key for files on the local store.
103const FILE_TYPE: &str = "file";
104/// Metrics type key for index files (puffin) on the local store.
105const INDEX_TYPE: &str = "index";
106/// Metrics type key for selector result cache.
107const SELECTOR_RESULT_TYPE: &str = "selector_result";
108/// Metrics type key for range scan result cache.
109const RANGE_RESULT_TYPE: &str = "range_result";
110/// Metrics type key for prefilter result cache.
111const 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    /// Number of acquisitions that found too few permits and parked. Test-only
121    /// signal to synchronize with a caller waiting inside [Self::acquire].
122    #[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    /// Never wait while holding a partial load: several partial loads could
167    /// otherwise exhaust the budget and wait on each other indefinitely.
168    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        // Nothing awaits between this check and the parking below, so an observed
191        // increment means the caller is about to wait for the missing permits.
192        #[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/// Cached SST metadata combines the parquet footer with the decoded region metadata.
211///
212/// The cached parquet footer strips the `greptime:metadata` JSON payload and stores the decoded
213/// [RegionMetadata] separately so readers can skip repeated deserialization work.
214#[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/// Compact, authoritative form of one SST's metadata.
224///
225/// The entry contains a zstd-compressed, self-contained Parquet metadata stream. The decoded
226/// representation is held weakly to coalesce concurrent decodes without retaining memory outside
227/// the cache capacity.
228#[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/// Both forms produced by decoding metadata after a cache miss.
238#[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/// Result of decoding SST metadata and attempting to encode its compact cache entry.
251#[derive(Debug)]
252pub(crate) enum SstMetaPreparation {
253    /// Both cache representations are ready for admission.
254    Prepared(PreparedSstMeta),
255    /// The metadata is usable by the reader, but its compact cache encoding failed.
256    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
320/// Decodes SST metadata on the given blocking runtime without preparing a compact cache entry.
321pub(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
344/// Decodes SST metadata and attempts to encode both cache representations on the given blocking runtime.
345pub(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
366/// Synchronously prepares SST metadata. Callers must run this on a blocking runtime.
367pub(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    // Defensively discard column indexes supplied by external metadata producers. Mito only
386    // consumes offset indexes.
387    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    // `zstd::bulk::compress` allocates for the compression upper bound and leaves the excess
407    // capacity in its `Vec`. Convert through a boxed slice so the retained allocation matches the
408    // cache weight.
409    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        // Keep the previous JSON-byte floor and charge the decoded structures as well.
507        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    /// Returns the immutable parquet metadata size computed when it was decoded.
525    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
598// `parquet::arrow::arrow_reader::RowSelector` does not implement `Hash`, but
599// prefilter cache keys must hash the upstream row-selection snapshot. Keep a
600// local hashable mirror of the two fields that define selector semantics.
601// TODO(yingwen): Remove this mirror if upstream `RowSelector` implements `Hash`.
602impl 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/// Key for a cached prefilter result.
612#[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/// Cache strategies that may only enable a subset of caches.
695#[derive(Clone)]
696pub enum CacheStrategy {
697    /// Strategy for normal operations.
698    /// Doesn't disable any cache.
699    EnableAll(CacheManagerRef),
700    /// Strategy for compaction.
701    /// Disables some caches during compaction to avoid affecting queries.
702    /// Enables the write cache so that the compaction can read files cached
703    /// in the write cache and write the compacted files back to the write cache.
704    Compaction(CacheManagerRef),
705    /// Do not use any cache.
706    Disabled,
707}
708
709impl CacheStrategy {
710    /// Returns the runtime for CPU-bound SST metadata work for this request.
711    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    /// Returns whether the SST metadata cache is enabled for this strategy.
721    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    /// Shares the resolved source on success and failure, without sharing I/O accounting.
748    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    /// Observes memory residency for this caller before sharing decode or I/O work.
791    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    /// Resolves an observed metadata entry and updates this caller's cache metrics.
808    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    /// Calls [CacheManager::get_sst_meta_data_from_mem_cache()].
835    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    /// Calls [CacheManager::get_parquet_meta_data_from_mem_cache()].
851    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    /// Calls [CacheManager::put_compact_sst_meta()].
861    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    /// Calls [CacheManager::put_parquet_meta_data()].
875    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    /// Calls [CacheManager::get_prefilter_result()].
892    /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled].
893    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    /// Calls [CacheManager::put_prefilter_result()].
901    /// It does nothing if the strategy isn't [CacheStrategy::EnableAll].
902    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    /// Calls [CacheManager::remove_parquet_meta_data()].
909    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    /// Calls [CacheManager::get_repeated_vector()].
922    /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled].
923    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    /// Calls [CacheManager::put_repeated_vector()].
937    /// It does nothing if the strategy isn't [CacheStrategy::EnableAll].
938    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    /// Returns the Parquet byte cache for normal query reads.
945    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    /// Calls [CacheManager::evict_puffin_cache()].
953    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    /// Calls [CacheManager::get_selector_result()].
966    /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled].
967    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    /// Calls [CacheManager::put_selector_result()].
980    /// It does nothing if the strategy isn't [CacheStrategy::EnableAll].
981    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    /// Calls [CacheManager::get_range_result()].
992    /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled].
993    #[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    /// Calls [CacheManager::put_range_result()].
1005    /// It does nothing if the strategy isn't [CacheStrategy::EnableAll].
1006    #[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    /// Returns true if the range result cache is enabled.
1025    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    /// Calls [CacheManager::write_cache()].
1038    /// It returns None if the strategy is [CacheStrategy::Disabled].
1039    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    /// Calls [CacheManager::index_cache()].
1048    /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled].
1049    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    /// Calls [CacheManager::bloom_filter_index_cache()].
1057    /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled].
1058    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    /// Calls [CacheManager::puffin_metadata_cache()].
1066    /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled].
1067    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    /// Calls [CacheManager::index_result_cache()].
1075    /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled].
1076    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    /// Triggers download if the strategy is [CacheStrategy::EnableAll] and write cache is available.
1084    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/// Manages cached data for the engine.
1105///
1106/// All caches are disabled by default.
1107#[derive(Default)]
1108pub struct CacheManager {
1109    /// Cache for compact, authoritative SST metadata.
1110    sst_meta_cache: Option<SstMetaCache>,
1111    /// Unregisters compact metadata bytes when this manager is dropped.
1112    _sst_meta_size_registration: Option<CacheSizeRegistration>,
1113    /// Skips output metadata prewarming for caches used only by the current compaction.
1114    disable_sst_meta_cache_on_write: bool,
1115    /// Cache for decoded SST metadata, used only as an acceleration tier.
1116    sst_decoded_meta_cache: Option<SstDecodedMetaCache>,
1117    /// Cache for vectors.
1118    vector_cache: Option<VectorCache>,
1119    /// Cache for SST byte ranges.
1120    page_cache: Option<Arc<ParquetPageCache>>,
1121    /// A Cache for writing files to object stores.
1122    write_cache: Option<WriteCacheRef>,
1123    /// Cache for inverted index.
1124    inverted_index_cache: Option<InvertedIndexCacheRef>,
1125    /// Cache for bloom filter index.
1126    bloom_filter_index_cache: Option<BloomFilterIndexCacheRef>,
1127    /// Puffin metadata cache.
1128    puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
1129    /// Cache for time series selectors.
1130    selector_result_cache: Option<SelectorResultCache>,
1131    /// Cache for range scan outputs in flat format.
1132    range_result_cache: Option<RangeResultCacheRef>,
1133    /// Configured capacity for range scan outputs in flat format.
1134    #[cfg(test)]
1135    range_result_cache_size: u64,
1136    /// Cache for index result.
1137    index_result_cache: Option<IndexResultCache>,
1138    /// Cache for prefilter result.
1139    prefilter_result_cache: Option<PrefilterResultCache>,
1140}
1141
1142pub type CacheManagerRef = Arc<CacheManager>;
1143
1144impl CacheManager {
1145    /// Returns a builder to build the cache.
1146    pub fn builder() -> CacheManagerBuilder {
1147        CacheManagerBuilder::default()
1148    }
1149
1150    /// Gets fused SST metadata with metrics tracking for direct cache tests.
1151    #[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    /// Gets cached fused SST metadata from in-memory cache.
1248    /// This method does not perform I/O.
1249    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    /// Gets cached [ParquetMetaData] from in-memory cache.
1269    /// This method does not perform I/O.
1270    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    /// Puts fused SST metadata into the cache.
1280    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    /// Puts a compact SST metadata entry into the cache.
1293    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            // The key describes the loading request, while compact metadata may
1320            // already contain stronger page-index guarantees.
1321            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    /// Gets and decodes only the compact tier without promoting into the decoded cache.
1331    ///
1332    /// This is used by startup preloading so inspecting already-cached metadata cannot consume the
1333    /// decoded reservation and prematurely stop compact preloading.
1334    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    /// Puts [ParquetMetaData] into the cache.
1354    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    /// Removes [ParquetMetaData] from the cache.
1382    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    /// Returns whether the authoritative SST metadata tier has reached its reservation.
1401    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    /// Returns true if the in-memory SST meta cache is enabled.
1412    pub(crate) fn sst_meta_cache_enabled(&self) -> bool {
1413        self.sst_meta_cache.is_some()
1414    }
1415
1416    /// Whether newly written SST metadata should be prepared for later readers.
1417    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    /// Gets a vector with repeated value for specific `key`.
1422    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    /// Puts a vector with repeated value into the cache.
1434    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    /// Evicts every puffin-related cache entry for the given file.
1445    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            // Evict this exact generation across all storage scopes and paths.
1460            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    /// Gets result of for the selector.
1477    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    /// Puts result of the selector into the cache.
1487    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    /// Retires decoded range results when an SST leaves the current version.
1501    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    /// Gets cached result for range scan.
1508    #[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    /// Puts range scan result into cache.
1519    #[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    /// Returns true if the range result cache is enabled.
1531    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    /// Gets the write cache.
1541    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
1577/// Increases selector cache miss metrics.
1578pub fn selector_result_cache_miss() {
1579    CACHE_MISS.with_label_values(&[SELECTOR_RESULT_TYPE]).inc()
1580}
1581
1582/// Increases selector cache hit metrics.
1583pub fn selector_result_cache_hit() {
1584    CACHE_HIT.with_label_values(&[SELECTOR_RESULT_TYPE]).inc()
1585}
1586
1587/// Builder to construct a [CacheManager].
1588#[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    /// Sets meta cache size.
1607    pub fn sst_meta_cache_size(mut self, bytes: u64) -> Self {
1608        self.sst_meta_cache_size = bytes;
1609        self
1610    }
1611
1612    /// Keeps read-side metadata caching without prewarming newly written SSTs.
1613    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    /// Sets vector cache size.
1619    pub fn vector_cache_size(mut self, bytes: u64) -> Self {
1620        self.vector_cache_size = bytes;
1621        self
1622    }
1623
1624    /// Sets page cache size.
1625    pub fn page_cache_size(mut self, bytes: u64) -> Self {
1626        self.page_cache_size = bytes;
1627        self
1628    }
1629
1630    /// Sets write cache.
1631    pub fn write_cache(mut self, cache: Option<WriteCacheRef>) -> Self {
1632        self.write_cache = cache;
1633        self
1634    }
1635
1636    /// Sets cache size for index metadata.
1637    pub fn index_metadata_size(mut self, bytes: u64) -> Self {
1638        self.index_metadata_size = bytes;
1639        self
1640    }
1641
1642    /// Sets cache size for index content.
1643    pub fn index_content_size(mut self, bytes: u64) -> Self {
1644        self.index_content_size = bytes;
1645        self
1646    }
1647
1648    /// Sets page size for index content.
1649    pub fn index_content_page_size(mut self, bytes: u64) -> Self {
1650        self.index_content_page_size = bytes;
1651        self
1652    }
1653
1654    /// Sets cache size for index result.
1655    pub fn index_result_cache_size(mut self, bytes: u64) -> Self {
1656        self.index_result_cache_size = bytes;
1657        self
1658    }
1659
1660    /// Sets cache size for prefilter result.
1661    pub fn prefilter_result_cache_size(mut self, bytes: u64) -> Self {
1662        self.prefilter_result_cache_size = bytes;
1663        self
1664    }
1665
1666    /// Sets cache size for puffin metadata.
1667    pub fn puffin_metadata_size(mut self, bytes: u64) -> Self {
1668        self.puffin_metadata_size = bytes;
1669        self
1670    }
1671
1672    /// Sets selector result cache size.
1673    pub fn selector_result_cache_size(mut self, bytes: u64) -> Self {
1674        self.selector_result_cache_size = bytes;
1675        self
1676    }
1677
1678    /// Sets range result cache size.
1679    pub fn range_result_cache_size(mut self, bytes: u64) -> Self {
1680        self.range_result_cache_size = bytes;
1681        self
1682    }
1683
1684    /// Builds the [CacheManager].
1685    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        // Reserve half the configured capacity for the compact authoritative tier. This prevents
1694        // decoded acceleration entries from evicting metadata that would require storage I/O to
1695        // recover. Both reservations together equal the configured limit.
1696        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        // TODO(ruihang): check if it's ok to reuse the same param with inverted index
1749        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    // We ignore the size of `Arc`. Region metadata is already present in the compressed Parquet
1801    // stream and is materialized only in the decoded acceleration tier.
1802    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    // We ignore the heap size of `Value`.
1812    (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
1823/// Updates cache hit/miss metrics.
1824fn 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/// Cache key (region id, file id) for SST meta.
1834#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1835struct SstMetaKey(RegionId, FileId);
1836
1837impl SstMetaKey {
1838    /// Returns memory used by the key (estimated).
1839    fn estimated_size(&self) -> usize {
1840        mem::size_of::<Self>()
1841    }
1842}
1843
1844/// Cache key for time series row selector result.
1845#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1846pub struct SelectorResultKey {
1847    /// Id of the SST file.
1848    pub file_id: FileId,
1849    /// Index of the row group.
1850    pub row_group_idx: usize,
1851    /// Time series row selector.
1852    pub selector: TimeSeriesRowSelector,
1853}
1854
1855/// Result stored in the selector result cache.
1856pub enum SelectorResult {
1857    /// Batches in the primary key format.
1858    PrimaryKey(Vec<Batch>),
1859    /// Record batches in the flat format.
1860    Flat(Vec<RecordBatch>),
1861}
1862
1863/// Cached result for time series row selector.
1864pub struct SelectorResultValue {
1865    /// Batches of rows selected by the selector.
1866    pub result: SelectorResult,
1867    /// The read columns of rows.
1868    pub read_cols: ParquetReadColumns,
1869    /// JSON2 target types used by flat-format reads.
1870    ///
1871    /// JSON2 projection is query-driven; the same parquet columns can produce
1872    /// different cached batches under different type hints.
1873    pub json_target_types: JsonTargetTypes,
1874}
1875
1876impl SelectorResultValue {
1877    /// Creates a new selector result value with primary key format.
1878    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    /// Creates a new selector result value with flat format.
1887    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    /// Returns memory used by the value (estimated).
1900    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
1912/// A caller's memory observation. Holding the entry avoids a second lookup after coordination.
1913pub(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/// Resolved cache source, not physical I/O. Also returned when the subsequent load fails.
1936/// A caller that observed a memory miss can find a completed value at the native recheck;
1937/// it reuses that value's source without executing its initializer. This does not prove
1938/// that the caller waited for an in-flight load or performed a file-cache lookup.
1939/// Ordinary memory hits return before this resolution is applied.
1940/// A default value means no file-cache resolution was needed (or known).
1941#[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
1952/// Metadata returned by a shared load. Cache hits ignore the original load's miss flag.
1953struct SstMetaLoadResult {
1954    metadata: Arc<CachedSstMeta>,
1955    /// Preserves background-download eligibility for all requests sharing a storage load.
1956    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/// Separates loads that require different page-index guarantees.
1983#[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
1998/// Maps (region id, file id) to compact SST metadata.
1999type SstMetaCache = Cache<SstMetaKey, Arc<CompactSstMeta>>;
2000/// Maps (region id, file id) to decoded SST metadata retained as an acceleration tier.
2001type SstDecodedMetaCache = common_cache::Cache<SstDecodedMetaKey, Arc<SstMetaLoadResult>>;
2002/// Maps [Value] to a vector that holds this value repeatedly.
2003///
2004/// e.g. `"hello" => ["hello", "hello", "hello"]`
2005type VectorCache = Cache<(ConcreteDataType, Value), VectorRef>;
2006/// Maps (file id, row group id, time series row selector) to [SelectorResultValue].
2007type SelectorResultCache = Cache<SelectorResultKey, Arc<SelectorResultValue>>;
2008/// Maps partition-range scan key to cached flat batches.
2009type 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        // Moka applies capacity limits during maintenance, not synchronously on insertion.
2127        manager
2128            .sst_decoded_meta_cache
2129            .as_ref()
2130            .unwrap()
2131            .run_pending_tasks()
2132            .await;
2133        // An entry larger than the decoded reservation must not be retained.
2134        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            // Losing the compact copy must not force another decode or storage load.
2206            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        // Populate all policy keys before removing the file, not just the last inserted key.
2294        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        // A different policy or file does not wait for this load.
2334        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        // Repeated remote compactions must not leave byte counts after releasing their caches.
2491        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(&region_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        // Retain only the compact form so the lookup must exercise decompression and decoding.
2651        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        // Budget for the actual keys and empty values used below, including all
3190        // three footer entries. Shared scopes are conservatively charged per key.
3191        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        // Retirement only has the immutable index handle, so matching scopes are
3261        // all invalidated while a different generation remains available.
3262        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        // Refill caches and evict via CacheStrategy to ensure delegation works.
3304        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}