This is an automated email from the ASF dual-hosted git repository. jark pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit 3e1018e0cbc6d28d59f2d0658c6a94b044484f67 Author: luoyuxia <[email protected]> AuthorDate: Tue Jun 28 20:58:06 2022 +0800 [FLINK-27619][sql] Support NTILE Function --- docs/data/sql_functions.yml | 5 + docs/data/sql_functions_zh.yml | 5 + .../planner/expressions/ExpressionBuilder.java | 5 + .../functions/aggfunctions/NTILEAggFunction.java | 112 +++++++++++++++++++++ .../functions/sql/FlinkSqlOperatorTable.java | 1 + .../exec/batch/BatchExecOverAggregateBase.java | 5 +- .../planner/plan/utils/AggFunctionFactory.scala | 11 ++ .../runtime/batch/sql/OverAggregateITCase.scala | 48 +++++++++ .../operators/over/frame/InsensitiveOverFrame.java | 2 +- 9 files changed, 192 insertions(+), 2 deletions(-) diff --git a/docs/data/sql_functions.yml b/docs/data/sql_functions.yml index 646beaae6f2..2a7ba09c844 100644 --- a/docs/data/sql_functions.yml +++ b/docs/data/sql_functions.yml @@ -1009,6 +1009,11 @@ aggregate: description: Return the cumulative distribution of a value in a group of values. The result is the number of rows preceding or equal to the current row in the ordering of the partition divided by the number of rows in the window partition. - sql: PERCENT_RANK() description: Return the percentage ranking of a value in a group of values. The result is the rank value minus one, divided by the number of rows in the parition minus one. If the partition only contains one row, the function will return 0. + - sql: NTILE(n) + description: | + Divides the rows for each window partition into `n` buckets ranging from 1 to at most `n`. + If the number of rows in the window partition doesn't divide evenly into the number of buckets, then the remainder values are distributed one per bucket, starting with the first bucket. + For example, with 6 rows and 4 buckets, the bucket values would be as follows: 1 1 2 2 3 4 catalog: - sql: CURRENT_DATABASE() diff --git a/docs/data/sql_functions_zh.yml b/docs/data/sql_functions_zh.yml index e43efd90282..ea47e7e0fe0 100644 --- a/docs/data/sql_functions_zh.yml +++ b/docs/data/sql_functions_zh.yml @@ -1100,6 +1100,11 @@ aggregate: description: 返回值在一组值的累积分布。结果是小于或等于当前行的值的行数除以窗口分区的总行数。 - sql: PERCENT_RANK() description: 返回值在一组值的百分比排名。结果是当前行在窗口分区中的排名减 1，然后除以窗口分区的总行数减 1。如果窗口分区的总行数为 1，则该函数返回 0。 + - sql: NTILE(n) + description: | + 将窗口分区中的所有数据按照顺序划分为 n 个分组，返回分配给各行数据的分组编号（从 1 开始）。 + 如果不能均匀划分为 n 个分组，则从第 1 个分组开始，为每一分组分配一个剩余值。 + 比如某个窗口分区有 6 行数据，划分为 4 个分组，则各行的分组编号为：1，1，2，2，3，4。 catalog: - sql: CURRENT_DATABASE() diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/expressions/ExpressionBuilder.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/expressions/ExpressionBuilder.java index 867ce097d66..213b8b51cef 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/expressions/ExpressionBuilder.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/expressions/ExpressionBuilder.java @@ -40,6 +40,7 @@ import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.GREATE import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.IF; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.IS_NULL; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.LESS_THAN; +import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.LESS_THAN_OR_EQUAL; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.MINUS; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.MOD; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.NOT; @@ -143,6 +144,10 @@ public class ExpressionBuilder { return call(EQUALS, input1, input2); } + public static UnresolvedCallExpression lessThanOrEqual(Expression input1, Expression input2) { + return call(LESS_THAN_OR_EQUAL, input1, input2); + } + public static UnresolvedCallExpression lessThan(Expression input1, Expression input2) { return call(LESS_THAN, input1, input2); } diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/functions/aggfunctions/NTILEAggFunction.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/functions/aggfunctions/NTILEAggFunction.java new file mode 100644 index 00000000000..e133c0fc675 --- /dev/null +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/functions/aggfunctions/NTILEAggFunction.java @@ -0,0 +1,112 @@ +/* + * 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.functions.aggfunctions; + +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.api.TableException; +import org.apache.flink.table.expressions.Expression; +import org.apache.flink.table.expressions.UnresolvedReferenceExpression; +import org.apache.flink.table.types.DataType; + +import static org.apache.flink.table.expressions.ApiExpressionUtils.unresolvedRef; +import static org.apache.flink.table.expressions.ApiExpressionUtils.valueLiteral; +import static org.apache.flink.table.planner.expressions.ExpressionBuilder.div; +import static org.apache.flink.table.planner.expressions.ExpressionBuilder.greaterThan; +import static org.apache.flink.table.planner.expressions.ExpressionBuilder.ifThenElse; +import static org.apache.flink.table.planner.expressions.ExpressionBuilder.lessThanOrEqual; +import static org.apache.flink.table.planner.expressions.ExpressionBuilder.literal; +import static org.apache.flink.table.planner.expressions.ExpressionBuilder.mod; +import static org.apache.flink.table.planner.expressions.ExpressionBuilder.plus; + +/** built-in NTILE agg function. */ +public class NTILEAggFunction extends DeclarativeAggregateFunction + implements SizeBasedWindowFunction { + + private final Expression bucketSize; + private final Expression bucketsWithPadding; + private final UnresolvedReferenceExpression rowNumber = unresolvedRef("rowNumber"); + private final UnresolvedReferenceExpression bucket = unresolvedRef("bucket"); + private final UnresolvedReferenceExpression bucketThreshold = unresolvedRef("bucketThreshold"); + + public NTILEAggFunction() { + this.bucketSize = div(windowSizeAttribute(), operand(0)); + this.bucketsWithPadding = mod(windowSizeAttribute(), operand(0)); + } + + @Override + public int operandCount() { + return 1; + } + + @Override + public UnresolvedReferenceExpression[] aggBufferAttributes() { + return new UnresolvedReferenceExpression[] {rowNumber, bucket, bucketThreshold}; + } + + @Override + public DataType[] getAggBufferTypes() { + return new DataType[] {DataTypes.BIGINT(), DataTypes.BIGINT(), DataTypes.BIGINT()}; + } + + @Override + public DataType getResultType() { + return DataTypes.BIGINT(); + } + + @Override + public Expression[] initialValuesExpressions() { + return new Expression[] {valueLiteral(0L), valueLiteral(0L), valueLiteral(0L)}; + } + + @Override + public Expression[] accumulateExpressions() { + return new Expression[] { + /* rowNumber = */ plus(rowNumber, literal(1L)), + /* bucket = */ plus(bucket, bucketOverflowThen(literal(1L))), + /* bucketThreshold = */ plus( + bucketThreshold, + bucketOverflowThen( + plus( + bucketSize, + ifThenElse( + lessThanOrEqual(bucket, bucketsWithPadding), + valueLiteral(1), + valueLiteral(0))))) + }; + } + + private Expression bucketOverflowThen(Expression e) { + return ifThenElse(greaterThan(rowNumber, bucketThreshold), e, valueLiteral(0L)); + } + + @Override + public Expression[] retractExpressions() { + throw new TableException("This function does not support retraction."); + } + + @Override + public Expression[] mergeExpressions() { + throw new TableException("This function does not support merge."); + } + + @Override + public Expression getValueExpression() { + return bucket; + } +} diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/functions/sql/FlinkSqlOperatorTable.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/functions/sql/FlinkSqlOperatorTable.java index 0161b3c2140..cfdf34e2e79 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/functions/sql/FlinkSqlOperatorTable.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/functions/sql/FlinkSqlOperatorTable.java @@ -1140,6 +1140,7 @@ public class FlinkSqlOperatorTable extends ReflectiveSqlOperatorTable { public static final SqlAggFunction ROW_NUMBER = SqlStdOperatorTable.ROW_NUMBER; public static final SqlAggFunction CUME_DIST = SqlStdOperatorTable.CUME_DIST; public static final SqlAggFunction PERCENT_RANK = SqlStdOperatorTable.PERCENT_RANK; + public static final SqlAggFunction NTILE = SqlStdOperatorTable.NTILE; public static final SqlAggFunction LEAD = SqlStdOperatorTable.LEAD; public static final SqlAggFunction LAG = SqlStdOperatorTable.LAG; diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/batch/BatchExecOverAggregateBase.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/batch/BatchExecOverAggregateBase.java index 437ab6a80d9..b639980610f 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/batch/BatchExecOverAggregateBase.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/batch/BatchExecOverAggregateBase.java @@ -115,6 +115,9 @@ public abstract class BatchExecOverAggregateBase extends ExecNodeBase<RowData> /** Infer the over window mode based on given group info. */ protected OverWindowMode inferGroupMode(GroupSpec group) { AggregateCall aggCall = group.getAggCalls().get(0); + if (aggCall.getAggregation().getKind() == NTILE) { + return OverWindowMode.INSENSITIVE; + } if (aggCall.getAggregation().allowsFraming()) { if (group.isRows()) { return OverWindowMode.ROW; @@ -147,7 +150,7 @@ public abstract class BatchExecOverAggregateBase extends ExecNodeBase<RowData> OFFSET, /** * The INSENSITIVE mode does not care the window framing without LEAD/LAG agg function, for - * example RANK/DENSE_RANK/PERCENT_RANK/ROW_NUMBER. + * example RANK/DENSE_RANK/PERCENT_RANK/ROW_NUMBER/NTILE. */ INSENSITIVE } diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/utils/AggFunctionFactory.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/utils/AggFunctionFactory.scala index 1def38cec15..a9599ce5d0b 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/utils/AggFunctionFactory.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/utils/AggFunctionFactory.scala @@ -112,6 +112,13 @@ class AggFunctionFactory( throw new TableException("PERCENT_RANK Function is not supported in stream mode.") } + case _: SqlNtileAggFunction => + if (isBounded) { + createNTILEAggFUnction(argTypes) + } else { + throw new TableException("NTILE Function is not supported in stream mode.") + } + case func: SqlLeadLagAggFunction => if (isBounded) { createBatchLeadLagAggFunction(argTypes, index) @@ -535,6 +542,10 @@ class AggFunctionFactory( new PercentRankAggFunction(argTypes) } + private def createNTILEAggFUnction(argTypes: Array[LogicalType]): UserDefinedFunction = { + new NTILEAggFunction + } + private def createFirstValueAggFunction( argTypes: Array[LogicalType], index: Int): UserDefinedFunction = { diff --git a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/runtime/batch/sql/OverAggregateITCase.scala b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/runtime/batch/sql/OverAggregateITCase.scala index 748f0a82379..979fa83e8e3 100644 --- a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/runtime/batch/sql/OverAggregateITCase.scala +++ b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/runtime/batch/sql/OverAggregateITCase.scala @@ -2855,6 +2855,54 @@ class OverAggregateITCase extends BatchTestBase { ) ) } + + @Test + def testNTILE(): Unit = { + checkResult( + "SELECT f, NTILE(4) over (order by e)," + + " NTILE(3) over (partition by d order by e desc) FROM Table5", + Seq( + row(0, 1, 1), + row(1, 1, 2), + row(2, 1, 1), + row(3, 1, 3), + row(4, 2, 2), + row(5, 2, 1), + row(6, 2, 3), + row(7, 2, 2), + row(8, 3, 1), + row(9, 3, 1), + row(10, 3, 3), + row(11, 3, 2), + row(12, 4, 2), + row(13, 4, 1), + row(14, 4, 1) + ) + ) + + // test values of order-key containing duplicates + checkResult( + "SELECT d, NTILE(4) over (order by d)" + + " FROM Table5", + Seq( + row(5, 3), + row(5, 3), + row(5, 4), + row(5, 4), + row(5, 4), + row(4, 2), + row(4, 2), + row(4, 3), + row(4, 3), + row(3, 1), + row(3, 2), + row(3, 2), + row(2, 1), + row(2, 1), + row(1, 1) + ) + ) + } } /** The initial accumulator for count aggregate function */ diff --git a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/over/frame/InsensitiveOverFrame.java b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/over/frame/InsensitiveOverFrame.java index 0f4383e9b78..2c8ba39ac14 100644 --- a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/over/frame/InsensitiveOverFrame.java +++ b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/over/frame/InsensitiveOverFrame.java @@ -27,7 +27,7 @@ import org.apache.flink.table.runtime.util.ResettableExternalBuffer; /** * The insensitive window frame calculates the statements which shouldn't care the window frame, for - * example RANK/DENSE_RANK/PERCENT_RANK/CUME_DIST/ROW_NUMBER. + * example RANK/DENSE_RANK/PERCENT_RANK/NTILE/ROW_NUMBER. */ public class InsensitiveOverFrame implements OverWindowFrame {
