Skip to content
This repository was archived by the owner on Aug 2, 2022. It is now read-only.

Commit d419e8e

Browse files
dai-chendai-chen
authored
Support NULLS FIRST/LAST ordering for window functions (#929)
* Skip sort items in window functions * Use sort option in window function AST node * Analyze sort option in window function analyzer * Add more UT * Add IT * Add doc test
1 parent 0855c53 commit d419e8e

File tree

12 files changed

+226
-43
lines changed

12 files changed

+226
-43
lines changed

‎core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzer.java

+15-4
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818

1919
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption.DEFAULT_ASC;
2020
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption.DEFAULT_DESC;
21+
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.ASC;
22+
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.DESC;
2123

2224
import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor;
2325
import com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias;
@@ -77,7 +79,7 @@ public LogicalPlan visitWindowFunction(WindowFunction node, AnalysisContext cont
7779
WindowDefinition windowDefinition = new WindowDefinition(partitionByList, sortList);
7880

7981
return new LogicalWindow(
80-
new LogicalSort(child,windowDefinition.getAllSortItems()),
82+
new LogicalSort(child, windowDefinition.getAllSortItems()),
8183
windowFunction,
8284
windowDefinition);
8385
}
@@ -94,13 +96,22 @@ private List<Pair<SortOption, Expression>> analyzeSortList(WindowFunction node,
9496
return node.getSortList()
9597
.stream()
9698
.map(pair -> ImmutablePair
97-
.of(getSortOption(pair.getLeft()),
99+
.of(analyzeSortOption(pair.getLeft()),
98100
expressionAnalyzer.analyze(pair.getRight(), context)))
99101
.collect(Collectors.toList());
100102
}
101103

102-
private SortOption getSortOption(String option) {
103-
return "ASC".equalsIgnoreCase(option) ? DEFAULT_ASC : DEFAULT_DESC;
104+
/**
105+
* Frontend creates sort option from query directly which means sort or null order may be null.
106+
* The final and default value for each is determined here during expression analysis.
107+
*/
108+
private SortOption analyzeSortOption(SortOption option) {
109+
if (option.getNullOrder() == null) {
110+
return (option.getSortOrder() == DESC) ? DEFAULT_DESC : DEFAULT_ASC;
111+
}
112+
return new SortOption(
113+
(option.getSortOrder() == DESC) ? DESC : ASC,
114+
option.getNullOrder());
104115
}
105116

106117
}

‎core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/dsl/AstDSL.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
import com.amazon.opendistroforelasticsearch.sql.ast.tree.RelationSubquery;
5555
import com.amazon.opendistroforelasticsearch.sql.ast.tree.Rename;
5656
import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort;
57+
import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption;
5758
import com.amazon.opendistroforelasticsearch.sql.ast.tree.UnresolvedPlan;
5859
import com.amazon.opendistroforelasticsearch.sql.ast.tree.Values;
5960
import java.util.Arrays;
@@ -232,7 +233,7 @@ public When when(UnresolvedExpression condition, UnresolvedExpression result) {
232233

233234
public UnresolvedExpression window(Function function,
234235
List<UnresolvedExpression> partitionByList,
235-
List<Pair<String, UnresolvedExpression>> sortList) {
236+
List<Pair<SortOption, UnresolvedExpression>> sortList) {
236237
return new WindowFunction(function, partitionByList, sortList);
237238
}
238239

‎core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/expression/WindowFunction.java

+4-1
Original file line numberDiff line numberDiff line change
@@ -18,23 +18,26 @@
1818

1919
import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor;
2020
import com.amazon.opendistroforelasticsearch.sql.ast.Node;
21+
import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption;
2122
import java.util.Collections;
2223
import java.util.List;
2324
import lombok.AllArgsConstructor;
2425
import lombok.EqualsAndHashCode;
2526
import lombok.Getter;
2627
import lombok.RequiredArgsConstructor;
28+
import lombok.ToString;
2729
import org.apache.commons.lang3.tuple.Pair;
2830

2931
@AllArgsConstructor
3032
@EqualsAndHashCode(callSuper = false)
3133
@Getter
3234
@RequiredArgsConstructor
35+
@ToString
3336
public class WindowFunction extends UnresolvedExpression {
3437

3538
private final Function function;
3639
private List<UnresolvedExpression> partitionByList;
37-
private List<Pair<String, UnresolvedExpression>> sortList;
40+
private List<Pair<SortOption, UnresolvedExpression>> sortList;
3841

3942
@Override
4043
public List<? extends Node> getChild() {

‎core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/AnalyzerTest.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -385,7 +385,7 @@ public void window_function() {
385385
AstDSL.function("row_number"),
386386
Collections.singletonList(AstDSL.qualifiedName("string_value")),
387387
Collections.singletonList(
388-
ImmutablePair.of("ASC", AstDSL.qualifiedName("integer_value")))))));
388+
ImmutablePair.of(DEFAULT_ASC, AstDSL.qualifiedName("integer_value")))))));
389389
}
390390

391391
/**

‎core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzerTest.java

+41-1
Original file line numberDiff line numberDiff line change
@@ -16,20 +16,29 @@
1616

1717
package com.amazon.opendistroforelasticsearch.sql.analysis;
1818

19+
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST;
20+
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.NullOrder.NULL_LAST;
1921
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption.DEFAULT_ASC;
2022
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption.DEFAULT_DESC;
23+
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.ASC;
24+
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.DESC;
2125
import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.INTEGER;
2226
import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.STRING;
2327
import static org.junit.jupiter.api.Assertions.assertEquals;
2428

2529
import com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL;
30+
import com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias;
31+
import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption;
2632
import com.amazon.opendistroforelasticsearch.sql.expression.DSL;
2733
import com.amazon.opendistroforelasticsearch.sql.expression.config.ExpressionConfig;
2834
import com.amazon.opendistroforelasticsearch.sql.expression.window.WindowDefinition;
2935
import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlan;
3036
import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlanDSL;
3137
import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRelation;
38+
import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalSort;
3239
import com.google.common.collect.ImmutableList;
40+
import com.google.common.collect.ImmutableMap;
41+
import java.util.Collections;
3342
import org.apache.commons.lang3.tuple.ImmutablePair;
3443
import org.junit.jupiter.api.BeforeEach;
3544
import org.junit.jupiter.api.DisplayNameGeneration;
@@ -76,7 +85,7 @@ void should_wrap_child_with_window_and_sort_operator_if_project_item_windowed()
7685
AstDSL.function("row_number"),
7786
ImmutableList.of(AstDSL.qualifiedName("string_value")),
7887
ImmutableList.of(
79-
ImmutablePair.of("DESC", AstDSL.qualifiedName("integer_value"))))),
88+
ImmutablePair.of(DEFAULT_DESC, AstDSL.qualifiedName("integer_value"))))),
8089
analysisContext));
8190
}
8291

@@ -91,4 +100,35 @@ void should_return_original_child_if_project_item_not_windowed() {
91100
analysisContext));
92101
}
93102

103+
@Test
104+
void can_analyze_sort_options() {
105+
// Mapping from input option to expected option after analysis
106+
ImmutableMap<SortOption, SortOption> expects =
107+
ImmutableMap.<SortOption, SortOption>builder()
108+
.put(new SortOption(null, null), DEFAULT_ASC)
109+
.put(new SortOption(ASC, null), DEFAULT_ASC)
110+
.put(new SortOption(DESC, null), DEFAULT_DESC)
111+
.put(new SortOption(null, NULL_FIRST), DEFAULT_ASC)
112+
.put(new SortOption(null, NULL_LAST), new SortOption(ASC, NULL_LAST))
113+
.put(new SortOption(ASC, NULL_FIRST), DEFAULT_ASC)
114+
.put(new SortOption(DESC, NULL_FIRST), new SortOption(DESC, NULL_FIRST))
115+
.put(new SortOption(DESC, NULL_LAST), DEFAULT_DESC)
116+
.build();
117+
118+
expects.forEach((option, expect) -> {
119+
Alias ast = AstDSL.alias(
120+
"row_number",
121+
AstDSL.window(
122+
AstDSL.function("row_number"),
123+
Collections.emptyList(),
124+
ImmutableList.of(
125+
ImmutablePair.of(option, AstDSL.qualifiedName("integer_value")))));
126+
127+
LogicalPlan plan = analyzer.analyze(ast, analysisContext);
128+
LogicalSort sort = (LogicalSort) plan.getChild().get(0);
129+
assertEquals(expect, sort.getSortList().get(0).getLeft(),
130+
"Assertion failed on input option: " + option);
131+
});
132+
}
133+
94134
}

‎docs/user/dql/window.rst

+19-1
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ The syntax of a window function is as follows in which both ``PARTITION BY`` and
3232
function_name (expression [, expression...])
3333
OVER (
3434
PARTITION BY expression [, expression...]
35-
ORDER BY expression [ASC | DESC] [, ...]
35+
ORDER BY expression [ASC | DESC] [NULLS {FIRST | LAST}] [, ...]
3636
)
3737

3838

@@ -62,6 +62,24 @@ ROW_NUMBER
6262
| M | 39225 | 3 |
6363
+----------+-----------+-------+
6464

65+
Similarly as regular ``ORDER BY`` clause, you can specify null ordering by ``NULLS FIRST`` or ``NULLS LAST`` which has exactly same behavior::
66+
67+
od> SELECT
68+
... employer,
69+
... ROW_NUMBER() OVER(
70+
... ORDER BY employer NULLS LAST
71+
... ) AS num
72+
... FROM accounts
73+
... ORDER BY employer NULLS LAST;
74+
fetched rows / total rows = 4/4
75+
+------------+-------+
76+
| employer | num |
77+
|------------+-------|
78+
| Netagy | 1 |
79+
| Pyrami | 2 |
80+
| Quility | 3 |
81+
| null | 4 |
82+
+------------+-------+
6583

6684
RANK
6785
----

‎integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/WindowFunctionIT.java

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*
15+
*/
16+
17+
package com.amazon.opendistroforelasticsearch.sql.sql;
18+
19+
import static com.amazon.opendistroforelasticsearch.sql.util.MatcherUtils.rows;
20+
import static com.amazon.opendistroforelasticsearch.sql.util.MatcherUtils.verifyDataRows;
21+
22+
import com.amazon.opendistroforelasticsearch.sql.legacy.SQLIntegTestCase;
23+
import com.amazon.opendistroforelasticsearch.sql.legacy.TestsConstants;
24+
import org.json.JSONObject;
25+
import org.junit.Test;
26+
27+
public class WindowFunctionIT extends SQLIntegTestCase {
28+
29+
@Override
30+
protected void init() throws Exception {
31+
loadIndex(Index.BANK_WITH_NULL_VALUES);
32+
}
33+
34+
@Test
35+
public void testOrderByNullFirst() {
36+
JSONObject response = new JSONObject(
37+
executeQuery("SELECT age, ROW_NUMBER() OVER(ORDER BY age DESC NULLS FIRST) "
38+
+ "FROM " + TestsConstants.TEST_INDEX_BANK_WITH_NULL_VALUES, "jdbc"));
39+
40+
verifyDataRows(response,
41+
rows(null, 1),
42+
rows(36, 2),
43+
rows(36, 3),
44+
rows(34, 4),
45+
rows(33, 5),
46+
rows(32, 6),
47+
rows(28, 7));
48+
}
49+
50+
@Test
51+
public void testOrderByNullLast() {
52+
JSONObject response = new JSONObject(
53+
executeQuery("SELECT age, ROW_NUMBER() OVER(ORDER BY age NULLS LAST) "
54+
+ "FROM " + TestsConstants.TEST_INDEX_BANK_WITH_NULL_VALUES, "jdbc"));
55+
56+
verifyDataRows(response,
57+
rows(28, 1),
58+
rows(32, 2),
59+
rows(33, 3),
60+
rows(34, 4),
61+
rows(36, 5),
62+
rows(36, 6),
63+
rows(null, 7));
64+
}
65+
66+
}

‎sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/AstExpressionBuilder.java

+5-7
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.MathExpressionAtomContext;
3535
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.NotExpressionContext;
3636
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.NullLiteralContext;
37-
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.OrderByElementContext;
3837
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.OverClauseContext;
3938
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.QualifiedNameContext;
4039
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.RankingWindowFunctionContext;
@@ -47,6 +46,7 @@
4746
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.TimeLiteralContext;
4847
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.TimestampLiteralContext;
4948
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.WindowFunctionContext;
49+
import static com.amazon.opendistroforelasticsearch.sql.sql.parser.ParserUtils.createSortOption;
5050

5151
import com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL;
5252
import com.amazon.opendistroforelasticsearch.sql.ast.expression.AggregateFunction;
@@ -63,6 +63,7 @@
6363
import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression;
6464
import com.amazon.opendistroforelasticsearch.sql.ast.expression.When;
6565
import com.amazon.opendistroforelasticsearch.sql.ast.expression.WindowFunction;
66+
import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption;
6667
import com.amazon.opendistroforelasticsearch.sql.common.utils.StringUtils;
6768
import com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser;
6869
import com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.AndExpressionContext;
@@ -171,12 +172,13 @@ public UnresolvedExpression visitWindowFunction(WindowFunctionContext ctx) {
171172
.collect(Collectors.toList());
172173
}
173174

174-
List<Pair<String, UnresolvedExpression>> sortList = Collections.emptyList();
175+
List<Pair<SortOption, UnresolvedExpression>> sortList = Collections.emptyList();
175176
if (overClause.orderByClause() != null) {
176177
sortList = overClause.orderByClause()
177178
.orderByElement()
178179
.stream()
179-
.map(item -> ImmutablePair.of(getOrder(item), visit(item.expression())))
180+
.map(item -> ImmutablePair.of(
181+
createSortOption(item), visit(item.expression())))
180182
.collect(Collectors.toList());
181183
}
182184
return new WindowFunction((Function) visit(ctx.function), partitionByList, sortList);
@@ -337,8 +339,4 @@ private QualifiedName visitIdentifiers(List<IdentContext> identifiers) {
337339
);
338340
}
339341

340-
private String getOrder(OrderByElementContext item) {
341-
return (item.order == null) ? "ASC" : item.order.getText();
342-
}
343-
344342
}

‎sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/ParserUtils.java

+39
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,15 @@
1717

1818
package com.amazon.opendistroforelasticsearch.sql.sql.parser;
1919

20+
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.NullOrder;
21+
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption;
22+
import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder;
23+
import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.OrderByElementContext;
24+
2025
import lombok.experimental.UtilityClass;
2126
import org.antlr.v4.runtime.ParserRuleContext;
2227
import org.antlr.v4.runtime.Token;
28+
import org.antlr.v4.runtime.tree.TerminalNode;
2329

2430
/**
2531
* Parser Utils Class.
@@ -35,4 +41,37 @@ public static String getTextInQuery(ParserRuleContext ctx, String queryString) {
3541
Token stop = ctx.getStop();
3642
return queryString.substring(start.getStartIndex(), stop.getStopIndex() + 1);
3743
}
44+
45+
/**
46+
* Create sort option from syntax tree node.
47+
*/
48+
public static SortOption createSortOption(OrderByElementContext orderBy) {
49+
return new SortOption(
50+
createSortOrder(orderBy.order),
51+
createNullOrder(orderBy.FIRST(), orderBy.LAST()));
52+
}
53+
54+
/**
55+
* Create sort order for sort option use from ASC/DESC token.
56+
*/
57+
public static SortOrder createSortOrder(Token ctx) {
58+
if (ctx == null) {
59+
return null;
60+
}
61+
return SortOrder.valueOf(ctx.getText().toUpperCase());
62+
}
63+
64+
/**
65+
* Create null order for sort option use from FIRST/LAST token.
66+
*/
67+
public static NullOrder createNullOrder(TerminalNode first, TerminalNode last) {
68+
if (first != null) {
69+
return NullOrder.NULL_FIRST;
70+
} else if (last != null) {
71+
return NullOrder.NULL_LAST;
72+
} else {
73+
return null;
74+
}
75+
}
76+
3877
}

0 commit comments

Comments
 (0)