Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
5 changes: 5 additions & 0 deletions core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}
12 changes: 11 additions & 1 deletion core/src/main/java/org/opensearch/sql/ast/tree/Union.java
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
3 changes: 2 additions & 1 deletion sql/src/main/antlr/OpenSearchSQLParser.g4
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,8 @@ dmlStatement

// Primary DML Statements
selectStatement
: querySpecification # simpleSelect
: querySpecification # simpleSelect
| querySpecification (UNION ALL querySpecification)+ # unionSelect
;

adminStatement
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Loading