feat(sql): Support UNION (distinct) for unified SQL

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

@@ -8,7 +8,8 @@
 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 static org.opensearch.sql.ast.dsl.AstDSL.unionAll;
+import static org.opensearch.sql.ast.dsl.AstDSL.unionDistinct;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -29,6 +30,7 @@
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.common.antlr.AstBuildGuard;
 import org.opensearch.sql.common.setting.Settings;
+import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.sql.antlr.SQLSyntaxParser;
 import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser;
 import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ExistsSubqueryExpressionAtomContext;
@@ -154,7 +156,14 @@ private JoinType toJoinType(JoinClauseContext ctx) {
     public UnresolvedPlan visitUnionSelect(UnionSelectContext ctx) {
       List<UnresolvedPlan> datasets =
           ctx.querySpecification().stream().map(this::visit).collect(Collectors.toList());
-      return union(datasets);
+      if (ctx.ALL().isEmpty()) {
+        return unionDistinct(datasets);
+      }
+      if (ctx.ALL().size() == ctx.UNION().size()) {
+        return unionAll(datasets);
+      }
+      throw new SemanticCheckException(
+          "Mixing UNION and UNION ALL in the same query is not supported");
     }
 
     /**

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

@@ -216,6 +216,22 @@ public void testUnionAll() {
             """);
   }
 
+  @Test
+  public void testUnionDistinct() {
+    givenQuery(
+            """
+            SELECT name FROM catalog.employees UNION SELECT dept_name FROM catalog.departments
+            """)
+        .assertPlan(
+            """
+            LogicalUnion(all=[false])
+              LogicalProject(name=[$1])
+                LogicalTableScan(table=[[catalog, employees]])
+              LogicalProject(dept_name=[$1])
+                LogicalTableScan(table=[[catalog, departments]])
+            """);
+  }
+
   @Test
   public void testMultiWayUnion() {
     givenQuery(
@@ -236,6 +252,26 @@ public void testMultiWayUnion() {
             """);
   }
 
+  @Test
+  public void testMultiWayUnionDistinct() {
+    givenQuery(
+            """
+            SELECT name FROM catalog.employees
+            UNION SELECT dept_name FROM catalog.departments
+            UNION SELECT name FROM catalog.employees
+            """)
+        .assertPlan(
+            """
+            LogicalUnion(all=[false])
+              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

@@ -786,4 +786,12 @@ public static ExistsSubquery existsSubquery(UnresolvedPlan query) {
   public static Union union(List<UnresolvedPlan> datasets) {
     return new Union(datasets);
   }
+
+  public static Union unionAll(List<UnresolvedPlan> datasets) {
+    return new Union(datasets, false);
+  }
+
+  public static Union unionDistinct(List<UnresolvedPlan> datasets) {
+    return new Union(datasets, true);
+  }
 }

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

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

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

@@ -3156,7 +3156,7 @@ public RelNode visitUnion(Union node, CalcitePlanContext context) {
     for (RelNode input : unifiedInputs) {
       context.relBuilder.push(input);
     }
-    context.relBuilder.union(true, unifiedInputs.size()); // true = UNION ALL
+    context.relBuilder.union(!node.isDistinct(), unifiedInputs.size()); // all = !distinct
 
     if (node.getMaxout() != null) {
       context.relBuilder.push(

integ-test/src/test/java/org/opensearch/sql/sql/SetOperationIT.java

@@ -0,0 +1,115 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.sql;
+
+import static org.junit.Assert.assertThrows;
+import static org.opensearch.sql.util.Capability.SET_OPERATION;
+import static org.opensearch.sql.util.MatcherUtils.rows;
+import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
+
+import org.json.JSONObject;
+import org.junit.Test;
+import org.opensearch.sql.legacy.SQLIntegTestCase;
+import org.opensearch.sql.legacy.TestsConstants;
+import org.opensearch.sql.util.RequiresCapability;
+
+/** SQL set operations: {@code UNION} and {@code UNION ALL}. */
+@RequiresCapability(SET_OPERATION)
+public class SetOperationIT extends SQLIntegTestCase {
+
+  @Override
+  protected void init() throws Exception {
+    loadIndex(Index.BANK);
+  }
+
+  @Test
+  public void testUnionAll() {
+    JSONObject response =
+        new JSONObject(
+            executeQuery(
+                """
+                SELECT age FROM %s WHERE age < 30\
+                 UNION ALL SELECT age FROM %s WHERE age > 35\
+                """
+                    .formatted(TestsConstants.TEST_INDEX_BANK, TestsConstants.TEST_INDEX_BANK),
+                "jdbc"));
+
+    verifyDataRows(response, rows(28), rows(36), rows(36), rows(39));
+  }
+
+  /** Same operands as {@link #testUnionAll()}: the duplicate 36 collapses to one row. */
+  @Test
+  public void testUnionDistinct() {
+    JSONObject response =
+        new JSONObject(
+            executeQuery(
+                """
+                SELECT age FROM %s WHERE age < 30\
+                 UNION SELECT age FROM %s WHERE age > 35\
+                """
+                    .formatted(TestsConstants.TEST_INDEX_BANK, TestsConstants.TEST_INDEX_BANK),
+                "jdbc"));
+
+    verifyDataRows(response, rows(28), rows(36), rows(39));
+  }
+
+  @Test
+  public void testMultiWayUnionAll() {
+    JSONObject response =
+        new JSONObject(
+            executeQuery(
+                """
+                SELECT age FROM %s WHERE age < 30\
+                 UNION ALL SELECT age FROM %s WHERE age > 35\
+                 UNION ALL SELECT age FROM %s WHERE age = 33\
+                """
+                    .formatted(
+                        TestsConstants.TEST_INDEX_BANK,
+                        TestsConstants.TEST_INDEX_BANK,
+                        TestsConstants.TEST_INDEX_BANK),
+                "jdbc"));
+
+    verifyDataRows(response, rows(28), rows(33), rows(36), rows(36), rows(39));
+  }
+
+  /** Same three operands as {@link #testMultiWayUnionAll()}: the duplicate 36 collapses. */
+  @Test
+  public void testMultiWayUnionDistinct() {
+    JSONObject response =
+        new JSONObject(
+            executeQuery(
+                """
+                SELECT age FROM %s WHERE age < 30\
+                 UNION SELECT age FROM %s WHERE age > 35\
+                 UNION SELECT age FROM %s WHERE age = 33\
+                """
+                    .formatted(
+                        TestsConstants.TEST_INDEX_BANK,
+                        TestsConstants.TEST_INDEX_BANK,
+                        TestsConstants.TEST_INDEX_BANK),
+                "jdbc"));
+
+    verifyDataRows(response, rows(28), rows(33), rows(36), rows(39));
+  }
+
+  @Test
+  public void testMixedUnionAndUnionAllIsRejected() {
+    assertThrows(
+        RuntimeException.class,
+        () ->
+            executeQuery(
+                """
+                SELECT age FROM %s WHERE age < 30\
+                 UNION ALL SELECT age FROM %s WHERE age > 35\
+                 UNION SELECT age FROM %s WHERE age = 33\
+                """
+                    .formatted(
+                        TestsConstants.TEST_INDEX_BANK,
+                        TestsConstants.TEST_INDEX_BANK,
+                        TestsConstants.TEST_INDEX_BANK),
+                "jdbc"));
+  }
+}

integ-test/src/test/java/org/opensearch/sql/util/Backend.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.util;
+
+import java.util.EnumSet;
+import java.util.Map;
+import java.util.Set;
+
+/** Query-execution backends, each paired with the {@link Capability} set it does not support. */
+public enum Backend {
+  OPENSEARCH,
+  ANALYTICS_ENGINE;
+
+  /**
+   * {@code OPENSEARCH} lacks only the capabilities listed here; every other capability in the
+   * registry is an analytics-engine gap, so {@code ANALYTICS_ENGINE} takes the complement.
+   */
+  private static final Map<Backend, Set<Capability>> UNSUPPORTED =
+      Map.of(
+          OPENSEARCH,
+          EnumSet.of(Capability.SET_OPERATION),
+          ANALYTICS_ENGINE,
+          EnumSet.complementOf(EnumSet.of(Capability.SET_OPERATION)));
+
+  /** Whether this backend can run a test declaring {@code capability}. */
+  public boolean supports(Capability capability) {
+    return !UNSUPPORTED.get(this).contains(capability);
+  }
+}

integ-test/src/test/java/org/opensearch/sql/util/BackendCapabilities.java

@@ -16,8 +16,13 @@ public static void requireCapability(Capability capability) {
   }
 
   public static void requireCapability(Capability capability, String note) {
-    // Today analytics-engine supports none of the defined capabilities, so all are skipped on it.
-    Assume.assumeTrue(skipMessage(capability, note), !TestUtils.AnalyticsIndexConfig.isEnabled());
+    Assume.assumeTrue(skipMessage(capability, note), activeBackend().supports(capability));
+  }
+
+  private static Backend activeBackend() {
+    return TestUtils.AnalyticsIndexConfig.isEnabled()
+        ? Backend.ANALYTICS_ENGINE
+        : Backend.OPENSEARCH;
   }
 
   private static String skipMessage(Capability capability, String note) {

integ-test/src/test/java/org/opensearch/sql/util/Capability.java

@@ -8,9 +8,10 @@
 /**
  * Backend-agnostic registry of execution capabilities a test may require. Each constant names a
  * behavior (nested fields, document mutation, stable head ordering, ...) and carries the reason it
- * is unavailable on a backend that lacks it — currently the analytics-engine route. Tests declare
- * the capability they need via {@code BackendCapabilities.requireCapability(...)} or the {@link
- * RequiresCapability} annotation rather than naming a backend.
+ * is unavailable on a backend that lacks it. Which backend lacks which capability lives in {@link
+ * Backend}. Tests declare the capability they need via {@code
+ * BackendCapabilities.requireCapability(...)} or the {@link RequiresCapability} annotation rather
+ * than naming a backend.
  *
  * <p>Keeping every reason here makes the full set of route gaps greppable in one place — both for
  * humans tracking what still needs fixing and as a single block of context to hand an agent for
@@ -582,7 +583,10 @@ public enum Capability {
   /** BACKEND: FILTER(WHERE) on aggregates can't be executed via Substrait streaming. */
   FILTERED_AGGREGATE(
       "FILTER(WHERE) on aggregates can't be executed on the analytics-engine route: the Substrait"
-          + " streaming path doesn't support filtered aggregates.");
+          + " streaming path doesn't support filtered aggregates."),
+
+  /** Combining the result rows of two or more queries with a SQL set operator. */
+  SET_OPERATION("SQL set operations are unsupported.");
 
   private final String reason;
 

sql/src/main/antlr/OpenSearchSQLParser.g4

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