From 34600156777554a74cd2dc373edb7842f164d486 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:17:11 +0800 Subject: [PATCH] fix(mito2): backport repartition count correctness (#9154) Backport upstream d9b97796a6c3fa884c5ffab37d6d192b2b5c1e99. Adapt regression fixtures to release APIs and assert complete row identities after repartition and SPLIT. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/read/scan_region.rs | 55 ++++++ src/mito2/src/read/seq_scan.rs | 3 +- src/mito2/src/read/series_scan.rs | 3 +- src/mito2/src/read/unordered_scan.rs | 3 +- src/store-api/src/region_engine.rs | 15 ++ src/table/src/table/scan.rs | 180 +++++++++++++++++--- tests-integration/tests/repartition.rs | 224 ++++++++++++++++++++++++- 7 files changed, 457 insertions(+), 26 deletions(-) diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 375e0059de7..74275844156 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -1437,6 +1437,27 @@ impl ScanInput { ); } + /// Whether physical source row counts equal the rows visible to this region, + /// before applying query predicates. + pub(crate) fn total_rows_is_exact(&self) -> bool { + if self.region_partition_expr.is_none() { + return true; + } + + // Extension ranges do not expose their partition expressions. + #[cfg(feature = "enterprise")] + if !self.extension_ranges.is_empty() { + return false; + } + + // Repartition flushes existing memtables before entering staging. New + // writes are routed under the target partition rules, whereas historical + // SSTs can be shared by regions and still need partition filtering. + self.files + .iter() + .all(|file| self.should_skip_region_partition(file)) + } + pub(crate) fn total_rows(&self) -> usize { let rows_in_files: usize = self.files.iter().map(|f| f.num_rows()).sum(); let rows_in_memtables: usize = self.memtables.iter().map(|m| m.stats().num_rows()).sum(); @@ -2126,6 +2147,40 @@ mod tests { .with_files(vec![file]) } + #[tokio::test] + async fn test_total_rows_is_exact_after_partition_filter() { + let expr = partition_col("k0").gt_eq(Value::String("foo".into())); + let other = partition_col("k0").gt_eq(Value::String("bar".into())); + for (region_expr, file_exprs, exact) in [ + (None, vec![None], true), + (None, vec![Some(expr.clone())], true), + (Some(expr.clone()), vec![], true), + (Some(expr.clone()), vec![Some(expr.clone())], true), + (Some(expr.clone()), vec![None], false), + (Some(expr.clone()), vec![Some(other)], false), + (Some(expr.clone()), vec![Some(expr), None], false), + ] { + let mut builder = + RegionMetadataBuilder::from_existing(metadata_with_primary_key(vec![0, 1], false)); + builder.partition_expr_json(region_expr.map(|expr| expr.as_json_str().unwrap())); + let metadata = Arc::new(builder.build_without_validation().unwrap()); + let files = file_exprs + .into_iter() + .map(|partition_expr| { + FileHandle::new( + FileMeta { + partition_expr, + ..Default::default() + }, + Arc::new(crate::sst::file_purger::NoopFilePurger), + ) + }) + .collect(); + let input = new_scan_input(metadata, vec![]).await.with_files(files); + assert_eq!(input.total_rows_is_exact(), exact); + } + } + /// Helper to create a timestamp millisecond literal. fn ts_lit(val: i64) -> datafusion_expr::Expr { lit(ScalarValue::TimestampMillisecond(Some(val), None)) diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 7253ba532f6..6803cbc5db5 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -76,7 +76,8 @@ impl SeqScan { pub(crate) fn new(input: ScanInput) -> Self { let mut properties = ScannerProperties::default() .with_append_mode(input.append_mode) - .with_total_rows(input.total_rows()); + .with_total_rows(input.total_rows()) + .with_total_rows_is_exact(input.append_mode && input.total_rows_is_exact()); if let Some(counters) = input.query_stat_counters.clone() { properties.set_query_stat_counters(counters); } diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index adefae91173..d9b48d704ff 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -86,7 +86,8 @@ impl SeriesScan { pub(crate) fn new(input: ScanInput) -> Self { let mut properties = ScannerProperties::default() .with_append_mode(input.append_mode) - .with_total_rows(input.total_rows()); + .with_total_rows(input.total_rows()) + .with_total_rows_is_exact(input.append_mode && input.total_rows_is_exact()); if let Some(counters) = input.query_stat_counters.clone() { properties.set_query_stat_counters(counters); } diff --git a/src/mito2/src/read/unordered_scan.rs b/src/mito2/src/read/unordered_scan.rs index a5c09e5cdef..883721534aa 100644 --- a/src/mito2/src/read/unordered_scan.rs +++ b/src/mito2/src/read/unordered_scan.rs @@ -61,7 +61,8 @@ impl UnorderedScan { pub(crate) fn new(input: ScanInput) -> Self { let mut properties = ScannerProperties::default() .with_append_mode(input.append_mode) - .with_total_rows(input.total_rows()); + .with_total_rows(input.total_rows()) + .with_total_rows_is_exact(input.append_mode && input.total_rows_is_exact()); if let Some(counters) = input.query_stat_counters.clone() { properties.set_query_stat_counters(counters); } diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index d2d7410257c..7dc8f685eca 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -306,6 +306,9 @@ pub struct ScannerProperties { /// [ScannerProperties::append_mode] is true. total_rows: usize, + /// Whether total_rows is exact before applying query predicates. + total_rows_is_exact: bool, + /// Whether to yield an empty batch to distinguish partition ranges. pub distinguish_partition_range: bool, @@ -334,12 +337,19 @@ impl ScannerProperties { self } + /// Sets whether the source row count is exact before query predicates. + pub fn with_total_rows_is_exact(mut self, exact: bool) -> Self { + self.total_rows_is_exact = exact; + self + } + /// Creates a new [`ScannerProperties`] with the given partitioning. pub fn new(partitions: Vec>, append_mode: bool, total_rows: usize) -> Self { Self { partitions, append_mode, total_rows, + total_rows_is_exact: false, distinguish_partition_range: false, target_partitions: 0, logical_region: false, @@ -374,6 +384,11 @@ impl ScannerProperties { self.total_rows } + /// Returns whether the source row count is exact before query predicates. + pub fn total_rows_is_exact(&self) -> bool { + self.total_rows_is_exact + } + /// Returns whether the scanner is scanning a logical region. pub fn is_logical_region(&self) -> bool { self.logical_region diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 9ce752553e5..0ba1d415299 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -447,30 +447,32 @@ impl ExecutionPlan for RegionScanExec { } fn partition_statistics(&self, partition: Option) -> DfResult { - if partition.is_some() { + if partition.is_some() || !self.append_mode { return Ok(Statistics::new_unknown(self.schema().as_ref())); } - let statistics = - if self.append_mode && !self.scanner.lock().unwrap().has_predicate_without_region() { - let column_statistics = self - .arrow_schema - .fields - .iter() - .map(|_| ColumnStatistics { - distinct_count: Precision::Exact(self.total_rows), - null_count: Precision::Exact(0), // all null rows are counted for append-only table - ..Default::default() - }) - .collect(); - Statistics { - num_rows: Precision::Exact(self.total_rows), - total_byte_size: Default::default(), - column_statistics, - } - } else { - Statistics::new_unknown(&self.arrow_schema) - }; + let scanner = self.scanner.lock().unwrap(); + let statistics = if scanner.properties().total_rows_is_exact() + && !scanner.has_predicate_without_region() + { + let column_statistics = self + .arrow_schema + .fields + .iter() + .map(|_| ColumnStatistics { + distinct_count: Precision::Exact(self.total_rows), + null_count: Precision::Exact(0), // all null rows are counted for append-only table + ..Default::default() + }) + .collect(); + Statistics { + num_rows: Precision::Exact(self.total_rows), + total_byte_size: Default::default(), + column_statistics, + } + } else { + Statistics::new_unknown(&self.arrow_schema) + }; Ok(statistics) } @@ -577,11 +579,145 @@ mod test { use datatypes::vectors::{Int32Vector, TimestampMillisecondVector}; use futures::TryStreamExt; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; - use store_api::region_engine::SinglePartitionScanner; + use store_api::region_engine::{RegionScanner, ScannerProperties, SinglePartitionScanner}; use store_api::storage::RegionId; use super::*; + #[derive(Debug)] + struct RepeatableScanner { + batches: RecordBatches, + properties: ScannerProperties, + metadata: store_api::metadata::RegionMetadataRef, + } + + impl DisplayAs for RepeatableScanner { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter<'_>, + ) -> std::fmt::Result { + write!(f, "RepeatableScanner") + } + } + + impl RegionScanner for RepeatableScanner { + fn name(&self) -> &str { + "RepeatableScanner" + } + + fn properties(&self) -> &ScannerProperties { + &self.properties + } + + fn schema(&self) -> SchemaRef { + self.batches.schema() + } + + fn metadata(&self) -> store_api::metadata::RegionMetadataRef { + self.metadata.clone() + } + + fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError> { + self.properties.prepare(request); + Ok(()) + } + + fn scan_partition( + &self, + _ctx: &QueryScanContext, + _metrics_set: &ExecutionPlanMetricsSet, + _partition: usize, + ) -> Result { + Ok(self.batches.as_stream()) + } + + fn has_predicate_without_region(&self) -> bool { + false + } + + fn add_dyn_filter_to_predicate( + &mut self, + filters: Vec>, + ) -> Vec { + vec![false; filters.len()] + } + + fn set_logical_region(&mut self, logical_region: bool) { + self.properties.set_logical_region(logical_region); + } + + fn set_query_load_region_id(&mut self, region_id: RegionId) { + self.properties.set_query_load_region_id(region_id); + } + } + + fn count_statistics_test_data() -> ( + SchemaRef, + RecordBatch, + store_api::metadata::RegionMetadataRef, + ) { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("a", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(Int32Vector::from_slice([1, 2, 3, 4])) as _, + Arc::new(TimestampMillisecondVector::from_slice([1, 2, 3, 4])) as _, + ], + ) + .unwrap(); + let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5685)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("a", ConcreteDataType::int32_datatype(), false), + semantic_type: SemanticType::Field, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 2, + }) + .primary_key(vec![]); + (schema, batch, Arc::new(builder.build().unwrap())) + } + + #[test] + fn test_count_statistics_require_exact_source_rows() { + for (append_mode, exact, expected) in [ + (true, false, Precision::Absent), + (true, true, Precision::Exact(4)), + (false, true, Precision::Absent), + ] { + let (schema, batch, metadata) = count_statistics_test_data(); + let scanner = Box::new(RepeatableScanner { + batches: RecordBatches::try_new(schema, vec![batch]).unwrap(), + properties: ScannerProperties::default() + .with_append_mode(append_mode) + .with_total_rows(4) + .with_total_rows_is_exact(exact), + metadata, + }); + let plan = RegionScanExec::new(scanner, ScanRequest::default(), None).unwrap(); + assert_eq!(plan.partition_statistics(None).unwrap().num_rows, expected); + assert_eq!( + plan.partition_statistics(Some(0)).unwrap().num_rows, + Precision::Absent, + ); + } + } + #[tokio::test] async fn test_simple_table_scan() { let ctx = SessionContext::new(); diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index cf64f40aa00..70e08bbbbbe 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -24,7 +24,7 @@ use common_procedure::{ProcedureContext, ProcedureWithId, watcher}; use common_query::Output; use common_telemetry::info; use common_test_util::recordbatch::check_output_stream; -use common_test_util::temp_dir::create_temp_dir; +use common_test_util::temp_dir::{TempDir, create_temp_dir}; use common_wal::config::DatanodeWalConfig; use frontend::instance::Instance; use meta_srv::gc::{self, BatchGcProcedure, GcSchedulerOptions, GcTickerRef}; @@ -117,6 +117,228 @@ macro_rules! repartition_tests { }; } +/// COUNT must use visible rows rather than the full row counts of shared SSTs. +#[tokio::test(flavor = "multi_thread")] +async fn test_repartition_append_count_file() { + let (cluster, _home_guard) = append_count_cluster("repartition_append_count").await; + let instance = cluster.fe_instance(); + let table = "count_repartition"; + prepare_append_count_table(instance, table, "").await; + assert_append_count(instance, table, 100, 1).await; + flush_append_count_table(instance, table).await; + assert_append_count(instance, table, 100, 1).await; + + let sql = "ALTER TABLE count_repartition PARTITION ON COLUMNS (device_id) \ + (device_id < 50, device_id >= 50)"; + run_sql(instance, sql, QueryContext::arc()).await.unwrap(); + wait_for_append_count_regions(instance, table, 2).await; + assert_append_count(instance, table, 100, 0).await; + check_append_count_after_write(instance, table, 0).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_split_append_count_file() { + let (cluster, _home_guard) = append_count_cluster("split_append_count").await; + let instance = cluster.fe_instance(); + let table = "count_split"; + prepare_append_count_table( + instance, + table, + "PARTITION ON COLUMNS (device_id) (device_id < 50, device_id >= 50)", + ) + .await; + assert_append_count(instance, table, 100, 2).await; + flush_append_count_table(instance, table).await; + assert_append_count(instance, table, 100, 2).await; + + let sql = "ALTER TABLE count_split SPLIT PARTITION (device_id < 50) INTO \ + (device_id < 25, device_id >= 25 AND device_id < 50)"; + run_sql(instance, sql, QueryContext::arc()).await.unwrap(); + wait_for_append_count_regions(instance, table, 3).await; + assert_append_count(instance, table, 100, 1).await; + check_append_count_after_write(instance, table, 1).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_repartition_append_count_memtable_file() { + let (cluster, home_guard) = append_count_cluster("repartition_append_count_memtable").await; + assert!(home_guard.path().is_dir()); + let instance = cluster.fe_instance(); + let table = "count_repartition_memtable"; + prepare_append_count_table(instance, table, "").await; + assert_append_count(instance, table, 100, 1).await; + + // Entering staging must flush the populated memtable before changing partitions. + let sql = "ALTER TABLE count_repartition_memtable PARTITION ON COLUMNS (device_id) \ + (device_id < 50, device_id >= 50)"; + run_sql(instance, sql, QueryContext::arc()).await.unwrap(); + wait_for_append_count_regions(instance, table, 2).await; + assert_append_count(instance, table, 100, 0).await; + check_append_count_after_write(instance, table, 0).await; + assert!(home_guard.path().is_dir()); +} + +async fn append_count_cluster(name: &str) -> (GreptimeDbCluster, Arc) { + common_telemetry::init_default_ut_logging(); + let (store_config, _guard) = get_test_store_config(&StorageType::File); + let home_dir = Arc::new(create_temp_dir(name)); + let cluster = GreptimeDbClusterBuilder::new(name) + .await + .with_shared_home_dir(Arc::clone(&home_dir)) + .with_datanodes(3) + .with_store_config(store_config) + .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + ..Default::default() + }) + .build(true) + .await; + (cluster, home_dir) +} + +async fn prepare_append_count_table(instance: &Arc, table: &str, partitions: &str) { + let sql = format!( + "CREATE TABLE {table} (ts TIMESTAMP TIME INDEX, device_id INT) \ + {partitions} ENGINE=mito WITH(append_mode='true')" + ); + run_sql(instance, &sql, QueryContext::arc()).await.unwrap(); + let sql = format!( + "INSERT INTO {table} SELECT to_timestamp_millis(value), CAST(value AS INT) \ + FROM generate_series(0, 99)" + ); + run_sql(instance, &sql, QueryContext::arc()).await.unwrap(); +} + +async fn flush_append_count_table(instance: &Arc, table: &str) { + run_sql( + instance, + &format!("ADMIN flush_table('{table}')"), + QueryContext::arc(), + ) + .await + .unwrap(); +} + +async fn check_append_count_after_write( + instance: &Arc, + table: &str, + statistics_regions: usize, +) { + run_sql( + instance, + &format!("INSERT INTO {table} VALUES (to_timestamp_millis(100), 100)"), + QueryContext::arc(), + ) + .await + .unwrap(); + assert_append_count(instance, table, 101, statistics_regions).await; + flush_append_count_table(instance, table).await; + assert_append_count(instance, table, 101, statistics_regions).await; +} + +async fn wait_for_append_count_regions(instance: &Arc, table: &str, expected: usize) { + // Wait for frontend routing, not for COUNT to become correct. + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let plan = + append_count_query(instance, &format!("EXPLAIN ANALYZE SELECT * FROM {table}")) + .await; + if plan.matches("UnorderedScan: region=").count() == expected { + return; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await + .expect("new partition routes must become visible"); +} + +async fn append_count_query(instance: &Arc, sql: &str) -> String { + let output = run_sql(instance, sql, QueryContext::arc()).await.unwrap(); + let batches = match output.data { + common_query::OutputData::Stream(stream) => { + common_recordbatch::RecordBatches::try_collect(stream) + .await + .unwrap() + } + common_query::OutputData::RecordBatches(batches) => batches, + _ => panic!("expected query output: {sql}"), + }; + batches.pretty_print().unwrap() +} + +async fn assert_append_count( + instance: &Arc, + table: &str, + rows: usize, + statistics_regions: usize, +) { + let expected = format!("+-----+\n| n |\n+-----+\n| {rows:<3} |\n+-----+"); + let count = append_count_query(instance, &format!("SELECT count(*) AS n FROM {table}")).await; + assert_eq!(count, expected, "unbounded COUNT for {table}"); + let scanned_count = append_count_query( + instance, + &format!("SELECT count(*) AS n FROM (SELECT * FROM {table} LIMIT 1000)"), + ) + .await; + assert_eq!(scanned_count, expected, "scanned COUNT for {table}"); + + let output = run_sql( + instance, + &format!("SELECT ts, device_id FROM {table} ORDER BY ts, device_id"), + QueryContext::arc(), + ) + .await + .unwrap(); + let batches = match output.data { + common_query::OutputData::Stream(stream) => { + common_recordbatch::RecordBatches::try_collect(stream) + .await + .unwrap() + } + common_query::OutputData::RecordBatches(batches) => batches, + _ => panic!("expected row output for {table}"), + }; + let actual = batches + .iter() + .flat_map(|batch| { + let timestamps = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let device_ids = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(move |index| (timestamps.value(index), device_ids.value(index))) + }) + .collect::>(); + let expected_rows = (0..rows as i64) + .map(|value| (value, value as i32)) + .collect::>(); + assert_eq!(actual, expected_rows, "ordered rows for {table}"); + + // Only regions whose source row counts are exact may use statistics. + let plan = append_count_query( + instance, + &format!("EXPLAIN ANALYZE SELECT count(*) FROM {table}"), + ) + .await; + assert_eq!( + plan.matches("PlaceholderRowExec").count(), + statistics_regions, + "{plan}" + ); +} + #[tokio::test(flavor = "multi_thread")] async fn test_repartition_multi_hop_fulltext_index_gc_file() { if StorageType::File.test_on() {