Skip to content

Commit 53958e8

Browse files
dai-chendai-chen
committed
fix: window functions with ORDER BY/LIMIT on unified SQL path
Scope the fix to the unified query API parser (ExtendedAstBuilder); fixing the shared AstBuilder directly would break the V2/legacy engine. Signed-off-by: Chen Dai <daichen@amazon.com>
1 parent cc65d75 commit 53958e8

4 files changed

Lines changed: 159 additions & 4 deletions

File tree

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

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,25 +10,36 @@
1010
import static org.opensearch.sql.ast.dsl.AstDSL.join;
1111
import static org.opensearch.sql.ast.dsl.AstDSL.union;
1212

13+
import java.util.ArrayList;
1314
import java.util.List;
1415
import java.util.Optional;
1516
import java.util.stream.Collectors;
1617
import org.antlr.v4.runtime.tree.ParseTree;
18+
import org.opensearch.sql.ast.expression.Field;
1719
import org.opensearch.sql.ast.expression.Not;
1820
import org.opensearch.sql.ast.expression.UnresolvedExpression;
21+
import org.opensearch.sql.ast.expression.WindowFunction;
1922
import org.opensearch.sql.ast.statement.Query;
2023
import org.opensearch.sql.ast.statement.Statement;
2124
import org.opensearch.sql.ast.tree.Join.JoinType;
25+
import org.opensearch.sql.ast.tree.Project;
26+
import org.opensearch.sql.ast.tree.Sort;
27+
import org.opensearch.sql.ast.tree.Sort.SortOption;
2228
import org.opensearch.sql.ast.tree.UnresolvedPlan;
2329
import org.opensearch.sql.sql.antlr.SQLSyntaxParser;
2430
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser;
2531
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ExistsSubqueryExpressionAtomContext;
32+
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.FromClauseContext;
2633
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.InSubqueryPredicateContext;
2734
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.JoinClauseContext;
35+
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.OrderByClauseContext;
36+
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.QuerySpecificationContext;
2837
import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.UnionSelectContext;
2938
import org.opensearch.sql.sql.parser.AstBuilder;
3039
import org.opensearch.sql.sql.parser.AstExpressionBuilder;
40+
import org.opensearch.sql.sql.parser.AstSortBuilder;
3141
import org.opensearch.sql.sql.parser.AstStatementBuilder;
42+
import org.opensearch.sql.sql.parser.context.QuerySpecification;
3243

3344
/** SQL query parser that produces {@link UnresolvedPlan} using the V2 ANTLR grammar. */
3445
public class SqlV2QueryParser implements UnifiedQueryParser<UnresolvedPlan> {
@@ -62,6 +73,53 @@ private static class ExtendedAstBuilder extends AstBuilder {
6273
super(query);
6374
}
6475

76+
@Override
77+
public UnresolvedPlan visitQuerySpecification(QuerySpecificationContext queryContext) {
78+
if (!hasWindowFunctionInProjectList(queryContext)) {
79+
return super.visitQuerySpecification(queryContext);
80+
}
81+
82+
context.push();
83+
context.peek().collect(queryContext, query);
84+
Project project = (Project) visit(queryContext.selectClause());
85+
UnresolvedPlan result = project.attach(visit(queryContext.fromClause()));
86+
87+
// Window output must be computed before ORDER BY/LIMIT, so build Limit(Sort(Project(from)))
88+
OrderByClauseContext orderByClause = queryContext.fromClause().orderByClause();
89+
if (orderByClause != null) {
90+
result = new ExtendedAstSortBuilder(context.peek()).visit(orderByClause).attach(result);
91+
}
92+
if (queryContext.limitClause() != null) {
93+
result = visit(queryContext.limitClause()).attach(result);
94+
}
95+
96+
context.pop();
97+
return result;
98+
}
99+
100+
@Override
101+
public UnresolvedPlan visitFromClause(FromClauseContext ctx) {
102+
UnresolvedPlan from = super.visitFromClause(ctx);
103+
if (hasWindowFunctionInProjectList(context.peek()) && from instanceof Sort sort) {
104+
// Drop the ORDER BY Sort for window queries; it is re-attached above the Project
105+
return sort.getChild().get(0);
106+
}
107+
return from;
108+
}
109+
110+
private boolean hasWindowFunctionInProjectList(QuerySpecificationContext queryContext) {
111+
if (queryContext.fromClause() == null) {
112+
return false;
113+
}
114+
QuerySpecification probe = new QuerySpecification();
115+
probe.collect(queryContext, query);
116+
return hasWindowFunctionInProjectList(probe);
117+
}
118+
119+
private static boolean hasWindowFunctionInProjectList(QuerySpecification querySpec) {
120+
return querySpec.getSelectItems().stream().anyMatch(item -> item instanceof WindowFunction);
121+
}
122+
65123
@Override
66124
protected AstExpressionBuilder createExpressionBuilder() {
67125
return new ExtendedAstExpressionBuilder();
@@ -114,4 +172,32 @@ public UnresolvedExpression visitExistsSubqueryExpressionAtom(
114172
}
115173
}
116174
}
175+
176+
/**
177+
* Keeps an ORDER BY window-alias as a column reference (Sort is above the Project) to avoid a
178+
* second RexOver.
179+
*/
180+
private static class ExtendedAstSortBuilder extends AstSortBuilder {
181+
182+
ExtendedAstSortBuilder(QuerySpecification querySpec) {
183+
super(querySpec);
184+
}
185+
186+
@Override
187+
public UnresolvedPlan visitOrderByClause(OrderByClauseContext ctx) {
188+
List<Field> fields = new ArrayList<>();
189+
List<UnresolvedExpression> items = querySpec.getOrderByItems();
190+
List<SortOption> options = querySpec.getOrderByOptions();
191+
for (int i = 0; i < items.size(); i++) {
192+
UnresolvedExpression item = items.get(i);
193+
UnresolvedExpression sortKey =
194+
(querySpec.isSelectAlias(item)
195+
&& querySpec.getSelectItemByAlias(item) instanceof WindowFunction)
196+
? item
197+
: querySpec.replaceIfAliasOrOrdinal(item);
198+
fields.add(new Field(sortKey, createSortArguments(options.get(i))));
199+
}
200+
return new Sort(fields);
201+
}
202+
}
117203
}

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -547,4 +547,73 @@ SELECT LENGTH(name) FROM catalog.employees GROUP BY LENGTH(name) ORDER BY LENGTH
547547
LogicalTableScan(table=[[catalog, employees]])
548548
""");
549549
}
550+
551+
@Test
552+
public void testWindowOverGroupByWithLimit() {
553+
givenQuery(
554+
"""
555+
SELECT department, COUNT(*) AS cnt, ROW_NUMBER() OVER (ORDER BY COUNT(*) DESC) AS rn
556+
FROM catalog.employees GROUP BY department LIMIT 3
557+
""")
558+
.assertPlan(
559+
"""
560+
LogicalSort(fetch=[3])
561+
LogicalProject(department=[$0], cnt=[$1], rn=[ROW_NUMBER() OVER (ORDER BY $1 DESC NULLS FIRST)])
562+
LogicalAggregate(group=[{0}], COUNT(*)=[COUNT()])
563+
LogicalProject(department=[$3])
564+
LogicalTableScan(table=[[catalog, employees]])
565+
""");
566+
}
567+
568+
@Test
569+
public void testWindowOverGroupByOrderByWindowAlias() {
570+
givenQuery(
571+
"""
572+
SELECT department, COUNT(*) AS cnt, ROW_NUMBER() OVER (ORDER BY COUNT(*) DESC) AS rn
573+
FROM catalog.employees GROUP BY department ORDER BY rn LIMIT 3
574+
""")
575+
.assertPlan(
576+
"""
577+
LogicalSort(sort0=[$2], dir0=[ASC-nulls-first], fetch=[3])
578+
LogicalProject(department=[$0], cnt=[$1], rn=[ROW_NUMBER() OVER (ORDER BY $1 DESC NULLS FIRST)])
579+
LogicalAggregate(group=[{0}], COUNT(*)=[COUNT()])
580+
LogicalProject(department=[$3])
581+
LogicalTableScan(table=[[catalog, employees]])
582+
""");
583+
}
584+
585+
@Test
586+
public void testWindowOverGroupByOrderByWindowAliasWithoutLimit() {
587+
givenQuery(
588+
"""
589+
SELECT department, COUNT(*) AS cnt, ROW_NUMBER() OVER (ORDER BY COUNT(*) DESC) AS rn
590+
FROM catalog.employees GROUP BY department ORDER BY rn
591+
""")
592+
.assertPlan(
593+
"""
594+
LogicalSort(sort0=[$2], dir0=[ASC-nulls-first])
595+
LogicalProject(department=[$0], cnt=[$1], rn=[ROW_NUMBER() OVER (ORDER BY $1 DESC NULLS FIRST)])
596+
LogicalAggregate(group=[{0}], COUNT(*)=[COUNT()])
597+
LogicalProject(department=[$3])
598+
LogicalTableScan(table=[[catalog, employees]])
599+
""");
600+
}
601+
602+
@Test
603+
public void testMultipleWindowFunctionsOrderByWindowAlias() {
604+
givenQuery(
605+
"""
606+
SELECT department, COUNT(*) AS cnt, ROW_NUMBER() OVER (ORDER BY COUNT(*) DESC) AS rn,
607+
ROW_NUMBER() OVER (ORDER BY department) AS rn2
608+
FROM catalog.employees GROUP BY department ORDER BY rn LIMIT 3
609+
""")
610+
.assertPlan(
611+
"""
612+
LogicalSort(sort0=[$2], dir0=[ASC-nulls-first], fetch=[3])
613+
LogicalProject(department=[$0], cnt=[$1], rn=[ROW_NUMBER() OVER (ORDER BY $1 DESC NULLS FIRST)], rn2=[ROW_NUMBER() OVER (ORDER BY $0 NULLS FIRST)])
614+
LogicalAggregate(group=[{0}], COUNT(*)=[COUNT()])
615+
LogicalProject(department=[$3])
616+
LogicalTableScan(table=[[catalog, employees]])
617+
""");
618+
}
550619
}

‎sql/src/main/java/org/opensearch/sql/sql/parser/AstBuilder.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,13 @@ public class AstBuilder extends OpenSearchSQLParserBaseVisitor<UnresolvedPlan> {
5454
private final AstExpressionBuilder expressionBuilder;
5555

5656
/** Parsing context stack that contains context for current query parsing. */
57-
private final ParsingContext context = new ParsingContext();
57+
protected final ParsingContext context = new ParsingContext();
5858

5959
/**
6060
* SQL query to get original token text. This is necessary because token.getText() returns text
6161
* without whitespaces or other characters discarded by lexer.
6262
*/
63-
private final String query;
63+
protected final String query;
6464

6565
public AstBuilder(String query) {
6666
this.query = query;

‎sql/src/main/java/org/opensearch/sql/sql/parser/AstSortBuilder.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
@RequiredArgsConstructor
3434
public class AstSortBuilder extends OpenSearchSQLParserBaseVisitor<UnresolvedPlan> {
3535

36-
private final QuerySpecification querySpec;
36+
protected final QuerySpecification querySpec;
3737

3838
@Override
3939
public UnresolvedPlan visitOrderByClause(OrderByClauseContext ctx) {
@@ -57,7 +57,7 @@ private List<Field> createSortFields() {
5757
* Argument "asc" is required. Argument "nullFirst" is optional and determined by Analyzer later
5858
* if absent.
5959
*/
60-
private List<Argument> createSortArguments(SortOption option) {
60+
protected List<Argument> createSortArguments(SortOption option) {
6161
SortOrder sortOrder = option.getSortOrder();
6262
NullOrder nullOrder = option.getNullOrder();
6363
ImmutableList.Builder<Argument> args = ImmutableList.builder();

0 commit comments

Comments
 (0)