Hint non-null in aggregateWithTrimming

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

Author: yuancu

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

@@ -896,12 +896,14 @@ private boolean isCountField(RexCall call) {
    * @param groupExprList group by expression list
    * @param aggExprList aggregate expression list
    * @param context CalcitePlanContext
+   * @param hintBucketNonNull adda bucket nullable hint on LogicalAggregate if set
    * @return Pair of (group-by list, field list, aggregate list)
    */
   private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
       List<UnresolvedExpression> groupExprList,
       List<UnresolvedExpression> aggExprList,
-      CalcitePlanContext context) {
+      CalcitePlanContext context,
+      boolean hintBucketNonNull) {
     Pair<List<RexNode>, List<AggCall>> resolved =
         resolveAttributesForAggregation(groupExprList, aggExprList, context);
     List<RexNode> resolvedGroupByList = resolved.getLeft();
@@ -1005,6 +1007,7 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
     List<String> intendedGroupKeyAliases = getGroupKeyNamesAfterAggregation(reResolved.getLeft());
     context.relBuilder.aggregate(
         context.relBuilder.groupKey(reResolved.getLeft()), reResolved.getRight());
+    if (hintBucketNonNull) hintBucketNonNullOnAggregate(context.relBuilder);
     // During aggregation, Calcite projects both input dependencies and output group-by fields.
     // When names conflict, Calcite adds numeric suffixes (e.g., "value0").
     // Apply explicit renaming to restore the intended aliases.
@@ -1013,6 +1016,24 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
     return Pair.of(reResolved.getLeft(), reResolved.getRight());
   }
 
+  private void hintBucketNonNullOnAggregate(RelBuilder relBuilder) {
+    final RelHint statHits =
+        RelHint.builder("stats_args").hintOption(Argument.BUCKET_NULLABLE, "false").build();
+    assert relBuilder.peek() instanceof LogicalAggregate
+        : "Stats hits should be added to LogicalAggregate";
+    relBuilder.hints(statHits);
+    relBuilder
+        .getCluster()
+        .setHintStrategies(
+            HintStrategyTable.builder()
+                .hintStrategy(
+                    "stats_args",
+                    (hint, rel) -> {
+                      return rel instanceof LogicalAggregate;
+                    })
+                .build());
+  }
+
   /**
    * Imitates {@code Registrar.registerExpression} of {@link RelBuilder} to derive the output order
    * of group-by keys after aggregation.
@@ -1114,25 +1135,7 @@ public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
     }
 
     Pair<List<RexNode>, List<AggCall>> aggregationAttributes =
-        aggregateWithTrimming(groupExprList, aggExprList, context);
-    if (toAddHintsOnAggregate) {
-      final RelHint statHits =
-          RelHint.builder("stats_args").hintOption(Argument.BUCKET_NULLABLE, "false").build();
-      assert context.relBuilder.peek() instanceof LogicalAggregate
-          : "Stats hits should be added to LogicalAggregate";
-      context.relBuilder.hints(statHits);
-      context
-          .relBuilder
-          .getCluster()
-          .setHintStrategies(
-              HintStrategyTable.builder()
-                  .hintStrategy(
-                      "stats_args",
-                      (hint, rel) -> {
-                        return rel instanceof LogicalAggregate;
-                      })
-                  .build());
-    }
+        aggregateWithTrimming(groupExprList, aggExprList, context, toAddHintsOnAggregate);
 
     // schema reordering
     // As an example, in command `stats count() by colA, colB`,
@@ -1869,7 +1872,7 @@ public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
     groupExprList.addAll(fieldList);
     List<UnresolvedExpression> aggExprList =
         List.of(AstDSL.alias(countFieldName, AstDSL.aggregate("count", null)));
-    aggregateWithTrimming(groupExprList, aggExprList, context);
+    aggregateWithTrimming(groupExprList, aggExprList, context, false);
 
     // 2. add a window column
     List<RexNode> partitionKeys = rexVisitor.analyze(node.getGroupExprList(), context);
@@ -2193,7 +2196,7 @@ public RelNode visitTimechart(
     try {
       // Step 1: Initial aggregation - IMPORTANT: order is [spanExpr, byField]
       groupExprList = Arrays.asList(spanExpr, byField);
-      aggregateWithTrimming(groupExprList, List.of(node.getAggregateFunction()), context);
+      aggregateWithTrimming(groupExprList, List.of(node.getAggregateFunction()), context, false);
 
       // First rename the timestamp field (2nd to last) to @timestamp
       List<String> fieldNames = context.relBuilder.peek().getRowType().getFieldNames();