Keep categories with max summed values when top k is set

Signed-off-by: Yuanchun Shen <[email protected]>

Author: yuancu

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -1131,15 +1131,16 @@ private void visitAggregation(
     UnresolvedExpression span = node.getSpan();
     if (Objects.nonNull(span)) {
       groupExprList.add(span);
-      if (getTimeSpanField(span).isPresent()){
-          nonNullGroupMask.set(0);
+      if (getTimeSpanField(span).isPresent()) {
+        nonNullGroupMask.set(0);
       }
     }
     groupExprList.addAll(node.getGroupExprList());
 
     // add stats hint to LogicalAggregation
     boolean toAddHintsOnAggregate =
-        nonNullGroupMask.nextClearBit(0) >= groupExprList.size() // This checks if all group-bys are nonnull
+        nonNullGroupMask.nextClearBit(0)
+                >= groupExprList.size() // This checks if all group-bys should be nonnull
             && !groupExprList.isEmpty()
             && !(groupExprList.size() == 1 && getTimeSpanField(span).isPresent());
     // add isNotNull filter before aggregation for non-nullable buckets
@@ -2413,10 +2414,6 @@ public RelNode visitChart(Chart node, CalcitePlanContext context) {
     }
     visitAggregation(aggregation, context, false, nonNullGroupMask);
     RelBuilder relBuilder = context.relBuilder;
-    String columnSplitName =
-        relBuilder.peek().getRowType().getFieldNames().size() > 2
-            ? relBuilder.peek().getRowType().getFieldNames().get(1)
-            : null;
 
     // If row or column split does not present or limit equals 0, this is the same as `stats agg
     // [group by col]` because all truncating is performed on the column split
@@ -2457,16 +2454,13 @@ public RelNode visitChart(Chart node, CalcitePlanContext context) {
     final String GRAND_TOTAL_COL = "__grand_total__";
     relBuilder.aggregate(
         relBuilder.groupKey(relBuilder.field(0)),
-        buildAggCall(context.relBuilder, aggFunction, relBuilder.field(1))
-            .as(GRAND_TOTAL_COL)); // results: group key, agg calls
+        // Top-K semantic: Retain categories whose summed values are among the greatest
+        relBuilder.sum(relBuilder.field(1)).as(GRAND_TOTAL_COL)); // results: group key, agg calls
     RexNode grandTotal = relBuilder.field(GRAND_TOTAL_COL);
-    // Apply sorting: for MIN/EARLIEST, reverse the top/bottom logic
-    boolean smallestFirst =
-        aggFunction == BuiltinFunctionName.MIN || aggFunction == BuiltinFunctionName.EARLIEST;
-    if (config.top != smallestFirst) {
+    // Apply sorting: keep the max values if top is set
+    if (config.top) {
       grandTotal = relBuilder.desc(grandTotal);
     }
-
     // Always set it to null last so that it does not interfere with top / bottom calculation
     grandTotal = relBuilder.nullsLast(grandTotal);
     RexNode rowNum =
@@ -2518,6 +2512,7 @@ public RelNode visitChart(Chart node, CalcitePlanContext context) {
               relBuilder.literal(config.otherStr));
     }
 
+    String columnSplitName = ((Alias) node.getColumnSplit()).getName();
     String aggFieldName = relBuilder.peek().getRowType().getFieldNames().get(2);
     relBuilder.project(
         relBuilder.field(0),

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

@@ -13,6 +13,7 @@
 import static org.opensearch.sql.util.MatcherUtils.rows;
 import static org.opensearch.sql.util.MatcherUtils.schema;
 import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
+import static org.opensearch.sql.util.MatcherUtils.verifyDataRowsInOrder;
 import static org.opensearch.sql.util.MatcherUtils.verifySchema;
 
 import java.io.IOException;
@@ -221,10 +222,11 @@ public void testChartLimitTopWithUseOther() throws IOException {
         schema("max(severityNumber)", "bigint"));
     verifyDataRows(
         result,
-        rows(1, "max_among_other", 17),
-        rows(0, "max_among_other", 22),
-        rows(0, "FATAL3", 23),
-        rows(0, "FATAL4", 24));
+        rows(0, "ERROR", 17),
+        rows(0, "FATAL4", 24),
+        rows(0, "max_among_other", 23),
+        rows(1, "ERROR", 17),
+        rows(1, "max_among_other", 9));
   }
 
   @Test
@@ -255,12 +257,13 @@ public void testChartLimitTopWithMinAgg() throws IOException {
         schema("flags", "bigint"),
         schema("severityText", "string"),
         schema("min(severityNumber)", "bigint"));
-    verifyDataRows(
+    verifyDataRowsInOrder(
         result,
-        rows(1, "OTHER", 9),
-        rows(1, "TRACE", 1),
-        rows(0, "OTHER", 3),
-        rows(0, "TRACE2", 2));
+        rows(0, "ERROR", 17),
+        rows(0, "FATAL4", 24),
+        rows(0, "OTHER", 2),
+        rows(1, "ERROR", 17),
+        rows(1, "OTHER", 1));
   }
 
   @Test
@@ -316,19 +319,17 @@ public void testChartUseNullFalseWithNullStr() throws IOException {
   }
 
   @Test
-    public void testChartNullsInRowSplitShouldBeIgnored() throws IOException {
-      JSONObject result =
-          executeQuery(
-                  "source=events_null | chart min(cpu_usage) by host region");
-      verifySchema(
-              result,
-              schema("host", "string"),
-              schema("region", "string"),
-              schema("min(cpu_usage)", "double"));
-      verifyDataRows(
-              result,
-              rows("db-01", "eu-west",  42.1),
-              rows("web-01", "us-east",  45.2),
-              rows("web-02", "us-west",  38.7));
+  public void testChartNullsInRowSplitShouldBeIgnored() throws IOException {
+    JSONObject result = executeQuery("source=events_null | chart min(cpu_usage) by host region");
+    verifySchema(
+        result,
+        schema("host", "string"),
+        schema("region", "string"),
+        schema("min(cpu_usage)", "double"));
+    verifyDataRows(
+        result,
+        rows("db-01", "eu-west", 42.1),
+        rows("web-01", "us-east", 45.2),
+        rows("web-02", "us-west", 38.7));
   }
 }

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

@@ -11,7 +11,7 @@ calcite:
                     LogicalFilter(condition=[IS NOT NULL($4)])
                       CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
               LogicalProject(age=[$0], __grand_total__=[$1], _row_number_chart_=[ROW_NUMBER() OVER (ORDER BY $1 DESC NULLS LAST)])
-                LogicalAggregate(group=[{0}], __grand_total__=[AVG($1)])
+                LogicalAggregate(group=[{0}], __grand_total__=[SUM($1)])
                   LogicalFilter(condition=[IS NOT NULL($0)])
                     LogicalProject(age=[SAFE_CAST($1)], avg(balance)=[$2])
                       LogicalAggregate(group=[{0, 1}], avg(balance)=[AVG($2)])
@@ -31,7 +31,6 @@ calcite:
                 EnumerableSort(sort0=[$0], dir0=[ASC])
                   EnumerableCalc(expr#0..2=[{inputs}], age=[$t0], $1=[$t2])
                     EnumerableWindow(window#0=[window(order by [1 DESC-nulls-last] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
-                      EnumerableCalc(expr#0..2=[{inputs}], expr#3=[0], expr#4=[=($t2, $t3)], expr#5=[null:DOUBLE], expr#6=[CASE($t4, $t5, $t1)], expr#7=[/($t6, $t2)], age=[$t0], __grand_total__=[$t7])
-                        EnumerableAggregate(group=[{0}], agg#0=[$SUM0($1)], agg#1=[COUNT($1)])
-                          EnumerableCalc(expr#0..1=[{inputs}], expr#2=[SAFE_CAST($t0)], expr#3=[IS NOT NULL($t2)], age=[$t2], avg(balance)=[$t1], $condition=[$t3])
-                            CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[gender, balance, age], FILTER->IS NOT NULL($0), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0, 1},avg(balance)=AVG($2)), PROJECT->[age, avg(balance)]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"exists":{"field":"gender","boost":1.0}},"_source":{"includes":["gender","balance","age"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}},{"age":{"terms":{"field":"age","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"avg(balance)":{"avg":{"field":"balance"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
+                      EnumerableAggregate(group=[{0}], __grand_total__=[SUM($1)])
+                        EnumerableCalc(expr#0..1=[{inputs}], expr#2=[SAFE_CAST($t0)], expr#3=[IS NOT NULL($t2)], age=[$t2], avg(balance)=[$t1], $condition=[$t3])
+                          CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[gender, balance, age], FILTER->IS NOT NULL($0), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0, 1},avg(balance)=AVG($2)), PROJECT->[age, avg(balance)]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"exists":{"field":"gender","boost":1.0}},"_source":{"includes":["gender","balance","age"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}},{"age":{"terms":{"field":"age","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"avg(balance)":{"avg":{"field":"balance"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

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

@@ -11,7 +11,7 @@ calcite:
                     LogicalFilter(condition=[IS NOT NULL($4)])
                       CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank_with_null_values]])
               LogicalProject(age=[$0], __grand_total__=[$1], _row_number_chart_=[ROW_NUMBER() OVER (ORDER BY $1 DESC NULLS LAST)])
-                LogicalAggregate(group=[{0}], __grand_total__=[AVG($1)])
+                LogicalAggregate(group=[{0}], __grand_total__=[SUM($1)])
                   LogicalFilter(condition=[IS NOT NULL($0)])
                     LogicalProject(age=[SAFE_CAST($1)], avg(balance)=[$2])
                       LogicalAggregate(group=[{0, 2}], avg(balance)=[AVG($1)])
@@ -33,9 +33,8 @@ calcite:
                 EnumerableSort(sort0=[$0], dir0=[ASC])
                   EnumerableCalc(expr#0..2=[{inputs}], age=[$t0], $1=[$t2])
                     EnumerableWindow(window#0=[window(order by [1 DESC-nulls-last] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
-                      EnumerableCalc(expr#0..2=[{inputs}], expr#3=[0], expr#4=[=($t2, $t3)], expr#5=[null:DOUBLE], expr#6=[CASE($t4, $t5, $t1)], expr#7=[/($t6, $t2)], age=[$t0], __grand_total__=[$t7])
-                        EnumerableAggregate(group=[{0}], agg#0=[$SUM0($1)], agg#1=[COUNT($1)])
-                          EnumerableCalc(expr#0..3=[{inputs}], expr#4=[SAFE_CAST($t1)], expr#5=[0], expr#6=[=($t3, $t5)], expr#7=[null:BIGINT], expr#8=[CASE($t6, $t7, $t2)], expr#9=[CAST($t8):DOUBLE], expr#10=[/($t9, $t3)], expr#11=[IS NOT NULL($t4)], age=[$t4], avg(balance)=[$t10], $condition=[$t11])
-                            EnumerableAggregate(group=[{0, 2}], agg#0=[$SUM0($1)], agg#1=[COUNT($1)])
-                              EnumerableCalc(expr#0..2=[{inputs}], expr#3=[10], expr#4=[null:NULL], expr#5=[SPAN($t2, $t3, $t4)], gender=[$t1], balance=[$t0], age0=[$t5])
-                                CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank_with_null_values]], PushDownContext=[[PROJECT->[balance, gender, age], FILTER->IS NOT NULL($1)], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","query":{"exists":{"field":"gender","boost":1.0}},"_source":{"includes":["balance","gender","age"],"excludes":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
+                      EnumerableAggregate(group=[{0}], __grand_total__=[SUM($1)])
+                        EnumerableCalc(expr#0..3=[{inputs}], expr#4=[SAFE_CAST($t1)], expr#5=[0], expr#6=[=($t3, $t5)], expr#7=[null:BIGINT], expr#8=[CASE($t6, $t7, $t2)], expr#9=[CAST($t8):DOUBLE], expr#10=[/($t9, $t3)], expr#11=[IS NOT NULL($t4)], age=[$t4], avg(balance)=[$t10], $condition=[$t11])
+                          EnumerableAggregate(group=[{0, 2}], agg#0=[$SUM0($1)], agg#1=[COUNT($1)])
+                            EnumerableCalc(expr#0..2=[{inputs}], expr#3=[10], expr#4=[null:NULL], expr#5=[SPAN($t2, $t3, $t4)], gender=[$t1], balance=[$t0], age0=[$t5])
+                              CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank_with_null_values]], PushDownContext=[[PROJECT->[balance, gender, age], FILTER->IS NOT NULL($1)], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","query":{"exists":{"field":"gender","boost":1.0}},"_source":{"includes":["balance","gender","age"],"excludes":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

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

@@ -11,7 +11,7 @@ calcite:
                     LogicalFilter(condition=[IS NOT NULL($3)])
                       CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_time_data]])
               LogicalProject(category=[$0], __grand_total__=[$1], _row_number_chart_=[ROW_NUMBER() OVER (ORDER BY $1 DESC NULLS LAST)])
-                LogicalAggregate(group=[{0}], __grand_total__=[MAX($1)])
+                LogicalAggregate(group=[{0}], __grand_total__=[SUM($1)])
                   LogicalFilter(condition=[IS NOT NULL($0)])
                     LogicalProject(category=[$0], max(value)=[$2])
                       LogicalAggregate(group=[{0, 2}], max(value)=[MAX($1)])
@@ -28,4 +28,5 @@ calcite:
               EnumerableSort(sort0=[$0], dir0=[ASC])
                 EnumerableCalc(expr#0..2=[{inputs}], category=[$t0], $1=[$t2])
                   EnumerableWindow(window#0=[window(order by [1 DESC-nulls-last] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
-                    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_time_data]], PushDownContext=[[PROJECT->[category, value, timestamp], FILTER->IS NOT NULL($2), FILTER->IS NOT NULL($0), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},__grand_total__=MAX($1))], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"bool":{"filter":[{"exists":{"field":"timestamp","boost":1.0}},{"exists":{"field":"category","boost":1.0}}],"adjust_pure_negative":true,"boost":1.0}},"_source":{"includes":["category","value","timestamp"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"category":{"terms":{"field":"category","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"__grand_total__":{"max":{"field":"value"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
+                    EnumerableAggregate(group=[{0}], __grand_total__=[SUM($1)])
+                      CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_time_data]], PushDownContext=[[PROJECT->[category, value, timestamp], FILTER->IS NOT NULL($2), FILTER->IS NOT NULL($0), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0, 2},max(value)=MAX($1)), PROJECT->[category, max(value)]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"bool":{"filter":[{"exists":{"field":"timestamp","boost":1.0}},{"exists":{"field":"category","boost":1.0}}],"adjust_pure_negative":true,"boost":1.0}},"_source":{"includes":["category","value","timestamp"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"category":{"terms":{"field":"category","missing_bucket":true,"missing_order":"first","order":"asc"}}},{"timestamp0":{"date_histogram":{"field":"timestamp","missing_bucket":false,"order":"asc","calendar_interval":"1w"}}}]},"aggregations":{"max(value)":{"max":{"field":"value"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

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

@@ -10,7 +10,7 @@ calcite:
                   LogicalFilter(condition=[IS NOT NULL($23)])
                     CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_otel_logs]])
               LogicalProject(severityText=[$0], __grand_total__=[$1], _row_number_chart_=[ROW_NUMBER() OVER (ORDER BY $1 DESC NULLS LAST)])
-                LogicalAggregate(group=[{0}], __grand_total__=[MAX($1)])
+                LogicalAggregate(group=[{0}], __grand_total__=[SUM($1)])
                   LogicalFilter(condition=[IS NOT NULL($0)])
                     LogicalProject(severityText=[$1], max(severityNumber)=[$2])
                       LogicalAggregate(group=[{0, 1}], max(severityNumber)=[MAX($2)])
@@ -27,4 +27,5 @@ calcite:
               EnumerableSort(sort0=[$0], dir0=[ASC])
                 EnumerableCalc(expr#0..2=[{inputs}], severityText=[$t0], $1=[$t2])
                   EnumerableWindow(window#0=[window(order by [1 DESC-nulls-last] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
-                    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_otel_logs]], PushDownContext=[[PROJECT->[severityText, flags, severityNumber], FILTER->IS NOT NULL($1), FILTER->IS NOT NULL($0), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={1},__grand_total__=MAX($2))], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"bool":{"filter":[{"exists":{"field":"flags","boost":1.0}},{"exists":{"field":"severityText","boost":1.0}}],"adjust_pure_negative":true,"boost":1.0}},"_source":{"includes":["severityText","flags","severityNumber"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"severityText":{"terms":{"field":"severityText","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"__grand_total__":{"max":{"field":"severityNumber"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
+                    EnumerableAggregate(group=[{0}], __grand_total__=[SUM($1)])
+                      CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_otel_logs]], PushDownContext=[[PROJECT->[severityText, flags, severityNumber], FILTER->IS NOT NULL($1), FILTER->IS NOT NULL($0), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0, 1},max(severityNumber)=MAX($2)), PROJECT->[severityText, max(severityNumber)]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"bool":{"filter":[{"exists":{"field":"flags","boost":1.0}},{"exists":{"field":"severityText","boost":1.0}}],"adjust_pure_negative":true,"boost":1.0}},"_source":{"includes":["severityText","flags","severityNumber"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"flags":{"terms":{"field":"flags","missing_bucket":true,"missing_order":"first","order":"asc"}}},{"severityText":{"terms":{"field":"severityText","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"max(severityNumber)":{"max":{"field":"severityNumber"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])