refactor: flatten union plan

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

@@ -15,6 +15,7 @@
 use async_channel::Receiver;
 use databend_common_exception::Result;
 use databend_common_expression::DataBlock;
+use databend_common_expression::DataSchemaRef;
 use databend_common_pipeline_core::processors::ProcessorPtr;
 use databend_common_pipeline_sinks::UnionReceiveSink;
 use databend_common_sql::executor::physical_plans::UnionAll;
@@ -26,17 +27,29 @@ use crate::sessions::QueryContext;
 
 impl PipelineBuilder {
     pub fn build_union_all(&mut self, union_all: &UnionAll) -> Result<()> {
-        self.build_pipeline(&union_all.left)?;
-        let union_all_receiver = self.expand_union_all(&union_all.right)?;
+        self.build_pipeline(union_all.children.last().unwrap())?;
+        let mut remain_children_receivers = vec![];
+        for (idx, remaining_child) in union_all
+            .children
+            .iter()
+            .take(union_all.children.len() - 1)
+            .enumerate()
+        {
+            remain_children_receivers.push((idx, self.expand_union_all(remaining_child)?));
+        }
+        let schemas: Vec<DataSchemaRef> = union_all
+            .children
+            .iter()
+            .map(|plan| plan.output_schema())
+            .collect::<Result<_>>()?;
         self.main_pipeline
             .add_transform(|transform_input_port, transform_output_port| {
                 Ok(ProcessorPtr::create(TransformMergeBlock::try_create(
                     transform_input_port,
                     transform_output_port,
-                    union_all.left.output_schema()?,
-                    union_all.right.output_schema()?,
-                    union_all.pairs.clone(),
-                    union_all_receiver.clone(),
+                    schemas.clone(),
+                    union_all.output_cols.clone(),
+                    remain_children_receivers.clone(),
                 )?))
             })?;
         Ok(())

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

@@ -27,6 +27,7 @@ use databend_common_pipeline_core::processors::Event;
 use databend_common_pipeline_core::processors::InputPort;
 use databend_common_pipeline_core::processors::OutputPort;
 use databend_common_pipeline_core::processors::Processor;
+use databend_common_sql::IndexType;
 
 pub struct TransformMergeBlock {
     finished: bool,
@@ -35,90 +36,96 @@ pub struct TransformMergeBlock {
 
     input_data: Option<DataBlock>,
     output_data: Option<DataBlock>,
-    left_schema: DataSchemaRef,
-    right_schema: DataSchemaRef,
-    pairs: Vec<(String, String)>,
+    schemas: Vec<DataSchemaRef>,
+    output_cols: Vec<Vec<IndexType>>,
 
-    receiver: Receiver<DataBlock>,
-    receiver_result: Option<DataBlock>,
+    receivers: Vec<(usize, Receiver<DataBlock>)>,
+    receiver_results: Vec<(usize, DataBlock)>,
 }
 
 impl TransformMergeBlock {
     pub fn try_create(
         input: Arc<InputPort>,
         output: Arc<OutputPort>,
-        left_schema: DataSchemaRef,
-        right_schema: DataSchemaRef,
-        pairs: Vec<(String, String)>,
-        receiver: Receiver<DataBlock>,
+        schemas: Vec<DataSchemaRef>,
+        output_cols: Vec<Vec<IndexType>>,
+        receivers: Vec<(usize, Receiver<DataBlock>)>,
     ) -> Result<Box<dyn Processor>> {
         Ok(Box::new(TransformMergeBlock {
             finished: false,
             input,
             output,
             input_data: None,
             output_data: None,
-            left_schema,
-            right_schema,
-            pairs,
-            receiver,
-            receiver_result: None,
+            schemas,
+            output_cols,
+            receivers,
+            receiver_results: vec![],
         }))
     }
 
-    fn project_block(&self, block: DataBlock, is_left: bool) -> Result<DataBlock> {
+    fn project_block(&self, block: DataBlock, idx: Option<usize>) -> Result<DataBlock> {
         let num_rows = block.num_rows();
-        let columns = self
-            .pairs
-            .iter()
-            .map(|(left, right)| {
-                if is_left {
+        let columns = if let Some(idx) = idx {
+            self.check_type(idx, &block)?
+        } else {
+            self.output_cols
+                .last()
+                .unwrap()
+                .iter()
+                .map(|idx| {
                     Ok(block
-                        .get_by_offset(self.left_schema.index_of(left)?)
+                        .get_by_offset(self.schemas.last().unwrap().index_of(&idx.to_string())?)
                         .clone())
-                } else {
-                    // If block from right, check if block schema matches self scheme(left schema)
-                    // If unmatched, covert block columns types or report error
-                    self.check_type(left, right, &block)
-                }
-            })
-            .collect::<Result<Vec<_>>>()?;
-        Ok(DataBlock::new(columns, num_rows))
+                })
+                .collect::<Result<Vec<_>>>()?
+        };
+        let block = DataBlock::new(columns, num_rows);
+        Ok(block)
     }
 
-    fn check_type(
-        &self,
-        left_name: &str,
-        right_name: &str,
-        block: &DataBlock,
-    ) -> Result<BlockEntry> {
-        let left_field = self.left_schema.field_with_name(left_name)?;
-        let left_data_type = left_field.data_type();
-
-        let right_field = self.right_schema.field_with_name(right_name)?;
-        let right_data_type = right_field.data_type();
-
-        let index = self.right_schema.index_of(right_name)?;
-
-        if left_data_type == right_data_type {
-            return Ok(block.get_by_offset(index).clone());
-        }
+    fn check_type(&self, idx: usize, block: &DataBlock) -> Result<Vec<BlockEntry>> {
+        let mut columns = vec![];
+        for (left_idx, right_idx) in self
+            .output_cols
+            .last()
+            .unwrap()
+            .iter()
+            .zip(self.output_cols[idx].iter())
+        {
+            let left_field = self
+                .schemas
+                .last()
+                .unwrap()
+                .field_with_name(&left_idx.to_string())?;
+            let left_data_type = left_field.data_type();
+
+            let right_field = self.schemas[idx].field_with_name(&right_idx.to_string())?;
+            let right_data_type = right_field.data_type();
+
+            let offset = self.schemas[idx].index_of(&right_idx.to_string())?;
+            if left_data_type == right_data_type {
+                columns.push(block.get_by_offset(offset).clone());
+                continue;
+            }
 
-        if left_data_type.remove_nullable() == right_data_type.remove_nullable() {
-            let origin_column = block.get_by_offset(index).clone();
-            let mut builder = ColumnBuilder::with_capacity(left_data_type, block.num_rows());
-            let value = origin_column.value.as_ref();
-            for idx in 0..block.num_rows() {
-                let scalar = value.index(idx).unwrap();
-                builder.push(scalar);
+            if left_data_type.remove_nullable() == right_data_type.remove_nullable() {
+                let origin_column = block.get_by_offset(offset).clone();
+                let mut builder = ColumnBuilder::with_capacity(left_data_type, block.num_rows());
+                let value = origin_column.value.as_ref();
+                for idx in 0..block.num_rows() {
+                    let scalar = value.index(idx).unwrap();
+                    builder.push(scalar);
+                }
+                let col = builder.build();
+                columns.push(BlockEntry::new(left_data_type.clone(), Value::Column(col)));
+            } else {
+                return Err(ErrorCode::IllegalDataType(
+                    "The data type on both sides of the union does not match",
+                ));
             }
-            let col = builder.build();
-            Ok(BlockEntry::new(left_data_type.clone(), Value::Column(col)))
-        } else {
-            Err(ErrorCode::IllegalDataType(
-                "The data type on both sides of the union does not match",
-            ))
         }
+        Ok(columns)
     }
 }
 
@@ -148,12 +155,7 @@ impl Processor for TransformMergeBlock {
             return Ok(Event::NeedConsume);
         }
 
-        if self.input_data.is_some() || self.receiver_result.is_some() {
-            return Ok(Event::Sync);
-        }
-
-        if let Ok(result) = self.receiver.try_recv() {
-            self.receiver_result = Some(result);
+        if self.input_data.is_some() || !self.receiver_results.is_empty() {
             return Ok(Event::Sync);
         }
 
@@ -175,28 +177,25 @@ impl Processor for TransformMergeBlock {
     }
 
     fn process(&mut self) -> Result<()> {
+        let mut blocks = vec![];
+        for (idx, receive_result) in self.receiver_results.iter() {
+            blocks.push(self.project_block(receive_result.clone(), Some(*idx))?);
+        }
+        self.receiver_results.clear();
         if let Some(input_data) = self.input_data.take() {
-            if let Some(receiver_result) = self.receiver_result.take() {
-                self.output_data = Some(DataBlock::concat(&[
-                    self.project_block(input_data, true)?,
-                    self.project_block(receiver_result, false)?,
-                ])?);
-            } else {
-                self.output_data = Some(self.project_block(input_data, true)?);
-            }
-        } else if let Some(receiver_result) = self.receiver_result.take() {
-            self.output_data = Some(self.project_block(receiver_result, false)?);
+            blocks.push(self.project_block(input_data, None)?);
         }
-
+        self.output_data = Some(DataBlock::concat(&blocks)?);
         Ok(())
     }
 
     #[async_backtrace::framed]
     async fn async_process(&mut self) -> Result<()> {
         if !self.finished {
-            if let Ok(result) = self.receiver.recv().await {
-                self.receiver_result = Some(result);
-                return Ok(());
+            for (idx, receiver) in self.receivers.iter() {
+                if let Ok(result) = receiver.recv().await {
+                    self.receiver_results.push((*idx, result));
+                }
             }
             self.finished = true;
         }

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

@@ -255,27 +255,24 @@ impl PhysicalPlanReplacer for Fragmenter {
 
     fn replace_union(&mut self, plan: &UnionAll) -> Result<PhysicalPlan> {
         let mut fragments = vec![];
-        let left_input = self.replace(plan.left.as_ref())?;
-        let left_state = self.state.clone();
-
-        // Consume current fragments to prevent them being consumed by `right_input`.
-        fragments.append(&mut self.fragments);
-        let right_input = self.replace(plan.right.as_ref())?;
-        let right_state = self.state.clone();
-
-        fragments.append(&mut self.fragments);
+        let mut children = vec![];
+        let mut states = vec![];
+        for child in plan.children.iter() {
+            children.push(Box::new(self.replace(child)?));
+            states.push(self.state.clone());
+            fragments.append(&mut self.fragments);
+        }
         self.fragments = fragments;
 
         // If any of the input is a source fragment, the union all is a source fragment.
-        if left_state == State::SelectLeaf || right_state == State::SelectLeaf {
+        if states.iter().any(|state| state == &State::SelectLeaf) {
             self.state = State::SelectLeaf;
         } else {
             self.state = State::Other;
         }
 
         Ok(PhysicalPlan::UnionAll(UnionAll {
-            left: Box::new(left_input),
-            right: Box::new(right_input),
+            children,
             ..plan.clone()
         }))
     }

src/query/sql/src/executor/format.rs

@@ -159,14 +159,14 @@ impl PhysicalPlan {
                 ))
             }
             PhysicalPlan::UnionAll(union_all) => {
-                let left_child = union_all.left.format_join(metadata)?;
-                let right_child = union_all.right.format_join(metadata)?;
-
-                let children = vec![
-                    FormatTreeNode::with_children("Left".to_string(), vec![left_child]),
-                    FormatTreeNode::with_children("Right".to_string(), vec![right_child]),
-                ];
-
+                let mut children = Vec::with_capacity(union_all.children.len());
+                for (idx, child) in union_all.children.iter().enumerate() {
+                    let child = child.format_join(metadata)?;
+                    children.push(FormatTreeNode::with_children(
+                        format!("#{:?}", idx + 1),
+                        vec![child],
+                    ))
+                }
                 Ok(FormatTreeNode::with_children(
                     "UnionAll".to_string(),
                     children,
@@ -1013,11 +1013,9 @@ fn union_all_to_format_tree(
     }
 
     append_profile_info(&mut children, profs, plan.plan_id);
-
-    children.extend(vec![
-        to_format_tree(&plan.left, metadata, profs)?,
-        to_format_tree(&plan.right, metadata, profs)?,
-    ]);
+    for child in plan.children.iter() {
+        children.push(to_format_tree(child, metadata, profs)?);
+    }
 
     Ok(FormatTreeNode::with_children(
         "UnionAll".to_string(),

src/query/sql/src/executor/physical_plan.rs

@@ -230,8 +230,9 @@ impl PhysicalPlan {
             PhysicalPlan::UnionAll(plan) => {
                 plan.plan_id = *next_id;
                 *next_id += 1;
-                plan.left.adjust_plan_id(next_id);
-                plan.right.adjust_plan_id(next_id);
+                for child in plan.children.iter_mut() {
+                    child.adjust_plan_id(next_id);
+                }
             }
             PhysicalPlan::CteScan(plan) => {
                 plan.plan_id = *next_id;
@@ -573,9 +574,7 @@ impl PhysicalPlan {
             ),
             PhysicalPlan::Exchange(plan) => Box::new(std::iter::once(plan.input.as_ref())),
             PhysicalPlan::ExchangeSink(plan) => Box::new(std::iter::once(plan.input.as_ref())),
-            PhysicalPlan::UnionAll(plan) => Box::new(
-                std::iter::once(plan.left.as_ref()).chain(std::iter::once(plan.right.as_ref())),
-            ),
+            PhysicalPlan::UnionAll(plan) => Box::new(plan.children.iter().map(|child| &(**child))),
             PhysicalPlan::DistributedInsertSelect(plan) => {
                 Box::new(std::iter::once(plan.input.as_ref()))
             }
@@ -823,11 +822,6 @@ impl PhysicalPlan {
             PhysicalPlan::CteScan(v) => {
                 format!("CTE index: {}, sub index: {}", v.cte_idx.0, v.cte_idx.1)
             }
-            PhysicalPlan::UnionAll(v) => v
-                .pairs
-                .iter()
-                .map(|(l, r)| format!("#{} <- #{}", l, r))
-                .join(", "),
             _ => String::new(),
         })
     }