fix: Honor PPL fetch_size on the analytics-engine route (#5567)

* fix: Honor PPL fetch_size on the analytics-engine route

PPL's fetch_size caps a response to N rows (no cursor) — the V2 path lowers it
to a top-level `head N` in AstStatementBuilder.visitPplStatement. The
analytics-engine route bypasses that builder: TransportPPLQueryAction forwards
only the query string to RestUnifiedQueryAction.execute(), so fetch_size was
dropped and the engine returned the full result set.

Thread the request's fetchSize through execute() and apply an equivalent
top-level limit on the planned RelNode (addFetchSizeLimit), using the same
relBuilder.limit primitive that `head` lowers to. fetchSize <= 0 keeps the
prior "system default" behavior; the SQL path (separate cursor-based fetch_size)
is unchanged.

Before/after (CalcitePPLFetchSizeIT-equivalent, analytics-engine route):
  before:  9/19 pass  (10 fail — fetch_size ignored, full set returned)
  after:  19/19 pass

Signed-off-by: Kai Huang <ahkcs@amazon.com>

* chore: spotlessApply on RestUnifiedQueryAction (javadoc reflow + signature wrap)

Signed-off-by: Kai Huang <ahkcs@amazon.com>

---------

Signed-off-by: Kai Huang <ahkcs@amazon.com>

Author: ahkcs

plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java

@@ -152,24 +152,26 @@ public void execute(
       QueryType queryType,
       boolean profiling,
       ActionListener<TransportPPLQueryResponse> listener) {
-    doExecute(query, queryType, profiling, null, listener);
+    doExecute(query, queryType, profiling, 0, null, listener);
   }
 
   /** Execute linked to {@code parentTask} so a front-end cancel propagates into the engine. */
   public void execute(
       String query,
       QueryType queryType,
       boolean profiling,
+      int fetchSize,
       Task parentTask,
       ActionListener<TransportPPLQueryResponse> listener) {
     assert parentTask != null : "parentTask required for cancellation propagation";
-    doExecute(query, queryType, profiling, parentTask, listener);
+    doExecute(query, queryType, profiling, fetchSize, parentTask, listener);
   }
 
   private void doExecute(
       String query,
       QueryType queryType,
       boolean profiling,
+      int fetchSize,
       Task parentTask,
       ActionListener<TransportPPLQueryResponse> listener) {
     client
@@ -193,6 +195,10 @@ private void doExecute(
                     UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
                     RelNode plan = planner.plan(query);
                     CalcitePlanContext planContext = context.getPlanContext();
+                    // PPL fetch_size caps the response to N rows (no cursor) — the V2 path attaches
+                    // a `head N` in AstStatementBuilder; the unified path parses only the query
+                    // string, so apply the equivalent top-level limit here before the system cap.
+                    plan = addFetchSizeLimit(plan, planContext, fetchSize);
                     plan = addQuerySizeLimit(plan, planContext);
                     if (profiling) {
                       analyticsEngine.executeWithProfile(
@@ -339,6 +345,21 @@ private static RelNode addQuerySizeLimit(RelNode plan, CalcitePlanContext contex
         context.relBuilder.literal(context.sysLimit.querySizeLimit()));
   }
 
+  /**
+   * Cap the result to {@code fetchSize} rows when {@code fetchSize > 0}, mirroring PPL's {@code
+   * fetch_size} (which the V2 path lowers to a top-level {@code head N} in AstStatementBuilder).
+   * {@code fetchSize <= 0} means "use system default", so no limit is added. Uses the same {@code
+   * relBuilder.limit} primitive that {@code head} lowers to, so the analytics backend sees an
+   * ordinary fetch limit.
+   */
+  private static RelNode addFetchSizeLimit(
+      RelNode plan, CalcitePlanContext context, int fetchSize) {
+    if (fetchSize <= 0) {
+      return plan;
+    }
+    return context.relBuilder.push(plan).limit(0, fetchSize).build();
+  }
+
   private ResponseListener<QueryResponse> createQueryListener(
       QueryType queryType, ActionListener<TransportPPLQueryResponse> transportListener) {
     ResponseFormatter<QueryResult> formatter = new SimpleJsonResponseFormatter(PRETTY);

plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java

@@ -189,6 +189,7 @@ protected void doExecute(
             transformedRequest.getRequest(),
             QueryType.PPL,
             transformedRequest.profile(),
+            transformedRequest.getFetchSize(),
             task,
             clearingListener);
       }