Skip to content

Commit 181fda8

Browse files
github-actions[bot]github-actions[bot]
committed
Enhance sort command in PPL (#3934)
* enhance sort command Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * update failing tests Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * fix formatting Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * add integ tests Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * update documentation Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * fix failing test Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * update default and tests Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * fix tests Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * update analyzer test Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * update reverse sort direction Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * update formatting Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * update docs Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * add javadoc Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * add tests Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * fix failing tests Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * fix failing tests Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * update integ tests for query size limit change Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * add explainit for desc and type cast Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * add tests for desc Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * fix formatting Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * make count optional Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * add cross cluster tests Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * fix tests Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * normalize count in AST node Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * default null count to 0 Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> * update logicalsort default constructor Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> --------- Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com> (cherry picked from commit c8d2694) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent ab211df commit 181fda8

File tree

33 files changed

+495
-34
lines changed

33 files changed

+495
-34
lines changed

‎core/src/main/java/org/opensearch/sql/analysis/Analyzer.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -512,7 +512,7 @@ public LogicalPlan visitPatterns(Patterns node, AnalysisContext context) {
512512
@Override
513513
public LogicalPlan visitSort(Sort node, AnalysisContext context) {
514514
LogicalPlan child = node.getChild().get(0).accept(this, context);
515-
return buildSort(child, context, node.getSortList());
515+
return buildSort(child, context, node.getCount(), node.getSortList());
516516
}
517517

518518
/** Build {@link LogicalDedupe}. */
@@ -672,7 +672,7 @@ public LogicalPlan visitTrendline(Trendline node, AnalysisContext context) {
672672
}
673673

674674
return new LogicalTrendline(
675-
buildSort(child, context, Collections.singletonList(node.getSortByField().get())),
675+
buildSort(child, context, 0, Collections.singletonList(node.getSortByField().get())),
676676
computationsAndTypes.build());
677677
}
678678

@@ -725,7 +725,7 @@ public LogicalPlan visitAppendCol(AppendCol node, AnalysisContext context) {
725725
}
726726

727727
private LogicalSort buildSort(
728-
LogicalPlan child, AnalysisContext context, List<Field> sortFields) {
728+
LogicalPlan child, AnalysisContext context, Integer count, List<Field> sortFields) {
729729
ExpressionReferenceOptimizer optimizer =
730730
new ExpressionReferenceOptimizer(expressionAnalyzer.getRepository(), child);
731731

@@ -742,7 +742,7 @@ private LogicalSort buildSort(
742742
return ImmutablePair.of(analyzeSortOption(sortField.getFieldArgs()), expression);
743743
})
744744
.collect(Collectors.toList());
745-
return new LogicalSort(child, sortList);
745+
return new LogicalSort(child, count, sortList);
746746
}
747747

748748
/**

‎core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -469,7 +469,11 @@ public static Span span(UnresolvedExpression field, UnresolvedExpression value,
469469
}
470470

471471
public static Sort sort(UnresolvedPlan input, Field... sorts) {
472-
return new Sort(input, Arrays.asList(sorts));
472+
return new Sort(Arrays.asList(sorts)).attach(input);
473+
}
474+
475+
public static Sort sort(UnresolvedPlan input, Integer count, Field... sorts) {
476+
return new Sort(count, Arrays.asList(sorts)).attach(input);
473477
}
474478

475479
public static Dedupe dedupe(UnresolvedPlan input, List<Argument> options, Field... fields) {

‎core/src/main/java/org/opensearch/sql/ast/tree/Sort.java‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,9 @@
1212

1313
import com.google.common.collect.ImmutableList;
1414
import java.util.List;
15-
import lombok.AllArgsConstructor;
1615
import lombok.Data;
1716
import lombok.EqualsAndHashCode;
1817
import lombok.Getter;
19-
import lombok.RequiredArgsConstructor;
2018
import lombok.ToString;
2119
import org.opensearch.sql.ast.AbstractNodeVisitor;
2220
import org.opensearch.sql.ast.expression.Field;
@@ -25,12 +23,25 @@
2523
@ToString
2624
@EqualsAndHashCode(callSuper = false)
2725
@Getter
28-
@RequiredArgsConstructor
29-
@AllArgsConstructor
3026
public class Sort extends UnresolvedPlan {
3127
private UnresolvedPlan child;
28+
29+
/**
30+
* The count value can be either 0 or a positive number. A value of 0 means return all documents.
31+
*/
32+
private final Integer count;
33+
3234
private final List<Field> sortList;
3335

36+
public Sort(List<Field> sortList) {
37+
this(0, sortList);
38+
}
39+
40+
public Sort(Integer count, List<Field> sortList) {
41+
this.count = count;
42+
this.sortList = sortList;
43+
}
44+
3445
@Override
3546
public Sort attach(UnresolvedPlan child) {
3647
this.child = child;

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -337,6 +337,11 @@ public RelNode visitSort(Sort node, CalcitePlanContext context) {
337337
})
338338
.collect(Collectors.toList());
339339
context.relBuilder.sort(sortList);
340+
// Apply count parameter as limit
341+
if (node.getCount() != 0) {
342+
context.relBuilder.limit(0, node.getCount());
343+
}
344+
340345
return context.relBuilder.peek();
341346
}
342347

‎core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,11 @@ public PhysicalPlan visitNested(LogicalNested node, C context) {
106106

107107
@Override
108108
public PhysicalPlan visitSort(LogicalSort node, C context) {
109-
return new SortOperator(visitChild(node, context), node.getSortList());
109+
PhysicalPlan child = visitChild(node, context);
110+
if (node.getCount() != 0) {
111+
return new TakeOrderedOperator(child, node.getCount(), 0, node.getSortList());
112+
}
113+
return new SortOperator(child, node.getSortList());
110114
}
111115

112116
@Override

‎core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,11 @@ public static LogicalPlan sort(LogicalPlan input, Pair<SortOption, Expression>..
101101
return new LogicalSort(input, Arrays.asList(sorts));
102102
}
103103

104+
public static LogicalPlan sort(
105+
LogicalPlan input, Integer count, Pair<SortOption, Expression>... sorts) {
106+
return new LogicalSort(input, count, Arrays.asList(sorts));
107+
}
108+
104109
public static LogicalPlan dedupe(LogicalPlan input, Expression... fields) {
105110
return dedupe(input, 1, false, false, fields);
106111
}

‎core/src/main/java/org/opensearch/sql/planner/logical/LogicalSort.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,20 @@
2020
@EqualsAndHashCode(callSuper = true)
2121
public class LogicalSort extends LogicalPlan {
2222

23+
/** Maximum number of results to return after sorting. */
24+
private final Integer count;
25+
2326
private final List<Pair<SortOption, Expression>> sortList;
2427

25-
/** Constructor of LogicalSort. */
2628
public LogicalSort(LogicalPlan child, List<Pair<SortOption, Expression>> sortList) {
29+
this(child, 0, sortList);
30+
}
31+
32+
/** Constructor of LogicalSort. */
33+
public LogicalSort(
34+
LogicalPlan child, Integer count, List<Pair<SortOption, Expression>> sortList) {
2735
super(Collections.singletonList(child));
36+
this.count = count;
2837
this.sortList = sortList;
2938
}
3039

‎core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushFilterUnderSort.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ public PushFilterUnderSort() {
4141
@Override
4242
public LogicalPlan apply(LogicalFilter filter, Captures captures) {
4343
LogicalSort sort = captures.get(capture);
44-
return new LogicalSort(filter.replaceChildPlans(sort.getChild()), sort.getSortList());
44+
return new LogicalSort(
45+
filter.replaceChildPlans(sort.getChild()), sort.getCount(), sort.getSortList());
4546
}
4647
}

‎core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1021,6 +1021,7 @@ public void window_function() {
10211021
LogicalPlanDSL.window(
10221022
LogicalPlanDSL.sort(
10231023
LogicalPlanDSL.relation("test", table),
1024+
0,
10241025
ImmutablePair.of(DEFAULT_ASC, DSL.ref("string_value", STRING)),
10251026
ImmutablePair.of(DEFAULT_ASC, DSL.ref("integer_value", INTEGER))),
10261027
DSL.named("window_function", DSL.rowNumber()),
@@ -1465,6 +1466,7 @@ public void trendline_with_sort() {
14651466
LogicalPlanDSL.trendline(
14661467
LogicalPlanDSL.sort(
14671468
LogicalPlanDSL.relation("schema", table),
1469+
0,
14681470
Pair.of(
14691471
new SortOption(SortOrder.ASC, NullOrder.NULL_FIRST),
14701472
DSL.ref("float_value", ExprCoreType.FLOAT))),

‎core/src/test/java/org/opensearch/sql/analysis/WindowExpressionAnalyzerTest.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ void should_wrap_child_with_window_and_sort_operator_if_project_item_windowed()
5050
LogicalPlanDSL.window(
5151
LogicalPlanDSL.sort(
5252
LogicalPlanDSL.relation("test", table),
53+
0,
5354
ImmutablePair.of(DEFAULT_ASC, DSL.ref("string_value", STRING)),
5455
ImmutablePair.of(DEFAULT_DESC, DSL.ref("integer_value", INTEGER))),
5556
DSL.named("row_number", DSL.rowNumber()),

0 commit comments

Comments
 (0)