rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 23:12:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(query): stop reporting region scans as hash partitioned

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-09 17:03:09 +08:00
parent 4ade13d9b6
commit a8159b29a8
7 changed files with 287 additions and 548 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/query/src/dist_plan/merge_scan.rs
+169 -175
View File
@@ -17,7 +17,7 @@ use std::cell::Cell;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use ahash::{HashMap, HashSet};
use ahash::HashMap;
use arrow_schema::{
ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema,
SchemaRef as ArrowSchemaRef, SortOptions,
@@ -44,10 +44,10 @@ use datafusion::physical_plan::{
};
use datafusion_common::stats::Precision;
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::{Column as ColumnExpr, DFSchemaRef, DataFusionError, Result, Statistics};
use datafusion_common::{DFSchemaRef, DataFusionError, Result, Statistics};
use datafusion_expr::{Expr, Extension, FetchType, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion_physical_expr::expressions::Column;
use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalSortExpr};
use datafusion_physical_expr::{EquivalenceProperties, PhysicalSortExpr};
use datatypes::extension::json::is_any_json_extension_type;
use futures_util::StreamExt;
use greptime_proto::v1::region::RegionRequestHeader;
@@ -489,7 +489,6 @@ pub struct MergeScanExec {
remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
captured_remote_dyn_filters: Arc<Mutex<Vec<CapturedDynFilter>>>,
target_partition: usize,
partition_cols: AliasMapping,
enable_per_region_metrics: bool,
}
@@ -514,7 +513,6 @@ impl MergeScanExec {
region_query_handler: RegionQueryHandlerRef,
query_ctx: QueryContextRef,
target_partition: usize,
partition_cols: AliasMapping,
remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
enable_per_region_metrics: bool,
) -> Result<Self> {
@@ -551,24 +549,8 @@ impl MergeScanExec {
EquivalenceProperties::new(arrow_schema.clone())
};
let partition_exprs = partition_cols
.iter()
.filter_map(|col| {
if let Some(first_alias) = col.1.first() {
session_state
.create_physical_expr(
Expr::Column(ColumnExpr::new_unqualified(
first_alias.name().to_string(),
)),
plan.schema(),
)
.ok()
} else {
None
}
})
.collect();
let partitioning = Partitioning::Hash(partition_exprs, output_partition_count);
// Region stripes are positional, not hash buckets; do not advertise a hash distribution.
let partitioning = Partitioning::UnknownPartitioning(output_partition_count);
let properties = Arc::new(PlanProperties::new(
eq_properties,
@@ -590,7 +572,6 @@ impl MergeScanExec {
remote_dyn_filter_producer_id,
captured_remote_dyn_filters: Arc::default(),
target_partition,
partition_cols,
enable_per_region_metrics,
})
}
@@ -912,78 +893,6 @@ impl MergeScanExec {
)))
}
pub fn try_with_new_distribution(&self, distribution: Distribution) -> Option<Self> {
let Distribution::KeyPartitioned(hash_exprs) = distribution else {
// not applicable
return None;
};
if let Partitioning::Hash(curr_dist, _) = &self.properties.partitioning
&& curr_dist == &hash_exprs
{
// No need to change the distribution
return None;
}
let hash_expr_col_names: HashSet<_> = hash_exprs
.iter()
.filter_map(|expr| {
expr.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()
.flat_map(|aliases| aliases.iter().map(|c| c.name()))
.collect();
let overlaps: Vec<_> = hash_exprs
.iter()
.filter(|expr| {
expr.downcast_ref::<Column>()
.is_some_and(|col_expr| all_partition_col_aliases.contains(col_expr.name()))
})
.cloned()
.collect();
if overlaps.is_empty() {
return None;
}
Some(Self {
table: self.table.clone(),
regions: self.regions.clone(),
plan: self.plan.clone(),
arrow_schema: self.arrow_schema.clone(),
region_query_handler: self.region_query_handler.clone(),
metric: self.metric.clone(),
properties: Arc::new(PlanProperties::new(
self.properties.eq_properties.clone(),
Partitioning::Hash(overlaps, self.partition_count()),
self.properties.emission_type,
self.properties.boundedness,
)),
sub_stage_metrics: self.sub_stage_metrics.clone(),
partition_metrics: self.partition_metrics.clone(),
query_ctx: self.query_ctx.clone(),
remote_dyn_filter_producer_id: self.remote_dyn_filter_producer_id,
captured_remote_dyn_filters: self.captured_remote_dyn_filters.clone(),
target_partition: self.target_partition,
partition_cols: self.partition_cols.clone(),
enable_per_region_metrics: self.enable_per_region_metrics,
})
}
fn captured_remote_dyn_filters(&self) -> Vec<CapturedDynFilter> {
self.captured_remote_dyn_filters.lock().unwrap().clone()
}
@@ -1420,7 +1329,7 @@ impl MergeScanMetric {
#[cfg(test)]
mod tests {
use std::collections::{BTreeSet, HashMap as StdHashMap};
use std::collections::HashMap as StdHashMap;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll};
@@ -1439,17 +1348,20 @@ mod tests {
use common_recordbatch::{
DfRecordBatch, EmptyRecordBatchStream, RecordBatch, RecordBatchStream,
};
use datafusion::common::NullEquality;
use datafusion::config::ConfigOptions;
use datafusion::execution::SessionStateBuilder;
use datafusion::execution::{SessionConfig, SessionStateBuilder};
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult;
use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::{ExecutionPlanProperties, Partitioning};
use datafusion::physical_plan::{StatisticsArgs, StatisticsContext};
use datafusion_common::TableReference;
use datafusion_expr::{LogicalPlanBuilder, col, lit};
use datafusion_expr::{JoinType, LogicalPlanBuilder, col, lit};
use datafusion_physical_expr::PhysicalExpr;
use datafusion_physical_expr::expressions::{
Column, DynamicFilterPhysicalExpr, lit as physical_lit,
};
use datafusion_physical_expr::{Distribution, PhysicalExpr};
use datatypes::prelude::{ConcreteDataType, VectorRef};
use datatypes::schema::{ColumnSchema, Schema};
use datatypes::vectors::{Int64Vector, StringVector, TimestampMillisecondVector};
@@ -1517,7 +1429,6 @@ mod tests {
Arc::new(TestRegionQueryHandler::default()),
QueryContext::arc(),
target_partition,
AliasMapping::new(),
None,
false,
)
@@ -1560,7 +1471,6 @@ mod tests {
handler,
query_ctx,
1,
AliasMapping::new(),
Some(RemoteDynFilterProducerId::new(42)),
false,
)
@@ -1661,9 +1571,165 @@ mod tests {
exec.properties().output_partitioning().partition_count(),
expected
);
assert!(matches!(
exec.properties().output_partitioning(),
Partitioning::UnknownPartitioning(count) if *count == expected
));
}
}
#[tokio::test]
async fn merge_scan_left_join_shuffles_unknown_partitioning_before_hash_join() {
use arrow::array::Array;
let schema = int64_schema(&["k", "n"]);
let region = |id| RegionId::new(2048, id);
let make_batch = |rows: &[(i64, i64)]| {
record_batch(
schema.clone(),
vec![
Arc::new(Int64Vector::from_slice(
rows.iter().map(|row| row.0).collect::<Vec<_>>(),
)) as _,
Arc::new(Int64Vector::from_slice(
rows.iter().map(|row| row.1).collect::<Vec<_>>(),
)) as _,
],
)
};
// The peer vectors are intentionally differently ordered. Left rows include
// key 99, which has no right-side match and must survive the LEFT JOIN.
let left_regions = vec![region(3), region(1), region(2)];
let right_regions = vec![region(2), region(3), region(1)];
let left_responses = vec![
(region(1), make_batch(&[(1, 10), (2, 20)])),
(region(2), make_batch(&[(3, 30)])),
(region(3), make_batch(&[(99, 90)])),
];
let right_responses = vec![
(region(1), make_batch(&[(1, 2)])),
(region(2), make_batch(&[(2, 3)])),
(region(3), make_batch(&[(3, 4)])),
];
let mut left_scan = merge_scan_exec_with_handler(
left_regions,
schema.arrow_schema().as_ref().clone(),
Arc::new(TestRegionQueryHandler::new(left_responses)),
2,
);
left_scan.remote_dyn_filter_producer_id = Some(RemoteDynFilterProducerId::new(42));
let left = Arc::new(left_scan) as Arc<dyn ExecutionPlan>;
let mut right_scan = merge_scan_exec_with_handler(
right_regions,
schema.arrow_schema().as_ref().clone(),
Arc::new(TestRegionQueryHandler::new(right_responses)),
2,
);
right_scan.remote_dyn_filter_producer_id = Some(RemoteDynFilterProducerId::new(43));
let right = Arc::new(right_scan) as Arc<dyn ExecutionPlan>;
assert!(matches!(
left.output_partitioning(),
Partitioning::UnknownPartitioning(2)
));
assert!(matches!(
right.output_partitioning(),
Partitioning::UnknownPartitioning(2)
));
let join = Arc::new(
HashJoinExec::try_new(
left,
right,
vec![(
Arc::new(Column::new("k", 0)) as _,
Arc::new(Column::new("k", 0)) as _,
)],
None,
&JoinType::Left,
None,
PartitionMode::Partitioned,
NullEquality::NullEqualsNothing,
false,
)
.unwrap(),
) as Arc<dyn ExecutionPlan>;
let session = SessionStateBuilder::new()
.with_config(SessionConfig::new().with_target_partitions(2))
.with_default_features()
.build();
let optimized =
datafusion::physical_optimizer::ensure_requirements::EnsureRequirements::new()
.optimize(join, &session.config_options())
.unwrap();
let repartition_count = optimized
.children()
.iter()
.filter(|child| child.as_any().is::<RepartitionExec>())
.count();
assert!(
repartition_count > 0,
"EnsureRequirements must insert actual shuffles"
);
let mut producer_ids = Vec::new();
for child in optimized.children() {
let repartition = child
.as_any()
.downcast_ref::<RepartitionExec>()
.expect("EnsureRequirements must wrap both join inputs in a shuffle");
let scan = repartition
.input()
.downcast_ref::<MergeScanExec>()
.expect("actual shuffle must retain the MergeScan child");
producer_ids.push(scan.remote_dyn_filter_producer_id().unwrap());
}
producer_ids.sort_unstable();
assert_eq!(
producer_ids,
vec![
RemoteDynFilterProducerId::new(42),
RemoteDynFilterProducerId::new(43)
]
);
let batches = datafusion::physical_plan::collect(optimized, session.task_ctx())
.await
.unwrap();
let mut actual = Vec::new();
for batch in batches {
let lk = batch
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let ln = batch
.column(1)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let rn = batch
.column(3)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
for row in 0..batch.num_rows() {
actual.push((
lk.value(row),
ln.value(row),
(!rn.is_null(row)).then(|| rn.value(row)),
));
}
}
actual.sort_unstable();
assert_eq!(
actual,
vec![
(1, 10, Some(2)),
(2, 20, Some(3)),
(3, 30, Some(4)),
(99, 90, None)
]
);
}
#[test]
fn merge_scan_reports_only_deterministic_plan_bounds() {
use datafusion::functions_aggregate::expr_fn::count;
@@ -1874,7 +1940,6 @@ mod tests {
handler.clone(),
query_ctx,
1,
AliasMapping::new(),
Some(RemoteDynFilterProducerId::new(42)),
false,
)
@@ -2682,7 +2747,6 @@ mod tests {
handler,
QueryContext::arc(),
target_partition,
AliasMapping::new(),
None,
false,
)
@@ -3288,73 +3352,6 @@ mod tests {
assert_int64_batch(&batches[0], (11, 12));
}
#[test]
fn try_with_new_distribution_preserves_remote_dyn_filter_producer_id() {
let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42);
// Build a plan whose schema contains "col1"
let plan = LogicalPlanBuilder::empty(true)
.project(vec![lit(1i32).alias("col1")])
.unwrap()
.build()
.unwrap();
let schema = plan.schema().as_arrow().clone();
let table = TableName::new("catalog", "schema", "table");
let regions = vec![RegionId::new(1024, 1)];
let query_ctx = QueryContext::arc();
// Non-empty partition_cols so try_with_new_distribution can detect an overlap
let mut partition_cols = AliasMapping::new();
partition_cols.insert(
"col1".to_string(),
BTreeSet::from([ColumnExpr::new(Some(TableReference::bare("table")), "col1")]),
);
let session_state = SessionStateBuilder::new().build();
let handler = Arc::new(TestRegionQueryHandler::default());
let target_partition = 2;
let exec = MergeScanExec::new(
&session_state,
table,
regions,
plan,
&schema,
handler,
query_ctx,
target_partition,
partition_cols,
Some(remote_dyn_filter_producer_id),
false,
)
.unwrap();
assert_eq!(
exec.remote_dyn_filter_producer_id(),
Some(remote_dyn_filter_producer_id)
);
// A distribution that differs from the current partitioning but shares a
// column name present in partition_cols, so try_with_new_distribution
// produces a clone instead of returning None.
let new_dist = Distribution::KeyPartitioned(vec![
Arc::new(Column::new("col1", 0)),
Arc::new(Column::new("col2", 1)),
]);
let cloned = exec
.try_with_new_distribution(new_dist)
.expect("expected a cloned exec with overlapping partition col");
assert_eq!(
cloned.remote_dyn_filter_producer_id(),
Some(remote_dyn_filter_producer_id),
"try_with_new_distribution must preserve remote dynamic filter producer id"
);
}
#[test]
fn merge_scan_apply_expressions_exposes_remote_dyn_filter_id() {
let query_ctx = QueryContext::arc();
@@ -3397,7 +3394,6 @@ mod tests {
handler,
query_ctx,
1,
AliasMapping::new(),
Some(remote_dyn_filter_producer_id),
false,
)
@@ -3503,7 +3499,6 @@ mod tests {
Arc::new(TestRegionQueryHandler::default()),
QueryContext::arc(),
1,
AliasMapping::new(),
None,
true,
)
@@ -3572,7 +3567,6 @@ mod tests {
Arc::new(TestRegionQueryHandler::default()),
QueryContext::arc(),
1,
AliasMapping::new(),
None,
true,
)
src/query/src/dist_plan/planner.rs
-1
View File
@@ -228,7 +228,6 @@ impl ExtensionPlanner for DistExtensionPlanner {
self.region_query_handler.clone(),
query_ctx,
session.config().target_partitions(),
merge_scan.partition_cols().clone(),
merge_scan.remote_dyn_filter_producer_id(),
self.enable_per_region_metrics,
)?;
src/query/src/metrics.rs
+1 -2
View File
@@ -407,7 +407,7 @@ fn finalize_region_watermarks(merged: BTreeMap<u64, MergeState>) -> Vec<RegionWa
#[cfg(test)]
mod tests {
use std::collections::{BTreeMap, BTreeSet};
use std::collections::BTreeMap;
use std::sync::Arc;
use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
@@ -493,7 +493,6 @@ mod tests {
Arc::new(NoopRegionQueryHandler),
query_ctx,
1,
BTreeMap::<String, BTreeSet<datafusion_common::Column>>::new(),
Some(RemoteDynFilterProducerId::new(0)),
false,
)
src/query/src/optimizer.rs
-1
View File
@@ -22,7 +22,6 @@ pub(crate) mod insert_assignment;
pub(crate) mod json_schema_concretize;
pub(crate) mod json_type_concretize;
pub mod parallelize_scan;
pub mod pass_distribution;
pub mod promql_tsid_narrow_join;
pub mod remove_duplicate;
pub mod scan_hint;
src/query/src/optimizer/pass_distribution.rs
-362
View File
@@ -1,362 +0,0 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::sync::Arc;
use datafusion::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::projection::ProjectionExec;
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::{
ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions,
};
use datafusion_common::Result as DfResult;
use datafusion_physical_expr::Distribution;
use datafusion_physical_expr::utils::map_columns_before_projection;
use crate::dist_plan::MergeScanExec;
/// This is a [`PhysicalOptimizerRule`] to pass distribution requirement to
/// [`MergeScanExec`] to avoid unnecessary shuffling.
///
/// This rule is expected to be run before [`EnforceDistribution`].
///
/// [`EnforceDistribution`]: datafusion::physical_optimizer::enforce_distribution::EnforceDistribution
/// [`MergeScanExec`]: crate::dist_plan::MergeScanExec
#[derive(Debug)]
pub struct PassDistribution;
impl PhysicalOptimizerRule for PassDistribution {
fn optimize(
&self,
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> DfResult<Arc<dyn ExecutionPlan>> {
Self::do_optimize(plan, config)
}
fn name(&self) -> &str {
"PassDistributionRule"
}
fn schema_check(&self) -> bool {
false
}
}
impl PassDistribution {
fn do_optimize(
plan: Arc<dyn ExecutionPlan>,
_config: &ConfigOptions,
) -> DfResult<Arc<dyn ExecutionPlan>> {
// Start from root with no requirement
Self::rewrite_with_distribution(plan, None)
}
/// Top-down rewrite that propagates distribution requirements to children.
fn rewrite_with_distribution(
plan: Arc<dyn ExecutionPlan>,
current_req: Option<Distribution>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
// If this is a MergeScanExec, try to apply the current requirement.
if let Some(merge_scan) = plan.downcast_ref::<MergeScanExec>()
&& let Some(Distribution::KeyPartitioned(hash_exprs)) = current_req.as_ref()
{
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::KeyPartitioned(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.
let children = plan.children();
if children.is_empty() {
return Ok(plan);
}
let required = plan.input_distribution_requirements();
let mut new_children = Vec::with_capacity(children.len());
for (idx, child) in children.into_iter().enumerate() {
let child_req = match required.child_distribution(idx) {
Some(Distribution::UnspecifiedDistribution) if idx == 0 => {
Self::map_hash_requirement_through_projection(plan.as_ref(), &current_req)
}
Some(Distribution::UnspecifiedDistribution) => None,
None => current_req.clone(),
Some(req) => Some(req.clone()),
};
let new_child = Self::rewrite_with_distribution(child.clone(), child_req)?;
new_children.push(new_child);
}
// Rebuild the node only if any child changed (pointer inequality)
let unchanged = plan
.children()
.into_iter()
.zip(new_children.iter())
.all(|(old, new)| Arc::ptr_eq(old, new));
if unchanged {
Ok(plan)
} else {
plan.replace_children(
new_children,
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
)
}
}
fn map_hash_requirement_through_projection(
plan: &dyn ExecutionPlan,
current_req: &Option<Distribution>,
) -> Option<Distribution> {
let Some(Distribution::KeyPartitioned(required_exprs)) = current_req else {
return None;
};
let projection = plan.downcast_ref::<ProjectionExec>()?;
let proj_exprs = projection
.expr()
.iter()
.map(|expr| (Arc::clone(&expr.expr), expr.alias.clone()))
.collect::<Vec<_>>();
let mapped = map_columns_before_projection(required_exprs, &proj_exprs);
(mapped.len() == required_exprs.len()).then_some(Distribution::KeyPartitioned(mapped))
}
}
#[cfg(test)]
mod tests {
use std::collections::{BTreeMap, BTreeSet};
use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit};
use async_trait::async_trait;
use common_query::request::QueryRequest;
use datafusion::common::NullEquality;
use datafusion::execution::SessionStateBuilder;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr};
use datafusion::physical_plan::{ExecutionPlanProperties, Partitioning};
use datafusion_expr::{JoinType, LogicalPlanBuilder};
use datafusion_physical_expr::PhysicalExpr;
use datafusion_physical_expr::expressions::Column as PhysicalColumn;
use session::ReadPreference;
use session::context::QueryContext;
use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
use store_api::storage::RegionId;
use table::table_name::TableName;
use super::*;
use crate::dist_plan::RemoteDynFilterProducerId;
use crate::error::Result as QueryResult;
use crate::region_query::RegionQueryHandler;
struct NoopRegionQueryHandler;
#[async_trait]
impl RegionQueryHandler for NoopRegionQueryHandler {
async fn select_target(
&self,
_read_preference: ReadPreference,
_region_id: RegionId,
) -> QueryResult<crate::region_query::RegionQueryTarget> {
unreachable!("pass distribution tests should not execute remote queries")
}
async fn do_get(
&self,
_target: &crate::region_query::RegionQueryTarget,
_request: QueryRequest,
) -> QueryResult<common_recordbatch::SendableRecordBatchStream> {
unreachable!("pass distribution tests should not execute remote queries")
}
async fn handle_remote_dyn_filter_update(
&self,
_target: &crate::region_query::RegionQueryTarget,
_query_id: String,
_update: RemoteDynFilterUpdate,
) -> QueryResult<()> {
unreachable!("pass distribution tests should not send remote dyn filter updates")
}
async fn handle_remote_dyn_filter_unregister(
&self,
_target: &crate::region_query::RegionQueryTarget,
_query_id: String,
_unregister: RemoteDynFilterUnregister,
) -> QueryResult<()> {
unreachable!("pass distribution tests should not send remote dyn filter unregisters")
}
}
#[test]
fn passes_hash_requirement_through_projection_to_merge_scan() {
let schema = test_schema();
let left_merge_scan = Arc::new(test_merge_scan_exec(schema.clone()));
let right_merge_scan = Arc::new(test_merge_scan_exec(schema.clone()));
let left_projection = Arc::new(
ProjectionExec::try_new(
vec![
ProjectionExpr::new(partition_column("greptime_value", 3), "greptime_value"),
ProjectionExpr::new(
partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1),
DATA_SCHEMA_TSID_COLUMN_NAME,
),
ProjectionExpr::new(
partition_column("greptime_timestamp", 2),
"greptime_timestamp",
),
],
left_merge_scan,
)
.unwrap(),
) as Arc<dyn datafusion::physical_plan::ExecutionPlan>;
let join = Arc::new(
HashJoinExec::try_new(
left_projection,
right_merge_scan,
vec![
(
partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1),
partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1),
),
(
partition_column("greptime_timestamp", 2),
partition_column("greptime_timestamp", 2),
),
],
None,
&JoinType::Inner,
None,
PartitionMode::Partitioned,
NullEquality::NullEqualsNull,
false,
)
.unwrap(),
) as Arc<dyn datafusion::physical_plan::ExecutionPlan>;
let optimized = PassDistribution
.optimize(join, &ConfigOptions::default())
.unwrap();
let hash_join = optimized.downcast_ref::<HashJoinExec>().unwrap();
let left_projection = hash_join.left().downcast_ref::<ProjectionExec>().unwrap();
let left_partitioning = left_projection.input().output_partitioning();
let right_partitioning = hash_join.right().output_partitioning();
let Partitioning::Hash(left_exprs, left_count) = left_partitioning else {
panic!("expected left merge scan hash partitioning");
};
let Partitioning::Hash(right_exprs, right_count) = right_partitioning else {
panic!("expected right merge scan hash partitioning");
};
assert_eq!(*left_count, 2);
assert_eq!(*right_count, 2);
assert_eq!(
column_names(left_exprs),
vec![DATA_SCHEMA_TSID_COLUMN_NAME, "greptime_timestamp"]
);
assert_eq!(
column_names(right_exprs),
vec![DATA_SCHEMA_TSID_COLUMN_NAME, "greptime_timestamp"]
);
}
#[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::KeyPartitioned(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([
(
DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
BTreeSet::from([datafusion_common::Column::from_name(
DATA_SCHEMA_TSID_COLUMN_NAME,
)]),
),
(
"greptime_timestamp".to_string(),
BTreeSet::from([datafusion_common::Column::from_name("greptime_timestamp")]),
),
]);
let plan = LogicalPlanBuilder::empty(false).build().unwrap();
MergeScanExec::new(
&session_state,
TableName::new("greptime", "public", "test"),
vec![RegionId::new(1, 0), RegionId::new(1, 1)],
plan,
schema.as_ref(),
Arc::new(NoopRegionQueryHandler),
QueryContext::arc(),
32,
partition_cols,
Some(RemoteDynFilterProducerId::new(1)),
false,
)
.unwrap()
}
fn test_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("host", DataType::Utf8, true),
Field::new(DATA_SCHEMA_TSID_COLUMN_NAME, DataType::UInt64, false),
Field::new(
"greptime_timestamp",
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
Field::new("greptime_value", DataType::Float64, true),
]))
}
fn partition_column(name: &str, index: usize) -> Arc<dyn PhysicalExpr> {
Arc::new(PhysicalColumn::new(name, index))
}
fn column_names(exprs: &[Arc<dyn PhysicalExpr>]) -> Vec<&str> {
exprs
.iter()
.map(|expr| expr.downcast_ref::<PhysicalColumn>().unwrap().name())
.collect()
}
}
src/query/src/query_engine/state.rs
+2 -7
View File
@@ -70,7 +70,6 @@ use crate::optimizer::global_limit::EnsureGlobalLimitForFetch;
use crate::optimizer::json_schema_concretize::JsonSchemaConcretizeRule;
use crate::optimizer::json_type_concretize::JsonTypeConcretizeRule;
use crate::optimizer::parallelize_scan::ParallelizeScan;
use crate::optimizer::pass_distribution::PassDistribution;
use crate::optimizer::promql_tsid_narrow_join::PromqlTsidNarrowJoin;
use crate::optimizer::remove_duplicate::RemoveDuplicate;
use crate::optimizer::scan_hint::ScanHintRule;
@@ -232,17 +231,13 @@ impl QueryEngineState {
physical_optimizer
.rules
.insert(5, Arc::new(ParallelizeScan));
// Pass distribution requirement to MergeScanExec to avoid unnecessary shuffling
physical_optimizer
.rules
.insert(6, Arc::new(PassDistribution));
// Prefer collecting narrow PromQL build sides over repartitioning wide label streams.
physical_optimizer
.rules
.insert(7, Arc::new(PromqlTsidNarrowJoin));
.insert(6, Arc::new(PromqlTsidNarrowJoin));
// Re-enforce sorting after custom rules update scan partitioning and distribution.
// Keep it immediately before DataFusion's default EnsureRequirements.
physical_optimizer.rules.insert(8, Arc::new(EnforceSorting));
physical_optimizer.rules.insert(7, Arc::new(EnforceSorting));
// Add rule for windowed sort
physical_optimizer
.rules
tests-integration/src/tests/instance_test.rs
+115
View File
@@ -257,6 +257,121 @@ WHERE greptime_timestamp = (
);
}
#[apply(both_instances_cases)]
async fn test_partitioned_merge_scan_left_join_repartitions(instance: Arc<dyn MockInstance>) {
let frontend = instance.frontend();
execute_sql(
&frontend,
r#"
CREATE TABLE join_repro (
k STRING,
kind STRING,
v DOUBLE,
ts TIMESTAMP NOT NULL,
TIME INDEX (ts)
)
PARTITION ON COLUMNS (k) (
k < '1',
k >= '1' AND k < '2',
k >= '2' AND k < '3',
k >= '3' AND k < '4',
k >= '4' AND k < '5',
k >= '5' AND k < '6',
k >= '6' AND k < '7',
k >= '7' AND k < '8',
k >= '8' AND k < '9',
k >= '9' AND k < 'a',
k >= 'a' AND k < 'b',
k >= 'b' AND k < 'c',
k >= 'c' AND k < 'd',
k >= 'd' AND k < 'e',
k >= 'e' AND k < 'f',
k >= 'f'
)
ENGINE = mito
WITH (append_mode = 'true')
"#,
)
.await;
let mut rows = Vec::new();
for key in [
'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f', 'z',
] {
rows.push(format!("('{key}', 'root', 1.0, '2024-01-01 00:00:00')"));
if key != 'z' {
for _ in 0..3 {
rows.push(format!("('{key}', 'child', 1.0, '2024-01-01 00:00:00')"));
}
}
}
execute_sql(
&frontend,
&format!("INSERT INTO join_repro VALUES {}", rows.join(",")),
)
.await;
let aggregate_sql = r#"SELECT count(*) roots, count(c.n) roots_with_children, sum(c.n) children
FROM (SELECT k FROM join_repro WHERE kind = 'root') r
LEFT JOIN (SELECT k, count(*) n FROM join_repro WHERE kind = 'child' GROUP BY k) c ON c.k = r.k"#;
for parallelism in [16, 8] {
let output = execute_sql_with_query_parallelism(&frontend, aggregate_sql, parallelism)
.await
.data
.pretty_print()
.await;
assert_eq!(
output,
"+-------+---------------------+----------+\n| roots | roots_with_children | children |\n+-------+---------------------+----------+\n| 17 | 16 | 48 |\n+-------+---------------------+----------+"
);
let explain = execute_sql_with_query_parallelism(
&frontend,
&format!("EXPLAIN {aggregate_sql}"),
parallelism,
)
.await
.data
.pretty_print()
.await;
assert!(
explain.contains("HashJoinExec: mode=Partitioned"),
"query_parallelism={parallelism} should use a partitioned hash join; explain:\n{explain}"
);
assert!(
explain.contains("RepartitionExec: partitioning=Hash([k@0]"),
"query_parallelism={parallelism} should repartition on the join key; explain:\n{explain}"
);
let ordered_sql = r#"SELECT r.k, c.n
FROM (SELECT k FROM join_repro WHERE kind = 'root') r
LEFT JOIN (SELECT k, count(*) n FROM join_repro WHERE kind = 'child' GROUP BY k) c ON c.k = r.k
ORDER BY r.k"#;
let ordered = execute_sql_with_query_parallelism(&frontend, ordered_sql, parallelism)
.await
.data
.pretty_print()
.await;
assert_eq!(
ordered,
"+---+------+\n| k | n |\n+---+------+\n| 0 | 3 |\n| 1 | 3 |\n| 2 | 3 |\n| 3 | 3 |\n| 4 | 3 |\n| 5 | 3 |\n| 6 | 3 |\n| 7 | 3 |\n| 8 | 3 |\n| 9 | 3 |\n| a | 3 |\n| b | 3 |\n| c | 3 |\n| d | 3 |\n| e | 3 |\n| f | 3 |\n| z | NULL |\n+---+------+"
);
let pruned_sql = r#"SELECT count(*) roots, count(c.n) roots_with_children, sum(c.n) children
FROM (SELECT k FROM join_repro WHERE kind = 'root' AND k < '8') r
LEFT JOIN (SELECT k, count(*) n FROM join_repro WHERE kind = 'child' AND k < 'c' GROUP BY k) c ON c.k = r.k"#;
let pruned = execute_sql_with_query_parallelism(&frontend, pruned_sql, parallelism)
.await
.data
.pretty_print()
.await;
assert_eq!(
pruned,
"+-------+---------------------+----------+\n| roots | roots_with_children | children |\n+-------+---------------------+----------+\n| 8 | 8 | 24 |\n+-------+---------------------+----------+"
);
}
}
#[apply(both_instances_cases)]
async fn test_show_create_table(instance: Arc<dyn MockInstance>) {
let frontend = instance.frontend();