fix(cubestore): Panic when nested Union All (#6010)

Author: waralexrom

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

@@ -0,0 +1,82 @@
+use datafusion::error::DataFusionError;
+use datafusion::execution::context::ExecutionProps;
+use datafusion::logical_plan::{DFSchema, LogicalPlan};
+use datafusion::optimizer::optimizer::OptimizerRule;
+use datafusion::optimizer::utils;
+use std::sync::Arc;
+
+pub struct FlattenUnion;
+impl OptimizerRule for FlattenUnion {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> Result<LogicalPlan, DataFusionError> {
+        match plan {
+            LogicalPlan::Union { inputs, schema, .. } => {
+                let new_inputs = inputs
+                    .iter()
+                    .map(|p| self.optimize(p, execution_props))
+                    .collect::<Result<Vec<_>, _>>()?;
+
+                let result_inputs = try_remove_sub_union(&new_inputs, schema.clone());
+
+                let expr = plan.expressions().clone();
+
+                utils::from_plan(plan, &expr, &result_inputs)
+            }
+            // Rest: recurse into plan, apply optimization where possible
+            LogicalPlan::Filter { .. }
+            | LogicalPlan::Projection { .. }
+            | LogicalPlan::Window { .. }
+            | LogicalPlan::Aggregate { .. }
+            | LogicalPlan::Repartition { .. }
+            | LogicalPlan::CreateExternalTable { .. }
+            | LogicalPlan::Extension { .. }
+            | LogicalPlan::Sort { .. }
+            | LogicalPlan::Explain { .. }
+            | LogicalPlan::Limit { .. }
+            | LogicalPlan::Skip { .. }
+            | LogicalPlan::Join { .. }
+            | LogicalPlan::CrossJoin { .. } => {
+                // apply the optimization to all inputs of the plan
+                let inputs = plan.inputs();
+                let new_inputs = inputs
+                    .iter()
+                    .map(|p| self.optimize(p, execution_props))
+                    .collect::<Result<Vec<_>, _>>()?;
+
+                let expr = plan.expressions().clone();
+
+                utils::from_plan(plan, &expr, &new_inputs)
+            }
+            LogicalPlan::TableScan { .. } | LogicalPlan::EmptyRelation { .. } => Ok(plan.clone()),
+        }
+    }
+
+    fn name(&self) -> &str {
+        "flatten_union"
+    }
+}
+
+fn try_remove_sub_union(
+    parent_inputs: &Vec<LogicalPlan>,
+    parent_schema: Arc<DFSchema>,
+) -> Vec<LogicalPlan> {
+    let mut may_be_result = Vec::new();
+    for inp in parent_inputs.iter() {
+        match inp {
+            LogicalPlan::Union { inputs, schema, .. } => {
+                if *schema == *&parent_schema {
+                    may_be_result.extend(inputs.iter().cloned());
+                } else {
+                    return parent_inputs.clone();
+                }
+            }
+            _ => {
+                return parent_inputs.clone();
+            }
+        }
+    }
+    return may_be_result;
+}

rust/cubestore/cubestore/src/queryplanner/mod.rs

@@ -11,6 +11,7 @@ mod topk;
 pub use topk::MIN_TOPK_STREAM_ROWS;
 mod coalesce;
 mod filter_by_key_range;
+mod flatten_union;
 pub mod info_schema;
 pub mod now;
 pub mod udfs;
@@ -21,6 +22,7 @@ use crate::config::ConfigObj;
 use crate::metastore::multi_index::MultiPartition;
 use crate::metastore::table::{Table, TablePath};
 use crate::metastore::{IdRow, MetaStore};
+use crate::queryplanner::flatten_union::FlattenUnion;
 use crate::queryplanner::info_schema::{
     SchemataInfoSchemaTableDef, SystemCacheTableDef, SystemChunksTableDef, SystemIndexesTableDef,
     SystemJobsTableDef, SystemPartitionsTableDef, SystemReplayHandlesTableDef,
@@ -171,7 +173,9 @@ impl QueryPlannerImpl {
 impl QueryPlannerImpl {
     async fn execution_context(&self) -> Result<Arc<ExecutionContext>, CubeError> {
         Ok(Arc::new(ExecutionContext::with_config(
-            ExecutionConfig::new().add_optimizer_rule(Arc::new(MaterializeNow {})),
+            ExecutionConfig::new()
+                .add_optimizer_rule(Arc::new(MaterializeNow {}))
+                .add_optimizer_rule(Arc::new(FlattenUnion {})),
         )))
     }
 }

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

@@ -2293,6 +2293,68 @@ mod tests {
         }).await;
     }
 
+    #[tokio::test]
+    async fn flatten_union() {
+        Config::test("flatten_union").start_test(async move |services| {
+            let service = services.sql_service;
+
+            let _ = service.exec_query("CREATE SCHEMA foo").await.unwrap();
+
+            let _ = service.exec_query("CREATE TABLE foo.a (a int, b int, c int)").await.unwrap();
+            let _ = service.exec_query("CREATE TABLE foo.b (a int, b int, c int)").await.unwrap();
+
+            let _ = service.exec_query("CREATE TABLE foo.a1 (a int, b int, c int)").await.unwrap();
+            let _ = service.exec_query("CREATE TABLE foo.b1 (a int, b int, c int)").await.unwrap();
+
+            service.exec_query(
+                "INSERT INTO foo.a (a, b, c) VALUES (1, 1, 1)"
+            ).await.unwrap();
+            service.exec_query(
+                "INSERT INTO foo.b (a, b, c) VALUES (2, 2, 1)"
+            ).await.unwrap();
+            service.exec_query(
+                "INSERT INTO foo.a1 (a, b, c) VALUES (1, 1, 2)"
+            ).await.unwrap();
+            service.exec_query(
+                "INSERT INTO foo.b1 (a, b, c) VALUES (2, 2, 2)"
+            ).await.unwrap();
+
+            let result = service.exec_query("EXPLAIN SELECT a, b, sum(c) from ( \
+                         select * from ( \
+                                        select * from foo.a \
+                                        union all \
+                                        select * from foo.b \
+                                        ) \
+                             union all 
+                             select * from 
+                                ( \
+                                        select * from foo.a1 \
+                                        union all \
+                                        select * from foo.b1 \
+                                        union all \
+                                        select * from foo.b \
+                                ) \
+                         ) group by 1, 2").await.unwrap();
+            match &result.get_rows()[0].values()[0] {
+                TableValue::String(s) => {
+                    assert_eq!(s,
+                                "Projection, [a, b, SUM(c)]\
+                                \n  Aggregate\
+                                \n    ClusterSend, indices: [[1, 2, 3, 4, 2]]\
+                                \n      Union\
+                                \n        Scan foo.a, source: CubeTable(index: default:1:[1]:sort_on[a, b]), fields: *\
+                                \n        Scan foo.b, source: CubeTable(index: default:2:[2]:sort_on[a, b]), fields: *\
+                                \n        Scan foo.a1, source: CubeTable(index: default:3:[3]:sort_on[a, b]), fields: *\
+                                \n        Scan foo.b1, source: CubeTable(index: default:4:[4]:sort_on[a, b]), fields: *\
+                                \n        Scan foo.b, source: CubeTable(index: default:2:[2]:sort_on[a, b]), fields: *"
+
+                               );
+                }
+                _ => assert!(false),
+            };
+        }).await;
+    }
+
     #[tokio::test]
     async fn over_10k_join() {
         Config::test("over_10k_join").update_config(|mut c| {