Fallback to sub-aggregation if composite aggregation doesn't support

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

@@ -941,4 +941,54 @@ public void testStatsWithBinsOnTimeField_Avg() throws IOException {
         rows(41.8, "2024-07-01 00:04:00"),
         rows(50.0, "2024-07-01 00:05:00"));
   }
+
+  @Test
+  public void testStatsWithBinsOnTimeAndTermField_Count() throws IOException {
+    // TODO: Remove this after addressing https://github.com/opensearch-project/sql/issues/4317
+    enabledOnlyWhenPushdownIsEnabled();
+
+    JSONObject result =
+        executeQuery(
+            "source=events_null | bin @timestamp bins=3 | stats bucket_nullable=false count() by"
+                + " region, @timestamp");
+    // TODO: @timestamp should keep date as its type, to be addressed by this issue:
+    // https://github.com/opensearch-project/sql/issues/4317
+    verifySchema(
+        result,
+        schema("count()", null, "bigint"),
+        schema("region", null, "string"),
+        schema("@timestamp", null, "string"));
+    // auto_date_histogram will choose span=5m for bins=3
+    verifyDataRows(
+        result,
+        rows(1, "eu-west", "2024-07-01 00:03:00"),
+        rows(2, "us-east", "2024-07-01 00:00:00"),
+        rows(1, "us-east", "2024-07-01 00:05:00"),
+        rows(2, "us-west", "2024-07-01 00:01:00"));
+  }
+
+  @Test
+  public void testStatsWithBinsOnTimeAndTermField_Avg() throws IOException {
+    // TODO: Remove this after addressing https://github.com/opensearch-project/sql/issues/4317
+    enabledOnlyWhenPushdownIsEnabled();
+
+    JSONObject result =
+        executeQuery(
+            "source=events_null | bin @timestamp bins=3 | stats bucket_nullable=false "
+                + " avg(cpu_usage) by region, @timestamp");
+    // TODO: @timestamp should keep date as its type, to be addressed by this issue:
+    // https://github.com/opensearch-project/sql/issues/4317
+    verifySchema(
+        result,
+        schema("avg(cpu_usage)", null, "double"),
+        schema("region", null, "string"),
+        schema("@timestamp", null, "string"));
+    // auto_date_histogram will choose span=5m for bins=3
+    verifyDataRows(
+        result,
+        rows(42.1, "eu-west", "2024-07-01 00:03:00"),
+        rows(50.25, "us-east", "2024-07-01 00:00:00"),
+        rows(50, "us-east", "2024-07-01 00:05:00"),
+        rows(40.25, "us-west", "2024-07-01 00:01:00"));
+  }
 }

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

@@ -358,6 +358,40 @@ public void testExplainStatsWithBinsOnTimeField() throws IOException {
             "source=events | bin @timestamp bins=3 | stats avg(cpu_usage) by @timestamp"));
   }
 
+  @Test
+  public void testExplainStatsWithSubAggregation() throws IOException {
+    enabledOnlyWhenPushdownIsEnabled();
+    String expected = loadExpectedPlan("explain_stats_bins_on_time_and_term.yaml");
+    assertYamlEqualsJsonIgnoreId(
+        expected,
+        explainQueryToString(
+            "source=events | bin @timestamp bins=3 | stats bucket_nullable=false count() by"
+                + " @timestamp, region"));
+
+    expected = loadExpectedPlan("explain_stats_bins_on_time_and_term2.yaml");
+    assertYamlEqualsJsonIgnoreId(
+        expected,
+        explainQueryToString(
+            "source=events | bin @timestamp bins=3 | stats bucket_nullable=false avg(cpu_usage) by"
+                + " @timestamp, region"));
+  }
+
+  @Test
+  public void bucketNullableNotSupportSubAggregation() throws IOException {
+    // TODO: Don't throw exception after addressing
+    // https://github.com/opensearch-project/sql/issues/4317
+    // When bucketNullable is true, sub aggregation is not supported. Hence we cannot pushdown the
+    // aggregation in this query. Caused by issue
+    // https://github.com/opensearch-project/sql/issues/4317,
+    // bin aggregation on timestamp field won't work if not been push down.
+    enabledOnlyWhenPushdownIsEnabled();
+    assertThrows(
+        Exception.class,
+        () ->
+            explainQueryToString(
+                "source=events | bin @timestamp bins=3 | stats count() by @timestamp, region"));
+  }
+
   @Test
   public void testExplainBinWithSpan() throws IOException {
     String expected = loadExpectedPlan("explain_bin_span.yaml");
@@ -673,15 +707,15 @@ public void testPushdownLimitIntoAggregation() throws IOException {
             "source=opensearch-sql_test_index_account | stats count() by state | sort state | head"
                 + " 100 | head 10 from 10 "));
 
-    expected = loadExpectedPlan("explain_limit_agg_pushdown_bucket_nullable1.json");
-    assertJsonEqualsIgnoreId(
+    expected = loadExpectedPlan("explain_limit_agg_pushdown_bucket_nullable1.yaml");
+    assertYamlEqualsJsonIgnoreId(
         expected,
         explainQueryToString(
             "source=opensearch-sql_test_index_account | stats bucket_nullable=false count() by"
                 + " state | head 100 | head 10 from 10 "));
 
-    expected = loadExpectedPlan("explain_limit_agg_pushdown_bucket_nullable2.json");
-    assertJsonEqualsIgnoreId(
+    expected = loadExpectedPlan("explain_limit_agg_pushdown_bucket_nullable2.yaml");
+    assertYamlEqualsJsonIgnoreId(
         expected,
         explainQueryToString(
             "source=opensearch-sql_test_index_account | stats bucket_nullable=false count() by"
@@ -853,15 +887,15 @@ public void testExplainCountsByAgg() throws IOException {
   public void testExplainSortOnMetricsNoBucketNullable() throws IOException {
     // TODO enhancement later: https://github.com/opensearch-project/sql/issues/4282
     enabledOnlyWhenPushdownIsEnabled();
-    String expected = loadExpectedPlan("explain_agg_sort_on_metrics1.json");
-    assertJsonEqualsIgnoreId(
+    String expected = loadExpectedPlan("explain_agg_sort_on_metrics1.yaml");
+    assertYamlEqualsJsonIgnoreId(
         expected,
         explainQueryToString(
             "source=opensearch-sql_test_index_account | stats bucket_nullable=false count() by"
                 + " state | sort `count()`"));
 
-    expected = loadExpectedPlan("explain_agg_sort_on_metrics2.json");
-    assertJsonEqualsIgnoreId(
+    expected = loadExpectedPlan("explain_agg_sort_on_metrics2.yaml");
+    assertYamlEqualsJsonIgnoreId(
         expected,
         explainQueryToString(
             "source=opensearch-sql_test_index_account | stats bucket_nullable=false count() by"

integ-test/src/test/resources/expectedOutput/calcite/explain_agg_sort_on_metrics1.yaml

@@ -0,0 +1,13 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(sort0=[$0], dir0=[ASC-nulls-first], fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalSort(sort0=[$0], dir0=[ASC-nulls-first])
+        LogicalProject(count()=[$1], state=[$0])
+          LogicalAggregate(group=[{0}], count()=[COUNT()])
+            LogicalProject(state=[$7])
+              LogicalFilter(condition=[IS NOT NULL($7)])
+                CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    EnumerableLimit(fetch=[10000])
+      EnumerableSort(sort0=[$0], dir0=[ASC-nulls-first])
+        CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT()), PROJECT->[count(), state]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"state":{"terms":{"field":"state.keyword","missing_bucket":false,"order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_agg_sort_on_metrics2.yaml

@@ -0,0 +1,13 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(sort0=[$0], dir0=[ASC-nulls-first], fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalSort(sort0=[$0], dir0=[ASC-nulls-first])
+        LogicalProject(count()=[$2], gender=[$0], state=[$1])
+          LogicalAggregate(group=[{0, 1}], count()=[COUNT()])
+            LogicalProject(gender=[$4], state=[$7])
+              LogicalFilter(condition=[AND(IS NOT NULL($4), IS NOT NULL($7))])
+                CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    EnumerableLimit(fetch=[10000])
+      EnumerableSort(sort0=[$0], dir0=[ASC-nulls-first])
+        CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0, 1},count()=COUNT()), PROJECT->[count(), gender, state]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":false,"order":"asc"}}},{"state":{"terms":{"field":"state.keyword","missing_bucket":false,"order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_limit_agg_pushdown_bucket_nullable1.yaml

@@ -0,0 +1,15 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalSort(offset=[10], fetch=[10])
+        LogicalSort(fetch=[100])
+          LogicalProject(count()=[$1], state=[$0])
+            LogicalAggregate(group=[{0}], count()=[COUNT()])
+              LogicalProject(state=[$7])
+                LogicalFilter(condition=[IS NOT NULL($7)])
+                  CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    EnumerableLimit(fetch=[10000])
+      EnumerableCalc(expr#0..1=[{inputs}], count()=[$t1], state=[$t0])
+        EnumerableLimit(offset=[10], fetch=[10])
+          CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT()), LIMIT->100, LIMIT->[10 from 10]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":20,"sources":[{"state":{"terms":{"field":"state.keyword","missing_bucket":false,"order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_limit_agg_pushdown_bucket_nullable2.yaml

@@ -0,0 +1,15 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalSort(offset=[10], fetch=[10])
+        LogicalSort(sort0=[$1], dir0=[ASC-nulls-first], fetch=[100])
+          LogicalProject(count()=[$1], state=[$0])
+            LogicalAggregate(group=[{0}], count()=[COUNT()])
+              LogicalProject(state=[$7])
+                LogicalFilter(condition=[IS NOT NULL($7)])
+                  CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    EnumerableLimit(fetch=[10000])
+      EnumerableCalc(expr#0..1=[{inputs}], count()=[$t1], state=[$t0])
+        EnumerableLimit(offset=[10], fetch=[10])
+          CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT()), SORT->[0 ASC FIRST], LIMIT->100, LIMIT->[10 from 10]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":20,"sources":[{"state":{"terms":{"field":"state.keyword","missing_bucket":false,"order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_stats_bins_on_time.yaml

@@ -8,4 +8,4 @@ calcite:
   physical: |
     EnumerableLimit(fetch=[10000])
       EnumerableCalc(expr#0..1=[{inputs}], expr#2=[0], expr#3=[>($t1, $t2)], count()=[$t1], @timestamp=[$t0], $condition=[$t3])
-        CalciteEnumerableIndexScan(table=[[OpenSearch, events]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT())], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"@timestamp":{"auto_date_histogram":{"field":"@timestamp","buckets":3,"minimum_interval":null},"aggregations":{"count()":{"value_count":{"field":"_index"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
+        CalciteEnumerableIndexScan(table=[[OpenSearch, events]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT())], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"@timestamp":{"auto_date_histogram":{"field":"@timestamp","buckets":3,"minimum_interval":null}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_stats_bins_on_time_and_term.yaml

@@ -0,0 +1,11 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(count()=[$2], @timestamp=[$0], region=[$1])
+        LogicalAggregate(group=[{0, 1}], count()=[COUNT()])
+          LogicalProject(@timestamp=[$15], region=[$7])
+            LogicalFilter(condition=[IS NOT NULL($7)])
+              LogicalProject(environment=[$0], status_code=[$2], service=[$3], host=[$4], memory_usage=[$5], response_time=[$6], cpu_usage=[$7], region=[$8], bytes_sent=[$9], _id=[$10], _index=[$11], _score=[$12], _maxscore=[$13], _sort=[$14], _routing=[$15], @timestamp=[WIDTH_BUCKET($1, 3, -(MAX($1) OVER (), MIN($1) OVER ()), MAX($1) OVER ())])
+                CalciteLogicalIndexScan(table=[[OpenSearch, events]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, events]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0, 1},count()=COUNT()), PROJECT->[count(), @timestamp, region], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"region":{"terms":{"field":"region","size":1000,"min_doc_count":1,"shard_min_doc_count":0,"show_term_doc_count_error":false,"order":{"_key":"asc"}},"aggregations":{"@timestamp":{"auto_date_histogram":{"field":"@timestamp","buckets":3,"minimum_interval":null}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_stats_bins_on_time_and_term2.yaml

@@ -0,0 +1,11 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(avg(cpu_usage)=[$2], @timestamp=[$0], region=[$1])
+        LogicalAggregate(group=[{0, 1}], avg(cpu_usage)=[AVG($2)])
+          LogicalProject(@timestamp=[$15], region=[$7], cpu_usage=[$6])
+            LogicalFilter(condition=[IS NOT NULL($7)])
+              LogicalProject(environment=[$0], status_code=[$2], service=[$3], host=[$4], memory_usage=[$5], response_time=[$6], cpu_usage=[$7], region=[$8], bytes_sent=[$9], _id=[$10], _index=[$11], _score=[$12], _maxscore=[$13], _sort=[$14], _routing=[$15], @timestamp=[WIDTH_BUCKET($1, 3, -(MAX($1) OVER (), MIN($1) OVER ()), MAX($1) OVER ())])
+                CalciteLogicalIndexScan(table=[[OpenSearch, events]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, events]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={1, 2},avg(cpu_usage)=AVG($0)), PROJECT->[avg(cpu_usage), @timestamp, region], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"region":{"terms":{"field":"region","size":1000,"min_doc_count":1,"shard_min_doc_count":0,"show_term_doc_count_error":false,"order":{"_key":"asc"}},"aggregations":{"@timestamp":{"auto_date_histogram":{"field":"@timestamp","buckets":3,"minimum_interval":null},"aggregations":{"avg(cpu_usage)":{"avg":{"field":"cpu_usage"}}}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])