Support composite aggregation paginating

core/src/main/java/org/opensearch/sql/calcite/plan/LogicalSystemLimit.java

@@ -39,17 +39,6 @@ public enum SystemLimitType {
 
   @Getter private final SystemLimitType type;
 
-  private LogicalSystemLimit(
-      SystemLimitType type,
-      RelOptCluster cluster,
-      RelTraitSet traitSet,
-      RelNode input,
-      RelCollation collation,
-      @Nullable RexNode offset,
-      @Nullable RexNode fetch) {
-    this(type, cluster, traitSet, Collections.emptyList(), input, collation, offset, fetch);
-  }
-
   private LogicalSystemLimit(
       SystemLimitType type,
       RelOptCluster cluster,
@@ -76,7 +65,8 @@ public static LogicalSystemLimit create(
     RelCollation collation = collations == null ? null : collations.get(0);
     collation = RelCollationTraitDef.INSTANCE.canonize(collation);
     RelTraitSet traitSet = input.getTraitSet().replace(Convention.NONE).replace(collation);
-    return new LogicalSystemLimit(type, cluster, traitSet, input, collation, offset, fetch);
+    return new LogicalSystemLimit(
+        type, cluster, traitSet, Collections.emptyList(), input, collation, offset, fetch);
   }
 
   @Override

core/src/main/java/org/opensearch/sql/calcite/utils/PPLHintStrategyTable.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.utils;
+
+import com.google.common.base.Suppliers;
+import java.util.function.Supplier;
+import lombok.experimental.UtilityClass;
+import org.apache.calcite.rel.hint.HintStrategyTable;
+import org.apache.calcite.rel.logical.LogicalAggregate;
+
+@UtilityClass
+public class PPLHintStrategyTable {
+
+  private static final Supplier<HintStrategyTable> HINT_STRATEGY_TABLE =
+      Suppliers.memoize(
+          () ->
+              HintStrategyTable.builder()
+                  .hintStrategy(
+                      "stats_args",
+                      (hint, rel) -> {
+                        return rel instanceof LogicalAggregate;
+                      })
+                  // add more here
+                  .build());
+
+  /** Update the HINT_STRATEGY_TABLE when you create a new hint. */
+  public static HintStrategyTable getHintStrategyTable() {
+    return HINT_STRATEGY_TABLE.get();
+  }
+}

core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java

@@ -29,7 +29,6 @@
 import org.apache.calcite.rel.core.Project;
 import org.apache.calcite.rel.core.Sort;
 import org.apache.calcite.rel.core.TableScan;
-import org.apache.calcite.rel.hint.HintStrategyTable;
 import org.apache.calcite.rel.hint.RelHint;
 import org.apache.calcite.rel.logical.LogicalAggregate;
 import org.apache.calcite.rel.logical.LogicalFilter;
@@ -605,16 +604,7 @@ static void addIgnoreNullBucketHintToAggregate(RelBuilder relBuilder) {
     assert relBuilder.peek() instanceof LogicalAggregate
         : "Stats hits should be added to LogicalAggregate";
     relBuilder.hints(statHits);
-    relBuilder
-        .getCluster()
-        .setHintStrategies(
-            HintStrategyTable.builder()
-                .hintStrategy(
-                    "stats_args",
-                    (hint, rel) -> {
-                      return rel instanceof LogicalAggregate;
-                    })
-                .build());
+    relBuilder.getCluster().setHintStrategies(PPLHintStrategyTable.getHintStrategyTable());
   }
 
   /** Extract the RexLiteral from the aggregate call if the aggregate call is a LITERAL_AGG. */

core/src/main/java/org/opensearch/sql/exception/NonFallbackCalciteException.java

@@ -11,4 +11,8 @@ public class NonFallbackCalciteException extends QueryEngineException {
   public NonFallbackCalciteException(String message) {
     super(message);
   }
+
+  public NonFallbackCalciteException(String message, Throwable cause) {
+    super(message, cause);
+  }
 }

docs/user/optimization/optimization.rst

@@ -44,7 +44,7 @@ The consecutive Filter operator will be merged as one Filter operator::
           {
             "name": "OpenSearchIndexScan",
             "description": {
-              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"bool\":{\"filter\":[{\"range\":{\"age\":{\"from\":null,\"to\":20,\"include_lower\":true,\"include_upper\":false,\"boost\":1.0}}},{\"range\":{\"age\":{\"from\":10,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}}],\"adjust_pure_negative\":true,\"boost\":1.0}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, searchDone=false)"
+              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"bool\":{\"filter\":[{\"range\":{\"age\":{\"from\":null,\"to\":20,\"include_lower\":true,\"include_upper\":false,\"boost\":1.0}}},{\"range\":{\"age\":{\"from\":10,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}}],\"adjust_pure_negative\":true,\"boost\":1.0}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}})"
             },
             "children": []
           }
@@ -71,7 +71,7 @@ The Filter operator should be push down under Sort operator::
           {
             "name": "OpenSearchIndexScan",
             "description": {
-              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":null,\"to\":20,\"include_lower\":true,\"include_upper\":false,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_first\"}}]}, searchDone=false)"
+              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":null,\"to\":20,\"include_lower\":true,\"include_upper\":false,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_first\"}}]})"
             },
             "children": []
           }
@@ -102,7 +102,7 @@ The Project list will push down to Query DSL to `filter the source <https://www.
           {
             "name": "OpenSearchIndexScan",
             "description": {
-              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, searchDone=false)"
+              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}})"
             },
             "children": []
           }
@@ -128,7 +128,7 @@ The Filter operator will merge into OpenSearch Query DSL::
           {
             "name": "OpenSearchIndexScan",
             "description": {
-              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, searchDone=false)"
+              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}})"
             },
             "children": []
           }
@@ -154,7 +154,7 @@ The Sort operator will merge into OpenSearch Query DSL::
           {
             "name": "OpenSearchIndexScan",
             "description": {
-              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_first\"}}]}, searchDone=false)"
+              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_first\"}}]})"
             },
             "children": []
           }
@@ -188,7 +188,7 @@ Because the OpenSearch Script Based Sorting can't handle NULL/MISSING value, the
               {
                 "name": "OpenSearchIndexScan",
                 "description": {
-                  "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\"}, searchDone=false)"
+                  "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\"})"
                 },
                 "children": []
               }
@@ -216,7 +216,7 @@ The Limit operator will merge in OpenSearch Query DSL::
               {
                 "name": "OpenSearchIndexScan",
                 "description": {
-                  "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":5,\"size\":10,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, searchDone=false)"
+                  "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":5,\"size\":10,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}})"
                 },
                 "children": []
               }
@@ -252,7 +252,7 @@ If sort that includes expression, which cannot be merged into query DSL, also ex
                   {
                     "name": "OpenSearchIndexScan",
                     "description": {
-                      "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\"}, searchDone=false)"
+                      "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\"})"
                     },
                     "children": []
                   }
@@ -280,7 +280,7 @@ The Aggregation operator will merge into OpenSearch Aggregation::
           {
             "name": "OpenSearchIndexScan",
             "description": {
-              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone=false)"
+              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}})"
             },
             "children": []
           }
@@ -306,7 +306,7 @@ The Sort operator will merge into OpenSearch Aggregation.::
           {
             "name": "OpenSearchIndexScan",
             "description": {
-              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"last\",\"order\":\"desc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone=false)"
+              "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"last\",\"order\":\"desc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}})"
             },
             "children": []
           }
@@ -341,7 +341,7 @@ Because the OpenSearch Composite Aggregation doesn't support order by metrics fi
               {
                 "name": "OpenSearchIndexScan",
                 "description": {
-                  "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone=false)"
+                  "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}})"
                 },
                 "children": []
               }

docs/user/ppl/cmd/explain.rst

@@ -48,7 +48,7 @@ Explain::
           {
             "name": "OpenSearchIndexScan",
             "description": {
-              "request": """OpenSearchQueryRequest(indexName=state_country, sourceBuilder={"from":0,"size":10000,"timeout":"1m","query":{"bool":{"should":[{"term":{"country":{"value":"USA","boost":1.0}}},{"term":{"country":{"value":"England","boost":1.0}}}],"adjust_pure_negative":true,"boost":1.0}},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"country":{"terms":{"field":"country","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"count()":{"value_count":{"field":"_index"}}}}}}, needClean=true, searchDone=false, pitId=null, cursorKeepAlive=null, searchAfter=null, searchResponse=null)"""
+              "request": """OpenSearchQueryRequest(indexName=state_country, sourceBuilder={"from":0,"size":10000,"timeout":"1m","query":{"bool":{"should":[{"term":{"country":{"value":"USA","boost":1.0}}},{"term":{"country":{"value":"England","boost":1.0}}}],"adjust_pure_negative":true,"boost":1.0}},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"country":{"terms":{"field":"country","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"count()":{"value_count":{"field":"_index"}}}}}}, pitId=null, cursorKeepAlive=null, searchAfter=null, searchResponse=null)"""
             },
             "children": []
           }

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java

@@ -1024,10 +1024,10 @@ public void testRexExplain() throws IOException {
 
   @Test
   public void testExplainAppendCommand() throws IOException {
-    String expected = loadExpectedPlan("explain_append_command.json");
-    assertJsonEqualsIgnoreId(
+    String expected = loadExpectedPlan("explain_append_command.yaml");
+    assertYamlEqualsIgnoreId(
         expected,
-        explainQueryToString(
+        explainQueryYaml(
             String.format(
                 Locale.ROOT,
                 "source=%s | stats count(balance) as cnt by gender | append [ source=%s | stats"
@@ -1273,6 +1273,119 @@ public void testExplainCountsByAgg() throws IOException {
                 TEST_INDEX_ACCOUNT)));
   }
 
+  @Test
+  public void testPaginatingAggForHaving() throws IOException {
+    enabledOnlyWhenPushdownIsEnabled();
+    try {
+      setQueryBucketSize(2);
+      String expected = loadExpectedPlan("explain_agg_paginating_having1.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              "source=opensearch-sql_test_index_account | stats count() as c by"
+                  + " state | where c > 10"));
+      expected = loadExpectedPlan("explain_agg_paginating_having2.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              "source=opensearch-sql_test_index_account | stats bucket_nullable = false count() by"
+                  + " state | where `count()` > 10"));
+      expected = loadExpectedPlan("explain_agg_paginating_having3.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              "source=opensearch-sql_test_index_account | stats avg(balance) as avg, count() as cnt"
+                  + " by state | eval new_avg = avg + 1000, new_cnt = cnt + 1 | where new_avg >"
+                  + " 1000 or new_cnt > 1"));
+    } finally {
+      resetQueryBucketSize();
+    }
+  }
+
+  @Test
+  public void testPaginatingAggForJoin() throws IOException {
+    enabledOnlyWhenPushdownIsEnabled();
+    try {
+      setQueryBucketSize(2);
+      String expected = loadExpectedPlan("explain_agg_paginating_join1.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              "source=opensearch-sql_test_index_account | stats count() as c by state | join left=l"
+                  + " right=r on l.state=r.state [ source=opensearch-sql_test_index_bank | stats"
+                  + " count() as c by state ]"));
+      expected = loadExpectedPlan("explain_agg_paginating_join2.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              "source=opensearch-sql_test_index_account | stats bucket_nullable = false count() as"
+                  + " c by state | join left=l right=r on l.state=r.state ["
+                  + " source=opensearch-sql_test_index_bank | stats bucket_nullable = false"
+                  + " count() as c by state ]"));
+      expected = loadExpectedPlan("explain_agg_paginating_join3.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              "source=opensearch-sql_test_index_account | stats count() as c by state | join"
+                  + " type=inner state [ source=opensearch-sql_test_index_bank | stats count()"
+                  + " as c by state ]"));
+      expected = loadExpectedPlan("explain_agg_paginating_join4.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              "source=opensearch-sql_test_index_account | stats count() as c by state | head 10"
+                  + " | join type=inner state [ source=opensearch-sql_test_index_account"
+                  + " | stats count() as c by state ]"));
+    } finally {
+      resetQueryBucketSize();
+    }
+  }
+
+  @Test
+  public void testPaginatingAggForHeadFrom() throws IOException {
+    enabledOnlyWhenPushdownIsEnabled();
+    try {
+      setQueryBucketSize(2);
+      String expected = loadExpectedPlan("explain_agg_paginating_head_from.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              "source=opensearch-sql_test_index_account | stats count() as c by state | head 10"
+                  + " from 2"));
+    } finally {
+      resetQueryBucketSize();
+    }
+  }
+
+  @Test
+  public void testPaginatingHeadSizeNoLessThanQueryBucketSize() throws IOException {
+    enabledOnlyWhenPushdownIsEnabled();
+    try {
+      setQueryBucketSize(2);
+      String expected =
+          loadExpectedPlan("explain_agg_paginating_head_size_query_bucket_size1.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              String.format(
+                  "source=%s | stats count() by age | sort -age | head 3", TEST_INDEX_BANK)));
+      expected = loadExpectedPlan("explain_agg_paginating_head_size_query_bucket_size2.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              String.format(
+                  "source=%s | stats count() by age | sort -age | head 2", TEST_INDEX_BANK)));
+      expected = loadExpectedPlan("explain_agg_paginating_head_size_query_bucket_size3.yaml");
+      assertYamlEqualsIgnoreId(
+          expected,
+          explainQueryYaml(
+              String.format(
+                  "source=%s | stats count() by age | sort -age | head 1", TEST_INDEX_BANK)));
+    } finally {
+      resetQueryBucketSize();
+    }
+  }
+
   @Test
   public void testExplainSortOnMeasure() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();