fix exchange

zhyass · zhyass · commit a116b11866bd · 2025-05-09T02:51:06.000+08:00
diff --git a/src/query/service/src/interpreters/interpreter_table_recluster.rs b/src/query/service/src/interpreters/interpreter_table_recluster.rs
@@ -467,7 +467,7 @@ impl ReclusterTableInterpreter {
             plan = Box::new(PhysicalPlan::Exchange(Exchange {
                 plan_id: 0,
                 input: plan,
-                kind: FragmentKind::Normal,
+                kind: FragmentKind::Modulo,
                 keys: vec![expr],
                 allow_adjust_parallelism: true,
                 ignore_exchange: false,
diff --git a/src/query/service/src/pipelines/builders/builder_hilbert_partition.rs b/src/query/service/src/pipelines/builders/builder_hilbert_partition.rs
@@ -99,6 +99,7 @@ impl PipelineBuilder {
                     self.ctx.clone(),
                     input,
                     output,
+                    MutationKind::Recluster,
                     table,
                     partition.table_meta_timestamps,
                     false,
diff --git a/src/query/service/src/pipelines/processors/transforms/aggregator/aggregate_exchange_injector.rs b/src/query/service/src/pipelines/processors/transforms/aggregator/aggregate_exchange_injector.rs
@@ -233,6 +233,7 @@ impl ExchangeInjector for AggregateInjector {
         match exchange {
             DataExchange::Merge(_) => unreachable!(),
             DataExchange::Broadcast(_) => unreachable!(),
+            DataExchange::Modulo(_) => unreachable!(),
             DataExchange::ShuffleDataExchange(exchange) => {
                 Ok(Arc::new(Box::new(HashTableHashScatter {
                     buckets: exchange.destination_ids.len(),
diff --git a/src/query/service/src/schedulers/fragments/fragmenter.rs b/src/query/service/src/schedulers/fragments/fragmenter.rs
@@ -15,6 +15,7 @@
 use std::sync::Arc;
 
 use databend_common_catalog::table_context::TableContext;
+use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_meta_types::NodeInfo;
 use databend_common_sql::executor::physical_plans::CompactSource;
@@ -40,6 +41,7 @@ use crate::schedulers::PlanFragment;
 use crate::servers::flight::v1::exchange::BroadcastExchange;
 use crate::servers::flight::v1::exchange::DataExchange;
 use crate::servers::flight::v1::exchange::MergeExchange;
+use crate::servers::flight::v1::exchange::ModuloExchange;
 use crate::servers::flight::v1::exchange::ShuffleDataExchange;
 use crate::sessions::QueryContext;
 use crate::sql::executor::PhysicalPlan;
@@ -116,6 +118,15 @@ impl Fragmenter {
                 FragmentKind::Expansive => {
                     Ok(Some(BroadcastExchange::create(Self::get_executors(ctx))))
                 }
+                FragmentKind::Modulo => {
+                    if plan.keys.len() != 1 {
+                        return Err(ErrorCode::Internal("Modulo exchange require one key"));
+                    }
+                    Ok(Some(ModuloExchange::create(
+                        Self::get_executors(ctx),
+                        plan.keys[0].clone(),
+                    )))
+                }
                 _ => Ok(None),
             },
             _ => Ok(None),
@@ -203,8 +214,12 @@ impl PhysicalPlanReplacer for Fragmenter {
     }
 
     fn replace_hilbert_serialize(&mut self, plan: &HilbertPartition) -> Result<PhysicalPlan> {
+        let input = self.replace(&plan.input)?;
         self.state = State::HilbertRecluster;
-        Ok(PhysicalPlan::HilbertPartition(Box::new(plan.clone())))
+        Ok(PhysicalPlan::HilbertPartition(Box::new(HilbertPartition {
+            input: Box::new(input),
+            ..plan.clone()
+        })))
     }
 
     fn replace_compact_source(&mut self, plan: &CompactSource) -> Result<PhysicalPlan> {
diff --git a/src/query/service/src/schedulers/fragments/plan_fragment.rs b/src/query/service/src/schedulers/fragments/plan_fragment.rs
@@ -602,7 +602,9 @@ struct ReplaceHilbert {
 
 impl PhysicalPlanReplacer for ReplaceHilbert {
     fn replace_hilbert_serialize(&mut self, plan: &HilbertPartition) -> Result<PhysicalPlan> {
+        let input = self.replace(&plan.input)?;
         Ok(PhysicalPlan::HilbertPartition(Box::new(HilbertPartition {
+            input: Box::new(input),
             range_width: self.range_width,
             range_start: self.range_start,
             ..plan.clone()
diff --git a/src/query/service/src/schedulers/fragments/query_fragment_actions_display.rs b/src/query/service/src/schedulers/fragments/query_fragment_actions_display.rs
@@ -72,6 +72,7 @@ impl Display for QueryFragmentActionsWrap<'_> {
                 DataExchange::Merge(_) => writeln!(f, "  DataExchange: Merge")?,
                 DataExchange::Broadcast(_) => writeln!(f, "  DataExchange: Broadcast")?,
                 DataExchange::ShuffleDataExchange(_) => writeln!(f, "  DataExchange: Shuffle")?,
+                DataExchange::Modulo(_) => writeln!(f, "  DataExchange: Modulo")?,
             }
         }
 
diff --git a/src/query/service/src/servers/flight/v1/exchange/data_exchange.rs b/src/query/service/src/servers/flight/v1/exchange/data_exchange.rs
@@ -19,6 +19,7 @@ pub enum DataExchange {
     Merge(MergeExchange),
     Broadcast(BroadcastExchange),
     ShuffleDataExchange(ShuffleDataExchange),
+    Modulo(ModuloExchange),
 }
 
 impl DataExchange {
@@ -27,6 +28,7 @@ impl DataExchange {
             DataExchange::Merge(exchange) => vec![exchange.destination_id.clone()],
             DataExchange::Broadcast(exchange) => exchange.destination_ids.clone(),
             DataExchange::ShuffleDataExchange(exchange) => exchange.destination_ids.clone(),
+            DataExchange::Modulo(exchange) => exchange.destination_ids.clone(),
         }
     }
 }
@@ -77,3 +79,18 @@ impl BroadcastExchange {
         DataExchange::Broadcast(BroadcastExchange { destination_ids })
     }
 }
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct ModuloExchange {
+    pub destination_ids: Vec<String>,
+    pub shuffle_key: RemoteExpr,
+}
+
+impl ModuloExchange {
+    pub fn create(destination_ids: Vec<String>, shuffle_key: RemoteExpr) -> DataExchange {
+        DataExchange::Modulo(ModuloExchange {
+            destination_ids,
+            shuffle_key,
+        })
+    }
+}
diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_injector.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_injector.rs
@@ -29,6 +29,7 @@ use crate::servers::flight::v1::exchange::ShuffleExchangeParams;
 use crate::servers::flight::v1::scatter::BroadcastFlightScatter;
 use crate::servers::flight::v1::scatter::FlightScatter;
 use crate::servers::flight::v1::scatter::HashFlightScatter;
+use crate::servers::flight::v1::scatter::ModFlightScatter;
 use crate::sessions::QueryContext;
 
 pub trait ExchangeInjector: Send + Sync + 'static {
@@ -100,6 +101,11 @@ impl ExchangeInjector for DefaultExchangeInjector {
                     local_pos,
                 )?
             }
+            DataExchange::Modulo(exchange) => ModFlightScatter::try_create(
+                ctx.get_function_context()?,
+                &exchange.shuffle_key,
+                exchange.destination_ids.len(),
+            )?,
         }))
     }
 
diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs
@@ -303,7 +303,7 @@ impl DataExchangeManager {
                         None,
                         Some(config.query.to_rpc_client_tls_config()),
                     )
-                    .await?,
+                        .await?,
                 ))),
                 false => Ok(FlightClient::new(FlightServiceClient::new(
                     ConnectionFactory::create_rpc_channel(address.to_owned(), None, None).await?,
@@ -1011,6 +1011,19 @@ impl FragmentCoordinator {
                         .flight_scatter(&info.query_ctx, data_exchange)?,
                 }),
             )),
+            DataExchange::Modulo(exchange) => {
+                Ok(Some(ExchangeParams::ShuffleExchange(ShuffleExchangeParams {
+                    exchange_injector: exchange_injector.clone(),
+                    schema: self.physical_plan.output_schema()?,
+                    fragment_id: self.fragment_id,
+                    query_id: info.query_id.to_string(),
+                    executor_id: info.current_executor.to_string(),
+                    destination_ids: exchange.destination_ids.to_owned(),
+                    shuffle_scatter: exchange_injector
+                        .flight_scatter(&info.query_ctx, data_exchange)?,
+                })
+                ))
+            }
         }
     }
 
diff --git a/src/query/service/src/servers/flight/v1/exchange/mod.rs b/src/query/service/src/servers/flight/v1/exchange/mod.rs
@@ -32,6 +32,7 @@ pub mod serde;
 pub use data_exchange::BroadcastExchange;
 pub use data_exchange::DataExchange;
 pub use data_exchange::MergeExchange;
+pub use data_exchange::ModuloExchange;
 pub use data_exchange::ShuffleDataExchange;
 pub use exchange_injector::DefaultExchangeInjector;
 pub use exchange_injector::ExchangeInjector;
diff --git a/src/query/service/src/servers/flight/v1/scatter/flight_scatter_mod.rs b/src/query/service/src/servers/flight/v1/scatter/flight_scatter_mod.rs
@@ -0,0 +1,92 @@
+// Copyright 2021 Datafuse Labs
+//
+// 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 databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
+use databend_common_expression::type_check::check_function;
+use databend_common_expression::types::DataType;
+use databend_common_expression::types::NumberDataType;
+use databend_common_expression::types::NumberScalar;
+use databend_common_expression::DataBlock;
+use databend_common_expression::Evaluator;
+use databend_common_expression::Expr;
+use databend_common_expression::FunctionContext;
+use databend_common_expression::RemoteExpr;
+use databend_common_expression::Scalar;
+use databend_common_functions::BUILTIN_FUNCTIONS;
+
+use crate::servers::flight::v1::scatter::FlightScatter;
+
+#[derive(Clone)]
+pub struct ModFlightScatter {
+    scatter_size: usize,
+    func_ctx: FunctionContext,
+    expr: Expr,
+}
+
+impl ModFlightScatter {
+    pub fn try_create(
+        func_ctx: FunctionContext,
+        expr: &RemoteExpr,
+        scatter_size: usize,
+    ) -> Result<Box<dyn FlightScatter>> {
+        let expr = check_function(
+            None,
+            "modulo",
+            &[],
+            &[
+                expr.as_expr(&BUILTIN_FUNCTIONS),
+                Expr::constant(
+                    Scalar::Number(NumberScalar::UInt64(scatter_size as u64)),
+                    Some(DataType::Number(NumberDataType::UInt64)),
+                ),
+            ],
+            &BUILTIN_FUNCTIONS,
+        )?;
+        let return_type = expr.data_type();
+        if !matches!(return_type, DataType::Number(NumberDataType::UInt64)) {
+            return Err(ErrorCode::Internal(format!(
+                "ModFlightScatter expects modulo expression to return UInt64, but got {:?}",
+                return_type
+            )));
+        }
+
+        Ok(Box::new(ModFlightScatter {
+            scatter_size,
+            func_ctx,
+            expr,
+        }))
+    }
+}
+
+impl FlightScatter for ModFlightScatter {
+    fn execute(&self, data_block: DataBlock) -> Result<Vec<DataBlock>> {
+        let evaluator = Evaluator::new(&data_block, &self.func_ctx, &BUILTIN_FUNCTIONS);
+        let num = data_block.num_rows();
+
+        let column = evaluator
+            .run(&self.expr)?
+            .into_full_column(&DataType::Number(NumberDataType::UInt64), num);
+        let indices = column.as_number().unwrap().as_u_int64().unwrap();
+        let data_blocks = DataBlock::scatter(&data_block, indices, self.scatter_size)?;
+
+        let block_meta = data_block.get_meta();
+        let mut res = Vec::with_capacity(data_blocks.len());
+        for data_block in data_blocks {
+            res.push(data_block.add_meta(block_meta.cloned())?);
+        }
+
+        Ok(res)
+    }
+}
diff --git a/src/query/service/src/servers/flight/v1/scatter/mod.rs b/src/query/service/src/servers/flight/v1/scatter/mod.rs
@@ -15,7 +15,9 @@
 mod flight_scatter;
 mod flight_scatter_broadcast;
 mod flight_scatter_hash;
+mod flight_scatter_mod;
 
 pub use flight_scatter::FlightScatter;
 pub use flight_scatter_broadcast::BroadcastFlightScatter;
 pub use flight_scatter_hash::HashFlightScatter;
+pub use flight_scatter_mod::ModFlightScatter;
diff --git a/src/query/sql/src/executor/format.rs b/src/query/sql/src/executor/format.rs
@@ -1639,6 +1639,14 @@ fn exchange_to_format_tree(
             ),
             FragmentKind::Expansive => "Broadcast".to_string(),
             FragmentKind::Merge => "Merge".to_string(),
+            FragmentKind::Modulo => format!(
+                "Modulo({})",
+                plan.keys
+                    .iter()
+                    .map(|key| { key.as_expr(&BUILTIN_FUNCTIONS).sql_display() })
+                    .collect::<Vec<_>>()
+                    .join(", ")
+            ),
         })),
         to_format_tree(&plan.input, metadata, profs, context)?,
     ]))
diff --git a/src/query/sql/src/executor/physical_plans/common.rs b/src/query/sql/src/executor/physical_plans/common.rs
@@ -67,6 +67,8 @@ pub enum FragmentKind {
     // Broadcast
     Expansive,
     Merge,
+    // Partitioned by a specified expression % node_nums
+    Modulo,
 }
 
 #[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Copy)]
diff --git a/src/query/sql/src/executor/physical_plans/physical_exchange.rs b/src/query/sql/src/executor/physical_plans/physical_exchange.rs
@@ -81,6 +81,14 @@ impl PhysicalPlanBuilder {
                 allow_adjust_parallelism = false;
                 FragmentKind::Merge
             }
+            crate::plans::Exchange::Modulo(scalar) => {
+                let expr = scalar
+                    .type_check(input_schema.as_ref())?
+                    .project_column_ref(|index| input_schema.index_of(&index.to_string()).unwrap());
+                let (expr, _) = ConstantFolder::fold(&expr, &self.func_ctx, &BUILTIN_FUNCTIONS);
+                keys.push(expr.as_remote_expr());
+                FragmentKind::Modulo
+            }
         };
         Ok(PhysicalPlan::Exchange(Exchange {
             plan_id: 0,
diff --git a/src/query/sql/src/planner/format/display_rel_operator.rs b/src/query/sql/src/planner/format/display_rel_operator.rs
@@ -397,6 +397,7 @@ fn exchange_to_format_tree<I: IdHumanizer>(id_humanizer: &I, op: &Exchange) -> F
         Exchange::Broadcast => "Exchange(Broadcast)",
         Exchange::Merge => "Exchange(Merge)",
         Exchange::MergeSort => "Exchange(MergeSort)",
+        Exchange::Modulo(_) => "Exchange(Modulo)",
     };
 
     match op {
diff --git a/src/query/sql/src/planner/optimizer/ir/format.rs b/src/query/sql/src/planner/optimizer/ir/format.rs
@@ -66,6 +66,7 @@ fn display_rel_op(rel_op: &RelOperator) -> String {
                 Exchange::Broadcast => "Broadcast".to_string(),
                 Exchange::Merge => "Merge".to_string(),
                 Exchange::MergeSort => "MergeSort".to_string(),
+                Exchange::Modulo(scalar) => format!("Modulo({})", scalar.as_raw_expr()),
             })
         }
         RelOperator::DummyTableScan(_) => "DummyTableScan".to_string(),
diff --git a/src/query/sql/src/planner/optimizer/ir/property/enforcer.rs b/src/query/sql/src/planner/optimizer/ir/property/enforcer.rs
@@ -73,6 +73,7 @@ impl Enforcer for DistributionEnforcer {
             Distribution::Random | Distribution::Any => Err(ErrorCode::Internal(
                 "Cannot enforce random or any distribution",
             )),
+            Distribution::Modulo(key) => Ok(Exchange::Modulo(key.clone()).into()),
         }
     }
 }
diff --git a/src/query/sql/src/planner/optimizer/ir/property/property.rs b/src/query/sql/src/planner/optimizer/ir/property/property.rs
@@ -92,6 +92,7 @@ pub enum Distribution {
     Serial,
     Broadcast,
     Hash(Vec<ScalarExpr>),
+    Modulo(Box<ScalarExpr>),
 }
 
 impl Default for Distribution {
@@ -110,11 +111,15 @@ impl Distribution {
             | (Distribution::Random, _)
             | (Distribution::Serial, Distribution::Serial)
             | (Distribution::Broadcast, Distribution::Broadcast)
-            | (Distribution::Hash(_), Distribution::Broadcast) => true,
+            | (Distribution::Hash(_), Distribution::Broadcast)
+            | (Distribution::Modulo(_), Distribution::Broadcast) => true,
 
             (Distribution::Hash(ref keys), Distribution::Hash(ref other_keys)) => {
                 keys == other_keys
             }
+            (Distribution::Modulo(ref key), Distribution::Modulo(ref other_key)) => {
+                key == other_key
+            }
             _ => false,
         }
     }
@@ -135,6 +140,7 @@ impl Display for Distribution {
                     .collect::<Vec<_>>()
                     .join(", ")
             ),
+            Distribution::Modulo(ref key) => write!(f, "Modulo({})", key.as_raw_expr()),
         }
     }
 }
diff --git a/src/query/sql/src/planner/optimizer/optimizers/cascades/cost/model.rs b/src/query/sql/src/planner/optimizer/optimizers/cascades/cost/model.rs
@@ -158,7 +158,7 @@ impl DefaultCostModel {
         let exchange: Exchange = (*m_expr.plan.clone()).clone().try_into()?;
         let group = memo.group(m_expr.group_index)?;
         let cost = match exchange {
-            Exchange::Hash(_) => {
+            Exchange::Hash(_) | Exchange::Modulo(_) => {
                 group.stat_info.cardinality * self.network_per_row
                     + group.stat_info.cardinality * self.compute_per_row
             }
diff --git a/src/query/sql/src/planner/plans/exchange.rs b/src/query/sql/src/planner/plans/exchange.rs
diff --git a/src/query/storages/fuse/src/operations/append.rs b/src/query/storages/fuse/src/operations/append.rs
diff --git a/src/query/storages/fuse/src/operations/common/processors/transform_block_writer.rs b/src/query/storages/fuse/src/operations/common/processors/transform_block_writer.rs

Original file line number	Diff line number	Diff line change
`@@ -72,6 +72,7 @@ impl Display for QueryFragmentActionsWrap<'_> {`
`72`	`72`	`DataExchange::Merge(_) => writeln!(f, " DataExchange: Merge")?,`
`73`	`73`	`DataExchange::Broadcast(_) => writeln!(f, " DataExchange: Broadcast")?,`
`74`	`74`	`DataExchange::ShuffleDataExchange(_) => writeln!(f, " DataExchange: Shuffle")?,`
	`75`	`+ DataExchange::Modulo(_) => writeln!(f, " DataExchange: Modulo")?,`
`75`	`76`	`}`
`76`	`77`	`}`
`77`	`78`
Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,7 @@ pub enum DataExchange {`
`19`	`19`	`Merge(MergeExchange),`
`20`	`20`	`Broadcast(BroadcastExchange),`
`21`	`21`	`ShuffleDataExchange(ShuffleDataExchange),`
	`22`	`+ Modulo(ModuloExchange),`
`22`	`23`	`}`
`23`	`24`
`24`	`25`	`impl DataExchange {`
`@@ -27,6 +28,7 @@ impl DataExchange {`
`27`	`28`	`DataExchange::Merge(exchange) => vec![exchange.destination_id.clone()],`
`28`	`29`	`DataExchange::Broadcast(exchange) => exchange.destination_ids.clone(),`
`29`	`30`	`DataExchange::ShuffleDataExchange(exchange) => exchange.destination_ids.clone(),`
	`31`	`+ DataExchange::Modulo(exchange) => exchange.destination_ids.clone(),`
`30`	`32`	`}`
`31`	`33`	`}`
`32`	`34`	`}`
`@@ -77,3 +79,18 @@ impl BroadcastExchange {`
`77`	`79`	`DataExchange::Broadcast(BroadcastExchange { destination_ids })`
`78`	`80`	`}`
`79`	`81`	`}`
	`82`	`+`
	`83`	`+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]`
	`84`	`+pub struct ModuloExchange {`
	`85`	`+ pub destination_ids: Vec<String>,`
	`86`	`+ pub shuffle_key: RemoteExpr,`
	`87`	`+}`
	`88`	`+`
	`89`	`+impl ModuloExchange {`
	`90`	`+ pub fn create(destination_ids: Vec<String>, shuffle_key: RemoteExpr) -> DataExchange {`
	`91`	`+ DataExchange::Modulo(ModuloExchange {`
	`92`	`+ destination_ids,`
	`93`	`+ shuffle_key,`
	`94`	`+ })`
	`95`	`+ }`
	`96`	`+}`