feat: share concurrent SST metadata loads (#9729)

Signed-off-by: WenyXu <wenymedia@gmail.com>
Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>
This commit is contained in:
Weny Xu
2026-10-29 10:03:35 +00:00
committed by GitHub
co-authored by Lei, HUANG
parent c7ea80558f
commit ee55c0c2fa
17 changed files with 1265 additions and 192 deletions
Generated
+1
View File
@@ -8798,6 +8798,7 @@ dependencies = [
"bytes",
"chrono",
"common-base",
"common-cache",
"common-datasource",
"common-error",
"common-function",
+1
View File
@@ -301,6 +301,7 @@ client = { path = "src/client" }
cmd = { path = "src/cmd", default-features = false }
common-base = { path = "src/common/base" }
common-batcher = { path = "src/common/batcher" }
common-cache = { path = "src/common/cache" }
common-catalog = { path = "src/common/catalog" }
common-config = { path = "src/common/config" }
common-datasource = { path = "src/common/datasource" }
+1
View File
@@ -28,6 +28,7 @@ bytemuck.workspace = true
bytes.workspace = true
chrono.workspace = true
common-base.workspace = true
common-cache.workspace = true
common-datasource.workspace = true
common-error.workspace = true
common-function.workspace = true
+16 -13
View File
@@ -446,26 +446,29 @@ impl AccessLayer {
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.
runtime.spawn_blocking(move || {
match prepare_sst_meta_sync(
&file_path,
Arc::unwrap_or_clone(parquet_metadata),
Some(region_metadata),
page_index_policy,
) {
// Prepare and cache metadata in the background without delaying the SST write.
let task_runtime = runtime.clone();
runtime.spawn(async move {
let preparation = task_runtime.spawn_blocking(move || {
prepare_sst_meta_sync(
&file_path,
Arc::unwrap_or_clone(parquet_metadata),
Some(region_metadata),
page_index_policy,
)
}).await.context(crate::error::JoinSnafu).and_then(|result| result);
match preparation {
Ok(SstMetaPreparation::Prepared(metadata)) => {
cache_manager.put_prepared_sst_meta(file_id, metadata, true);
cache_manager.put_compact_sst_meta(file_id, metadata.compact);
cache_manager.put_sst_meta_data(file_id, metadata.decoded).await;
}
Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!(
encoding_error;
"Failed to encode parquet metadata for cache, file: {}",
file_path
file_id.file_id()
),
Err(err) => {
warn!(err; "Failed to cache parquet metadata for {}", file_path);
warn!(err; "Failed to cache parquet metadata for {}", file_id.file_id());
}
}
});
+743 -136
View File
File diff suppressed because it is too large Load Diff
+32 -1
View File
@@ -15,11 +15,13 @@
//! Utilities for testing cache.
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use bytes::Bytes;
use datatypes::arrow::array::{ArrayRef, Int64Array};
use datatypes::arrow::record_batch::RecordBatch;
use object_store::ObjectStore;
use object_store::layers::mock;
use object_store::services::Fs;
use parquet::arrow::ArrowWriter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
@@ -31,6 +33,33 @@ use store_api::metadata::RegionMetadataRef;
use crate::sst::parquet::PARQUET_METADATA_KEY;
use crate::test_util::sst_util::sst_region_metadata;
/// Holds storage reads until every concurrent caller has been polled.
pub(crate) struct GatedStorageReader {
pub(crate) inner: mock::oio::Reader,
pub(crate) gate: Arc<tokio::sync::Semaphore>,
pub(crate) reads: Arc<AtomicUsize>,
}
impl mock::oio::Read for GatedStorageReader {
async fn open(
&self,
range: mock::BytesRange,
) -> object_store::Result<(mock::RpRead, Box<dyn mock::oio::ReadStreamDyn>)> {
self.reads.fetch_add(1, Ordering::Relaxed);
self.gate.acquire().await.unwrap().forget();
self.inner.open(range).await
}
async fn read(
&self,
range: mock::BytesRange,
) -> object_store::Result<(mock::RpRead, object_store::Buffer)> {
self.reads.fetch_add(1, Ordering::Relaxed);
self.gate.acquire().await.unwrap().forget();
self.inner.read(range).await
}
}
/// Returns a parquet meta data.
pub(crate) fn parquet_meta() -> Arc<ParquetMetaData> {
let file_data = parquet_file_data();
@@ -60,7 +89,9 @@ fn parquet_file_data() -> Vec<u8> {
parquet_file_data_inner(None)
}
fn parquet_file_data_with_region_metadata(region_metadata: &RegionMetadataRef) -> Vec<u8> {
pub(crate) fn parquet_file_data_with_region_metadata(
region_metadata: &RegionMetadataRef,
) -> Vec<u8> {
let json = region_metadata.to_json().unwrap();
let key_value = KeyValue::new(PARQUET_METADATA_KEY.to_string(), json);
parquet_file_data_inner(Some(vec![key_value]))
+3
View File
@@ -21,6 +21,7 @@ use common_time::range::TimestampRange;
use common_time::timestamp::TimeUnit;
use datafusion_common::ScalarValue;
use datafusion_expr::Expr;
use futures::FutureExt;
use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
use snafu::OptionExt;
use store_api::metadata::RegionMetadataRef;
@@ -67,8 +68,10 @@ impl CompactionSstReaderBuilder<'_> {
let schema = scan_input.mapper.output_schema();
let schema = rewrite_json2_schema(schema.arrow_schema(), &plans);
// Erase the nested scan future before embedding it in the compaction task.
let stream = SeqScan::new(scan_input)
.build_flat_reader_for_compaction()
.boxed()
.await?;
Ok(FlatSource::new_stream(schema, stream))
}
+35 -1
View File
@@ -198,6 +198,13 @@ pub enum Error {
location: Location,
},
#[snafu(display("Failed to load SST metadata"))]
LoadSstMeta {
source: Arc<Error>,
#[snafu(implicit)]
location: Location,
},
#[snafu(display("Failed to write parquet file"))]
WriteParquet {
#[snafu(source)]
@@ -1455,6 +1462,7 @@ impl Error {
match self {
Error::OpenDal { error, .. } => error.kind() == ErrorKind::NotFound,
Error::ManifestDeltaNotFound { .. } => true,
Error::LoadSstMeta { source, .. } => source.is_object_not_found(),
_ => false,
}
}
@@ -1577,7 +1585,7 @@ impl ErrorExt for Error {
OpenRegion { source, .. } => source.status_code(),
WriteParquet { .. } => StatusCode::StorageUnavailable,
WriteGroup { source, .. } => source.status_code(),
WriteGroup { source, .. } | LoadSstMeta { source, .. } => source.status_code(),
InvalidBatch { .. } => StatusCode::InvalidArguments,
InvalidRecordBatch { .. } => StatusCode::InvalidArguments,
ConvertVector { source, .. } => source.status_code(),
@@ -1732,6 +1740,7 @@ impl ErrorExt for Error {
| External { source, .. } => source.retry_hint(),
OpenRegion { source, .. }
| LoadSstMeta { source, .. }
| WriteGroup { source, .. }
| FlushRegion { source, .. }
| BuildIndexAsync { source, .. }
@@ -1861,6 +1870,31 @@ mod tests {
}
}
#[test]
fn test_shared_sst_metadata_preserves_error_semantics() {
for source in [
OpenDalSnafu {}
.into_error(object_store::Error::new(ErrorKind::NotFound, "missing SST")),
OpenDalSnafu {}.into_error(object_store::Error::new(
ErrorKind::Unexpected,
"read failed",
)),
RejectWriteSnafu {
region_id: RegionId::new(1, 1),
}
.build(),
] {
let expected_not_found = source.is_object_not_found();
let expected_status = source.status_code();
let expected_retry = source.retry_hint();
let source = Arc::new(source);
let wrapped = LoadSstMetaSnafu.into_error(source);
assert_eq!(expected_not_found, wrapped.is_object_not_found());
assert_eq!(expected_status, wrapped.status_code());
assert_eq!(expected_retry, wrapped.retry_hint());
}
}
#[test]
fn test_manifest_update_persistence() {
let rejected_kinds = [
+44 -6
View File
@@ -272,8 +272,44 @@ lazy_static! {
).unwrap();
}
pub(crate) fn cache_loading_metrics(
cache_type: &'static str,
) -> std::sync::Arc<common_cache::metrics::CacheMetrics> {
use common_cache::metrics::{CacheMetrics, OperationMetrics};
let operation = |stage| OperationMetrics {
started: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "started"]),
succeeded: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "success"]),
failed: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "error"]),
cancelled: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "cancelled"]),
active: CACHE_NATIVE_LOAD_ACTIVE.with_label_values(&[cache_type, stage]),
elapsed: CACHE_NATIVE_LOAD_ELAPSED.with_label_values(&[cache_type, stage]),
};
std::sync::Arc::new(CacheMetrics {
requests: operation("request"),
loads: operation("loader"),
})
}
// Cache metrics.
lazy_static! {
/// Native cache request and actual loader lifecycle; request is not a waiter count.
pub static ref CACHE_NATIVE_LOAD_EVENTS: IntCounterVec = register_int_counter_vec!(
"greptime_mito_cache_native_load_events_total",
"Native cache request and executed loader outcomes; requests include hits and shared loads",
&[TYPE_LABEL, STAGE_LABEL, "event"]
).unwrap();
pub static ref CACHE_NATIVE_LOAD_ACTIVE: IntGaugeVec = register_int_gauge_vec!(
"greptime_mito_cache_native_load_active",
"Active native cache requests or executed loaders, not waiting requests",
&[TYPE_LABEL, STAGE_LABEL]
).unwrap();
pub static ref CACHE_NATIVE_LOAD_ELAPSED: HistogramVec = register_histogram_vec!(
"greptime_mito_cache_native_load_seconds",
"Native cache request or executed loader wall time including suspension and cancellation, not CPU time",
&[TYPE_LABEL, STAGE_LABEL], exponential_buckets(0.0001, 10.0, 6).unwrap()
).unwrap();
/// Cache hit counter.
pub static ref CACHE_HIT: IntCounterVec = register_int_counter_vec!(
"greptime_mito_cache_hit",
@@ -289,12 +325,14 @@ lazy_static! {
)
.unwrap();
/// Cache size in bytes.
pub static ref CACHE_BYTES: IntGaugeVec = register_int_gauge_vec!(
"greptime_mito_cache_bytes",
"mito cache bytes",
&[TYPE_LABEL]
)
.unwrap();
pub static ref CACHE_BYTES: common_cache::metrics::CacheSizeMetrics = {
let metrics = common_cache::metrics::CacheSizeMetrics::new(
prometheus::Opts::new("greptime_mito_cache_bytes", "mito cache bytes; backend cache weights are approximate"),
&[TYPE_LABEL],
).unwrap();
prometheus::register(Box::new(metrics.clone())).unwrap();
metrics
};
/// Download bytes counter in the write cache.
pub static ref WRITE_CACHE_DOWNLOAD_BYTES_TOTAL: IntCounter = register_int_counter!(
"mito_write_cache_download_bytes_total",
+15 -3
View File
@@ -810,7 +810,7 @@ mod tests {
use super::*;
use crate::read::flat_projection::FlatProjectionMapper;
use crate::read::range::RowGroupIndex;
use crate::read::range::{RangeMeta, RowGroupIndex};
use crate::read::scan_region::{PredicateGroup, ScanInput};
use crate::read::scan_util::PartitionMetrics;
use crate::sst::file::{FileHandle, FileMeta};
@@ -853,7 +853,13 @@ mod tests {
.with_files(files)
.with_append_mode(true)
.build();
let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input));
let query_start = input.query_start.unwrap_or_else(Instant::now);
let ranges = RangeMeta::unordered_scan_ranges(&input).await;
let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(
input,
ranges,
query_start,
));
let pruner = Arc::new(Pruner::new_with_options(
stream_ctx,
1,
@@ -912,7 +918,13 @@ mod tests {
.with_predicate(predicate)
.with_append_mode(true)
.build();
let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input));
let query_start = input.query_start.unwrap_or_else(Instant::now);
let ranges = RangeMeta::unordered_scan_ranges(&input).await;
let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(
input,
ranges,
query_start,
));
let pruner = Arc::new(Pruner::new(stream_ctx, 1));
(env, pruner)
}
+21 -14
View File
@@ -15,6 +15,7 @@
//! Structs for partition ranges.
use common_time::Timestamp;
use futures::future::BoxFuture;
use smallvec::{SmallVec, smallvec};
use store_api::region_engine::PartitionRange;
use store_api::storage::TimeSeriesDistribution;
@@ -105,20 +106,24 @@ impl RangeMeta {
}
/// Creates a list of ranges from the `input` for unordered scan.
pub(crate) fn unordered_scan_ranges(input: &ScanInput) -> Vec<RangeMeta> {
let mut ranges = Vec::with_capacity(input.memtables.len() + input.files.len());
Self::push_unordered_mem_ranges(&input.memtables, &mut ranges);
Self::push_unordered_file_ranges(
input.memtables.len(),
&input.files,
&input.cache_strategy,
&mut ranges,
);
pub(crate) fn unordered_scan_ranges(input: &ScanInput) -> BoxFuture<'_, Vec<RangeMeta>> {
// Keep Moka's async lookup type out of callers' nested scanner futures.
Box::pin(async move {
let mut ranges = Vec::with_capacity(input.memtables.len() + input.files.len());
Self::push_unordered_mem_ranges(&input.memtables, &mut ranges);
Self::push_unordered_file_ranges(
input.memtables.len(),
&input.files,
&input.cache_strategy,
&mut ranges,
)
.await;
#[cfg(feature = "enterprise")]
Self::push_extension_ranges(input, &mut ranges);
#[cfg(feature = "enterprise")]
Self::push_extension_ranges(input, &mut ranges);
ranges
ranges
})
}
/// Returns true if the time range of given `meta` overlaps with the time range of this meta.
@@ -203,7 +208,7 @@ impl RangeMeta {
}
}
fn push_unordered_file_ranges(
async fn push_unordered_file_ranges(
num_memtables: usize,
files: &[FileHandle],
cache: &CacheStrategy,
@@ -213,7 +218,9 @@ impl RangeMeta {
for (i, file) in files.iter().enumerate() {
let file_index = num_memtables + i;
// Get parquet meta from the cache.
let parquet_meta = cache.get_parquet_meta_data_from_mem_cache(file.file_id());
let parquet_meta = cache
.get_parquet_meta_data_from_mem_cache(file.file_id())
.await;
if let Some(parquet_meta) = parquet_meta {
// Scans each row group.
for row_group_index in 0..file.meta_ref().num_row_groups {
+9 -4
View File
@@ -435,7 +435,10 @@ impl ScanRegion {
#[tracing::instrument(skip_all, fields(region_id = %self.region_id()))]
pub(crate) async fn unordered_scan(self) -> Result<UnorderedScan> {
let input = self.scan_input().await?;
Ok(UnorderedScan::new(input))
let query_start = input.query_start.unwrap_or_else(Instant::now);
let ranges = RangeMeta::unordered_scan_ranges(&input).await;
let stream_ctx = StreamContext::unordered_scan_ctx(input, ranges, query_start);
Ok(UnorderedScan::new(stream_ctx))
}
/// Scans by series.
@@ -2041,9 +2044,11 @@ impl StreamContext {
}
/// Creates a new [StreamContext] for [UnorderedScan].
pub(crate) fn unordered_scan_ctx(input: ScanInput) -> Self {
let query_start = input.query_start.unwrap_or_else(Instant::now);
let ranges = RangeMeta::unordered_scan_ranges(&input);
pub(crate) fn unordered_scan_ctx(
input: ScanInput,
ranges: Vec<RangeMeta>,
query_start: Instant,
) -> Self {
READ_SST_COUNT.observe(input.num_files() as f64);
Self {
input,
+6 -3
View File
@@ -35,7 +35,9 @@ use store_api::region_engine::{
use crate::error::{PartitionOutOfRangeSnafu, Result};
use crate::read::pruner::{PartitionPruner, Pruner};
use crate::read::scan_region::{ScanInput, StreamContext};
#[cfg(test)]
use crate::read::scan_region::ScanInput;
use crate::read::scan_region::StreamContext;
use crate::read::scan_util::{
PartitionMetrics, PartitionMetricsList, prepare_flat_mem_stream, scan_flat_file_ranges,
scan_flat_mem_ranges,
@@ -59,7 +61,8 @@ pub struct UnorderedScan {
impl UnorderedScan {
/// Creates a new [UnorderedScan].
pub(crate) fn new(input: ScanInput) -> Self {
pub(crate) fn new(stream_ctx: StreamContext) -> Self {
let input = &stream_ctx.input;
let mut properties = ScannerProperties::default()
.with_append_mode(input.append_mode)
.with_total_rows(input.total_rows())
@@ -67,7 +70,7 @@ impl UnorderedScan {
if let Some(counters) = input.query_stat_counters.clone() {
properties.set_query_stat_counters(counters);
}
let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input));
let stream_ctx = Arc::new(stream_ctx);
properties.partitions = vec![stream_ctx.partition_ranges()];
// Create the shared pruner with number of workers equal to CPU cores.
+6 -2
View File
@@ -1299,7 +1299,7 @@ async fn preload_parquet_meta_cache_for_files(
}
match metadata {
SstMetaPreparation::Prepared(metadata) => {
cache_manager.put_prepared_sst_meta(file_id, metadata, false);
cache_manager.put_compact_sst_meta(file_id, metadata.compact);
loaded += 1;
}
SstMetaPreparation::DecodedOnly { encoding_error, .. } => warn!(
@@ -1342,7 +1342,7 @@ async fn preload_parquet_meta_cache_for_files(
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);
cache_manager.put_compact_sst_meta(file_id, metadata.compact);
loaded += 1;
}
Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!(
@@ -1801,6 +1801,7 @@ mod tests {
assert!(
cache_manager
.get_parquet_meta_data_from_mem_cache(region_file_id)
.await
.is_none()
);
@@ -1877,6 +1878,7 @@ mod tests {
assert!(
cache_manager
.get_parquet_meta_data_from_mem_cache(region_file_id)
.await
.is_none()
);
@@ -1896,6 +1898,7 @@ mod tests {
assert!(
cache_manager
.get_parquet_meta_data_from_mem_cache(region_file_id)
.await
.is_none()
);
}
@@ -1950,6 +1953,7 @@ mod tests {
assert!(
cache_manager
.get_parquet_meta_data_from_mem_cache(region_file_id)
.await
.is_none()
);
+3 -1
View File
@@ -873,7 +873,9 @@ pub async fn delete_files(
if let Some(cache) = &cache_manager {
for (file_id, _) in file_ids {
cache.invalidate_range_results(*file_id);
cache.remove_parquet_meta_data(RegionFileId::new(region_id, *file_id));
cache
.remove_parquet_meta_data(RegionFileId::new(region_id, *file_id))
.await;
}
}
let mut attempted_files = Vec::with_capacity(file_ids.len());
@@ -636,7 +636,8 @@ mod tests {
move |expr| {
let _d = &d;
let cache = Arc::new(InvertedIndexCache::new(10, 10, 100));
let puffin_metadata_cache = Arc::new(PuffinMetadataCache::new(10, &CACHE_BYTES));
let puffin_metadata_cache =
Arc::new(PuffinMetadataCache::new(10, CACHE_BYTES.gauges()));
let applier = InvertedIndexApplierBuilder::new(
table_dir.clone(),
PathType::Bare,
+327 -7
View File
@@ -53,7 +53,10 @@ use store_api::storage::{ColumnId, FileId};
use table::predicate::Predicate;
use crate::cache::index::result_cache::PredicateKey;
use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta};
use crate::cache::{
CacheStrategy, CachedSstMeta, SstMetaObservation, SstMetaPreparation, SstMetaResolution,
prepare_sst_meta,
};
use crate::error::{
ParquetToArrowSchemaSnafu, ReadDataPartSnafu, Result, SerializePartitionExprSnafu,
UnexpectedSnafu,
@@ -822,14 +825,66 @@ impl ParquetReaderBuilder {
.with_label_values(&["read_parquet_metadata"])
.start_timer();
let observation = self
.cache_strategy
.observe_sst_meta(self.file_handle.file_id(), page_index_policy, cache_metrics)
.await;
if let SstMetaObservation::Decoded(metadata) = observation {
cache_metrics.metadata_load_cost += start.elapsed();
return Ok((metadata, false));
}
let mut loader_metrics = MetadataCacheMetrics::default();
// Keep the shared loader future off the stack of nested series-index builders.
// Decoded metadata hits above do not need this allocation.
let (result, resolution) = Box::pin(self.cache_strategy.get_or_load_observed_sst_meta(
self.file_handle.file_id(),
page_index_policy,
async {
let result = self
.read_parquet_metadata_inner(
file_path,
file_size,
&mut loader_metrics,
page_index_policy,
observation,
)
.await;
(
result,
SstMetaResolution {
file_cache_hit: loader_metrics.file_cache_hit,
cache_miss: loader_metrics.cache_miss,
},
)
},
))
.await;
// These are resolution counts, including a known miss before a failed remote read.
// Only the selected loader contributes actual I/O counts and bytes.
cache_metrics.file_cache_hit += resolution.file_cache_hit;
cache_metrics.cache_miss += resolution.cache_miss;
cache_metrics.num_reads += loader_metrics.num_reads;
cache_metrics.bytes_read += loader_metrics.bytes_read;
cache_metrics.metadata_load_cost += start.elapsed();
result
}
async fn read_parquet_metadata_inner(
&self,
file_path: &str,
file_size: u64,
cache_metrics: &mut MetadataCacheMetrics,
page_index_policy: PageIndexPolicy,
observation: SstMetaObservation,
) -> Result<(Arc<CachedSstMeta>, bool)> {
let file_id = self.file_handle.file_id();
// Tries to get from cache with metrics tracking.
if let Some(metadata) = self
.cache_strategy
.get_sst_meta_data(file_id, cache_metrics, page_index_policy)
.resolve_sst_meta(observation, file_id, cache_metrics, page_index_policy)
.await
{
cache_metrics.metadata_load_cost += start.elapsed();
return Ok((metadata, false));
}
@@ -852,7 +907,7 @@ impl ParquetReaderBuilder {
match metadata {
SstMetaPreparation::Prepared(metadata) => {
self.cache_strategy
.put_prepared_sst_meta(file_id, metadata, true);
.put_compact_sst_meta(file_id, metadata.compact);
}
SstMetaPreparation::DecodedOnly { encoding_error, .. } => {
warn!(
@@ -872,7 +927,6 @@ impl ParquetReaderBuilder {
)?)
};
cache_metrics.metadata_load_cost += start.elapsed();
Ok((decoded, true))
}
@@ -1635,13 +1689,13 @@ impl ReaderFilterMetrics {
pub struct MetadataCacheMetrics {
/// Number of memory cache hits for parquet metadata.
pub mem_cache_hit: usize,
/// Number of file cache hits for parquet metadata.
/// Successful file-cache resolutions, including results shared with concurrent callers.
pub file_cache_hit: usize,
/// Number of cache misses for parquet metadata.
pub cache_miss: usize,
/// Duration to load parquet metadata.
pub metadata_load_cost: Duration,
/// Number of read operations performed.
/// Physical reads executed by this caller; shared waiters do not repeat them.
pub num_reads: usize,
/// Total bytes read from storage.
pub bytes_read: u64,
@@ -2470,10 +2524,12 @@ mod tests {
use datatypes::extension::json::Json2ExtensionType;
use datatypes::prelude::ConcreteDataType;
use datatypes::schema::ColumnSchema;
use object_store::layers::mock::MockLayerBuilder;
use object_store::services::Memory;
use parquet::arrow::ArrowWriter;
use parquet::arrow::arrow_reader::RowSelector;
use parquet::file::properties::WriterProperties;
use std::sync::atomic::{AtomicUsize, Ordering};
use store_api::codec::PrimaryKeyEncoding;
use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
use store_api::region_request::PathType;
@@ -2482,12 +2538,276 @@ mod tests {
use super::*;
use crate::cache::CacheManager;
use crate::cache::test_util::{GatedStorageReader, parquet_file_data_with_region_metadata};
use crate::sst::parquet::metadata::MetadataLoader;
use crate::sst::parquet::prefilter::{build_reader_filter_plan, execute_prefilter};
use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns};
use crate::sst::parquet::row_group::ParquetFetchMetrics;
use crate::test_util::sst_util::{sst_file_handle, sst_region_metadata};
#[tokio::test]
async fn test_shared_file_metadata_resolution_counts_without_duplicate_io() {
use crate::cache::file_cache::{FileType, IndexKey, IndexValue};
use crate::cache::write_cache::WriteCache;
use crate::test_util::TestEnv;
use common_base::readable_size::ReadableSize;
let env = TestEnv::new().await;
let metadata = Arc::new(sst_region_metadata());
let bytes = parquet_file_data_with_region_metadata(&metadata);
let handle = sst_file_handle(0, 3);
let reads = Arc::new(AtomicUsize::new(0));
let gate = Arc::new(tokio::sync::Semaphore::new(0));
let layer = MockLayerBuilder::default()
.reader_factory(Arc::new({
let reads = reads.clone();
let gate = gate.clone();
move |_, _, inner| {
Box::new(GatedStorageReader {
inner,
reads: reads.clone(),
gate: gate.clone(),
})
}
}))
.build()
.unwrap();
let local = ObjectStore::new(Memory::default()).unwrap().layer(layer);
let write_cache = Arc::new(
WriteCache::new(
local.clone(),
ReadableSize::mb(10),
None,
None,
false,
env.get_puffin_manager(),
env.get_intermediate_manager(),
None,
)
.await
.unwrap(),
);
let file_cache = write_cache.file_cache();
let key = IndexKey::new(
handle.region_id(),
handle.file_id().file_id(),
FileType::Parquet,
);
local
.write(&file_cache.cache_file_path(key), bytes.clone())
.await
.unwrap();
file_cache
.put(
key,
IndexValue {
file_size: bytes.len() as _,
},
)
.await;
let cache = Arc::new(
CacheManager::builder()
.sst_meta_cache_size(1024 * 1024)
.write_cache(Some(write_cache))
.build(),
);
// The remote store is empty: success must come from the local file cache.
let builder = ParquetReaderBuilder::new(
"file_resolution".to_string(),
PathType::Bare,
handle.clone(),
ObjectStore::new(Memory::default()).unwrap(),
)
.cache(CacheStrategy::EnableAll(cache));
let path = handle.file_path("file_resolution", PathType::Bare);
let mut requests = (0..8)
.map(|_| {
Box::pin(async {
let mut metrics = MetadataCacheMetrics::default();
let result = builder
.read_parquet_metadata(
&path,
bytes.len() as u64,
&mut metrics,
PageIndexPolicy::Skip,
)
.await;
(result, metrics)
})
})
.collect::<Vec<_>>();
for request in &mut requests {
assert!(futures::poll!(request).is_pending());
}
assert_eq!(1, reads.load(Ordering::Relaxed));
gate.add_permits(32);
let results =
tokio::time::timeout(Duration::from_secs(10), futures::future::join_all(requests))
.await
.unwrap();
assert_eq!(1, reads.load(Ordering::Relaxed));
assert_eq!(
1,
results
.iter()
.map(|(_, metrics)| metrics.num_reads)
.sum::<usize>()
);
for (result, metrics) in results {
assert!(!result.unwrap().1);
assert_eq!(0, metrics.mem_cache_hit);
assert_eq!(1, metrics.file_cache_hit);
assert_eq!(0, metrics.cache_miss);
}
let mut warm = MetadataCacheMetrics::default();
builder
.read_parquet_metadata(&path, bytes.len() as u64, &mut warm, PageIndexPolicy::Skip)
.await
.unwrap();
assert_eq!(1, warm.mem_cache_hit);
assert_eq!(
0, warm.file_cache_hit,
"hot hits must not replay the original file source"
);
assert_eq!(0, warm.num_reads);
}
#[tokio::test]
#[allow(
clippy::print_stdout,
reason = "records the controlled storage-read comparison"
)]
async fn test_metadata_load_coalescing_storage_reads() {
const CONCURRENCY: usize = 8;
let metadata = Arc::new(sst_region_metadata());
let bytes = parquet_file_data_with_region_metadata(&metadata);
let file_handle = sst_file_handle(0, 3);
let file_path = file_handle.file_path("metadata_coalescing", PathType::Bare);
// Same bytes, file, cache budget and concurrency; only bypass the coordinator
// in the control arm. Repeat with fresh caches to verify cold behavior.
for _ in 0..3 {
for coalesce in [false, true] {
let reads = Arc::new(AtomicUsize::new(0));
let gate = Arc::new(tokio::sync::Semaphore::new(0));
let layer = MockLayerBuilder::default()
.reader_factory(Arc::new({
let reads = reads.clone();
let gate = gate.clone();
move |_, _, inner| {
Box::new(GatedStorageReader {
inner,
reads: reads.clone(),
gate: gate.clone(),
})
}
}))
.build()
.unwrap();
let store = ObjectStore::new(Memory::default()).unwrap().layer(layer);
store.write(&file_path, bytes.clone()).await.unwrap();
let cache = Arc::new(
CacheManager::builder()
.sst_meta_cache_size(1024 * 1024)
.build(),
);
let builder = ParquetReaderBuilder::new(
"metadata_coalescing".to_string(),
PathType::Bare,
file_handle.clone(),
store,
)
.cache(CacheStrategy::EnableAll(cache.clone()));
let mut requests = (0..CONCURRENCY)
.map(|_| {
Box::pin(async {
let mut metrics = MetadataCacheMetrics::default();
let result = if coalesce {
builder
.read_parquet_metadata(
&file_path,
bytes.len() as u64,
&mut metrics,
PageIndexPolicy::Skip,
)
.await
} else {
let observation = builder
.cache_strategy
.observe_sst_meta(
file_handle.file_id(),
PageIndexPolicy::Skip,
&mut metrics,
)
.await;
builder
.read_parquet_metadata_inner(
&file_path,
bytes.len() as u64,
&mut metrics,
PageIndexPolicy::Skip,
observation,
)
.await
};
(result.unwrap(), metrics)
})
})
.collect::<Vec<_>>();
for request in &mut requests {
assert!(futures::poll!(request).is_pending());
}
let expected_reads = if coalesce { 1 } else { CONCURRENCY };
assert_eq!(expected_reads, reads.load(Ordering::Relaxed));
gate.add_permits(CONCURRENCY * 4);
let results = tokio::time::timeout(
Duration::from_secs(10),
futures::future::join_all(requests),
)
.await
.unwrap();
assert_eq!(expected_reads, reads.load(Ordering::Relaxed));
assert_eq!(
expected_reads,
results.iter().map(|(_, m)| m.num_reads).sum::<usize>()
);
for ((value, miss), request_metrics) in &results {
assert_eq!(1, request_metrics.cache_miss);
assert_eq!(0, request_metrics.mem_cache_hit);
assert_eq!(0, request_metrics.file_cache_hit);
assert!(
*miss,
"each request must preserve background-download eligibility"
);
assert_eq!(3, value.parquet_metadata().file_metadata().num_rows());
if coalesce {
assert!(Arc::ptr_eq(value, &results[0].0.0));
}
}
let mut warm_metrics = MetadataCacheMetrics::default();
let (_, miss) = tokio::time::timeout(
Duration::from_secs(5),
builder.read_parquet_metadata(
&file_path,
bytes.len() as u64,
&mut warm_metrics,
PageIndexPolicy::Skip,
),
)
.await
.unwrap()
.unwrap();
assert!(!miss);
assert_eq!(1, warm_metrics.mem_cache_hit);
assert_eq!(0, warm_metrics.cache_miss);
assert_eq!(0, warm_metrics.num_reads);
assert_eq!(expected_reads, reads.load(Ordering::Relaxed));
println!(
"coalesce={coalesce} concurrency={CONCURRENCY} cold_storage_reads={expected_reads} warm_storage_reads=0"
);
}
}
}
async fn prefilter_test_builder(
object_store: ObjectStore,
predicate: Predicate,