fix: scope post-merge last row selection

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

src/metric-engine/src/engine/read.rs

@@ -318,7 +318,7 @@ mod test {
     use mito2::config::MitoConfig;
     use store_api::region_engine::{PrepareRequest, QueryScanContext};
     use store_api::region_request::{RegionFlushRequest, RegionPutRequest, RegionRequest};
-    use store_api::storage::TimeSeriesDistribution;
+    use store_api::storage::{TimeSeriesDistribution, TimeSeriesRowSelector};
 
     use super::*;
     use crate::config::EngineConfig;
@@ -429,6 +429,8 @@ mod test {
         let scan_req = ScanRequest {
             projection,
             filters: vec![],
+            series_row_selector: Some(TimeSeriesRowSelector::LastRow),
+            series_row_selector_after_merge: true,
             ..Default::default()
         };
 
@@ -444,6 +446,11 @@ mod test {
             &[11, 10, 9, 8, 0, 1, 4]
         );
         assert_eq!(scan_req.filters.len(), 1);
+        assert_eq!(
+            scan_req.series_row_selector,
+            Some(TimeSeriesRowSelector::LastRow)
+        );
+        assert!(scan_req.series_row_selector_after_merge);
         assert_eq!(
             scan_req.filters[0],
             logical_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)

src/mito2/src/cache.rs

@@ -2770,6 +2770,7 @@ mod tests {
                 filters: vec!["tag_0 = 1".to_string()],
                 time_filters: vec![],
                 series_row_selector: None,
+                series_row_selector_after_merge: false,
                 append_mode: false,
                 filter_deleted: true,
                 merge_mode: crate::region::options::MergeMode::LastRow,

src/mito2/src/engine/row_selector_test.rs

@@ -442,7 +442,17 @@ async fn test_last_row_returns_older_put_after_newest_delete_across_ssts() {
         .await;
         flush_region(&engine, region_id, None).await;
 
-        let scanner = last_row_scanner(&engine, region_id).await;
+        let scanner = engine
+            .scanner(
+                region_id,
+                ScanRequest {
+                    series_row_selector: Some(TimeSeriesRowSelector::LastRow),
+                    series_row_selector_after_merge: true,
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
         assert_eq!(2, scanner.num_files());
         assert_eq!(0, scanner.num_memtables());
         assert_eq!(

src/mito2/src/read/range_cache.rs

@@ -60,6 +60,7 @@ pub(crate) struct ScanRequestFingerprint {
     /// Filters with the time index column.
     time_filters: Option<Arc<Vec<String>>>,
     series_row_selector: Option<TimeSeriesRowSelector>,
+    series_row_selector_after_merge: bool,
     append_mode: bool,
     filter_deleted: bool,
     merge_mode: MergeMode,
@@ -87,6 +88,7 @@ pub(crate) struct ScanRequestFingerprintBuilder {
     pub(crate) filters: Vec<String>,
     pub(crate) time_filters: Vec<String>,
     pub(crate) series_row_selector: Option<TimeSeriesRowSelector>,
+    pub(crate) series_row_selector_after_merge: bool,
     pub(crate) append_mode: bool,
     pub(crate) filter_deleted: bool,
     pub(crate) merge_mode: MergeMode,
@@ -102,6 +104,7 @@ impl ScanRequestFingerprintBuilder {
             filters,
             time_filters,
             series_row_selector,
+            series_row_selector_after_merge,
             append_mode,
             filter_deleted,
             merge_mode,
@@ -117,6 +120,7 @@ impl ScanRequestFingerprintBuilder {
             }),
             time_filters: (!time_filters.is_empty()).then(|| Arc::new(time_filters)),
             series_row_selector,
+            series_row_selector_after_merge,
             append_mode,
             filter_deleted,
             merge_mode,
@@ -168,6 +172,7 @@ impl ScanRequestFingerprint {
             inner: Arc::clone(&self.inner),
             time_filters: None,
             series_row_selector: self.series_row_selector,
+            series_row_selector_after_merge: self.series_row_selector_after_merge,
             append_mode: self.append_mode,
             filter_deleted: self.filter_deleted,
             merge_mode: self.merge_mode,
@@ -182,6 +187,7 @@ impl ScanRequestFingerprint {
             inner: Arc::clone(&self.inner),
             time_filters: self.time_filters.clone(),
             series_row_selector: self.series_row_selector,
+            series_row_selector_after_merge: self.series_row_selector_after_merge,
             append_mode: self.append_mode,
             filter_deleted: self.filter_deleted,
             merge_mode: self.merge_mode,
@@ -196,6 +202,7 @@ impl ScanRequestFingerprint {
             inner: Arc::clone(&self.inner),
             time_filters: self.time_filters.clone(),
             series_row_selector: self.series_row_selector,
+            series_row_selector_after_merge: self.series_row_selector_after_merge,
             append_mode: self.append_mode,
             filter_deleted: self.filter_deleted,
             merge_mode: self.merge_mode,
@@ -868,6 +875,7 @@ pub fn bench_cache_flat_range_stream(
         filters: vec![],
         time_filters: vec![],
         series_row_selector: None,
+        series_row_selector_after_merge: false,
         append_mode: false,
         filter_deleted: false,
         merge_mode: MergeMode::LastRow,
@@ -925,6 +933,7 @@ mod tests {
         filters: Vec<String>,
         time_filters: Vec<String>,
         series_row_selector: Option<TimeSeriesRowSelector>,
+        series_row_selector_after_merge: bool,
         filter_deleted: bool,
         partition_expr_version: u64,
     ) -> ScanRequestFingerprint {
@@ -935,6 +944,7 @@ mod tests {
             filters,
             time_filters,
             series_row_selector,
+            series_row_selector_after_merge,
             append_mode: false,
             filter_deleted,
             merge_mode: MergeMode::LastRow,
@@ -949,7 +959,7 @@ mod tests {
         let key = RangeScanCacheKey {
             region_id,
             row_groups: vec![],
-            scan: test_scan_fingerprint(vec![], vec![], None, false, 0),
+            scan: test_scan_fingerprint(vec![], vec![], None, false, false, 0),
         };
 
         let metrics_set = ExecutionPlanMetricsSet::new();
@@ -1441,9 +1451,31 @@ mod tests {
     }
 
     #[test]
-    fn normalizes_and_clears_time_filters() {
+    fn selector_after_merge_changes_fingerprint() {
+        let ordinary = test_scan_fingerprint(
+            vec!["k0 = 'foo'".to_string()],
+            vec![],
+            Some(TimeSeriesRowSelector::LastRow),
+            false,
+            true,
+            0,
+        );
+        let after_merge = test_scan_fingerprint(
+            vec!["k0 = 'foo'".to_string()],
+            vec![],
+            Some(TimeSeriesRowSelector::LastRow),
+            true,
+            true,
+            0,
+        );
+
+        assert_ne!(ordinary, after_merge);
+    }
+
+    #[test]
+    fn true_selector_after_merge_is_preserved_by_fingerprint_transforms() {
         let normalized =
-            test_scan_fingerprint(vec!["k0 = 'foo'".to_string()], vec![], None, true, 0);
+            test_scan_fingerprint(vec!["k0 = 'foo'".to_string()], vec![], None, false, true, 0);
 
         assert!(normalized.time_filters().is_empty());
 
@@ -1452,16 +1484,23 @@ mod tests {
             vec!["ts >= 1000".to_string()],
             Some(TimeSeriesRowSelector::LastRow),
             true,
+            true,
             7,
         );
 
         let reset = fingerprint.without_time_filters();
+        let candidate = fingerprint.for_candidate_series();
+        let series_data = fingerprint.for_series_data(SeriesRange::new(0, 1).unwrap());
 
         assert_eq!(reset.read_columns(), fingerprint.read_columns());
         assert_eq!(reset.read_column_types(), fingerprint.read_column_types());
         assert_eq!(reset.filters(), fingerprint.filters());
         assert!(reset.time_filters().is_empty());
         assert_eq!(reset.series_row_selector, fingerprint.series_row_selector);
+        assert!(fingerprint.series_row_selector_after_merge);
+        assert!(reset.series_row_selector_after_merge);
+        assert!(candidate.series_row_selector_after_merge);
+        assert!(series_data.series_row_selector_after_merge);
         assert_eq!(reset.append_mode, fingerprint.append_mode);
         assert_eq!(reset.filter_deleted, fingerprint.filter_deleted);
         assert_eq!(reset.merge_mode, fingerprint.merge_mode);

src/mito2/src/read/scan_region.rs

@@ -574,6 +574,7 @@ impl ScanRegion {
             .with_filter_deleted(self.filter_deleted)
             .with_merge_mode(self.version.options.merge_mode())
             .with_series_row_selector(self.request.series_row_selector)
+            .with_series_row_selector_after_merge(self.request.series_row_selector_after_merge)
             .with_distribution(self.request.distribution)
             .with_explain_flat_format(
                 self.version.options.sst_format == Some(crate::sst::FormatType::Flat),
@@ -977,6 +978,8 @@ pub struct ScanInput {
     pub(crate) merge_mode: MergeMode,
     /// Hint to select rows from time series.
     pub(crate) series_row_selector: Option<TimeSeriesRowSelector>,
+    /// Whether the selector must run after cross-source merge and deduplication.
+    pub(crate) series_row_selector_after_merge: bool,
     /// Hint for the required distribution of the scanner.
     pub(crate) distribution: Option<TimeSeriesDistribution>,
     /// Whether the region's configured SST format is flat.
@@ -1051,6 +1054,7 @@ impl ScanInput {
                 filter_deleted: true,
                 merge_mode: MergeMode::default(),
                 series_row_selector: None,
+                series_row_selector_after_merge: false,
                 distribution: None,
                 explain_flat_format: false,
                 snapshot_sequence: None,
@@ -1199,6 +1203,7 @@ impl ScanInputBuilder {
                 filters,
                 time_filters,
                 series_row_selector: input.series_row_selector,
+                series_row_selector_after_merge: input.series_row_selector_after_merge,
                 append_mode: input.append_mode,
                 filter_deleted: input.filter_deleted,
                 merge_mode: input.merge_mode,
@@ -1381,6 +1386,13 @@ impl ScanInputBuilder {
         self
     }
 
+    /// Sets whether the row selector must run after merge and deduplication.
+    #[must_use]
+    pub(crate) fn with_series_row_selector_after_merge(mut self, after_merge: bool) -> Self {
+        self.input.series_row_selector_after_merge = after_merge;
+        self
+    }
+
     #[must_use]
     pub(crate) fn with_snapshot_sequence(
         mut self,
@@ -2504,6 +2516,7 @@ mod tests {
         .await
         .with_distribution(Some(TimeSeriesDistribution::PerSeries))
         .with_series_row_selector(Some(TimeSeriesRowSelector::LastRow))
+        .with_series_row_selector_after_merge(true)
         .with_merge_mode(MergeMode::LastNonNull)
         .with_filter_deleted(false)
         .build();
@@ -2529,6 +2542,7 @@ mod tests {
             ],
             time_filters: vec![col("ts").gt_eq(ts_lit(1000)).to_string()],
             series_row_selector: Some(TimeSeriesRowSelector::LastRow),
+            series_row_selector_after_merge: true,
             append_mode: false,
             filter_deleted: false,
             merge_mode: MergeMode::LastNonNull,
@@ -2537,6 +2551,7 @@ mod tests {
         }
         .build();
         assert_eq!(&expected, fingerprint);
+        assert!(input.series_row_selector_after_merge);
     }
 
     #[tokio::test]
@@ -2611,6 +2626,7 @@ mod tests {
             filters: vec![col("k0").eq(lit("foo")).to_string()],
             time_filters: vec![],
             series_row_selector: None,
+            series_row_selector_after_merge: false,
             append_mode: false,
             filter_deleted: true,
             merge_mode: MergeMode::LastRow,

src/mito2/src/read/scan_util.rs

@@ -1571,8 +1571,13 @@ pub fn build_flat_file_range_scan_stream(
             let build_reader_start = Instant::now();
             let Some(mut reader) = range
                 .flat_reader(
-                    // LastRow must run after cross-source merge and deduplication.
-                    None,
+                    if stream_ctx.input.sequence_range.is_none()
+                        && !stream_ctx.input.series_row_selector_after_merge
+                    {
+                        stream_ctx.input.series_row_selector
+                    } else {
+                        None
+                    },
                     fetch_metrics.as_deref(),
                 )
                 .await?

src/mito2/src/read/series_scan.rs

@@ -180,7 +180,7 @@ impl SeriesScan {
     }
 
     fn supports_two_phase(input: &ScanInput) -> bool {
-        if !is_sparse_metric_metadata(input.region_metadata()) {
+        if input.sequence_range.is_some() || !is_sparse_metric_metadata(input.region_metadata()) {
             return false;
         }
         #[cfg(feature = "enterprise")]
@@ -1214,6 +1214,39 @@ mod tests {
     use datatypes::arrow::record_batch::RecordBatch;
 
     use super::*;
+    use crate::read::flat_projection::FlatProjectionMapper;
+    use crate::read::scan_region::PredicateGroup;
+    use crate::test_util::scheduler_util::SchedulerEnv;
+    use crate::test_util::sst_util::sst_region_metadata_with_encoding;
+
+    #[tokio::test]
+    async fn two_phase_eligibility_rejects_exact_sequence_range() {
+        let env = SchedulerEnv::new().await;
+        let metadata = Arc::new(sst_region_metadata_with_encoding(
+            store_api::codec::PrimaryKeyEncoding::Sparse,
+        ));
+        let predicate = PredicateGroup::default();
+
+        let eligible = ScanInput::builder(
+            env.access_layer.clone(),
+            FlatProjectionMapper::new(&metadata, [0]).unwrap(),
+        )
+        .with_predicate(predicate.clone())
+        .build();
+        assert!(SeriesScan::supports_two_phase(&eligible));
+
+        let exact_sequence = ScanInput::builder(
+            env.access_layer.clone(),
+            FlatProjectionMapper::new(&metadata, [0]).unwrap(),
+        )
+        .with_predicate(predicate)
+        .with_sequence_range(Some(store_api::storage::SequenceRange::GtLtEq {
+            min: 1,
+            max: 2,
+        }))
+        .build();
+        assert!(!SeriesScan::supports_two_phase(&exact_sequence));
+    }
 
     #[test]
     fn candidate_distributor_stops_after_all_receivers_close() {

src/query/src/dummy_catalog.rs

@@ -286,6 +286,13 @@ impl DummyTableProvider {
         self.scan_request.lock().unwrap().series_row_selector = Some(selector);
     }
 
+    /// Sets the instant-derived LastRow hint, which must run after merge and deduplication.
+    pub fn with_last_row_selector_after_merge_hint(&self) {
+        let mut request = self.scan_request.lock().unwrap();
+        request.series_row_selector = Some(TimeSeriesRowSelector::LastRow);
+        request.series_row_selector_after_merge = true;
+    }
+
     /// Clones this provider for one logical table-scan use-site.
     ///
     /// The optimizer may attach different hints to scans that share the same

src/query/src/optimizer/scan_hint.rs

@@ -124,7 +124,7 @@ impl ScanHintRule {
             // Apply the instant-derived hint after the aggregate hint. Both
             // select LastRow today, and this ordering preserves the existing
             // aggregate selector when the instant guard rejects a scan.
-            adapter.with_time_series_selector_hint(TimeSeriesRowSelector::LastRow);
+            adapter.with_last_row_selector_after_merge_hint();
         }
         table_scan.source =
             std::sync::Arc::new(DefaultTableSource::new(std::sync::Arc::new(adapter)));
@@ -636,6 +636,8 @@ mod test {
             scan_requests(&rewritten)[0].series_row_selector,
             Some(TimeSeriesRowSelector::LastRow)
         );
+        assert!(scan_requests(&rewritten)[0].series_row_selector_after_merge);
+        assert!(!provider.scan_request().series_row_selector_after_merge);
         assert_eq!(provider.scan_request().series_row_selector, None);
     }
 
@@ -649,7 +651,9 @@ mod test {
             .data;
 
         assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
+        assert!(!scan_requests(&rewritten)[0].series_row_selector_after_merge);
         assert_eq!(provider.scan_request().series_row_selector, None);
+        assert!(!provider.scan_request().series_row_selector_after_merge);
     }
 
     #[test]
@@ -665,6 +669,7 @@ mod test {
             scan_requests(&rewritten)[0].series_row_selector,
             Some(TimeSeriesRowSelector::LastRow)
         );
+        assert!(scan_requests(&rewritten)[0].series_row_selector_after_merge);
     }
 
     #[test]
@@ -680,6 +685,7 @@ mod test {
             scan_requests(&rewritten)[0].series_row_selector,
             Some(TimeSeriesRowSelector::LastRow)
         );
+        assert!(scan_requests(&rewritten)[0].series_row_selector_after_merge);
     }
 
     #[test]
@@ -693,6 +699,7 @@ mod test {
             .data;
 
         assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
+        assert!(!scan_requests(&rewritten)[0].series_row_selector_after_merge);
     }
 
     #[test]
@@ -705,6 +712,7 @@ mod test {
             .data;
 
         assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
+        assert!(!scan_requests(&rewritten)[0].series_row_selector_after_merge);
     }
 
     #[test]
@@ -717,6 +725,7 @@ mod test {
             .data;
 
         assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
+        assert!(!scan_requests(&rewritten)[0].series_row_selector_after_merge);
     }
 
     #[test]
@@ -729,6 +738,7 @@ mod test {
             .data;
 
         assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
+        assert!(!scan_requests(&rewritten)[0].series_row_selector_after_merge);
     }
 
     #[test]
@@ -830,6 +840,7 @@ mod test {
             scan_requests(&rewritten)[0].series_row_selector,
             Some(TimeSeriesRowSelector::LastRow)
         );
+        assert!(scan_requests(&rewritten)[0].series_row_selector_after_merge);
     }
 
     #[test]
@@ -852,7 +863,10 @@ mod test {
             Some(TimeSeriesRowSelector::LastRow)
         );
         assert_eq!(requests[1].series_row_selector, None);
+        assert!(requests[0].series_row_selector_after_merge);
+        assert!(!requests[1].series_row_selector_after_merge);
         assert_eq!(provider.scan_request().series_row_selector, None);
+        assert!(!provider.scan_request().series_row_selector_after_merge);
     }
 
     #[test]
@@ -916,7 +930,11 @@ mod test {
         let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data;
 
         let scan_req = scan_requests(&rewritten)[0].clone();
-        let _ = scan_req.series_row_selector.unwrap();
+        assert_eq!(
+            scan_req.series_row_selector,
+            Some(TimeSeriesRowSelector::LastRow)
+        );
+        assert!(!scan_req.series_row_selector_after_merge);
     }
 
     #[test]

src/store-api/src/storage/requests.rs

@@ -112,6 +112,11 @@ pub struct ScanRequest {
     pub limit: Option<usize>,
     /// Optional hint to select rows from time-series.
     pub series_row_selector: Option<TimeSeriesRowSelector>,
+    /// Whether the row selector must run only after cross-source merge and deduplication.
+    ///
+    /// This is required by instant-derived LastRow scans. Ordinary selector users
+    /// keep the default `false` and may use source-local selector optimization.
+    pub series_row_selector_after_merge: bool,
     /// Optional constraint on the sequence number of the rows to read.
     /// If set, only rows with a sequence number **lesser or equal** to this value
     /// will be returned.
@@ -205,6 +210,13 @@ impl Display for ScanRequest {
                 series_row_selector
             )?;
         }
+        if self.series_row_selector_after_merge {
+            write!(
+                f,
+                "{}series_row_selector_after_merge: true",
+                delimiter.as_str()
+            )?;
+        }
         if let Some(sequence) = &self.memtable_max_sequence {
             write!(f, "{}sequence: {}", delimiter.as_str(), sequence)?;
         }
@@ -324,12 +336,13 @@ mod tests {
 
         let request = ScanRequest {
             snapshot_on_scan: true,
+            series_row_selector_after_merge: true,
             exact_sequence_range: true,
             ..Default::default()
         };
         assert_eq!(
             request.to_string(),
-            "ScanRequest { snapshot_on_scan: true, exact_sequence_range: true }"
+            "ScanRequest { series_row_selector_after_merge: true, snapshot_on_scan: true, exact_sequence_range: true }"
         );
 
         let request = ScanRequest {