Skip to content

Commit 81b3db0

Browse files
SwiddisSwiddis
authored
Merge branch 'main' into increment-version-3.8.0
2 parents 38a2688 + 6199484 commit 81b3db0

75 files changed

Lines changed: 4205 additions & 279 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.apache.calcite.rel.RelRoot;
1515
import org.apache.calcite.rel.core.Sort;
1616
import org.apache.calcite.rel.logical.LogicalSort;
17+
import org.apache.calcite.runtime.CalciteException;
1718
import org.apache.calcite.sql.SqlKind;
1819
import org.apache.calcite.sql.SqlNode;
1920
import org.apache.calcite.sql.util.SqlVisitor;
@@ -25,6 +26,7 @@
2526
import org.opensearch.sql.ast.tree.UnresolvedPlan;
2627
import org.opensearch.sql.calcite.CalciteRelNodeVisitor;
2728
import org.opensearch.sql.common.antlr.SyntaxCheckException;
29+
import org.opensearch.sql.common.error.ErrorReport;
2830
import org.opensearch.sql.exception.QueryEngineException;
2931
import org.opensearch.sql.exception.SemanticCheckException;
3032

@@ -74,13 +76,18 @@ public RelNode plan(String query) {
7476
} catch (SyntaxCheckException
7577
| QueryEngineException
7678
| UnsupportedOperationException
77-
| IllegalArgumentException e) {
78-
LOG.error("Failed to plan query: {}", e.getMessage());
79+
| IllegalArgumentException
80+
| ErrorReport e) {
81+
LOG.warn("Failed to plan query: {}", e.getMessage());
7982
throw e;
83+
} catch (CalciteException e) {
84+
// Calcite validation errors (e.g. table not found) indicate an invalid query.
85+
LOG.warn("Failed to plan query, invalid query: {}", e.getMessage());
86+
throw new SemanticCheckException(e.getMessage(), e);
8087
} catch (AssertionError e) {
8188
// Calcite throws assertion error directly when building bad RelNode
8289
String message = "Failed to plan query: invalid plan structure";
83-
LOG.error(message, e);
90+
LOG.warn(message, e);
8491
throw new SemanticCheckException(message, e);
8592
} catch (Exception e) {
8693
String message = "Failed to plan query: unexpected error";

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,11 @@
88
import static org.opensearch.sql.ast.dsl.AstDSL.existsSubquery;
99
import static org.opensearch.sql.ast.dsl.AstDSL.inSubquery;
1010
import static org.opensearch.sql.ast.dsl.AstDSL.join;
11+
import static org.opensearch.sql.ast.dsl.AstDSL.union;
1112

13+
import java.util.List;
1214
import java.util.Optional;
15+
import java.util.stream.Collectors;
1316
import org.antlr.v4.runtime.tree.ParseTree;
1417
import org.opensearch.sql.ast.expression.Not;
1518
import org.opensearch.sql.ast.expression.UnresolvedExpression;
@@ -22,6 +25,7 @@
2225
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ExistsSubqueryExpressionAtomContext;
2326
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.InSubqueryPredicateContext;
2427
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.JoinClauseContext;
28+
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.UnionSelectContext;
2529
import org.opensearch.sql.sql.parser.AstBuilder;
2630
import org.opensearch.sql.sql.parser.AstExpressionBuilder;
2731
import org.opensearch.sql.sql.parser.AstStatementBuilder;
@@ -81,6 +85,13 @@ private JoinType toJoinType(JoinClauseContext ctx) {
8185
};
8286
}
8387

88+
@Override
89+
public UnresolvedPlan visitUnionSelect(UnionSelectContext ctx) {
90+
List<UnresolvedPlan> datasets =
91+
ctx.querySpecification().stream().map(this::visit).collect(Collectors.toList());
92+
return union(datasets);
93+
}
94+
8495
/**
8596
* Expression builder with IN/EXISTS subquery support. Accesses the enclosing AstBuilder to
8697
* visit subquery plan nodes. Must be created via {@link #createExpressionBuilder()} because the

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

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,42 @@ WHERE age NOT IN (SELECT age FROM catalog.departments WHERE dept_name = 'Enginee
200200
""");
201201
}
202202

203+
@Test
204+
public void testUnionAll() {
205+
givenQuery(
206+
"""
207+
SELECT name FROM catalog.employees UNION ALL SELECT dept_name FROM catalog.departments
208+
""")
209+
.assertPlan(
210+
"""
211+
LogicalUnion(all=[true])
212+
LogicalProject(name=[$1])
213+
LogicalTableScan(table=[[catalog, employees]])
214+
LogicalProject(dept_name=[$1])
215+
LogicalTableScan(table=[[catalog, departments]])
216+
""");
217+
}
218+
219+
@Test
220+
public void testMultiWayUnion() {
221+
givenQuery(
222+
"""
223+
SELECT name FROM catalog.employees
224+
UNION ALL SELECT dept_name FROM catalog.departments
225+
UNION ALL SELECT name FROM catalog.employees
226+
""")
227+
.assertPlan(
228+
"""
229+
LogicalUnion(all=[true])
230+
LogicalProject(name=[$1])
231+
LogicalTableScan(table=[[catalog, employees]])
232+
LogicalProject(dept_name=[$1])
233+
LogicalTableScan(table=[[catalog, departments]])
234+
LogicalProject(name=[$1])
235+
LogicalTableScan(table=[[catalog, employees]])
236+
""");
237+
}
238+
203239
@Test
204240
public void testNotExistsSubquery() {
205241
givenQuery(
@@ -241,6 +277,57 @@ public void selectExpressionWithoutFrom() {
241277
""");
242278
}
243279

280+
@Test
281+
public void testGroupByAggregateAlias() {
282+
givenQuery(
283+
"""
284+
SELECT department, SUM(age) AS total FROM catalog.employees GROUP BY department
285+
""")
286+
.assertPlan(
287+
"""
288+
LogicalProject(department=[$0], total=[$1])
289+
LogicalAggregate(group=[{0}], SUM(age)=[SUM($1)])
290+
LogicalProject(department=[$3], age=[$2])
291+
LogicalTableScan(table=[[catalog, employees]])
292+
""");
293+
}
294+
295+
@Test
296+
public void testOrderByAggregateAlias() {
297+
givenQuery(
298+
"""
299+
SELECT department, COUNT(*) AS cnt FROM catalog.employees
300+
GROUP BY department ORDER BY cnt DESC LIMIT 3
301+
""")
302+
.assertPlan(
303+
"""
304+
LogicalSort(sort0=[$1], dir0=[DESC-nulls-last])
305+
LogicalProject(department=[$1], cnt=[$0])
306+
LogicalSort(sort0=[$0], dir0=[DESC-nulls-last], fetch=[3])
307+
LogicalProject(COUNT(*)=[$1], department=[$0])
308+
LogicalAggregate(group=[{0}], COUNT(*)=[COUNT()])
309+
LogicalProject(department=[$3])
310+
LogicalTableScan(table=[[catalog, employees]])
311+
""");
312+
}
313+
314+
@Test
315+
public void testAliasPreservedInOutputSchema() {
316+
givenQuery("SELECT COUNT(*) AS cnt FROM catalog.employees").assertFields("cnt");
317+
318+
givenQuery("SELECT department, COUNT(*) AS cnt FROM catalog.employees GROUP BY department")
319+
.assertFields("department", "cnt");
320+
321+
givenQuery("SELECT department, COUNT(*) FROM catalog.employees GROUP BY department")
322+
.assertFields("department", "COUNT(*)");
323+
324+
givenQuery("SELECT MAX(age) + MIN(age) AS range_sum FROM catalog.employees")
325+
.assertFields("range_sum");
326+
327+
givenQuery("SELECT id, name, age AS years, department FROM catalog.employees")
328+
.assertFields("id", "name", "years", "department");
329+
}
330+
244331
@Test
245332
public void testHavingMaxCol() {
246333
givenQuery(
@@ -259,6 +346,47 @@ GROUP BY department HAVING MAX(age) > 30
259346
""");
260347
}
261348

349+
@Test
350+
public void testCountStarWithFilter() {
351+
givenQuery("SELECT COUNT(*) FILTER(WHERE age > 30) FROM catalog.employees")
352+
.assertPlan(
353+
"""
354+
LogicalAggregate(group=[{}], COUNT(*) FILTER(WHERE age > 30)=[COUNT() FILTER $0])
355+
LogicalProject($f1=[>($2, 30)])
356+
LogicalTableScan(table=[[catalog, employees]])
357+
""");
358+
}
359+
360+
@Test
361+
public void testFilteredAggregateWithGroupBy() {
362+
givenQuery(
363+
"""
364+
SELECT department, SUM(age) FILTER(WHERE age > 30) FROM catalog.employees
365+
GROUP BY department
366+
""")
367+
.assertPlan(
368+
"""
369+
LogicalAggregate(group=[{0}], SUM(age) FILTER(WHERE age > 30)=[SUM($1) FILTER $2])
370+
LogicalProject(department=[$3], age=[$2], $f3=[>($2, 30)])
371+
LogicalTableScan(table=[[catalog, employees]])
372+
""");
373+
}
374+
375+
@Test
376+
public void testMultipleFilteredAggregates() {
377+
givenQuery(
378+
"""
379+
SELECT MAX(age) FILTER(WHERE age > 30), MIN(age) FILTER(WHERE age < 50)
380+
FROM catalog.employees
381+
""")
382+
.assertPlan(
383+
"""
384+
LogicalAggregate(group=[{}], MAX(age) FILTER(WHERE age > 30)=[MAX($0) FILTER $1], MIN(age) FILTER(WHERE age < 50)=[MIN($0) FILTER $2])
385+
LogicalProject(age=[$2], $f4=[>($2, 30)], $f5=[<($2, 50)])
386+
LogicalTableScan(table=[[catalog, employees]])
387+
""");
388+
}
389+
262390
@Test
263391
public void testScalarFnOverAggregate() {
264392
givenQuery("SELECT ABS(MAX(age)) FROM catalog.employees")
@@ -337,6 +465,33 @@ GROUP BY department HAVING MAX(age) > 30 AND MIN(age) < 50
337465
""");
338466
}
339467

468+
@Test
469+
public void testCountDistinctWindowWithOrderBy() {
470+
// No frame printed: RANGE .. CURRENT ROW is Calcite's default for ORDER BY.
471+
givenQuery(
472+
"""
473+
SELECT department, COUNT(DISTINCT name) OVER(ORDER BY department) FROM catalog.employees
474+
""")
475+
.assertPlan(
476+
"""
477+
LogicalProject(department=[$3], COUNT(DISTINCT name) OVER(ORDER BY department)=[COUNT(DISTINCT $1) OVER (ORDER BY $3 NULLS FIRST)])
478+
LogicalTableScan(table=[[catalog, employees]])
479+
""");
480+
}
481+
482+
@Test
483+
public void testSumWindowWithPartitionAndOrderBy() {
484+
givenQuery(
485+
"""
486+
SELECT name, SUM(age) OVER(PARTITION BY department ORDER BY age) FROM catalog.employees
487+
""")
488+
.assertPlan(
489+
"""
490+
LogicalProject(name=[$1], SUM(age) OVER(PARTITION BY department ORDER BY age)=[SUM($2) OVER (PARTITION BY $3 ORDER BY $2 NULLS FIRST)])
491+
LogicalTableScan(table=[[catalog, employees]])
492+
""");
493+
}
494+
340495
@Test
341496
public void testWindowOrderByDefaultsNullsFirst() {
342497
// Window function ORDER BY without explicit NULLS FIRST/LAST defaults to NULLS FIRST,

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,12 @@
1010

1111
import java.util.Map;
1212
import org.apache.calcite.rel.RelNode;
13+
import org.apache.calcite.runtime.CalciteException;
1314
import org.apache.calcite.schema.Schema;
1415
import org.apache.calcite.schema.impl.AbstractSchema;
1516
import org.junit.Test;
1617
import org.opensearch.sql.common.antlr.SyntaxCheckException;
18+
import org.opensearch.sql.common.error.ErrorReport;
1719
import org.opensearch.sql.exception.SemanticCheckException;
1820
import org.opensearch.sql.executor.QueryType;
1921

@@ -72,7 +74,7 @@ public void testPPLQueryPlanningWithDefaultNamespaceMultiLevel() {
7274

7375
// This is valid in SparkSQL, but Calcite requires "catalog" as the default root schema to
7476
// resolve it
75-
assertThrows(IllegalStateException.class, () -> planner.plan("source = opensearch.employees"));
77+
assertThrows(SemanticCheckException.class, () -> planner.plan("source = opensearch.employees"));
7678
}
7779

7880
@Test
@@ -131,6 +133,20 @@ public void semanticErrorIsRethrownAsSemanticCheckException() {
131133
.assertErrorMessageEquals("Source and target patterns have different wildcard counts");
132134
}
133135

136+
@Test
137+
public void fieldNotFoundIsRethrownAsErrorReport() {
138+
givenInvalidQuery("source = catalog.employees | where unknown_field = 1")
139+
.assertErrorType(ErrorReport.class)
140+
.assertErrorMessageContains("Field [unknown_field] not found");
141+
}
142+
143+
@Test
144+
public void invalidTableIsRethrownAsSemanticCheckException() {
145+
givenInvalidQuery("source = catalog.nonexistent_table")
146+
.assertErrorType(SemanticCheckException.class)
147+
.assertCauseType(CalciteException.class);
148+
}
149+
134150
@Test
135151
public void assertionErrorIsWrappedAsSemanticCheckException() {
136152
// Remove when the underlying Calcite assertion is fixed.

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@
8484
import org.opensearch.sql.ast.tree.SubqueryAlias;
8585
import org.opensearch.sql.ast.tree.TableFunction;
8686
import org.opensearch.sql.ast.tree.Trendline;
87+
import org.opensearch.sql.ast.tree.Union;
8788
import org.opensearch.sql.ast.tree.UnresolvedPlan;
8889
import org.opensearch.sql.ast.tree.Values;
8990
import org.opensearch.sql.calcite.plan.OpenSearchConstants;
@@ -781,4 +782,8 @@ public static InSubquery inSubquery(UnresolvedPlan query, UnresolvedExpression..
781782
public static ExistsSubquery existsSubquery(UnresolvedPlan query) {
782783
return new ExistsSubquery(query);
783784
}
785+
786+
public static Union union(List<UnresolvedPlan> datasets) {
787+
return new Union(datasets);
788+
}
784789
}

‎core/src/main/java/org/opensearch/sql/ast/expression/WindowBound.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
package org.opensearch.sql.ast.expression;
77

8+
import lombok.EqualsAndHashCode;
89
import lombok.Getter;
910

1011
public abstract class WindowBound {
@@ -25,6 +26,7 @@ public boolean isPreceding() {
2526
}
2627
}
2728

29+
@EqualsAndHashCode(callSuper = false)
2830
public static class CurrentRowWindowBound extends WindowBound {
2931
CurrentRowWindowBound() {}
3032

‎core/src/main/java/org/opensearch/sql/ast/expression/WindowFrame.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,13 @@ public static WindowFrame toCurrentRow() {
4040
AstDSL.stringLiteral("CURRENT ROW"));
4141
}
4242

43+
public static WindowFrame rangeToCurrentRow() {
44+
return WindowFrame.of(
45+
FrameType.RANGE,
46+
AstDSL.stringLiteral("UNBOUNDED PRECEDING"),
47+
AstDSL.stringLiteral("CURRENT ROW"));
48+
}
49+
4350
public static WindowFrame of(FrameType type, String lower, String upper) {
4451
return WindowFrame.of(type, AstDSL.stringLiteral(lower), AstDSL.stringLiteral(upper));
4552
}

‎core/src/main/java/org/opensearch/sql/ast/tree/Union.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,15 +21,25 @@
2121
@RequiredArgsConstructor
2222
@AllArgsConstructor
2323
public class Union extends UnresolvedPlan {
24+
/** Input subplans (operands) combined by this UNION ALL. */
2425
private final List<UnresolvedPlan> datasets;
2526

27+
/** Whether inputs are unified to a common schema by name (PPL) vs combined positionally (SQL). */
28+
private boolean unifySchema;
29+
30+
/** Optional cap on output rows (PPL {@code maxout}); {@code null} if unbounded. */
2631
private Integer maxout;
2732

33+
/** PPL constructor: UNION ALL with schema unification. */
34+
public Union(List<UnresolvedPlan> datasets, Integer maxout) {
35+
this(datasets, true, maxout);
36+
}
37+
2838
@Override
2939
public UnresolvedPlan attach(UnresolvedPlan child) {
3040
List<UnresolvedPlan> newDatasets =
3141
ImmutableList.<UnresolvedPlan>builder().add(child).addAll(datasets).build();
32-
return new Union(newDatasets, maxout);
42+
return new Union(newDatasets, unifySchema, maxout);
3343
}
3444

3545
@Override

0 commit comments

Comments
 (0)