disable EliminateNestedUnion

wren-core/core/src/mdl/context.rs

@@ -29,7 +29,6 @@ use datafusion::optimizer::eliminate_filter::EliminateFilter;
 use datafusion::optimizer::eliminate_group_by_constant::EliminateGroupByConstant;
 use datafusion::optimizer::eliminate_join::EliminateJoin;
 use datafusion::optimizer::eliminate_limit::EliminateLimit;
-use datafusion::optimizer::eliminate_nested_union::EliminateNestedUnion;
 use datafusion::optimizer::eliminate_one_union::EliminateOneUnion;
 use datafusion::optimizer::eliminate_outer_join::EliminateOuterJoin;
 use datafusion::optimizer::extract_equijoin_predicate::ExtractEquijoinPredicate;
@@ -153,7 +152,9 @@ fn analyze_rule_for_unparsing(
 /// Optimizer rules for unparse
 fn optimize_rule_for_unparsing() -> Vec<Arc<dyn OptimizerRule + Send + Sync>> {
     vec![
-        Arc::new(EliminateNestedUnion::new()),
+        // Disable EliminateNestedUnion because unparser only support unparsing an union with two inputs
+        // see https://github.com/apache/datafusion/issues/13621 for details
+        // Arc::new(EliminateNestedUnion::new()),
         // Disable SimplifyExpressions to avoid apply some function locally
         // Arc::new(SimplifyExpressions::new()),
         Arc::new(UnwrapCastInComparison::new()),

wren-core/core/src/mdl/mod.rs

@@ -1202,6 +1202,26 @@ mod test {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_disable_eliminate_nested_union() -> Result<()> {
+        let ctx = SessionContext::new();
+        let sql = r#"SELECT * FROM (SELECT 1 x, 'a' y UNION ALL
+    SELECT 1 x, 'b' y UNION ALL
+    SELECT 2 x, 'a' y UNION ALL
+    SELECT 2 x, 'c' y)"#;
+        let result =
+            transform_sql_with_ctx(&ctx, Arc::new(AnalyzedWrenMDL::default()), &[], sql)
+                .await?;
+        assert_eq!(
+            result,
+            "SELECT x, y FROM (SELECT 1 AS x, 'a' AS y \
+        UNION ALL SELECT 1 AS x, 'b' AS y \
+        UNION ALL SELECT 2 AS x, 'a' AS y \
+        UNION ALL SELECT 2 AS x, 'c' AS y)"
+        );
+        Ok(())
+    }
+
     /// Return a RecordBatch with made up data about customer
     fn customer() -> RecordBatch {
         let custkey: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));