Ignore rows without a row split in chart command

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

Author: yuancu

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

@@ -30,6 +30,7 @@
 import com.google.common.collect.Streams;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.BitSet;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashSet;
@@ -39,6 +40,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
 import org.apache.calcite.plan.RelOptTable;
@@ -1097,7 +1099,15 @@ private Pair<List<RexNode>, List<AggCall>> resolveAttributesForAggregation(
 
   @Override
   public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
-    visitAggregation(node, context, true);
+    Argument.ArgumentMap statsArgs = Argument.ArgumentMap.of(node.getArgExprList());
+    Boolean bucketNullable =
+        (Boolean) statsArgs.getOrDefault(Argument.BUCKET_NULLABLE, Literal.TRUE).getValue();
+    int nGroup = node.getGroupExprList().size() + (Objects.nonNull(node.getSpan()) ? 1 : 0);
+    BitSet nonNullGroupMask = new BitSet(nGroup);
+    if (!bucketNullable) {
+      nonNullGroupMask.set(0, nGroup);
+    }
+    visitAggregation(node, context, true, nonNullGroupMask);
     return context.relBuilder.peek();
   }
 
@@ -1108,8 +1118,10 @@ public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
    * @param context the Calcite plan context for building RelNodes
    * @param aggFirst if true, aggregation results (metrics) appear first in output schema (agg,
    *     group-by fields); if false, group expressions appear first (group-by fields, agg).
+   * @param nonNullGroupMask bit set indicating group by fields that need to be non-null
    */
-  private void visitAggregation(Aggregation node, CalcitePlanContext context, boolean aggFirst) {
+  private void visitAggregation(
+      Aggregation node, CalcitePlanContext context, boolean aggFirst, BitSet nonNullGroupMask) {
     visitChildren(node, context);
 
     List<UnresolvedExpression> aggExprList = node.getAggExprList();
@@ -1119,36 +1131,30 @@ private void visitAggregation(Aggregation node, CalcitePlanContext context, bool
     UnresolvedExpression span = node.getSpan();
     if (Objects.nonNull(span)) {
       groupExprList.add(span);
-      List<RexNode> timeSpanFilters =
-          getTimeSpanField(span).stream()
-              .map(f -> rexVisitor.analyze(f, context))
-              .map(context.relBuilder::isNotNull)
-              .toList();
-      if (!timeSpanFilters.isEmpty()) {
-        // add isNotNull filter before aggregation for time span
-        context.relBuilder.filter(timeSpanFilters);
+      if (getTimeSpanField(span).isPresent()){
+          nonNullGroupMask.set(0);
       }
     }
     groupExprList.addAll(node.getGroupExprList());
 
     // add stats hint to LogicalAggregation
-    Argument.ArgumentMap statsArgs = Argument.ArgumentMap.of(node.getArgExprList());
-    Boolean bucketNullable =
-        (Boolean) statsArgs.getOrDefault(Argument.BUCKET_NULLABLE, Literal.TRUE).getValue();
-    boolean toAddHintsOnAggregate = false;
-    if (!bucketNullable
-        && !groupExprList.isEmpty()
-        && !(groupExprList.size() == 1 && getTimeSpanField(span).isPresent())) {
-      toAddHintsOnAggregate = true;
-      // add isNotNull filter before aggregation for non-nullable buckets
-      List<RexNode> groupByList =
-          groupExprList.stream().map(expr -> rexVisitor.analyze(expr, context)).toList();
-      context.relBuilder.filter(
-          PlanUtils.getSelectColumns(groupByList).stream()
-              .map(context.relBuilder::field)
-              .map(context.relBuilder::isNotNull)
-              .toList());
-    }
+    boolean toAddHintsOnAggregate =
+        nonNullGroupMask.nextClearBit(0) >= groupExprList.size() // This checks if all group-bys are nonnull
+            && !groupExprList.isEmpty()
+            && !(groupExprList.size() == 1 && getTimeSpanField(span).isPresent());
+    // add isNotNull filter before aggregation for non-nullable buckets
+    List<RexNode> groupByList =
+        groupExprList.stream().map(expr -> rexVisitor.analyze(expr, context)).toList();
+    List<RexNode> nonNullGroupBys =
+        IntStream.range(0, groupByList.size())
+            .filter(nonNullGroupMask::get)
+            .mapToObj(groupByList::get)
+            .toList();
+    context.relBuilder.filter(
+        PlanUtils.getSelectColumns(nonNullGroupBys).stream()
+            .map(context.relBuilder::field)
+            .map(context.relBuilder::isNotNull)
+            .toList());
 
     Pair<List<RexNode>, List<AggCall>> aggregationAttributes =
         aggregateWithTrimming(groupExprList, aggExprList, context, toAddHintsOnAggregate);
@@ -2397,12 +2403,15 @@ public RelNode visitChart(Chart node, CalcitePlanContext context) {
     ChartConfig config = ChartConfig.fromArguments(argMap);
     Aggregation aggregation =
         new Aggregation(
-            List.of(node.getAggregationFunction()),
-            List.of(),
-            groupExprList,
-            null,
-            List.of(new Argument(Argument.BUCKET_NULLABLE, AstDSL.booleanLiteral(config.useNull))));
-    visitAggregation(aggregation, context, false);
+            List.of(node.getAggregationFunction()), List.of(), groupExprList, null, List.of());
+    BitSet nonNullGroupMask = new BitSet(groupExprList.size());
+    // Rows without a row-split are always ignored
+    if (config.useNull) {
+      nonNullGroupMask.set(0);
+    } else {
+      nonNullGroupMask.set(0, groupExprList.size());
+    }
+    visitAggregation(aggregation, context, false, nonNullGroupMask);
     RelBuilder relBuilder = context.relBuilder;
     String columnSplitName =
         relBuilder.peek().getRowType().getFieldNames().size() > 2

docs/user/ppl/cmd/chart.rst

@@ -30,28 +30,29 @@ Syntax
 
 **Parameters:**
 
-* **limit**: optional. Specifies the number of distinct values to display when using column split.
+* **limit**: optional. Specifies the number of categories to display when using column split. Each unique value in the column split field represents a category.
 
   * Default: top10
   * Syntax: ``limit=(top|bottom)<number>`` or ``limit=<number>`` (defaults to top)
-  * When ``limit=K`` is set, the top or bottom K distinct column split values are retained; the additional values are grouped into an "OTHER" category if ``useother`` is not set to false.
-  * Set limit to 0 to show all distinct values without any limit.
-  * Use ``limit=topK`` or ``limit=bottomK`` to specify whether to retain the top or bottom K column categories. The ranking is based on the aggregated values for each distinct column-split value.  For example, ``chart limit=top3 count() by a b`` retains the 3 most common b categories; ``chart limit=top5 min(value) by a b`` selects the 5 b categories that contains smallest aggregated values. If not specified, top is used by default.
-  * Only applies when column split presents (by 2 fields or over...by... coexists).
+  * When ``limit=K`` is set, the top or bottom K categories from the column split field are retained; the remaining categories are grouped into an "OTHER" category if ``useother`` is not set to false.
+  * Set limit to 0 to show all categories without any limit.
+  * Use ``limit=topK`` or ``limit=bottomK`` to specify whether to retain the top or bottom K column categories. The ranking is based on the aggregated values for each category.  For example, ``chart limit=top3 count() by a b`` retains the 3 most common b categories; ``chart limit=top5 min(value) by a b`` selects the 5 b categories that contain the smallest aggregated values. If not specified, top is used by default.
+  * Only applies when column split is present (by 2 fields or over...by... coexists).
 
-* **useother**: optional. Controls whether to create an "OTHER" category for distinct column-split values beyond the limit.
+* **useother**: optional. Controls whether to create an "OTHER" category for categories beyond the limit.
 
   * Default: true
-  * When set to false, only the top/bottom N distinct column-split values (based on limit) are shown without an "OTHER" category.
-  * When set to true, distinct values beyond the limit are grouped into an "OTHER" category.
-  * Only applies when using column split and when there are more distinct column-split values than the limit.
+  * When set to false, only the top/bottom N categories (based on limit) are shown without an "OTHER" category.
+  * When set to true, categories beyond the limit are grouped into an "OTHER" category.
+  * Only applies when using column split and when there are more categories than the limit.
 
 * **usenull**: optional. Controls whether to group events without a column split (i.e. whose column split is null) into a separate "NULL" category.
 
   * Default: true
+  * ``usenull`` only applies to column split.
+  * Row split should always be non-null value. Events with null values in row split will be ignored.
   * When ``usenull=false``, events with a null column split are excluded from results.
   * When ``usenull=true``, events with a null column split are grouped into a separate "NULL" category.
-  * ``usenull`` only applies to column split. Null values in the row split are handled in the same way as normal aggregations.
 
 * **nullstr**: optional. Specifies the category name for rows that do not contain the column split value.
 

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

@@ -314,4 +314,21 @@ public void testChartUseNullFalseWithNullStr() throws IOException {
         result, schema("gender", "string"), schema("age", "string"), schema("count()", "bigint"));
     verifyDataRows(result, rows("M", "30", 4), rows("F", "30", 1), rows("F", "20", 1));
   }
+
+  @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));
+  }
 }

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

@@ -8,14 +8,16 @@ calcite:
               LogicalProject(gender=[$0], age=[SAFE_CAST($1)], avg(balance)=[$2])
                 LogicalAggregate(group=[{0, 1}], avg(balance)=[AVG($2)])
                   LogicalProject(gender=[$4], age=[$10], balance=[$7])
-                    CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
+                    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)])
                   LogicalFilter(condition=[IS NOT NULL($0)])
                     LogicalProject(age=[SAFE_CAST($1)], avg(balance)=[$2])
                       LogicalAggregate(group=[{0, 1}], avg(balance)=[AVG($2)])
                         LogicalProject(gender=[$4], age=[$10], balance=[$7])
-                          CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
+                          LogicalFilter(condition=[IS NOT NULL($4)])
+                            CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
   physical: |
     EnumerableLimit(fetch=[10000])
       EnumerableSort(sort0=[$0], sort1=[$1], dir0=[ASC], dir1=[ASC])
@@ -25,11 +27,11 @@ calcite:
               EnumerableMergeJoin(condition=[=($1, $3)], joinType=[left])
                 EnumerableSort(sort0=[$1], dir0=[ASC])
                   EnumerableCalc(expr#0..2=[{inputs}], expr#3=[SAFE_CAST($t1)], gender=[$t0], age=[$t3], avg(balance)=[$t2])
-                    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0, 1},avg(balance)=AVG($2))], OpenSearchRequestBuilder(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"}}},{"age":{"terms":{"field":"age","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"avg(balance)":{"avg":{"field":"balance"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
+                    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))], 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)])
                 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)], $f0=[$t2], avg(balance)=[$t1], $condition=[$t3])
-                            CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[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","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)])
+                          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)])