perf(mito2): postpone covered time index filters (#8998)

* perf(mito2): postpone covered time index filters

Signed-off-by: evenyag <realevenyag@gmail.com>

* perf(mito2): reuse implied time range for prefilter

Signed-off-by: evenyag <realevenyag@gmail.com>

* refactor(mito2): build finalized scan inputs

Signed-off-by: evenyag <realevenyag@gmail.com>

* fix(mito2): reject empty implied time filters

Signed-off-by: evenyag <realevenyag@gmail.com>

* fix(mito2): guard last row shortcut with remaining filters

Signed-off-by: evenyag <realevenyag@gmail.com>

---------

Signed-off-by: evenyag <realevenyag@gmail.com>

src/mito2/src/compaction/reader.rs

@@ -127,7 +127,7 @@ impl CompactionSstReaderBuilder<'_> {
             &targets,
         )?;
 
-        let mut scan_input = ScanInput::new(self.sst_layer, mapper)
+        let mut scan_input = ScanInput::builder(self.sst_layer, mapper)
             .with_json2_rewrite_targets(targets)
             .with_files(self.inputs.to_vec())
             .with_compaction(true)
@@ -147,7 +147,7 @@ impl CompactionSstReaderBuilder<'_> {
                 scan_input.with_predicate(time_range_to_predicate(time_range, &self.metadata)?);
         }
 
-        Ok(scan_input)
+        Ok(scan_input.build())
     }
 
     async fn collect_parquet_metadata(&self) -> Result<Vec<Arc<ParquetMetaData>>> {

src/mito2/src/engine/row_selector_test.rs

@@ -15,6 +15,7 @@
 use api::v1::Rows;
 use common_base::readable_size::ReadableSize;
 use common_recordbatch::RecordBatches;
+use datafusion_common::ScalarValue;
 use datafusion_expr::{col, lit};
 use store_api::region_engine::RegionEngine;
 use store_api::region_request::RegionRequest;
@@ -131,6 +132,67 @@ async fn scan_last_row(
     sort_batches_and_print(&batches, &["tag_0", "ts"])
 }
 
+async fn new_flat_last_row_engine(
+    selector_result_cache_size: ReadableSize,
+) -> (TestEnv, MitoEngine, RegionId) {
+    let mut env = TestEnv::new().await;
+    let engine = env
+        .create_engine(MitoConfig {
+            selector_result_cache_size,
+            ..Default::default()
+        })
+        .await;
+    let region_id = RegionId::new(1, 1);
+
+    env.get_schema_metadata_manager()
+        .register_region_table_info(
+            region_id.table_id(),
+            "test_table",
+            "test_catalog",
+            "test_schema",
+            None,
+            env.get_kv_backend(),
+        )
+        .await;
+
+    let request = CreateRequestBuilder::new()
+        .insert_option("sst_format", "flat")
+        .build();
+    let column_schemas = rows_schema(&request);
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+
+    let rows = Rows {
+        schema: column_schemas,
+        rows: build_rows_for_key("a", 0, 11, 0),
+    };
+    put_rows(&engine, region_id, rows).await;
+    flush_region(&engine, region_id, Some(16)).await;
+
+    (env, engine, region_id)
+}
+
+fn mixed_time_filters() -> Vec<datafusion_expr::Expr> {
+    let ts_lit = |value| lit(ScalarValue::TimestampMillisecond(Some(value), None));
+    vec![col("ts").gt_eq(ts_lit(0)), col("ts").not_eq(ts_lit(10_000))]
+}
+
+const LAST_ROW_AT_NINE: &str = "\
++-------+---------+---------------------+
+| tag_0 | field_0 | ts                  |
++-------+---------+---------------------+
+| a     | 9.0     | 1970-01-01T00:00:09 |
++-------+---------+---------------------+";
+
+const LAST_ROW_AT_TEN: &str = "\
++-------+---------+---------------------+
+| tag_0 | field_0 | ts                  |
++-------+---------+---------------------+
+| a     | 10.0    | 1970-01-01T00:00:10 |
++-------+---------+---------------------+";
+
 #[tokio::test]
 async fn test_last_row_append_mode_disabled() {
     test_last_row(false, false).await;
@@ -151,6 +213,28 @@ async fn test_last_row_flat_format_append_mode_enabled() {
     test_last_row(true, true).await;
 }
 
+#[tokio::test]
+async fn test_last_row_flat_format_non_tag_filter_without_selector_cache() {
+    let (_env, engine, region_id) = new_flat_last_row_engine(ReadableSize(0)).await;
+
+    let filtered = scan_last_row(&engine, region_id, mixed_time_filters()).await;
+    assert_eq!(LAST_ROW_AT_NINE, filtered);
+}
+
+#[tokio::test]
+async fn test_last_row_flat_format_non_tag_filter_does_not_reuse_selector_cache() {
+    let (_env, engine, region_id) = new_flat_last_row_engine(ReadableSize::mb(1)).await;
+
+    let unfiltered = scan_last_row(&engine, region_id, vec![]).await;
+    assert_eq!(LAST_ROW_AT_TEN, unfiltered);
+
+    let filtered = scan_last_row(&engine, region_id, mixed_time_filters()).await;
+    assert_eq!(LAST_ROW_AT_NINE, filtered);
+
+    let unfiltered = scan_last_row(&engine, region_id, vec![]).await;
+    assert_eq!(LAST_ROW_AT_TEN, unfiltered);
+}
+
 #[tokio::test]
 async fn test_last_row_flat_format_prefilter_does_not_poison_selector_cache() {
     let mut env = TestEnv::new().await;

src/mito2/src/read/pruner.rs

@@ -841,9 +841,10 @@ mod tests {
             })
             .collect();
 
-        let input = ScanInput::new(env.access_layer.clone(), mapper)
+        let input = ScanInput::builder(env.access_layer.clone(), mapper)
             .with_files(files)
-            .with_append_mode(true);
+            .with_append_mode(true)
+            .build();
         let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input));
         let pruner = Arc::new(Pruner::new_with_options(
             stream_ctx,
@@ -898,10 +899,11 @@ mod tests {
             })
             .collect();
 
-        let input = ScanInput::new(env.access_layer.clone(), mapper)
+        let input = ScanInput::builder(env.access_layer.clone(), mapper)
             .with_files(files)
             .with_predicate(predicate)
-            .with_append_mode(true);
+            .with_append_mode(true)
+            .build();
         let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input));
         let pruner = Arc::new(Pruner::new(stream_ctx, 1));
         (env, pruner)

src/mito2/src/read/range_cache.rs

@@ -321,8 +321,9 @@ pub(crate) fn collect_partition_range_row_groups(
 
 /// Returns the timestamp range where all time-only predicates are guaranteed true.
 ///
-/// Returns `Some(min_to_max)` for empty input (vacuously true everywhere).
-/// Returns `None` if any expression contains an unsupported shape: `OR`, `NOT`,
+/// Returns `None` for empty input because there is no time-filter implication
+/// that cache-key normalization or prefilter postponement can use. It also
+/// returns `None` if any expression contains an unsupported shape: `OR`, `NOT`,
 /// `IN`, non-literal RHS, unsupported operator, column-name mismatch, an `=`
 /// literal that cannot be represented exactly in the column unit, or overflow
 /// during bound adjustment.
@@ -340,6 +341,10 @@ pub(crate) fn implied_time_range_from_exprs(
     ts_col_unit: TimeUnit,
     exprs: &[&Expr],
 ) -> Option<TimestampRange> {
+    if exprs.is_empty() {
+        return None;
+    }
+
     let mut acc = TimestampRange::min_to_max();
     for expr in exprs {
         let r = implied_time_range_from_expr(ts_col_name, ts_col_unit, expr)?;
@@ -532,7 +537,7 @@ fn build_range_cache_key_inner(
         return None;
     }
 
-    let fingerprint = stream_ctx.scan_fingerprint.as_ref()?;
+    let fingerprint = stream_ctx.input.scan_fingerprint()?;
 
     // Dyn filters can change at runtime, so we can't cache when they're present.
     let has_dyn_filters = stream_ctx
@@ -551,12 +556,13 @@ fn build_range_cache_key_inner(
 
     // If the implied range covers this partition's `FileTimeRange`, drop
     // time-only predicates from the cache key so that queries with different
-    // but equally-covering time bounds share an entry. `None` means some
-    // time-only predicate had an unsupported shape (e.g. `OR`), so we keep
-    // them in the key.
+    // but equally-covering time bounds share an entry. `None` means there is no
+    // analyzable time-only predicate or some predicate had an unsupported shape
+    // (e.g. `OR`), so we keep the fingerprint unchanged. When `time_filters` is
+    // already empty, cloning the fingerprint is equivalent to stripping them.
     let range_meta = &stream_ctx.ranges[part_range.identifier];
     let (file_min, file_max) = range_meta.time_range;
-    let covers = match &stream_ctx.scan_implied_time_range {
+    let covers = match stream_ctx.input.implied_time_range() {
         // An empty implied range can never cover a non-empty file range, so
         // short-circuit.
         Some(implied) if !implied.is_empty() => {
@@ -982,11 +988,12 @@ mod tests {
             partition_time_range.0.value(),
             partition_time_range.1.value(),
         );
-        let input = ScanInput::new(env.access_layer.clone(), mapper)
+        let input = ScanInput::builder(env.access_layer.clone(), mapper)
             .with_predicate(predicate)
             .with_time_range(query_time_range)
             .with_files(vec![file])
-            .with_cache(test_cache_strategy());
+            .with_cache(test_cache_strategy())
+            .build();
         let range_meta = RangeMeta {
             time_range: partition_time_range,
             indices: smallvec![SourceIndex {
@@ -1000,16 +1007,9 @@ mod tests {
             num_rows: 10,
         };
         let partition_range = range_meta.new_partition_range(0);
-        let (scan_fingerprint, scan_implied_time_range) =
-            match crate::read::scan_region::build_scan_fingerprint(&input) {
-                Some(b) => (Some(b.fingerprint), b.implied_time_range),
-                None => (None, None),
-            };
         let stream_ctx = StreamContext {
             input,
             ranges: vec![range_meta],
-            scan_fingerprint,
-            scan_implied_time_range,
             query_start: Instant::now(),
         };
 
@@ -1219,7 +1219,7 @@ mod tests {
         )
         .await;
 
-        assert!(ctx_a.scan_implied_time_range.is_none());
+        assert!(ctx_a.input.implied_time_range().is_none());
         let key_a = build_range_cache_key(&ctx_a, &part_a).unwrap();
         let key_b = build_range_cache_key(&ctx_b, &part_b).unwrap();
         assert_ne!(key_a.scan, key_b.scan);
@@ -1241,13 +1241,17 @@ mod tests {
         );
 
         let (mut ctx, part_range) = new_stream_context(
-            vec![col("ts").gt_eq(ts_lit(1500)), col("k0").eq(lit("foo"))],
+            vec![
+                col("ts").gt_eq(ts_lit(1500)),
+                col("ts").lt(ts_lit(1500)),
+                col("k0").eq(lit("foo")),
+            ],
             TimestampRange::with_unit(1500, 3000, TimeUnit::Millisecond),
             partition,
         )
         .await;
 
-        ctx.scan_implied_time_range = Some(TimestampRange::empty());
+        assert!(ctx.input.implied_time_range().unwrap().is_empty());
         ctx.ranges[0].time_range = (
             Timestamp::new(1_000_000_000, TimeUnit::Nanosecond),
             Timestamp::new(2_000_000_000, TimeUnit::Nanosecond),
@@ -1370,7 +1374,7 @@ mod tests {
 
         assert_eq!(
             implied_time_range_from_exprs("ts", TimeUnit::Millisecond, &[]),
-            Some(TimestampRange::min_to_max())
+            None
         );
     }
 

src/mito2/src/read/scan_util.rs

@@ -1379,13 +1379,13 @@ mod split_tests {
         let env = SchedulerEnv::new().await;
         let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false));
         let mapper = FlatProjectionMapper::new(&metadata, [0, 2, 3]).unwrap();
-        let input = ScanInput::new(env.access_layer.clone(), mapper).with_files(files);
+        let input = ScanInput::builder(env.access_layer.clone(), mapper)
+            .with_files(files)
+            .build();
 
         StreamContext {
             input,
             ranges: vec![],
-            scan_fingerprint: None,
-            scan_implied_time_range: None,
             query_start: std::time::Instant::now(),
         }
     }
@@ -1823,16 +1823,15 @@ mod tests {
         let env = SchedulerEnv::new().await;
         let metadata = metadata_for_test();
         let mapper = FlatProjectionMapper::new(&metadata, [0, 2, 3]).unwrap();
-        let input = ScanInput::new(env.access_layer.clone(), mapper)
+        let input = ScanInput::builder(env.access_layer.clone(), mapper)
             .with_cache(CacheStrategy::Disabled)
             .with_memtables(memtables)
-            .with_files(files);
+            .with_files(files)
+            .build();
 
         Arc::new(StreamContext {
             input,
             ranges: Vec::new(),
-            scan_fingerprint: None,
-            scan_implied_time_range: None,
             query_start: Instant::now(),
         })
     }

src/mito2/src/read/series_candidate.rs

@@ -586,10 +586,9 @@ mod tests {
         ));
         let mapper =
             FlatProjectionMapper::new(&metadata, 0..metadata.column_metadatas.len()).unwrap();
-        let stream_ctx = Arc::new(StreamContext::seq_scan_ctx(ScanInput::new(
-            env.access_layer.clone(),
-            mapper,
-        )));
+        let stream_ctx = Arc::new(StreamContext::seq_scan_ctx(
+            ScanInput::builder(env.access_layer.clone(), mapper).build(),
+        ));
         let pruner = Arc::new(Pruner::new(stream_ctx.clone(), 1));
         let metrics_set = ExecutionPlanMetricsSet::new();
         let part_metrics = PartitionMetrics::new(

src/mito2/src/sst/parquet/file_range.rs

@@ -206,8 +206,11 @@ impl FileRange {
             .map(|s| s == TimeSeriesRowSelector::LastRow)
             .unwrap_or(false)
         {
-            // Only use LastRowReader if row group does not contain DELETE
-            // and all rows are selected.
+            // Only use LastRowReader if row group does not contain DELETE, all
+            // rows are selected, and filters that still run after this reader
+            // cannot change which row is last. Tag filters are safe because a
+            // tag is constant within a series. Timestamp and field filters are
+            // not safe for this shortcut.
             let put_only = !self
                 .context
                 .contains_delete(self.row_group_idx)
@@ -215,7 +218,7 @@ impl FileRange {
                     error!(e; "Failed to decode min value of op_type, fallback to FlatRowGroupReader");
                 })
                 .unwrap_or(true);
-            put_only && self.select_all()
+            put_only && self.select_all() && self.context.remaining_filters_preserve_last_row()
         } else {
             false
         };
@@ -223,7 +226,7 @@ impl FileRange {
         let flat_prune_reader = if use_last_row_reader {
             let flat_row_group_reader =
                 FlatRowGroupReader::new(self.context.clone(), parquet_reader);
-            // Flat PK prefilter makes the input stream predicate-dependent, so cached
+            // Predicate prefiltering makes the input stream predicate-dependent, so cached
             // selector results are not reusable across queries with different filters.
             let cache_strategy = if self.context.reader_builder.has_predicate_prefilter() {
                 CacheStrategy::Disabled
@@ -413,6 +416,16 @@ impl FileRangeContext {
         self.base.partition_filter.is_some()
     }
 
+    /// Returns true if applying the remaining precise filters after selecting
+    /// the last row cannot change which row is selected for a series.
+    fn remaining_filters_preserve_last_row(&self) -> bool {
+        !self.has_partition_filter()
+            && self
+                .filters()
+                .iter()
+                .all(|filter| filter.semantic_type() == SemanticType::Tag)
+    }
+
     /// Returns the format helper.
     pub(crate) fn read_format(&self) -> &FlatReadFormat {
         &self.base.read_format

src/mito2/src/sst/parquet/prefilter.rs

@@ -395,13 +395,15 @@ pub(crate) fn build_bulk_filter_plan(
 ///
 /// With predicate prefiltering enabled, tag and timestamp predicates that lower to
 /// [`SimpleFilterEvaluator`] are an exception — the engine enforces them precisely in
-/// the prefilter pass. When it is disabled, all simple filters remain on the normal
-/// precise-filter path instead.
+/// the prefilter pass. A caller can postpone simple timestamp filters to the normal
+/// precise-filter path when the scan time range covers the SST. When predicate
+/// prefiltering is disabled, all simple filters remain on the normal path instead.
 pub(crate) fn build_reader_filter_plan(
     predicate: Option<&Predicate>,
     expected_metadata: Option<&RegionMetadata>,
     pre_filter_mode: PreFilterMode,
     enable_predicate_prefilter: bool,
+    postpone_time_index_filter: bool,
     read_format: &FlatReadFormat,
     codec: &Arc<dyn PrimaryKeyCodec>,
 ) -> ReaderFilterPlan {
@@ -455,6 +457,11 @@ pub(crate) fn build_reader_filter_plan(
                 continue;
             };
 
+            if postpone_time_index_filter && filter_ctx.semantic_type() == SemanticType::Timestamp {
+                remaining_simple_filters.push(filter_ctx);
+                continue;
+            }
+
             // If the column is stored as a separate parquet column and is already projected in the main read,
             // we can evaluate the simple filter directly during prefilter.
             let direct_prefilter = can_direct_prefilter(filter_ctx.semantic_type());
@@ -1592,6 +1599,7 @@ mod tests {
             None,
             PreFilterMode::SkipFields,
             true,
+            false,
             &full_read_format,
             &codec,
         );
@@ -1601,6 +1609,42 @@ mod tests {
             vec!["field_0"]
         );
 
+        let postponed_time_plan = build_reader_filter_plan(
+            Some(&Predicate::new(vec![
+                col("tag_0").eq(lit("a")),
+                col("field_0").gt(lit(1_u64)),
+                col("ts").gt_eq(lit(ScalarValue::TimestampMillisecond(Some(1), None))),
+            ])),
+            None,
+            PreFilterMode::SkipFields,
+            true,
+            true,
+            &full_read_format,
+            &codec,
+        );
+        assert!(postponed_time_plan.prefilter_builder.is_some());
+        assert_eq!(
+            remaining_simple_filter_columns(&postponed_time_plan.remaining_simple_filters),
+            vec!["field_0", "ts"]
+        );
+
+        let postponed_time_only_plan = build_reader_filter_plan(
+            Some(&Predicate::new(vec![col("ts").gt_eq(lit(
+                ScalarValue::TimestampMillisecond(Some(1), None),
+            ))])),
+            None,
+            PreFilterMode::All,
+            true,
+            true,
+            &full_read_format,
+            &codec,
+        );
+        assert!(postponed_time_only_plan.prefilter_builder.is_none());
+        assert_eq!(
+            remaining_simple_filter_columns(&postponed_time_only_plan.remaining_simple_filters),
+            vec!["ts"]
+        );
+
         let metric_metadata: RegionMetadataRef = Arc::new(sst_region_metadata_with_encoding(
             PrimaryKeyEncoding::Sparse,
         ));
@@ -1620,6 +1664,7 @@ mod tests {
             None,
             PreFilterMode::All,
             true,
+            false,
             &projected_read_format,
             &metric_codec,
         );
@@ -1643,6 +1688,7 @@ mod tests {
             None,
             PreFilterMode::All,
             false,
+            true,
             &projected_read_format,
             &metric_codec,
         );
@@ -1675,6 +1721,7 @@ mod tests {
             None,
             PreFilterMode::All,
             true,
+            false,
             &read_format,
             &codec,
         );
@@ -1683,6 +1730,7 @@ mod tests {
             None,
             PreFilterMode::All,
             true,
+            false,
             &read_format,
             &codec,
         );

src/mito2/src/sst/parquet/reader.rs

@@ -228,6 +228,8 @@ pub struct ParquetReaderBuilder {
     pre_filter_mode: PreFilterMode,
     /// Whether to run the reduced-column predicate prefilter pass.
     enable_predicate_prefilter: bool,
+    /// Whether to apply simple time index filters during the normal precise-filter pass.
+    postpone_time_index_filter: bool,
     /// Whether to decode primary key values eagerly when reading primary key format SSTs.
     decode_primary_key_values: bool,
     page_index_policy: PageIndexPolicy,
@@ -264,6 +266,7 @@ impl ParquetReaderBuilder {
             compaction: false,
             pre_filter_mode: PreFilterMode::All,
             enable_predicate_prefilter: true,
+            postpone_time_index_filter: false,
             decode_primary_key_values: false,
             page_index_policy: Default::default(),
             defer_optional_page_index: false,
@@ -379,6 +382,13 @@ impl ParquetReaderBuilder {
         self
     }
 
+    /// Sets whether to postpone simple time index filters to precise filtering.
+    #[must_use]
+    pub(crate) fn postpone_time_index_filter(mut self, postpone: bool) -> Self {
+        self.postpone_time_index_filter = postpone;
+        self
+    }
+
     /// Decodes primary key values eagerly when reading primary key format SSTs.
     #[must_use]
     pub(crate) fn decode_primary_key_values(mut self, decode: bool) -> Self {
@@ -585,6 +595,7 @@ impl ParquetReaderBuilder {
             self.expected_metadata.as_deref(),
             self.pre_filter_mode,
             self.enable_predicate_prefilter,
+            self.postpone_time_index_filter,
             &read_format,
             &codec,
         );