Add query size limit to RelNode plan instead of post-processing

Add LogicalSystemLimit to the RelNode plan before passing it to the
analytics engine, consistent with PPL V3 (QueryService.convertToCalcitePlan).
This ensures the analytics engine enforces the limit during execution
rather than returning all rows for post-processing truncation.

Remove post-processing querySizeLimit truncation from
AnalyticsExecutionEngine -- the limit is now part of the plan the
executor receives.

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

Author: ahkcs

core/src/main/java/org/opensearch/sql/executor/analytics/AnalyticsExecutionEngine.java

@@ -65,11 +65,10 @@ public void explain(PhysicalPlan plan, ResponseListener<ExplainResponse> listene
   public void execute(
       RelNode plan, CalcitePlanContext context, ResponseListener<QueryResponse> listener) {
     try {
-      Integer querySizeLimit = context.sysLimit.querySizeLimit();
       Iterable<Object[]> rows = planExecutor.execute(plan, null);
 
       List<RelDataTypeField> fields = plan.getRowType().getFieldList();
-      List<ExprValue> results = convertRows(rows, fields, querySizeLimit);
+      List<ExprValue> results = convertRows(rows, fields);
       Schema schema = buildSchema(fields);
 
       listener.onResponse(new QueryResponse(schema, results, Cursor.None));
@@ -78,13 +77,9 @@ public void execute(
     }
   }
 
-  private List<ExprValue> convertRows(
-      Iterable<Object[]> rows, List<RelDataTypeField> fields, Integer querySizeLimit) {
+  private List<ExprValue> convertRows(Iterable<Object[]> rows, List<RelDataTypeField> fields) {
     List<ExprValue> results = new ArrayList<>();
     for (Object[] row : rows) {
-      if (querySizeLimit != null && results.size() >= querySizeLimit) {
-        break;
-      }
       Map<String, ExprValue> valueMap = new LinkedHashMap<>();
       for (int i = 0; i < fields.size(); i++) {
         String columnName = fields.get(i).getName();

core/src/test/java/org/opensearch/sql/executor/analytics/AnalyticsExecutionEngineTest.java

@@ -12,10 +12,8 @@
 import static org.mockito.Mockito.when;
 
 import java.lang.reflect.Field;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import org.apache.calcite.rel.RelNode;
@@ -148,24 +146,8 @@ void executeRelNode_temporalTypes() {
     assertEquals(0, response.getResults().size(), "Should have 0 rows. " + dump);
   }
 
-  @Test
-  void executeRelNode_querySizeLimit() throws Exception {
-    RelNode relNode = mockRelNode("id", SqlTypeName.INTEGER);
-    List<Object[]> manyRows = new ArrayList<>();
-    for (int i = 0; i < 100; i++) {
-      manyRows.add(new Object[] {i});
-    }
-    when(mockExecutor.execute(relNode, null)).thenReturn(manyRows);
-    setSysLimit(mockContext, new SysLimit(10, 10000, 50000));
-
-    QueryResponse response = executeAndCapture(relNode);
-    String dump = dumpResponse(response);
-
-    assertEquals(
-        10,
-        response.getResults().size(),
-        "Should truncate to querySizeLimit=10, got " + response.getResults().size() + ". " + dump);
-  }
+  // Query size limit is now enforced in the RelNode plan (LogicalSystemLimit) before it reaches
+  // AnalyticsExecutionEngine. The engine trusts the executor to honor the limit.
 
   @Test
   void executeRelNode_emptyResults() {

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

@@ -24,6 +24,7 @@
 import org.opensearch.sql.api.UnifiedQueryContext;
 import org.opensearch.sql.api.UnifiedQueryPlanner;
 import org.opensearch.sql.calcite.CalcitePlanContext;
+import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.datasources.exceptions.DataSourceClientException;
@@ -129,13 +130,18 @@ private void doExecuteViaTransport(
 
         UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
         RelNode plan = planner.plan(query);
+
+        // Add query size limit to the plan so the analytics engine can enforce it
+        // during execution, consistent with PPL V3 (see QueryService.convertToCalcitePlan)
+        CalcitePlanContext planContext = context.getPlanContext();
+        plan = addQuerySizeLimit(plan, planContext);
+
         long planTime = System.nanoTime();
         LOG.info(
             "[unified] Planning completed in {}ms for {} query",
             (planTime - startTime) / 1_000_000,
             queryType);
 
-        CalcitePlanContext planContext = context.getPlanContext();
         analyticsEngine.execute(
             plan, planContext, createTransportQueryListener(queryType, planTime, listener));
       }
@@ -144,6 +150,17 @@ private void doExecuteViaTransport(
     }
   }
 
+  /**
+   * Add a system-level query size limit to the plan. This ensures the analytics engine enforces the
+   * limit during execution rather than returning all rows for post-processing truncation.
+   */
+  private static RelNode addQuerySizeLimit(RelNode plan, CalcitePlanContext context) {
+    return LogicalSystemLimit.create(
+        LogicalSystemLimit.SystemLimitType.QUERY_SIZE_LIMIT,
+        plan,
+        context.relBuilder.literal(context.sysLimit.querySizeLimit()));
+  }
+
   private ResponseListener<QueryResponse> createTransportQueryListener(
       QueryType queryType,
       long planEndTime,
@@ -188,13 +205,16 @@ private void doExecute(String query, QueryType queryType, RestChannel channel) {
 
         UnifiedQueryPlanner planner = new UnifiedQueryPlanner(context);
         RelNode plan = planner.plan(query);
+
+        CalcitePlanContext planContext = context.getPlanContext();
+        plan = addQuerySizeLimit(plan, planContext);
+
         long planTime = System.nanoTime();
         LOG.info(
             "[unified] Planning completed in {}ms for {} query",
             (planTime - startTime) / 1_000_000,
             queryType);
 
-        CalcitePlanContext planContext = context.getPlanContext();
         analyticsEngine.execute(
             plan, planContext, createQueryListener(queryType, channel, planTime));
       }