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

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<Pair<SortOption, Expression>> 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());
   }
 
 }

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<UnresolvedExpression> partitionByList,
-                                     List<Pair<String, UnresolvedExpression>> sortList) {
+                                     List<Pair<SortOption, UnresolvedExpression>> sortList) {
     return new WindowFunction(function, partitionByList, sortList);
   }
 

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<UnresolvedExpression> partitionByList;
-  private List<Pair<String, UnresolvedExpression>> sortList;
+  private List<Pair<SortOption, UnresolvedExpression>> sortList;
 
   @Override
   public List<? extends Node> getChild() {

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")))))));
   }
 
   /**

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<SortOption, SortOption> expects =
+        ImmutableMap.<SortOption, SortOption>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);
+    });
+  }
+
 }

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
 ----

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));
+  }
+
+}

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<Pair<String, UnresolvedExpression>> sortList = Collections.emptyList();
+    List<Pair<SortOption, UnresolvedExpression>> 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<IdentContext> identifiers) {
     );
   }
 
-  private String getOrder(OrderByElementContext item) {
-    return (item.order == null) ? "ASC" : item.order.getText();
-  }
-
 }

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;
+    }
+  }
+
 }