rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-07-13 17:30:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: repartition subset partition key joins (#8460)

Browse Source 
Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-07-10 14:10:28 +08:00
committed by GitHub
parent 74bdb1f2a8
commit ff2fa71d50
4 changed files with 223 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
src/query/src/dist_plan/merge_scan.rs
View File

@@ -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::<Column>()
.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()

40
src/query/src/optimizer/pass_distribution.rs
View File

@@ -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<Arc<dyn ExecutionPlan>> {
// If this is a MergeScanExec, try to apply the current requirement.
if let Some(merge_scan) = plan.as_any().downcast_ref::<MergeScanExec>()
&& 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([

118
tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result Normal file
View File

@@ -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

52
tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql Normal file
View File

@@ -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;