[Bug Fix] Explicit head command should override fetch_size in PPL

When a user's PPL query contains an explicit head command, the
fetch_size parameter should not inject an additional Head node on top.
Previously, fetch_size always wrapped the plan with Head(fetchSize),
creating a double-Head that could cap results unexpectedly (e.g.,
head 100 with fetch_size=5 would only return 5 rows).

This adds a containsHead() check in AstStatementBuilder to skip the
Head injection when the user's query already has a head command,
letting the user's explicit intent take precedence.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Kai Huang <ahkcs@amazon.com>

Author: ahkcs

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

@@ -172,9 +172,9 @@ public void testFetchSizeWithStats() throws IOException {
   }
 
   @Test
-  public void testFetchSizeWithHead() throws IOException {
-    // Both head command and fetch_size - the smaller limit should win
-    // head 3 limits to 3, fetch_size 10 would allow 10, so we get 3
+  public void testHeadOverridesFetchSizeWhenSmaller() throws IOException {
+    // Explicit head takes precedence over fetch_size
+    // head 3 returns 3 rows regardless of fetch_size=10
     JSONObject result =
         executeQueryWithFetchSize(
             String.format("source=%s | head 3 | fields firstname", TEST_INDEX_ACCOUNT), 10);
@@ -183,16 +183,59 @@ public void testFetchSizeWithHead() throws IOException {
   }
 
   @Test
-  public void testFetchSizeSmallerThanHead() throws IOException {
-    // fetch_size smaller than head - fetch_size should further limit
-    // head 100 would return 100, but fetch_size 5 limits to 5
+  public void testHeadOverridesFetchSizeWhenLarger() throws IOException {
+    // Explicit head takes precedence over fetch_size
+    // head 100 should return 100 rows even though fetch_size=5
     JSONObject result =
         executeQueryWithFetchSize(
             String.format("source=%s | head 100 | fields firstname", TEST_INDEX_ACCOUNT), 5);
     JSONArray dataRows = result.getJSONArray("datarows");
+    assertEquals(100, dataRows.length());
+  }
+
+  @Test
+  public void testHeadOverridesFetchSizeWithOffset() throws IOException {
+    // Explicit head with offset takes precedence over fetch_size
+    // head 3 from 2 should skip 2 rows and return 3 rows, ignoring fetch_size=100
+    JSONObject result =
+        executeQueryWithFetchSize(
+            String.format("source=%s | head 3 from 2 | fields firstname", TEST_INDEX_ACCOUNT), 100);
+    JSONArray dataRows = result.getJSONArray("datarows");
+    assertEquals(3, dataRows.length());
+  }
+
+  @Test
+  public void testHeadOverridesFetchSizeWithFilter() throws IOException {
+    // Explicit head after filter takes precedence over fetch_size
+    // Even with fetch_size=2, head 5 should return 5 matching rows
+    JSONObject result =
+        executeQueryWithFetchSize(
+            String.format(
+                "source=%s | where age > 30 | head 5 | fields firstname, age", TEST_INDEX_ACCOUNT),
+            2);
+    JSONArray dataRows = result.getJSONArray("datarows");
     assertEquals(5, dataRows.length());
   }
 
+  @Test
+  public void testHeadEqualToFetchSize() throws IOException {
+    // When head and fetch_size are the same value, head takes precedence (no double-Head)
+    JSONObject result =
+        executeQueryWithFetchSize(
+            String.format("source=%s | head 7 | fields firstname", TEST_INDEX_ACCOUNT), 7);
+    JSONArray dataRows = result.getJSONArray("datarows");
+    assertEquals(7, dataRows.length());
+  }
+
+  @Test
+  public void testHeadLargerThanDatasetWithFetchSize() throws IOException {
+    // head 1000 on a 7-row index with fetch_size=3: head takes precedence, returns all 7 rows
+    JSONObject result =
+        executeQueryWithFetchSize(String.format("source=%s | head 1000", TEST_INDEX_BANK), 3);
+    JSONArray dataRows = result.getJSONArray("datarows");
+    assertEquals(7, dataRows.length());
+  }
+
   @Test
   public void testFetchSizeAsUrlParameter() throws IOException {
     // fetch_size specified as URL parameter instead of JSON body

ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java

@@ -32,7 +32,7 @@ public class AstStatementBuilder extends OpenSearchPPLParserBaseVisitor<Statemen
   @Override
   public Statement visitPplStatement(OpenSearchPPLParser.PplStatementContext ctx) {
     UnresolvedPlan rawPlan = astBuilder.visit(ctx);
-    if (context.getFetchSize() > 0) {
+    if (context.getFetchSize() > 0 && !containsHead(rawPlan)) {
       rawPlan = new Head(context.getFetchSize(), 0).attach(rawPlan);
     }
     UnresolvedPlan plan = addSelectAll(rawPlan);
@@ -69,6 +69,23 @@ public static class StatementBuilderContext {
     private final String explainMode;
   }
 
+  /**
+   * Recursively checks if the AST contains a {@link Head} node. When the user's query already
+   * includes an explicit {@code head} command, we should not inject an additional Head for
+   * fetch_size so that the user's explicit limit takes precedence.
+   */
+  private boolean containsHead(UnresolvedPlan plan) {
+    if (plan instanceof Head) {
+      return true;
+    }
+    for (var child : plan.getChild()) {
+      if (child instanceof UnresolvedPlan && containsHead((UnresolvedPlan) child)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   private UnresolvedPlan addSelectAll(UnresolvedPlan plan) {
     if ((plan instanceof Project) && !((Project) plan).isExcluded()) {
       return plan;

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

@@ -87,33 +87,31 @@ public void buildQueryStatementWithLargeFetchSize() {
   @Test
   public void buildQueryStatementWithFetchSizeAndSmallerHead() {
     // User query has head 3, fetchSize=10
-    // Head(10) wraps Head(3), then Project(*) wraps on top
-    // The inner head 3 limits first, so only 3 rows are returned
+    // Explicit head takes precedence over fetch_size, so no outer Head(10) is injected
     assertEqualWithFetchSize(
         "source=t | head 3",
         10,
-        new Query(project(head(head(relation("t"), 3, 0), 10, 0), AllFields.of()), 0, PPL));
+        new Query(project(head(relation("t"), 3, 0), AllFields.of()), 0, PPL));
   }
 
   @Test
   public void buildQueryStatementWithFetchSizeSmallerThanHead() {
     // User query has head 100, fetchSize=5
-    // Head(5) wraps Head(100), then Project(*) wraps on top
-    // The outer head 5 limits, so only 5 rows are returned
+    // Explicit head takes precedence over fetch_size, so no outer Head(5) is injected
     assertEqualWithFetchSize(
         "source=t | head 100",
         5,
-        new Query(project(head(head(relation("t"), 100, 0), 5, 0), AllFields.of()), 0, PPL));
+        new Query(project(head(relation("t"), 100, 0), AllFields.of()), 0, PPL));
   }
 
   @Test
   public void buildQueryStatementWithFetchSizeAndHeadWithOffset() {
     // User query has head 3 from 1 (with offset), fetchSize=10
-    // The inner head offset is preserved, outer Head always has offset 0
+    // Explicit head takes precedence over fetch_size, so no outer Head(10) is injected
     assertEqualWithFetchSize(
         "source=t | head 3 from 1",
         10,
-        new Query(project(head(head(relation("t"), 3, 1), 10, 0), AllFields.of()), 0, PPL));
+        new Query(project(head(relation("t"), 3, 1), AllFields.of()), 0, PPL));
   }
 
   private void assertEqualWithFetchSize(String query, int fetchSize, Statement expectedStatement) {