From d419e8ebe3edd1e578de7b0277afb138d96e2b9a Mon Sep 17 00:00:00 2001 From: Chen Dai <46505291+dai-chen@users.noreply.github.com> Date: Tue, 5 Jan 2021 09:55:35 -0800 Subject: [PATCH] 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 --- .../analysis/WindowExpressionAnalyzer.java | 19 ++++-- .../sql/ast/dsl/AstDSL.java | 3 +- .../sql/ast/expression/WindowFunction.java | 5 +- .../sql/analysis/AnalyzerTest.java | 2 +- .../WindowExpressionAnalyzerTest.java | 42 +++++++++++- docs/user/dql/window.rst | 20 +++++- .../sql/sql/WindowFunctionIT.java | 66 +++++++++++++++++++ .../sql/sql/parser/AstExpressionBuilder.java | 12 ++-- .../sql/sql/parser/ParserUtils.java | 39 +++++++++++ .../parser/context/QuerySpecification.java | 34 +++------- .../sql/parser/AstExpressionBuilderTest.java | 19 +++++- .../context/QuerySpecificationTest.java | 8 +++ 12 files changed, 226 insertions(+), 43 deletions(-) create mode 100644 integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/WindowFunctionIT.java diff --git a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzer.java b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzer.java index 67440045ca..d5fbe1b19b 100644 --- a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzer.java +++ b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzer.java @@ -18,6 +18,8 @@ import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption.DEFAULT_ASC; import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption.DEFAULT_DESC; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.ASC; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.DESC; import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor; import com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias; @@ -77,7 +79,7 @@ public LogicalPlan visitWindowFunction(WindowFunction node, AnalysisContext cont WindowDefinition windowDefinition = new WindowDefinition(partitionByList, sortList); return new LogicalWindow( - new LogicalSort(child,windowDefinition.getAllSortItems()), + new LogicalSort(child, windowDefinition.getAllSortItems()), windowFunction, windowDefinition); } @@ -94,13 +96,22 @@ private List> analyzeSortList(WindowFunction node, return node.getSortList() .stream() .map(pair -> ImmutablePair - .of(getSortOption(pair.getLeft()), + .of(analyzeSortOption(pair.getLeft()), expressionAnalyzer.analyze(pair.getRight(), context))) .collect(Collectors.toList()); } - private SortOption getSortOption(String option) { - return "ASC".equalsIgnoreCase(option) ? DEFAULT_ASC : DEFAULT_DESC; + /** + * Frontend creates sort option from query directly which means sort or null order may be null. + * The final and default value for each is determined here during expression analysis. + */ + private SortOption analyzeSortOption(SortOption option) { + if (option.getNullOrder() == null) { + return (option.getSortOrder() == DESC) ? DEFAULT_DESC : DEFAULT_ASC; + } + return new SortOption( + (option.getSortOrder() == DESC) ? DESC : ASC, + option.getNullOrder()); } } diff --git a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/dsl/AstDSL.java b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/dsl/AstDSL.java index fc0e79622f..ae5ae3ab55 100644 --- a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/dsl/AstDSL.java +++ b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/dsl/AstDSL.java @@ -54,6 +54,7 @@ import com.amazon.opendistroforelasticsearch.sql.ast.tree.RelationSubquery; import com.amazon.opendistroforelasticsearch.sql.ast.tree.Rename; import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort; +import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption; import com.amazon.opendistroforelasticsearch.sql.ast.tree.UnresolvedPlan; import com.amazon.opendistroforelasticsearch.sql.ast.tree.Values; import java.util.Arrays; @@ -232,7 +233,7 @@ public When when(UnresolvedExpression condition, UnresolvedExpression result) { public UnresolvedExpression window(Function function, List partitionByList, - List> sortList) { + List> sortList) { return new WindowFunction(function, partitionByList, sortList); } diff --git a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/expression/WindowFunction.java b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/expression/WindowFunction.java index 2e323f937f..976be0c48f 100644 --- a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/expression/WindowFunction.java +++ b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/expression/WindowFunction.java @@ -18,23 +18,26 @@ import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor; import com.amazon.opendistroforelasticsearch.sql.ast.Node; +import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption; import java.util.Collections; import java.util.List; import lombok.AllArgsConstructor; import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.RequiredArgsConstructor; +import lombok.ToString; import org.apache.commons.lang3.tuple.Pair; @AllArgsConstructor @EqualsAndHashCode(callSuper = false) @Getter @RequiredArgsConstructor +@ToString public class WindowFunction extends UnresolvedExpression { private final Function function; private List partitionByList; - private List> sortList; + private List> sortList; @Override public List getChild() { diff --git a/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/AnalyzerTest.java b/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/AnalyzerTest.java index 8b0f3cebf5..b8cd8ede2f 100644 --- a/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/AnalyzerTest.java +++ b/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/AnalyzerTest.java @@ -385,7 +385,7 @@ public void window_function() { AstDSL.function("row_number"), Collections.singletonList(AstDSL.qualifiedName("string_value")), Collections.singletonList( - ImmutablePair.of("ASC", AstDSL.qualifiedName("integer_value"))))))); + ImmutablePair.of(DEFAULT_ASC, AstDSL.qualifiedName("integer_value"))))))); } /** diff --git a/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzerTest.java b/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzerTest.java index a13be922ce..3292690b9a 100644 --- a/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzerTest.java +++ b/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/WindowExpressionAnalyzerTest.java @@ -16,20 +16,29 @@ package com.amazon.opendistroforelasticsearch.sql.analysis; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.NullOrder.NULL_LAST; import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption.DEFAULT_ASC; import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption.DEFAULT_DESC; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.ASC; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.DESC; import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.INTEGER; import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.STRING; import static org.junit.jupiter.api.Assertions.assertEquals; import com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL; +import com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias; +import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption; import com.amazon.opendistroforelasticsearch.sql.expression.DSL; import com.amazon.opendistroforelasticsearch.sql.expression.config.ExpressionConfig; import com.amazon.opendistroforelasticsearch.sql.expression.window.WindowDefinition; import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlan; import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlanDSL; import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRelation; +import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalSort; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import java.util.Collections; import org.apache.commons.lang3.tuple.ImmutablePair; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; @@ -76,7 +85,7 @@ void should_wrap_child_with_window_and_sort_operator_if_project_item_windowed() AstDSL.function("row_number"), ImmutableList.of(AstDSL.qualifiedName("string_value")), ImmutableList.of( - ImmutablePair.of("DESC", AstDSL.qualifiedName("integer_value"))))), + ImmutablePair.of(DEFAULT_DESC, AstDSL.qualifiedName("integer_value"))))), analysisContext)); } @@ -91,4 +100,35 @@ void should_return_original_child_if_project_item_not_windowed() { analysisContext)); } + @Test + void can_analyze_sort_options() { + // Mapping from input option to expected option after analysis + ImmutableMap expects = + ImmutableMap.builder() + .put(new SortOption(null, null), DEFAULT_ASC) + .put(new SortOption(ASC, null), DEFAULT_ASC) + .put(new SortOption(DESC, null), DEFAULT_DESC) + .put(new SortOption(null, NULL_FIRST), DEFAULT_ASC) + .put(new SortOption(null, NULL_LAST), new SortOption(ASC, NULL_LAST)) + .put(new SortOption(ASC, NULL_FIRST), DEFAULT_ASC) + .put(new SortOption(DESC, NULL_FIRST), new SortOption(DESC, NULL_FIRST)) + .put(new SortOption(DESC, NULL_LAST), DEFAULT_DESC) + .build(); + + expects.forEach((option, expect) -> { + Alias ast = AstDSL.alias( + "row_number", + AstDSL.window( + AstDSL.function("row_number"), + Collections.emptyList(), + ImmutableList.of( + ImmutablePair.of(option, AstDSL.qualifiedName("integer_value"))))); + + LogicalPlan plan = analyzer.analyze(ast, analysisContext); + LogicalSort sort = (LogicalSort) plan.getChild().get(0); + assertEquals(expect, sort.getSortList().get(0).getLeft(), + "Assertion failed on input option: " + option); + }); + } + } \ No newline at end of file diff --git a/docs/user/dql/window.rst b/docs/user/dql/window.rst index 59b2f024cf..004e3a42df 100644 --- a/docs/user/dql/window.rst +++ b/docs/user/dql/window.rst @@ -32,7 +32,7 @@ The syntax of a window function is as follows in which both ``PARTITION BY`` and function_name (expression [, expression...]) OVER ( PARTITION BY expression [, expression...] - ORDER BY expression [ASC | DESC] [, ...] + ORDER BY expression [ASC | DESC] [NULLS {FIRST | LAST}] [, ...] ) @@ -62,6 +62,24 @@ ROW_NUMBER | M | 39225 | 3 | +----------+-----------+-------+ +Similarly as regular ``ORDER BY`` clause, you can specify null ordering by ``NULLS FIRST`` or ``NULLS LAST`` which has exactly same behavior:: + + od> SELECT + ... employer, + ... ROW_NUMBER() OVER( + ... ORDER BY employer NULLS LAST + ... ) AS num + ... FROM accounts + ... ORDER BY employer NULLS LAST; + fetched rows / total rows = 4/4 + +------------+-------+ + | employer | num | + |------------+-------| + | Netagy | 1 | + | Pyrami | 2 | + | Quility | 3 | + | null | 4 | + +------------+-------+ RANK ---- diff --git a/integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/WindowFunctionIT.java b/integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/WindowFunctionIT.java new file mode 100644 index 0000000000..b9f981971f --- /dev/null +++ b/integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/WindowFunctionIT.java @@ -0,0 +1,66 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + * + */ + +package com.amazon.opendistroforelasticsearch.sql.sql; + +import static com.amazon.opendistroforelasticsearch.sql.util.MatcherUtils.rows; +import static com.amazon.opendistroforelasticsearch.sql.util.MatcherUtils.verifyDataRows; + +import com.amazon.opendistroforelasticsearch.sql.legacy.SQLIntegTestCase; +import com.amazon.opendistroforelasticsearch.sql.legacy.TestsConstants; +import org.json.JSONObject; +import org.junit.Test; + +public class WindowFunctionIT extends SQLIntegTestCase { + + @Override + protected void init() throws Exception { + loadIndex(Index.BANK_WITH_NULL_VALUES); + } + + @Test + public void testOrderByNullFirst() { + JSONObject response = new JSONObject( + executeQuery("SELECT age, ROW_NUMBER() OVER(ORDER BY age DESC NULLS FIRST) " + + "FROM " + TestsConstants.TEST_INDEX_BANK_WITH_NULL_VALUES, "jdbc")); + + verifyDataRows(response, + rows(null, 1), + rows(36, 2), + rows(36, 3), + rows(34, 4), + rows(33, 5), + rows(32, 6), + rows(28, 7)); + } + + @Test + public void testOrderByNullLast() { + JSONObject response = new JSONObject( + executeQuery("SELECT age, ROW_NUMBER() OVER(ORDER BY age NULLS LAST) " + + "FROM " + TestsConstants.TEST_INDEX_BANK_WITH_NULL_VALUES, "jdbc")); + + verifyDataRows(response, + rows(28, 1), + rows(32, 2), + rows(33, 3), + rows(34, 4), + rows(36, 5), + rows(36, 6), + rows(null, 7)); + } + +} diff --git a/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/AstExpressionBuilder.java b/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/AstExpressionBuilder.java index 4c5947a223..197c1ce1b2 100644 --- a/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/AstExpressionBuilder.java +++ b/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/AstExpressionBuilder.java @@ -34,7 +34,6 @@ import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.MathExpressionAtomContext; import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.NotExpressionContext; import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.NullLiteralContext; -import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.OrderByElementContext; import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.OverClauseContext; import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.QualifiedNameContext; import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.RankingWindowFunctionContext; @@ -47,6 +46,7 @@ import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.TimeLiteralContext; import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.TimestampLiteralContext; import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.WindowFunctionContext; +import static com.amazon.opendistroforelasticsearch.sql.sql.parser.ParserUtils.createSortOption; import com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL; import com.amazon.opendistroforelasticsearch.sql.ast.expression.AggregateFunction; @@ -63,6 +63,7 @@ import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression; import com.amazon.opendistroforelasticsearch.sql.ast.expression.When; import com.amazon.opendistroforelasticsearch.sql.ast.expression.WindowFunction; +import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption; import com.amazon.opendistroforelasticsearch.sql.common.utils.StringUtils; import com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser; import com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.AndExpressionContext; @@ -171,12 +172,13 @@ public UnresolvedExpression visitWindowFunction(WindowFunctionContext ctx) { .collect(Collectors.toList()); } - List> sortList = Collections.emptyList(); + List> sortList = Collections.emptyList(); if (overClause.orderByClause() != null) { sortList = overClause.orderByClause() .orderByElement() .stream() - .map(item -> ImmutablePair.of(getOrder(item), visit(item.expression()))) + .map(item -> ImmutablePair.of( + createSortOption(item), visit(item.expression()))) .collect(Collectors.toList()); } return new WindowFunction((Function) visit(ctx.function), partitionByList, sortList); @@ -337,8 +339,4 @@ private QualifiedName visitIdentifiers(List identifiers) { ); } - private String getOrder(OrderByElementContext item) { - return (item.order == null) ? "ASC" : item.order.getText(); - } - } diff --git a/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/ParserUtils.java b/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/ParserUtils.java index 5af75cb19e..d4b8d5a841 100644 --- a/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/ParserUtils.java +++ b/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/ParserUtils.java @@ -17,9 +17,15 @@ package com.amazon.opendistroforelasticsearch.sql.sql.parser; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.NullOrder; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder; +import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.OrderByElementContext; + import lombok.experimental.UtilityClass; import org.antlr.v4.runtime.ParserRuleContext; import org.antlr.v4.runtime.Token; +import org.antlr.v4.runtime.tree.TerminalNode; /** * Parser Utils Class. @@ -35,4 +41,37 @@ public static String getTextInQuery(ParserRuleContext ctx, String queryString) { Token stop = ctx.getStop(); return queryString.substring(start.getStartIndex(), stop.getStopIndex() + 1); } + + /** + * Create sort option from syntax tree node. + */ + public static SortOption createSortOption(OrderByElementContext orderBy) { + return new SortOption( + createSortOrder(orderBy.order), + createNullOrder(orderBy.FIRST(), orderBy.LAST())); + } + + /** + * Create sort order for sort option use from ASC/DESC token. + */ + public static SortOrder createSortOrder(Token ctx) { + if (ctx == null) { + return null; + } + return SortOrder.valueOf(ctx.getText().toUpperCase()); + } + + /** + * Create null order for sort option use from FIRST/LAST token. + */ + public static NullOrder createNullOrder(TerminalNode first, TerminalNode last) { + if (first != null) { + return NullOrder.NULL_FIRST; + } else if (last != null) { + return NullOrder.NULL_LAST; + } else { + return null; + } + } + } diff --git a/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/context/QuerySpecification.java b/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/context/QuerySpecification.java index 1124f0150a..f89bfba826 100644 --- a/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/context/QuerySpecification.java +++ b/sql/src/main/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/context/QuerySpecification.java @@ -21,6 +21,8 @@ import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.SelectClauseContext; import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.SelectElementContext; import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.SubqueryAsRelationContext; +import static com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.WindowFunctionContext; +import static com.amazon.opendistroforelasticsearch.sql.sql.parser.ParserUtils.createSortOption; import static com.amazon.opendistroforelasticsearch.sql.sql.parser.ParserUtils.getTextInQuery; import com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL; @@ -28,9 +30,7 @@ import com.amazon.opendistroforelasticsearch.sql.ast.expression.Literal; import com.amazon.opendistroforelasticsearch.sql.ast.expression.QualifiedName; import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression; -import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.NullOrder; import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption; -import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder; import com.amazon.opendistroforelasticsearch.sql.common.utils.StringUtils; import com.amazon.opendistroforelasticsearch.sql.exception.SemanticCheckException; import com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLParser.AggregateFunctionCallContext; @@ -47,9 +47,7 @@ import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.ToString; -import org.antlr.v4.runtime.Token; import org.antlr.v4.runtime.tree.ParseTree; -import org.antlr.v4.runtime.tree.TerminalNode; /** * Query specification domain that collects basic info for a simple query. @@ -183,6 +181,12 @@ public Void visitSubqueryAsRelation(SubqueryAsRelationContext ctx) { return null; } + @Override + public Void visitWindowFunction(WindowFunctionContext ctx) { + // skip collecting sort items in window functions + return null; + } + @Override public Void visitSelectClause(SelectClauseContext ctx) { super.visitSelectClause(ctx); @@ -215,10 +219,7 @@ public Void visitGroupByElement(GroupByElementContext ctx) { @Override public Void visitOrderByElement(OrderByElementContext ctx) { orderByItems.add(visitAstExpression(ctx.expression())); - orderByOptions.add( - new SortOption( - visitSortOrder(ctx.order), - visitNullOrderClause(ctx.FIRST(), ctx.LAST()))); + orderByOptions.add(createSortOption(ctx)); return super.visitOrderByElement(ctx); } @@ -232,23 +233,6 @@ private boolean isDistinct(SelectSpecContext ctx) { return (ctx != null) && (ctx.DISTINCT() != null); } - private SortOrder visitSortOrder(Token ctx) { - if (ctx == null) { - return null; - } - return SortOrder.valueOf(ctx.getText().toUpperCase()); - } - - private NullOrder visitNullOrderClause(TerminalNode first, TerminalNode last) { - if (first != null) { - return NullOrder.NULL_FIRST; - } else if (last != null) { - return NullOrder.NULL_LAST; - } else { - return null; - } - } - private UnresolvedExpression visitAstExpression(ParseTree tree) { return expressionBuilder.visit(tree); } diff --git a/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/AstExpressionBuilderTest.java b/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/AstExpressionBuilderTest.java index 622a5be712..c306625dbc 100644 --- a/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/AstExpressionBuilderTest.java +++ b/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/AstExpressionBuilderTest.java @@ -33,11 +33,15 @@ import static com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL.timestampLiteral; import static com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL.when; import static com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL.window; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.NullOrder.NULL_LAST; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.ASC; +import static com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOrder.DESC; import static org.junit.jupiter.api.Assertions.assertEquals; import com.amazon.opendistroforelasticsearch.sql.ast.Node; import com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL; import com.amazon.opendistroforelasticsearch.sql.ast.expression.DataType; +import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption; import com.amazon.opendistroforelasticsearch.sql.common.antlr.CaseInsensitiveCharStream; import com.amazon.opendistroforelasticsearch.sql.common.antlr.SyntaxAnalysisErrorListener; import com.amazon.opendistroforelasticsearch.sql.sql.antlr.parser.OpenDistroSQLLexer; @@ -255,7 +259,7 @@ public void canBuildWindowFunction() { window( function("RANK"), ImmutableList.of(qualifiedName("state")), - ImmutableList.of(ImmutablePair.of("ASC", qualifiedName("age")))), + ImmutableList.of(ImmutablePair.of(new SortOption(null, null), qualifiedName("age")))), buildExprAst("RANK() OVER (PARTITION BY state ORDER BY age)")); } @@ -265,10 +269,21 @@ public void canBuildWindowFunctionWithoutPartitionBy() { window( function("DENSE_RANK"), ImmutableList.of(), - ImmutableList.of(ImmutablePair.of("DESC", qualifiedName("age")))), + ImmutableList.of(ImmutablePair.of(new SortOption(DESC, null), qualifiedName("age")))), buildExprAst("DENSE_RANK() OVER (ORDER BY age DESC)")); } + @Test + public void canBuildWindowFunctionWithNullOrderSpecified() { + assertEquals( + window( + function("DENSE_RANK"), + ImmutableList.of(), + ImmutableList.of(ImmutablePair.of( + new SortOption(ASC, NULL_LAST), qualifiedName("age")))), + buildExprAst("DENSE_RANK() OVER (ORDER BY age ASC NULLS LAST)")); + } + @Test public void canBuildWindowFunctionWithoutOrderBy() { assertEquals( diff --git a/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/context/QuerySpecificationTest.java b/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/context/QuerySpecificationTest.java index 9ccee1aecc..5e417bf44a 100644 --- a/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/context/QuerySpecificationTest.java +++ b/sql/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/parser/context/QuerySpecificationTest.java @@ -125,6 +125,14 @@ void can_collect_sort_options_in_order_by_clause() { collect("SELECT name FROM test ORDER BY name DESC NULLS FIRST").getOrderByOptions()); } + @Test + void should_skip_sort_items_in_window_function() { + assertEquals(1, + collect("SELECT name, RANK() OVER(ORDER BY age) " + + "FROM test ORDER BY name" + ).getOrderByOptions().size()); + } + private QuerySpecification collect(String query) { QuerySpecification querySpec = new QuerySpecification(); querySpec.collect(parse(query), query);