diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/calcite/RelTimeIndicatorConverter.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/calcite/RelTimeIndicatorConverter.java index 262665155789d5..e447e45b1c67ee 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/calcite/RelTimeIndicatorConverter.java +++ b/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(), diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/hint/FlinkHintStrategies.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/hint/FlinkHintStrategies.java index 33e6d4346dd913..5978f98017d91d 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/hint/FlinkHintStrategies.java +++ b/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( 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 813034adf290cd..089e542f4c9563 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 @@ -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. + * + *

This rule supports both {@link LogicalAggregate} and {@link LogicalWindowAggregate}. */ @Internal @Value.Enclosing public class WrapJsonAggFunctionArgumentsRule extends RelRule { - 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 aggCallList = new ArrayList<>(aggregate.getAggCallList()); // This map is a mapping relationship between jsonObjectAggCall and the argument index @@ -175,19 +183,25 @@ 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 jsonAggPredicate = + default Config onAggregateClass(Class aggregateClass) { + final Predicate jsonAggPredicate = aggregate -> aggregate.getAggCallList().stream() .anyMatch(WrapJsonAggFunctionArgumentsRule::isJsonAggregation); @@ -195,7 +209,7 @@ default Config onJsonAggregateFunctions() { final RelRule.OperandTransform aggTransform = operandBuilder -> operandBuilder - .operand(LogicalAggregate.class) + .operand(aggregateClass) .predicate(jsonAggPredicate) .anyInputs(); diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/calcite/LogicalWindowAggregate.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/calcite/LogicalWindowAggregate.scala index 0d9e223b26867e..16eb483a901484 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/calcite/LogicalWindowAggregate.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/calcite/LogicalWindowAggregate.scala @@ -24,6 +24,7 @@ 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 @@ -31,12 +32,33 @@ 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, diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/calcite/WindowAggregate.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/calcite/WindowAggregate.scala index 333d80961d000e..fd50f547871595 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/calcite/WindowAggregate.scala +++ b/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), diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/logical/FlinkLogicalWindowAggregate.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/logical/FlinkLogicalWindowAggregate.scala index c483d660f0a38a..1f2ddab0806dcd 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/logical/FlinkLogicalWindowAggregate.scala +++ b/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, diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkBatchRuleSets.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkBatchRuleSets.scala index e840808b5e1c24..3fb057bc08d3eb 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkBatchRuleSets.scala +++ b/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) diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkStreamRuleSets.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkStreamRuleSets.scala index bc06d2d688baac..cf7654a2ce52bd 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkStreamRuleSets.scala +++ b/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 ) diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/BuiltInAggregateFunctionTestBase.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/BuiltInAggregateFunctionTestBase.java index 4c5f85a057ca56..f6b748948f4cd5 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/BuiltInAggregateFunctionTestBase.java +++ b/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 rows) { + protected static Table asTable( + TableEnvironment tEnv, + DataType sourceRowType, + List 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 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 sourceRows) { return this; } + TestSpec withWatermark(String rowtimeColumn, String watermarkExpression) { + this.rowtimeColumn = rowtimeColumn; + this.watermarkExpression = watermarkExpression; + return this; + } + TestSpec testSqlResult( Function sqlSpec, DataType expectedRowType, List 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); }; 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 34b00813d759d5..32827b79d5a5d6 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 @@ -22,6 +22,7 @@ import org.apache.flink.table.functions.BuiltInFunctionDefinitions; import org.apache.flink.types.Row; +import java.time.LocalDateTime; import java.util.Arrays; import java.util.Collections; import java.util.stream.Stream; @@ -29,6 +30,7 @@ import static org.apache.flink.table.api.DataTypes.INT; import static org.apache.flink.table.api.DataTypes.ROW; import static org.apache.flink.table.api.DataTypes.STRING; +import static org.apache.flink.table.api.DataTypes.TIMESTAMP; import static org.apache.flink.table.api.DataTypes.VARCHAR; import static org.apache.flink.table.api.Expressions.$; import static org.apache.flink.table.api.Expressions.jsonArrayAgg; @@ -240,6 +242,102 @@ public Stream getTestCaseSpecs() { ROW(INT(), STRING(), STRING().notNull()), Arrays.asList( Row.of(1, "A", "[\"A\"]"), - Row.of(2, "D", "[\"C\",\"D\"]")))); + Row.of(2, "D", "[\"C\",\"D\"]"))), + + // JSON_OBJECTAGG - Window (group window) Aggregation + TestSpec.forFunction(BuiltInFunctionDefinitions.JSON_OBJECTAGG_NULL_ON_NULL) + .withDescription("Window Aggregation") + .withSource( + ROW(STRING(), INT(), TIMESTAMP(3)), + Arrays.asList( + Row.ofKind( + INSERT, + "A", + 1, + LocalDateTime.parse("2020-01-01T00:00:01")), + Row.ofKind( + INSERT, + "B", + 2, + LocalDateTime.parse("2020-01-01T00:00:02")), + Row.ofKind( + INSERT, + "C", + 3, + LocalDateTime.parse("2020-01-01T00:00:06")))) + .withWatermark("f2", "f2 - INTERVAL '1' SECOND") + .testSqlResult( + source -> + "SELECT JSON_OBJECTAGG(f0 VALUE f1) FROM " + + source + + " 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 Group Aggregation With Other Aggs") + .withSource( + ROW(INT(), STRING(), INT(), TIMESTAMP(3)), + Arrays.asList( + Row.ofKind( + INSERT, + 1, + "A", + 1, + LocalDateTime.parse("2020-01-01T00:00:01")), + Row.ofKind( + INSERT, + 1, + "B", + 3, + LocalDateTime.parse("2020-01-01T00:00:02")), + Row.ofKind( + INSERT, + 2, + "A", + 2, + LocalDateTime.parse("2020-01-01T00:00:06")), + Row.ofKind( + INSERT, + 2, + "C", + 5, + LocalDateTime.parse("2020-01-01T00:00:07")))) + .withWatermark("f3", "f3 - INTERVAL '1' SECOND") + .testSqlResult( + source -> + "SELECT f0, JSON_OBJECTAGG(f1 VALUE f2), max(f2) FROM " + + source + + " GROUP BY f0, TUMBLE(f3, INTERVAL '5' SECOND)", + ROW(INT(), VARCHAR(2000).notNull(), INT()), + Arrays.asList( + Row.of(1, "{\"A\":1,\"B\":3}", 3), + Row.of(2, "{\"A\":2,\"C\":5}", 5))), + + // JSON_ARRAYAGG - Window (group window) Aggregation + TestSpec.forFunction(BuiltInFunctionDefinitions.JSON_ARRAYAGG_ABSENT_ON_NULL) + .withDescription("Window Aggregation") + .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")), + Row.ofKind( + INSERT, + "C", + LocalDateTime.parse("2020-01-01T00:00:06")))) + .withWatermark("f1", "f1 - INTERVAL '1' SECOND") + .testSqlResult( + source -> + "SELECT JSON_ARRAYAGG(f0) FROM " + + source + + " GROUP BY TUMBLE(f1, INTERVAL '5' SECOND)", + ROW(VARCHAR(2000).notNull()), + Arrays.asList(Row.of("[\"A\",\"B\"]"), Row.of("[\"C\"]")))); } } diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRuleTest.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRuleTest.java index e512d3b26d5451..e424cd2aec1f32 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRuleTest.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/WrapJsonAggFunctionArgumentsRuleTest.java @@ -24,27 +24,35 @@ import org.apache.flink.testutils.junit.extensions.parameterized.ParameterizedTestExtension; import org.apache.flink.testutils.junit.extensions.parameterized.Parameters; +import org.junit.jupiter.api.Assumptions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.TestTemplate; import org.junit.jupiter.api.extension.ExtendWith; +import java.util.ArrayList; import java.util.Arrays; -import java.util.Collection; +import java.util.List; /** Tests for {@link WrapJsonAggFunctionArgumentsRule}. */ @ExtendWith(ParameterizedTestExtension.class) class WrapJsonAggFunctionArgumentsRuleTest extends TableTestBase { private final boolean batchMode; + private final boolean isGroupWindowAgg; private TableTestUtil util; - @Parameters(name = "batchMode = {0}") - private static Collection data() { - return Arrays.asList(true, false); + @Parameters(name = "batchMode = {0}, isGroupWindowAgg = {1}") + private static List data() { + return Arrays.asList( + new Boolean[] {true, false}, + new Boolean[] {true, true}, + new Boolean[] {false, false}, + new Boolean[] {false, true}); } - WrapJsonAggFunctionArgumentsRuleTest(boolean batchMode) { + WrapJsonAggFunctionArgumentsRuleTest(boolean batchMode, boolean isGroupWindowAgg) { this.batchMode = batchMode; + this.isGroupWindowAgg = isGroupWindowAgg; } @BeforeEach @@ -67,70 +75,89 @@ void setup() { + batchMode + "'\n)"); - util.tableEnv() - .executeSql( - "CREATE TABLE T1(\n" - + " f0 INTEGER,\n" - + " f1 VARCHAR,\n" - + " f2 BIGINT\n" - + ") WITH (\n" - + " 'connector' = 'values'\n" - + " ,'bounded' = '" - + batchMode - + "'\n)"); + if (isGroupWindowAgg) { + util.tableEnv() + .executeSql( + "ALTER TABLE T add (rt TIMESTAMP(3), WATERMARK FOR rt as rt - interval '1' second)"); + } } @TestTemplate void testJsonObjectAgg() { - util.verifyRelPlan("SELECT JSON_OBJECTAGG(f1 VALUE f1) FROM T"); + util.verifyRelPlan(renderAggregateSQL("JSON_OBJECTAGG(f1 VALUE f1)")); } @TestTemplate void testJsonObjectAggInGroupWindow() { - util.verifyRelPlan("SELECT f0, JSON_OBJECTAGG(f1 VALUE f0) FROM T GROUP BY f0"); + util.verifyRelPlan(renderAggregateSQL("f0, JSON_OBJECTAGG(f1 VALUE f0)", "f0")); } @TestTemplate void testJsonArrayAgg() { - util.verifyRelPlan("SELECT JSON_ARRAYAGG(f0) FROM T"); + util.verifyRelPlan(renderAggregateSQL("JSON_ARRAYAGG(f0)")); } @TestTemplate void testJsonArrayAggInGroupWindow() { - util.verifyRelPlan("SELECT f0, JSON_ARRAYAGG(f0) FROM T GROUP BY f0"); + util.verifyRelPlan(renderAggregateSQL("f0, JSON_ARRAYAGG(f0)", "f0")); } @TestTemplate void testJsonObjectAggWithOtherAggs() { - util.verifyRelPlan("SELECT COUNT(*), JSON_OBJECTAGG(f1 VALUE f1) FROM T"); + util.verifyRelPlan(renderAggregateSQL("COUNT(*), JSON_OBJECTAGG(f1 VALUE f1)")); } @TestTemplate void testGroupJsonObjectAggWithOtherAggs() { util.verifyRelPlan( - "SELECT f0, COUNT(*), JSON_OBJECTAGG(f1 VALUE f0), SUM(f2) FROM T GROUP BY f0"); + renderAggregateSQL("f0, COUNT(*), JSON_OBJECTAGG(f1 VALUE f0), SUM(f2)", "f0")); } @TestTemplate void testJsonArrayAggWithOtherAggs() { - util.verifyRelPlan("SELECT COUNT(*), JSON_ARRAYAGG(f0) FROM T"); + util.verifyRelPlan(renderAggregateSQL("COUNT(*), JSON_ARRAYAGG(f0)")); } @TestTemplate void testGroupJsonArrayAggInWithOtherAggs() { - util.verifyRelPlan("SELECT f0, COUNT(*), JSON_ARRAYAGG(f0), SUM(f2) FROM T GROUP BY f0"); + util.verifyRelPlan(renderAggregateSQL("f0, COUNT(*), JSON_ARRAYAGG(f0), SUM(f2)", "f0")); } @TestTemplate void testJsonArrayAggAndJsonObjectAggWithOtherAggs() { util.verifyRelPlan( - "SELECT MAX(f0), JSON_OBJECTAGG(f1 VALUE f0), JSON_ARRAYAGG(f1), JSON_ARRAYAGG(f0) FROM T"); + renderAggregateSQL( + "MAX(f0), JSON_OBJECTAGG(f1 VALUE f0), JSON_ARRAYAGG(f1), JSON_ARRAYAGG(f0)")); } @TestTemplate void testGroupJsonArrayAggAndJsonObjectAggWithOtherAggs() { util.verifyRelPlan( - "SELECT f0, JSON_OBJECTAGG(f1 VALUE f2), JSON_ARRAYAGG(f1), JSON_ARRAYAGG(f2)," - + " SUM(f2) FROM T GROUP BY f0"); + renderAggregateSQL( + "f0, JSON_OBJECTAGG(f1 VALUE f2), JSON_ARRAYAGG(f1), JSON_ARRAYAGG(f2)," + + " SUM(f2)", + "f0")); + } + + @TestTemplate + void testJsonObjectAggWithWindowTVF() { + Assumptions.assumeTrue(isGroupWindowAgg); + util.verifyRelPlan( + "SELECT JSON_OBJECTAGG(f1 VALUE f1) " + + "FROM TABLE(TUMBLE(TABLE T, DESCRIPTOR(rt), INTERVAL '5' SECOND))"); + } + + private String renderAggregateSQL(String projects, String... groupKeys) { + List newGroupKeys = new ArrayList<>(Arrays.asList(groupKeys)); + if (isGroupWindowAgg) { + newGroupKeys.add("TUMBLE(rt, INTERVAL '5' SECOND)"); + } + String groupKeyStr; + if (newGroupKeys.isEmpty()) { + groupKeyStr = ""; + } else { + groupKeyStr = " GROUP BY " + String.join(",", newGroupKeys); + } + return String.format("SELECT %s FROM T%s", projects, groupKeyStr); } } 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 de4f8bb7341cb6..e5480a1e27bcfd 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 @@ -16,7 +16,7 @@ See the License for the specific language governing permissions and limitations under the License. --> - + @@ -35,7 +35,30 @@ GroupAggregate(groupBy=[f0], select=[f0, JSON_OBJECTAGG_NULL_ON_NULL(f1, $f4) AS ]]> - + + + + + + + + + + + + @@ -55,7 +78,29 @@ SortAggregate(isMerge=[false], groupBy=[f0], select=[f0, JSON_OBJECTAGG_NULL_ON_ ]]> - + + + + + + + + + + + + @@ -75,7 +120,30 @@ GroupAggregate(groupBy=[f0], select=[f0, COUNT(*) AS EXPR$1, JSON_ARRAYAGG_ABSEN ]]> - + + + + + + + + + + + + @@ -96,7 +164,29 @@ SortAggregate(isMerge=[false], groupBy=[f0], select=[f0, COUNT(*) AS EXPR$1, JSO ]]> - + + + + + + + + + + + + @@ -115,7 +205,30 @@ GroupAggregate(groupBy=[f0], select=[f0, COUNT(*) AS EXPR$1, JSON_OBJECTAGG_NULL ]]> - + + + + + + + + + + + + @@ -137,7 +250,31 @@ SortAggregate(isMerge=[true], groupBy=[f0], select=[f0, Final_COUNT(count1$0) AS ]]> - + + + + + + + + + + + + @@ -157,7 +294,30 @@ GroupAggregate(select=[MAX(f0) AS EXPR$0, JSON_OBJECTAGG_NULL_ON_NULL(f1, $f2) A ]]> - + + + + + + + + + + + + @@ -177,7 +337,29 @@ SortAggregate(isMerge=[false], select=[MAX(f0) AS EXPR$0, JSON_OBJECTAGG_NULL_ON ]]> - + + + + + + + + + + + + @@ -197,7 +379,30 @@ GroupAggregate(groupBy=[f0], select=[f0, JSON_ARRAYAGG_ABSENT_ON_NULL($f1) AS EX ]]> - + + + + + + + + + + + + @@ -218,7 +423,29 @@ SortAggregate(isMerge=[false], groupBy=[f0], select=[f0, JSON_ARRAYAGG_ABSENT_ON ]]> - + + + + + + + + + + + + @@ -238,7 +465,30 @@ GroupAggregate(select=[COUNT(*) AS EXPR$0, JSON_ARRAYAGG_ABSENT_ON_NULL($f1) AS ]]> - + + + + + + + + + + + + @@ -258,7 +508,29 @@ SortAggregate(isMerge=[false], select=[COUNT(*) AS EXPR$0, JSON_ARRAYAGG_ABSENT_ ]]> - + + + + + + + + + + + + @@ -278,7 +550,30 @@ GroupAggregate(select=[JSON_ARRAYAGG_ABSENT_ON_NULL($f1) AS EXPR$0]) ]]> - + + + + + + + + + + + + @@ -298,7 +593,29 @@ SortAggregate(isMerge=[false], select=[JSON_ARRAYAGG_ABSENT_ON_NULL($f1) AS EXPR ]]> - + + + + + + + + + + + + @@ -318,7 +635,30 @@ GroupAggregate(groupBy=[f0], select=[f0, JSON_OBJECTAGG_NULL_ON_NULL(f1, $f2) AS ]]> - + + + + + + + + + + + + @@ -341,7 +681,31 @@ SortAggregate(isMerge=[true], groupBy=[f0], select=[f0, Final_JSON_OBJECTAGG_NUL ]]> - + + + + + + + + + + + + @@ -361,7 +725,30 @@ GroupAggregate(select=[COUNT(*) AS EXPR$0, JSON_OBJECTAGG_NULL_ON_NULL($f1, $f1) ]]> - + + + + + + + + + + + + @@ -382,7 +769,80 @@ SortAggregate(isMerge=[true], select=[Final_COUNT(count1$0) AS EXPR$0, Final_JSO ]]> - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -402,7 +862,30 @@ GroupAggregate(select=[JSON_OBJECTAGG_NULL_ON_NULL($f1, $f1) AS EXPR$0]) ]]> - + + + + + + + + + + + + @@ -420,6 +903,30 @@ SortAggregate(isMerge=[true], select=[Final_JSON_OBJECTAGG_NULL_ON_NULL(EXPR$0) +- LocalSortAggregate(select=[Partial_JSON_OBJECTAGG_NULL_ON_NULL($f1, $f1) AS EXPR$0]) +- Calc(select=[f1, JSON_STRING(f1) AS $f1]) +- TableSourceScan(table=[[default_catalog, default_database, T, project=[f1], metadata=[]]], fields=[f1]) +]]> + + + + + + + + + + + diff --git a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/metadata/FlinkRelMdHandlerTestBase.scala b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/metadata/FlinkRelMdHandlerTestBase.scala index d210e06b43ff25..0837f65b3e2263 100644 --- a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/metadata/FlinkRelMdHandlerTestBase.scala +++ b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/metadata/FlinkRelMdHandlerTestBase.scala @@ -1834,6 +1834,7 @@ class FlinkRelMdHandlerTestBase { val logicalWindowAgg = new LogicalWindowAggregate( ts.getCluster, ts.getTraitSet, + util.List.of(), project, ImmutableBitSet.of(0, 1), aggCallOfWindowAgg, @@ -1845,6 +1846,7 @@ class FlinkRelMdHandlerTestBase { val flinkLogicalWindowAgg = new FlinkLogicalWindowAggregate( ts.getCluster, logicalTraits, + util.List.of(), new FlinkLogicalCalc(ts.getCluster, flinkLogicalTraits, flinkLogicalTs, program), ImmutableBitSet.of(0, 1), aggCallOfWindowAgg, @@ -2004,6 +2006,7 @@ class FlinkRelMdHandlerTestBase { val logicalWindowAgg = new LogicalWindowAggregate( ts.getCluster, ts.getTraitSet, + util.List.of(), project, ImmutableBitSet.of(1), aggCallOfWindowAgg, @@ -2015,6 +2018,7 @@ class FlinkRelMdHandlerTestBase { val flinkLogicalWindowAgg = new FlinkLogicalWindowAggregate( ts.getCluster, logicalTraits, + util.List.of(), new FlinkLogicalCalc(ts.getCluster, flinkLogicalTraits, flinkLogicalTs, program), ImmutableBitSet.of(1), aggCallOfWindowAgg, @@ -2187,6 +2191,7 @@ class FlinkRelMdHandlerTestBase { val logicalWindowAggWithAuxGroup = new LogicalWindowAggregate( ts.getCluster, ts.getTraitSet, + util.List.of(), project, ImmutableBitSet.of(0), aggCallOfWindowAgg, @@ -2198,6 +2203,7 @@ class FlinkRelMdHandlerTestBase { val flinkLogicalWindowAggWithAuxGroup = new FlinkLogicalWindowAggregate( ts.getCluster, logicalTraits, + util.List.of(), new FlinkLogicalCalc(ts.getCluster, flinkLogicalTraits, flinkLogicalTs, program), ImmutableBitSet.of(0), aggCallOfWindowAgg, diff --git a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/metadata/FlinkRelMdRowCountTest.scala b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/metadata/FlinkRelMdRowCountTest.scala index ad44ac902db6fc..8985ae6e69ca7b 100644 --- a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/metadata/FlinkRelMdRowCountTest.scala +++ b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/metadata/FlinkRelMdRowCountTest.scala @@ -28,6 +28,8 @@ import org.apache.calcite.util.ImmutableBitSet import org.junit.jupiter.api.Assertions._ import org.junit.jupiter.api.Test +import java.util + class FlinkRelMdRowCountTest extends FlinkRelMdHandlerTestBase { @Test @@ -183,6 +185,7 @@ class FlinkRelMdRowCountTest extends FlinkRelMdHandlerTestBase { val windowAgg = new LogicalWindowAggregate( ts.getCluster, ts.getTraitSet, + util.List.of(), ts, ImmutableBitSet.of(0, 1), aggCallOfWindowAgg,