Support PPL queries when having trailing pipes and/or empty pipes

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

@@ -0,0 +1,172 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ppl;
+
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ACCOUNT;
+
+import java.io.IOException;
+import org.json.JSONObject;
+import org.junit.jupiter.api.Test;
+
+/** Integration tests for trailing pipe and middle empty pipe functionality in PPL queries. */
+public class TrailingPipeIT extends PPLIntegTestCase {
+
+  /**
+   * Initializes the test environment by loading the account index.
+   *
+   * @throws Exception if initialization fails
+   */
+  @Override
+  public void init() throws Exception {
+    super.init();
+    loadIndex(Index.ACCOUNT);
+  }
+
+  /**
+   * Tests that a trailing pipe after a source command produces identical results to a query without
+   * the trailing pipe.
+   *
+   * @throws IOException if query execution fails
+   */
+  @Test
+  public void testTrailingPipeAfterSource() throws IOException {
+    // Query with trailing pipe should produce same results as without
+    JSONObject resultWithout = executeQuery(String.format("source=%s", TEST_INDEX_ACCOUNT));
+    JSONObject resultWith = executeQuery(String.format("source=%s |", TEST_INDEX_ACCOUNT));
+
+    // Both should return the same data
+    assertTrue(resultWithout.similar(resultWith));
+  }
+
+  /**
+   * Tests that a trailing pipe after a fields command produces identical results to a query without
+   * the trailing pipe.
+   *
+   * @throws IOException if query execution fails
+   */
+  @Test
+  public void testTrailingPipeAfterFields() throws IOException {
+    JSONObject resultWithout =
+        executeQuery(
+            String.format(
+                "source=%s | where age > 30 | fields firstname, age", TEST_INDEX_ACCOUNT));
+    JSONObject resultWith =
+        executeQuery(
+            String.format(
+                "source=%s | where age > 30 | fields firstname, age |", TEST_INDEX_ACCOUNT));
+
+    assertTrue(resultWithout.similar(resultWith));
+  }
+
+  /**
+   * Tests that a trailing pipe after a head command produces identical results to a query without
+   * the trailing pipe.
+   *
+   * @throws IOException if query execution fails
+   */
+  @Test
+  public void testTrailingPipeAfterHead() throws IOException {
+    JSONObject resultWithout =
+        executeQuery(
+            String.format("source=%s | fields firstname, age | head 3", TEST_INDEX_ACCOUNT));
+    JSONObject resultWith =
+        executeQuery(
+            String.format("source=%s | fields firstname, age | head 3 |", TEST_INDEX_ACCOUNT));
+
+    assertTrue(resultWithout.similar(resultWith));
+  }
+
+  /**
+   * Tests that a trailing pipe after a complex query with multiple commands (where, fields, stats,
+   * sort) produces identical results to a query without the trailing pipe.
+   *
+   * @throws IOException if query execution fails
+   */
+  @Test
+  public void testTrailingPipeWithComplexQuery() throws IOException {
+    JSONObject resultWithout =
+        executeQuery(
+            String.format(
+                "source=%s | where age > 25 | fields firstname, age, state | stats avg(age) by"
+                    + " state | sort state",
+                TEST_INDEX_ACCOUNT));
+    JSONObject resultWith =
+        executeQuery(
+            String.format(
+                "source=%s | where age > 25 | fields firstname, age, state | stats avg(age) by"
+                    + " state | sort state |",
+                TEST_INDEX_ACCOUNT));
+
+    assertTrue(resultWithout.similar(resultWith));
+  }
+
+  /**
+   * Tests that an empty pipe in the middle of a query pipeline is properly ignored and produces
+   * identical results to a query without the empty pipe.
+   *
+   * @throws IOException if query execution fails
+   */
+  @Test
+  public void testEmptyPipeInMiddle() throws IOException {
+    // Empty pipe in middle should be ignored
+    JSONObject resultNormal =
+        executeQuery(
+            String.format(
+                "source=%s | where age > 30 | fields firstname, age", TEST_INDEX_ACCOUNT));
+    JSONObject resultWithEmpty =
+        executeQuery(
+            String.format(
+                "source=%s | | where age > 30 | fields firstname, age", TEST_INDEX_ACCOUNT));
+
+    assertTrue(resultNormal.similar(resultWithEmpty));
+  }
+
+  /**
+   * Tests that multiple empty pipes scattered throughout a query pipeline are properly ignored and
+   * produce identical results to a query without the empty pipes.
+   *
+   * @throws IOException if query execution fails
+   */
+  @Test
+  public void testMultipleEmptyPipes() throws IOException {
+    // Multiple empty pipes should be ignored
+    JSONObject resultNormal =
+        executeQuery(
+            String.format(
+                "source=%s | where age > 30 | fields firstname, age | sort age",
+                TEST_INDEX_ACCOUNT));
+    JSONObject resultWithEmpty =
+        executeQuery(
+            String.format(
+                "source=%s | | where age > 30 | | fields firstname, age | sort age",
+                TEST_INDEX_ACCOUNT));
+
+    assertTrue(resultNormal.similar(resultWithEmpty));
+  }
+
+  /**
+   * Tests that a combination of empty pipes in the middle and a trailing pipe at the end are
+   * properly handled and produce identical results to a query without these extraneous pipes.
+   *
+   * @throws IOException if query execution fails
+   */
+  @Test
+  public void testEmptyPipesAndTrailingPipe() throws IOException {
+    // Multiple empty pipes should be ignored
+    JSONObject resultNormal =
+        executeQuery(
+            String.format(
+                "source=%s | where age > 30 | fields firstname, age | sort age",
+                TEST_INDEX_ACCOUNT));
+    JSONObject resultWithEmpty =
+        executeQuery(
+            String.format(
+                "source=%s | | where age > 30 | fields firstname, age | sort age |",
+                TEST_INDEX_ACCOUNT));
+
+    assertTrue(resultNormal.similar(resultWithEmpty));
+  }
+}

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -20,11 +20,11 @@ pplStatement
    ;
 
 subPipeline
-   : PIPE? commands (PIPE commands)*
+   : PIPE? commands (PIPE commands?)*
    ;
 
 queryStatement
-   : (PIPE)? pplCommands (PIPE commands)*
+   : (PIPE)? pplCommands (PIPE commands?)*
    ;
 
 explainStatement
@@ -39,7 +39,7 @@ explainMode
     ;
 
 subSearch
-   : searchCommand (PIPE commands)*
+   : searchCommand (PIPE commands?)*
    ;
 
 // commands

ppl/src/test/java/org/opensearch/sql/ppl/antlr/PPLSyntaxParserTest.java

@@ -946,4 +946,55 @@ public void testAddTotalsCommandWithAllOptionsShouldPass() {
             .parse("source=t | addtotals price, quantity label='Total' labelfield='type'");
     assertNotEquals(null, tree);
   }
+
+  @Test
+  public void testQueryWithMultipleTrailingPipesShouldPass() {
+    // Multiple consecutive trailing pipes should be handled gracefully
+    ParseTree tree = new PPLSyntaxParser().parse("search source=t a=1 b=2 | fields a,b | |");
+    assertNotEquals(null, tree);
+  }
+
+  @Test
+  public void testQueryWithTrailingPipeAndWhitespaceShouldPass() {
+    ParseTree tree = new PPLSyntaxParser().parse("search source=t a=1 b=2 | fields a,b |   ");
+    assertNotEquals(null, tree);
+  }
+
+  @Test
+  public void testQueryWithMiddleEmptyPipe() {
+    ParseTree tree = new PPLSyntaxParser().parse("search source=t a=1 b=2 | | fields a,b");
+    assertNotEquals(null, tree);
+  }
+
+  @Test
+  public void testQueryWithMiddleEmptyPipeAndTrailingPipe() {
+    ParseTree tree = new PPLSyntaxParser().parse("search source=t a=1 b=2 | | fields a,b |   ");
+    assertNotEquals(null, tree);
+  }
+
+  @Test
+  public void testComplexQueryWithTrailingPipeShouldPass() {
+    ParseTree tree =
+        new PPLSyntaxParser()
+            .parse("source=t | where x > 5 | stats count() by status | sort -count |");
+    assertNotEquals(null, tree);
+  }
+
+  @Test
+  public void testSubSearchWithTrailingPipeShouldPass() {
+    ParseTree tree =
+        new PPLSyntaxParser().parse("source=outer | join a [source=inner | fields x,y |]");
+    assertNotEquals(null, tree);
+  }
+
+  /**
+   * Tests that the parser correctly rejects queries with invalid command tokens after a pipe,
+   * ensuring proper error detection for malformed queries.
+   */
+  @Test
+  public void testPipeWithInvalidCommandShouldFail() {
+    assertThrows(
+        SyntaxCheckException.class,
+        () -> new PPLSyntaxParser().parse("source=t | | 123invalidcommand"));
+  }
 }

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

@@ -1726,4 +1726,86 @@ public void testGraphLookupCommand() {
         SemanticCheckException.class,
         () -> plan("source=t | graphLookup employees fromField=manager as reportingHierarchy"));
   }
+
+  @Test
+  public void testTrailingPipeAfterSource() {
+    // Test that trailing pipe after source produces same AST
+    assertEqual("source=t |", relation("t"));
+  }
+
+  @Test
+  public void testTrailingPipeAfterStats() {
+    // Test trailing pipe after stats command
+    assertEqual(
+        "source=t | stats count(a) by b |",
+        agg(
+            relation("t"),
+            exprList(alias("count(a)", aggregate("count", field("a")))),
+            emptyList(),
+            exprList(alias("b", field("b"))),
+            defaultStatsArgs()));
+  }
+
+  @Test
+  public void testTrailingPipeWithComplexQuery() {
+    // Test trailing pipe with complex query including where, stats, and sort
+    assertEqual(
+        "source=t | where a > 1 | stats count(b) by c | sort c |",
+        sort(
+            agg(
+                filter(relation("t"), compare(">", field("a"), intLiteral(1))),
+                exprList(alias("count(b)", aggregate("count", field("b")))),
+                emptyList(),
+                exprList(alias("c", field("c"))),
+                defaultStatsArgs()),
+            field("c", defaultSortFieldArgs())));
+  }
+
+  @Test
+  public void testEmptyPipeAfterSource() {
+    // Test that empty pipe after source is ignored
+    assertEqual("source=t | |", relation("t"));
+  }
+
+  @Test
+  public void testEmptyPipeInMiddle() {
+    // Test that empty pipe in middle is ignored
+    assertEqual(
+        "source=t | | where a=1", filter(relation("t"), compare("=", field("a"), intLiteral(1))));
+  }
+
+  @Test
+  public void testMultipleEmptyPipes() {
+    // Test multiple empty pipes are ignored
+    assertEqual(
+        "source=t | | where a=1 | | fields b | |",
+        projectWithArg(
+            filter(relation("t"), compare("=", field("a"), intLiteral(1))),
+            defaultFieldsArgs(),
+            field("b")));
+  }
+
+  /**
+   * Tests that a combination of empty pipes in the middle and a trailing pipe at the end are
+   * properly handled and produce the same AST as a query without these extraneous pipes.
+   */
+  @Test
+  public void testEmptyPipeAndTrailingPipeTogether() {
+    // Test both empty pipe in middle and trailing pipe at end
+    assertEqual(
+        "source=t | | where a=1 | fields b |",
+        projectWithArg(
+            filter(relation("t"), compare("=", field("a"), intLiteral(1))),
+            defaultFieldsArgs(),
+            field("b")));
+  }
+
+  /**
+   * Tests that the parser correctly rejects queries with invalid command tokens after a pipe,
+   * ensuring proper error detection for malformed queries.
+   */
+  @Test(expected = org.opensearch.sql.common.antlr.SyntaxCheckException.class)
+  public void testMalformedPipeProducesSyntaxError() {
+    plan("source=t | invalidCmd |");
+  }
 }