top and head: allow limit keywords

Signed-off-by: Simeon Widdis <[email protected]>

Author: Swiddis

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,25 @@
+/*
+ * 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 {
+    public final Settings settings = Mockito.mock(Settings.class);
+    public final PPLSyntaxParser parser = new PPLSyntaxParser();
+
+    public 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

@@ -76,17 +76,14 @@
 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.AstPlanningTest;
 import org.opensearch.sql.ppl.antlr.PPLSyntaxParser;
 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 +1062,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,28 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache 2.0
+ */
+
+package org.opensearch.sql.ppl.parser;
+
+import org.junit.Test;
+import org.opensearch.sql.ppl.AstPlanningTest;
+
+import static org.junit.Assert.assertEquals;
+
+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

@@ -20,18 +20,11 @@
 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.AstPlanningTest;
 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));
-  }
-
+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 +52,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"));
-  }
 }