diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/batch/BatchExecRank.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/batch/BatchExecRank.java index f5a958ec052679..95b78098293576 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/batch/BatchExecRank.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/batch/BatchExecRank.java @@ -35,13 +35,17 @@ import org.apache.flink.table.planner.plan.nodes.exec.SingleTransformationTranslator; import org.apache.flink.table.planner.plan.nodes.exec.utils.ExecNodeUtil; import org.apache.flink.table.planner.plan.utils.SortUtil; +import org.apache.flink.table.runtime.operators.rank.RankType; import org.apache.flink.table.runtime.operators.sort.RankOperator; import org.apache.flink.table.runtime.typeutils.InternalTypeInfo; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.annotation.JsonCreator; +import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.annotation.JsonInclude; import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.annotation.JsonProperty; +import javax.annotation.Nullable; + import java.util.Collections; import java.util.List; @@ -66,6 +70,7 @@ public class BatchExecRank extends ExecNodeBase public static final String FIELD_NAME_RANK_START = "rankStart"; public static final String FIELD_NAME_RANK_END = "rankEnd"; public static final String FIELD_NAME_OUTPUT_RANK_NUMBER = "outputRowNumber"; + public static final String FIELD_NAME_RANK_TYPE = "rankType"; @JsonProperty(FIELD_NAME_PARTITION_FIELDS) private final int[] partitionFields; @@ -82,6 +87,10 @@ public class BatchExecRank extends ExecNodeBase @JsonProperty(FIELD_NAME_OUTPUT_RANK_NUMBER) private final boolean outputRankNumber; + @JsonProperty(FIELD_NAME_RANK_TYPE) + @JsonInclude(JsonInclude.Include.NON_NULL) + private final RankType rankType; + public BatchExecRank( ReadableConfig tableConfig, int[] partitionFields, @@ -89,6 +98,7 @@ public BatchExecRank( long rankStart, long rankEnd, boolean outputRankNumber, + RankType rankType, InputProperty inputProperty, RowType outputType, String description) { @@ -104,6 +114,7 @@ public BatchExecRank( this.rankStart = rankStart; this.rankEnd = rankEnd; this.outputRankNumber = outputRankNumber; + this.rankType = rankType == null ? RankType.RANK : rankType; } @JsonCreator @@ -116,6 +127,7 @@ public BatchExecRank( @JsonProperty(FIELD_NAME_RANK_START) long rankStart, @JsonProperty(FIELD_NAME_RANK_END) long rankEnd, @JsonProperty(FIELD_NAME_OUTPUT_RANK_NUMBER) boolean outputRankNumber, + @Nullable @JsonProperty(FIELD_NAME_RANK_TYPE) RankType rankType, @JsonProperty(FIELD_NAME_INPUT_PROPERTIES) List inputProperties, @JsonProperty(FIELD_NAME_OUTPUT_TYPE) RowType outputType, @JsonProperty(FIELD_NAME_DESCRIPTION) String description) { @@ -125,6 +137,7 @@ public BatchExecRank( this.rankStart = rankStart; this.rankEnd = rankEnd; this.outputRankNumber = outputRankNumber; + this.rankType = rankType == null ? RankType.RANK : rankType; } @SuppressWarnings("unchecked") @@ -154,6 +167,7 @@ protected Transformation translateToPlanInternal( "OrderByComparator", inputType, SortUtil.getAscendingSortSpec(sortFields)), + rankType, rankStart, rankEnd, outputRankNumber); diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/physical/batch/BatchPhysicalRank.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/physical/batch/BatchPhysicalRank.scala index c4442405f4f49b..0ab12c01020196 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/physical/batch/BatchPhysicalRank.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/physical/batch/BatchPhysicalRank.scala @@ -68,7 +68,9 @@ class BatchPhysicalRank( outputRankNumber) with BatchPhysicalRel { - require(rankType == RankType.RANK, "Only RANK is supported now") + require( + rankType == RankType.RANK || rankType == RankType.ROW_NUMBER, + "Only RANK and ROW_NUMBER are supported now") val (rankStart, rankEnd) = rankRange match { case r: ConstantRankRange => (r.getRankStart, r.getRankEnd) case o => throw new TableException(s"$o is not supported now") @@ -240,6 +242,7 @@ class BatchPhysicalRank( rankStart, rankEnd, outputRankNumber, + rankType, InputProperty.builder().requiredDistribution(requiredDistribution).build(), FlinkTypeFactory.toLogicalRowType(getRowType), getRelDetailedDescription) diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRule.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRule.scala index 8cc198f7c65cc3..ac5a80893f504b 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRule.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRule.scala @@ -231,8 +231,8 @@ class FlinkLogicalRankRuleForConstantRange extends FlinkLogicalRankRuleBase { } val agg = group.aggCalls.get(0) - if (agg.getOperator.kind != SqlKind.RANK) { - // only accept RANK function + if (agg.getOperator.kind != SqlKind.RANK && agg.getOperator.kind != SqlKind.ROW_NUMBER) { + // only accept RANK and ROW_NUMBER functions return false } diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/physical/batch/BatchPhysicalRankRule.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/physical/batch/BatchPhysicalRankRule.scala index 57a557ea0153d8..923455e9740a53 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/physical/batch/BatchPhysicalRankRule.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/physical/batch/BatchPhysicalRankRule.scala @@ -46,8 +46,9 @@ class BatchPhysicalRankRule(config: Config) extends ConverterRule(config) { override def matches(call: RelOptRuleCall): Boolean = { val rank: FlinkLogicalRank = call.rel(0) - // Only support rank() now - rank.rankType == RankType.RANK && rank.rankRange.isInstanceOf[ConstantRankRange] + // Support RANK and ROW_NUMBER with a constant rank range + (rank.rankType == RankType.RANK || rank.rankType == RankType.ROW_NUMBER) && + rank.rankRange.isInstanceOf[ConstantRankRange] } def convert(rel: RelNode): RelNode = { diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/batch/RankBatchRestoreTest.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/batch/RankBatchRestoreTest.java index 571a936fd1de47..6e9b424c0c02ee 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/batch/RankBatchRestoreTest.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/batch/RankBatchRestoreTest.java @@ -43,6 +43,7 @@ public List programs() { // RankTestPrograms.RANK_TEST_RETRACT_STRATEGY, RankTestPrograms.RANK_TEST_UPDATE_FAST_STRATEGY, RankTestPrograms.RANK_N_TEST, - RankTestPrograms.RANK_2_TEST); + RankTestPrograms.RANK_2_TEST, + RankTestPrograms.ROW_NUMBER_TOP_N); } } diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/common/RankTestPrograms.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/common/RankTestPrograms.java index bc91f70af8d90e..13d9561b4a71fd 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/common/RankTestPrograms.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/common/RankTestPrograms.java @@ -208,6 +208,34 @@ private static TableTestProgram getTableTestProgram( + " where c <= 2") .build(); + public static final TableTestProgram ROW_NUMBER_TOP_N = + TableTestProgram.of("row-number-top-n", "validates batch ROW_NUMBER top-n rank") + .setupTableSource( + SourceTestStep.newBuilder("MyTable") + .addSchema( + "a INT", "b VARCHAR", "c INT primary key not enforced") + .addOption(CHANGELOG_MODE, "I") + .producedBeforeRestore( + Row.of(2, "a", 6), + Row.of(4, "b", 8), + Row.of(6, "c", 10), + Row.of(1, "a", 5), + Row.of(3, "b", 7), + Row.of(5, "c", 9)) + .producedAfterRestore(Row.of(4, "d", 7), Row.of(3, "e", 8)) + .build()) + .setupTableSink( + SinkTestStep.newBuilder("sink_t") + .addSchema("a INT", "b VARCHAR") + .consumedBeforeRestore("+I[2, a]", "+I[4, b]", "+I[6, c]") + .consumedAfterRestore("+I[4, d]", "+I[3, e]") + .build()) + .runSql( + "INSERT INTO sink_t SELECT a, b FROM (" + + " SELECT a, b, ROW_NUMBER() OVER (PARTITION BY b ORDER BY a DESC) rn" + + " FROM MyTable) t WHERE rn = 1") + .build(); + public static final TableTestProgram RANK_2_TEST = TableTestProgram.of("rank-2-test", "validates rank node can handle multiple outputs") .setupTableSource( diff --git a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/batch/sql/RankTest.xml b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/batch/sql/RankTest.xml index 384214cb8fd0da..59f11ce209d564 100644 --- a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/batch/sql/RankTest.xml +++ b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/batch/sql/RankTest.xml @@ -34,14 +34,16 @@ LogicalSink(table=[default_catalog.default_database.sink], fields=[name, eat, cn @@ -255,17 +257,22 @@ LogicalProject(rn1=[CAST($3):INTEGER NOT NULL], rn2=[CAST($4):INTEGER NOT NULL]) @@ -345,6 +352,35 @@ Calc(select=[CONCAT('http://txmov2.a.yximgs.com', uri) AS url, reqcount AS downl +- Exchange(distribution=[forward]) +- Sort(orderBy=[start_time ASC, bucket_id ASC, reqcount DESC]) +- BoundedStreamScan(table=[[default_catalog, default_database, MyTable1]], fields=[uri, reqcount, start_time, bucket_id]) +]]> + + + + + + + + + + + diff --git a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRuleForConstantRangeTest.xml b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRuleForConstantRangeTest.xml index da506419030463..940071d568a118 100644 --- a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRuleForConstantRangeTest.xml +++ b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRuleForConstantRangeTest.xml @@ -331,10 +331,9 @@ LogicalProject(a=[$0], b=[$1], rn=[$2]) diff --git a/flink-table/flink-table-planner/src/test/resources/restore-tests/batch-exec-rank_1/row-number-top-n/plan/row-number-top-n.json b/flink-table/flink-table-planner/src/test/resources/restore-tests/batch-exec-rank_1/row-number-top-n/plan/row-number-top-n.json new file mode 100644 index 00000000000000..8c308a9f4454f8 --- /dev/null +++ b/flink-table/flink-table-planner/src/test/resources/restore-tests/batch-exec-rank_1/row-number-top-n/plan/row-number-top-n.json @@ -0,0 +1,278 @@ +{ + "flinkVersion" : "2.4", + "nodes" : [ { + "id" : 1, + "type" : "batch-exec-table-source-scan_1", + "scanTableSource" : { + "table" : { + "identifier" : "`default_catalog`.`default_database`.`MyTable`", + "resolvedTable" : { + "schema" : { + "columns" : [ { + "name" : "a", + "dataType" : "INT" + }, { + "name" : "b", + "dataType" : "VARCHAR(2147483647)" + }, { + "name" : "c", + "dataType" : "INT NOT NULL" + } ], + "primaryKey" : { + "name" : "PK_c", + "type" : "PRIMARY_KEY", + "columns" : [ "c" ] + } + } + } + }, + "abilities" : [ { + "type" : "ProjectPushDown", + "projectedFields" : [ [ 0 ], [ 1 ] ], + "producedType" : "ROW<`a` INT, `b` VARCHAR(2147483647)> NOT NULL" + }, { + "type" : "ReadingMetadata", + "metadataKeys" : [ ], + "producedType" : "ROW<`a` INT, `b` VARCHAR(2147483647)> NOT NULL" + } ] + }, + "outputType" : "ROW<`a` INT, `b` VARCHAR(2147483647)>", + "description" : "TableSourceScan(table=[[default_catalog, default_database, MyTable, project=[a, b], metadata=[]]], fields=[a, b])", + "dynamicFilteringDataListenerID" : "4c016fba-ee43-4867-a180-603c7dca303d" + }, { + "id" : 2, + "type" : "batch-exec-sort_1", + "configuration" : { + "table.exec.resource.sort.memory" : "128 mb", + "table.exec.sort.async-merge-enabled" : "true", + "table.exec.sort.max-num-file-handles" : "128", + "table.exec.spill-compression.block-size" : "64 kb", + "table.exec.spill-compression.enabled" : "true" + }, + "sortSpec" : { + "fields" : [ { + "index" : 1, + "isAscending" : true, + "nullIsLast" : false + }, { + "index" : 0, + "isAscending" : false, + "nullIsLast" : true + } ] + }, + "inputProperties" : [ { + "requiredDistribution" : { + "type" : "UNKNOWN" + }, + "damBehavior" : "END_INPUT", + "priority" : 0 + } ], + "outputType" : "ROW<`a` INT, `b` VARCHAR(2147483647)>", + "description" : "Sort(orderBy=[b ASC, a DESC])" + }, { + "id" : 8, + "type" : "batch-exec-exchange_1", + "inputProperties" : [ { + "requiredDistribution" : { + "type" : "KEEP_INPUT_AS_IS", + "inputDistribution" : { + "type" : "UNKNOWN" + }, + "isStrict" : true + }, + "damBehavior" : "PIPELINED", + "priority" : 0 + } ], + "outputType" : "ROW<`a` INT, `b` VARCHAR(2147483647)>", + "description" : "Exchange(distribution=[forward])", + "requiredExchangeMode" : "UNDEFINED" + }, { + "id" : 3, + "type" : "batch-exec-rank_1", + "partitionFields" : [ 1 ], + "sortFields" : [ 0 ], + "rankStart" : 1, + "rankEnd" : 1, + "outputRowNumber" : false, + "rankType" : "ROW_NUMBER", + "inputProperties" : [ { + "requiredDistribution" : { + "type" : "UNKNOWN" + }, + "damBehavior" : "PIPELINED", + "priority" : 0 + } ], + "outputType" : "ROW<`a` INT, `b` VARCHAR(2147483647)>", + "description" : "Rank(rankType=[ROW_NUMBER], rankRange=[rankStart=1, rankEnd=1], partitionBy=[b], orderBy=[a DESC], global=[false], select=[a, b])" + }, { + "id" : 4, + "type" : "batch-exec-exchange_1", + "inputProperties" : [ { + "requiredDistribution" : { + "type" : "HASH", + "keys" : [ 1 ] + }, + "damBehavior" : "BLOCKING", + "priority" : 0 + } ], + "outputType" : "ROW<`a` INT, `b` VARCHAR(2147483647)>", + "description" : "Exchange(distribution=[hash[b]])", + "requiredExchangeMode" : "UNDEFINED" + }, { + "id" : 5, + "type" : "batch-exec-sort_1", + "configuration" : { + "table.exec.resource.sort.memory" : "128 mb", + "table.exec.sort.async-merge-enabled" : "true", + "table.exec.sort.max-num-file-handles" : "128", + "table.exec.spill-compression.block-size" : "64 kb", + "table.exec.spill-compression.enabled" : "true" + }, + "sortSpec" : { + "fields" : [ { + "index" : 1, + "isAscending" : true, + "nullIsLast" : false + }, { + "index" : 0, + "isAscending" : false, + "nullIsLast" : true + } ] + }, + "inputProperties" : [ { + "requiredDistribution" : { + "type" : "UNKNOWN" + }, + "damBehavior" : "END_INPUT", + "priority" : 0 + } ], + "outputType" : "ROW<`a` INT, `b` VARCHAR(2147483647)>", + "description" : "Sort(orderBy=[b ASC, a DESC])" + }, { + "id" : 9, + "type" : "batch-exec-exchange_1", + "inputProperties" : [ { + "requiredDistribution" : { + "type" : "KEEP_INPUT_AS_IS", + "inputDistribution" : { + "type" : "HASH", + "keys" : [ 1 ] + }, + "isStrict" : true + }, + "damBehavior" : "PIPELINED", + "priority" : 0 + } ], + "outputType" : "ROW<`a` INT, `b` VARCHAR(2147483647)>", + "description" : "Exchange(distribution=[forward])", + "requiredExchangeMode" : "UNDEFINED" + }, { + "id" : 6, + "type" : "batch-exec-rank_1", + "partitionFields" : [ 1 ], + "sortFields" : [ 0 ], + "rankStart" : 1, + "rankEnd" : 1, + "outputRowNumber" : false, + "rankType" : "ROW_NUMBER", + "inputProperties" : [ { + "requiredDistribution" : { + "type" : "HASH", + "keys" : [ 1 ] + }, + "damBehavior" : "PIPELINED", + "priority" : 0 + } ], + "outputType" : "ROW<`a` INT, `b` VARCHAR(2147483647)>", + "description" : "Rank(rankType=[ROW_NUMBER], rankRange=[rankStart=1, rankEnd=1], partitionBy=[b], orderBy=[a DESC], global=[true], select=[a, b])" + }, { + "id" : 7, + "type" : "batch-exec-sink_1", + "configuration" : { + "table.exec.sink.not-null-enforcer" : "ERROR", + "table.exec.sink.type-length-enforcer" : "IGNORE" + }, + "dynamicTableSink" : { + "table" : { + "identifier" : "`default_catalog`.`default_database`.`sink_t`", + "resolvedTable" : { + "schema" : { + "columns" : [ { + "name" : "a", + "dataType" : "INT" + }, { + "name" : "b", + "dataType" : "VARCHAR(2147483647)" + } ] + } + } + } + }, + "inputProperties" : [ { + "requiredDistribution" : { + "type" : "UNKNOWN" + }, + "damBehavior" : "BLOCKING", + "priority" : 0 + } ], + "outputType" : "ROW<`a` INT, `b` VARCHAR(2147483647)>", + "description" : "Sink(table=[default_catalog.default_database.sink_t], fields=[a, b])" + } ], + "edges" : [ { + "source" : 1, + "target" : 2, + "shuffle" : { + "type" : "FORWARD" + }, + "shuffleMode" : "PIPELINED" + }, { + "source" : 2, + "target" : 8, + "shuffle" : { + "type" : "FORWARD" + }, + "shuffleMode" : "PIPELINED" + }, { + "source" : 8, + "target" : 3, + "shuffle" : { + "type" : "FORWARD" + }, + "shuffleMode" : "PIPELINED" + }, { + "source" : 3, + "target" : 4, + "shuffle" : { + "type" : "FORWARD" + }, + "shuffleMode" : "PIPELINED" + }, { + "source" : 4, + "target" : 5, + "shuffle" : { + "type" : "FORWARD" + }, + "shuffleMode" : "PIPELINED" + }, { + "source" : 5, + "target" : 9, + "shuffle" : { + "type" : "FORWARD" + }, + "shuffleMode" : "PIPELINED" + }, { + "source" : 9, + "target" : 6, + "shuffle" : { + "type" : "FORWARD" + }, + "shuffleMode" : "PIPELINED" + }, { + "source" : 6, + "target" : 7, + "shuffle" : { + "type" : "FORWARD" + }, + "shuffleMode" : "PIPELINED" + } ] +} \ No newline at end of file diff --git a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/batch/sql/RankTest.scala b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/batch/sql/RankTest.scala index 06d7c7688a8fb6..dbef10bba846ce 100644 --- a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/batch/sql/RankTest.scala +++ b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/batch/sql/RankTest.scala @@ -71,6 +71,17 @@ class RankTest extends TableTestBase { .isThrownBy(() => util.verifyExecPlan(sqlQuery)) } + @Test + def testRowNumberTopNConvertedToRank(): Unit = { + val sqlQuery = + """ + |SELECT a, b, c FROM ( + | SELECT a, b, c, ROW_NUMBER() OVER (PARTITION BY b ORDER BY c DESC) rn FROM MyTable) t + |WHERE rn = 1 + """.stripMargin + util.verifyExecPlan(sqlQuery) + } + @Test def testRankWithoutOrderBy(): Unit = { val sqlQuery = diff --git a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRuleForConstantRangeTest.scala b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRuleForConstantRangeTest.scala index f33d040457f36c..762326a42253fa 100644 --- a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRuleForConstantRangeTest.scala +++ b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/FlinkLogicalRankRuleForConstantRangeTest.scala @@ -32,7 +32,7 @@ class FlinkLogicalRankRuleForConstantRangeTest extends TableTestBase { @Test def testRowNumberFunc(): Unit = { - // can not be converted to Rank + // ROW_NUMBER Top-N with a constant rank range is converted to Rank val sqlQuery = """ |SELECT * FROM ( diff --git a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/runtime/batch/sql/RowNumberITCase.scala b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/runtime/batch/sql/RowNumberITCase.scala new file mode 100644 index 00000000000000..7f69fcf28c493e --- /dev/null +++ b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/runtime/batch/sql/RowNumberITCase.scala @@ -0,0 +1,67 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.table.planner.runtime.batch.sql + +import org.apache.flink.api.common.typeinfo.BasicTypeInfo.{INT_TYPE_INFO, LONG_TYPE_INFO, STRING_TYPE_INFO} +import org.apache.flink.api.java.typeutils.RowTypeInfo +import org.apache.flink.table.planner.runtime.utils.BatchTestBase +import org.apache.flink.table.planner.runtime.utils.BatchTestBase.row + +import org.junit.jupiter.api.{BeforeEach, Test} + +/** Correctness for batch `ROW_NUMBER() OVER (PARTITION BY ... ORDER BY ...) WHERE rn <= N`. */ +class RowNumberITCase extends BatchTestBase { + + @BeforeEach + override def before(): Unit = { + super.before() + val data = + Seq(row(1, "a", 10L), row(2, "a", 30L), row(3, "a", 30L), row(4, "b", 5L), row(5, "b", 7L)) + val tType = new RowTypeInfo(INT_TYPE_INFO, STRING_TYPE_INFO, LONG_TYPE_INFO) + registerCollection("T", data, tType, "id, grp, v") + } + + @Test + def testKeepFirstPerGroup(): Unit = { + checkResult( + "SELECT grp, v FROM (" + + "SELECT grp, v, ROW_NUMBER() OVER (PARTITION BY grp ORDER BY v ASC) rn FROM T) t " + + "WHERE rn = 1", + Seq(row("a", 10L), row("b", 5L)) + ) + } + + @Test + def testKeepLastPerGroup(): Unit = { + checkResult( + "SELECT grp FROM (" + + "SELECT grp, v, ROW_NUMBER() OVER (PARTITION BY grp ORDER BY v DESC) rn FROM T) t " + + "WHERE rn = 1", + Seq(row("a"), row("b"))) + } + + @Test + def testTopNKeepsExactlyNAcrossTies(): Unit = { + checkResult( + "SELECT grp, v FROM (" + + "SELECT grp, v, ROW_NUMBER() OVER (PARTITION BY grp ORDER BY v ASC) rn FROM T) t " + + "WHERE rn <= 2", + Seq(row("a", 10L), row("a", 30L), row("b", 5L), row("b", 7L)) + ) + } +} diff --git a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/sort/RankOperator.java b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/sort/RankOperator.java index 91fe0fa8bbd142..fb65af11eb0973 100644 --- a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/sort/RankOperator.java +++ b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/sort/RankOperator.java @@ -26,6 +26,7 @@ import org.apache.flink.table.runtime.generated.GeneratedRecordComparator; import org.apache.flink.table.runtime.generated.RecordComparator; import org.apache.flink.table.runtime.operators.TableStreamOperator; +import org.apache.flink.table.runtime.operators.rank.RankType; import org.apache.flink.table.runtime.typeutils.AbstractRowDataSerializer; import org.apache.flink.table.runtime.util.StreamRecordCollector; @@ -35,6 +36,7 @@ public class RankOperator extends TableStreamOperator private GeneratedRecordComparator partitionByGenComp; private GeneratedRecordComparator orderByGenComp; + private final RankType rankType; private final long rankStart; private final long rankEnd; private final boolean outputRankFunColumn; @@ -52,11 +54,13 @@ public class RankOperator extends TableStreamOperator public RankOperator( GeneratedRecordComparator partitionByGenComp, GeneratedRecordComparator orderByGenComp, + RankType rankType, long rankStart, long rankEnd, boolean outputRankFunColumn) { this.partitionByGenComp = partitionByGenComp; this.orderByGenComp = orderByGenComp; + this.rankType = rankType; this.rankStart = rankStart; this.rankEnd = rankEnd; this.outputRankFunColumn = outputRankFunColumn; @@ -102,9 +106,10 @@ public void processElement(StreamRecord element) throws Exception { } private void emitInternal(RowData element) { - if (rank >= rankStart && rank <= rankEnd) { + long rankValue = (rankType == RankType.ROW_NUMBER) ? rowNum : rank; + if (rankValue >= rankStart && rankValue <= rankEnd) { if (outputRankFunColumn) { - rankValueRow.setField(0, rank); + rankValueRow.setField(0, rankValue); collector.collect(joinedRow.replace(element, rankValueRow)); } else { collector.collect(element);