Add UNION ALL support to V3 SQL engine (opensearch-project#5506)

Add UNION ALL set operations to the Calcite-based unified
query path. The legacy V2 SQL engine rejects UNION with a
SyntaxCheckException so queries fall back to the legacy engine
appropriately.

Signed-off-by: Chen Dai <daichen@amazon.com>

Author: dai-chen

api/src/main/java/org/opensearch/sql/api/parser/SqlV2QueryParser.java

@@ -8,8 +8,11 @@
 import static org.opensearch.sql.ast.dsl.AstDSL.existsSubquery;
 import static org.opensearch.sql.ast.dsl.AstDSL.inSubquery;
 import static org.opensearch.sql.ast.dsl.AstDSL.join;
+import static org.opensearch.sql.ast.dsl.AstDSL.union;
 
+import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
 import org.antlr.v4.runtime.tree.ParseTree;
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
@@ -22,6 +25,7 @@
 import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ExistsSubqueryExpressionAtomContext;
 import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.InSubqueryPredicateContext;
 import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.JoinClauseContext;
+import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.UnionSelectContext;
 import org.opensearch.sql.sql.parser.AstBuilder;
 import org.opensearch.sql.sql.parser.AstExpressionBuilder;
 import org.opensearch.sql.sql.parser.AstStatementBuilder;
@@ -81,6 +85,13 @@ private JoinType toJoinType(JoinClauseContext ctx) {
       };
     }
 
+    @Override
+    public UnresolvedPlan visitUnionSelect(UnionSelectContext ctx) {
+      List<UnresolvedPlan> datasets =
+          ctx.querySpecification().stream().map(this::visit).collect(Collectors.toList());
+      return union(datasets);
+    }
+
     /**
      * Expression builder with IN/EXISTS subquery support. Accesses the enclosing AstBuilder to
      * visit subquery plan nodes. Must be created via {@link #createExpressionBuilder()} because the

api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java

@@ -200,6 +200,42 @@ WHERE age NOT IN (SELECT age FROM catalog.departments WHERE dept_name = 'Enginee
             """);
   }
 
+  @Test
+  public void testUnionAll() {
+    givenQuery(
+            """
+            SELECT name FROM catalog.employees UNION ALL SELECT dept_name FROM catalog.departments
+            """)
+        .assertPlan(
+            """
+            LogicalUnion(all=[true])
+              LogicalProject(name=[$1])
+                LogicalTableScan(table=[[catalog, employees]])
+              LogicalProject(dept_name=[$1])
+                LogicalTableScan(table=[[catalog, departments]])
+            """);
+  }
+
+  @Test
+  public void testMultiWayUnion() {
+    givenQuery(
+            """
+            SELECT name FROM catalog.employees
+            UNION ALL SELECT dept_name FROM catalog.departments
+            UNION ALL SELECT name FROM catalog.employees
+            """)
+        .assertPlan(
+            """
+            LogicalUnion(all=[true])
+              LogicalProject(name=[$1])
+                LogicalTableScan(table=[[catalog, employees]])
+              LogicalProject(dept_name=[$1])
+                LogicalTableScan(table=[[catalog, departments]])
+              LogicalProject(name=[$1])
+                LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
   @Test
   public void testNotExistsSubquery() {
     givenQuery(

core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java

@@ -84,6 +84,7 @@
 import org.opensearch.sql.ast.tree.SubqueryAlias;
 import org.opensearch.sql.ast.tree.TableFunction;
 import org.opensearch.sql.ast.tree.Trendline;
+import org.opensearch.sql.ast.tree.Union;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
 import org.opensearch.sql.calcite.plan.OpenSearchConstants;
@@ -781,4 +782,8 @@ public static InSubquery inSubquery(UnresolvedPlan query, UnresolvedExpression..
   public static ExistsSubquery existsSubquery(UnresolvedPlan query) {
     return new ExistsSubquery(query);
   }
+
+  public static Union union(List<UnresolvedPlan> datasets) {
+    return new Union(datasets);
+  }
 }

core/src/main/java/org/opensearch/sql/ast/tree/Union.java

@@ -21,15 +21,25 @@
 @RequiredArgsConstructor
 @AllArgsConstructor
 public class Union extends UnresolvedPlan {
+  /** Input subplans (operands) combined by this UNION ALL. */
   private final List<UnresolvedPlan> datasets;
 
+  /** Whether inputs are unified to a common schema by name (PPL) vs combined positionally (SQL). */
+  private boolean unifySchema;
+
+  /** Optional cap on output rows (PPL {@code maxout}); {@code null} if unbounded. */
   private Integer maxout;
 
+  /** PPL constructor: UNION ALL with schema unification. */
+  public Union(List<UnresolvedPlan> datasets, Integer maxout) {
+    this(datasets, true, maxout);
+  }
+
   @Override
   public UnresolvedPlan attach(UnresolvedPlan child) {
     List<UnresolvedPlan> newDatasets =
         ImmutableList.<UnresolvedPlan>builder().add(child).addAll(datasets).build();
-    return new Union(newDatasets, maxout);
+    return new Union(newDatasets, unifySchema, maxout);
   }
 
   @Override

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -3031,8 +3031,10 @@ public RelNode visitUnion(Union node, CalcitePlanContext context) {
           "Union command requires at least two datasets. Provided: " + inputNodes.size());
     }
 
-    List<RelNode> unifiedInputs =
-        SchemaUnifier.buildUnifiedSchemaWithTypeCoercion(inputNodes, context);
+    List<RelNode> unifiedInputs = inputNodes;
+    if (node.isUnifySchema()) {
+      unifiedInputs = SchemaUnifier.buildUnifiedSchemaWithTypeCoercion(inputNodes, context);
+    }
 
     for (RelNode input : unifiedInputs) {
       context.relBuilder.push(input);

sql/src/main/antlr/OpenSearchSQLParser.g4

@@ -68,7 +68,8 @@ dmlStatement
 
 // Primary DML Statements
 selectStatement
-   : querySpecification # simpleSelect
+   : querySpecification                                          # simpleSelect
+   | querySpecification (UNION ALL querySpecification)+          # unionSelect
    ;
 
 adminStatement

sql/src/main/java/org/opensearch/sql/sql/parser/AstBuilder.java

@@ -263,6 +263,12 @@ public UnresolvedPlan visitJoinClause(OpenSearchSQLParser.JoinClauseContext ctx)
         "JOIN is not supported in the V2 SQL engine. Falling back to legacy engine.");
   }
 
+  @Override
+  public UnresolvedPlan visitUnionSelect(OpenSearchSQLParser.UnionSelectContext ctx) {
+    throw new SyntaxCheckException(
+        "UNION is not supported in the V2 SQL engine. Falling back to legacy engine.");
+  }
+
   @Override
   public UnresolvedPlan visitHavingClause(HavingClauseContext ctx) {
     AstHavingFilterBuilder builder = new AstHavingFilterBuilder(context.peek());

sql/src/test/java/org/opensearch/sql/sql/parser/AstBuilderTest.java

@@ -757,6 +757,13 @@ public UnresolvedPlan visitJoinClause(OpenSearchSQLParser.JoinClauseContext ctx)
     assertNotNull(new SQLSyntaxParser().parse(query).accept(builder));
   }
 
+  @Test
+  public void union_throws_syntax_check_exception() {
+    assertThrows(
+        SyntaxCheckException.class,
+        () -> buildAST("SELECT name FROM t1 UNION ALL SELECT name FROM t2"));
+  }
+
   @Test
   public void in_subquery_throws_syntax_check_exception() {
     assertThrows(