diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 27a0d9196b..39bae738c6 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -554,6 +554,24 @@ impl MergeScanExec { return None; } + let hash_expr_col_names: HashSet<_> = hash_exprs + .iter() + .filter_map(|expr| { + expr.as_any() + .downcast_ref::() + .map(|col_expr| col_expr.name()) + }) + .collect(); + + let covers_all_partition_cols = self.partition_cols.values().all(|aliases| { + aliases + .iter() + .any(|col| hash_expr_col_names.contains(col.name())) + }); + if !covers_all_partition_cols { + return None; + } + let all_partition_col_aliases: HashSet<_> = self .partition_cols .values() diff --git a/src/query/src/optimizer/pass_distribution.rs b/src/query/src/optimizer/pass_distribution.rs index 22e2b90359..8ff21e046d 100644 --- a/src/query/src/optimizer/pass_distribution.rs +++ b/src/query/src/optimizer/pass_distribution.rs @@ -16,8 +16,9 @@ use std::sync::Arc; use datafusion::config::ConfigOptions; use datafusion::physical_optimizer::PhysicalOptimizerRule; -use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use datafusion_common::Result as DfResult; use datafusion_physical_expr::Distribution; use datafusion_physical_expr::utils::map_columns_before_projection; @@ -68,11 +69,26 @@ impl PassDistribution { ) -> DfResult> { // If this is a MergeScanExec, try to apply the current requirement. if let Some(merge_scan) = plan.as_any().downcast_ref::() - && let Some(distribution) = current_req.as_ref() - && let Some(new_plan) = merge_scan.try_with_new_distribution(distribution.clone()) + && let Some(Distribution::HashPartitioned(hash_exprs)) = current_req.as_ref() { - // Leaf node; no children to process - return Ok(Arc::new(new_plan) as _); + if let Partitioning::Hash(current_hash_exprs, _) = &merge_scan.properties().partitioning + && *current_hash_exprs == *hash_exprs + { + return Ok(plan); + } + + if let Some(new_plan) = merge_scan + .try_with_new_distribution(Distribution::HashPartitioned(hash_exprs.clone())) + { + // Leaf node; no children to process + return Ok(Arc::new(new_plan) as _); + } + + let partitioning = Partitioning::Hash( + hash_exprs.clone(), + merge_scan.properties().partitioning.partition_count(), + ); + return Ok(Arc::new(RepartitionExec::try_new(plan, partitioning)?) as _); } // Compute per-child requirements from the current node. @@ -266,6 +282,20 @@ mod tests { ); } + #[test] + fn merge_scan_rejects_hash_requirement_on_partition_key_subset() { + let merge_scan = test_merge_scan_exec(test_schema()); + + let new_plan = merge_scan.try_with_new_distribution(Distribution::HashPartitioned(vec![ + partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1), + ])); + + assert!( + new_plan.is_none(), + "partitioning by a subset of multi-column partition keys is not sufficient" + ); + } + fn test_merge_scan_exec(schema: SchemaRef) -> MergeScanExec { let session_state = SessionStateBuilder::new().with_default_features().build(); let partition_cols = BTreeMap::from([ diff --git a/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result new file mode 100644 index 0000000000..91c0e3e851 --- /dev/null +++ b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result @@ -0,0 +1,118 @@ +CREATE TABLE pd_sub_l( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + lv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +Affected Rows: 0 + +CREATE TABLE pd_sub_r( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + rv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +Affected Rows: 0 + +INSERT INTO pd_sub_l VALUES + (1, 1, 'a', 10), + (2, 1, 'z', 20); + +Affected Rows: 2 + +INSERT INTO pd_sub_r VALUES + (1, 1, 'a', 100), + (2, 1, 'z', 200); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('pd_sub_l'); + ++-------------------------------+ +| ADMIN FLUSH_TABLE('pd_sub_l') | ++-------------------------------+ +| 0 | ++-------------------------------+ + +ADMIN FLUSH_TABLE('pd_sub_r'); + ++-------------------------------+ +| ADMIN FLUSH_TABLE('pd_sub_r') | ++-------------------------------+ +| 0 | ++-------------------------------+ + +-- Joining on only a subset of a multi-column partition key must still compare +-- rows across all matching partitions for that subset key. +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED_REGION +-- SQLNESS REPLACE Hash\(\[device_id@0\],\s+\d+\) Hash([device_id@0],REDACTED_PARTITIONS) +-- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED +-- SQLNESS REPLACE input_partitions=REDACTED\s+\| input_partitions=REDACTED| +EXPLAIN SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + ++---------------+-------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: l.area ASC NULLS LAST, r.area ASC NULLS LAST | +| | Projection: l.area, r.area | +| | Inner Join: l.device_id = r.device_id | +| | Projection: l.device_id, l.area | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: l | +| | Filter: pd_sub_l.device_id IS NOT NULL | +| | TableScan: pd_sub_l, partial_filters=[pd_sub_l.device_id IS NOT NULL] | +| | ]] | +| | Projection: r.device_id, r.area | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: r | +| | Filter: pd_sub_r.device_id IS NOT NULL | +| | TableScan: pd_sub_r, partial_filters=[pd_sub_r.device_id IS NOT NULL] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [area@0 ASC NULLS LAST, area@1 ASC NULLS LAST] | +| | SortExec: expr=[area@0 ASC NULLS LAST, area@1 ASC NULLS LAST], preserve_partitioning=[true] | +| | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(device_id@0, device_id@0)], projection=[area@1, area@3] | +| | RepartitionExec: partitioning=Hash([device_id@0],REDACTED_PARTITIONS), input_partitions=REDACTED| +| | ProjectionExec: expr=[device_id@1 as device_id, area@2 as area] | +| | MergeScanExec: peers=[REDACTED_REGION, REDACTED_REGION, REDACTED_REGION, ] | +| | RepartitionExec: partitioning=Hash([device_id@0],REDACTED_PARTITIONS), input_partitions=REDACTED| +| | ProjectionExec: expr=[device_id@1 as device_id, area@2 as area] | +| | MergeScanExec: peers=[REDACTED_REGION, REDACTED_REGION, REDACTED_REGION, ] | +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------+ + +SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + ++------+------+ +| area | area | ++------+------+ +| a | a | +| a | z | +| z | a | +| z | z | ++------+------+ + +DROP TABLE pd_sub_l; + +Affected Rows: 0 + +DROP TABLE pd_sub_r; + +Affected Rows: 0 + diff --git a/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql new file mode 100644 index 0000000000..41337e74a0 --- /dev/null +++ b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql @@ -0,0 +1,52 @@ +CREATE TABLE pd_sub_l( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + lv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +CREATE TABLE pd_sub_r( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + rv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +INSERT INTO pd_sub_l VALUES + (1, 1, 'a', 10), + (2, 1, 'z', 20); + +INSERT INTO pd_sub_r VALUES + (1, 1, 'a', 100), + (2, 1, 'z', 200); + +ADMIN FLUSH_TABLE('pd_sub_l'); +ADMIN FLUSH_TABLE('pd_sub_r'); + +-- Joining on only a subset of a multi-column partition key must still compare +-- rows across all matching partitions for that subset key. +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED_REGION +-- SQLNESS REPLACE Hash\(\[device_id@0\],\s+\d+\) Hash([device_id@0],REDACTED_PARTITIONS) +-- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED +-- SQLNESS REPLACE input_partitions=REDACTED\s+\| input_partitions=REDACTED| +EXPLAIN SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + +SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + +DROP TABLE pd_sub_l; +DROP TABLE pd_sub_r;