Enhance sort command in PPL

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

@@ -510,7 +510,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());
   }
 
@@ -719,7 +719,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);
 
@@ -736,7 +736,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

@@ -466,7 +466,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

@@ -333,6 +333,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()),

docs/user/ppl/cmd/sort.rst

@@ -16,11 +16,13 @@ Description
 
 Syntax
 ============
-sort <[+|-] sort-field>...
+sort [count] <[+|-] sort-field>... [desc|d]
 
 
+* count: optional. The number of results to return. **Default:** returns all results. Specifying a count of 0 or less than 0 also returns all results.
 * [+|-]: optional. The plus [+] stands for ascending order and NULL/MISSING first and a minus [-] stands for descending order and NULL/MISSING last. **Default:** ascending order and NULL/MISSING first.
-* sort-field: mandatory. The field used to sort.
+* sort-field: mandatory. The field used to sort. Can use ``auto(field)``, ``str(field)``, ``ip(field)``, or ``num(field)`` to specify how to interpret field values.
+* [desc|d]: optional. Reverses the sort results. If multiple fields are specified, reverses order of the first field then for all duplicate values of the first field, reverses the order of the values of the second field and so on.
 
 
 Example 1: Sort by one field
@@ -49,7 +51,7 @@ The example show sort all the document with age field in ascending order.
 
 PPL query::
 
-    os> source=accounts | sort age | fields account_number, age;
+    os> source=accounts | sort 0 age | fields account_number, age;
     fetched rows / total rows = 4/4
     +----------------+-----+
     | account_number | age |
@@ -114,3 +116,74 @@ PPL query::
     | Pyrami   |
     | Quility  |
     +----------+
+
+Example 5: Specify the number of sorted documents to return
+============================================================
+
+The example shows sorting all the document and returning 2 documents.
+
+PPL query::
+
+    os> source=accounts | sort 2 age | fields account_number, age;
+    fetched rows / total rows = 2/2
+    +----------------+-----+
+    | account_number | age |
+    |----------------+-----|
+    | 13             | 28  |
+    | 1              | 32  |
+    +----------------+-----+
+
+Example 6: Sort with desc modifier
+===================================
+
+The example shows sorting with the desc modifier to reverse sort order.
+
+PPL query::
+
+    os> source=accounts | sort age desc | fields account_number, age;
+    fetched rows / total rows = 4/4
+    +----------------+-----+
+    | account_number | age |
+    |----------------+-----|
+    | 6              | 36  |
+    | 18             | 33  |
+    | 1              | 32  |
+    | 13             | 28  |
+    +----------------+-----+
+
+Example 7: Sort by multiple fields with desc modifier
+======================================================
+
+The example shows sorting by multiple fields using desc, which reverses the sort order for all specified fields. Gender is reversed from ascending to descending, and the descending age sort is reversed to ascending within each gender group.
+
+PPL query::
+
+    os> source=accounts | sort gender, -age desc | fields account_number, gender, age;
+    fetched rows / total rows = 4/4
+    +----------------+--------+-----+
+    | account_number | gender | age |
+    |----------------+--------+-----|
+    | 1              | M      | 32  |
+    | 18             | M      | 33  |
+    | 6              | M      | 36  |
+    | 13             | F      | 28  |
+    +----------------+--------+-----+
+
+
+Example 8: Sort with specifying field type
+==================================
+
+The example shows sorting with str() to sort numeric values lexicographically.
+
+PPL query::
+
+    os> source=accounts | sort str(account_number) | fields account_number;
+    fetched rows / total rows = 4/4
+    +----------------+
+    | account_number |
+    |----------------|
+    | 1              |
+    | 13             |
+    | 18             |
+    | 6              |
+    +----------------+

integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java

@@ -150,6 +150,33 @@ public void testSortPushDownExplain() throws IOException {
                 + "| fields age"));
   }
 
+  @Test
+  public void testSortWithCountPushDownExplain() throws IOException {
+    String expected = loadExpectedPlan("explain_sort_count_push.json");
+    assertJsonEqualsIgnoreId(
+        expected,
+        explainQueryToString("source=opensearch-sql_test_index_account | sort 5 age | fields age"));
+  }
+
+  @Test
+  public void testSortWithDescPushDownExplain() throws IOException {
+    String expected = loadExpectedPlan("explain_sort_desc_push.json");
+    assertJsonEqualsIgnoreId(
+        expected,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account | sort age, - firstname desc | fields age,"
+                + " firstname"));
+  }
+
+  @Test
+  public void testSortWithTypePushDownExplain() throws IOException {
+    String expected = loadExpectedPlan("explain_sort_type_push.json");
+    assertJsonEqualsIgnoreId(
+        expected,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account | sort num(age) | fields age"));
+  }
+
   @Test
   public void testSortWithAggregationExplain() throws IOException {
     // Sorts whose by fields are aggregators should not be pushed down

integ-test/src/test/java/org/opensearch/sql/ppl/SortCommandIT.java

@@ -161,4 +161,60 @@ public void testSortThenHead() throws IOException {
         executeQuery(String.format("source=%s | sort age | head 2 | fields age", TEST_INDEX_BANK));
     verifyOrder(result, rows(28), rows(32));
   }
+
+  @Test
+  public void testSortWithCountLimit() throws IOException {
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | sort 3 - account_number | fields account_number", TEST_INDEX_BANK));
+    verifyOrder(result, rows(32), rows(25), rows(20));
+  }
+
+  @Test
+  public void testSortWithCountZero() throws IOException {
+    // count=0 should return all results
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | sort 0 account_number | fields account_number", TEST_INDEX_BANK));
+    verifyOrder(result, rows(1), rows(6), rows(13), rows(18), rows(20), rows(25), rows(32));
+  }
+
+  @Test
+  public void testSortWithDesc() throws IOException {
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | sort account_number desc | fields account_number", TEST_INDEX_BANK));
+    verifyOrder(result, rows(32), rows(25), rows(20), rows(18), rows(13), rows(6), rows(1));
+  }
+
+  @Test
+  public void testSortWithDescMultipleFields() throws IOException {
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | sort 4 age, - account_number desc | fields age, account_number",
+                TEST_INDEX_BANK));
+    verifyOrder(result, rows(39, 25), rows(36, 6), rows(36, 20), rows(34, 32));
+  }
+
+  @Test
+  public void testSortWithStrCast() throws IOException {
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | sort str(account_number) | fields account_number", TEST_INDEX_BANK));
+    verifyOrder(result, rows(1), rows(13), rows(18), rows(20), rows(25), rows(32), rows(6));
+  }
+
+  @Test
+  public void testSortWithNumCast() throws IOException {
+    JSONObject result =
+        executeQuery(
+            String.format("source=%s | sort num(bytes) | fields bytes", TEST_INDEX_WEBLOGS));
+    verifyOrder(
+        result, rows("1234"), rows("3985"), rows("4085"), rows("4321"), rows("6245"), rows("9876"));
+  }
 }