Skip to content

Commit 3a256d2

Browse files
dai-chendai-chen
committed
feat(sql): add IN/EXISTS subquery support in unified query path (opensearch-project#5448)
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 <daichen@amazon.com>
1 parent abb3ca4 commit 3a256d2

5 files changed

Lines changed: 142 additions & 0 deletions

File tree

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,25 @@
55

66
package org.opensearch.sql.api.parser;
77

8+
import static org.opensearch.sql.ast.dsl.AstDSL.existsSubquery;
9+
import static org.opensearch.sql.ast.dsl.AstDSL.inSubquery;
810
import static org.opensearch.sql.ast.dsl.AstDSL.join;
911

1012
import java.util.Optional;
1113
import org.antlr.v4.runtime.tree.ParseTree;
14+
import org.opensearch.sql.ast.expression.Not;
1215
import org.opensearch.sql.ast.expression.UnresolvedExpression;
1316
import org.opensearch.sql.ast.statement.Query;
1417
import org.opensearch.sql.ast.statement.Statement;
1518
import org.opensearch.sql.ast.tree.Join.JoinType;
1619
import org.opensearch.sql.ast.tree.UnresolvedPlan;
1720
import org.opensearch.sql.sql.antlr.SQLSyntaxParser;
1821
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser;
22+
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ExistsSubqueryExpressionAtomContext;
23+
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.InSubqueryPredicateContext;
1924
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.JoinClauseContext;
2025
import org.opensearch.sql.sql.parser.AstBuilder;
26+
import org.opensearch.sql.sql.parser.AstExpressionBuilder;
2127
import org.opensearch.sql.sql.parser.AstStatementBuilder;
2228

2329
/** SQL query parser that produces {@link UnresolvedPlan} using the V2 ANTLR grammar. */
@@ -52,6 +58,11 @@ private static class ExtendedAstBuilder extends AstBuilder {
5258
super(query);
5359
}
5460

61+
@Override
62+
protected AstExpressionBuilder createExpressionBuilder() {
63+
return new ExtendedAstExpressionBuilder();
64+
}
65+
5566
@Override
5667
public UnresolvedPlan visitJoinClause(JoinClauseContext ctx) {
5768
JoinType joinType = toJoinType(ctx);
@@ -69,5 +80,27 @@ private JoinType toJoinType(JoinClauseContext ctx) {
6980
default -> JoinType.INNER;
7081
};
7182
}
83+
84+
/**
85+
* Expression builder with IN/EXISTS subquery support. Accesses the enclosing AstBuilder to
86+
* visit subquery plan nodes. Must be created via {@link #createExpressionBuilder()} because the
87+
* enclosing {@code this} reference is not available during {@code super()} construction.
88+
*/
89+
private class ExtendedAstExpressionBuilder extends AstExpressionBuilder {
90+
91+
@Override
92+
public UnresolvedExpression visitInSubqueryPredicate(InSubqueryPredicateContext ctx) {
93+
UnresolvedPlan subquery = ExtendedAstBuilder.this.visit(ctx.querySpecification());
94+
UnresolvedExpression inExpr = inSubquery(subquery, visit(ctx.predicate()));
95+
return (ctx.NOT() != null) ? new Not(inExpr) : inExpr;
96+
}
97+
98+
@Override
99+
public UnresolvedExpression visitExistsSubqueryExpressionAtom(
100+
ExistsSubqueryExpressionAtomContext ctx) {
101+
UnresolvedPlan subquery = ExtendedAstBuilder.this.visit(ctx.querySpecification());
102+
return existsSubquery(subquery);
103+
}
104+
}
72105
}
73106
}

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,4 +142,80 @@ public void testJoinWithFilterAndOrderBy() {
142142
LogicalTableScan(table=[[catalog, departments]])
143143
""");
144144
}
145+
146+
@Test
147+
public void testInSubquery() {
148+
givenQuery(
149+
"""
150+
SELECT name FROM catalog.employees
151+
WHERE age IN (SELECT age FROM catalog.departments WHERE dept_name = 'Engineering')
152+
""")
153+
.assertPlan(
154+
"""
155+
LogicalProject(name=[$1])
156+
LogicalFilter(condition=[IN($2, {
157+
LogicalProject(age=[$cor0.age])
158+
LogicalFilter(condition=[=($1, 'Engineering')])
159+
LogicalTableScan(table=[[catalog, departments]])
160+
})], variablesSet=[[$cor0]])
161+
LogicalTableScan(table=[[catalog, employees]])
162+
""");
163+
}
164+
165+
@Test
166+
public void testExistsSubquery() {
167+
givenQuery(
168+
"""
169+
SELECT name FROM catalog.employees
170+
WHERE EXISTS (SELECT 1 FROM catalog.departments WHERE dept_id = age)
171+
""")
172+
.assertPlan(
173+
"""
174+
LogicalProject(name=[$1])
175+
LogicalFilter(condition=[EXISTS({
176+
LogicalProject(1=[1])
177+
LogicalFilter(condition=[=($0, $cor0.age)])
178+
LogicalTableScan(table=[[catalog, departments]])
179+
})], variablesSet=[[$cor0]])
180+
LogicalTableScan(table=[[catalog, employees]])
181+
""");
182+
}
183+
184+
@Test
185+
public void testNotInSubquery() {
186+
givenQuery(
187+
"""
188+
SELECT name FROM catalog.employees
189+
WHERE age NOT IN (SELECT age FROM catalog.departments WHERE dept_name = 'Engineering')
190+
""")
191+
.assertPlan(
192+
"""
193+
LogicalProject(name=[$1])
194+
LogicalFilter(condition=[NOT(IN($2, {
195+
LogicalProject(age=[$cor0.age])
196+
LogicalFilter(condition=[=($1, 'Engineering')])
197+
LogicalTableScan(table=[[catalog, departments]])
198+
}))], variablesSet=[[$cor0]])
199+
LogicalTableScan(table=[[catalog, employees]])
200+
""");
201+
}
202+
203+
@Test
204+
public void testNotExistsSubquery() {
205+
givenQuery(
206+
"""
207+
SELECT name FROM catalog.employees
208+
WHERE NOT EXISTS (SELECT 1 FROM catalog.departments WHERE dept_id = age)
209+
""")
210+
.assertPlan(
211+
"""
212+
LogicalProject(name=[$1])
213+
LogicalFilter(condition=[NOT(EXISTS({
214+
LogicalProject(1=[1])
215+
LogicalFilter(condition=[=($0, $cor0.age)])
216+
LogicalTableScan(table=[[catalog, departments]])
217+
}))], variablesSet=[[$cor0]])
218+
LogicalTableScan(table=[[catalog, employees]])
219+
""");
220+
}
145221
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@
4848
import org.opensearch.sql.ast.expression.When;
4949
import org.opensearch.sql.ast.expression.WindowFunction;
5050
import org.opensearch.sql.ast.expression.Xor;
51+
import org.opensearch.sql.ast.expression.subquery.ExistsSubquery;
52+
import org.opensearch.sql.ast.expression.subquery.InSubquery;
5153
import org.opensearch.sql.ast.tree.Aggregation;
5254
import org.opensearch.sql.ast.tree.AppendPipe;
5355
import org.opensearch.sql.ast.tree.Bin;
@@ -771,4 +773,12 @@ public static UnresolvedPlan join(
771773
Optional.empty(),
772774
Argument.ArgumentMap.empty());
773775
}
776+
777+
public static InSubquery inSubquery(UnresolvedPlan query, UnresolvedExpression... values) {
778+
return new InSubquery(List.of(values), query);
779+
}
780+
781+
public static ExistsSubquery existsSubquery(UnresolvedPlan query) {
782+
return new ExistsSubquery(query);
783+
}
774784
}

‎integ-test/src/test/java/org/opensearch/sql/legacy/SqlLegacyEngineSanityIT.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,4 +45,13 @@ public void testLeftJoinFallback() throws IOException {
4545
.formatted(TEST_INDEX_PEOPLE, TEST_INDEX_DOG));
4646
verifyDataRows(result, rows("Daenerys", "rex"));
4747
}
48+
49+
@Test
50+
public void testInSubqueryFallback() throws IOException {
51+
JSONObject result =
52+
executeQuery(
53+
"SELECT a.firstname FROM %s a WHERE a.firstname IN (SELECT holdersName FROM %s)"
54+
.formatted(TEST_INDEX_PEOPLE, TEST_INDEX_DOG));
55+
verifyDataRows(result, rows("Daenerys"), rows("Hattie"));
56+
}
4857
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -756,4 +756,18 @@ public UnresolvedPlan visitJoinClause(OpenSearchSQLParser.JoinClauseContext ctx)
756756
};
757757
assertNotNull(new SQLSyntaxParser().parse(query).accept(builder));
758758
}
759+
760+
@Test
761+
public void in_subquery_throws_syntax_check_exception() {
762+
assertThrows(
763+
SyntaxCheckException.class,
764+
() -> buildAST("SELECT * FROM t WHERE age IN (SELECT age FROM t2)"));
765+
}
766+
767+
@Test
768+
public void exists_subquery_throws_syntax_check_exception() {
769+
assertThrows(
770+
SyntaxCheckException.class,
771+
() -> buildAST("SELECT * FROM t WHERE EXISTS (SELECT 1 FROM t2)"));
772+
}
759773
}

0 commit comments

Comments
 (0)