fix(mito2): restore source splitting for cross-range series merges (#9818)

* fix: restore batched series merges across ranges Signed-off-by: WenyXu <wenymedia@gmail.com> * fix(mito2): split range sources before last-row selection Signed-off-by: WenyXu <wenymedia@gmail.com> * fix(mito2): avoid tiny splits for single-source cross-range merges Signed-off-by: WenyXu <wenymedia@gmail.com> --------- Signed-off-by: WenyXu <wenymedia@gmail.com>
2026-10-23 04:12:32 +00:00 · 2026-10-21 14:58:15 +00:00
parent ec4ef45f40
commit 27c531336a
5 changed files with 259 additions and 14 deletions
@@ -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::<Vec<_>>()
+        } 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] {
@@ -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>()
            + 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);
@@ -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;
@@ -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<StreamContext>,
        part_range: &PartitionRange,
-        compaction: bool,
+        read_mode: FlatReadMode,
        part_metrics: &PartitionMetrics,
        partition_pruner: Arc<PartitionPruner>,
        file_scan_semaphore: Option<Arc<Semaphore>>,
        merge_semaphore: Option<Arc<Semaphore>>,
    ) -> 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<StreamContext>,
    part_range: &PartitionRange,
-    compaction: bool,
+    read_mode: FlatReadMode,
    part_metrics: &PartitionMetrics,
    partition_pruner: Arc<PartitionPruner>,
    sources: &mut Vec<BoxedRecordBatchStream>,
    semaphore: Option<Arc<Semaphore>>,
 ) -> Result<Option<usize>> {
+    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
@@ -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,