Skip to main content

mito2/
cache.rs

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