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

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

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

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): use current DataFusion plan downcasts

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): match null rendering in join regression

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): refresh union plan for unknown scan partitioning

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): refresh shared plans after distribution rule removal

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): cover semi joins with remote dynamic filtering

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): strengthen partitioned merge scan regressions

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): remove unused merge scan test helpers

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(perf): measure merge scan hash repartition cost

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(query): adapt merge scan deserialization to two-arg constructor

Adapt the production and test serializer callers added by #9206 and remove the obsolete partition_cols assertion.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): adapt nested merge scan integration constructor

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

---------

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-10 11:18:46 +00:00
committed by GitHub
parent ad83232ad1
commit 32e8467cda
21 changed files with 851 additions and 673 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/query/src/dist_plan/analyzer.rs
+1 -8
View File
@@ -768,14 +768,7 @@ impl PlanRewriter {
);
// add merge scan as the new root
let mut node = MergeScanLogicalPlan::new(
on_node.clone(),
false,
// at this stage, the partition cols should be set
// treat it as non-partitioned if None
self.partition_cols.clone().unwrap_or_default(),
)
.into_logical_plan();
let mut node = MergeScanLogicalPlan::new(on_node.clone(), false).into_logical_plan();
// expand stages
for new_stage in self.stage.drain(..) {
src/query/src/dist_plan/analyzer/fallback.rs
+2 -70
View File
@@ -17,18 +17,11 @@
//! This is a temporary solution, and will be removed once we have a more robust plan rewriter
//!
use std::collections::BTreeSet;
use common_telemetry::debug;
use datafusion::datasource::DefaultTableSource;
use datafusion_common::Result as DfResult;
use datafusion_common::tree_node::{Transformed, TreeNodeRewriter};
use datafusion_expr::LogicalPlan;
use table::metadata::TableType;
use table::table::adapter::DfTableProviderAdapter;
use crate::dist_plan::MergeScanLogicalPlan;
use crate::dist_plan::analyzer::{AliasMapping, OTHER_PHY_PART_COL_PLACEHOLDER};
/// FallbackPlanRewriter is a plan rewriter that will only push down table scan node
/// This is used when `PlanRewriter` produce errors when trying to rewrite the plan
@@ -44,69 +37,8 @@ impl TreeNodeRewriter for FallbackPlanRewriter {
&mut self,
plan: Self::Node,
) -> DfResult<datafusion_common::tree_node::Transformed<Self::Node>> {
if let LogicalPlan::TableScan(table_scan) = &plan {
let partition_cols = if let Some(source) =
table_scan.source.downcast_ref::<DefaultTableSource>()
{
if let Some(provider) = source
.table_provider
.downcast_ref::<DfTableProviderAdapter>()
{
if provider.table().table_type() == TableType::Base {
let info = provider.table().table_info();
let partition_key_indices = info.meta.partition_key_indices.clone();
let schema = info.meta.schema.clone();
let partition_cols = partition_key_indices
.iter()
.map(|index| schema.column_name_by_index(*index).to_string())
.collect::<Vec<String>>();
debug!(
"FallbackPlanRewriter: loaded table partition metadata, table: {}, table_id: {}, partition_key_indices: {:?}, partition_columns: {:?}",
info.name,
info.ident.table_id,
info.meta.partition_key_indices,
partition_cols,
);
Some(partition_cols
.into_iter()
.map(|c| {
if c == OTHER_PHY_PART_COL_PLACEHOLDER {
// for placeholder, just return a empty alias
return Ok((c.clone(), BTreeSet::new()));
}
let index =
plan.schema().index_of_column_by_name(None, &c).ok_or_else(|| {
datafusion_common::DataFusionError::Internal(
format!(
"PlanRewriter: maybe_set_partitions: column {c} not found in schema of plan: {plan}"
),
)
})?;
let column = plan.schema().columns().get(index).cloned().ok_or_else(|| {
datafusion_common::DataFusionError::Internal(format!(
"PlanRewriter: maybe_set_partitions: column index {index} out of bounds in schema of plan: {plan}"
))
})?;
Ok((c.clone(), BTreeSet::from([column])))
})
.collect::<DfResult<AliasMapping>>()?)
} else {
None
}
} else {
None
}
} else {
None
};
let node = MergeScanLogicalPlan::new(
plan,
false,
// at this stage, the partition cols should be set
// treat it as non-partitioned if None
partition_cols.clone().unwrap_or_default(),
)
.into_logical_plan();
if matches!(plan, LogicalPlan::TableScan(_)) {
let node = MergeScanLogicalPlan::new(plan, false).into_logical_plan();
Ok(Transformed::yes(node))
} else {
Ok(Transformed::no(plan))
src/query/src/dist_plan/merge_scan.rs
+18 -203
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,9 @@ 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;
@@ -64,7 +63,6 @@ use tokio::time;
use tokio::time::Instant;
use tracing::{Instrument, Span};
use crate::dist_plan::analyzer::AliasMapping;
use crate::dist_plan::analyzer::utils::patch_batch_timezone;
use crate::dist_plan::dyn_filter_bridge::{
CapturedDynFilter, capture_remote_dyn_filters_for_pushdown,
@@ -339,7 +337,6 @@ pub struct MergeScanLogicalPlan {
output_schema: DFSchemaRef,
/// If this plan is a placeholder
is_placeholder: bool,
partition_cols: AliasMapping,
/// Assigned after dist-plan rewriting so rewriters only deal with plan shape.
remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
}
@@ -350,29 +347,20 @@ impl PartialOrd for MergeScanLogicalPlan {
input,
output_schema,
is_placeholder,
partition_cols,
remote_dyn_filter_producer_id,
} = self;
let Self {
input: other_input,
output_schema: other_output_schema,
is_placeholder: other_is_placeholder,
partition_cols: other_partition_cols,
remote_dyn_filter_producer_id: other_remote_dyn_filter_producer_id,
} = other;
let ordering = (
input,
is_placeholder,
partition_cols,
remote_dyn_filter_producer_id,
)
.partial_cmp(&(
other_input,
other_is_placeholder,
other_partition_cols,
other_remote_dyn_filter_producer_id,
));
let ordering = (input, is_placeholder, remote_dyn_filter_producer_id).partial_cmp(&(
other_input,
other_is_placeholder,
other_remote_dyn_filter_producer_id,
));
match ordering {
Some(std::cmp::Ordering::Equal) if output_schema != other_output_schema => None,
ordering => ordering,
@@ -419,12 +407,11 @@ impl UserDefinedLogicalNodeCore for MergeScanLogicalPlan {
impl MergeScanLogicalPlan {
/// Creates a merge scan with the input plan's schema.
pub fn new(input: LogicalPlan, is_placeholder: bool, partition_cols: AliasMapping) -> Self {
pub fn new(input: LogicalPlan, is_placeholder: bool) -> Self {
Self {
output_schema: input.schema().clone(),
input,
is_placeholder,
partition_cols,
remote_dyn_filter_producer_id: None,
}
}
@@ -462,10 +449,6 @@ impl MergeScanLogicalPlan {
&self.input
}
pub fn partition_cols(&self) -> &AliasMapping {
&self.partition_cols
}
pub fn remote_dyn_filter_producer_id(&self) -> Option<RemoteDynFilterProducerId> {
self.remote_dyn_filter_producer_id
}
@@ -489,7 +472,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 +496,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 +532,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 +555,6 @@ impl MergeScanExec {
remote_dyn_filter_producer_id,
captured_remote_dyn_filters: Arc::default(),
target_partition,
partition_cols,
enable_per_region_metrics,
})
}
@@ -912,78 +876,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()
}
@@ -1037,13 +929,6 @@ impl MergeScanExec {
}
}
#[cfg(test)]
impl MergeScanExec {
fn remote_dyn_filter_producer_id(&self) -> Option<RemoteDynFilterProducerId> {
self.remote_dyn_filter_producer_id
}
}
/// Metrics for a region of a partition.
#[derive(Debug, Clone)]
struct RegionMetrics {
@@ -1420,7 +1305,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};
@@ -1441,15 +1326,15 @@ mod tests {
};
use datafusion::config::ConfigOptions;
use datafusion::execution::SessionStateBuilder;
use datafusion::physical_plan::Partitioning;
use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult;
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::{StatisticsArgs, StatisticsContext};
use datafusion_common::TableReference;
use datafusion_expr::{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 +1402,6 @@ mod tests {
Arc::new(TestRegionQueryHandler::default()),
QueryContext::arc(),
target_partition,
AliasMapping::new(),
None,
false,
)
@@ -1560,7 +1444,6 @@ mod tests {
handler,
query_ctx,
1,
AliasMapping::new(),
Some(RemoteDynFilterProducerId::new(42)),
false,
)
@@ -1661,6 +1544,10 @@ mod tests {
exec.properties().output_partitioning().partition_count(),
expected
);
assert!(matches!(
exec.properties().output_partitioning(),
Partitioning::UnknownPartitioning(count) if *count == expected
));
}
}
@@ -1874,7 +1761,6 @@ mod tests {
handler.clone(),
query_ctx,
1,
AliasMapping::new(),
Some(RemoteDynFilterProducerId::new(42)),
false,
)
@@ -2682,7 +2568,6 @@ mod tests {
handler,
QueryContext::arc(),
target_partition,
AliasMapping::new(),
None,
false,
)
@@ -3288,73 +3173,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 +3215,6 @@ mod tests {
handler,
query_ctx,
1,
AliasMapping::new(),
Some(remote_dyn_filter_producer_id),
false,
)
@@ -3503,7 +3320,6 @@ mod tests {
Arc::new(TestRegionQueryHandler::default()),
QueryContext::arc(),
1,
AliasMapping::new(),
None,
true,
)
@@ -3572,7 +3388,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/json_schema_concretize.rs
+1 -2
View File
@@ -192,8 +192,7 @@ mod tests {
produce_one_row: false,
schema,
});
let merge_scan =
MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan();
let merge_scan = MergeScanLogicalPlan::new(input, false).into_logical_plan();
let plan = LogicalPlanBuilder::from(merge_scan)
.project(vec![col("j")])?
.build()?;
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/default_serializer.rs
+3 -10
View File
@@ -220,11 +220,7 @@ impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> {
let input = self.decode_payload(merge_scan.input).await?;
// `PbMergeScan` lacks `partition_cols`; decoded plans lose this optimization and may repartition.
Ok(
MergeScanLogicalPlan::new(input, merge_scan.is_placeholder, Default::default())
.into_logical_plan(),
)
Ok(MergeScanLogicalPlan::new(input, merge_scan.is_placeholder).into_logical_plan())
}
// Keep child dispatch on this consumer; delegating the parent bypasses async MergeScan decoding.
@@ -739,8 +735,6 @@ mod tests {
.expect("Expect a MergeScan plan node");
assert!(merge_scan.is_placeholder());
assert_eq!(merge_scan.input().to_string(), numbers_scan().to_string());
// `PbMergeScan` doesn't carry `partition_cols`, so decoded nodes have an empty mapping.
assert!(merge_scan.partition_cols().is_empty());
}
/// Payload table binding uses the engine catalog, not the request catalog.
@@ -979,8 +973,7 @@ mod tests {
.unwrap()
.build()
.unwrap();
let plan =
MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan();
let plan = MergeScanLogicalPlan::new(input, false).into_logical_plan();
let bytes = DFLogicalSubstraitConvertor
.encode(&plan, DefaultSerializer)
@@ -1041,7 +1034,7 @@ mod tests {
}
fn merge_scan(input: LogicalPlan, is_placeholder: bool) -> LogicalPlan {
MergeScanLogicalPlan::new(input, is_placeholder, Default::default()).into_logical_plan()
MergeScanLogicalPlan::new(input, is_placeholder).into_logical_plan()
}
fn encode_plan(plan: &LogicalPlan) -> Bytes {
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