test: guard instant last row correctness

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

src/mito2/src/engine/row_selector_test.rs

@@ -12,20 +12,25 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use api::v1::Rows;
+use api::v1::value::ValueData;
+use api::v1::{Rows, Value};
 use common_base::readable_size::ReadableSize;
+use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
 use common_recordbatch::RecordBatches;
 use datafusion_common::ScalarValue;
 use datafusion_expr::{col, lit};
+use datatypes::arrow::array::{Float64Array, StringArray, TimestampMillisecondArray};
 use store_api::region_engine::RegionEngine;
 use store_api::region_request::RegionRequest;
 use store_api::storage::{RegionId, ScanRequest, TimeSeriesRowSelector};
 
 use crate::config::MitoConfig;
 use crate::engine::MitoEngine;
+use crate::read::scan_region::Scanner;
 use crate::test_util::batch_util::sort_batches_and_print;
 use crate::test_util::{
-    CreateRequestBuilder, TestEnv, build_rows_for_key, flush_region, put_rows, rows_schema,
+    CreateRequestBuilder, TestEnv, build_delete_rows_for_key, build_rows_for_key, delete_rows,
+    flush_region, put_rows, rows_schema,
 };
 
 async fn test_last_row(append_mode: bool, flat_format: bool) {
@@ -193,6 +198,231 @@ const LAST_ROW_AT_TEN: &str = "\
 | a     | 10.0    | 1970-01-01T00:00:10 |
 +-------+---------+---------------------+";
 
+async fn new_merge_last_row_engine(
+    flat_format: bool,
+) -> (
+    TestEnv,
+    MitoEngine,
+    RegionId,
+    Vec<api::v1::ColumnSchema>,
+    Vec<api::v1::ColumnSchema>,
+) {
+    let mut env = TestEnv::new().await;
+    let engine = env.create_engine(MitoConfig::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 sst_format = if flat_format { "flat" } else { "primary_key" };
+    let request = CreateRequestBuilder::new()
+        .insert_option("sst_format", sst_format)
+        .build();
+    let schema = rows_schema(&request);
+    let delete_schema = crate::test_util::delete_rows_schema(&request);
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+    (env, engine, region_id, schema, delete_schema)
+}
+
+fn value_row(key: &str, value: f64, timestamp: i64) -> api::v1::Row {
+    api::v1::Row {
+        values: vec![
+            Value {
+                value_data: Some(ValueData::StringValue(key.to_string())),
+            },
+            Value {
+                value_data: Some(ValueData::F64Value(value)),
+            },
+            Value {
+                value_data: Some(ValueData::TimestampMillisecondValue(timestamp)),
+            },
+        ],
+    }
+}
+
+async fn last_row_scanner(engine: &MitoEngine, region_id: RegionId) -> Scanner {
+    engine
+        .scanner(
+            region_id,
+            ScanRequest {
+                series_row_selector: Some(TimeSeriesRowSelector::LastRow),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap()
+}
+
+async fn scan_last_row_batches(scanner: &Scanner) -> RecordBatches {
+    RecordBatches::try_collect(scanner.scan().await.unwrap())
+        .await
+        .unwrap()
+}
+
+fn last_row_values(batches: &RecordBatches) -> Vec<(String, f64, i64)> {
+    let mut rows = Vec::new();
+    for batch in batches.iter() {
+        let batch = batch.df_record_batch();
+        let tags = batch
+            .column_by_name("tag_0")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<StringArray>()
+            .unwrap();
+        let fields = batch
+            .column_by_name("field_0")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<Float64Array>()
+            .unwrap();
+        let timestamps = batch
+            .column_by_name("ts")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<TimestampMillisecondArray>()
+            .unwrap();
+        for index in 0..batch.num_rows() {
+            rows.push((
+                tags.value(index).to_string(),
+                fields.value(index),
+                timestamps.value(index),
+            ));
+        }
+    }
+    rows.sort_by(|left, right| left.0.cmp(&right.0));
+    rows
+}
+
+#[tokio::test]
+async fn test_last_row_merge_deduplicates_same_timestamp_across_ssts() {
+    for flat_format in [false, true] {
+        let (_env, engine, region_id, schema, _delete_schema) =
+            new_merge_last_row_engine(flat_format).await;
+
+        put_rows(
+            &engine,
+            region_id,
+            Rows {
+                schema: schema.clone(),
+                rows: vec![value_row("a", 1.0, 1000)],
+            },
+        )
+        .await;
+        flush_region(&engine, region_id, None).await;
+        put_rows(
+            &engine,
+            region_id,
+            Rows {
+                schema,
+                rows: vec![value_row("a", 2.0, 1000)],
+            },
+        )
+        .await;
+        flush_region(&engine, region_id, None).await;
+
+        let scanner = last_row_scanner(&engine, region_id).await;
+        assert_eq!(2, scanner.num_files());
+        assert_eq!(0, scanner.num_memtables());
+        assert_eq!(
+            vec![("a".to_string(), 2.0, 1000)],
+            last_row_values(&scan_last_row_batches(&scanner).await)
+        );
+    }
+}
+
+#[tokio::test]
+async fn test_last_row_returns_stale_marker_and_preserves_ordinary_nan() {
+    for flat_format in [false, true] {
+        let (_env, engine, region_id, schema, _delete_schema) =
+            new_merge_last_row_engine(flat_format).await;
+        put_rows(
+            &engine,
+            region_id,
+            Rows {
+                schema: schema.clone(),
+                rows: vec![value_row("a", 1.0, 1000), value_row("b", f64::NAN, 1000)],
+            },
+        )
+        .await;
+        flush_region(&engine, region_id, None).await;
+        put_rows(
+            &engine,
+            region_id,
+            Rows {
+                schema,
+                rows: vec![value_row(
+                    "a",
+                    f64::from_bits(PROMETHEUS_STALE_NAN_BITS),
+                    1000,
+                )],
+            },
+        )
+        .await;
+        flush_region(&engine, region_id, None).await;
+
+        let scanner = last_row_scanner(&engine, region_id).await;
+        assert_eq!(2, scanner.num_files());
+        assert_eq!(0, scanner.num_memtables());
+        let values = last_row_values(&scan_last_row_batches(&scanner).await);
+        assert_eq!(2, values.len(), "unexpected LastRow values: {values:?}");
+        assert_eq!("a", values[0].0);
+        assert_eq!(PROMETHEUS_STALE_NAN_BITS, values[0].1.to_bits());
+        assert_eq!(1000, values[0].2);
+        assert_eq!("b", values[1].0);
+        assert!(values[1].1.is_nan());
+        assert_ne!(PROMETHEUS_STALE_NAN_BITS, values[1].1.to_bits());
+    }
+}
+
+#[tokio::test]
+async fn test_last_row_delete_wins_across_ssts() {
+    for flat_format in [false, true] {
+        let (_env, engine, region_id, schema, delete_schema) =
+            new_merge_last_row_engine(flat_format).await;
+        put_rows(
+            &engine,
+            region_id,
+            Rows {
+                schema,
+                rows: vec![value_row("a", 1.0, 1000)],
+            },
+        )
+        .await;
+        flush_region(&engine, region_id, None).await;
+        delete_rows(
+            &engine,
+            region_id,
+            Rows {
+                schema: delete_schema,
+                rows: build_delete_rows_for_key("a", 1, 2),
+            },
+        )
+        .await;
+        flush_region(&engine, region_id, None).await;
+
+        let scanner = last_row_scanner(&engine, region_id).await;
+        assert_eq!(2, scanner.num_files());
+        assert_eq!(0, scanner.num_memtables());
+        assert!(last_row_values(&scan_last_row_batches(&scanner).await).is_empty());
+    }
+}
+
+#[tokio::test]
+async fn test_last_row_empty_region_returns_empty() {
+    let (_env, engine, region_id, _schema, _delete_schema) = new_merge_last_row_engine(false).await;
+    let scanner = last_row_scanner(&engine, region_id).await;
+    assert!(last_row_values(&scan_last_row_batches(&scanner).await).is_empty());
+}
+
 #[tokio::test]
 async fn test_last_row_append_mode_disabled() {
     test_last_row(false, false).await;

src/query/src/optimizer/scan_hint.rs

@@ -18,6 +18,7 @@ use api::v1::SemanticType;
 use arrow_schema::SortOptions;
 use common_function::aggrs::aggr_wrapper::aggr_state_func_name;
 use common_recordbatch::OrderOption;
+use common_recordbatch::filter::SimpleFilterEvaluator;
 use datafusion::datasource::DefaultTableSource;
 use datafusion_common::tree_node::{Transformed, TreeNodeRewriter};
 use datafusion_common::{Column, Result};
@@ -90,13 +91,24 @@ impl ScanHintRule {
             return Ok(Transformed::no(LogicalPlan::TableScan(table_scan)));
         };
 
+        // LastRow is only sound when every predicate already attached to this
+        // scan is restricted to columns that preserve the series' last row.
+        // In particular, do not infer this from a Filter node above the scan:
+        // only the filters that DataFusion actually pushed to this scan matter.
+        let filters_preserve_last_row = if rewriter.inside_single_evaluation {
+            Self::filters_preserve_last_row(&table_scan, original)
+        } else {
+            true
+        };
+        let use_last_row = rewriter.inside_single_evaluation && filters_preserve_last_row;
+
         #[cfg(feature = "vector_index")]
         let has_vector_hint = rewriter.vector_search.need_rewrite();
         #[cfg(not(feature = "vector_index"))]
         let has_vector_hint = false;
         let has_hint = rewriter.order_expr.is_some()
             || rewriter.ts_row_selector.is_some()
-            || rewriter.inside_single_evaluation
+            || use_last_row
             || has_vector_hint;
         if !has_hint {
             return Ok(Transformed::no(LogicalPlan::TableScan(table_scan)));
@@ -107,8 +119,11 @@ impl ScanHintRule {
         // the shared catalog provider. This keeps order/vector/legacy hints
         // local as well as the new LastRow hint.
         let adapter = original.clone_for_scan();
-        Self::apply_hints(&adapter, rewriter, &table_scan);
-        if rewriter.inside_single_evaluation {
+        Self::apply_hints(&adapter, rewriter, &table_scan, use_last_row);
+        if use_last_row {
+            // 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);
         }
         table_scan.source =
@@ -116,13 +131,36 @@ impl ScanHintRule {
         Ok(Transformed::yes(LogicalPlan::TableScan(table_scan)))
     }
 
+    fn filters_preserve_last_row(
+        table_scan: &datafusion_expr::logical_plan::TableScan,
+        provider: &DummyTableProvider,
+    ) -> bool {
+        let metadata = provider.region_metadata();
+        for filter in &table_scan.filters {
+            let Some(filter) = SimpleFilterEvaluator::try_new(filter) else {
+                return false;
+            };
+            let Some(column_metadata) = metadata.column_by_name(filter.column_name()) else {
+                return false;
+            };
+            if !matches!(
+                column_metadata.semantic_type,
+                SemanticType::Tag | SemanticType::Timestamp
+            ) {
+                return false;
+            }
+        }
+        true
+    }
+
     fn apply_hints(
         adapter: &DummyTableProvider,
         rewriter: &mut ScanHintRewriter,
         table_scan: &datafusion_expr::logical_plan::TableScan,
+        use_last_row: bool,
     ) {
         #[cfg(not(feature = "vector_index"))]
-        let _ = table_scan;
+        let _ = (table_scan, use_last_row);
         if let Some(order_expr) = &rewriter.order_expr {
             Self::set_order_hint(adapter, order_expr);
         }
@@ -130,7 +168,7 @@ impl ScanHintRule {
             Self::set_time_series_row_selector_hint(adapter, group_by_cols, order_by_col);
         }
         #[cfg(feature = "vector_index")]
-        if rewriter.inside_single_evaluation {
+        if use_last_row {
             // LastRow and vector search are mutually exclusive for one scan:
             // vector search would bypass the ordinary sort/limit path needed by
             // the single-evaluation semantics. Still consume the queued hint so
@@ -453,10 +491,11 @@ mod test {
 
     use datafusion::functions_aggregate::first_last::last_value_udaf;
     use datafusion_common::tree_node::TreeNodeRecursion;
-    use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams};
+    use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams, Cast};
     use datafusion_expr::expr_fn::scalar_subquery;
-    use datafusion_expr::{Extension, LogicalPlan, LogicalPlanBuilder, col};
+    use datafusion_expr::{Extension, LogicalPlan, LogicalPlanBuilder, col, lit};
     use datafusion_optimizer::OptimizerContext;
+    use datatypes::arrow::datatypes::DataType;
     use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
     use store_api::storage::{RegionId, TimeSeriesRowSelector};
 
@@ -493,6 +532,29 @@ mod test {
         instant_plan_named(provider, "t", end)
     }
 
+    fn instant_plan_with_filters(
+        provider: Arc<DummyTableProvider>,
+        filters: Vec<Expr>,
+    ) -> LogicalPlan {
+        let scan = scan_plan(provider, "t");
+        let LogicalPlan::TableScan(mut scan) = scan else {
+            unreachable!();
+        };
+        scan.filters = filters;
+        LogicalPlan::Extension(Extension {
+            node: Arc::new(InstantManipulate::new(
+                1000,
+                1000,
+                1000,
+                1000,
+                "ts".to_string(),
+                vec![],
+                Some("v0".to_string()),
+                LogicalPlan::TableScan(scan),
+            )),
+        })
+    }
+
     fn scan_plan(provider: Arc<DummyTableProvider>, table_name: &str) -> LogicalPlan {
         LogicalPlanBuilder::scan(
             table_name,
@@ -590,6 +652,85 @@ mod test {
         assert_eq!(provider.scan_request().series_row_selector, None);
     }
 
+    #[test]
+    fn single_evaluation_with_tag_filter_sets_last_row() {
+        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
+        let plan = instant_plan_with_filters(provider, vec![col("k0").eq(lit("tag"))]);
+        let rewritten = ScanHintRule
+            .rewrite(plan, &OptimizerContext::default())
+            .unwrap()
+            .data;
+
+        assert_eq!(
+            scan_requests(&rewritten)[0].series_row_selector,
+            Some(TimeSeriesRowSelector::LastRow)
+        );
+    }
+
+    #[test]
+    fn single_evaluation_with_timestamp_filter_sets_last_row() {
+        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
+        let plan = instant_plan_with_filters(provider, vec![col("ts").gt_eq(lit(1_i64))]);
+        let rewritten = ScanHintRule
+            .rewrite(plan, &OptimizerContext::default())
+            .unwrap()
+            .data;
+
+        assert_eq!(
+            scan_requests(&rewritten)[0].series_row_selector,
+            Some(TimeSeriesRowSelector::LastRow)
+        );
+    }
+
+    #[test]
+    fn single_evaluation_with_cast_timestamp_filter_does_not_set_last_row() {
+        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
+        let filter = Expr::Cast(Cast::new(Box::new(col("ts")), DataType::Int64)).gt(lit(1_i64));
+        let plan = instant_plan_with_filters(provider, vec![filter]);
+        let rewritten = ScanHintRule
+            .rewrite(plan, &OptimizerContext::default())
+            .unwrap()
+            .data;
+
+        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
+    }
+
+    #[test]
+    fn single_evaluation_with_multi_column_timestamp_filter_does_not_set_last_row() {
+        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
+        let plan = instant_plan_with_filters(provider, vec![col("ts").gt(col("k0"))]);
+        let rewritten = ScanHintRule
+            .rewrite(plan, &OptimizerContext::default())
+            .unwrap()
+            .data;
+
+        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
+    }
+
+    #[test]
+    fn single_evaluation_with_field_filter_does_not_set_last_row() {
+        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
+        let plan = instant_plan_with_filters(provider, vec![col("v0").gt(lit(1.0_f64))]);
+        let rewritten = ScanHintRule
+            .rewrite(plan, &OptimizerContext::default())
+            .unwrap()
+            .data;
+
+        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
+    }
+
+    #[test]
+    fn single_evaluation_with_unknown_filter_does_not_set_last_row() {
+        let provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));
+        let plan = instant_plan_with_filters(provider, vec![col("unknown").eq(lit(1_i64))]);
+        let rewritten = ScanHintRule
+            .rewrite(plan, &OptimizerContext::default())
+            .unwrap()
+            .data;
+
+        assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None);
+    }
+
     #[test]
     fn expression_subquery_isolated_from_outer_single_evaluation() {
         let outer_provider = Arc::new(mock_table_provider(RegionId::new(1, 1)));