Skip to content

Commit 7d8613d

Browse files
dai-chendai-chen
committed
test(api): Add SQL planner tests and refactor test base for multi-language support
- Refactor UnifiedQueryTestBase with queryType() hook for subclass override - Add UnifiedSqlQueryPlannerTest covering SELECT, WHERE, GROUP BY, JOIN, ORDER BY, subquery, case sensitivity, namespaces, and error handling - Update UnifiedQueryContextTest to verify SQL context creation Signed-off-by: Chen Dai <daichen@amazon.com>
1 parent d6335d2 commit 7d8613d

3 files changed

Lines changed: 269 additions & 9 deletions

File tree

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -63,14 +63,15 @@ public void testMissingQueryType() {
6363
UnifiedQueryContext.builder().catalog("opensearch", testSchema).build();
6464
}
6565

66-
@Test(expected = IllegalArgumentException.class)
67-
public void testUnsupportedQueryType() {
66+
@Test
67+
public void testSqlQueryType() {
6868
UnifiedQueryContext context =
6969
UnifiedQueryContext.builder()
70-
.language(QueryType.SQL) // only PPL is supported for now
70+
.language(QueryType.SQL)
7171
.catalog("opensearch", testSchema)
7272
.build();
73-
new UnifiedQueryPlanner(context);
73+
UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
74+
assertNotNull("SQL planner should be created", planner);
7475
}
7576

7677
@Test(expected = IllegalArgumentException.class)

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

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.api;
7+
8+
import static org.junit.Assert.assertNotNull;
9+
import static org.junit.Assert.assertThrows;
10+
11+
import java.util.Map;
12+
import org.apache.calcite.schema.Schema;
13+
import org.apache.calcite.schema.impl.AbstractSchema;
14+
import org.junit.Test;
15+
import org.opensearch.sql.executor.QueryType;
16+
17+
public class UnifiedQueryPlannerSqlTest extends UnifiedQueryTestBase {
18+
19+
private final AbstractSchema testDeepSchema =
20+
new AbstractSchema() {
21+
@Override
22+
protected Map<String, Schema> getSubSchemaMap() {
23+
return Map.of("opensearch", testSchema);
24+
}
25+
};
26+
27+
@Override
28+
protected QueryType queryType() {
29+
return QueryType.SQL;
30+
}
31+
32+
@Test
33+
public void testSqlQueryPlanning() {
34+
givenQuery(
35+
"""
36+
SELECT *
37+
FROM catalog.employees\
38+
""")
39+
.assertPlan(
40+
"""
41+
LogicalProject(id=[$0], name=[$1], age=[$2], department=[$3])
42+
LogicalTableScan(table=[[catalog, employees]])
43+
""");
44+
}
45+
46+
@Test
47+
public void testSqlSelectSpecificColumns() {
48+
givenQuery(
49+
"""
50+
SELECT id, name
51+
FROM catalog.employees\
52+
""")
53+
.assertPlan(
54+
"""
55+
LogicalProject(id=[$0], name=[$1])
56+
LogicalTableScan(table=[[catalog, employees]])
57+
""");
58+
}
59+
60+
@Test
61+
public void testSqlFilterQueryPlanning() {
62+
givenQuery(
63+
"""
64+
SELECT name
65+
FROM catalog.employees
66+
WHERE age > 30\
67+
""")
68+
.assertPlan(
69+
"""
70+
LogicalProject(name=[$1])
71+
LogicalFilter(condition=[>($2, 30)])
72+
LogicalTableScan(table=[[catalog, employees]])
73+
""");
74+
}
75+
76+
@Test
77+
public void testSqlAggregateQueryPlanning() {
78+
givenQuery(
79+
"""
80+
SELECT department, count(*) AS cnt
81+
FROM catalog.employees
82+
GROUP BY department\
83+
""")
84+
.assertPlan(
85+
"""
86+
LogicalAggregate(group=[{0}], cnt=[COUNT()])
87+
LogicalProject(department=[$3])
88+
LogicalTableScan(table=[[catalog, employees]])
89+
""");
90+
}
91+
92+
@Test
93+
public void testSqlJoinQueryPlanning() {
94+
givenQuery(
95+
"""
96+
SELECT a.id, b.name
97+
FROM catalog.employees a
98+
JOIN catalog.employees b ON a.id = b.age\
99+
""")
100+
.assertPlan(
101+
"""
102+
LogicalProject(id=[$0], name=[$5])
103+
LogicalJoin(condition=[=($0, $6)], joinType=[inner])
104+
LogicalTableScan(table=[[catalog, employees]])
105+
LogicalTableScan(table=[[catalog, employees]])
106+
""");
107+
}
108+
109+
@Test
110+
public void testSqlOrderByQueryPlanning() {
111+
givenQuery(
112+
"""
113+
SELECT name
114+
FROM catalog.employees
115+
ORDER BY age DESC\
116+
""")
117+
.assertPlan(
118+
"""
119+
LogicalProject(name=[$0])
120+
LogicalSort(sort0=[$1], dir0=[DESC])
121+
LogicalProject(name=[$1], age=[$2])
122+
LogicalTableScan(table=[[catalog, employees]])
123+
""");
124+
}
125+
126+
@Test
127+
public void testSqlSubqueryPlanning() {
128+
// Calcite represents scalar subqueries as $SCALAR_QUERY{...} with embedded plan text whose
129+
// formatting (whitespace, line breaks) may vary across versions. Assert output fields only.
130+
givenQuery(
131+
"""
132+
SELECT name
133+
FROM catalog.employees
134+
WHERE age > (SELECT avg(age) FROM catalog.employees)\
135+
""")
136+
.assertFields("name");
137+
}
138+
139+
@Test
140+
public void testSqlCteQueryPlanning() {
141+
// CTE is inlined by Calcite — same plan as a direct filter query
142+
givenQuery(
143+
"""
144+
WITH seniors AS (
145+
SELECT name, age FROM catalog.employees WHERE age > 30
146+
)
147+
SELECT name
148+
FROM seniors\
149+
""")
150+
.assertPlan(
151+
"""
152+
LogicalProject(name=[$1])
153+
LogicalFilter(condition=[>($2, 30)])
154+
LogicalTableScan(table=[[catalog, employees]])
155+
""");
156+
}
157+
158+
@Test
159+
public void testSqlQueryPlanningWithDefaultNamespace() {
160+
UnifiedQueryContext sqlContext =
161+
UnifiedQueryContext.builder()
162+
.language(QueryType.SQL)
163+
.catalog("opensearch", testSchema)
164+
.defaultNamespace("opensearch")
165+
.build();
166+
UnifiedQueryPlanner sqlPlanner = new UnifiedQueryPlanner(sqlContext);
167+
168+
assertNotNull("Plan should be created", sqlPlanner.plan("SELECT * FROM opensearch.employees"));
169+
assertNotNull("Plan should be created", sqlPlanner.plan("SELECT * FROM employees"));
170+
}
171+
172+
@Test
173+
public void testSqlQueryPlanningWithDefaultNamespaceMultiLevel() {
174+
UnifiedQueryContext sqlContext =
175+
UnifiedQueryContext.builder()
176+
.language(QueryType.SQL)
177+
.catalog("catalog", testDeepSchema)
178+
.defaultNamespace("catalog.opensearch")
179+
.build();
180+
UnifiedQueryPlanner sqlPlanner = new UnifiedQueryPlanner(sqlContext);
181+
182+
assertNotNull(
183+
"Plan should be created", sqlPlanner.plan("SELECT * FROM catalog.opensearch.employees"));
184+
assertNotNull("Plan should be created", sqlPlanner.plan("SELECT * FROM employees"));
185+
186+
assertThrows(
187+
IllegalStateException.class, () -> sqlPlanner.plan("SELECT * FROM opensearch.employees"));
188+
}
189+
190+
@Test
191+
public void testSqlQueryPlanningWithMultipleCatalogs() {
192+
UnifiedQueryContext sqlContext =
193+
UnifiedQueryContext.builder()
194+
.language(QueryType.SQL)
195+
.catalog("catalog1", testSchema)
196+
.catalog("catalog2", testSchema)
197+
.build();
198+
UnifiedQueryPlanner sqlPlanner = new UnifiedQueryPlanner(sqlContext);
199+
200+
assertNotNull(
201+
"Plan should be created",
202+
sqlPlanner.plan(
203+
"""
204+
SELECT a.id
205+
FROM catalog1.employees a
206+
JOIN catalog2.employees b ON a.id = b.id\
207+
"""));
208+
}
209+
210+
@Test
211+
public void testInvalidSqlThrowsException() {
212+
assertThrows(IllegalStateException.class, () -> planner.plan("SELECT FROM"));
213+
}
214+
}

‎api/src/testFixtures/java/org/opensearch/sql/api/UnifiedQueryTestBase.java‎

Lines changed: 50 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import static org.apache.calcite.sql.type.SqlTypeName.INTEGER;
99
import static org.apache.calcite.sql.type.SqlTypeName.VARCHAR;
10+
import static org.junit.Assert.assertEquals;
1011

1112
import java.util.List;
1213
import java.util.Map;
@@ -15,6 +16,8 @@
1516
import org.apache.calcite.DataContext;
1617
import org.apache.calcite.linq4j.Enumerable;
1718
import org.apache.calcite.linq4j.Linq4j;
19+
import org.apache.calcite.plan.RelOptUtil;
20+
import org.apache.calcite.rel.RelNode;
1821
import org.apache.calcite.rel.type.RelDataType;
1922
import org.apache.calcite.rel.type.RelDataTypeFactory;
2023
import org.apache.calcite.schema.ScannableTable;
@@ -55,14 +58,25 @@ protected Map<String, Table> getTableMap() {
5558
}
5659
};
5760

58-
context =
59-
UnifiedQueryContext.builder()
60-
.language(QueryType.PPL)
61-
.catalog(DEFAULT_CATALOG, testSchema)
62-
.build();
61+
context = buildContext(queryType());
6362
planner = new UnifiedQueryPlanner(context);
6463
}
6564

65+
/**
66+
* Returns the query type for this test class. Subclasses override to test different languages.
67+
*/
68+
protected QueryType queryType() {
69+
return QueryType.PPL;
70+
}
71+
72+
/** Builds a UnifiedQueryContext with the test schema for the given query type. */
73+
protected UnifiedQueryContext buildContext(QueryType queryType) {
74+
return UnifiedQueryContext.builder()
75+
.language(queryType)
76+
.catalog(DEFAULT_CATALOG, testSchema)
77+
.build();
78+
}
79+
6680
@After
6781
public void tearDown() throws Exception {
6882
if (context != null) {
@@ -128,4 +142,35 @@ public boolean rolledUpColumnValidInsideAgg(
128142
return false;
129143
}
130144
}
145+
146+
/** Fluent helper for asserting query plan results. */
147+
protected QueryAssert givenQuery(String query) {
148+
return new QueryAssert(planner.plan(query));
149+
}
150+
151+
/** Fluent assertion on a query's logical plan. */
152+
protected static class QueryAssert {
153+
private final RelNode plan;
154+
155+
QueryAssert(RelNode plan) {
156+
this.plan = plan;
157+
}
158+
159+
/** Assert the logical plan matches the expected tree string. */
160+
public QueryAssert assertPlan(String expected) {
161+
assertEquals(expected, RelOptUtil.toString(plan));
162+
return this;
163+
}
164+
165+
/** Assert the output field names match. */
166+
public QueryAssert assertFields(String... names) {
167+
assertEquals(List.of(names), plan.getRowType().getFieldNames());
168+
return this;
169+
}
170+
171+
/** Access the underlying plan for custom assertions. */
172+
public RelNode plan() {
173+
return plan;
174+
}
175+
}
131176
}

0 commit comments

Comments
 (0)