Apply a non-null filter on fields referred by aggregations

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

Author: yuancu

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

@@ -1097,6 +1097,7 @@ private Pair<List<RexNode>, List<AggCall>> resolveAttributesForAggregation(
     return Pair.of(groupByList, aggCallList);
   }
 
+  /** Visits an aggregation for stats command */
   @Override
   public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
     Argument.ArgumentMap statsArgs = Argument.ArgumentMap.of(node.getArgExprList());
@@ -1107,7 +1108,7 @@ public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
     if (!bucketNullable) {
       nonNullGroupMask.set(0, nGroup);
     }
-    visitAggregation(node, context, true, nonNullGroupMask);
+    visitAggregation(node, context, nonNullGroupMask, true, false);
     return context.relBuilder.peek();
   }
 
@@ -1116,12 +1117,19 @@ public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
    *
    * @param node the aggregation node containing group expressions and aggregation functions
    * @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
+   * @param metricsFirst if true, aggregation results (metrics) appear first in output schema
+   *     (metrics, group-by fields); if false, group expressions appear first (group-by fields,
+   *     metrics).
+   * @param includeAggFieldsInNullFilter if true, also applies non-null filters to aggregation input
+   *     fields in addition to group-by fields
    */
   private void visitAggregation(
-      Aggregation node, CalcitePlanContext context, boolean aggFirst, BitSet nonNullGroupMask) {
+      Aggregation node,
+      CalcitePlanContext context,
+      BitSet nonNullGroupMask,
+      boolean metricsFirst,
+      boolean includeAggFieldsInNullFilter) {
     visitChildren(node, context);
 
     List<UnresolvedExpression> aggExprList = node.getAggExprList();
@@ -1143,15 +1151,23 @@ private void visitAggregation(
             // This checks if all group-bys should be nonnull
             && nonNullGroupMask.nextClearBit(0) >= groupExprList.size();
     // 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())
+    List<RexNode> nonNullCandidates =
+        groupExprList.stream()
+            .map(expr -> rexVisitor.analyze(expr, context))
+            .collect(Collectors.toCollection(ArrayList::new));
+    if (includeAggFieldsInNullFilter) {
+      nonNullCandidates.addAll(
+          PlanUtils.getInputRefsFromAggCall(
+              aggExprList.stream().map(expr -> aggVisitor.analyze(expr, context)).toList()));
+      nonNullGroupMask.set(groupExprList.size(), nonNullCandidates.size());
+    }
+    List<RexNode> nonNullFields =
+        IntStream.range(0, nonNullCandidates.size())
             .filter(nonNullGroupMask::get)
-            .mapToObj(groupByList::get)
+            .mapToObj(nonNullCandidates::get)
             .toList();
     context.relBuilder.filter(
-        PlanUtils.getSelectColumns(nonNullGroupBys).stream()
+        PlanUtils.getSelectColumns(nonNullFields).stream()
             .map(context.relBuilder::field)
             .map(context.relBuilder::isNotNull)
             .toList());
@@ -1175,7 +1191,7 @@ private void visitAggregation(
             .map(context.relBuilder::field)
             .map(f -> (RexNode) f)
             .toList();
-    if (aggFirst) {
+    if (metricsFirst) {
       // As an example, in command `stats count() by colA, colB`,
       // the sequence of output schema is "count, colA, colB".
       reordered.addAll(aggRexList);
@@ -2411,11 +2427,11 @@ public RelNode visitChart(Chart node, CalcitePlanContext context) {
     } else {
       nonNullGroupMask.set(0, groupExprList.size());
     }
-    visitAggregation(aggregation, context, false, nonNullGroupMask);
+    visitAggregation(aggregation, context, nonNullGroupMask, false, true);
     RelBuilder relBuilder = context.relBuilder;
 
-    // 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
+    // If a second split does not present or limit equals 0, we go no further for limit, nullstr,
+    // otherstr parameters because all truncating & renaming is performed on the column split
     if (node.getRowSplit() == null
         || node.getColumnSplit() == null
         || Objects.equals(config.limit, 0)) {

docs/user/ppl/cmd/chart.rst

@@ -36,7 +36,7 @@ Syntax
   * Syntax: ``limit=(top|bottom)<number>`` or ``limit=<number>`` (defaults to top)
   * 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.
+  * Use ``limit=topK`` or ``limit=bottomK`` to specify whether to retain the top or bottom K column categories. The ranking is based on the sum of aggregated values for each column category. For example, ``chart limit=top3 count() by region, product`` keeps the 3 products with the highest total counts across all regions. 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 categories beyond the limit.
@@ -50,7 +50,7 @@ Syntax
 
   * 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.
+  * Row split should always be non-null value. Documents 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.
 
@@ -84,7 +84,8 @@ Notes
 =====
 
 * The fields generated by column splitting are converted to strings so that they are compatible with ``nullstr`` and ``otherstr`` and can be used as column names once pivoted.
-* The aggregation metric appears as the last column in the result. Result columns are ordered as: [row-split] [column-split] [aggregation-metrics]
+* Documents with null values in fields used by the aggregation function are excluded from aggregation. For example, in ``chart avg(balance) over deptno, group``, documents where ``balance`` is null are excluded from the average calculation.
+* The aggregation metric appears as the last column in the result. Result columns are ordered as: [row-split] [column-split] [aggregation-metrics].
 
 Examples
 ========

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

@@ -279,11 +279,7 @@ public void testChartUseNullTrueWithNullStr() throws IOException {
         schema("age", "string"),
         schema("avg(balance)", "double"));
     verifyDataRows(
-        result,
-        rows("M", "30", 21702.5),
-        rows("F", "30", 48086.0),
-        rows("F", "20", 32838.0),
-        rows("F", "nil", null));
+        result, rows("M", "30", 21702.5), rows("F", "30", 48086.0), rows("F", "20", 32838.0));
   }
 
   @Test

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

@@ -1409,43 +1409,48 @@ public void testPushDownMinOrMaxAggOnDerivedField() throws IOException {
   @Test
   public void testExplainChartWithSingleGroupKey() throws IOException {
     assertYamlEqualsIgnoreId(
-        loadExpectedPlan("explain_chart_single_group_key.yaml"),
+        loadExpectedPlan("chart_single_group_key.yaml"),
         explainQueryYaml(
             String.format("source=%s | chart avg(balance) by gender", TEST_INDEX_BANK)));
 
     assertYamlEqualsIgnoreId(
-        loadExpectedPlan("explain_chart_with_span.yaml"),
+        loadExpectedPlan("chart_with_integer_span.yaml"),
         explainQueryYaml(
             String.format("source=%s | chart max(balance) by age span=10", TEST_INDEX_BANK)));
 
     assertYamlEqualsIgnoreId(
-        loadExpectedPlan("explain_chart_timestamp_span.yaml"),
+        loadExpectedPlan("chart_with_timestamp_span.yaml"),
         explainQueryYaml(
             String.format(
-                "source=%s | chart max(value) over timestamp span=1week by category",
-                TEST_INDEX_TIME_DATA)));
+                "source=%s | chart count by @timestamp span=1day", TEST_INDEX_TIME_DATA)));
   }
 
   @Test
   public void testExplainChartWithMultipleGroupKeys() throws IOException {
-    String expected = loadExpectedPlan("explain_chart_multiple_group_keys.yaml");
     assertYamlEqualsIgnoreId(
-        expected,
+        loadExpectedPlan("chart_multiple_group_keys.yaml"),
         explainQueryYaml(
             String.format("source=%s | chart avg(balance) over gender by age", TEST_INDEX_BANK)));
+
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("chart_timestamp_span_and_category.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | chart max(value) over timestamp span=1week by category",
+                TEST_INDEX_TIME_DATA)));
   }
 
   @Test
   public void testExplainChartWithLimits() throws IOException {
-    String expected = loadExpectedPlan("explain_chart_with_limit.yaml");
+    String expected = loadExpectedPlan("chart_with_limit.yaml");
     assertYamlEqualsIgnoreId(
         expected,
         explainQueryYaml(
             String.format(
                 "source=%s | chart limit=0 avg(balance) over state by gender", TEST_INDEX_BANK)));
 
     assertYamlEqualsIgnoreId(
-        loadExpectedPlan("explain_chart_use_other.yaml"),
+        loadExpectedPlan("chart_use_other.yaml"),
         explainQueryYaml(
             String.format(
                 "source=%s | chart limit=2 useother=true otherstr='max_among_other'"
@@ -1455,7 +1460,7 @@ public void testExplainChartWithLimits() throws IOException {
 
   @Test
   public void testExplainChartWithNullStr() throws IOException {
-    String expected = loadExpectedPlan("explain_chart_null_str.yaml");
+    String expected = loadExpectedPlan("chart_null_str.yaml");
     assertYamlEqualsIgnoreId(
         expected,
         explainQueryYaml(

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

@@ -0,0 +1,10 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(sort0=[$0], dir0=[ASC], fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalSort(sort0=[$0], dir0=[ASC])
+        LogicalAggregate(group=[{0}], avg(balance)=[AVG($1)])
+          LogicalProject(gender=[$4], balance=[$7])
+            LogicalFilter(condition=[AND(IS NOT NULL($4), IS NOT NULL($7))])
+              CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[FILTER->AND(IS NOT NULL($4), IS NOT NULL($7)), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},avg(balance)=AVG($1)), SORT->[0], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"bool":{"must":[{"exists":{"field":"gender","boost":1.0}},{"exists":{"field":"balance","boost":1.0}}],"adjust_pure_negative":true,"boost":1.0}},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":false,"order":"asc"}}}]},"aggregations":{"avg(balance)":{"avg":{"field":"balance"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

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

@@ -8,25 +8,25 @@ calcite:
               LogicalProject(timestamp=[$1], category=[$0], max(value)=[$2])
                 LogicalAggregate(group=[{0, 2}], max(value)=[MAX($1)])
                   LogicalProject(category=[$1], value=[$2], timestamp0=[SPAN($3, 1, 'w')])
-                    LogicalFilter(condition=[IS NOT NULL($3)])
+                    LogicalFilter(condition=[AND(IS NOT NULL($3), IS NOT NULL($2))])
                       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__=[SUM($1)])
                   LogicalFilter(condition=[IS NOT NULL($0)])
                     LogicalProject(category=[$0], max(value)=[$2])
                       LogicalAggregate(group=[{0, 2}], max(value)=[MAX($1)])
                         LogicalProject(category=[$1], value=[$2], timestamp0=[SPAN($3, 1, 'w')])
-                          LogicalFilter(condition=[IS NOT NULL($3)])
+                          LogicalFilter(condition=[AND(IS NOT NULL($3), IS NOT NULL($2))])
                             CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_time_data]])
   physical: |
     EnumerableLimit(fetch=[10000])
       EnumerableSort(sort0=[$0], sort1=[$1], dir0=[ASC], dir1=[ASC])
         EnumerableAggregate(group=[{0, 1}], max(value)=[MAX($2)])
           EnumerableCalc(expr#0..4=[{inputs}], expr#5=[IS NULL($t1)], expr#6=['NULL'], expr#7=[10], expr#8=[<=($t4, $t7)], expr#9=['OTHER'], expr#10=[CASE($t5, $t6, $t8, $t1, $t9)], timestamp=[$t0], category=[$t10], max(value)=[$t2])
             EnumerableMergeJoin(condition=[=($1, $3)], joinType=[left])
-              CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_time_data]], PushDownContext=[[PROJECT->[category, value, timestamp], FILTER->IS NOT NULL($2), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0, 2},max(value)=MAX($1)), PROJECT->[timestamp0, category, max(value)], SORT->[1]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"exists":{"field":"timestamp","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":"last","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)])
+              CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_time_data]], PushDownContext=[[PROJECT->[category, value, timestamp], FILTER->AND(IS NOT NULL($2), IS NOT NULL($1)), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0, 2},max(value)=MAX($1)), PROJECT->[timestamp0, category, max(value)], SORT->[1]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"bool":{"must":[{"exists":{"field":"timestamp","boost":1.0}},{"exists":{"field":"value","boost":1.0}}],"filter":[{"exists":{"field":"category","boost":1.0}},{"exists":{"field":"category","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":"last","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)])
               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()])])
                     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)])
+                      CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_time_data]], PushDownContext=[[PROJECT->[category, value, timestamp], FILTER->AND(IS NOT NULL($2), IS NOT NULL($1)), 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":{"must":[{"exists":{"field":"timestamp","boost":1.0}},{"exists":{"field":"value","boost":1.0}}],"filter":[{"exists":{"field":"category","boost":1.0}},{"exists":{"field":"category","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)])