[FLINK-40095][table-planner] Fix JSON aggregation functions failing in group window aggregate

flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/calcite/RelTimeIndicatorConverter.java

@@ -544,6 +544,7 @@ private FlinkLogicalWindowAggregate visitWindowAggregate(FlinkLogicalWindowAggre
         return new FlinkLogicalWindowAggregate(
                 agg.getCluster(),
                 agg.getTraitSet(),
+                agg.getHints(),
                 newInput,
                 agg.getGroupSet(),
                 updatedAggCalls,
@@ -556,6 +557,7 @@ private RelNode visitWindowTableAggregate(FlinkLogicalWindowTableAggregate table
                 new FlinkLogicalWindowAggregate(
                         tableAgg.getCluster(),
                         tableAgg.getTraitSet(),
+                        List.of(),
                         tableAgg.getInput(),
                         tableAgg.getGroupSet(),
                         tableAgg.getAggCallList(),

flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/hint/FlinkHintStrategies.java

@@ -82,7 +82,9 @@ public static HintStrategyTable createHintStrategyTable() {
                 .hintStrategy(
                         FlinkHints.HINT_NAME_JSON_AGGREGATE_WRAPPED,
                         HintStrategy.builder(HintPredicates.AGGREGATE)
-                                .excludedRules(WrapJsonAggFunctionArgumentsRule.INSTANCE)
+                                .excludedRules(
+                                        WrapJsonAggFunctionArgumentsRule.AGGREGATE_INSTANCE,
+                                        WrapJsonAggFunctionArgumentsRule.WINDOW_AGGREGATE_INSTANCE)
                                 .build())
                 // internal query hint used for alias
                 .hintStrategy(

flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRule.java

@@ -22,12 +22,14 @@
 import org.apache.flink.table.functions.BuiltInFunctionDefinitions;
 import org.apache.flink.table.planner.functions.bridging.BridgingSqlFunction;
 import org.apache.flink.table.planner.hint.FlinkHints;
+import org.apache.flink.table.planner.plan.nodes.calcite.LogicalWindowAggregate;
 
 import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.plan.RelRule;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Aggregate;
 import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.rel.hint.RelHint;
 import org.apache.calcite.rel.logical.LogicalAggregate;
@@ -63,13 +65,19 @@
  * their correct representation, and the actual aggregation function's implementation can simply
  * insert the values as raw nodes instead. This avoids having to re-implement the logic for all
  * supported types in the aggregation function again.
+ *
+ * <p>This rule supports both {@link LogicalAggregate} and {@link LogicalWindowAggregate}.
  */
 @Internal
 @Value.Enclosing
 public class WrapJsonAggFunctionArgumentsRule
         extends RelRule<WrapJsonAggFunctionArgumentsRule.Config> {
 
-    public static final RelOptRule INSTANCE = new WrapJsonAggFunctionArgumentsRule(Config.DEFAULT);
+    public static final RelOptRule AGGREGATE_INSTANCE =
+            new WrapJsonAggFunctionArgumentsRule(Config.AGGREGATE_DEFAULT);
+
+    public static final RelOptRule WINDOW_AGGREGATE_INSTANCE =
+            new WrapJsonAggFunctionArgumentsRule(Config.WINDOW_AGGREGATE_DEFAULT);
 
     /** Marker hint that a call has already been transformed. */
     private static final RelHint MARKER_HINT =
@@ -81,15 +89,15 @@ public WrapJsonAggFunctionArgumentsRule(Config config) {
 
     @Override
     public void onMatch(RelOptRuleCall call) {
-        final LogicalAggregate aggregate = call.rel(0);
+        final Aggregate aggregate = call.rel(0);
         final RelNode aggInput = aggregate.getInput();
         final RelBuilder relBuilder = call.builder().push(aggInput);
 
-        final LogicalAggregate wrappedAggregate = wrapJsonAggregate(aggregate, relBuilder);
+        final Aggregate wrappedAggregate = wrapJsonAggregate(aggregate, relBuilder);
         call.transformTo(wrappedAggregate.withHints(Collections.singletonList(MARKER_HINT)));
     }
 
-    private LogicalAggregate wrapJsonAggregate(LogicalAggregate aggregate, RelBuilder relBuilder) {
+    private Aggregate wrapJsonAggregate(Aggregate aggregate, RelBuilder relBuilder) {
         final int inputCount = aggregate.getInput().getRowType().getFieldCount();
         List<AggregateCall> aggCallList = new ArrayList<>(aggregate.getAggCallList());
         // This map is a mapping relationship between jsonObjectAggCall and the argument index
@@ -175,27 +183,33 @@ private static boolean isJsonAggregation(AggregateCall aggCall) {
     /** Configuration for {@link WrapJsonAggFunctionArgumentsRule}. */
     @Value.Immutable(singleton = false)
     public interface Config extends RelRule.Config {
-        Config DEFAULT =
+        Config AGGREGATE_DEFAULT =
+                ImmutableWrapJsonAggFunctionArgumentsRule.Config.builder()
+                        .build()
+                        .as(Config.class)
+                        .onAggregateClass(LogicalAggregate.class);
+
+        Config WINDOW_AGGREGATE_DEFAULT =
                 ImmutableWrapJsonAggFunctionArgumentsRule.Config.builder()
                         .build()
                         .as(Config.class)
-                        .onJsonAggregateFunctions();
+                        .onAggregateClass(LogicalWindowAggregate.class);
 
         @Override
         default RelOptRule toRule() {
             return new WrapJsonAggFunctionArgumentsRule(this);
         }
 
-        default Config onJsonAggregateFunctions() {
-            final Predicate<LogicalAggregate> jsonAggPredicate =
+        default <T extends Aggregate> Config onAggregateClass(Class<T> aggregateClass) {
+            final Predicate<T> jsonAggPredicate =
                     aggregate ->
                             aggregate.getAggCallList().stream()
                                     .anyMatch(WrapJsonAggFunctionArgumentsRule::isJsonAggregation);
 
             final RelRule.OperandTransform aggTransform =
                     operandBuilder ->
                             operandBuilder
-                                    .operand(LogicalAggregate.class)
+                                    .operand(aggregateClass)
                                     .predicate(jsonAggPredicate)
                                     .anyInputs();
 

flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/calcite/LogicalWindowAggregate.scala

@@ -24,19 +24,41 @@ import org.apache.calcite.plan.{Convention, RelOptCluster, RelTraitSet}
 import org.apache.calcite.rel.RelNode
 import org.apache.calcite.rel.core.{Aggregate, AggregateCall}
 import org.apache.calcite.rel.core.Aggregate.Group
+import org.apache.calcite.rel.hint.RelHint
 import org.apache.calcite.util.ImmutableBitSet
 
 import java.util
 
 final class LogicalWindowAggregate(
     cluster: RelOptCluster,
     traitSet: RelTraitSet,
+    hints: util.List[RelHint],
     child: RelNode,
     groupSet: ImmutableBitSet,
     aggCalls: util.List[AggregateCall],
     window: LogicalWindow,
     namedProperties: util.List[NamedWindowProperty])
-  extends WindowAggregate(cluster, traitSet, child, groupSet, aggCalls, window, namedProperties) {
+  extends WindowAggregate(
+    cluster,
+    traitSet,
+    hints,
+    child,
+    groupSet,
+    aggCalls,
+    window,
+    namedProperties) {
+
+  override def withHints(hintList: util.List[RelHint]): RelNode = {
+    new LogicalWindowAggregate(
+      cluster,
+      traitSet,
+      hintList,
+      input,
+      getGroupSet,
+      aggCalls,
+      window,
+      namedProperties)
+  }
 
   override def copy(
       traitSet: RelTraitSet,
@@ -47,6 +69,7 @@ final class LogicalWindowAggregate(
     new LogicalWindowAggregate(
       cluster,
       traitSet,
+      getHints,
       input,
       groupSet,
       aggCalls,
@@ -58,6 +81,7 @@ final class LogicalWindowAggregate(
     new LogicalWindowAggregate(
       cluster,
       traitSet,
+      getHints,
       input,
       getGroupSet,
       aggCalls,
@@ -76,6 +100,7 @@ final class LogicalWindowAggregate(
     new LogicalWindowAggregate(
       cluster,
       traitSet,
+      getHints,
       input,
       groupSet,
       aggCalls,
@@ -97,6 +122,7 @@ object LogicalWindowAggregate {
     new LogicalWindowAggregate(
       cluster,
       traitSet,
+      agg.getHints,
       agg.getInput,
       agg.getGroupSet,
       agg.getAggCallList,

flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/calcite/WindowAggregate.scala

@@ -41,6 +41,7 @@ import scala.collection.JavaConverters._
 abstract class WindowAggregate(
     cluster: RelOptCluster,
     traitSet: RelTraitSet,
+    hints: util.List[RelHint],
     child: RelNode,
     groupSet: ImmutableBitSet,
     aggCalls: util.List[AggregateCall],
@@ -49,7 +50,7 @@ abstract class WindowAggregate(
   extends Aggregate(
     cluster,
     traitSet,
-    new util.ArrayList[RelHint],
+    hints,
     child,
     groupSet,
     ImmutableList.of(groupSet),

flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/logical/FlinkLogicalWindowAggregate.scala

@@ -28,6 +28,7 @@ import org.apache.calcite.rel.convert.ConverterRule
 import org.apache.calcite.rel.convert.ConverterRule.Config
 import org.apache.calcite.rel.core.{Aggregate, AggregateCall}
 import org.apache.calcite.rel.core.Aggregate.Group
+import org.apache.calcite.rel.hint.RelHint
 import org.apache.calcite.rel.metadata.RelMetadataQuery
 import org.apache.calcite.sql.SqlKind
 import org.apache.calcite.util.ImmutableBitSet
@@ -39,12 +40,21 @@ import scala.collection.JavaConverters._
 class FlinkLogicalWindowAggregate(
     cluster: RelOptCluster,
     traitSet: RelTraitSet,
+    hints: util.List[RelHint],
     child: RelNode,
     groupSet: ImmutableBitSet,
     aggCalls: util.List[AggregateCall],
     window: LogicalWindow,
     namedProperties: util.List[NamedWindowProperty])
-  extends WindowAggregate(cluster, traitSet, child, groupSet, aggCalls, window, namedProperties)
+  extends WindowAggregate(
+    cluster,
+    traitSet,
+    hints,
+    child,
+    groupSet,
+    aggCalls,
+    window,
+    namedProperties)
   with FlinkLogicalRel {
 
   override def copy(
@@ -56,6 +66,7 @@ class FlinkLogicalWindowAggregate(
     new FlinkLogicalWindowAggregate(
       cluster,
       traitSet,
+      getHints,
       input,
       groupSet,
       aggCalls,
@@ -73,6 +84,17 @@ class FlinkLogicalWindowAggregate(
     planner.getCostFactory.makeCost(rowCnt, cpuCost, rowCnt * rowSize)
   }
 
+  override def withHints(hintList: util.List[RelHint]): RelNode = {
+    new FlinkLogicalWindowAggregate(
+      cluster,
+      traitSet,
+      hintList,
+      input,
+      getGroupSet,
+      aggCalls,
+      window,
+      namedProperties)
+  }
 }
 
 class FlinkLogicalWindowAggregateConverter(config: Config) extends ConverterRule(config) {
@@ -100,6 +122,7 @@ class FlinkLogicalWindowAggregateConverter(config: Config) extends ConverterRule
     new FlinkLogicalWindowAggregate(
       rel.getCluster,
       traitSet,
+      agg.getHints,
       newInput,
       agg.getGroupSet,
       agg.getAggCallList,

flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkBatchRuleSets.scala

@@ -139,7 +139,8 @@ object FlinkBatchRuleSets {
         // vector search rule.
         ConstantVectorSearchCallToCorrelateRule.INSTANCE,
         // Wrap arguments for JSON aggregate functions
-        WrapJsonAggFunctionArgumentsRule.INSTANCE,
+        WrapJsonAggFunctionArgumentsRule.AGGREGATE_INSTANCE,
+        WrapJsonAggFunctionArgumentsRule.WINDOW_AGGREGATE_INSTANCE,
         // prune COUNT(*) input to project a constant before aggregation
         PruneCountStarInputRule.INSTANCE
       )).asJava)

flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkStreamRuleSets.scala

@@ -139,7 +139,8 @@ object FlinkStreamRuleSets {
           // rewrite constant table function scan to correlate
           JoinTableFunctionScanToCorrelateRule.INSTANCE,
           // Wrap arguments for JSON aggregate functions
-          WrapJsonAggFunctionArgumentsRule.INSTANCE,
+          WrapJsonAggFunctionArgumentsRule.AGGREGATE_INSTANCE,
+          WrapJsonAggFunctionArgumentsRule.WINDOW_AGGREGATE_INSTANCE,
           // prune COUNT(*) input to project a constant before aggregation
           PruneCountStarInputRule.INSTANCE
         )

flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/BuiltInAggregateFunctionTestBase.java

@@ -106,10 +106,19 @@ final void test(TestCase testCase, @InjectMiniCluster MiniCluster miniCluster)
         testCase.execute(new MiniClusterClient(miniCluster.getConfiguration(), miniCluster));
     }
 
-    protected static Table asTable(TableEnvironment tEnv, DataType sourceRowType, List<Row> rows) {
+    protected static Table asTable(
+            TableEnvironment tEnv,
+            DataType sourceRowType,
+            List<Row> rows,
+            @Nullable String rowtimeColumn,
+            @Nullable String watermarkExpression) {
+        final Schema.Builder schemaBuilder = Schema.newBuilder().fromRowDataType(sourceRowType);
+        if (rowtimeColumn != null) {
+            schemaBuilder.watermark(rowtimeColumn, watermarkExpression);
+        }
         final TableDescriptor descriptor =
                 TableFactoryHarness.newBuilder()
-                        .schema(Schema.newBuilder().fromRowDataType(sourceRowType).build())
+                        .schema(schemaBuilder.build())
                         .source(asSource(rows, sourceRowType))
                         .build();
 
@@ -204,6 +213,8 @@ protected static class TestSpec {
         private final Configuration configuration = new Configuration();
         private DataType sourceRowType;
         private List<Row> sourceRows;
+        private @Nullable String rowtimeColumn;
+        private @Nullable String watermarkExpression;
 
         private TestSpec(BuiltInFunctionDefinition definition) {
             this.definition = definition;
@@ -238,6 +249,12 @@ TestSpec withSource(DataType sourceRowType, List<Row> sourceRows) {
             return this;
         }
 
+        TestSpec withWatermark(String rowtimeColumn, String watermarkExpression) {
+            this.rowtimeColumn = rowtimeColumn;
+            this.watermarkExpression = watermarkExpression;
+            return this;
+        }
+
         TestSpec testSqlResult(
                 Function<Table, String> sqlSpec, DataType expectedRowType, List<Row> expectedRows) {
             this.testItems.add(new SqlTestItem(sqlSpec, expectedRowType, expectedRows));
@@ -343,7 +360,13 @@ private TestCaseWithClusterClient createTestItemExecutable(
                                         .inStreamingMode()
                                         .withConfiguration(configuration)
                                         .build());
-                final Table sourceTable = asTable(tEnv, sourceRowType, sourceRows);
+                final Table sourceTable =
+                        asTable(
+                                tEnv,
+                                sourceRowType,
+                                sourceRows,
+                                rowtimeColumn,
+                                watermarkExpression);
 
                 testItem.execute(tEnv, sourceTable, clusterClient);
             };