Fix bug that Streamstats command incorrectly treats null as a valid group (#4777)

* use null to cover the aggregate value when group value is null

Signed-off-by: Xinyu Hao <[email protected]>

* fix test

Signed-off-by: Xinyu Hao <[email protected]>

* fix some problems

Signed-off-by: Xinyu Hao <[email protected]>

* add IT case

Signed-off-by: Xinyu Hao <[email protected]>

---------

Signed-off-by: Xinyu Hao <[email protected]>

Author: ishaoxy

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

@@ -1703,7 +1703,10 @@ public RelNode visitStreamWindow(StreamWindow node, CalcitePlanContext context)
           new String[] {ROW_NUMBER_COLUMN_FOR_STREAMSTATS});
     }
 
-    // Default
+    // Default: first get rawExpr
+    List<RexNode> overExpressions =
+        node.getWindowFunctionList().stream().map(w -> rexVisitor.analyze(w, context)).toList();
+
     if (hasGroup) {
       // only build sequence when there is by condition
       RexNode streamSeq =
@@ -1714,21 +1717,54 @@ public RelNode visitStreamWindow(StreamWindow node, CalcitePlanContext context)
               .rowsTo(RexWindowBounds.CURRENT_ROW)
               .as(ROW_NUMBER_COLUMN_FOR_STREAMSTATS);
       context.relBuilder.projectPlus(streamSeq);
-    }
 
-    List<RexNode> overExpressions =
-        node.getWindowFunctionList().stream().map(w -> rexVisitor.analyze(w, context)).toList();
-    context.relBuilder.projectPlus(overExpressions);
+      // construct groupNotNull predicate
+      List<RexNode> groupByList =
+          groupList.stream().map(expr -> rexVisitor.analyze(expr, context)).toList();
+      List<RexNode> notNullList =
+          PlanUtils.getSelectColumns(groupByList).stream()
+              .map(context.relBuilder::field)
+              .map(context.relBuilder::isNotNull)
+              .toList();
+      RexNode groupNotNull = context.relBuilder.and(notNullList);
 
-    // resort when there is by condition
-    if (hasGroup) {
+      // wrap each expr: CASE WHEN groupNotNull THEN rawExpr ELSE CAST(NULL AS rawType) END
+      List<RexNode> wrappedOverExprs =
+          wrapWindowFunctionsWithGroupNotNull(overExpressions, groupNotNull, context);
+      context.relBuilder.projectPlus(wrappedOverExprs);
+      // resort when there is by condition
       context.relBuilder.sort(context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_STREAMSTATS));
       context.relBuilder.projectExcept(context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_STREAMSTATS));
+    } else {
+      context.relBuilder.projectPlus(overExpressions);
     }
 
     return context.relBuilder.peek();
   }
 
+  private List<RexNode> wrapWindowFunctionsWithGroupNotNull(
+      List<RexNode> overExpressions, RexNode groupNotNull, CalcitePlanContext context) {
+    List<RexNode> wrappedOverExprs = new ArrayList<>(overExpressions.size());
+    for (RexNode overExpr : overExpressions) {
+      RexNode rawExpr = overExpr;
+      String aliasName = null;
+      if (overExpr instanceof RexCall rc && rc.getOperator() == SqlStdOperatorTable.AS) {
+        rawExpr = rc.getOperands().get(0);
+        if (rc.getOperands().size() >= 2 && rc.getOperands().get(1) instanceof RexLiteral lit) {
+          aliasName = lit.getValueAs(String.class);
+        }
+      }
+      RexNode nullLiteral = context.rexBuilder.makeNullLiteral(rawExpr.getType());
+      RexNode caseExpr =
+          context.rexBuilder.makeCall(SqlStdOperatorTable.CASE, groupNotNull, rawExpr, nullLiteral);
+      if (aliasName != null) {
+        caseExpr = context.relBuilder.alias(caseExpr, aliasName);
+      }
+      wrappedOverExprs.add(caseExpr);
+    }
+    return wrappedOverExprs;
+  }
+
   private RelNode buildStreamWindowJoinPlan(
       CalcitePlanContext context,
       RelNode leftWithHelpers,

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

@@ -745,6 +745,7 @@ public void testExplainOnFirstLast() throws IOException {
   }
 
   // Only for Calcite
+  @Test
   public void testExplainOnEventstatsEarliestLatest() throws IOException {
     String expected = loadExpectedPlan("explain_eventstats_earliest_latest.json");
     assertJsonEqualsIgnoreId(
@@ -782,6 +783,7 @@ public void testExplainOnEventstatsEarliestLatestNoGroupBy() throws IOException
                 TEST_INDEX_LOGS)));
   }
 
+  @Test
   public void testExplainOnStreamstatsEarliestLatest() throws IOException {
     String expected = loadExpectedPlan("explain_streamstats_earliest_latest.yaml");
     assertYamlEqualsIgnoreId(