Handle common agg functions for OTHER category for timechart

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

Author: yuancu

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

@@ -1926,7 +1926,7 @@ public RelNode visitFlatten(Flatten node, CalcitePlanContext context) {
   }
 
   /** Helper method to get the function name for proper column naming */
-  private String getValueFunctionName(UnresolvedExpression aggregateFunction) {
+  private String getAggFieldAlias(UnresolvedExpression aggregateFunction) {
     if (aggregateFunction instanceof Alias) {
       return ((Alias) aggregateFunction).getName();
     }
@@ -1976,15 +1976,15 @@ public RelNode visitTimechart(
 
     // Handle no by field case
     if (node.getByField() == null) {
-      String valueFunctionName = getValueFunctionName(node.getAggregateFunction());
+      String aggFieldAlias = getAggFieldAlias(node.getAggregateFunction());
 
       // Create group expression list with just the timestamp span but use a different alias
       // to avoid @timestamp naming conflict
       List<UnresolvedExpression> simpleGroupExprList = new ArrayList<>();
       simpleGroupExprList.add(new Alias("timestamp", spanExpr));
       // Create agg expression list with the aggregate function
       List<UnresolvedExpression> simpleAggExprList =
-          List.of(new Alias(valueFunctionName, node.getAggregateFunction()));
+          List.of(new Alias(aggFieldAlias, node.getAggregateFunction()));
       // Create an Aggregation object
       Aggregation aggregation =
           new Aggregation(
@@ -1999,9 +1999,9 @@ public RelNode visitTimechart(
       context.relBuilder.push(result);
       // Reorder fields: timestamp first, then count
       context.relBuilder.project(
-          context.relBuilder.field("timestamp"), context.relBuilder.field(valueFunctionName));
+          context.relBuilder.field("timestamp"), context.relBuilder.field(aggFieldAlias));
       // Rename timestamp to @timestamp
-      context.relBuilder.rename(List.of("@timestamp", valueFunctionName));
+      context.relBuilder.rename(List.of("@timestamp", aggFieldAlias));
 
       context.relBuilder.sort(context.relBuilder.field(0));
       return context.relBuilder.peek();
@@ -2010,7 +2010,7 @@ public RelNode visitTimechart(
     // Extract parameters for byField case
     UnresolvedExpression byField = node.getByField();
     String byFieldName = ((Field) byField).getField().toString();
-    String valueFunctionName = getValueFunctionName(node.getAggregateFunction());
+    String aggFieldAlias = getAggFieldAlias(node.getAggregateFunction());
 
     int limit = Optional.ofNullable(node.getLimit()).orElse(10);
     boolean useOther = Optional.ofNullable(node.getUseOther()).orElse(true);
@@ -2037,11 +2037,11 @@ public RelNode visitTimechart(
 
       // Handle no limit case - just sort and return with proper field aliases
       if (limit == 0) {
-        // Add final projection with proper aliases: [@timestamp, byField, valueFunctionName]
+        // Add final projection with proper aliases: [@timestamp, byField, aggFieldAlias]
         context.relBuilder.project(
             context.relBuilder.alias(context.relBuilder.field(0), "@timestamp"),
             context.relBuilder.alias(context.relBuilder.field(1), byFieldName),
-            context.relBuilder.alias(context.relBuilder.field(2), valueFunctionName));
+            context.relBuilder.alias(context.relBuilder.field(2), aggFieldAlias));
         context.relBuilder.sort(context.relBuilder.field(0), context.relBuilder.field(1));
         return context.relBuilder.peek();
       }
@@ -2051,32 +2051,61 @@ public RelNode visitTimechart(
 
       // Step 2: Find top N categories using window function approach (more efficient than separate
       // aggregation)
-      RelNode topCategories = buildTopCategoriesQuery(completeResults, limit, context);
+      String aggFunctionName = getAggFunctionName(node.getAggregateFunction());
+      Optional<BuiltinFunctionName> aggFuncNameOptional = BuiltinFunctionName.of(aggFunctionName);
+      if (aggFuncNameOptional.isEmpty()) {
+        throw new IllegalArgumentException(
+            StringUtils.format("Unrecognized aggregation function: %s", aggFunctionName));
+      }
+      BuiltinFunctionName aggFunction = aggFuncNameOptional.get();
+      RelNode topCategories = buildTopCategoriesQuery(completeResults, limit, aggFunction, context);
 
       // Step 3: Apply OTHER logic with single pass
       return buildFinalResultWithOther(
-          completeResults, topCategories, byFieldName, valueFunctionName, useOther, limit, context);
+          completeResults,
+          topCategories,
+          byFieldName,
+          aggFunction,
+          aggFieldAlias,
+          useOther,
+          limit,
+          context);
 
     } catch (Exception e) {
       throw new RuntimeException("Error in visitTimechart: " + e.getMessage(), e);
     }
   }
 
+  private String getAggFunctionName(UnresolvedExpression aggregateFunction) {
+    if (aggregateFunction instanceof Alias alias) {
+      return getAggFunctionName(alias.getDelegated());
+    }
+    return ((AggregateFunction) aggregateFunction).getFuncName();
+  }
+
   /** Build top categories query - simpler approach that works better with OTHER handling */
   private RelNode buildTopCategoriesQuery(
-      RelNode completeResults, int limit, CalcitePlanContext context) {
+      RelNode completeResults,
+      int limit,
+      BuiltinFunctionName aggFunction,
+      CalcitePlanContext context) {
     context.relBuilder.push(completeResults);
 
     // Filter out null values when determining top categories - null should not count towards limit
     context.relBuilder.filter(context.relBuilder.isNotNull(context.relBuilder.field(1)));
 
     // Get totals for non-null categories - field positions: 0=@timestamp, 1=byField, 2=value
+    RexInputRef valueField = context.relBuilder.field(2);
+    AggCall call = buildAggCall(context.relBuilder, aggFunction, valueField);
+
     context.relBuilder.aggregate(
-        context.relBuilder.groupKey(context.relBuilder.field(1)),
-        context.relBuilder.sum(context.relBuilder.field(2)).as("grand_total"));
+        context.relBuilder.groupKey(context.relBuilder.field(1)), call.as("grand_total"));
 
     // Apply sorting and limit to non-null categories only
-    context.relBuilder.sort(context.relBuilder.desc(context.relBuilder.field("grand_total")));
+    RexNode sortField = context.relBuilder.field("grand_total");
+    sortField =
+        aggFunction == BuiltinFunctionName.MIN ? sortField : context.relBuilder.desc(sortField);
+    context.relBuilder.sort(sortField);
     if (limit > 0) {
       context.relBuilder.limit(0, limit);
     }
@@ -2089,18 +2118,25 @@ private RelNode buildFinalResultWithOther(
       RelNode completeResults,
       RelNode topCategories,
       String byFieldName,
-      String valueFunctionName,
+      BuiltinFunctionName aggFunction,
+      String aggFieldAlias,
       boolean useOther,
       int limit,
       CalcitePlanContext context) {
 
     // Use zero-filling for count aggregations, standard result for others
-    if (valueFunctionName.equals("count")) {
+    if (aggFieldAlias.equals("count")) {
       return buildZeroFilledResult(
-          completeResults, topCategories, byFieldName, valueFunctionName, useOther, limit, context);
+          completeResults, topCategories, byFieldName, aggFieldAlias, useOther, limit, context);
     } else {
       return buildStandardResult(
-          completeResults, topCategories, byFieldName, valueFunctionName, useOther, context);
+          completeResults,
+          topCategories,
+          byFieldName,
+          aggFunction,
+          aggFieldAlias,
+          useOther,
+          context);
     }
   }
 
@@ -2109,7 +2145,8 @@ private RelNode buildStandardResult(
       RelNode completeResults,
       RelNode topCategories,
       String byFieldName,
-      String valueFunctionName,
+      BuiltinFunctionName aggFunctionName,
+      String aggFieldAlias,
       boolean useOther,
       CalcitePlanContext context) {
 
@@ -2132,11 +2169,13 @@ private RelNode buildStandardResult(
     context.relBuilder.project(
         context.relBuilder.alias(context.relBuilder.field(0), "@timestamp"),
         context.relBuilder.alias(categoryExpr, byFieldName),
-        context.relBuilder.alias(context.relBuilder.field(2), valueFunctionName));
+        context.relBuilder.alias(context.relBuilder.field(2), aggFieldAlias));
 
+    RexInputRef valueField = context.relBuilder.field(2);
+    AggCall aggCall = buildAggCall(context.relBuilder, aggFunctionName, valueField);
     context.relBuilder.aggregate(
         context.relBuilder.groupKey(context.relBuilder.field(0), context.relBuilder.field(1)),
-        context.relBuilder.sum(context.relBuilder.field(2)).as(valueFunctionName));
+        aggCall.as(aggFieldAlias));
 
     applyFiltersAndSort(useOther, context);
     return context.relBuilder.peek();
@@ -2171,7 +2210,7 @@ private RelNode buildZeroFilledResult(
       RelNode completeResults,
       RelNode topCategories,
       String byFieldName,
-      String valueFunctionName,
+      String aggFieldAlias,
       boolean useOther,
       int limit,
       CalcitePlanContext context) {
@@ -2210,7 +2249,7 @@ private RelNode buildZeroFilledResult(
             context.relBuilder.cast(context.relBuilder.field(0), SqlTypeName.TIMESTAMP),
             "@timestamp"),
         context.relBuilder.alias(context.relBuilder.field(1), byFieldName),
-        context.relBuilder.alias(context.relBuilder.literal(0), valueFunctionName));
+        context.relBuilder.alias(context.relBuilder.literal(0), aggFieldAlias));
     RelNode zeroFilledCombinations = context.relBuilder.build();
 
     // Get actual results with OTHER logic applied
@@ -2232,7 +2271,7 @@ private RelNode buildZeroFilledResult(
             context.relBuilder.cast(context.relBuilder.field(0), SqlTypeName.TIMESTAMP),
             "@timestamp"),
         context.relBuilder.alias(actualCategoryExpr, byFieldName),
-        context.relBuilder.alias(context.relBuilder.field(2), valueFunctionName));
+        context.relBuilder.alias(context.relBuilder.field(2), aggFieldAlias));
 
     context.relBuilder.aggregate(
         context.relBuilder.groupKey(context.relBuilder.field(0), context.relBuilder.field(1)),
@@ -2247,12 +2286,30 @@ private RelNode buildZeroFilledResult(
     // Aggregate to combine actual and zero-filled data
     context.relBuilder.aggregate(
         context.relBuilder.groupKey(context.relBuilder.field(0), context.relBuilder.field(1)),
-        context.relBuilder.sum(context.relBuilder.field(2)).as(valueFunctionName));
+        context.relBuilder.sum(context.relBuilder.field(2)).as(aggFieldAlias));
 
     applyFiltersAndSort(useOther, context);
     return context.relBuilder.peek();
   }
 
+  /**
+   * Aggregate a field based on a given built-in aggregation function name.
+   *
+   * <p>It is intended for secondary aggregations in timechart and chart commands. Using it
+   * elsewhere may lead to unintended results. It handles explicitly only MIN, MAX, AVG, COUNT,
+   * DISTINCT_COUNT, EARLIEST, and LATEST. It sums the results for the rest aggregation types,
+   * assuming them to be accumulative.
+   */
+  private AggCall buildAggCall(
+      RelBuilder relBuilder, BuiltinFunctionName aggFunction, RexNode node) {
+    return switch (aggFunction) {
+      case MIN, EARLIEST -> relBuilder.min(node);
+      case MAX, LATEST -> relBuilder.max(node);
+      case AVG -> relBuilder.avg(node);
+      default -> relBuilder.sum(node);
+    };
+  }
+
   @Override
   public RelNode visitTrendline(Trendline node, CalcitePlanContext context) {
     visitChildren(node, context);

integ-test/src/yamlRestTest/resources/rest-api-spec/test/issues/4582.yml

@@ -0,0 +1,120 @@
+setup:
+  - do:
+      query.settings:
+        body:
+          transient:
+            plugins.calcite.enabled : true
+  - do:
+      indices.create:
+        index: test_timechart_4582
+        body:
+          mappings:
+            properties:
+              "@timestamp":
+                type: date_nanos
+              severityNumber:
+                type: long
+              severityText:
+                type: keyword
+              body:
+                type: text
+  - do:
+      bulk:
+        index: test_timechart_4582
+        refresh: true
+        body:
+          - '{"index": {}}'
+          - '{"@timestamp": "2024-01-15T10:30:04.567890123Z", "severityNumber": 9, "severityText": "INFO", "body": "Info message"}'
+          - '{"index": {}}'
+          - '{"@timestamp": "2024-01-15T10:30:05.567890123Z", "severityNumber": 13, "severityText": "WARN", "body": "Warning message"}'
+          - '{"index": {}}'
+          - '{"@timestamp": "2024-01-15T10:30:06.567890123Z", "severityNumber": 17, "severityText": "ERROR", "body": "Error message"}'
+          - '{"index": {}}'
+          - '{"@timestamp": "2024-01-15T10:30:07.567890123Z", "severityNumber": 21, "severityText": "FATAL", "body": "Fatal message"}'
+          - '{"index": {}}'
+          - '{"@timestamp": "2024-01-15T10:30:08.567890123Z", "severityNumber": 24, "severityText": "FATAL4", "body": "Fatal4 message"}'
+          - '{"index": {}}'
+          - '{"@timestamp": "2024-01-15T10:30:09.567890123Z", "severityNumber": 23, "severityText": "DEBUG", "body": "Debug message"}'
+          - '{"index": {}}'
+          - '{"@timestamp": "2024-01-15T10:30:10.567890123Z", "severityNumber": 20, "severityText": "TRACE", "body": "Trace message"}'
+          - '{"index": {}}'
+          - '{"@timestamp": "2024-01-15T10:30:11.567890123Z", "severityNumber": 22, "severityText": "CUSTOM", "body": "Custom message"}'
+
+---
+teardown:
+  - do:
+      query.settings:
+        body:
+          transient:
+            plugins.calcite.enabled : false
+
+---
+"timechart max aggregation with limit should not sum OTHER values":
+  - skip:
+      features:
+        - headers
+        - allowed_warnings
+  - do:
+      headers:
+        Content-Type: 'application/json'
+      ppl:
+        body:
+          query: source=test_timechart_4582 | timechart limit=1 span=10seconds max(severityNumber) by severityText
+
+  - match: { total: 3 }
+  - match: { "schema": [{"name": "@timestamp", "type": "timestamp"}, {"name": "severityText", "type": "string"}, {"name": "max(severityNumber)", "type": "bigint"}] }
+  - match: { "datarows": [["2024-01-15 10:30:00", "FATAL4", 24], ["2024-01-15 10:30:00", "OTHER", 23], ["2024-01-15 10:30:10", "OTHER",22]] }
+
+---
+"timechart min aggregation with limit should not sum OTHER values":
+  - skip:
+      features:
+        - headers
+        - allowed_warnings
+  - do:
+      headers:
+        Content-Type: 'application/json'
+      ppl:
+        body:
+          query: source=test_timechart_4582 | timechart limit=2 span=1d min(severityNumber) by severityText
+
+  - match: { total: 3 }
+  - match: { "schema": [{"name": "@timestamp", "type": "timestamp"}, {"name": "severityText", "type": "string"}, {"name": "min(severityNumber)", "type": "bigint"}] }
+  - match: { "datarows": [["2024-01-15 00:00:00", "INFO", 9], ["2024-01-15 00:00:00", "OTHER", 17], ["2024-01-15 00:00:00", "WARN", 13]] }
+
+---
+"timechart earliest aggregation with limit should not sum OTHER values":
+  - skip:
+      features:
+        - headers
+        - allowed_warnings
+  - do:
+      headers:
+        Content-Type: 'application/json'
+      ppl:
+        body:
+          query: source=test_timechart_4582 | timechart limit=2 span=30seconds earliest(@timestamp) by severityText
+
+  - match: { total: 3 }
+  - match: { "schema": [{"name": "@timestamp", "type": "timestamp"}, {"name": "severityText", "type": "string"}, {"name": "earliest(@timestamp)", "type": "timestamp"}] }
+  - match: { "datarows": [
+    ["2024-01-15 10:30:00", "CUSTOM", "2024-01-15 10:30:11.567890123"],
+    ["2024-01-15 10:30:00", "OTHER", "2024-01-15 10:30:04.567890123"],
+    ["2024-01-15 10:30:00", "TRACE", "2024-01-15 10:30:10.567890123"]] }
+
+---
+"timechart count aggregation with limit should sum OTHER values":
+  - skip:
+      features:
+        - headers
+        - allowed_warnings
+  - do:
+      headers:
+        Content-Type: 'application/json'
+      ppl:
+        body:
+          query: source=test_timechart_4582 | timechart limit=3 span=1min count() by severityText
+
+  - match: { total: 4 }
+  - match: { "schema": [{"name": "@timestamp", "type": "timestamp"}, {"name": "severityText", "type": "string"}, {"name": "count", "type": "bigint"}] }
+  - match: { "datarows": [["2024-01-15 10:30:00", "CUSTOM", 1], ["2024-01-15 10:30:00", "DEBUG", 1], ["2024-01-15 10:30:00", "ERROR", 1], ["2024-01-15 10:30:00", "OTHER", 5]] }