From a8159b29a8471f228c40c54371ed7dffcec04bdd Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:03:09 +0800 Subject: [PATCH] fix(query): stop reporting region scans as hash partitioned Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/merge_scan.rs | 344 +++++++++--------- src/query/src/dist_plan/planner.rs | 1 - src/query/src/metrics.rs | 3 +- src/query/src/optimizer.rs | 1 - src/query/src/optimizer/pass_distribution.rs | 362 ------------------- src/query/src/query_engine/state.rs | 9 +- tests-integration/src/tests/instance_test.rs | 115 ++++++ 7 files changed, 287 insertions(+), 548 deletions(-) delete mode 100644 src/query/src/optimizer/pass_distribution.rs diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index ee4e9574280..d7f5e6ffda2 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -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, captured_remote_dyn_filters: Arc>>, 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, enable_per_region_metrics: bool, ) -> Result { @@ -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 { - 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::() - .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::() - .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 { 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::>(), + )) as _, + Arc::new(Int64Vector::from_slice( + rows.iter().map(|row| row.1).collect::>(), + )) 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; + 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; + 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; + 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::()) + .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::() + .expect("EnsureRequirements must wrap both join inputs in a shuffle"); + let scan = repartition + .input() + .downcast_ref::() + .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::() + .unwrap(); + let ln = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let rn = batch + .column(3) + .as_any() + .downcast_ref::() + .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, ) diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index b2a2ade9b4f..88b13caf388 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -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, )?; diff --git a/src/query/src/metrics.rs b/src/query/src/metrics.rs index d6d62c99c4f..3c0043dcc09 100644 --- a/src/query/src/metrics.rs +++ b/src/query/src/metrics.rs @@ -407,7 +407,7 @@ fn finalize_region_watermarks(merged: BTreeMap) -> Vec>::new(), Some(RemoteDynFilterProducerId::new(0)), false, ) diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index fec001e7506..03d6ce3aa46 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -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; diff --git a/src/query/src/optimizer/pass_distribution.rs b/src/query/src/optimizer/pass_distribution.rs deleted file mode 100644 index 732168e962f..00000000000 --- a/src/query/src/optimizer/pass_distribution.rs +++ /dev/null @@ -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, - config: &ConfigOptions, - ) -> DfResult> { - Self::do_optimize(plan, config) - } - - fn name(&self) -> &str { - "PassDistributionRule" - } - - fn schema_check(&self) -> bool { - false - } -} - -impl PassDistribution { - fn do_optimize( - plan: Arc, - _config: &ConfigOptions, - ) -> DfResult> { - // 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, - current_req: Option, - ) -> DfResult> { - // If this is a MergeScanExec, try to apply the current requirement. - if let Some(merge_scan) = plan.downcast_ref::() - && 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(), ¤t_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, - ) -> Option { - let Some(Distribution::KeyPartitioned(required_exprs)) = current_req else { - return None; - }; - - let projection = plan.downcast_ref::()?; - let proj_exprs = projection - .expr() - .iter() - .map(|expr| (Arc::clone(&expr.expr), expr.alias.clone())) - .collect::>(); - 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 { - unreachable!("pass distribution tests should not execute remote queries") - } - - async fn do_get( - &self, - _target: &crate::region_query::RegionQueryTarget, - _request: QueryRequest, - ) -> QueryResult { - 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; - 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; - - let optimized = PassDistribution - .optimize(join, &ConfigOptions::default()) - .unwrap(); - let hash_join = optimized.downcast_ref::().unwrap(); - let left_projection = hash_join.left().downcast_ref::().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 { - Arc::new(PhysicalColumn::new(name, index)) - } - - fn column_names(exprs: &[Arc]) -> Vec<&str> { - exprs - .iter() - .map(|expr| expr.downcast_ref::().unwrap().name()) - .collect() - } -} diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index aa3b9417b14..e6db7db10f2 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -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 diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index a1ccac27328..c7ac8ae4c37 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -257,6 +257,121 @@ WHERE greptime_timestamp = ( ); } +#[apply(both_instances_cases)] +async fn test_partitioned_merge_scan_left_join_repartitions(instance: Arc) { + 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) { let frontend = instance.frontend();