refactor(query): refactor physical plan organization in pipeline builder for better readability (#17866)

refactor(query): reorganize physical plans in pipeline builder

Author: BohuTANG

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

@@ -35,82 +35,42 @@ use crate::pipelines::PipelineBuilder;
 use crate::sessions::QueryContext;
 
 impl PipelineBuilder {
-    pub(crate) fn build_range_join(&mut self, range_join: &RangeJoin) -> Result<()> {
-        let state = Arc::new(RangeJoinState::new(self.ctx.clone(), range_join));
-        self.expand_right_side_pipeline(range_join, state.clone())?;
-        self.build_left_side(range_join, state)?;
-        Ok(())
-    }
-
-    fn build_left_side(
-        &mut self,
-        range_join: &RangeJoin,
-        state: Arc<RangeJoinState>,
-    ) -> Result<()> {
-        self.build_pipeline(&range_join.left)?;
-        let max_threads = self.settings.get_max_threads()? as usize;
-        self.main_pipeline.try_resize(max_threads)?;
-        self.main_pipeline.add_transform(|input, output| {
-            Ok(ProcessorPtr::create(TransformRangeJoinLeft::create(
-                input,
-                output,
-                state.clone(),
-            )))
-        })?;
-        Ok(())
-    }
-
-    fn expand_right_side_pipeline(
-        &mut self,
-        range_join: &RangeJoin,
-        state: Arc<RangeJoinState>,
-    ) -> Result<()> {
-        let right_side_context = QueryContext::create_from(self.ctx.as_ref());
-        let mut right_side_builder = PipelineBuilder::create(
+    // Create a new pipeline builder with the same context as the current builder
+    fn create_sub_pipeline_builder(&self) -> PipelineBuilder {
+        let sub_context = QueryContext::create_from(self.ctx.as_ref());
+        let mut sub_builder = PipelineBuilder::create(
             self.func_ctx.clone(),
             self.settings.clone(),
-            right_side_context,
+            sub_context,
             self.main_pipeline.get_scopes(),
         );
-        right_side_builder.hash_join_states = self.hash_join_states.clone();
-
-        let mut right_res = right_side_builder.finalize(&range_join.right)?;
-        right_res.main_pipeline.add_sink(|input| {
-            Ok(ProcessorPtr::create(
-                Sinker::<TransformRangeJoinRight>::create(
-                    input,
-                    TransformRangeJoinRight::create(state.clone()),
-                ),
-            ))
-        })?;
-        self.pipelines.push(right_res.main_pipeline.finalize());
-        self.pipelines.extend(right_res.sources_pipelines);
-        Ok(())
+        sub_builder.hash_join_states = self.hash_join_states.clone();
+        sub_builder
     }
 
-    pub(crate) fn build_join(&mut self, join: &HashJoin) -> Result<()> {
-        // for merge into target table as build side.
-        let (enable_merge_into_optimization, merge_into_is_distributed) =
-            self.merge_into_get_optimization_flag(join);
+    pub(crate) fn build_hash_join(&mut self, join: &HashJoin) -> Result<()> {
+        // Get optimization flags for merge-into operations
+        let (enable_optimization, is_distributed) = self.merge_into_get_optimization_flag(join);
 
-        let state = self.build_join_state(
-            join,
-            merge_into_is_distributed,
-            enable_merge_into_optimization,
-        )?;
+        // Create the join state with optimization flags
+        let state = self.build_hash_join_state(join, is_distributed, enable_optimization)?;
         if let Some((build_cache_index, _)) = join.build_side_cache_info {
             self.hash_join_states
                 .insert(build_cache_index, state.clone());
         }
-        self.expand_build_side_pipeline(&join.build, join, state.clone())?;
-        self.build_join_probe(join, state)?;
 
-        // In the case of spilling, we need to share state among multiple threads. Quickly fetch all data from this round to quickly start the next round.
+        // Build both phases of the Hash Join
+        self.build_hash_join_build_side(&join.build, join, state.clone())?;
+        self.build_hash_join_probe_side(join, state)?;
+
+        // In the case of spilling, we need to share state among multiple threads
+        // Quickly fetch all data from this round to quickly start the next round
         self.main_pipeline
             .resize(self.main_pipeline.output_len(), true)
     }
 
-    fn build_join_state(
+    // Create the Hash Join state
+    fn build_hash_join_state(
         &mut self,
         join: &HashJoin,
         merge_into_is_distributed: bool,
@@ -128,20 +88,14 @@ impl PipelineBuilder {
         )
     }
 
-    fn expand_build_side_pipeline(
+    // Build the build-side pipeline for Hash Join
+    fn build_hash_join_build_side(
         &mut self,
         build: &PhysicalPlan,
         hash_join_plan: &HashJoin,
         join_state: Arc<HashJoinState>,
     ) -> Result<()> {
-        let build_side_context = QueryContext::create_from(self.ctx.as_ref());
-        let mut build_side_builder = PipelineBuilder::create(
-            self.func_ctx.clone(),
-            self.settings.clone(),
-            build_side_context,
-            self.main_pipeline.get_scopes(),
-        );
-        build_side_builder.hash_join_states = self.hash_join_states.clone();
+        let build_side_builder = self.create_sub_pipeline_builder();
         let mut build_res = build_side_builder.finalize(build)?;
 
         assert!(build_res.main_pipeline.is_pulling_pipeline()?);
@@ -162,7 +116,7 @@ impl PipelineBuilder {
                 build_state.clone(),
             )?))
         };
-        // for distributed merge into when source as build side.
+        // For distributed merge-into when source as build side
         if hash_join_plan.need_hold_hash_table {
             self.join_state = Some(build_state.clone())
         }
@@ -173,7 +127,12 @@ impl PipelineBuilder {
         Ok(())
     }
 
-    fn build_join_probe(&mut self, join: &HashJoin, state: Arc<HashJoinState>) -> Result<()> {
+    // Build the probe-side pipeline for Hash Join
+    fn build_hash_join_probe_side(
+        &mut self,
+        join: &HashJoin,
+        state: Arc<HashJoinState>,
+    ) -> Result<()> {
         self.build_pipeline(&join.probe)?;
 
         let max_block_size = self.settings.get_max_block_size()? as usize;
@@ -203,16 +162,66 @@ impl PipelineBuilder {
             )?))
         })?;
 
+        // For merge-into operations that need to hold the hash table
         if join.need_hold_hash_table {
-            let mut projected_probe_fields = vec![];
+            // Extract projected fields from probe schema
+            let mut projected_fields = vec![];
             for (i, field) in probe_state.probe_schema.fields().iter().enumerate() {
                 if probe_state.probe_projections.contains(&i) {
-                    projected_probe_fields.push(field.clone());
+                    projected_fields.push(field.clone());
                 }
             }
-            self.merge_into_probe_data_fields = Some(projected_probe_fields);
+            self.merge_into_probe_data_fields = Some(projected_fields);
         }
 
         Ok(())
     }
+
+    pub(crate) fn build_range_join(&mut self, range_join: &RangeJoin) -> Result<()> {
+        let state = Arc::new(RangeJoinState::new(self.ctx.clone(), range_join));
+        self.build_range_join_right_side(range_join, state.clone())?;
+        self.build_range_join_left_side(range_join, state)?;
+        Ok(())
+    }
+
+    // Build the left-side pipeline for Range Join
+    fn build_range_join_left_side(
+        &mut self,
+        range_join: &RangeJoin,
+        state: Arc<RangeJoinState>,
+    ) -> Result<()> {
+        self.build_pipeline(&range_join.left)?;
+        let max_threads = self.settings.get_max_threads()? as usize;
+        self.main_pipeline.try_resize(max_threads)?;
+        self.main_pipeline.add_transform(|input, output| {
+            Ok(ProcessorPtr::create(TransformRangeJoinLeft::create(
+                input,
+                output,
+                state.clone(),
+            )))
+        })?;
+        Ok(())
+    }
+
+    // Build the right-side pipeline for Range Join
+    fn build_range_join_right_side(
+        &mut self,
+        range_join: &RangeJoin,
+        state: Arc<RangeJoinState>,
+    ) -> Result<()> {
+        let right_side_builder = self.create_sub_pipeline_builder();
+
+        let mut right_res = right_side_builder.finalize(&range_join.right)?;
+        right_res.main_pipeline.add_sink(|input| {
+            Ok(ProcessorPtr::create(
+                Sinker::<TransformRangeJoinRight>::create(
+                    input,
+                    TransformRangeJoinRight::create(state.clone()),
+                ),
+            ))
+        })?;
+        self.pipelines.push(right_res.main_pipeline.finalize());
+        self.pipelines.extend(right_res.sources_pipelines);
+        Ok(())
+    }
 }

src/query/service/src/pipelines/pipeline_builder.rs

@@ -163,50 +163,79 @@ impl PipelineBuilder {
         self.is_exchange_neighbor |= self.is_exchange_neighbor(plan);
 
         match plan {
+            // ==============================
+            // 1. Data Source Plans
+            // ==============================
+            // Basic table scans - retrieve data from tables
             PhysicalPlan::TableScan(scan) => self.build_table_scan(scan),
             PhysicalPlan::ConstantTableScan(scan) => self.build_constant_table_scan(scan),
+            PhysicalPlan::CacheScan(cache_scan) => self.build_cache_scan(cache_scan),
+            PhysicalPlan::ExpressionScan(expression_scan) => {
+                self.build_expression_scan(expression_scan)
+            }
+            PhysicalPlan::RecursiveCteScan(scan) => self.build_recursive_cte_scan(scan),
+
+            // Special source operations
+            PhysicalPlan::MutationSource(mutation_source) => {
+                self.build_mutation_source(mutation_source)
+            }
+
+            // ==============================
+            // 2. Relational Operators
+            // ==============================
+            // Filtering and projection
             PhysicalPlan::Filter(filter) => self.build_filter(filter),
             PhysicalPlan::EvalScalar(eval_scalar) => self.build_eval_scalar(eval_scalar),
+            PhysicalPlan::ProjectSet(project_set) => self.build_project_set(project_set),
+
+            // Sorting and limiting
+            PhysicalPlan::Sort(sort) => self.build_sort(sort),
+            PhysicalPlan::Limit(limit) => self.build_limit(limit),
+            PhysicalPlan::RowFetch(row_fetch) => self.build_row_fetch(row_fetch),
+
+            // Join operations
+            PhysicalPlan::HashJoin(join) => self.build_hash_join(join),
+            PhysicalPlan::RangeJoin(range_join) => self.build_range_join(range_join),
+
+            // Aggregation operations
             PhysicalPlan::AggregateExpand(aggregate) => self.build_aggregate_expand(aggregate),
             PhysicalPlan::AggregatePartial(aggregate) => self.build_aggregate_partial(aggregate),
             PhysicalPlan::AggregateFinal(aggregate) => self.build_aggregate_final(aggregate),
+
+            // Window functions
             PhysicalPlan::Window(window) => self.build_window(window),
             PhysicalPlan::WindowPartition(window_partition) => {
                 self.build_window_partition(window_partition)
             }
-            PhysicalPlan::Sort(sort) => self.build_sort(sort),
-            PhysicalPlan::Limit(limit) => self.build_limit(limit),
-            PhysicalPlan::RowFetch(row_fetch) => self.build_row_fetch(row_fetch),
-            PhysicalPlan::HashJoin(join) => self.build_join(join),
+
+            PhysicalPlan::UnionAll(union_all) => self.build_union_all(union_all),
+
+            // ==============================
+            // 3. Data Distribution
+            // ==============================
             PhysicalPlan::ExchangeSink(sink) => self.build_exchange_sink(sink),
             PhysicalPlan::ExchangeSource(source) => self.build_exchange_source(source),
-            PhysicalPlan::UnionAll(union_all) => self.build_union_all(union_all),
             PhysicalPlan::DistributedInsertSelect(insert_select) => {
                 self.build_distributed_insert_select(insert_select)
             }
-            PhysicalPlan::ProjectSet(project_set) => self.build_project_set(project_set),
-            PhysicalPlan::Udf(udf) => self.build_udf(udf),
-            PhysicalPlan::Exchange(_) => Err(ErrorCode::Internal(
-                "Invalid physical plan with PhysicalPlan::Exchange",
-            )),
-            PhysicalPlan::RangeJoin(range_join) => self.build_range_join(range_join),
-            PhysicalPlan::CacheScan(cache_scan) => self.build_cache_scan(cache_scan),
-            PhysicalPlan::ExpressionScan(expression_scan) => {
-                self.build_expression_scan(expression_scan)
-            }
+            PhysicalPlan::Shuffle(shuffle) => self.build_shuffle(shuffle),
+            PhysicalPlan::Duplicate(duplicate) => self.build_duplicate(duplicate),
 
-            // Copy into.
+            // ==============================
+            // 4. Data Modification Operations
+            // ==============================
+            // Copy operations
             PhysicalPlan::CopyIntoTable(copy) => self.build_copy_into_table(copy),
             PhysicalPlan::CopyIntoLocation(copy) => self.build_copy_into_location(copy),
 
-            // Replace.
+            // Replace operations
             PhysicalPlan::ReplaceAsyncSourcer(async_sourcer) => {
                 self.build_async_sourcer(async_sourcer)
             }
             PhysicalPlan::ReplaceDeduplicate(deduplicate) => self.build_deduplicate(deduplicate),
             PhysicalPlan::ReplaceInto(replace) => self.build_replace_into(replace),
 
-            // Mutation.
+            // Mutation operations (DELETE/UPDATE)
             PhysicalPlan::Mutation(mutation) => self.build_mutation(mutation),
             PhysicalPlan::MutationSplit(mutation_split) => {
                 self.build_mutation_split(mutation_split)
@@ -220,19 +249,14 @@ impl PipelineBuilder {
             PhysicalPlan::AddStreamColumn(add_stream_column) => {
                 self.build_add_stream_column(add_stream_column)
             }
+            PhysicalPlan::ColumnMutation(column_mutation) => {
+                self.build_column_mutation(column_mutation)
+            }
 
-            // Commit.
+            // Commit operations
             PhysicalPlan::CommitSink(plan) => self.build_commit_sink(plan),
 
-            // Compact.
-            PhysicalPlan::CompactSource(compact) => self.build_compact_source(compact),
-
-            // Recluster.
-            PhysicalPlan::Recluster(recluster) => self.build_recluster(recluster),
-            PhysicalPlan::HilbertPartition(partition) => self.build_hilbert_partition(partition),
-
-            PhysicalPlan::Duplicate(duplicate) => self.build_duplicate(duplicate),
-            PhysicalPlan::Shuffle(shuffle) => self.build_shuffle(shuffle),
+            // MERGE INTO chunk processing operations
             PhysicalPlan::ChunkFilter(chunk_filter) => self.build_chunk_filter(chunk_filter),
             PhysicalPlan::ChunkEvalScalar(chunk_project) => {
                 self.build_chunk_eval_scalar(chunk_project)
@@ -250,14 +274,27 @@ impl PipelineBuilder {
             PhysicalPlan::ChunkCommitInsert(chunk_commit_insert) => {
                 self.build_chunk_commit_insert(chunk_commit_insert)
             }
+
+            // ==============================
+            // 5. Data Maintenance Operations
+            // ==============================
+            PhysicalPlan::CompactSource(compact) => self.build_compact_source(compact),
+            PhysicalPlan::Recluster(recluster) => self.build_recluster(recluster),
+            PhysicalPlan::HilbertPartition(partition) => self.build_hilbert_partition(partition),
+
+            // ==============================
+            // 6. Special Processing Operations
+            // ==============================
+            // User-defined functions and async operations
+            PhysicalPlan::Udf(udf) => self.build_udf(udf),
             PhysicalPlan::AsyncFunction(async_func) => self.build_async_function(async_func),
-            PhysicalPlan::RecursiveCteScan(scan) => self.build_recursive_cte_scan(scan),
-            PhysicalPlan::MutationSource(mutation_source) => {
-                self.build_mutation_source(mutation_source)
-            }
-            PhysicalPlan::ColumnMutation(column_mutation) => {
-                self.build_column_mutation(column_mutation)
-            }
+
+            // ==============================
+            // 7. Invalid Plans
+            // ==============================
+            PhysicalPlan::Exchange(_) => Err(ErrorCode::Internal(
+                "Invalid physical plan with PhysicalPlan::Exchange",
+            )),
         }?;
 
         self.is_exchange_neighbor = is_exchange_neighbor;