diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRule.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRule.java index 089e542f4c956..ed8f1785757d4 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRule.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRule.java @@ -38,9 +38,6 @@ import org.apache.calcite.sql.fun.SqlJsonArrayAggAggFunction; import org.apache.calcite.sql.fun.SqlJsonObjectAggAggFunction; import org.apache.calcite.tools.RelBuilder; -import org.apache.calcite.util.mapping.MappingType; -import org.apache.calcite.util.mapping.Mappings; -import org.apache.calcite.util.mapping.Mappings.TargetMapping; import org.immutables.value.Value; import java.util.ArrayList; @@ -126,14 +123,16 @@ private Aggregate wrapJsonAggregate(Aggregate aggregate, RelBuilder relBuilder) valueIndicesAfterProjection); List newWrappedArgCallList = new ArrayList<>(aggCallList); - final int newInputCount = inputCount + valueIndicesAfterProjection.size(); for (Integer jsonAggCallIndex : wrapIndicesMap.keySet()) { - final TargetMapping argsMapping = - Mappings.create(MappingType.BIJECTION, newInputCount, newInputCount); - Integer valueIndex = wrapIndicesMap.get(jsonAggCallIndex); - argsMapping.set(valueIndex, valueIndicesAfterProjection.get(valueIndex)); - final AggregateCall newAggregateCall = - newWrappedArgCallList.get(jsonAggCallIndex).transform(argsMapping); + final AggregateCall aggregateCall = newWrappedArgCallList.get(jsonAggCallIndex); + final List newArgList = new ArrayList<>(aggregateCall.getArgList()); + // AggregateCall argument positions are zero-based: JSON_OBJECTAGG has (key, value), + // so its value is at position 1; JSON_ARRAYAGG has only (value), at position 0. + final int valueArgPosition = + aggregateCall.getAggregation() instanceof SqlJsonObjectAggAggFunction ? 1 : 0; + final Integer valueIndex = wrapIndicesMap.get(jsonAggCallIndex); + newArgList.set(valueArgPosition, valueIndicesAfterProjection.get(valueIndex)); + final AggregateCall newAggregateCall = aggregateCall.withArgList(newArgList); newWrappedArgCallList.set(jsonAggCallIndex, newAggregateCall); } diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/JsonAggregationFunctionsITCase.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/JsonAggregationFunctionsITCase.java index 32827b79d5a5d..4c352943b8d24 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/JsonAggregationFunctionsITCase.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/JsonAggregationFunctionsITCase.java @@ -60,6 +60,15 @@ public Stream getTestCaseSpecs() { ROW(VARCHAR(2000).notNull()), ROW(STRING().notNull()), Collections.singletonList(Row.of("{\"A\":1,\"B\":null,\"C\":3}"))), + TestSpec.forFunction(BuiltInFunctionDefinitions.JSON_OBJECTAGG_NULL_ON_NULL) + .withDescription("Aggregation With Same Key And Value") + .withSource( + ROW(STRING()), + Arrays.asList(Row.ofKind(INSERT, "A"), Row.ofKind(INSERT, "B"))) + .testSqlResult( + source -> "SELECT JSON_OBJECTAGG(f0 VALUE f0) FROM " + source, + ROW(VARCHAR(2000).notNull()), + Collections.singletonList(Row.of("{\"A\":\"A\",\"B\":\"B\"}"))), TestSpec.forFunction(BuiltInFunctionDefinitions.JSON_OBJECTAGG_ABSENT_ON_NULL) .withDescription("Omits NULLs") .withSource( @@ -273,6 +282,27 @@ public Stream getTestCaseSpecs() { + " GROUP BY TUMBLE(f2, INTERVAL '5' SECOND)", ROW(VARCHAR(2000).notNull()), Arrays.asList(Row.of("{\"A\":1,\"B\":2}"), Row.of("{\"C\":3}"))), + TestSpec.forFunction(BuiltInFunctionDefinitions.JSON_OBJECTAGG_NULL_ON_NULL) + .withDescription("Window Aggregation With Same Key And Value") + .withSource( + ROW(STRING(), TIMESTAMP(3)), + Arrays.asList( + Row.ofKind( + INSERT, + "A", + LocalDateTime.parse("2020-01-01T00:00:01")), + Row.ofKind( + INSERT, + "B", + LocalDateTime.parse("2020-01-01T00:00:02")))) + .withWatermark("f1", "f1 - INTERVAL '1' SECOND") + .testSqlResult( + source -> + "SELECT JSON_OBJECTAGG(f0 VALUE f0) FROM " + + source + + " GROUP BY TUMBLE(f1, INTERVAL '5' SECOND)", + ROW(VARCHAR(2000).notNull()), + Collections.singletonList(Row.of("{\"A\":\"A\",\"B\":\"B\"}"))), TestSpec.forFunction(BuiltInFunctionDefinitions.JSON_OBJECTAGG_NULL_ON_NULL) .withDescription("Window Group Aggregation With Other Aggs") .withSource( diff --git a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRuleTest.xml b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRuleTest.xml index e5480a1e27bcf..5f1a839794a81 100644 --- a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRuleTest.xml +++ b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRuleTest.xml @@ -718,7 +718,7 @@ LogicalAggregate(group=[{}], EXPR$0=[COUNT()], EXPR$1=[JSON_OBJECTAGG_NULL_ON_NU @@ -786,7 +786,7 @@ LogicalProject(EXPR$0=[$1], EXPR$1=[$2]) SortWindowAggregate(window=[TumblingGroupWindow('w$, rt, 5000)], select=[Final_COUNT(count1$0) AS EXPR$0, Final_JSON_OBJECTAGG_NULL_ON_NULL(EXPR$1) AS EXPR$1]) +- Sort(orderBy=[assignedWindow$ ASC]) +- Exchange(distribution=[single]) - +- LocalSortWindowAggregate(window=[TumblingGroupWindow('w$, rt, 5000)], select=[Partial_COUNT(*) AS count1$0, Partial_JSON_OBJECTAGG_NULL_ON_NULL($f2, $f2) AS EXPR$1]) + +- LocalSortWindowAggregate(window=[TumblingGroupWindow('w$, rt, 5000)], select=[Partial_COUNT(*) AS count1$0, Partial_JSON_OBJECTAGG_NULL_ON_NULL(f1, $f2) AS EXPR$1]) +- Calc(select=[rt, f1, JSON_STRING(f1) AS $f2]) +- Sort(orderBy=[rt ASC]) +- TableSourceScan(table=[[default_catalog, default_database, T, project=[rt, f1], metadata=[]]], fields=[rt, f1]) @@ -809,7 +809,7 @@ LogicalAggregate(group=[{}], EXPR$0=[JSON_OBJECTAGG_NULL_ON_NULL($0, $0)]) @@ -923,7 +923,7 @@ LogicalProject(EXPR$0=[$1]) SortWindowAggregate(window=[TumblingGroupWindow('w$, rt, 5000)], select=[Final_JSON_OBJECTAGG_NULL_ON_NULL(EXPR$0) AS EXPR$0]) +- Sort(orderBy=[assignedWindow$ ASC]) +- Exchange(distribution=[single]) - +- LocalSortWindowAggregate(window=[TumblingGroupWindow('w$, rt, 5000)], select=[Partial_JSON_OBJECTAGG_NULL_ON_NULL($f2, $f2) AS EXPR$0]) + +- LocalSortWindowAggregate(window=[TumblingGroupWindow('w$, rt, 5000)], select=[Partial_JSON_OBJECTAGG_NULL_ON_NULL(f1, $f2) AS EXPR$0]) +- Calc(select=[rt, f1, JSON_STRING(f1) AS $f2]) +- Sort(orderBy=[rt ASC]) +- TableSourceScan(table=[[default_catalog, default_database, T, project=[rt, f1], metadata=[]]], fields=[rt, f1])