Fix multiple limit pushdown with shirinking caused by offset

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

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

@@ -164,6 +164,18 @@ public void testMultipleLimitExplain() throws Exception {
                 + "| head 5 "
                 + "| fields age"));
 
+    String expected10from1then10from2 =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_limit_10from1_10from2_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_10from1_10from2_push.json");
+    assertJsonEqualsIgnoreId(
+        expected10from1then10from2,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| head 10 from 1 "
+                + "| head 10 from 2 "
+                + "| fields age"));
+
     // The second limit should not be pushed down for limit-filter-limit queries
     String expected10ThenFilterThen5 =
         isCalciteEnabled()

integ-test/src/test/resources/expectedOutput/calcite/explain_limit_10from1_10from2_push.json

@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "logical": "LogicalProject(age=[$8])\n  LogicalSort(offset=[2], fetch=[10])\n    LogicalSort(offset=[1], fetch=[10])\n      CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical": "CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[LIMIT->10, LIMIT->10, PROJECT->[age]], OpenSearchRequestBuilder(sourceBuilder={\"from\":3,\"size\":8,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, requestedTotalSize=8, pageSize=null, startFrom=3)])\n"
+  }
+}

integ-test/src/test/resources/expectedOutput/ppl/explain_limit_10from1_10from2_push.json

@@ -0,0 +1,17 @@
+{
+  "root": {
+    "name": "ProjectOperator",
+    "description": {
+      "fields": "[age]"
+    },
+    "children": [
+      {
+        "name": "OpenSearchIndexScan",
+        "description": {
+          "request": "OpenSearchQueryRequest(indexName=opensearch-sql_test_index_account, sourceBuilder={\"from\":3,\"size\":8,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, needClean=true, searchDone=false, pitId=null, cursorKeepAlive=null, searchAfter=null, searchResponse=null)"
+        },
+        "children": []
+      }
+    ]
+  }
+}

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java

@@ -215,7 +215,10 @@ public void pushDownLimit(Integer limit, Integer offset) {
     // E.g. for `source=t | head 10 | head 5`, we take 5
     // This also ensures that the limit won't exceed the initial default value. (set to
     // Settings.Key.QUERY_SIZE_LIMIT in OpenSearchIndex)
-    requestedTotalSize = Math.min(limit, requestedTotalSize);
+    // Besides, there may be cases when the existing requestedTotalSize does not satisfy the
+    // new limit and offset. E.g. for `head 11 | head 10 from 2`, the new requested total size
+    // is 9. We need to adjust it accordingly.
+    requestedTotalSize = Math.min(limit, requestedTotalSize - offset);
     // If there are multiple offset, we aggregate the offset
     // E.g. for `head 10 from 1 | head 5 from 2` equals to `head 5 from 3`
     startFrom += offset;