top and head: allow limit keywords (#4249)

* top and head: allow limit keywords

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>

* Add docs

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>

* Fix style + Ryan's comments

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>

* Update docs

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>

* Fix version header type

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>

* Maybe with no version header it's better?

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>

* Move some stuff around

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>

* Fmt

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>

* Fix doctest

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>

---------

Signed-off-by: Simeon Widdis <sawiddis@amazon.com>
Signed-off-by: Simeon Widdis <sawiddis@gmail.com>

Author: Swiddis

docs/user/ppl/cmd/head.md

@@ -10,7 +10,7 @@ The `head` command returns the first N lines from a search result.
 The `head` command has the following syntax:
 
 ```syntax
-head [<size>] [from <offset>]
+head [limit=][<size>] [from <offset>]
 ```
 
 ## Parameters

docs/user/ppl/cmd/top.md

@@ -10,7 +10,7 @@ The `top` command finds the most common combination of values across all fields
 The `top` command has the following syntax:
 
 ```syntax
-top [N] [top-options] <field-list> [by-clause]
+top [limit=][N] [top-options] <field-list> [by-clause]
 ```
 
 ## Parameters

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -380,7 +380,7 @@ fieldformatCommand
    ;
 
 headCommand
-   : HEAD (number = integerLiteral)? (FROM from = integerLiteral)?
+   : HEAD ((LIMIT EQUAL)? number = integerLiteral)? (FROM from = integerLiteral)?
    ;
 
 binCommand
@@ -412,7 +412,7 @@ logSpanValue
    ;
 
 rareTopCommand
-   : (TOP | RARE) (number = integerLiteral)? rareTopOption* fieldList (byClause)?
+   : (TOP | RARE) ((LIMIT EQUAL)? number = integerLiteral)? rareTopOption* fieldList (byClause)?
    ;
 
 rareTopOption
@@ -1656,6 +1656,7 @@ searchableKeyWord
    | FREQUENCY_THRESHOLD_PERCENTAGE
    | MAX_SAMPLE_COUNT
    | BUFFER_LIMIT
+   | LIMIT
    | SHOW_NUMBERED_TOKEN
    | WITH
    | REGEX

ppl/src/test/java/org/opensearch/sql/ppl/AstPlanningTestBase.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 AstPlanningTestBase {
+  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

@@ -65,7 +65,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.AllFields;
@@ -82,20 +81,15 @@
 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.exception.SemanticCheckException;
-import org.opensearch.sql.ppl.antlr.PPLSyntaxParser;
+import org.opensearch.sql.ppl.AstPlanningTestBase;
 import org.opensearch.sql.utils.SystemIndexUtils;
 
-public class AstBuilderTest {
+public class AstBuilderTest extends AstPlanningTestBase {
 
   @Rule public ExpectedException exceptionRule = ExpectedException.none();
 
-  private final Settings settings = Mockito.mock(Settings.class);
-
-  private final PPLSyntaxParser parser = new PPLSyntaxParser();
-
   @Test
   public void testDynamicSourceClauseThrowsUnsupportedException() {
     String query = "source=[myindex, logs, fieldIndex=\"test\"]";
@@ -1405,11 +1399,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.AstPlanningTestBase;
+
+public class AstEquivalenceTest extends AstPlanningTestBase {
+  @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"));
+  }
+}

…arch/sql/ppl/utils/SPathRewriteTest.java …ch/sql/ppl/rewrite/SpathRewriteTest.javappl/src/test/java/org/opensearch/sql/ppl/utils/SPathRewriteTest.java renamed to ppl/src/test/java/org/opensearch/sql/ppl/rewrite/SpathRewriteTest.java ppl/src/test/java/org/opensearch/sql/ppl/utils/SPathRewriteTest.java renamed to ppl/src/test/java/org/opensearch/sql/ppl/rewrite/SpathRewriteTest.java

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.ppl.utils;
+package org.opensearch.sql.ppl.rewrite;
 
 import static org.junit.Assert.assertEquals;
 import static org.opensearch.sql.ast.dsl.AstDSL.eval;
@@ -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.AstPlanningTestBase;
 
+public class SpathRewriteTest extends AstPlanningTestBase {
   // Control test to make sure something fundamental hasn't changed about the json_extract parsing
   @Test
   public void testEvalControl() {