top and head: allow limit keywords

docs/user/ppl/cmd/head.rst

@@ -19,6 +19,7 @@ Syntax
 head [<size>] [from <offset>]
 
 * <size>: optional integer. number of results to return. **Default:** 10
+  * For readability or convenience, this can be prefixed by `limit=`.
 * <offset>: integer after optional ``from``. number of results to skip. **Default:** 0
 
 Example 1: Get first 10 results

docs/user/ppl/cmd/top.rst

@@ -21,6 +21,7 @@ top [N] <field-list> [by-clause]
 top [N] <field-list> [top-options] [by-clause] ``(available from 3.1.0+)``
 
 * N: number of results to return. **Default**: 10
+  * For readability or convenience, this can be prefixed by `limit=`.
 * field-list: mandatory. comma-delimited list of field names.
 * by-clause: optional. one or more fields to group the results by.
 * top-options: optional. options for the top command. Supported syntax is [countfield=<string>] [showcount=<bool>].

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -198,7 +198,7 @@ evalCommand
    ;
 
 headCommand
-   : HEAD (number = integerLiteral)? (FROM from = integerLiteral)?
+   : HEAD ((LIMIT EQUAL)? number = integerLiteral)? (FROM from = integerLiteral)?
    ;
 
 binCommand
@@ -232,7 +232,7 @@ logSpanValue
    ;
 
 topCommand
-   : TOP (number = integerLiteral)? (COUNTFIELD EQUAL countfield = stringLiteral)? (SHOWCOUNT EQUAL showcount = booleanLiteral)? fieldList (byClause)?
+   : TOP ((LIMIT EQUAL)? number = integerLiteral)? (COUNTFIELD EQUAL countfield = stringLiteral)? (SHOWCOUNT EQUAL showcount = booleanLiteral)? fieldList (byClause)?
    ;
 
 rareCommand
@@ -1295,6 +1295,7 @@ keywordsCanBeId
    | FREQUENCY_THRESHOLD_PERCENTAGE
    | MAX_SAMPLE_COUNT
    | BUFFER_LIMIT
+   | LIMIT
    | WITH
    | REGEX
    | PUNCT

ppl/src/test/java/org/opensearch/sql/ppl/AstPlanningTest.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache 2.0
+ */
+
+package org.opensearch.sql.ppl;
+
+import org.mockito.Mockito;
+import org.opensearch.sql.ast.Node;
+import org.opensearch.sql.common.setting.Settings;
+import org.opensearch.sql.ppl.antlr.PPLSyntaxParser;
+import org.opensearch.sql.ppl.parser.AstBuilder;
+
+/** Base class for tests for the AST query planner. */
+public class AstPlanningTest {
+  protected final Settings settings = Mockito.mock(Settings.class);
+  protected final PPLSyntaxParser parser = new PPLSyntaxParser();
+
+  protected Node plan(String query) {
+    AstBuilder astBuilder = new AstBuilder(query, settings);
+    return astBuilder.visit(parser.parse(query));
+  }
+}

ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java

@@ -60,7 +60,6 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
-import org.mockito.Mockito;
 import org.opensearch.sql.ast.Node;
 import org.opensearch.sql.ast.dsl.AstDSL;
 import org.opensearch.sql.ast.expression.DataType;
@@ -74,19 +73,14 @@
 import org.opensearch.sql.ast.tree.ML;
 import org.opensearch.sql.ast.tree.RareTopN.CommandType;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
-import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.common.setting.Settings.Key;
-import org.opensearch.sql.ppl.antlr.PPLSyntaxParser;
+import org.opensearch.sql.ppl.AstPlanningTest;
 import org.opensearch.sql.utils.SystemIndexUtils;
 
-public class AstBuilderTest {
+public class AstBuilderTest extends AstPlanningTest {
 
   @Rule public ExpectedException exceptionRule = ExpectedException.none();
 
-  private final PPLSyntaxParser parser = new PPLSyntaxParser();
-
-  private final Settings settings = Mockito.mock(Settings.class);
-
   @Test
   public void testSearchCommand() {
     assertEqual(
@@ -1065,11 +1059,6 @@ protected void assertEqual(String query, String expected) {
     assertEqual(query, expectedPlan);
   }
 
-  private Node plan(String query) {
-    AstBuilder astBuilder = new AstBuilder(query, settings);
-    return astBuilder.visit(parser.parse(query));
-  }
-
   private String mappingTable(String indexName) {
     return SystemIndexUtils.mappingTable(indexName, PPL_SPEC);
   }

ppl/src/test/java/org/opensearch/sql/ppl/parser/AstEquivalenceTest.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache 2.0
+ */
+
+package org.opensearch.sql.ppl.parser;
+
+import static org.junit.Assert.assertEquals;
+
+import org.junit.Test;
+import org.opensearch.sql.ppl.AstPlanningTest;
+
+public class AstEquivalenceTest extends AstPlanningTest {
+  @Test
+  public void testSpathArgumentDeshuffle() {
+    assertEquals(plan("source = t | spath path=a input=a"), plan("source = t | spath input=a a"));
+  }
+
+  @Test
+  public void testHeadLimitEquivalent() {
+    assertEquals(plan("source = t | head limit=50"), plan("source = t | head 50"));
+  }
+
+  @Test
+  public void testTopLimitEquivalent() {
+    assertEquals(
+        plan("source = t | top limit=50 field_name"), plan("source = t | top 50 field_name"));
+  }
+}

ppl/src/test/java/org/opensearch/sql/ppl/utils/SPathRewriteTest.java

@@ -15,23 +15,11 @@
 import static org.opensearch.sql.ast.dsl.AstDSL.stringLiteral;
 
 import org.junit.Test;
-import org.mockito.Mockito;
-import org.opensearch.sql.ast.Node;
 import org.opensearch.sql.ast.tree.Eval;
 import org.opensearch.sql.ast.tree.SPath;
-import org.opensearch.sql.common.setting.Settings;
-import org.opensearch.sql.ppl.antlr.PPLSyntaxParser;
-import org.opensearch.sql.ppl.parser.AstBuilder;
-
-public class SPathRewriteTest {
-  private final Settings settings = Mockito.mock(Settings.class);
-  private final PPLSyntaxParser parser = new PPLSyntaxParser();
-
-  private Node plan(String query) {
-    AstBuilder astBuilder = new AstBuilder(query, settings);
-    return astBuilder.visit(parser.parse(query));
-  }
+import org.opensearch.sql.ppl.AstPlanningTest;
 
+public class SPathRewriteTest extends AstPlanningTest {
   // Control test to make sure something fundamental hasn't changed about the json_extract parsing
   @Test
   public void testEvalControl() {
@@ -59,9 +47,4 @@ public void testSpathMissingInputArgumentHandling() {
   public void testSpathMissingPathArgumentHandling() {
     plan("source = t | spath input=a output=a");
   }
-
-  @Test
-  public void testSpathArgumentDeshuffle() {
-    assertEquals(plan("source = t | spath path=a input=a"), plan("source = t | spath input=a a"));
-  }
 }