From 27c531336a3b7dece7712c682c53e277e4626d28 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 21 Oct 2026 14:58:15 +0000 Subject: [PATCH] fix(mito2): restore source splitting for cross-range series merges (#9818) * fix: restore batched series merges across ranges Signed-off-by: WenyXu * fix(mito2): split range sources before last-row selection Signed-off-by: WenyXu * fix(mito2): avoid tiny splits for single-source cross-range merges Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/mito2/src/engine/scan_test.rs | 192 ++++++++++++++++++++++++++++++ src/mito2/src/read/range_cache.rs | 20 ++++ src/mito2/src/read/scan_util.rs | 2 +- src/mito2/src/read/seq_scan.rs | 50 ++++++-- src/mito2/src/read/series_scan.rs | 9 +- 5 files changed, 259 insertions(+), 14 deletions(-) diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 39613efb383..81b9b016772 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -1489,6 +1489,198 @@ async fn test_series_scan_with_format(flat_format: bool) { ); } +/// Cross-range merges split long series but preserve larger batches for short series. +#[tokio::test] +async fn test_series_scan_single_source_range_batches() { + for rows_per_series in [16, 49, 50, 64] { + check_series_scan_single_source_range_batches(rows_per_series).await; + } +} + +async fn check_series_scan_single_source_range_batches(rows_per_series: usize) { + let mut env = TestEnv::with_prefix("test_series_scan_single_source_range_batches").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: true, + range_result_cache_size: ReadableSize::mb(64), + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.time_window", "1h") + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let mut expected = Vec::new(); + for start in [0, 3600] { + let mut rows = Vec::new(); + for key in ["a", "b", "c", "d"] { + rows.extend(test_util::build_rows_for_key( + key, + start, + start + rows_per_series, + 0, + )); + expected.extend( + (0..rows_per_series) + .map(|i| (key.to_string(), i as f64, (start + i) as i64 * 1000)), + ); + } + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows, + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + } + expected.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2))); + + // Exercise both newly read and cached range streams at the same merge boundary. + for last_row in [false, true] { + let expected = if last_row { + expected + .iter() + .filter(|(_, _, ts)| *ts == (3600 + rows_per_series as i64 - 1) * 1000) + .cloned() + .collect::>() + } else { + expected.clone() + }; + for with_cache in [false, true] { + let request = ScanRequest { + series_row_selector: last_row + .then_some(TimeSeriesRowSelector::LastRow { after_merge: true }), + // Range caching requires a tag predicate; this one retains every row. + filters: if with_cache { + vec![col("tag_0").gt_eq(lit("a"))] + } else { + vec![] + }, + ..Default::default() + }; + // Populate the ordinary scan cache first. Its unsplit batches must not + // replace the batches prepared for a later cross-range merge. + if with_cache { + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let Scanner::Seq(scanner) = + engine.scanner(region_id, request.clone()).await.unwrap() + else { + panic!("expected sequential scan"); + }; + let ctx = store_api::region_engine::QueryScanContext { + explain_verbose: true, + ..Default::default() + }; + let stream = scanner + .scan_partition(&ctx, &ExecutionPlanMetricsSet::default(), 0) + .unwrap(); + RecordBatches::try_collect(stream).await.unwrap(); + if format!("{}", VerboseDisplay(scanner)).contains("\"range_cache_hit\":") { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + let mut first_scan = true; + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let scanner = engine + .scanner( + region_id, + ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..request.clone() + }, + ) + .await + .unwrap(); + let Scanner::Series(scanner) = scanner else { + panic!("expected series scan"); + }; + assert_eq!(2, scanner.properties().partitions[0].len()); + + let ctx = store_api::region_engine::QueryScanContext { + explain_verbose: true, + ..Default::default() + }; + let metrics = ExecutionPlanMetricsSet::default(); + let stream = scanner.scan_partition(&ctx, &metrics, 0).unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let actual = collect_and_assert_partition_rows( + batches.schema(), + vec![batches.iter().cloned().collect()], + ); + assert_eq!(expected, actual); + let display = format!("{}", VerboseDisplay(scanner)); + let details: serde_json::Value = + serde_json::from_str(&display[display.find('{').unwrap()..]).unwrap(); + let partitions = details["metrics_per_partition"].as_array().unwrap(); + let rows: u64 = partitions + .iter() + .map(|p| { + p["metrics"]["merge_metrics"]["num_fetch_by_rows"] + .as_u64() + .unwrap_or(0) + }) + .sum(); + if last_row { + // LastRow coalesces the source splits before the cross-range + // merge. Do not turn its output into one batch per series. + let batches: u64 = partitions + .iter() + .map(|p| { + p["metrics"]["merge_metrics"]["num_fetch_by_batches"] + .as_u64() + .unwrap_or(0) + }) + .sum(); + assert!(batches <= 2, "{display}"); + } else if rows_per_series >= 50 { + assert_eq!(0, rows, "{display}"); + } else { + // Short series retain batches spanning multiple series instead + // of sending one tiny batch per series through the distributor. + assert!(rows > 0, "{display}"); + } + if with_cache && first_scan { + assert!( + partitions.iter().all(|p| { + p["metrics"]["range_cache_hit"].as_u64().unwrap_or(0) == 0 + }), + "ordinary scan cache was reused: {display}" + ); + } + first_scan = false; + if !with_cache + || partitions + .iter() + .any(|p| p["metrics"]["range_cache_hit"].as_u64().unwrap_or(0) > 0) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + } +} + #[tokio::test] async fn test_two_phase_series_scan() { for use_index in [false, true] { diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 1108b074697..081efa7bec8 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -76,6 +76,9 @@ pub(crate) struct ScanRequestFingerprint { #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] enum RangeScanStage { Data, + /// Source batches may be split for a later cross-range merge. Keep their + /// cached batch boundaries separate from ordinary sequential scans. + DataForCrossRangeMerge, CandidateSeries, SeriesData(SeriesRange), } @@ -237,6 +240,12 @@ pub(crate) struct RangeScanCacheKey { } impl RangeScanCacheKey { + pub(crate) fn for_cross_range_merge(mut self) -> Self { + debug_assert_eq!(self.scan.stage, RangeScanStage::Data); + self.scan.stage = RangeScanStage::DataForCrossRangeMerge; + self + } + pub(crate) fn estimated_size(&self) -> usize { mem::size_of::() + self.row_groups.capacity() * mem::size_of::<(FileId, i64)>() @@ -583,6 +592,11 @@ fn build_range_cache_key_inner( let scan = match stage { Some(RangeScanStage::CandidateSeries) => scan.for_candidate_series(), Some(RangeScanStage::SeriesData(assignment)) => scan.for_series_data(assignment), + Some(RangeScanStage::DataForCrossRangeMerge) => { + let mut scan = scan; + scan.stage = RangeScanStage::DataForCrossRangeMerge; + scan + } Some(RangeScanStage::Data) | None => scan, }; @@ -1208,6 +1222,12 @@ mod tests { let series_key_0 = build_series_range_cache_key(&ctx, &part_range, range_0).unwrap(); let series_key_1 = build_series_range_cache_key(&ctx, &part_range, range_1).unwrap(); + let merge_key = data_key.clone().for_cross_range_merge(); + assert_ne!(merge_key, data_key); + assert_ne!(merge_key.scan, candidate_key.scan); + assert_ne!(merge_key.scan, series_key_0.scan); + assert_eq!(merge_key.row_groups, data_key.row_groups); + assert_eq!(merge_key, data_key.clone().for_cross_range_merge()); assert_ne!(data_key.scan, candidate_key.scan); assert_ne!(data_key.scan, series_key_0.scan); assert_ne!(candidate_key.scan, series_key_0.scan); diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index 53967949387..b122704a71e 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -1203,7 +1203,7 @@ const SPLIT_ROW_THRESHOLD: u64 = DEFAULT_ROW_GROUP_SIZE as u64; const NUM_SERIES_THRESHOLD: u64 = 10240; /// Minimum batch size after splitting. The batch size is less than 60 because a series may only have /// 60 samples per hour. -const BATCH_SIZE_THRESHOLD: u64 = 50; +pub(crate) const BATCH_SIZE_THRESHOLD: u64 = 50; /// Minimum estimated rows per series for compaction to split batches. Shorter series /// would make compaction write one small batch per series. pub(crate) const MIN_COMPACTION_SPLIT_ROWS: usize = DEFAULT_READ_BATCH_SIZE / 2; diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index ecfa41bcf2e..8db83ed1845 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -46,15 +46,26 @@ use crate::read::range_cache::{ }; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ - MIN_COMPACTION_SPLIT_ROWS, PartitionMetrics, PartitionMetricsList, SplitRecordBatchStream, - compute_parallel_channel_size, scan_flat_file_ranges, scan_flat_mem_ranges, - should_split_flat_batches_for_merge, + BATCH_SIZE_THRESHOLD, MIN_COMPACTION_SPLIT_ROWS, PartitionMetrics, PartitionMetricsList, + SplitRecordBatchStream, compute_parallel_channel_size, scan_flat_file_ranges, + scan_flat_mem_ranges, should_split_flat_batches_for_merge, }; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; use crate::read::{BoxedRecordBatchStream, ScannerMetrics, scan_util}; use crate::region::options::MergeMode; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; +/// Describes where flat source batches will be merged. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FlatReadMode { + /// Each range is consumed independently. + Sequential, + /// Range outputs feed a later SeriesScan merge. + CrossRangeMerge, + /// Sources from all ranges feed the compaction merge. + Compaction, +} + /// Scans a region and returns rows in a sorted sequence. /// /// The output order is always `order by primary keys, time index` inside every @@ -168,7 +179,7 @@ impl SeqScan { build_flat_sources( stream_ctx, part_range, - true, + FlatReadMode::Compaction, part_metrics, partition_pruner.clone(), &mut sources, @@ -277,13 +288,19 @@ impl SeqScan { pub(crate) async fn build_flat_partition_range_read( stream_ctx: &Arc, part_range: &PartitionRange, - compaction: bool, + read_mode: FlatReadMode, part_metrics: &PartitionMetrics, partition_pruner: Arc, file_scan_semaphore: Option>, merge_semaphore: Option>, ) -> Result<(BoxedRecordBatchStream, usize)> { - let cache_key = build_range_cache_key(stream_ctx, part_range); + let cache_key = build_range_cache_key(stream_ctx, part_range).map(|key| { + if read_mode == FlatReadMode::CrossRangeMerge { + key.for_cross_range_merge() + } else { + key + } + }); if let Some(key) = cache_key.as_ref() { if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { @@ -297,7 +314,7 @@ impl SeqScan { let split_batch_size = build_flat_sources( stream_ctx, part_range, - compaction, + read_mode, part_metrics, partition_pruner, &mut sources, @@ -410,7 +427,11 @@ impl SeqScan { let (mut reader, _) = Self::build_flat_partition_range_read( &stream_ctx, &part_range, - compaction, + if compaction { + FlatReadMode::Compaction + } else { + FlatReadMode::Sequential + }, &part_metrics, partition_pruner.clone(), file_scan_semaphore.clone(), @@ -633,12 +654,13 @@ impl fmt::Debug for SeqScan { pub(crate) async fn build_flat_sources( stream_ctx: &Arc, part_range: &PartitionRange, - compaction: bool, + read_mode: FlatReadMode, part_metrics: &PartitionMetrics, partition_pruner: Arc, sources: &mut Vec, semaphore: Option>, ) -> Result> { + let compaction = read_mode == FlatReadMode::Compaction; // Gets range meta. let range_meta = &stream_ctx.ranges[part_range.identifier]; #[cfg(debug_assertions)] @@ -738,11 +760,17 @@ pub(crate) async fn build_flat_sources( } } - // Queries merge one range at a time; compaction merges sources from all ranges. + // SeriesScan also merges range outputs. Split at the source, before deduplication + // and LastRow selection, so their output batching is preserved by the range cache. let needs_merge = ordered_sources.iter().flatten().count() > 1 || (compaction && stream_ctx.ranges.len() > 1); - let split_batch_size = if needs_merge { + let split_batch_size = if needs_merge || read_mode == FlatReadMode::CrossRangeMerge { should_split_flat_batches_for_merge(stream_ctx, range_meta).filter(|rows_per_series| { + // A single-source range only benefits at the downstream merge. Tiny + // splits add cache replay, channel and distribution work per batch. + if !needs_merge && *rows_per_series < BATCH_SIZE_THRESHOLD as usize { + return false; + } // Compaction writes every merged batch out, so splitting short series // would hand the SST writer and indexer one tiny batch per series. !compaction || *rows_per_series >= MIN_COMPACTION_SPLIT_ROWS diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index a77cbebb557..64d5ceb07f8 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -53,7 +53,7 @@ use crate::read::scan_util::{ PartitionMetrics, PartitionMetricsList, SeriesDistributorMetrics, compute_average_batch_size, compute_parallel_channel_size, }; -use crate::read::seq_scan::SeqScan; +use crate::read::seq_scan::{FlatReadMode, SeqScan}; use crate::read::series_candidate::{SeriesCandidateScanner, is_sparse_metric_metadata}; use crate::read::series_reader::{AssignedSeriesBatch, SeriesBatchCollector, SeriesReader}; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; @@ -719,6 +719,7 @@ impl SeriesDistributor { // Create PartitionPruner covering all partitions let all_partition_ranges: Vec<_> = self.partitions.iter().flatten().cloned().collect(); + let needs_merge = all_partition_ranges.len() > 1; let partition_pruner = Arc::new(PartitionPruner::new( self.pruner.clone(), &all_partition_ranges, @@ -751,7 +752,11 @@ impl SeriesDistributor { SeqScan::build_flat_partition_range_read( &stream_ctx, &part_range, - false, + if needs_merge { + FlatReadMode::CrossRangeMerge + } else { + FlatReadMode::Sequential + }, &part_metrics, partition_pruner, file_scan_semaphore,