perf(mito2): read compaction input footers concurrently (#10000)

* perf(mito2): read compaction input footers concurrently Compaction read the parquet metadata of its inputs one file at a time before merging. It now reads up to 8 footers at once, keeping the input order. Remote compaction also caches SST metadata now, so the merge reuses those footers instead of reading each one a second time. With 20ms object store reads, a remote compaction of 8 inputs issued 57 read requests instead of 65, and one output took 0.51s instead of 0.75s with `experimental_compaction_output_parallelism` set to 8. Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(mito2): stop accounting metadata bytes when caches drop Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * test(mito2): cover compaction with a missing input SST Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix(mito2): skip output metadata warmup for remote compaction Per-job compaction caches are not reused by readers of newly written SSTs. Keep input footer reuse, but preserve output prewarming only for long-lived engine caches. Cover both cache policies and the unchanged flush path with real SST reads. Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> --------- Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
2026-11-06 19:06:58 +00:00 · 2026-11-04 11:09:24 +00:00
parent 351f4207b4
commit f4ebe0b34b
4 changed files with 353 additions and 35 deletions
@@ -425,7 +425,7 @@ impl AccessLayer {
        };

        // Put parquet metadata to cache manager.
-        if !sst_info.is_empty() && cache_manager.sst_meta_cache_enabled() {
+        if !sst_info.is_empty() && cache_manager.sst_meta_cache_on_write_enabled() {
            let runtime = match op_type {
                OperationType::Compact => common_runtime::compact_runtime(),
                OperationType::Flush => common_runtime::global_runtime(),
@@ -918,6 +918,116 @@ mod tests {
        }
    }

+    #[rstest::rstest]
+    #[case::flush(OperationType::Flush, false)]
+    #[case::local_compaction(OperationType::Compact, false)]
+    #[case::input_only_compaction(OperationType::Compact, true)]
+    #[tokio::test]
+    async fn test_written_sst_metadata_cache_policy(
+        #[case] op_type: OperationType,
+        #[case] input_only: bool,
+    ) {
+        use crate::cache::CacheStrategy;
+        use crate::sst::file::FileMeta;
+        use crate::sst::file_purger::NoopFilePurger;
+        use crate::sst::parquet::reader::MetadataCacheMetrics;
+        use crate::test_util::sst_util::{
+            new_flat_source_from_record_batches, new_record_batch_by_range, sst_region_metadata,
+        };
+
+        let mut env = TestEnv::new().await;
+        let store = env.init_object_store_manager();
+        let access_layer = AccessLayer::new(
+            "test",
+            PathType::Bare,
+            store,
+            env.get_puffin_manager(),
+            env.get_intermediate_manager(),
+        );
+        let builder = CacheManager::builder().sst_meta_cache_size(1024 * 1024);
+        let cache = Arc::new(if input_only {
+            builder.disable_sst_meta_cache_on_write().build()
+        } else {
+            builder.build()
+        });
+        let metadata = Arc::new(sst_region_metadata());
+        let region_id = metadata.region_id;
+        let infos =
+            access_layer
+                .write_sst(
+                    SstWriteRequest {
+                        op_type,
+                        metadata,
+                        source: new_flat_source_from_record_batches(vec![
+                            new_record_batch_by_range(&["a", "b"], 0, 16),
+                        ]),
+                        cache_manager: cache.clone(),
+                        max_sequence: None,
+                        sst_write_format: FormatType::Flat,
+                        preserve_row_sequence: false,
+                        index_options: Default::default(),
+                        index_config: Default::default(),
+                        inverted_index_config: Default::default(),
+                        fulltext_index_config: Default::default(),
+                        bloom_filter_index_config: Default::default(),
+                    },
+                    &WriteOptions::default(),
+                    &mut Metrics::new(WriteType::Compaction),
+                )
+                .await
+                .unwrap();
+        assert_eq!(1, infos.len());
+        // A prewarming task owns a cache reference until both tiers are populated.
+        // Drain it before checking absence so the negative assertion cannot race it.
+        tokio::time::timeout(Duration::from_secs(10), async {
+            while Arc::strong_count(&cache) > 1 {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .unwrap();
+        let info = &infos[0];
+        let file_id = RegionFileId::new(region_id, info.file_id);
+        let cached = cache
+            .get_sst_meta_data(
+                file_id,
+                &mut MetadataCacheMetrics::default(),
+                PageIndexPolicy::Skip,
+                &common_runtime::compact_runtime(),
+            )
+            .await;
+        assert_eq!(!input_only, cached.is_some());
+
+        let file = FileHandle::new(
+            FileMeta {
+                region_id,
+                file_id: info.file_id,
+                file_size: info.file_size,
+                ..Default::default()
+            },
+            Arc::new(NoopFilePurger),
+        );
+        let path = file.file_path(access_layer.table_dir(), access_layer.path_type());
+        let reader = access_layer
+            .read_sst(file)
+            .cache(CacheStrategy::Compaction(cache));
+        let mut cold = MetadataCacheMetrics::default();
+        let (first, _) = reader
+            .read_parquet_metadata(&path, info.file_size, &mut cold, PageIndexPolicy::Skip)
+            .await
+            .unwrap();
+        assert_eq!(input_only, cold.num_reads > 0);
+        let mut warm = MetadataCacheMetrics::default();
+        let (second, _) = reader
+            .read_parquet_metadata(&path, info.file_size, &mut warm, PageIndexPolicy::Skip)
+            .await
+            .unwrap();
+        assert_eq!(16, second.parquet_metadata().file_metadata().num_rows());
+        assert!(Arc::ptr_eq(&first, &second));
+        assert_eq!(0, warm.num_reads);
+        assert_eq!(1, warm.mem_cache_hit);
+    }
+
    #[tokio::test]
    async fn test_new_fs_cache_store() {
        let root = common_test_util::temp_dir::create_temp_dir("fs-cache-store");
@@ -30,6 +30,7 @@ use std::{io, mem};

 use bytes::Bytes;
 use common_base::readable_size::ReadableSize;
+use common_cache::metrics::{CacheSizeMetric, CacheSizeRegistration};
 use common_datasource::compression::CompressionType;
 use common_runtime::Runtime;
 use common_runtime::runtime::RuntimeTrait;
@@ -1107,6 +1108,10 @@ impl CacheStrategy {
 pub struct CacheManager {
    /// Cache for compact, authoritative SST metadata.
    sst_meta_cache: Option<SstMetaCache>,
+    /// Unregisters compact metadata bytes when this manager is dropped.
+    _sst_meta_size_registration: Option<CacheSizeRegistration>,
+    /// Skips output metadata prewarming for caches used only by the current compaction.
+    disable_sst_meta_cache_on_write: bool,
    /// Cache for decoded SST metadata, used only as an acceleration tier.
    sst_decoded_meta_cache: Option<SstDecodedMetaCache>,
    /// Cache for vectors.
@@ -1292,9 +1297,6 @@ impl CacheManager {
    ) {
        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).into());
            cache.insert(key, metadata);
        }
    }
@@ -1411,6 +1413,11 @@ impl CacheManager {
        self.sst_meta_cache.is_some()
    }

+    /// Whether newly written SST metadata should be prepared for later readers.
+    pub(crate) fn sst_meta_cache_on_write_enabled(&self) -> bool {
+        self.sst_meta_cache_enabled() && !self.disable_sst_meta_cache_on_write
+    }
+
    /// Gets a vector with repeated value for specific `key`.
    pub fn get_repeated_vector(
        &self,
@@ -1578,6 +1585,7 @@ pub fn selector_result_cache_hit() {
 #[derive(Default)]
 pub struct CacheManagerBuilder {
    sst_meta_cache_size: u64,
+    disable_sst_meta_cache_on_write: bool,
    vector_cache_size: u64,
    page_cache_size: u64,
    index_metadata_size: u64,
@@ -1598,6 +1606,12 @@ impl CacheManagerBuilder {
        self
    }

+    /// Keeps read-side metadata caching without prewarming newly written SSTs.
+    pub(crate) fn disable_sst_meta_cache_on_write(mut self) -> Self {
+        self.disable_sst_meta_cache_on_write = true;
+        self
+    }
+
    /// Sets vector cache size.
    pub fn vector_cache_size(mut self, bytes: u64) -> Self {
        self.vector_cache_size = bytes;
@@ -1666,6 +1680,13 @@ impl CacheManagerBuilder {

    /// Builds the [CacheManager].
    pub fn build(self) -> CacheManager {
+        self.build_with_sst_meta_size_metric(CACHE_BYTES.bind(&[SST_META_TYPE]))
+    }
+
+    fn build_with_sst_meta_size_metric(
+        self,
+        sst_meta_size_metric: CacheSizeMetric,
+    ) -> CacheManager {
        // 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.
@@ -1675,17 +1696,16 @@ impl CacheManagerBuilder {
            Cache::builder()
                .max_capacity(compact_meta_capacity)
                .weigher(meta_cache_weight)
-                .eviction_listener(|k, v, cause| {
-                    let size = meta_cache_weight(&k, &v);
-                    CACHE_BYTES
-                        .with_label_values(&[SST_META_TYPE])
-                        .sub(size.into());
+                .eviction_listener(|_k, _v, cause| {
                    CACHE_EVICTION
                        .with_label_values(&[SST_META_TYPE, removal_cause_str(cause)])
                        .inc();
                })
                .build()
        });
+        let sst_meta_size_registration = sst_meta_cache
+            .clone()
+            .map(|cache| sst_meta_size_metric.register(move || cache.weighted_size()));
        let sst_decoded_meta_cache = (decoded_meta_capacity != 0).then(|| {
            common_cache::Cache::new(
                "sst_decoded_metadata",
@@ -1754,6 +1774,8 @@ impl CacheManagerBuilder {
            .then(|| range_cache::RangeResultCache::new(self.range_result_cache_size));
        CacheManager {
            sst_meta_cache,
+            _sst_meta_size_registration: sst_meta_size_registration,
+            disable_sst_meta_cache_on_write: self.disable_sst_meta_cache_on_write,
            sst_decoded_meta_cache,
            vector_cache,
            page_cache,
@@ -2417,6 +2439,65 @@ mod tests {
        }
    }

+    #[test]
+    fn test_compact_sst_meta_bytes_follow_cache_lifetime() {
+        use common_cache::metrics::CacheSizeMetrics;
+        use prometheus::Opts;
+        use prometheus::core::Collector;
+
+        let metrics =
+            CacheSizeMetrics::new(Opts::new("test_compact_sst_meta_bytes", "test"), &["type"])
+                .unwrap();
+        let (metadata, region) = sst_parquet_meta();
+        let SstMetaPreparation::Prepared(prepared) = prepare_sst_meta_sync(
+            "test.parquet",
+            Arc::unwrap_or_clone(metadata),
+            None,
+            PageIndexPolicy::Skip,
+        )
+        .unwrap() else {
+            panic!("valid metadata should produce a compact entry");
+        };
+        let file_id = RegionFileId::new(region.region_id, FileId::random());
+        let expected_bytes = f64::from(meta_cache_weight(
+            &SstMetaKey(file_id.region_id(), file_id.file_id()),
+            &prepared.compact,
+        ));
+        let make_cache = || {
+            let manager = Arc::new(
+                CacheManager::builder()
+                    .sst_meta_cache_size(1024 * 1024)
+                    .build_with_sst_meta_size_metric(metrics.bind(&[SST_META_TYPE])),
+            );
+            manager.put_compact_sst_meta(file_id, prepared.compact.clone());
+            manager.sst_meta_cache.as_ref().unwrap().run_pending_tasks();
+            manager
+        };
+        let sample = || {
+            metrics.collect()[0].get_metric()[0]
+                .get_gauge()
+                .as_ref()
+                .unwrap()
+                .value()
+        };
+        let resident = make_cache();
+        assert!(expected_bytes > 0.0);
+        assert_eq!(expected_bytes, sample());
+
+        // Repeated remote compactions must not leave byte counts after releasing their caches.
+        for _ in 0..3 {
+            let transient = make_cache();
+            assert_eq!(2.0 * expected_bytes, sample());
+            let reader = transient.clone();
+            drop(transient);
+            assert_eq!(2.0 * expected_bytes, sample());
+            drop(reader);
+            assert_eq!(expected_bytes, sample());
+        }
+        drop(resident);
+        assert_eq!(0.0, sample());
+    }
+
    #[test]
    fn test_sst_meta_cache_splits_capacity() {
        let cache = CacheManager::builder().sst_meta_cache_size(101).build();
@@ -344,7 +344,14 @@ impl CompactionRegionOpener {
            region_options: req.region_options.clone(),
            engine_config: self.mito_config.clone(),
            region_metadata: region_metadata.clone(),
-            cache_manager: Arc::new(CacheManager::default()),
+            // Lets the scan reuse the input footers read while planning the merge.
+            cache_manager: Arc::new(
+                CacheManager::builder()
+                    .sst_meta_cache_size(mito_config.sst_meta_cache_size.as_bytes())
+                    // This per-job cache has no readers for the newly written outputs.
+                    .disable_sst_meta_cache_on_write()
+                    .build(),
+            ),
            access_layer,
            manifest_ctx,
            current_version,
@@ -1310,6 +1317,112 @@ mod tests {
        }
    }

+    #[rstest::rstest]
+    #[case::primary_key(false)]
+    #[case::flat(true)]
+    #[tokio::test]
+    async fn test_compaction_skips_missing_input(#[case] flat_format: bool) {
+        use datatypes::arrow::array::AsArray;
+        use datatypes::arrow::datatypes::TimestampMillisecondType;
+        use store_api::region_engine::RegionEngine;
+        use store_api::region_request::RegionRequest;
+
+        use crate::engine::compaction_test::put_and_flush;
+        use crate::test_util::{CreateRequestBuilder, TestEnv, rows_schema};
+
+        let mut env = TestEnv::new().await;
+        let config = MitoConfig {
+            default_flat_format: flat_format,
+            min_compaction_interval: Duration::from_secs(3600),
+            ..Default::default()
+        };
+        let engine = env.create_engine(config.clone()).await;
+        let region_id = RegionId::new(1, 1);
+        let request = CreateRequestBuilder::new().build();
+        let columns = rows_schema(&request);
+        engine
+            .handle_request(region_id, RegionRequest::Create(request))
+            .await
+            .unwrap();
+        // Nine files cross the eight-footer concurrency limit.
+        for file in 0..9 {
+            put_and_flush(&engine, region_id, &columns, file * 2..file * 2 + 2).await;
+        }
+        let region = engine.get_region(region_id).unwrap();
+        let version = region.version();
+        let mut inputs: Vec<_> = version.ssts.levels()[0].files().cloned().collect();
+        assert_eq!(9, inputs.len());
+        inputs.sort_unstable_by_key(|file| file.meta_ref().sequence);
+        let missing = inputs[4].clone();
+        let output = CompactionOutput {
+            output_level: 1,
+            inputs,
+            filter_deleted: true,
+            output_time_range: None,
+        };
+        // Delete after selecting the inputs; a fresh cache avoids hiding the missing
+        // footer behind the metadata populated by flush.
+        let context = CompactionContext {
+            region_id,
+            region_options: version.options.clone(),
+            version_options: version.options.clone(),
+            engine_config: Arc::new(config),
+            region_metadata: version.metadata.clone(),
+            cache_manager: Arc::new(
+                CacheManager::builder()
+                    .sst_meta_cache_size(1024 * 1024)
+                    .build(),
+            ),
+            access_layer: region.access_layer.clone(),
+            manifest_ctx: region.manifest_ctx.clone(),
+            plugins: Plugins::new(),
+        };
+        region
+            .access_layer
+            .object_store()
+            .delete(&missing.file_path(
+                region.access_layer.table_dir(),
+                region.access_layer.path_type(),
+            ))
+            .await
+            .unwrap();
+        let (files, _) = DefaultSstMerger
+            .merge_single_output(context, output, WriteOptions::default())
+            .await
+            .unwrap();
+        assert_eq!(1, files.len());
+        let outputs: Vec<_> = files.into_iter().map(new_file_handle).collect();
+        let mut reader = CompactionSstReaderBuilder {
+            metadata: version.metadata.clone(),
+            sst_layer: region.access_layer.clone(),
+            cache: Arc::new(CacheManager::default()),
+            inputs: &outputs,
+            append_mode: false,
+            filter_deleted: true,
+            time_range: None,
+            merge_mode: MergeMode::LastRow,
+            read_parallelism: 1,
+        }
+        .build_flat_sst_reader()
+        .await
+        .unwrap();
+        let mut timestamps = Vec::new();
+        while let Some(batch) = reader.next_batch().await.unwrap() {
+            timestamps.extend_from_slice(
+                batch
+                    .column_by_name("ts")
+                    .unwrap()
+                    .as_primitive::<TimestampMillisecondType>()
+                    .values(),
+            );
+        }
+        let expected: Vec<i64> = (0..18)
+            .filter(|ts| !(8..10).contains(ts))
+            .map(|ts| ts * 1000)
+            .collect();
+        assert_eq!(expected, timestamps);
+    }
+
    /// An [`SstMerger`] that returns pre-configured results per call index.
    ///
    /// Call 0 gets `results[0]`, call 1 gets `results[1]`, etc.
@@ -21,7 +21,7 @@ use common_time::range::TimestampRange;
 use common_time::timestamp::TimeUnit;
 use datafusion_common::ScalarValue;
 use datafusion_expr::Expr;
-use futures::FutureExt;
+use futures::{FutureExt, StreamExt, TryStreamExt};
 use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
 use snafu::OptionExt;
 use store_api::metadata::RegionMetadataRef;
@@ -157,36 +157,50 @@ impl CompactionSstReaderBuilder<'_> {
        Ok(scan_input.build())
    }

+    /// Reads the footers of all inputs concurrently, keeping the input order.
    async fn collect_parquet_metadata(&self) -> Result<Vec<Arc<ParquetMetaData>>> {
-        let mut metadata = Vec::with_capacity(self.inputs.len());
+        let reads: Vec<_> = self
+            .inputs
+            .iter()
+            .map(|file_handle| self.read_parquet_metadata(file_handle))
+            .collect();
+        let metadata: Vec<_> = futures::stream::iter(reads)
+            .buffered(MAX_CONCURRENT_METADATA_READS)
+            .try_collect()
+            .await?;
+        Ok(metadata.into_iter().flatten().collect())
+    }

-        for file_handle in self.inputs {
-            let file_path =
-                file_handle.file_path(self.sst_layer.table_dir(), self.sst_layer.path_type());
-            let file_size = file_handle.meta_ref().file_size;
-            let parquet_metadata = match self
-                .sst_layer
-                .read_sst(file_handle.clone())
-                .cache(CacheStrategy::Compaction(self.cache.clone()))
-                .read_parquet_metadata(
-                    &file_path,
-                    file_size,
-                    &mut MetadataCacheMetrics::default(),
-                    PageIndexPolicy::default(),
-                )
-                .await
-                .map(|x| x.0.parquet_metadata())
-            {
-                Ok(x) => x,
-                Err(e) if e.is_object_not_found() => continue,
-                Err(e) => return Err(e),
-            };
-            metadata.push(parquet_metadata);
+    /// Reads the parquet metadata of `file_handle`, or `None` if the file is gone.
+    async fn read_parquet_metadata(
+        &self,
+        file_handle: &FileHandle,
+    ) -> Result<Option<Arc<ParquetMetaData>>> {
+        let file_path =
+            file_handle.file_path(self.sst_layer.table_dir(), self.sst_layer.path_type());
+        let file_size = file_handle.meta_ref().file_size;
+        match self
+            .sst_layer
+            .read_sst(file_handle.clone())
+            .cache(CacheStrategy::Compaction(self.cache.clone()))
+            .read_parquet_metadata(
+                &file_path,
+                file_size,
+                &mut MetadataCacheMetrics::default(),
+                PageIndexPolicy::default(),
+            )
+            .await
+        {
+            Ok(metadata) => Ok(Some(metadata.0.parquet_metadata())),
+            Err(e) if e.is_object_not_found() => Ok(None),
+            Err(e) => Err(e),
        }
-        Ok(metadata)
    }
 }

+/// Max number of input footers read at the same time.
+const MAX_CONCURRENT_METADATA_READS: usize = 8;
+
 /// Number of merged batches buffered between the merge task and the SST writer.
 const MERGED_BATCH_CHANNEL_SIZE: usize = 2;