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