fix(table): require complete primary keys for hash partitioning (#9507)

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
(cherry picked from commit 28ec3ff5aa)
Signed-off-by: WenyXu <wenymedia@gmail.com>

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 [

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
+

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;