Enhance sort command in PPL (#3934) (#4074)

* enhance sort command



* update failing tests



* fix formatting



* add integ tests



* update documentation



* fix failing test



* update default and tests



* fix tests



* update analyzer test



* update reverse sort direction



* update formatting



* update docs



* add javadoc



* add tests



* fix failing tests



* fix failing tests



* update integ tests for query size limit change



* add explainit for desc and type cast



* add tests for desc



* fix formatting



* make count optional



* add cross cluster tests



* fix tests



* normalize count in AST node



* default null count to 0



* update logicalsort default constructor



---------


(cherry picked from commit c8d2694)

Signed-off-by: Ritvi Bhatt <ribhatt@amazon.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: opensearch-trigger-bot[bot]

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

@@ -512,7 +512,7 @@ public LogicalPlan visitPatterns(Patterns node, AnalysisContext context) {
   @Override
   public LogicalPlan visitSort(Sort node, AnalysisContext context) {
     LogicalPlan child = node.getChild().get(0).accept(this, context);
-    return buildSort(child, context, node.getSortList());
+    return buildSort(child, context, node.getCount(), node.getSortList());
   }
 
   /** Build {@link LogicalDedupe}. */
@@ -672,7 +672,7 @@ public LogicalPlan visitTrendline(Trendline node, AnalysisContext context) {
     }
 
     return new LogicalTrendline(
-        buildSort(child, context, Collections.singletonList(node.getSortByField().get())),
+        buildSort(child, context, 0, Collections.singletonList(node.getSortByField().get())),
         computationsAndTypes.build());
   }
 
@@ -725,7 +725,7 @@ public LogicalPlan visitAppendCol(AppendCol node, AnalysisContext context) {
   }
 
   private LogicalSort buildSort(
-      LogicalPlan child, AnalysisContext context, List<Field> sortFields) {
+      LogicalPlan child, AnalysisContext context, Integer count, List<Field> sortFields) {
     ExpressionReferenceOptimizer optimizer =
         new ExpressionReferenceOptimizer(expressionAnalyzer.getRepository(), child);
 
@@ -742,7 +742,7 @@ private LogicalSort buildSort(
                   return ImmutablePair.of(analyzeSortOption(sortField.getFieldArgs()), expression);
                 })
             .collect(Collectors.toList());
-    return new LogicalSort(child, sortList);
+    return new LogicalSort(child, count, sortList);
   }
 
   /**

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

@@ -469,7 +469,11 @@ public static Span span(UnresolvedExpression field, UnresolvedExpression value,
   }
 
   public static Sort sort(UnresolvedPlan input, Field... sorts) {
-    return new Sort(input, Arrays.asList(sorts));
+    return new Sort(Arrays.asList(sorts)).attach(input);
+  }
+
+  public static Sort sort(UnresolvedPlan input, Integer count, Field... sorts) {
+    return new Sort(count, Arrays.asList(sorts)).attach(input);
   }
 
   public static Dedupe dedupe(UnresolvedPlan input, List<Argument> options, Field... fields) {

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

@@ -12,11 +12,9 @@
 
 import com.google.common.collect.ImmutableList;
 import java.util.List;
-import lombok.AllArgsConstructor;
 import lombok.Data;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
-import lombok.RequiredArgsConstructor;
 import lombok.ToString;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.expression.Field;
@@ -25,12 +23,25 @@
 @ToString
 @EqualsAndHashCode(callSuper = false)
 @Getter
-@RequiredArgsConstructor
-@AllArgsConstructor
 public class Sort extends UnresolvedPlan {
   private UnresolvedPlan child;
+
+  /**
+   * The count value can be either 0 or a positive number. A value of 0 means return all documents.
+   */
+  private final Integer count;
+
   private final List<Field> sortList;
 
+  public Sort(List<Field> sortList) {
+    this(0, sortList);
+  }
+
+  public Sort(Integer count, List<Field> sortList) {
+    this.count = count;
+    this.sortList = sortList;
+  }
+
   @Override
   public Sort attach(UnresolvedPlan child) {
     this.child = child;

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

@@ -337,6 +337,11 @@ public RelNode visitSort(Sort node, CalcitePlanContext context) {
                 })
             .collect(Collectors.toList());
     context.relBuilder.sort(sortList);
+    // Apply count parameter as limit
+    if (node.getCount() != 0) {
+      context.relBuilder.limit(0, node.getCount());
+    }
+
     return context.relBuilder.peek();
   }
 

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

@@ -106,7 +106,11 @@ public PhysicalPlan visitNested(LogicalNested node, C context) {
 
   @Override
   public PhysicalPlan visitSort(LogicalSort node, C context) {
-    return new SortOperator(visitChild(node, context), node.getSortList());
+    PhysicalPlan child = visitChild(node, context);
+    if (node.getCount() != 0) {
+      return new TakeOrderedOperator(child, node.getCount(), 0, node.getSortList());
+    }
+    return new SortOperator(child, node.getSortList());
   }
 
   @Override

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

@@ -101,6 +101,11 @@ public static LogicalPlan sort(LogicalPlan input, Pair<SortOption, Expression>..
     return new LogicalSort(input, Arrays.asList(sorts));
   }
 
+  public static LogicalPlan sort(
+      LogicalPlan input, Integer count, Pair<SortOption, Expression>... sorts) {
+    return new LogicalSort(input, count, Arrays.asList(sorts));
+  }
+
   public static LogicalPlan dedupe(LogicalPlan input, Expression... fields) {
     return dedupe(input, 1, false, false, fields);
   }

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

@@ -20,11 +20,20 @@
 @EqualsAndHashCode(callSuper = true)
 public class LogicalSort extends LogicalPlan {
 
+  /** Maximum number of results to return after sorting. */
+  private final Integer count;
+
   private final List<Pair<SortOption, Expression>> sortList;
 
-  /** Constructor of LogicalSort. */
   public LogicalSort(LogicalPlan child, List<Pair<SortOption, Expression>> sortList) {
+    this(child, 0, sortList);
+  }
+
+  /** Constructor of LogicalSort. */
+  public LogicalSort(
+      LogicalPlan child, Integer count, List<Pair<SortOption, Expression>> sortList) {
     super(Collections.singletonList(child));
+    this.count = count;
     this.sortList = sortList;
   }
 

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

@@ -41,6 +41,7 @@ public PushFilterUnderSort() {
   @Override
   public LogicalPlan apply(LogicalFilter filter, Captures captures) {
     LogicalSort sort = captures.get(capture);
-    return new LogicalSort(filter.replaceChildPlans(sort.getChild()), sort.getSortList());
+    return new LogicalSort(
+        filter.replaceChildPlans(sort.getChild()), sort.getCount(), sort.getSortList());
   }
 }

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

@@ -1021,6 +1021,7 @@ public void window_function() {
             LogicalPlanDSL.window(
                 LogicalPlanDSL.sort(
                     LogicalPlanDSL.relation("test", table),
+                    0,
                     ImmutablePair.of(DEFAULT_ASC, DSL.ref("string_value", STRING)),
                     ImmutablePair.of(DEFAULT_ASC, DSL.ref("integer_value", INTEGER))),
                 DSL.named("window_function", DSL.rowNumber()),
@@ -1465,6 +1466,7 @@ public void trendline_with_sort() {
         LogicalPlanDSL.trendline(
             LogicalPlanDSL.sort(
                 LogicalPlanDSL.relation("schema", table),
+                0,
                 Pair.of(
                     new SortOption(SortOrder.ASC, NullOrder.NULL_FIRST),
                     DSL.ref("float_value", ExprCoreType.FLOAT))),

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

@@ -50,6 +50,7 @@ void should_wrap_child_with_window_and_sort_operator_if_project_item_windowed()
         LogicalPlanDSL.window(
             LogicalPlanDSL.sort(
                 LogicalPlanDSL.relation("test", table),
+                0,
                 ImmutablePair.of(DEFAULT_ASC, DSL.ref("string_value", STRING)),
                 ImmutablePair.of(DEFAULT_DESC, DSL.ref("integer_value", INTEGER))),
             DSL.named("row_number", DSL.rowNumber()),