Merge pull request #8638 from xudong963/fix_11_4

fix(union): union's pairs are handled incorrectly
databendlabs · Nov 5, 2022 · 5ae5272 · 5ae5272 · vercel · Nov 5, 2022
2 parents 14d314a + 7a3acf2
commit 5ae5272
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 9 deletions.
diff --git a/src/query/service/src/pipelines/processors/transforms/transform_merge_block.rs b/src/query/service/src/pipelines/processors/transforms/transform_merge_block.rs
@@ -60,16 +60,17 @@ impl TransformMergeBlock {
         })))
     }
 
-    fn project_block(&self, block: DataBlock) -> Result<DataBlock> {
+    fn project_block(&self, block: DataBlock, is_left: bool) -> Result<DataBlock> {
         tracing::info!("Processing block: {:?}", &block);
         let columns = self
             .pairs
             .iter()
             .map(|(left, right)| {
-                Ok(block
-                    .try_column_by_name(left)
-                    .or_else(|_| block.try_column_by_name(right))?
-                    .clone())
+                if is_left {
+                    Ok(block.try_column_by_name(left)?.clone())
+                } else {
+                    Ok(block.try_column_by_name(right)?.clone())
+                }
             })
             .collect::<Result<Vec<_>>>()?;
         Ok(DataBlock::create(self.schema.clone(), columns))
@@ -132,14 +133,14 @@ impl Processor for TransformMergeBlock {
         if let Some(input_data) = self.input_data.take() {
             if let Some(receiver_result) = self.receiver_result.take() {
                 self.output_data = Some(DataBlock::concat_blocks(&[
-                    self.project_block(input_data)?,
-                    self.project_block(receiver_result)?,
+                    self.project_block(input_data, true)?,
+                    self.project_block(receiver_result, false)?,
                 ])?);
             } else {
-                self.output_data = Some(self.project_block(input_data)?);
+                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)?);
+            self.output_data = Some(self.project_block(receiver_result, false)?);
         }
 
         Ok(())

diff --git a/tests/logictest/suites/query/union.test b/tests/logictest/suites/query/union.test
@@ -114,3 +114,34 @@ SELECT 'Кирилл' as a, 'Müller' as b, 1.0 as c, 2 as id UNION SELECT NULL
 NULL NULL 1.0 1
 Кирилл Müller 1.0 2
 
+statement query TI
+WITH tbl AS
+  (SELECT '12-17' AS age,
+          'a' AS device_id
+   UNION ALL SELECT '17-23' AS age,
+                    'a' AS device_id
+   UNION ALL SELECT '17-23' AS age,
+                    'b' AS device_id
+   UNION ALL SELECT '24-30' AS age,
+                    'c' AS device_id
+   UNION ALL SELECT '24-30' AS age,
+                    'd' AS device_id)
+SELECT age,
+       count(DISTINCT device_id) AS cnt
+FROM
+  (SELECT age,
+          device_id
+   FROM tbl
+   UNION ALL SELECT if(age IN ('12-17', '17-23'), '年轻', '非年轻') AS age,
+                    device_id
+   FROM tbl) AS t
+GROUP BY age
+ORDER BY age,
+         cnt;
+
+----
+12-17  1
+17-23  2
+24-30  2
+年轻  2
+非年轻  2