perf: reduce parquet metadata cache footprint (#8527)

* perf(mito2): reuse parquet metadata size accounting

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* perf(mito2): skip unused parquet column indexes

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* perf(mito2): cache compact parquet metadata

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix(mito2): tolerate metadata cache encoding failures

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix(mito2): skip disabled metadata cache encoding

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* refactor(mito2): avoid nested metadata cache task

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
This commit is contained in:
Ruihang Xia
2026-07-16 11:20:13 +00:00
committed by GitHub
parent 3ae8d7d46a
commit c969fb0a5e
10 changed files with 836 additions and 147 deletions
Generated
+1
View File
@@ -8519,6 +8519,7 @@ dependencies = [
"toml 0.8.23",
"tracing",
"uuid",
"zstd",
]
[[package]]
+1
View File
@@ -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
+41 -7
View File
@@ -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);
}
}
});
}
}
}
+588 -73
View File
@@ -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<ParquetMetaData>,
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<RegionMetadata>,
page_index_policy: PageIndexPolicy,
decoded: tokio::sync::Mutex<Weak<CachedSstMeta>>,
}
/// Both forms produced by decoding metadata after a cache miss.
#[derive(Debug)]
pub(crate) struct PreparedSstMeta {
compact: Arc<CompactSstMeta>,
decoded: Arc<CachedSstMeta>,
}
impl PreparedSstMeta {
pub(crate) fn decoded(&self) -> Arc<CachedSstMeta> {
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<CachedSstMeta>,
encoding_error: crate::error::Error,
},
}
impl SstMetaPreparation {
pub(crate) fn decoded(&self) -> Arc<CachedSstMeta> {
match self {
Self::Prepared(metadata) => metadata.decoded(),
Self::DecodedOnly { decoded, .. } => decoded.clone(),
}
}
}
impl CompactSstMeta {
async fn decode(&self) -> Result<Arc<CachedSstMeta>> {
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<RegionMetadataRef>,
page_index_policy: PageIndexPolicy,
) -> Result<Arc<CachedSstMeta>> {
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<RegionMetadataRef>,
page_index_policy: PageIndexPolicy,
) -> Result<SstMetaPreparation> {
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<RegionMetadataRef>,
page_index_policy: PageIndexPolicy,
) -> Result<SstMetaPreparation> {
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<RegionMetadataRef>,
page_index_policy: PageIndexPolicy,
cache_encoding: Result<(Bytes, usize)>,
) -> Result<SstMetaPreparation> {
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<Self> {
@@ -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<CachedSstMeta>) {
/// 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<SstMetaCache>,
/// Cache for decoded SST metadata, used only as an acceleration tier.
sst_decoded_meta_cache: Option<SstDecodedMetaCache>,
/// Cache for vectors.
vector_cache: Option<VectorCache>,
/// Cache for SST byte ranges.
@@ -777,40 +1035,82 @@ impl CacheManager {
metrics: &mut MetadataCacheMetrics,
page_index_policy: PageIndexPolicy,
) -> Option<Arc<CachedSstMeta>> {
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<Arc<ParquetMetaData>> {
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<Arc<CachedSstMeta>> {
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<CachedSstMeta>) {
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<Arc<CompactSstMeta>> {
self.sst_meta_cache.as_ref()?.get(key)
}
fn get_decoded_sst_meta(&self, key: &SstMetaKey) -> Option<Arc<CachedSstMeta>> {
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<Arc<CachedSstMeta>> {
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<ParquetMetaData>,
region_metadata: Option<RegionMetadataRef>,
) {
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<CachedSstMeta>) -> 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<CompactSstMeta>) -> 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::<CompactSstMeta>() + v.encoded_metadata.len();
u32::try_from(size).unwrap_or(u32::MAX)
}
fn decoded_meta_cache_weight(k: &SstMetaKey, v: &Arc<CachedSstMeta>) -> 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<SstMetaKey, Arc<CachedSstMeta>>;
type SstMetaCache = Cache<SstMetaKey, Arc<CompactSstMeta>>;
/// Maps (region id, file id) to decoded SST metadata retained as an acceleration tier.
type SstDecodedMetaCache = Cache<SstMetaKey, Arc<CachedSstMeta>>;
/// 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(&region_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]
+46 -23
View File
@@ -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<SstMetaPreparation> {
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<Arc<CachedSstMeta>> {
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<Option<Reader>> {
+2 -1
View File
@@ -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,
+75 -23
View File
@@ -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, &region_metadata)
&& let Some(primary_key_range) = extract_primary_key_range(
metadata.parquet_metadata().as_ref(),
&region_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(),
&region_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()
);
}
+1 -1
View File
@@ -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()
}
}
+45 -8
View File
@@ -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<u8> {
let total_rows = row_group_sizes.iter().sum::<usize>();
let mut fields = vec![Field::new("field", ArrowDataType::Int64, true)];
let mut columns: Vec<ArrayRef> =
@@ -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]
+36 -11
View File
@@ -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<ParquetMetaData>,
/// 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
}