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

Support NULLS FIRST/LAST ordering for window functions #929

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
Expand All @@ -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());
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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")))))));
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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));
}

Expand All @@ -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);
});
}

}
20 changes: 19 additions & 1 deletion docs/user/dql/window.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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}] [, ...]
)


Expand Down Expand Up @@ -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
----
Expand Down
Original file line number Diff line number Diff line change
@@ -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));
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -337,8 +339,4 @@ private QualifiedName visitIdentifiers(List<IdentContext> identifiers) {
);
}

private String getOrder(OrderByElementContext item) {
return (item.order == null) ? "ASC" : item.order.getText();
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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;
}
}

}
Loading