fix: filter two-phase SST reads by exact sequence range

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/mito2/src/engine/scan_test.rs

@@ -1090,6 +1090,12 @@ async fn test_two_phase_series_scan() {
     request
         .options
         .insert("sst_format".to_string(), "flat".to_string());
+    request
+        .options
+        .insert("append_mode".to_string(), "true".to_string());
+    request
+        .options
+        .insert("preserve_row_sequence".to_string(), "true".to_string());
     let full_row_schema = test_util::rows_schema(&request);
     let mut encoded_primary_key_schema = full_row_schema[0].clone();
     encoded_primary_key_schema.column_name = PRIMARY_KEY_COLUMN_NAME.to_string();
@@ -1166,6 +1172,9 @@ async fn test_two_phase_series_scan() {
             ScanRequest {
                 // Internal metric identifiers must still be read for candidate discovery.
                 projection: Some(vec![2, 4, 5]),
+                memtable_min_sequence: Some(0),
+                memtable_max_sequence: Some(1),
+                exact_sequence_range: true,
                 distribution: Some(TimeSeriesDistribution::PerSeries),
                 ..Default::default()
             },
@@ -1229,22 +1238,8 @@ async fn test_two_phase_series_scan() {
         }
     }
     actual_rows.sort();
-    assert_eq!(
-        vec![
-            ("a".to_string(), 11, 1000),
-            ("a".to_string(), 12, 2000),
-            ("b".to_string(), 20, 1000),
-            ("b".to_string(), 21, 2000),
-            ("c".to_string(), 30, 1000),
-            ("d".to_string(), 40, 1000),
-        ],
-        actual_rows
-    );
-    assert_eq!(4, series_to_partition.len());
-    assert_eq!(Some(&0), series_to_partition.get("a"));
-    assert_eq!(Some(&0), series_to_partition.get("c"));
-    assert_eq!(Some(&2), series_to_partition.get("b"));
-    assert_eq!(Some(&2), series_to_partition.get("d"));
+    assert_eq!(vec![("a".to_string(), 10, 1000)], actual_rows);
+    assert_eq!(1, series_to_partition.len());
 
     scanner.reset_state();
     assert_eq!("two_phase", scanner.mode());
@@ -3089,6 +3084,7 @@ async fn test_exact_sequence_read_series_scan_per_series() {
     let engine = env
         .create_engine(MitoConfig {
             default_flat_format: true,
+            experimental_series_scan_v2: true,
             ..Default::default()
         })
         .await;
@@ -3136,22 +3132,31 @@ async fn test_exact_sequence_read_series_scan_per_series() {
         batches.pretty_print().unwrap()
     };
 
-    let result = scan_exact_series(Some(2), Some(6)).await;
-    let mut rows = result
-        .lines()
-        .filter(|l| l.starts_with("| ") && !l.starts_with("| tag_0 "))
-        .collect::<Vec<_>>();
-    rows.sort_unstable();
-    assert_eq!(
-        vec![
-            "| 2     | 2.0     | 1970-01-01T00:00:02 |",
-            "| 3     | 3.0     | 1970-01-01T00:00:03 |",
-            "| 4     | 4.0     | 1970-01-01T00:00:04 |",
-            "| 5     | 5.0     | 1970-01-01T00:00:05 |",
-        ],
-        rows,
-        "unexpected set for (2, 6] on PerSeries path:\n{result}"
-    );
+    for (min, max) in [(Some(2), Some(6)), (Some(0), Some(2))] {
+        let result = scan_exact_series(min, max).await;
+        let mut rows = result
+            .lines()
+            .filter(|l| l.starts_with("| ") && !l.starts_with("| tag_0 "))
+            .collect::<Vec<_>>();
+        rows.sort_unstable();
+        let expected = if min == Some(2) {
+            vec![
+                "| 2     | 2.0     | 1970-01-01T00:00:02 |",
+                "| 3     | 3.0     | 1970-01-01T00:00:03 |",
+                "| 4     | 4.0     | 1970-01-01T00:00:04 |",
+                "| 5     | 5.0     | 1970-01-01T00:00:05 |",
+            ]
+        } else {
+            vec![
+                "| 0     | 0.0     | 1970-01-01T00:00:00 |",
+                "| 1     | 1.0     | 1970-01-01T00:00:01 |",
+            ]
+        };
+        assert_eq!(
+            expected, rows,
+            "unexpected set for ({min:?}, {max:?}]:\n{result}"
+        );
+    }
 }
 
 /// Range-cache fingerprint: identical files and filters with different (C, H]

src/mito2/src/read/scan_util.rs

@@ -1513,7 +1513,7 @@ pub(crate) async fn scan_flat_file_ranges(
 /// batches have already been virtualized to the target-local file barrier, so
 /// they are filtered using the effective batch sequence. Local untrusted files
 /// pass through because exact capability excludes them before row filtering.
-fn filter_flat_batch_by_sequence(
+pub(crate) fn filter_flat_batch_by_sequence(
     record_batch: RecordBatch,
     sequence_range: Option<SequenceRange>,
     file_sequence_trusted: bool,

src/mito2/src/read/series_reader.rs

@@ -23,6 +23,7 @@ use futures::TryStreamExt;
 use mito_codec::row_converter::{PrimaryKeyFilter, SparsePrimaryKeyCodec};
 use snafu::ResultExt;
 use store_api::region_engine::PartitionRange;
+use store_api::storage::SequenceRange;
 use tokio::sync::Semaphore;
 
 #[cfg(feature = "enterprise")]
@@ -36,8 +37,8 @@ use crate::read::range_cache::{
 use crate::read::scan_region::StreamContext;
 use crate::read::scan_util::{
     PartitionMetrics, SplitRecordBatchStream, compute_average_batch_size,
-    compute_parallel_channel_size, new_filter_metrics, scan_flat_mem_ranges,
-    should_split_flat_batches_for_merge,
+    compute_parallel_channel_size, filter_flat_batch_by_sequence, new_filter_metrics,
+    scan_flat_mem_ranges, should_split_flat_batches_for_merge,
 };
 use crate::read::seq_scan::SeqScan;
 use crate::read::series_candidate::validate_metric_metadata;
@@ -447,6 +448,8 @@ async fn build_series_partition_range(
                 ranges,
                 filter.clone(),
                 codec.clone(),
+                stream_ctx.input.sequence_range,
+                stream_ctx.input.region_metadata().region_id,
             );
             sources.push(Box::pin(stream) as BoxedRecordBatchStream);
             continue;
@@ -494,6 +497,8 @@ fn scan_series_file_ranges(
     ranges: smallvec::SmallVec<[crate::sst::parquet::file_range::FileRange; 2]>,
     filter: MetricSeriesFilter,
     codec: SparsePrimaryKeyCodec,
+    sequence_range: Option<SequenceRange>,
+    region_id: store_api::storage::RegionId,
 ) -> impl futures::Stream<Item = Result<datatypes::arrow::record_batch::RecordBatch>> {
     try_stream! {
         let fetch_metrics = part_metrics
@@ -521,11 +526,22 @@ fn scan_series_file_ranges(
             part_metrics.inc_build_reader_cost(build_cost);
 
             let scan_start = Instant::now();
+            let file_sequence_trusted = range
+                .file_handle()
+                .is_effective_target_sequence_trusted(region_id);
             while let Some(record_batch) = reader.next_batch().await? {
                 reader_metrics.num_record_batches += 1;
                 reader_metrics.num_batches += 1;
                 reader_metrics.num_rows += record_batch.num_rows();
 
+                let Some(record_batch) = filter_flat_batch_by_sequence(
+                    record_batch,
+                    sequence_range,
+                    file_sequence_trusted,
+                )? else {
+                    continue;
+                };
+
                 let num_rows_before_filter = record_batch.num_rows();
                 let Some(record_batch) = range.precise_filter_flat(
                     record_batch,

src/mito2/src/read/series_scan.rs

@@ -180,7 +180,7 @@ impl SeriesScan {
     }
 
     fn supports_two_phase(input: &ScanInput) -> bool {
-        if input.sequence_range.is_some() || !is_sparse_metric_metadata(input.region_metadata()) {
+        if !is_sparse_metric_metadata(input.region_metadata()) {
             return false;
         }
         #[cfg(feature = "enterprise")]
@@ -1239,7 +1239,7 @@ mod tests {
     use crate::test_util::sst_util::sst_region_metadata_with_encoding;
 
     #[tokio::test]
-    async fn two_phase_eligibility_rejects_exact_sequence_range() {
+    async fn two_phase_eligibility_allows_exact_sequence_range() {
         let env = SchedulerEnv::new().await;
         let metadata = Arc::new(sst_region_metadata_with_encoding(
             store_api::codec::PrimaryKeyEncoding::Sparse,
@@ -1264,7 +1264,7 @@ mod tests {
             max: 2,
         }))
         .build();
-        assert!(!SeriesScan::supports_two_phase(&exact_sequence));
+        assert!(SeriesScan::supports_two_phase(&exact_sequence));
     }
 
     #[test]