From 1cdd4b54ed64f332c1d04d79f413a6ed436fe6b0 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 10 Oct 2026 12:28:47 +0000 Subject: [PATCH] fix(table): require complete primary keys for hash partitioning (#9507) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commit 28ec3ff5aad8d274d83ae42f96dc54f838fb6b20) Signed-off-by: WenyXu --- src/table/src/table/scan.rs | 100 +++++++++++++++++- .../aggregate/projected_pk_group_by.result | 69 ++++++++++++ .../aggregate/projected_pk_group_by.sql | 33 ++++++ 3 files changed, 200 insertions(+), 2 deletions(-) create mode 100644 tests/cases/standalone/common/aggregate/projected_pk_group_by.result create mode 100644 tests/cases/standalone/common/aggregate/projected_pk_group_by.sql diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 2db051a1e9c..3e0eeab84a9 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -213,7 +213,10 @@ impl RegionScanExec { Ok(tsid_col) => { Partitioning::Hash(vec![Arc::new(tsid_col) as _], num_output_partition) } - Err(_) => Partitioning::Hash(pk_columns.clone(), num_output_partition), + Err(_) if pk_columns.len() == pk_names.len() => { + Partitioning::Hash(pk_columns.clone(), num_output_partition) + } + Err(_) => Partitioning::UnknownPartitioning(num_output_partition), } } Some(TimeSeriesDistribution::TimeWindowed) | None => { @@ -694,7 +697,7 @@ mod test { use store_api::region_engine::{ PrepareRequest, QueryScanContext, RegionScanner, ScannerProperties, SinglePartitionScanner, }; - use store_api::storage::RegionId; + use store_api::storage::{RegionId, TimeSeriesDistribution}; use super::*; @@ -857,6 +860,99 @@ mod test { ) } + #[test] + fn test_per_series_partitioning_requires_all_primary_keys() { + let full_schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("marker", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ])); + let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5683)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "marker", + ConcreteDataType::int32_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "host", + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 3, + }) + .primary_key(vec![1, 2]); + let metadata = Arc::new(builder.build().unwrap()); + + for (schema, expected) in [ + ( + full_schema.clone(), + Partitioning::Hash( + vec![ + Arc::new( + Column::new_with_schema("marker", full_schema.arrow_schema()).unwrap(), + ) as _, + Arc::new( + Column::new_with_schema("host", full_schema.arrow_schema()).unwrap(), + ) as _, + ], + 2, + ), + ), + ( + Arc::new(Schema::new(vec![ + ColumnSchema::new("marker", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ])), + Partitioning::UnknownPartitioning(2), + ), + ] { + let batches = RecordBatches::try_new(schema, vec![]).unwrap(); + let scanner = Box::new(RepeatableScanner { + batches, + properties: ScannerProperties::new(vec![vec![], vec![]], false, 0), + metadata: metadata.clone(), + dynamic_filters: Arc::new(Mutex::new(Vec::new())), + }); + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }; + let plan = RegionScanExec::new(scanner, request, None).unwrap(); + let actual = plan.properties().output_partitioning(); + match (&actual, &expected) { + ( + Partitioning::UnknownPartitioning(actual_count), + Partitioning::UnknownPartitioning(expected_count), + ) => assert_eq!(actual_count, expected_count), + _ => assert_eq!(actual, &expected), + } + } + } + #[test] fn test_count_statistics_require_exact_source_rows() { for (append_mode, exact, expected) in [ diff --git a/tests/cases/standalone/common/aggregate/projected_pk_group_by.result b/tests/cases/standalone/common/aggregate/projected_pk_group_by.result new file mode 100644 index 00000000000..7ca3f1e5826 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/projected_pk_group_by.result @@ -0,0 +1,69 @@ +CREATE TABLE projected_pk_group_by ( + marker TINYINT UNSIGNED, + host STRING, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (marker, host) +); + +Affected Rows: 0 + +INSERT INTO projected_pk_group_by VALUES + (0, 'a', 1000), + (0, 'a', 2000), + (0, 'b', 1000), + (0, 'b', 2000), + (0, 'c', 1000), + (1, NULL, 1000); + +Affected Rows: 6 + +SELECT marker, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker +ORDER BY marker; + ++--------+-----------+ +| marker | row_count | ++--------+-----------+ +| 0 | 5 | +| 1 | 1 | ++--------+-----------+ + +SELECT marker, host, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker, host +ORDER BY marker, host; + ++--------+------+-----------+ +| marker | host | row_count | ++--------+------+-----------+ +| 0 | a | 2 | +| 0 | b | 2 | +| 0 | c | 1 | +| 1 | | 1 | ++--------+------+-----------+ + +ADMIN FLUSH_TABLE('projected_pk_group_by'); + ++--------------------------------------------+ +| ADMIN FLUSH_TABLE('projected_pk_group_by') | ++--------------------------------------------+ +| 0 | ++--------------------------------------------+ + +SELECT marker, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker +ORDER BY marker; + ++--------+-----------+ +| marker | row_count | ++--------+-----------+ +| 0 | 5 | +| 1 | 1 | ++--------+-----------+ + +DROP TABLE projected_pk_group_by; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/projected_pk_group_by.sql b/tests/cases/standalone/common/aggregate/projected_pk_group_by.sql new file mode 100644 index 00000000000..d525e16bc55 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/projected_pk_group_by.sql @@ -0,0 +1,33 @@ +CREATE TABLE projected_pk_group_by ( + marker TINYINT UNSIGNED, + host STRING, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (marker, host) +); + +INSERT INTO projected_pk_group_by VALUES + (0, 'a', 1000), + (0, 'a', 2000), + (0, 'b', 1000), + (0, 'b', 2000), + (0, 'c', 1000), + (1, NULL, 1000); + +SELECT marker, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker +ORDER BY marker; + +SELECT marker, host, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker, host +ORDER BY marker, host; + +ADMIN FLUSH_TABLE('projected_pk_group_by'); + +SELECT marker, count(*) AS row_count +FROM projected_pk_group_by +GROUP BY marker +ORDER BY marker; + +DROP TABLE projected_pk_group_by;