Skip to content

Commit c8d2694

Browse files
ritvibhattritvibhatt
authored
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>
1 parent bb1a644 commit c8d2694

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
@@ -511,7 +511,7 @@ public LogicalPlan visitPatterns(Patterns node, AnalysisContext context) {
511511
@Override
512512
public LogicalPlan visitSort(Sort node, AnalysisContext context) {
513513
LogicalPlan child = node.getChild().get(0).accept(this, context);
514-
return buildSort(child, context, node.getSortList());
514+
return buildSort(child, context, node.getCount(), node.getSortList());
515515
}
516516

517517
/** Build {@link LogicalDedupe}. */
@@ -673,7 +673,7 @@ public LogicalPlan visitTrendline(Trendline node, AnalysisContext context) {
673673
}
674674

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

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

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

@@ -743,7 +743,7 @@ private LogicalSort buildSort(
743743
return ImmutablePair.of(analyzeSortOption(sortField.getFieldArgs()), expression);
744744
})
745745
.collect(Collectors.toList());
746-
return new LogicalSort(child, sortList);
746+
return new LogicalSort(child, count, sortList);
747747
}
748748

749749
/**

‎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
@@ -333,6 +333,11 @@ public RelNode visitSort(Sort node, CalcitePlanContext context) {
333333
})
334334
.collect(Collectors.toList());
335335
context.relBuilder.sort(sortList);
336+
// Apply count parameter as limit
337+
if (node.getCount() != 0) {
338+
context.relBuilder.limit(0, node.getCount());
339+
}
340+
336341
return context.relBuilder.peek();
337342
}
338343

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