Accumulate offsets when pushing down limit

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

Author: yuancu

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

@@ -142,8 +142,8 @@ public void testLimitWithFilterPushdownExplain() throws Exception {
   public void testMultipleLimitExplain() throws Exception {
     String expected5Then10 =
         isCalciteEnabled()
-            ? loadFromFile("expectedOutput/calcite/explain_limit_5_10.json")
-            : loadFromFile("expectedOutput/ppl/explain_limit_5_10.json");
+            ? loadFromFile("expectedOutput/calcite/explain_limit_5_10_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_5_10_push.json");
     assertJsonEqualsIgnoreRelId(
         expected5Then10,
         explainQueryToString(
@@ -154,8 +154,8 @@ public void testMultipleLimitExplain() throws Exception {
 
     String expected10Then5 =
         isCalciteEnabled()
-            ? loadFromFile("expectedOutput/calcite/explain_limit_10_5.json")
-            : loadFromFile("expectedOutput/ppl/explain_limit_10_5.json");
+            ? loadFromFile("expectedOutput/calcite/explain_limit_10_5_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_10_5_push.json");
     assertJsonEqualsIgnoreRelId(
         expected10Then5,
         explainQueryToString(
@@ -167,8 +167,8 @@ public void testMultipleLimitExplain() throws Exception {
     // The second limit should not be pushed down for limit-filter-limit queries
     String expected10ThenFilterThen5 =
         isCalciteEnabled()
-            ? loadFromFile("expectedOutput/calcite/explain_limit_10_filter_5.json")
-            : loadFromFile("expectedOutput/ppl/explain_limit_10_filter_5.json");
+            ? loadFromFile("expectedOutput/calcite/explain_limit_10_filter_5_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_10_filter_5_push.json");
     assertJsonEqualsIgnoreRelId(
         expected10ThenFilterThen5,
         explainQueryToString(
@@ -179,6 +179,22 @@ public void testMultipleLimitExplain() throws Exception {
                 + "| fields age"));
   }
 
+  @Test
+  public void testLimitWithMultipleOffsetPushdownExplain() throws Exception {
+    String expected =
+        isCalciteEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_limit_offsets_push.json")
+            : loadFromFile("expectedOutput/ppl/explain_limit_offsets_push.json");
+
+    assertJsonEqualsIgnoreRelId(
+        expected,
+        explainQueryToString(
+            "source=opensearch-sql_test_index_account"
+                + "| head 10 from 1 "
+                + "| head 5 from 2 "
+                + "| fields age"));
+  }
+
   @Test
   public void testFillNullPushDownExplain() throws Exception {
     String expected = loadFromFile("expectedOutput/ppl/explain_fillnull_push.json");

…edOutput/calcite/explain_limit_10_5.json …put/calcite/explain_limit_10_5_push.jsoninteg-test/src/test/resources/expectedOutput/calcite/explain_limit_10_5.json renamed to integ-test/src/test/resources/expectedOutput/calcite/explain_limit_10_5_push.json integ-test/src/test/resources/expectedOutput/calcite/explain_limit_10_5.json renamed to integ-test/src/test/resources/expectedOutput/calcite/explain_limit_10_5_push.json

@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "logical": "LogicalProject(age=[$8])\n  LogicalSort(offset=[2], fetch=[5])\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->5, PROJECT->[age]], OpenSearchRequestBuilder(sourceBuilder={\"from\":3,\"size\":5,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, requestedTotalSize=5, pageSize=null, startFrom=3)])\n"
+  }
+}

…t/calcite/explain_limit_10_filter_5.json …cite/explain_limit_10_filter_5_push.jsoninteg-test/src/test/resources/expectedOutput/calcite/explain_limit_10_filter_5.json renamed to integ-test/src/test/resources/expectedOutput/calcite/explain_limit_10_filter_5_push.json integ-test/src/test/resources/expectedOutput/calcite/explain_limit_10_filter_5.json renamed to integ-test/src/test/resources/expectedOutput/calcite/explain_limit_10_filter_5_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\":5,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, needClean=true, searchDone=false, pitId=null, cursorKeepAlive=null, searchAfter=null, searchResponse=null)"
+        },
+        "children": []
+      }
+    ]
+  }
+}

…edOutput/calcite/explain_limit_5_10.json …put/calcite/explain_limit_5_10_push.jsoninteg-test/src/test/resources/expectedOutput/calcite/explain_limit_5_10.json renamed to integ-test/src/test/resources/expectedOutput/calcite/explain_limit_5_10_push.json integ-test/src/test/resources/expectedOutput/calcite/explain_limit_5_10.json renamed to integ-test/src/test/resources/expectedOutput/calcite/explain_limit_5_10_push.json

@@ -202,8 +202,10 @@ public void pushDownLimit(Integer limit, Integer offset) {
     // 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);
-    startFrom = offset;
-    sourceBuilder.from(offset).size(requestedTotalSize);
+    // 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;
+    sourceBuilder.from(startFrom).size(requestedTotalSize);
   }
 
   public void pushDownTrackedScore(boolean trackScores) {