try to improve

Author: zhyass

scripts/ci/deploy/databend-query-cluster-3-nodes.sh

@@ -6,7 +6,7 @@ set -e
 
 SCRIPT_PATH="$(cd "$(dirname "$0")" >/dev/null 2>&1 && pwd)"
 cd "$SCRIPT_PATH/../../.." || exit
-BUILD_PROFILE=${BUILD_PROFILE:-debug}
+BUILD_PROFILE=${BUILD_PROFILE:-release}
 
 # Caveat: has to kill query first.
 # `query` tries to remove its liveness record from meta before shutting down.

src/query/service/src/interpreters/interpreter_table_recluster.rs

@@ -42,14 +42,12 @@ use databend_common_sql::executor::physical_plans::Recluster;
 use databend_common_sql::executor::PhysicalPlan;
 use databend_common_sql::executor::PhysicalPlanBuilder;
 use databend_common_sql::plans::set_update_stream_columns;
-use databend_common_sql::plans::BoundColumnRef;
 use databend_common_sql::plans::Plan;
 use databend_common_sql::plans::ReclusterPlan;
 use databend_common_sql::query_executor::QueryExecutor;
 use databend_common_sql::IdentifierNormalizer;
 use databend_common_sql::NameResolutionContext;
 use databend_common_sql::Planner;
-use databend_common_sql::ScalarExpr;
 use databend_common_sql::TypeChecker;
 use databend_enterprise_hilbert_clustering::get_hilbert_clustering_handler;
 use databend_storages_common_table_meta::table::ClusterType;
@@ -61,7 +59,6 @@ use log::warn;
 use crate::interpreters::hook::vacuum_hook::hook_clear_m_cte_temp_table;
 use crate::interpreters::hook::vacuum_hook::hook_disk_temp_dir;
 use crate::interpreters::hook::vacuum_hook::hook_vacuum_temp_files;
-use crate::interpreters::interpreter_insert_multi_table::scalar_expr_to_remote_expr;
 use crate::interpreters::Interpreter;
 use crate::interpreters::InterpreterClusteringHistory;
 use crate::pipelines::executor::ExecutorSettings;
@@ -267,11 +264,11 @@ impl ReclusterTableInterpreter {
                         acc
                     },
                 );
-                
+
                 let query_str = self.ctx.get_query_str();
                 let write_progress = self.ctx.get_write_progress();
                 let write_progress_value = write_progress.as_ref().get_values();
-                
+
                 let subquery_executor = Arc::new(ServiceQueryExecutor::new(
                     QueryContext::create_from(self.ctx.as_ref()),
                 ));
@@ -353,8 +350,7 @@ impl ReclusterTableInterpreter {
                     .unwrap_or_default();
                 let (stmt, _) = parse_sql(&tokens, sql_dialect)?;
 
-                let mut planner =
-                    Planner::new_with_query_executor(self.ctx.clone(), subquery_executor);
+                let mut planner = Planner::new(self.ctx.clone());
                 let plan = planner.plan_stmt(&stmt, false).await?;
                 let Plan::Query {
                     mut s_expr,
@@ -369,8 +365,6 @@ impl ReclusterTableInterpreter {
                     *s_expr = set_update_stream_columns(&s_expr)?;
                 }
 
-                write_progress.set(&write_progress_value);
-                self.ctx.attach_query_str(QueryKind::Other, query_str);
                 let mut builder = PhysicalPlanBuilder::new(metadata, self.ctx.clone(), false);
                 let mut plan = Box::new(builder.build(&s_expr, bind_context.column_set()).await?);
                 let mut is_exchange = false;
@@ -384,33 +378,22 @@ impl ReclusterTableInterpreter {
                     plan = input.clone();
                 }
 
+                let mut nodes_num = 1;
                 let cluster = self.ctx.get_cluster();
-                let is_distributed = is_exchange || !cluster.is_empty();
-                if is_distributed {
-                    let expr = scalar_expr_to_remote_expr(
-                        &ScalarExpr::BoundColumnRef(BoundColumnRef {
-                            span: None,
-                            column: bind_context.columns.last().unwrap().clone(),
-                        }),
-                        plan.output_schema()?.as_ref(),
-                    )?;
-                    plan = Box::new(PhysicalPlan::Exchange(Exchange {
-                        plan_id: 0,
-                        input: plan,
-                        kind: FragmentKind::Normal,
-                        keys: vec![expr],
-                        allow_adjust_parallelism: true,
-                        ignore_exchange: false,
-                    }));
+                if !cluster.is_empty() {
+                    nodes_num = cluster.nodes.len();
                 }
-
+                let is_distributed = is_exchange || !cluster.is_empty();
                 let mut plan = PhysicalPlan::HilbertPartition(Box::new(HilbertPartition {
                     plan_id: 0,
                     input: plan,
                     table_info: table_info.clone(),
                     range_id: bind_context.columns.last().unwrap().index,
-                    num_partitions: total_partitions,
+                    num_partitions: total_partitions.div_ceil(nodes_num),
                     rows_per_block,
+                    node_id: None,
+                    projections: bind_context.column_set(),
+                    nodes_num,
                 }));
 
                 if is_distributed {
@@ -424,6 +407,8 @@ impl ReclusterTableInterpreter {
                     });
                 }
 
+                write_progress.set(&write_progress_value);
+                self.ctx.attach_query_str(QueryKind::Other, query_str);
                 PhysicalPlan::CommitSink(Box::new(CommitSink {
                     input: Box::new(plan),
                     table_info,

src/query/service/src/pipelines/builders/builder_hilbert_partition.rs

@@ -17,9 +17,20 @@ use std::sync::atomic::AtomicUsize;
 
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
+use databend_common_expression::types::DataType;
+use databend_common_expression::types::NumberDataType;
+use databend_common_expression::types::NumberScalar;
+use databend_common_expression::Scalar;
+use databend_common_expression::PREDICATE_COLUMN_NAME;
 use databend_common_pipeline_core::processors::ProcessorPtr;
 use databend_common_sql::executor::physical_plans::HilbertPartition;
 use databend_common_sql::executor::physical_plans::MutationKind;
+use databend_common_sql::plans::BoundColumnRef;
+use databend_common_sql::plans::ConstantExpr;
+use databend_common_sql::plans::FunctionCall;
+use databend_common_sql::ColumnBindingBuilder;
+use databend_common_sql::ScalarExpr;
+use databend_common_sql::Visibility;
 use databend_common_storages_fuse::operations::TransformSerializeBlock;
 use databend_common_storages_fuse::statistics::ClusterStatsGenerator;
 use databend_common_storages_fuse::FuseTable;
@@ -36,6 +47,50 @@ use crate::spillers::SpillerDiskConfig;
 impl PipelineBuilder {
     pub(crate) fn build_hilbert_partition(&mut self, partition: &HilbertPartition) -> Result<()> {
         self.build_pipeline(&partition.input)?;
+        if let Some(node_id) = partition.node_id {
+            let input_schema = partition.input.output_schema()?;
+            let column = ColumnBindingBuilder::new(
+                PREDICATE_COLUMN_NAME.to_string(),
+                input_schema.num_fields() - 1,
+                Box::new(DataType::Number(NumberDataType::UInt64)),
+                Visibility::Visible,
+            )
+            .build();
+            let expr = ScalarExpr::FunctionCall(FunctionCall {
+                span: None,
+                func_name: "eq".to_string(),
+                params: vec![],
+                arguments: vec![
+                    ScalarExpr::FunctionCall(FunctionCall {
+                        span: None,
+                        func_name: "mod".to_string(),
+                        params: vec![],
+                        arguments: vec![
+                            ScalarExpr::BoundColumnRef(BoundColumnRef { span: None, column }),
+                            ScalarExpr::ConstantExpr(ConstantExpr {
+                                span: None,
+                                value: Scalar::Number(NumberScalar::UInt8(
+                                    partition.nodes_num as u8,
+                                )),
+                            }),
+                        ],
+                    }),
+                    ScalarExpr::ConstantExpr(ConstantExpr {
+                        span: None,
+                        value: Scalar::Number(NumberScalar::UInt8(node_id)),
+                    }),
+                ],
+            });
+
+            let remote_expr = expr
+                .as_expr()?
+                .project_column_ref(|col| col.index)
+                .as_remote_expr();
+            self.main_pipeline.add_transform(
+                self.filter_transform_builder(&[remote_expr], partition.projections.clone())?,
+            )?;
+        }
+
         let num_processors = self.main_pipeline.output_len();
         let table = self
             .ctx
@@ -44,7 +99,7 @@ impl PipelineBuilder {
 
         self.main_pipeline.exchange(
             num_processors,
-            HilbertPartitionExchange::create(partition.num_partitions),
+            HilbertPartitionExchange::create(partition.num_partitions, partition.nodes_num),
         );
 
         let settings = self.ctx.get_settings();

src/query/service/src/pipelines/processors/transforms/transform_partition_collect.rs

@@ -41,11 +41,15 @@ use crate::spillers::SpillerType;
 
 pub struct HilbertPartitionExchange {
     num_partitions: usize,
+    nod_nums: usize,
 }
 
 impl HilbertPartitionExchange {
-    pub fn create(num_partitions: usize) -> Arc<HilbertPartitionExchange> {
-        Arc::new(HilbertPartitionExchange { num_partitions })
+    pub fn create(num_partitions: usize, nod_nums: usize) -> Arc<HilbertPartitionExchange> {
+        Arc::new(HilbertPartitionExchange {
+            num_partitions,
+            nod_nums,
+        })
     }
 }
 
@@ -64,15 +68,20 @@ impl Exchange for HilbertPartitionExchange {
         // Scatter the data block to different partitions.
         let indices = range_ids
             .iter()
-            .map(|&id| (id % self.num_partitions as u64) as u16)
+            .map(|&id| ((id / self.nod_nums as u64) % self.num_partitions as u64) as u16)
             .collect::<Vec<_>>();
         data_block.pop_columns(1);
-        let scatter_blocks = DataBlock::scatter(&data_block, &indices, self.num_partitions)?;
-
+        let scatter_indices =
+            DataBlock::divide_indices_by_scatter_size(&indices, self.num_partitions);
         // Partition the data blocks to different processors.
         let mut output_data_blocks = vec![vec![]; n];
-        for (partition_id, data_block) in scatter_blocks.into_iter().enumerate() {
-            output_data_blocks[partition_id % n].push((partition_id, data_block));
+        for (partition_id, indices) in scatter_indices.iter().take(self.num_partitions).enumerate()
+        {
+            if indices.is_empty() {
+                continue;
+            }
+            let block = data_block.take_with_optimize_size(indices)?;
+            output_data_blocks[partition_id % n].push((partition_id, block));
         }
 
         // Union data blocks for each processor.

src/query/service/src/schedulers/fragments/fragmenter.rs

@@ -25,6 +25,7 @@ use databend_common_sql::executor::physical_plans::ExchangeSink;
 use databend_common_sql::executor::physical_plans::ExchangeSource;
 use databend_common_sql::executor::physical_plans::FragmentKind;
 use databend_common_sql::executor::physical_plans::HashJoin;
+use databend_common_sql::executor::physical_plans::HilbertPartition;
 use databend_common_sql::executor::physical_plans::MutationSource;
 use databend_common_sql::executor::physical_plans::Recluster;
 use databend_common_sql::executor::physical_plans::ReplaceInto;
@@ -65,6 +66,7 @@ enum State {
     ReplaceInto,
     Compact,
     Recluster,
+    Hilbert,
     Other,
 }
 
@@ -201,6 +203,11 @@ impl PhysicalPlanReplacer for Fragmenter {
         Ok(PhysicalPlan::Recluster(Box::new(plan.clone())))
     }
 
+    fn replace_hilbert_recluster(&mut self, plan: &HilbertPartition) -> Result<PhysicalPlan> {
+        self.state = State::Hilbert;
+        Ok(PhysicalPlan::HilbertPartition(Box::new(plan.clone())))
+    }
+
     fn replace_compact_source(&mut self, plan: &CompactSource) -> Result<PhysicalPlan> {
         self.state = State::Compact;
         Ok(PhysicalPlan::CompactSource(Box::new(plan.clone())))
@@ -301,6 +308,7 @@ impl PhysicalPlanReplacer for Fragmenter {
             State::Other => FragmentType::Intermediate,
             State::ReplaceInto => FragmentType::ReplaceInto,
             State::Compact => FragmentType::Compact,
+            State::Hilbert => FragmentType::Hilbert,
             State::Recluster => FragmentType::Recluster,
         };
         self.state = State::Other;

src/query/service/src/schedulers/fragments/plan_fragment.rs

@@ -29,6 +29,7 @@ use databend_common_sql::executor::physical_plans::CompactSource;
 use databend_common_sql::executor::physical_plans::ConstantTableScan;
 use databend_common_sql::executor::physical_plans::CopyIntoTable;
 use databend_common_sql::executor::physical_plans::CopyIntoTableSource;
+use databend_common_sql::executor::physical_plans::HilbertPartition;
 use databend_common_sql::executor::physical_plans::MutationSource;
 use databend_common_sql::executor::physical_plans::Recluster;
 use databend_common_sql::executor::physical_plans::ReplaceDeduplicate;
@@ -64,6 +65,7 @@ pub enum FragmentType {
     ReplaceInto,
     Compact,
     Recluster,
+    Hilbert,
     MutationSource,
 }
 
@@ -137,6 +139,9 @@ impl PlanFragment {
             FragmentType::Recluster => {
                 self.redistribute_recluster(ctx, &mut fragment_actions)?;
             }
+            FragmentType::Hilbert => {
+                self.redistribute_hilbert(ctx, &mut fragment_actions)?;
+            }
         }
 
         if let Some(ref exchange) = self.exchange {
@@ -346,6 +351,24 @@ impl PlanFragment {
         Ok(())
     }
 
+    fn redistribute_hilbert(
+        &self,
+        ctx: Arc<QueryContext>,
+        fragment_actions: &mut QueryFragmentActions,
+    ) -> Result<()> {
+        let executors = Fragmenter::get_executors(ctx.clone());
+        for (executor_idx, executor) in executors.into_iter().enumerate() {
+            let mut plan = self.plan.clone();
+            let mut replace_hilbert = ReplaceHilbertRecluster {
+                node: executor_idx as u8,
+            };
+            plan = replace_hilbert.replace(&plan)?;
+
+            fragment_actions.add_action(QueryFragmentAction::create(executor, plan));
+        }
+        Ok(())
+    }
+
     fn redistribute_recluster(
         &self,
         ctx: Arc<QueryContext>,
@@ -569,6 +592,19 @@ impl PhysicalPlanReplacer for ReplaceRecluster {
     }
 }
 
+struct ReplaceHilbertRecluster {
+    pub node: u8,
+}
+
+impl PhysicalPlanReplacer for ReplaceHilbertRecluster {
+    fn replace_hilbert_recluster(&mut self, plan: &HilbertPartition) -> Result<PhysicalPlan> {
+        Ok(PhysicalPlan::HilbertPartition(Box::new(HilbertPartition {
+            node_id: Some(self.node),
+            ..plan.clone()
+        })))
+    }
+}
+
 struct ReplaceMutationSource {
     pub partitions: Partitions,
 }

src/query/sql/src/executor/physical_plan_visitor.rs

@@ -108,7 +108,7 @@ pub trait PhysicalPlanReplacer {
             PhysicalPlan::ExpressionScan(plan) => self.replace_expression_scan(plan),
             PhysicalPlan::CacheScan(plan) => self.replace_cache_scan(plan),
             PhysicalPlan::Recluster(plan) => self.replace_recluster(plan),
-            PhysicalPlan::HilbertPartition(plan) => self.replace_hilbert_serialize(plan),
+            PhysicalPlan::HilbertPartition(plan) => self.replace_hilbert_recluster(plan),
             PhysicalPlan::Udf(plan) => self.replace_udf(plan),
             PhysicalPlan::AsyncFunction(plan) => self.replace_async_function(plan),
             PhysicalPlan::Duplicate(plan) => self.replace_duplicate(plan),
@@ -127,7 +127,7 @@ pub trait PhysicalPlanReplacer {
         Ok(PhysicalPlan::Recluster(Box::new(plan.clone())))
     }
 
-    fn replace_hilbert_serialize(&mut self, plan: &HilbertPartition) -> Result<PhysicalPlan> {
+    fn replace_hilbert_recluster(&mut self, plan: &HilbertPartition) -> Result<PhysicalPlan> {
         let input = self.replace(&plan.input)?;
         Ok(PhysicalPlan::HilbertPartition(Box::new(HilbertPartition {
             plan_id: plan.plan_id,
@@ -136,6 +136,9 @@ pub trait PhysicalPlanReplacer {
             range_id: plan.range_id,
             num_partitions: plan.num_partitions,
             rows_per_block: plan.rows_per_block,
+            node_id: plan.node_id,
+            projections: plan.projections.clone(),
+            nodes_num: plan.nodes_num,
         })))
     }
 

src/query/sql/src/executor/physical_plans/physical_recluster.rs

@@ -16,6 +16,7 @@ use databend_common_catalog::plan::ReclusterTask;
 use databend_common_meta_app::schema::TableInfo;
 
 use crate::executor::PhysicalPlan;
+use crate::ColumnSet;
 use crate::IndexType;
 
 #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
@@ -33,4 +34,8 @@ pub struct HilbertPartition {
     pub range_id: IndexType,
     pub num_partitions: usize,
     pub rows_per_block: usize,
+
+    pub node_id: Option<u8>,
+    pub projections: ColumnSet,
+    pub nodes_num: usize,
 }