From bd3074d2a77ca2c5e9a2c004372304638309df56 Mon Sep 17 00:00:00 2001 From: Chen Dai Date: Fri, 15 May 2026 14:34:06 -0700 Subject: [PATCH] feat(sql): add IN/EXISTS subquery support in unified query path Add grammar rules (inSubqueryPredicate, existsSubqueryExpressionAtom) and wire them through ExtendedAstExpressionBuilder to produce InSubquery and ExistsSubquery AST nodes for the Calcite-based unified query path. Base AstExpressionBuilder throws SyntaxCheckException to preserve legacy engine fallback. AstBuilder now uses createExpressionBuilder() factory method to allow subclass customization. Also add Alias handling in CalciteRelNodeVisitor.expandProjectFields required for any non-SELECT * query in the unified path. Signed-off-by: Chen Dai --- .../sql/api/parser/SqlV2QueryParser.java | 33 ++++++++ .../sql/api/UnifiedQueryPlannerSqlV2Test.java | 76 +++++++++++++++++++ .../org/opensearch/sql/ast/dsl/AstDSL.java | 10 +++ .../sql/calcite/CalciteRelNodeVisitor.java | 3 + .../sql/legacy/SqlLegacyEngineSanityIT.java | 9 +++ sql/src/main/antlr/OpenSearchSQLParser.g4 | 2 + .../opensearch/sql/sql/parser/AstBuilder.java | 14 +++- .../sql/sql/parser/AstExpressionBuilder.java | 16 ++++ .../sql/sql/parser/AstBuilderTest.java | 14 ++++ 9 files changed, 174 insertions(+), 3 deletions(-) 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 30e88e5ce3a..d6280b829c7 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 @@ -5,10 +5,13 @@ package org.opensearch.sql.api.parser; +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 java.util.Optional; import org.antlr.v4.runtime.tree.ParseTree; +import org.opensearch.sql.ast.expression.Not; import org.opensearch.sql.ast.expression.UnresolvedExpression; import org.opensearch.sql.ast.statement.Query; import org.opensearch.sql.ast.statement.Statement; @@ -16,8 +19,11 @@ import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.sql.antlr.SQLSyntaxParser; import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser; +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.parser.AstBuilder; +import org.opensearch.sql.sql.parser.AstExpressionBuilder; import org.opensearch.sql.sql.parser.AstStatementBuilder; /** SQL query parser that produces {@link UnresolvedPlan} using the V2 ANTLR grammar. */ @@ -52,6 +58,11 @@ private static class ExtendedAstBuilder extends AstBuilder { super(query); } + @Override + protected AstExpressionBuilder createExpressionBuilder() { + return new ExtendedAstExpressionBuilder(); + } + @Override public UnresolvedPlan visitJoinClause(JoinClauseContext ctx) { JoinType joinType = toJoinType(ctx); @@ -69,5 +80,27 @@ private JoinType toJoinType(JoinClauseContext ctx) { default -> JoinType.INNER; }; } + + /** + * Expression builder with IN/EXISTS subquery support. Accesses the enclosing AstBuilder to + * visit subquery plan nodes. Must be created via {@link #createExpressionBuilder()} because the + * enclosing {@code this} reference is not available during {@code super()} construction. + */ + private class ExtendedAstExpressionBuilder extends AstExpressionBuilder { + + @Override + public UnresolvedExpression visitInSubqueryPredicate(InSubqueryPredicateContext ctx) { + UnresolvedPlan subquery = ExtendedAstBuilder.this.visit(ctx.querySpecification()); + UnresolvedExpression inExpr = inSubquery(subquery, visit(ctx.predicate())); + return (ctx.NOT() != null) ? new Not(inExpr) : inExpr; + } + + @Override + public UnresolvedExpression visitExistsSubqueryExpressionAtom( + ExistsSubqueryExpressionAtomContext ctx) { + UnresolvedPlan subquery = ExtendedAstBuilder.this.visit(ctx.querySpecification()); + return existsSubquery(subquery); + } + } } } 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 afe08e3a2ad..5d7135ec7e7 100644 --- a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java +++ b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java @@ -142,4 +142,80 @@ public void testJoinWithFilterAndOrderBy() { LogicalTableScan(table=[[catalog, departments]]) """); } + + @Test + public void testInSubquery() { + givenQuery( + """ + SELECT name FROM catalog.employees + WHERE age IN (SELECT age FROM catalog.departments WHERE dept_name = 'Engineering') + """) + .assertPlan( + """ + LogicalProject(name=[$1]) + LogicalFilter(condition=[IN($2, { + LogicalProject(age=[$cor0.age]) + LogicalFilter(condition=[=($1, 'Engineering')]) + LogicalTableScan(table=[[catalog, departments]]) + })], variablesSet=[[$cor0]]) + LogicalTableScan(table=[[catalog, employees]]) + """); + } + + @Test + public void testExistsSubquery() { + givenQuery( + """ + SELECT name FROM catalog.employees + WHERE EXISTS (SELECT 1 FROM catalog.departments WHERE dept_id = age) + """) + .assertPlan( + """ + LogicalProject(name=[$1]) + LogicalFilter(condition=[EXISTS({ + LogicalProject(1=[1]) + LogicalFilter(condition=[=($0, $cor0.age)]) + LogicalTableScan(table=[[catalog, departments]]) + })], variablesSet=[[$cor0]]) + LogicalTableScan(table=[[catalog, employees]]) + """); + } + + @Test + public void testNotInSubquery() { + givenQuery( + """ + SELECT name FROM catalog.employees + WHERE age NOT IN (SELECT age FROM catalog.departments WHERE dept_name = 'Engineering') + """) + .assertPlan( + """ + LogicalProject(name=[$1]) + LogicalFilter(condition=[NOT(IN($2, { + LogicalProject(age=[$cor0.age]) + LogicalFilter(condition=[=($1, 'Engineering')]) + LogicalTableScan(table=[[catalog, departments]]) + }))], variablesSet=[[$cor0]]) + LogicalTableScan(table=[[catalog, employees]]) + """); + } + + @Test + public void testNotExistsSubquery() { + givenQuery( + """ + SELECT name FROM catalog.employees + WHERE NOT EXISTS (SELECT 1 FROM catalog.departments WHERE dept_id = age) + """) + .assertPlan( + """ + LogicalProject(name=[$1]) + LogicalFilter(condition=[NOT(EXISTS({ + LogicalProject(1=[1]) + LogicalFilter(condition=[=($0, $cor0.age)]) + LogicalTableScan(table=[[catalog, departments]]) + }))], variablesSet=[[$cor0]]) + LogicalTableScan(table=[[catalog, employees]]) + """); + } } 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 f38e46377dc..f00a87ab40f 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 @@ -48,6 +48,8 @@ import org.opensearch.sql.ast.expression.When; import org.opensearch.sql.ast.expression.WindowFunction; import org.opensearch.sql.ast.expression.Xor; +import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; +import org.opensearch.sql.ast.expression.subquery.InSubquery; import org.opensearch.sql.ast.tree.Aggregation; import org.opensearch.sql.ast.tree.AppendPipe; import org.opensearch.sql.ast.tree.Bin; @@ -771,4 +773,12 @@ public static UnresolvedPlan join( Optional.empty(), Argument.ArgumentMap.empty()); } + + public static InSubquery inSubquery(UnresolvedPlan query, UnresolvedExpression... values) { + return new InSubquery(List.of(values), query); + } + + public static ExistsSubquery existsSubquery(UnresolvedPlan query) { + return new ExistsSubquery(query); + } } 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 1251f51b131..3d5bf4f5a8e 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -541,6 +541,9 @@ private List expandProjectFields( .filter(addedFields::add) .forEach(field -> expandedFields.add(context.relBuilder.field(field))); } + case Alias alias -> { + expandedFields.add(rexVisitor.analyze(alias, context)); + } default -> throw new IllegalStateException( "Unexpected expression type in project list: " + expr.getClass().getSimpleName()); diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/SqlLegacyEngineSanityIT.java b/integ-test/src/test/java/org/opensearch/sql/legacy/SqlLegacyEngineSanityIT.java index d13df7e2fe7..b2d22808a49 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/SqlLegacyEngineSanityIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/SqlLegacyEngineSanityIT.java @@ -45,4 +45,13 @@ public void testLeftJoinFallback() throws IOException { .formatted(TEST_INDEX_PEOPLE, TEST_INDEX_DOG)); verifyDataRows(result, rows("Daenerys", "rex")); } + + @Test + public void testInSubqueryFallback() throws IOException { + JSONObject result = + executeQuery( + "SELECT a.firstname FROM %s a WHERE a.firstname IN (SELECT holdersName FROM %s)" + .formatted(TEST_INDEX_PEOPLE, TEST_INDEX_DOG)); + verifyDataRows(result, rows("Daenerys"), rows("Hattie")); + } } diff --git a/sql/src/main/antlr/OpenSearchSQLParser.g4 b/sql/src/main/antlr/OpenSearchSQLParser.g4 index f0fbec498f9..492f6dee9c6 100644 --- a/sql/src/main/antlr/OpenSearchSQLParser.g4 +++ b/sql/src/main/antlr/OpenSearchSQLParser.g4 @@ -322,6 +322,7 @@ predicate | left = predicate NOT? LIKE right = predicate # likePredicate | left = predicate REGEXP right = predicate # regexpPredicate | predicate NOT? IN '(' expressions ')' # inPredicate + | predicate NOT? IN '(' querySpecification ')' # inSubqueryPredicate ; expressions @@ -333,6 +334,7 @@ expressionAtom | columnName # fullColumnNameExpressionAtom | functionCall # functionCallExpressionAtom | LR_BRACKET expression RR_BRACKET # nestedExpressionAtom + | EXISTS LR_BRACKET querySpecification RR_BRACKET # existsSubqueryExpressionAtom | left = expressionAtom mathOperator = (STAR | SLASH | MODULE) right = expressionAtom # mathExpressionAtom | left = expressionAtom mathOperator = (PLUS | MINUS) right = expressionAtom # mathExpressionAtom ; diff --git a/sql/src/main/java/org/opensearch/sql/sql/parser/AstBuilder.java b/sql/src/main/java/org/opensearch/sql/sql/parser/AstBuilder.java index 3575aa8919a..ee532a10ed9 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/parser/AstBuilder.java +++ b/sql/src/main/java/org/opensearch/sql/sql/parser/AstBuilder.java @@ -23,7 +23,6 @@ import java.util.Collections; import java.util.Locale; import java.util.Optional; -import lombok.RequiredArgsConstructor; import org.antlr.v4.runtime.tree.ParseTree; import org.opensearch.sql.ast.expression.Alias; import org.opensearch.sql.ast.expression.AllFields; @@ -50,10 +49,9 @@ import org.opensearch.sql.sql.parser.context.ParsingContext; /** Abstract syntax tree (AST) builder. */ -@RequiredArgsConstructor public class AstBuilder extends OpenSearchSQLParserBaseVisitor { - private final AstExpressionBuilder expressionBuilder = new AstExpressionBuilder(); + private final AstExpressionBuilder expressionBuilder; /** Parsing context stack that contains context for current query parsing. */ private final ParsingContext context = new ParsingContext(); @@ -64,6 +62,11 @@ public class AstBuilder extends OpenSearchSQLParserBaseVisitor { */ private final String query; + public AstBuilder(String query) { + this.query = query; + this.expressionBuilder = createExpressionBuilder(); + } + @Override public UnresolvedPlan visitShowStatement(OpenSearchSQLParser.ShowStatementContext ctx) { final UnresolvedExpression tableFilter = visitAstExpression(ctx.tableFilter()); @@ -279,6 +282,11 @@ protected UnresolvedExpression visitAstExpression(ParseTree tree) { return expressionBuilder.visit(tree); } + /** Override to provide a custom expression builder (e.g., with subquery support). */ + protected AstExpressionBuilder createExpressionBuilder() { + return new AstExpressionBuilder(); + } + private UnresolvedExpression visitSelectItem(SelectElementContext ctx) { String name = StringUtils.unquoteIdentifier(getTextInQuery(ctx.expression(), query)); UnresolvedExpression expr = visitAstExpression(ctx.expression()); diff --git a/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java b/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java index 346ef6660d7..273076af40f 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java +++ b/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java @@ -28,6 +28,7 @@ import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.DataTypeFunctionCallContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.DateLiteralContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.DistinctCountFunctionCallContext; +import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ExistsSubqueryExpressionAtomContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ExtractFunctionCallContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.FilterClauseContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.FilteredAggregationFunctionCallContext; @@ -35,6 +36,7 @@ import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.GetFormatFunctionCallContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.HighlightFunctionCallContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.InPredicateContext; +import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.InSubqueryPredicateContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.IsNullPredicateContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.LikePredicateContext; import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.MathExpressionAtomContext; @@ -82,6 +84,7 @@ import org.opensearch.sql.ast.dsl.AstDSL; import org.opensearch.sql.ast.expression.*; import org.opensearch.sql.ast.tree.Sort.SortOption; +import org.opensearch.sql.common.antlr.SyntaxCheckException; import org.opensearch.sql.common.utils.StringUtils; import org.opensearch.sql.expression.function.BuiltinFunctionName; import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser; @@ -668,4 +671,17 @@ private List getExtractFunctionArguments(ExtractFunctionCa visitFunctionArg(ctx.extractFunction().functionArg())); return args; } + + @Override + public UnresolvedExpression visitInSubqueryPredicate(InSubqueryPredicateContext ctx) { + throw new SyntaxCheckException( + "IN subquery is not supported in the V2 SQL engine. Falling back to legacy engine."); + } + + @Override + public UnresolvedExpression visitExistsSubqueryExpressionAtom( + ExistsSubqueryExpressionAtomContext ctx) { + throw new SyntaxCheckException( + "EXISTS subquery is not supported in the V2 SQL engine. Falling back to legacy engine."); + } } diff --git a/sql/src/test/java/org/opensearch/sql/sql/parser/AstBuilderTest.java b/sql/src/test/java/org/opensearch/sql/sql/parser/AstBuilderTest.java index d6897230b40..7869ba5cdad 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/parser/AstBuilderTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/parser/AstBuilderTest.java @@ -756,4 +756,18 @@ public UnresolvedPlan visitJoinClause(OpenSearchSQLParser.JoinClauseContext ctx) }; assertNotNull(new SQLSyntaxParser().parse(query).accept(builder)); } + + @Test + public void in_subquery_throws_syntax_check_exception() { + assertThrows( + SyntaxCheckException.class, + () -> buildAST("SELECT * FROM t WHERE age IN (SELECT age FROM t2)")); + } + + @Test + public void exists_subquery_throws_syntax_check_exception() { + assertThrows( + SyntaxCheckException.class, + () -> buildAST("SELECT * FROM t WHERE EXISTS (SELECT 1 FROM t2)")); + } }