From c969fb0a5ea132a2a3ba761b6643960d8a4b9748 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Thu, 16 Jul 2026 19:20:13 +0800 Subject: [PATCH] perf: reduce parquet metadata cache footprint (#8527) * perf(mito2): reuse parquet metadata size accounting Signed-off-by: Ruihang Xia * perf(mito2): skip unused parquet column indexes Signed-off-by: Ruihang Xia * perf(mito2): cache compact parquet metadata Signed-off-by: Ruihang Xia * fix(mito2): tolerate metadata cache encoding failures Signed-off-by: Ruihang Xia * fix(mito2): skip disabled metadata cache encoding Signed-off-by: Ruihang Xia * refactor(mito2): avoid nested metadata cache task Signed-off-by: Ruihang Xia --------- Signed-off-by: Ruihang Xia --- Cargo.lock | 1 + src/mito2/Cargo.toml | 1 + src/mito2/src/access_layer.rs | 48 +- src/mito2/src/cache.rs | 661 +++++++++++++++++++++--- src/mito2/src/cache/file_cache.rs | 69 ++- src/mito2/src/cache/write_cache.rs | 3 +- src/mito2/src/region/opener.rs | 98 +++- src/mito2/src/sst/parquet/file_range.rs | 2 +- src/mito2/src/sst/parquet/metadata.rs | 53 +- src/mito2/src/sst/parquet/reader.rs | 47 +- 10 files changed, 836 insertions(+), 147 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 29640722ba1..87827905fe6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8519,6 +8519,7 @@ dependencies = [ "toml 0.8.23", "tracing", "uuid", + "zstd", ] [[package]] diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 7ae38da9f2a..e48a5758eab 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -93,6 +93,7 @@ tokio-stream.workspace = true tokio-util.workspace = true tracing.workspace = true uuid.workspace = true +zstd.workspace = true [dev-dependencies] common-function.workspace = true diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index b6adb6eb595..5622acacefa 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -16,11 +16,13 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use async_stream::try_stream; +use common_telemetry::warn; use common_time::Timestamp; use futures::{Stream, TryStreamExt}; use object_store::services::Fs; use object_store::util::{join_dir, with_instrument_layers}; use object_store::{ATOMIC_WRITE_DIR, ErrorKind, OLD_ATOMIC_WRITE_DIR, ObjectStore}; +use parquet::file::metadata::PageIndexPolicy; use smallvec::SmallVec; use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; @@ -28,9 +30,9 @@ use store_api::region_request::PathType; use store_api::sst_entry::StorageSstEntry; use store_api::storage::{FileId, RegionId, SequenceNumber}; -use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; use crate::cache::write_cache::SstUploadRequest; +use crate::cache::{CacheManagerRef, SstMetaPreparation, prepare_sst_meta_sync}; use crate::config::{BloomFilterConfig, FulltextIndexConfig, IndexConfig, InvertedIndexConfig}; use crate::error::{ CleanDirSnafu, DeleteIndexSnafu, DeleteIndexesSnafu, DeleteSstsSnafu, OpenDalSnafu, Result, @@ -410,14 +412,46 @@ impl AccessLayer { }; // Put parquet metadata to cache manager. - if !sst_info.is_empty() { + if !sst_info.is_empty() && cache_manager.sst_meta_cache_enabled() { for sst in &sst_info { if let Some(parquet_metadata) = &sst.file_metadata { - cache_manager.put_parquet_meta_data( - RegionFileId::new(region_id, sst.file_id), - parquet_metadata.clone(), - Some(region_metadata.clone()), - ) + let file_id = RegionFileId::new(region_id, sst.file_id); + let file_path = format!( + "region_id={}, file_id={}", + file_id.region_id(), + file_id.file_id() + ); + let page_index_policy = if parquet_metadata.offset_index().is_some() { + PageIndexPolicy::Optional + } else { + PageIndexPolicy::Skip + }; + let parquet_metadata = parquet_metadata.clone(); + let region_metadata = region_metadata.clone(); + let cache_manager = cache_manager.clone(); + // Compact cache preparation is best-effort. Run the entire operation in one + // detached blocking task so it neither blocks an async worker nor delays the + // SST write. + common_runtime::spawn_blocking_global(move || { + match prepare_sst_meta_sync( + &file_path, + Arc::unwrap_or_clone(parquet_metadata), + Some(region_metadata), + page_index_policy, + ) { + Ok(SstMetaPreparation::Prepared(metadata)) => { + cache_manager.put_prepared_sst_meta(file_id, metadata, true); + } + Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( + encoding_error; + "Failed to encode parquet metadata for cache, file: {}", + file_path + ), + Err(err) => { + warn!(err; "Failed to cache parquet metadata for {}", file_path); + } + } + }); } } } diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 5915aa32083..27db7daa5ab 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -26,10 +26,11 @@ pub(crate) mod write_cache; use std::collections::{BTreeMap, HashMap}; use std::mem; use std::ops::Range; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, RwLock, Weak}; use bytes::Bytes; use common_base::readable_size::ReadableSize; +use common_datasource::compression::CompressionType; use common_telemetry::warn; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::record_batch::RecordBatch; @@ -41,11 +42,13 @@ use moka::notification::RemovalCause; use moka::sync::Cache; use object_store::ObjectStore; use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; -use parquet::file::metadata::{FileMetaData, PageIndexPolicy, ParquetMetaData}; +use parquet::file::metadata::{ + FileMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, ParquetMetaDataWriter, +}; use puffin::puffin_manager::cache::{PuffinMetadataCache, PuffinMetadataCacheRef}; use smallvec::SmallVec; use snafu::{OptionExt, ResultExt}; -use store_api::metadata::RegionMetadataRef; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::{ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; use crate::cache::cache_size::parquet_meta_size; @@ -54,7 +57,10 @@ use crate::cache::index::inverted_index::{InvertedIndexCache, InvertedIndexCache #[cfg(feature = "vector_index")] use crate::cache::index::vector_index::{VectorIndexCache, VectorIndexCacheRef}; use crate::cache::write_cache::WriteCacheRef; -use crate::error::{InvalidMetadataSnafu, InvalidParquetSnafu, Result, UnexpectedSnafu}; +use crate::error::{ + CompressObjectSnafu, DecompressObjectSnafu, InvalidMetadataSnafu, InvalidParquetSnafu, + JoinSnafu, ReadParquetSnafu, Result, UnexpectedSnafu, WriteParquetSnafu, +}; use crate::memtable::record_batch_estimated_size; use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS}; use crate::read::Batch; @@ -66,6 +72,8 @@ use crate::sst::parquet::reader::MetadataCacheMetrics; /// Metrics type key for sst meta. const SST_META_TYPE: &str = "sst_meta"; +/// Metrics type key for the optional decoded SST metadata acceleration tier. +const SST_META_DECODED_TYPE: &str = "sst_meta_decoded"; /// Metrics type key for vector. const VECTOR_TYPE: &str = "vector"; /// Metrics type key for pages. @@ -152,11 +160,235 @@ impl RangeResultMemoryLimiter { #[derive(Debug)] pub(crate) struct CachedSstMeta { parquet_metadata: Arc, + parquet_metadata_size: usize, region_metadata: RegionMetadataRef, region_metadata_weight: usize, page_index_policy: PageIndexPolicy, } +/// Compact, authoritative form of one SST's metadata. +/// +/// The entry contains a zstd-compressed, self-contained Parquet metadata stream. The decoded +/// representation is held weakly to coalesce concurrent decodes without retaining memory outside +/// the cache capacity. +#[derive(Debug)] +pub(crate) struct CompactSstMeta { + encoded_metadata: Bytes, + decoded_size: usize, + region_metadata: Weak, + page_index_policy: PageIndexPolicy, + decoded: tokio::sync::Mutex>, +} + +/// Both forms produced by decoding metadata after a cache miss. +#[derive(Debug)] +pub(crate) struct PreparedSstMeta { + compact: Arc, + decoded: Arc, +} + +impl PreparedSstMeta { + pub(crate) fn decoded(&self) -> Arc { + self.decoded.clone() + } +} + +/// Result of decoding SST metadata and attempting to encode its compact cache entry. +#[derive(Debug)] +pub(crate) enum SstMetaPreparation { + /// Both cache representations are ready for admission. + Prepared(PreparedSstMeta), + /// The metadata is usable by the reader, but its compact cache encoding failed. + DecodedOnly { + decoded: Arc, + encoding_error: crate::error::Error, + }, +} + +impl SstMetaPreparation { + pub(crate) fn decoded(&self) -> Arc { + match self { + Self::Prepared(metadata) => metadata.decoded(), + Self::DecodedOnly { decoded, .. } => decoded.clone(), + } + } +} + +impl CompactSstMeta { + async fn decode(&self) -> Result> { + let mut decoded_guard = self.decoded.lock().await; + if let Some(decoded) = decoded_guard.upgrade() { + return Ok(decoded); + } + + let encoded_metadata = self.encoded_metadata.clone(); + let decoded_size = self.decoded_size; + let region_metadata = self.region_metadata.upgrade(); + let page_index_policy = self.page_index_policy; + let decoded = common_runtime::spawn_blocking_global(move || { + let bytes = zstd::bulk::decompress(&encoded_metadata, decoded_size).context( + DecompressObjectSnafu { + compress_type: CompressionType::Zstd, + path: "cached SST metadata", + }, + )?; + let bytes = Bytes::from(bytes); + let mut reader = ParquetMetaDataReader::new() + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(page_index_policy); + reader.try_parse(&bytes).context(ReadParquetSnafu { + path: "cached SST metadata", + })?; + let metadata = reader.finish().context(ReadParquetSnafu { + path: "cached SST metadata", + })?; + CachedSstMeta::try_new_with_page_index_policy( + "cached SST metadata", + metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)??; + + *decoded_guard = Arc::downgrade(&decoded); + Ok(decoded) + } + + fn satisfies_page_index_policy(&self, requested: PageIndexPolicy) -> bool { + satisfies_page_index_policy(self.page_index_policy, requested) + } +} + +/// Decodes SST metadata without preparing a compact cache entry. +pub(crate) async fn decode_sst_meta( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, +) -> Result> { + let file_path = file_path.to_string(); + common_runtime::spawn_blocking_global(move || { + let parquet_metadata = strip_column_indexes(parquet_metadata); + CachedSstMeta::try_new_with_page_index_policy( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)? +} + +/// Decodes SST metadata and attempts to encode both cache representations on the blocking runtime. +pub(crate) async fn prepare_sst_meta( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, +) -> Result { + let file_path = file_path.to_string(); + common_runtime::spawn_blocking_global(move || { + prepare_sst_meta_sync( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + }) + .await + .context(JoinSnafu)? +} + +/// Synchronously prepares SST metadata. Callers must run this on a blocking runtime. +pub(crate) fn prepare_sst_meta_sync( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, +) -> Result { + let parquet_metadata = strip_column_indexes(parquet_metadata); + let cache_encoding = encode_compact_sst_meta(file_path, &parquet_metadata); + finish_sst_meta_preparation( + file_path, + parquet_metadata, + region_metadata, + page_index_policy, + cache_encoding, + ) +} + +fn strip_column_indexes(parquet_metadata: ParquetMetaData) -> ParquetMetaData { + // Defensively discard column indexes supplied by external metadata producers. Mito only + // consumes offset indexes. + let mut builder = parquet_metadata.into_builder(); + builder.take_column_index(); + builder.build() +} + +fn encode_compact_sst_meta( + file_path: &str, + parquet_metadata: &ParquetMetaData, +) -> Result<(Bytes, usize)> { + let mut encoded = Vec::new(); + ParquetMetaDataWriter::new(&mut encoded, parquet_metadata) + .finish() + .context(WriteParquetSnafu)?; + let decoded_size = encoded.len(); + let encoded_metadata = zstd::bulk::compress(&encoded, 3).context(CompressObjectSnafu { + compress_type: CompressionType::Zstd, + path: file_path, + })?; + + // `zstd::bulk::compress` allocates for the compression upper bound and leaves the excess + // capacity in its `Vec`. Convert through a boxed slice so the retained allocation matches the + // cache weight. + Ok(( + Bytes::from(encoded_metadata.into_boxed_slice()), + decoded_size, + )) +} + +fn finish_sst_meta_preparation( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, + cache_encoding: Result<(Bytes, usize)>, +) -> Result { + let decoded = Arc::new(CachedSstMeta::try_new_with_page_index_policy( + file_path, + parquet_metadata, + region_metadata, + page_index_policy, + )?); + let (encoded_metadata, decoded_size) = match cache_encoding { + Ok(encoded) => encoded, + Err(encoding_error) => { + return Ok(SstMetaPreparation::DecodedOnly { + decoded, + encoding_error, + }); + } + }; + let compact = Arc::new(CompactSstMeta { + encoded_metadata, + decoded_size, + region_metadata: Arc::downgrade(&decoded.region_metadata), + page_index_policy, + decoded: tokio::sync::Mutex::new(Arc::downgrade(&decoded)), + }); + + Ok(SstMetaPreparation::Prepared(PreparedSstMeta { + compact, + decoded, + })) +} + impl CachedSstMeta { #[cfg(test)] pub(crate) fn try_new(file_path: &str, parquet_metadata: ParquetMetaData) -> Result { @@ -215,9 +447,11 @@ impl CachedSstMeta { // Keep the previous JSON-byte floor and charge the decoded structures as well. let region_metadata_weight = region_metadata.estimated_size().max(json.len()); let parquet_metadata = Arc::new(strip_region_metadata_from_parquet(parquet_metadata)); + let parquet_metadata_size = parquet_meta_size(&parquet_metadata); Ok(Self { parquet_metadata, + parquet_metadata_size, region_metadata, region_metadata_weight, page_index_policy, @@ -228,21 +462,30 @@ impl CachedSstMeta { self.parquet_metadata.clone() } + /// Returns the immutable parquet metadata size computed when it was decoded. + pub(crate) fn parquet_metadata_size(&self) -> usize { + self.parquet_metadata_size + } + pub(crate) fn region_metadata(&self) -> RegionMetadataRef { self.region_metadata.clone() } fn satisfies_page_index_policy(&self, requested: PageIndexPolicy) -> bool { - match requested { - PageIndexPolicy::Skip => true, - PageIndexPolicy::Optional => self.page_index_policy != PageIndexPolicy::Skip, - PageIndexPolicy::Required => self.page_index_policy == PageIndexPolicy::Required, - } + satisfies_page_index_policy(self.page_index_policy, requested) + } +} + +fn satisfies_page_index_policy(cached: PageIndexPolicy, requested: PageIndexPolicy) -> bool { + match requested { + PageIndexPolicy::Skip => true, + PageIndexPolicy::Optional => cached != PageIndexPolicy::Skip, + PageIndexPolicy::Required => cached == PageIndexPolicy::Required, } } fn infer_loaded_page_index_policy(parquet_metadata: &ParquetMetaData) -> PageIndexPolicy { - if parquet_metadata.column_index().is_some() || parquet_metadata.offset_index().is_some() { + if parquet_metadata.offset_index().is_some() { PageIndexPolicy::Optional } else { PageIndexPolicy::Skip @@ -270,12 +513,10 @@ fn strip_region_metadata_from_parquet(parquet_metadata: ParquetMetaData) -> Parq let mut builder = parquet_metadata.into_builder(); let row_groups = builder.take_row_groups(); - let column_index = builder.take_column_index(); let offset_index = builder.take_offset_index(); parquet::file::metadata::ParquetMetaDataBuilder::new(stripped_file_metadata) .set_row_groups(row_groups) - .set_column_index(column_index) .set_offset_index(offset_index) .build() } @@ -407,6 +648,16 @@ pub enum CacheStrategy { } impl CacheStrategy { + /// Returns whether the SST metadata cache is enabled for this strategy. + pub(crate) fn sst_meta_cache_enabled(&self) -> bool { + match self { + CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { + cache_manager.sst_meta_cache_enabled() + } + CacheStrategy::Disabled => false, + } + } + /// Gets fused SST metadata with cache metrics tracking. pub(crate) async fn get_sst_meta_data( &self, @@ -450,11 +701,16 @@ impl CacheStrategy { .map(|metadata| metadata.parquet_metadata()) } - /// Calls [CacheManager::put_sst_meta_data()]. - pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { + /// Puts compact and decoded forms of SST metadata into the shared cache capacity. + pub(crate) fn put_prepared_sst_meta( + &self, + file_id: RegionFileId, + metadata: PreparedSstMeta, + retain_decoded: bool, + ) { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager.put_sst_meta_data(file_id, metadata); + cache_manager.put_prepared_sst_meta(file_id, metadata, retain_decoded); } CacheStrategy::Disabled => {} } @@ -730,8 +986,10 @@ impl CacheStrategy { /// All caches are disabled by default. #[derive(Default)] pub struct CacheManager { - /// Cache for SST metadata. + /// Cache for compact, authoritative SST metadata. sst_meta_cache: Option, + /// Cache for decoded SST metadata, used only as an acceleration tier. + sst_decoded_meta_cache: Option, /// Cache for vectors. vector_cache: Option, /// Cache for SST byte ranges. @@ -777,40 +1035,82 @@ impl CacheManager { metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, ) -> Option> { - if let Some(metadata) = self.get_sst_meta_data_from_mem_cache(file_id, page_index_policy) { + let cache_key = SstMetaKey(file_id.region_id(), file_id.file_id()); + let compact = self + .get_compact_sst_meta(&cache_key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + let decoded = self + .get_decoded_sst_meta(&cache_key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + + if let Some(compact) = compact { + CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); metrics.mem_cache_hit += 1; - return Some(metadata); + if let Some(decoded) = decoded { + CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); + return Some(decoded); + } + + CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc(); + match compact.decode().await { + Ok(decoded) => { + self.put_sst_meta_data(file_id, decoded.clone()); + return Some(decoded); + } + Err(err) => { + warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); + self.remove_parquet_meta_data(file_id); + } + } + } else if let Some(decoded) = decoded { + // Metadata produced by a new SST writer can enter the decoded tier before its compact + // representation is prepared. + CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); + CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); + metrics.mem_cache_hit += 1; + return Some(decoded); + } else { + CACHE_MISS.with_label_values(&[SST_META_TYPE]).inc(); } let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); - if let Some(write_cache) = &self.write_cache - && let Some(metadata) = write_cache - .file_cache() - .get_sst_meta_data(key, metrics, page_index_policy) + if let Some(write_cache) = &self.write_cache { + let file_cache = write_cache.file_cache(); + if self.sst_meta_cache_enabled() { + if let Some(metadata) = file_cache + .get_sst_meta_data(key, metrics, page_index_policy) + .await + { + metrics.file_cache_hit += 1; + let decoded = metadata.decoded(); + match metadata { + SstMetaPreparation::Prepared(metadata) => { + self.put_prepared_sst_meta(file_id, metadata, true); + } + SstMetaPreparation::DecodedOnly { encoding_error, .. } => { + warn!( + encoding_error; + "Failed to encode file-cached SST metadata for memory cache, region_id: {}, file_id: {}", + file_id.region_id(), + file_id.file_id() + ); + } + } + return Some(decoded); + } + } else if let Some(decoded) = file_cache + .get_decoded_sst_meta_data(key, metrics, page_index_policy) .await - { - metrics.file_cache_hit += 1; - self.put_sst_meta_data(file_id, metadata.clone()); - return Some(metadata); + { + metrics.file_cache_hit += 1; + return Some(decoded); + } } metrics.cache_miss += 1; None } - /// Gets cached [ParquetMetaData] with metrics tracking. - /// Tries in-memory cache first, then file cache, updating metrics accordingly. - pub(crate) async fn get_parquet_meta_data( - &self, - file_id: RegionFileId, - metrics: &mut MetadataCacheMetrics, - page_index_policy: PageIndexPolicy, - ) -> Option> { - self.get_sst_meta_data(file_id, metrics, page_index_policy) - .await - .map(|metadata| metadata.parquet_metadata()) - } - /// Gets cached fused SST metadata from in-memory cache. /// This method does not perform I/O. pub(crate) fn get_sst_meta_data_from_mem_cache( @@ -818,12 +1118,11 @@ impl CacheManager { file_id: RegionFileId, page_index_policy: PageIndexPolicy, ) -> Option> { - self.sst_meta_cache.as_ref().and_then(|sst_meta_cache| { - let value = sst_meta_cache.get(&SstMetaKey(file_id.region_id(), file_id.file_id())); - let value = - value.filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); - update_hit_miss(value, SST_META_TYPE) - }) + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); + let value = self + .get_decoded_sst_meta(&key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + update_hit_miss(value, SST_META_DECODED_TYPE) } /// Gets cached [ParquetMetaData] from in-memory cache. @@ -838,15 +1137,65 @@ impl CacheManager { /// Puts fused SST metadata into the cache. pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { - if let Some(cache) = &self.sst_meta_cache { + if let Some(cache) = &self.sst_decoded_meta_cache { let key = SstMetaKey(file_id.region_id(), file_id.file_id()); CACHE_BYTES - .with_label_values(&[SST_META_TYPE]) - .add(meta_cache_weight(&key, &metadata).into()); + .with_label_values(&[SST_META_DECODED_TYPE]) + .add(decoded_meta_cache_weight(&key, &metadata).into()); cache.insert(key, metadata); } } + /// Puts a compact metadata entry and optionally retains its decoded acceleration entry. + pub(crate) fn put_prepared_sst_meta( + &self, + file_id: RegionFileId, + metadata: PreparedSstMeta, + retain_decoded: bool, + ) { + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); + if let Some(cache) = &self.sst_meta_cache { + CACHE_BYTES + .with_label_values(&[SST_META_TYPE]) + .add(meta_cache_weight(&key, &metadata.compact).into()); + cache.insert(key.clone(), metadata.compact); + } + if retain_decoded { + self.put_sst_meta_data(file_id, metadata.decoded); + } + } + + fn get_compact_sst_meta(&self, key: &SstMetaKey) -> Option> { + self.sst_meta_cache.as_ref()?.get(key) + } + + fn get_decoded_sst_meta(&self, key: &SstMetaKey) -> Option> { + self.sst_decoded_meta_cache.as_ref()?.get(key) + } + + /// Gets and decodes only the compact tier without promoting into the decoded cache. + /// + /// This is used by startup preloading so inspecting already-cached metadata cannot consume the + /// decoded reservation and prematurely stop compact preloading. + pub(crate) async fn get_compact_sst_meta_data( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + ) -> Option> { + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); + let compact = self + .get_compact_sst_meta(&key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy))?; + match compact.decode().await { + Ok(metadata) => Some(metadata), + Err(err) => { + warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); + self.remove_parquet_meta_data(file_id); + None + } + } + } + /// Puts [ParquetMetaData] into the cache. pub fn put_parquet_meta_data( &self, @@ -854,7 +1203,7 @@ impl CacheManager { metadata: Arc, region_metadata: Option, ) { - if self.sst_meta_cache.is_some() { + if self.sst_decoded_meta_cache.is_some() { let file_path = format!( "region_id={}, file_id={}", file_id.region_id(), @@ -877,17 +1226,24 @@ impl CacheManager { /// Removes [ParquetMetaData] from the cache. pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) { + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); if let Some(cache) = &self.sst_meta_cache { - cache.remove(&SstMetaKey(file_id.region_id(), file_id.file_id())); + cache.remove(&key); + } + if let Some(cache) = &self.sst_decoded_meta_cache { + cache.remove(&key); } } - /// Returns the total weighted size of the in-memory SST meta cache. - pub(crate) fn sst_meta_cache_weighted_size(&self) -> u64 { - self.sst_meta_cache - .as_ref() - .map(|cache| cache.weighted_size()) - .unwrap_or(0) + /// Returns whether the authoritative SST metadata tier has reached its reservation. + pub(crate) fn sst_meta_cache_is_full(&self) -> bool { + let Some(cache) = &self.sst_meta_cache else { + return true; + }; + cache + .policy() + .max_capacity() + .is_some_and(|capacity| cache.weighted_size() >= capacity) } /// Returns true if the in-memory SST meta cache is enabled. @@ -1190,9 +1546,14 @@ impl CacheManagerBuilder { /// Builds the [CacheManager]. pub fn build(self) -> CacheManager { - let sst_meta_cache = (self.sst_meta_cache_size != 0).then(|| { + // Reserve half the configured capacity for the compact authoritative tier. This prevents + // decoded acceleration entries from evicting metadata that would require storage I/O to + // recover. Both reservations together equal the configured limit. + let compact_meta_capacity = self.sst_meta_cache_size.div_ceil(2); + let decoded_meta_capacity = self.sst_meta_cache_size / 2; + let sst_meta_cache = (compact_meta_capacity != 0).then(|| { Cache::builder() - .max_capacity(self.sst_meta_cache_size) + .max_capacity(compact_meta_capacity) .weigher(meta_cache_weight) .eviction_listener(|k, v, cause| { let size = meta_cache_weight(&k, &v); @@ -1205,6 +1566,21 @@ impl CacheManagerBuilder { }) .build() }); + let sst_decoded_meta_cache = (decoded_meta_capacity != 0).then(|| { + Cache::builder() + .max_capacity(decoded_meta_capacity) + .weigher(decoded_meta_cache_weight) + .eviction_listener(|k, v, cause| { + let size = decoded_meta_cache_weight(&k, &v); + CACHE_BYTES + .with_label_values(&[SST_META_DECODED_TYPE]) + .sub(size.into()); + CACHE_EVICTION + .with_label_values(&[SST_META_DECODED_TYPE, removal_cause_str(cause)]) + .inc(); + }) + .build() + }); let vector_cache = (self.vector_cache_size != 0).then(|| { Cache::builder() .max_capacity(self.vector_cache_size) @@ -1274,6 +1650,7 @@ impl CacheManagerBuilder { }); CacheManager { sst_meta_cache, + sst_decoded_meta_cache, vector_cache, page_cache, write_cache: self.write_cache, @@ -1295,10 +1672,15 @@ impl CacheManagerBuilder { } } -fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { - // We ignore the size of `Arc`. - let size = - k.estimated_size() + parquet_meta_size(&v.parquet_metadata) + v.region_metadata_weight; +fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { + // We ignore the size of `Arc`. Region metadata is already present in the compressed Parquet + // stream and is materialized only in the decoded acceleration tier. + let size = k.estimated_size() + mem::size_of::() + v.encoded_metadata.len(); + u32::try_from(size).unwrap_or(u32::MAX) +} + +fn decoded_meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { + let size = k.estimated_size() + v.parquet_metadata_size + v.region_metadata_weight; u32::try_from(size).unwrap_or(u32::MAX) } @@ -1677,7 +2059,9 @@ impl SelectorResultValue { } /// Maps (region id, file id) to fused SST metadata. -type SstMetaCache = Cache>; +type SstMetaCache = Cache>; +/// Maps (region id, file id) to decoded SST metadata retained as an acceleration tier. +type SstDecodedMetaCache = Cache>; /// Maps [Value] to a vector that holds this value repeatedly. /// /// e.g. `"hello" => ["hello", "hello", "hello"]` @@ -1695,6 +2079,7 @@ mod tests { use api::v1::index::{BloomFilterMeta, InvertedIndexMetas}; use datatypes::schema::ColumnSchema; use datatypes::vectors::Int64Vector; + use parquet::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation}; use puffin::file_metadata::FileMetadata; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::storage::ColumnId; @@ -1715,6 +2100,7 @@ mod tests { async fn test_disable_cache() { let cache = CacheManager::default(); assert!(cache.sst_meta_cache.is_none()); + assert!(cache.sst_decoded_meta_cache.is_none()); assert!(cache.vector_cache.is_none()); assert!(cache.page_cache.is_none()); @@ -1725,7 +2111,7 @@ mod tests { cache.put_parquet_meta_data(file_id, metadata, None); assert!( cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data(file_id, &mut metrics, Default::default()) .await .is_none() ); @@ -1754,6 +2140,30 @@ mod tests { assert!(cache.write_cache().is_none()); } + #[test] + fn test_sst_meta_cache_splits_capacity() { + let cache = CacheManager::builder().sst_meta_cache_size(101).build(); + + assert_eq!( + Some(51), + cache + .sst_meta_cache + .as_ref() + .unwrap() + .policy() + .max_capacity() + ); + assert_eq!( + Some(50), + cache + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .policy() + .max_capacity() + ); + } + #[tokio::test] async fn test_parquet_meta_cache() { let cache = CacheManager::builder().sst_meta_cache_size(2000).build(); @@ -1762,7 +2172,7 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); assert!( cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data(file_id, &mut metrics, Default::default()) .await .is_none() ); @@ -1787,7 +2197,7 @@ mod tests { cache.remove_parquet_meta_data(file_id); assert!( cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data(file_id, &mut metrics, Default::default()) .await .is_none() ); @@ -1810,6 +2220,109 @@ mod tests { assert!(Arc::ptr_eq(®ion_metadata, &cached.region_metadata())); } + #[tokio::test] + async fn test_compact_sst_meta_round_trip() { + let cache = CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(); + let region_id = RegionId::new(1, 1); + let file_id = RegionFileId::new(region_id, FileId::random()); + let (metadata, expected_region_metadata) = sst_parquet_meta(); + let metadata = Arc::unwrap_or_clone(metadata); + let offset_indexes = metadata + .row_groups() + .iter() + .map(|row_group| { + row_group + .columns() + .iter() + .map(|_| OffsetIndexMetaData { + page_locations: vec![PageLocation { + offset: 42, + compressed_page_size: 128, + first_row_index: 0, + }], + unencoded_byte_array_data_bytes: None, + }) + .collect() + }) + .collect(); + let metadata = metadata + .into_builder() + .set_offset_index(Some(offset_indexes)) + .build(); + let expected_rows = metadata.file_metadata().num_rows(); + let prepared = prepare_sst_meta("test.parquet", metadata, None, PageIndexPolicy::Required) + .await + .unwrap(); + let SstMetaPreparation::Prepared(prepared) = prepared else { + panic!("valid metadata should produce a compact cache entry"); + }; + + // Retain only the compact form so the lookup must exercise decompression and decoding. + cache.put_prepared_sst_meta(file_id, prepared, false); + let mut metrics = MetadataCacheMetrics::default(); + let cached = cache + .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Required) + .await + .unwrap(); + + assert_eq!(1, metrics.mem_cache_hit); + assert_eq!(0, metrics.cache_miss); + assert_eq!( + expected_rows, + cached.parquet_metadata.file_metadata().num_rows() + ); + assert!(cached.parquet_metadata.offset_index().is_some()); + assert_eq!(expected_region_metadata, cached.region_metadata()); + assert!( + cached + .parquet_metadata + .file_metadata() + .key_value_metadata() + .is_none_or(|key_values| key_values + .iter() + .all(|key_value| key_value.key != PARQUET_METADATA_KEY)) + ); + } + + #[test] + fn test_cache_encoding_failure_preserves_decoded_metadata() { + let (metadata, expected_region_metadata) = sst_parquet_meta(); + let expected_rows = metadata.file_metadata().num_rows(); + let cache_encoding: Result<(Bytes, usize)> = Err(UnexpectedSnafu { + reason: "injected compact metadata encoding failure", + } + .build()); + + let preparation = finish_sst_meta_preparation( + "test.parquet", + Arc::unwrap_or_clone(metadata), + None, + PageIndexPolicy::Skip, + cache_encoding, + ) + .unwrap(); + let SstMetaPreparation::DecodedOnly { + decoded, + encoding_error, + } = preparation + else { + panic!("cache encoding failure should return decoded-only metadata"); + }; + + assert_eq!( + expected_rows, + decoded.parquet_metadata.file_metadata().num_rows() + ); + assert_eq!(expected_region_metadata, decoded.region_metadata()); + assert!( + encoding_error + .to_string() + .contains("injected compact metadata encoding failure") + ); + } + #[tokio::test] async fn test_parquet_meta_cache_respects_page_index_policy() { let cache = CacheManager::builder().sst_meta_cache_size(2000).build(); @@ -1868,7 +2381,7 @@ mod tests { } #[test] - fn test_meta_cache_weight_accounts_for_decoded_region_metadata() { + fn test_decoded_meta_cache_weight_accounts_for_region_metadata() { let region_metadata = Arc::new(wide_region_metadata(128)); let json_len = region_metadata.to_json().unwrap().len(); let metadata = sst_parquet_meta_with_region_metadata(region_metadata.clone()); @@ -1879,15 +2392,17 @@ mod tests { assert!(cached.region_metadata_weight > json_len); assert_eq!( - meta_cache_weight(&key, &cached) as usize, - key.estimated_size() - + parquet_meta_size(&cached.parquet_metadata) - + cached.region_metadata_weight + decoded_meta_cache_weight(&key, &cached) as usize, + key.estimated_size() + cached.parquet_metadata_size + cached.region_metadata_weight + ); + assert_eq!( + cached.parquet_metadata_size, + parquet_meta_size(&cached.parquet_metadata) ); } #[test] - fn test_meta_cache_weight_saturates_on_overflow() { + fn test_decoded_meta_cache_weight_saturates_on_overflow() { let region_metadata = Arc::new(wide_region_metadata(1)); let metadata = sst_parquet_meta_with_region_metadata(region_metadata.clone()); let mut cached = @@ -1896,7 +2411,7 @@ mod tests { let cached = Arc::new(cached); let key = SstMetaKey(region_metadata.region_id, FileId::random()); - assert_eq!(u32::MAX, meta_cache_weight(&key, &cached)); + assert_eq!(u32::MAX, decoded_meta_cache_weight(&key, &cached)); } #[test] diff --git a/src/mito2/src/cache/file_cache.rs b/src/mito2/src/cache/file_cache.rs index 681b2f90e3f..64f2fad6f1b 100644 --- a/src/mito2/src/cache/file_cache.rs +++ b/src/mito2/src/cache/file_cache.rs @@ -34,7 +34,9 @@ use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc::{Sender, UnboundedReceiver}; use crate::access_layer::TempFileCleaner; -use crate::cache::{CachedSstMeta, FILE_TYPE, INDEX_TYPE}; +use crate::cache::{ + CachedSstMeta, FILE_TYPE, INDEX_TYPE, SstMetaPreparation, decode_sst_meta, prepare_sst_meta, +}; use crate::error::{self, OpenDalSnafu, Result}; use crate::metrics::{ CACHE_BYTES, CACHE_HIT, CACHE_MISS, WRITE_CACHE_DOWNLOAD_BYTES_TOTAL, @@ -619,35 +621,56 @@ impl FileCache { /// Get fused SST metadata from the file cache. /// If the file is not in the cache, or metadata loading/decoding fails, return None. + /// Compact cache encoding failures return decoded-only metadata to the caller. pub(crate) async fn get_sst_meta_data( &self, key: IndexKey, cache_metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, + ) -> Option { + let file_path = self.inner.cache_file_path(key); + let metadata = self + .get_parquet_meta_data(key, cache_metrics, page_index_policy) + .await?; + match prepare_sst_meta(&file_path, metadata, None, page_index_policy).await { + Ok(metadata) => Some(metadata), + Err(err) => { + CACHE_MISS + .with_label_values(&[key.file_type.metric_label()]) + .inc(); + warn!( + err; "Failed to prepare cached parquet metadata for key {:?}", + key + ); + None + } + } + } + + /// Gets decoded SST metadata without preparing an in-memory cache entry. + pub(crate) async fn get_decoded_sst_meta_data( + &self, + key: IndexKey, + cache_metrics: &mut MetadataCacheMetrics, + page_index_policy: PageIndexPolicy, ) -> Option> { let file_path = self.inner.cache_file_path(key); - self.get_parquet_meta_data(key, cache_metrics, page_index_policy) - .await - .and_then(|metadata| { - match CachedSstMeta::try_new_with_page_index_policy( - &file_path, - metadata, - None, - page_index_policy, - ) { - Ok(metadata) => Some(Arc::new(metadata)), - Err(err) => { - CACHE_MISS - .with_label_values(&[key.file_type.metric_label()]) - .inc(); - warn!( - err; "Failed to decode cached parquet metadata for key {:?}", - key - ); - None - } - } - }) + let metadata = self + .get_parquet_meta_data(key, cache_metrics, page_index_policy) + .await?; + match decode_sst_meta(&file_path, metadata, None, page_index_policy).await { + Ok(metadata) => Some(metadata), + Err(err) => { + CACHE_MISS + .with_label_values(&[key.file_type.metric_label()]) + .inc(); + warn!( + err; "Failed to decode cached parquet metadata for key {:?}", + key + ); + None + } + } } async fn get_reader(&self, file_path: &str) -> object_store::Result> { diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 6b02a92811a..d186307c9b1 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -638,6 +638,7 @@ mod tests { .write_cache(Some(write_cache.clone())) .build(), ); + assert!(!cache_manager.sst_meta_cache_enabled()); // Create source let metadata = Arc::new(sst_region_metadata()); @@ -682,7 +683,7 @@ mod tests { let sst_info = sst_infos.remove(0); let write_parquet_metadata = sst_info.file_metadata.unwrap(); - // Read metadata from write cache + // Read metadata from write cache without preparing an in-memory metadata cache entry. let handle = sst_file_handle_with_file_id(sst_info.file_id, 0, 1000); let builder = ParquetReaderBuilder::new( data_home, diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index ad59e4fb94f..9388f943ceb 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -43,8 +43,8 @@ use store_api::storage::{ColumnId, RegionId}; use tokio::sync::Semaphore; use crate::access_layer::AccessLayer; -use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileCache, FileType, IndexKey}; +use crate::cache::{CacheManagerRef, SstMetaPreparation, prepare_sst_meta}; use crate::config::MitoConfig; use crate::engine::region_hook::RegionHookRef; use crate::error; @@ -1092,7 +1092,7 @@ async fn preload_parquet_meta_cache_for_files( ) -> usize { if !cache_manager.sst_meta_cache_enabled() || sst_meta_cache_capacity == 0 - || cache_manager.sst_meta_cache_weighted_size() >= sst_meta_cache_capacity + || cache_manager.sst_meta_cache_is_full() { return 0; } @@ -1104,30 +1104,60 @@ async fn preload_parquet_meta_cache_for_files( let mut loaded = 0usize; for file_handle in files { - // Stop when the shared SST meta cache is full. - if cache_manager.sst_meta_cache_weighted_size() >= sst_meta_cache_capacity { + // Stop when the protected authoritative tier is full. + if cache_manager.sst_meta_cache_is_full() { break; } let file_id = file_handle.file_id(); let mut cache_metrics = MetadataCacheMetrics::default(); if let Some(metadata) = cache_manager - .get_parquet_meta_data(file_id, &mut cache_metrics, PageIndexPolicy::Optional) + .get_compact_sst_meta_data(file_id, PageIndexPolicy::Optional) .await { if file_handle.primary_key_range().is_none() - && let Some(primary_key_range) = - extract_primary_key_range(&metadata, ®ion_metadata) + && let Some(primary_key_range) = extract_primary_key_range( + metadata.parquet_metadata().as_ref(), + ®ion_metadata, + ) { file_handle.set_primary_key_range(primary_key_range); } - // Metadata is either already in memory or loaded from file cache. - if cache_metrics.mem_cache_hit == 0 { - loaded += 1; - } continue; } + if let Some(write_cache) = cache_manager.write_cache() { + let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); + if let Some(metadata) = write_cache + .file_cache() + .get_sst_meta_data(key, &mut cache_metrics, PageIndexPolicy::Optional) + .await + { + let decoded = metadata.decoded(); + if file_handle.primary_key_range().is_none() + && let Some(primary_key_range) = extract_primary_key_range( + decoded.parquet_metadata().as_ref(), + ®ion_metadata, + ) + { + file_handle.set_primary_key_range(primary_key_range); + } + match metadata { + SstMetaPreparation::Prepared(metadata) => { + cache_manager.put_prepared_sst_meta(file_id, metadata, false); + loaded += 1; + } + SstMetaPreparation::DecodedOnly { encoding_error, .. } => warn!( + encoding_error; + "Failed to encode file-cached SST metadata during preload, region: {}, file: {}", + region_id, + file_id.file_id() + ), + } + continue; + } + } + if !allow_direct_load { continue; } @@ -1143,8 +1173,35 @@ async fn preload_parquet_meta_cache_for_files( { file_handle.set_primary_key_range(primary_key_range); } - cache_manager.put_parquet_meta_data(file_id, Arc::new(metadata), None); - loaded += 1; + match prepare_sst_meta( + &file_path, + metadata, + // The SST can predate schema changes, so decode its embedded region metadata + // instead of substituting the region's current schema. + None, + PageIndexPolicy::Optional, + ) + .await + { + Ok(SstMetaPreparation::Prepared(metadata)) => { + // Preload the compact canonical entry only. Query traffic promotes decoded + // entries into the separate acceleration tier according to actual demand. + cache_manager.put_prepared_sst_meta(file_id, metadata, false); + loaded += 1; + } + Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( + encoding_error; + "Failed to encode preloaded parquet metadata, region: {}, file: {}", + region_id, + file_path + ), + Err(err) => { + warn!( + err; "Failed to prepare preloaded parquet metadata, region: {}, file: {}", + region_id, file_path + ); + } + } } Err(err) => { // Preloading is best-effort. Failure shouldn't affect region open. @@ -1257,7 +1314,7 @@ mod tests { use object_store::ObjectStore; use object_store::services::{Fs, Memory, S3}; use parquet::arrow::ArrowWriter; - use parquet::file::metadata::KeyValue; + use parquet::file::metadata::{KeyValue, PageIndexPolicy}; use parquet::file::properties::WriterProperties; use store_api::region_request::PathType; use store_api::storage::{FileId, RegionId}; @@ -1490,17 +1547,11 @@ mod tests { assert_eq!(loaded, 1); assert!( cache_manager - .get_parquet_meta_data_from_mem_cache(region_file_id) - .is_some() - ); - // The cached entry must carry the page index so that later `Optional` queries hit - // the in-memory cache instead of reloading metadata on demand. - assert!( - cache_manager - .get_sst_meta_data_from_mem_cache( + .get_compact_sst_meta_data( region_file_id, parquet::file::metadata::PageIndexPolicy::Optional, ) + .await .is_some() ); assert!(file_handle.primary_key_range().is_some()); @@ -1645,7 +1696,8 @@ mod tests { assert_eq!(loaded, 1); assert!( cache_manager - .get_parquet_meta_data_from_mem_cache(region_file_id) + .get_compact_sst_meta_data(region_file_id, PageIndexPolicy::Optional) + .await .is_some() ); } diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index db2ce371991..52c74ef70a6 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -332,7 +332,7 @@ impl FileRangeContext { /// Returns the estimated memory size of this context. /// Mainly accounts for the parquet metadata size. pub(crate) fn memory_size(&self) -> usize { - crate::cache::cache_size::parquet_meta_size(self.reader_builder.parquet_metadata()) + self.reader_builder.parquet_metadata_size() } } diff --git a/src/mito2/src/sst/parquet/metadata.rs b/src/mito2/src/sst/parquet/metadata.rs index 441be373b0c..9d38d6e2167 100644 --- a/src/mito2/src/sst/parquet/metadata.rs +++ b/src/mito2/src/sst/parquet/metadata.rs @@ -81,7 +81,10 @@ impl<'a> MetadataLoader<'a> { let file_size = self.get_file_size().await?; let reader = ParquetMetaDataReader::new() .with_prefetch_hint(Some(DEFAULT_PREFETCH_SIZE as usize)) - .with_page_index_policy(self.page_index_policy); + // Mito uses offset indexes to translate row selections into byte ranges. It does not + // consume Parquet column indexes, so decoding them only bloats the metadata cache. + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(self.page_index_policy); let num_reads = AtomicUsize::new(0); let bytes_read = AtomicU64::new(0); @@ -190,6 +193,7 @@ mod tests { }; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use datatypes::arrow::record_batch::RecordBatch; + use object_store::services::Memory; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use parquet::file::metadata::{KeyValue, ParquetMetaData}; @@ -199,12 +203,12 @@ mod tests { use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::test_util::sst_util::sst_region_metadata; - fn build_test_metadata( + fn build_test_parquet_bytes( include_primary_key: bool, primary_keys: &[&[u8]], row_group_sizes: &[usize], stats_enabled: EnabledStatistics, - ) -> ParquetMetaData { + ) -> Vec { let total_rows = row_group_sizes.iter().sum::(); let mut fields = vec![Field::new("field", ArrowDataType::Int64, true)]; let mut columns: Vec = @@ -253,11 +257,44 @@ mod tests { } writer.close().unwrap(); - ParquetRecordBatchReaderBuilder::try_new(Bytes::from(parquet_bytes)) - .unwrap() - .metadata() - .as_ref() - .clone() + parquet_bytes + } + + fn build_test_metadata( + include_primary_key: bool, + primary_keys: &[&[u8]], + row_group_sizes: &[usize], + stats_enabled: EnabledStatistics, + ) -> ParquetMetaData { + ParquetRecordBatchReaderBuilder::try_new(Bytes::from(build_test_parquet_bytes( + include_primary_key, + primary_keys, + row_group_sizes, + stats_enabled, + ))) + .unwrap() + .metadata() + .as_ref() + .clone() + } + + #[tokio::test] + async fn test_metadata_loader_only_loads_offset_indexes() { + let parquet_bytes = build_test_parquet_bytes(false, &[], &[4], EnabledStatistics::Page); + let file_size = parquet_bytes.len() as u64; + let file_path = "test.parquet"; + let object_store = ObjectStore::new(Memory::default()).unwrap().finish(); + object_store.write(file_path, parquet_bytes).await.unwrap(); + + let mut loader = MetadataLoader::new(object_store, file_path, file_size); + loader.with_page_index_policy(PageIndexPolicy::Required); + let metadata = loader + .load(&mut MetadataCacheMetrics::default()) + .await + .unwrap(); + + assert!(metadata.column_index().is_none()); + assert!(metadata.offset_index().is_some()); } #[test] diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 6459a9877e8..87c521181c1 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -52,7 +52,7 @@ use table::predicate::Predicate; use self::stream::{NestedSchemaAligner, ProjectedRecordBatchStream}; use crate::cache::index::result_cache::PredicateKey; -use crate::cache::{CacheStrategy, CachedSstMeta}; +use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta}; #[cfg(feature = "vector_index")] use crate::error::ApplyVectorIndexSnafu; use crate::error::{ @@ -393,6 +393,7 @@ impl ParquetReaderBuilder { ) .await?; let mut parquet_meta = sst_meta.parquet_metadata(); + let mut parquet_metadata_size = sst_meta.parquet_metadata_size(); let region_meta = sst_meta.region_metadata(); let region_partition_expr_str = self .expected_metadata @@ -504,6 +505,7 @@ impl ParquetReaderBuilder { ) .await?; parquet_meta = sst_meta.parquet_metadata(); + parquet_metadata_size = sst_meta.parquet_metadata_size(); cache_miss |= page_index_cache_miss; } @@ -551,6 +553,7 @@ impl ParquetReaderBuilder { file_handle: self.file_handle.clone(), file_path, parquet_meta, + parquet_metadata_size, arrow_metadata, output_schema, object_store: self.object_store.clone(), @@ -695,18 +698,34 @@ impl ParquetReaderBuilder { metadata_loader.with_page_index_policy(page_index_policy); let metadata = metadata_loader.load(cache_metrics).await?; - let metadata = Arc::new(CachedSstMeta::try_new_with_page_index_policy( - file_path, - metadata, - None, - page_index_policy, - )?); - // Cache the metadata. - self.cache_strategy - .put_sst_meta_data(file_id, metadata.clone()); + let decoded = if self.cache_strategy.sst_meta_cache_enabled() { + let metadata = prepare_sst_meta(file_path, metadata, None, page_index_policy).await?; + let decoded = metadata.decoded(); + match metadata { + SstMetaPreparation::Prepared(metadata) => { + self.cache_strategy + .put_prepared_sst_meta(file_id, metadata, true); + } + SstMetaPreparation::DecodedOnly { encoding_error, .. } => { + warn!( + encoding_error; + "Failed to encode SST metadata for cache, using decoded metadata for {}", + file_path + ); + } + } + decoded + } else { + Arc::new(CachedSstMeta::try_new_with_page_index_policy( + file_path, + metadata, + None, + page_index_policy, + )?) + }; cache_metrics.metadata_load_cost += start.elapsed(); - Ok((metadata, true)) + Ok((decoded, true)) } /// Computes row groups to read, along with their respective row selections. @@ -1735,6 +1754,8 @@ pub(crate) struct RowGroupReaderBuilder { file_path: String, /// Metadata of the parquet file. parquet_meta: Arc, + /// Immutable metadata size, computed once when the footer is decoded. + parquet_metadata_size: usize, /// Arrow reader metadata for building async stream. arrow_metadata: ArrowReaderMetadata, /// Projected output schema aligned with `projection.projected_root_presence`. @@ -1777,6 +1798,10 @@ impl RowGroupReaderBuilder { &self.parquet_meta } + pub(crate) fn parquet_metadata_size(&self) -> usize { + self.parquet_metadata_size + } + pub(crate) fn cache_strategy(&self) -> &CacheStrategy { &self.cache_strategy }