diff --git a/api/src/main/java/org/opensearch/sql/api/parser/SqlV2QueryParser.java b/api/src/main/java/org/opensearch/sql/api/parser/SqlV2QueryParser.java index b1f61e732ff..d20097f8f88 100644 --- a/api/src/main/java/org/opensearch/sql/api/parser/SqlV2QueryParser.java +++ b/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 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"); } /** diff --git a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java index ec3c36daa89..b06d1bba62a 100644 --- a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java +++ b/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( diff --git a/core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java b/core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java index 3a96137f42d..4b3d2b20ed3 100644 --- a/core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java +++ b/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 datasets) { return new Union(datasets); } + + public static Union unionAll(List datasets) { + return new Union(datasets, false); + } + + public static Union unionDistinct(List datasets) { + return new Union(datasets, true); + } } diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Union.java b/core/src/main/java/org/opensearch/sql/ast/tree/Union.java index 7eefe89ce53..27dc17e6556 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Union.java +++ b/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 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 datasets, Integer maxout) { - this(datasets, true, maxout); + this(datasets, true, maxout, false); + } + + /** SQL constructor: inputs combined positionally, with or without deduplication. */ + public Union(List datasets, boolean distinct) { + this(datasets, false, null, distinct); } @Override public UnresolvedPlan attach(UnresolvedPlan child) { List newDatasets = ImmutableList.builder().add(child).addAll(datasets).build(); - return new Union(newDatasets, unifySchema, maxout); + return new Union(newDatasets, unifySchema, maxout, distinct); } @Override diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index 2c5a6aadcf7..7280c6c755f 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/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( diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/SetOperationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/SetOperationIT.java new file mode 100644 index 00000000000..9a955cc7c8e --- /dev/null +++ b/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")); + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/util/Backend.java b/integ-test/src/test/java/org/opensearch/sql/util/Backend.java new file mode 100644 index 00000000000..51425484a27 --- /dev/null +++ b/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> 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); + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/util/BackendCapabilities.java b/integ-test/src/test/java/org/opensearch/sql/util/BackendCapabilities.java index 1fd18571a8f..5505aad6764 100644 --- a/integ-test/src/test/java/org/opensearch/sql/util/BackendCapabilities.java +++ b/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) { diff --git a/integ-test/src/test/java/org/opensearch/sql/util/Capability.java b/integ-test/src/test/java/org/opensearch/sql/util/Capability.java index bb99dc85115..efdc8e1f1e3 100644 --- a/integ-test/src/test/java/org/opensearch/sql/util/Capability.java +++ b/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. * *

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; diff --git a/sql/src/main/antlr/OpenSearchSQLParser.g4 b/sql/src/main/antlr/OpenSearchSQLParser.g4 index e372382805c..f33a086738e 100644 --- a/sql/src/main/antlr/OpenSearchSQLParser.g4 +++ b/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