From c756212c8634fbd540b93ec4db16e5413947abce Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 7 Sep 2026 12:30:03 +0800 Subject: [PATCH] test: guard instant last row correctness Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/engine/row_selector_test.rs | 234 +++++++++++++++++++++- src/query/src/optimizer/scan_hint.rs | 155 +++++++++++++- 2 files changed, 380 insertions(+), 9 deletions(-) diff --git a/src/mito2/src/engine/row_selector_test.rs b/src/mito2/src/engine/row_selector_test.rs index 7c66a5562fe..0366211117d 100644 --- a/src/mito2/src/engine/row_selector_test.rs +++ b/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, + Vec, +) { + 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::() + .unwrap(); + let fields = batch + .column_by_name("field_0") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .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; diff --git a/src/query/src/optimizer/scan_hint.rs b/src/query/src/optimizer/scan_hint.rs index ecfd9293c55..6e9c570e113 100644 --- a/src/query/src/optimizer/scan_hint.rs +++ b/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, + filters: Vec, + ) -> 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, 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)));