Skip to content

Commit 3b5dc95

Browse files
ritvibhattaaarone90
ritvibhatt
authored and
aaarone90
committed
Enhance sort command in PPL (opensearch-project#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> Signed-off-by: Aaron Alvarez <aaarone@amazon.com>
1 parent 0323880 commit 3b5dc95

File tree

33 files changed

+494
-34
lines changed

33 files changed

+494
-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
@@ -558,7 +558,7 @@ public LogicalPlan visitPatterns(Patterns node, AnalysisContext context) {
558558
@Override
559559
public LogicalPlan visitSort(Sort node, AnalysisContext context) {
560560
LogicalPlan child = node.getChild().get(0).accept(this, context);
561-
return buildSort(child, context, node.getSortList());
561+
return buildSort(child, context, node.getCount(), node.getSortList());
562562
}
563563

564564
/** Build {@link LogicalDedupe}. */
@@ -720,7 +720,7 @@ public LogicalPlan visitTrendline(Trendline node, AnalysisContext context) {
720720
}
721721

722722
return new LogicalTrendline(
723-
buildSort(child, context, Collections.singletonList(node.getSortByField().get())),
723+
buildSort(child, context, 0, Collections.singletonList(node.getSortByField().get())),
724724
computationsAndTypes.build());
725725
}
726726

@@ -773,7 +773,7 @@ public LogicalPlan visitAppendCol(AppendCol node, AnalysisContext context) {
773773
}
774774

775775
private LogicalSort buildSort(
776-
LogicalPlan child, AnalysisContext context, List<Field> sortFields) {
776+
LogicalPlan child, AnalysisContext context, Integer count, List<Field> sortFields) {
777777
ExpressionReferenceOptimizer optimizer =
778778
new ExpressionReferenceOptimizer(expressionAnalyzer.getRepository(), child);
779779

@@ -790,7 +790,7 @@ private LogicalSort buildSort(
790790
return ImmutablePair.of(analyzeSortOption(sortField.getFieldArgs()), expression);
791791
})
792792
.collect(Collectors.toList());
793-
return new LogicalSort(child, sortList);
793+
return new LogicalSort(child, count, sortList);
794794
}
795795

796796
private SortOption analyzeSortOption(List<Argument> fieldArgs) {

‎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
@@ -466,7 +466,11 @@ public static Span span(UnresolvedExpression field, UnresolvedExpression value,
466466
}
467467

468468
public static Sort sort(UnresolvedPlan input, Field... sorts) {
469-
return new Sort(input, Arrays.asList(sorts));
469+
return new Sort(Arrays.asList(sorts)).attach(input);
470+
}
471+
472+
public static Sort sort(UnresolvedPlan input, Integer count, Field... sorts) {
473+
return new Sort(count, Arrays.asList(sorts)).attach(input);
470474
}
471475

472476
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
@@ -422,6 +422,11 @@ public RelNode visitSort(Sort node, CalcitePlanContext context) {
422422
})
423423
.collect(Collectors.toList());
424424
context.relBuilder.sort(sortList);
425+
// Apply count parameter as limit
426+
if (node.getCount() != 0) {
427+
context.relBuilder.limit(0, node.getCount());
428+
}
429+
425430
return context.relBuilder.peek();
426431
}
427432

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