fix(cubestore): Panic on combination of Union All and single select in root Union All (#6012)

Author: waralexrom

rust/cubestore/cubestore/src/queryplanner/flatten_union.rs

@@ -63,20 +63,20 @@ fn try_remove_sub_union(
     parent_inputs: &Vec<LogicalPlan>,
     parent_schema: Arc<DFSchema>,
 ) -> Vec<LogicalPlan> {
-    let mut may_be_result = Vec::new();
+    let mut result = Vec::new();
     for inp in parent_inputs.iter() {
         match inp {
             LogicalPlan::Union { inputs, schema, .. } => {
                 if schema.to_schema_ref() == parent_schema.to_schema_ref() {
-                    may_be_result.extend(inputs.iter().cloned());
+                    result.extend(inputs.iter().cloned());
                 } else {
                     return parent_inputs.clone();
                 }
             }
             _ => {
-                return parent_inputs.clone();
+                result.push(inp.clone());
             }
         }
     }
-    return may_be_result;
+    return result;
 }

rust/cubestore/cubestore/src/sql/mod.rs

@@ -2358,6 +2358,80 @@ mod tests {
                 }
                 _ => assert!(false),
             };
+
+            let result = service.exec_query("EXPLAIN SELECT a `sel__a`, b `sel__b`, sum(c) `sel__c` from ( \
+                         select * from ( \
+                                        select * from foo.a\
+                                        ) \
+                             union all 
+                             select * from 
+                                ( \
+                                        select * from foo.a1 \
+                                        union all \
+                                        select * from foo.b1 \
+                                ) \
+                            union all
+                            select * from foo.b \
+                         ) AS `lambda` where a = 1 group by 1, 2 order by 3 desc").await.unwrap();
+            match &result.get_rows()[0].values()[0] {
+                TableValue::String(s) => {
+                    println!("!! s {}", s);
+                    assert_eq!(s,
+                                "Sort\
+                                \n  Projection, [sel__a, sel__b, sel__c]\
+                                \n    Aggregate\
+                                \n      ClusterSend, indices: [[1, 3, 4, 2]]\
+                                \n        Union\
+                                \n          Filter\
+                                \n            Scan foo.a, source: CubeTable(index: default:1:[1]:sort_on[a, b]), fields: *\
+                                \n          Filter\
+                                \n            Scan foo.a1, source: CubeTable(index: default:3:[3]:sort_on[a, b]), fields: *\
+                                \n          Filter\
+                                \n            Scan foo.b1, source: CubeTable(index: default:4:[4]:sort_on[a, b]), fields: *\
+                                \n          Filter\
+                                \n            Scan foo.b, source: CubeTable(index: default:2:[2]:sort_on[a, b]), fields: *"
+
+                               );
+                }
+                _ => assert!(false),
+            };
+            let result = service.exec_query("EXPLAIN SELECT a `sel__a`, b `sel__b`, sum(c) `sel__c` from ( \
+                         select * from ( \
+                                        select * from foo.a where 1 = 0\
+                                        ) \
+                             union all 
+                             select * from 
+                                ( \
+                                        select * from foo.a1 \
+                                        union all \
+                                        select * from foo.b1 \
+                                ) \
+                            union all
+                            select * from foo.b \
+                         ) AS `lambda` where a = 1 group by 1, 2 order by 3 desc").await.unwrap();
+            match &result.get_rows()[0].values()[0] {
+                TableValue::String(s) => {
+                    println!("!! s {}", s);
+                    assert_eq!(s,
+                                "Sort\
+                                \n  Projection, [sel__a, sel__b, sel__c]\
+                                \n    Aggregate\
+                                \n      ClusterSend, indices: [[1, 3, 4, 2]]\
+                                \n        Union\
+                                \n          Filter\
+                                \n            Filter\
+                                \n              Scan foo.a, source: CubeTable(index: default:1:[1]:sort_on[a, b]), fields: *\
+                                \n          Filter\
+                                \n            Scan foo.a1, source: CubeTable(index: default:3:[3]:sort_on[a, b]), fields: *\
+                                \n          Filter\
+                                \n            Scan foo.b1, source: CubeTable(index: default:4:[4]:sort_on[a, b]), fields: *\
+                                \n          Filter\
+                                \n            Scan foo.b, source: CubeTable(index: default:2:[2]:sort_on[a, b]), fields: *"
+
+                               );
+                }
+                _ => assert!(false),
+            };
         }).await;
     }