| /* |
| * 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.codegen.agg |
| |
| import org.apache.calcite.tools.RelBuilder |
| import org.apache.flink.table.api.functions.DeclarativeAggregateFunction |
| import org.apache.flink.table.api.types.{DataType, DataTypes, InternalType} |
| import org.apache.flink.table.codegen.CodeGenUtils.primitiveTypeTermForType |
| import org.apache.flink.table.codegen.{CodeGeneratorContext, ExprCodeGenerator, GeneratedExpression} |
| import org.apache.flink.table.codegen.agg.AggsHandlerCodeGenerator._ |
| import org.apache.flink.table.expressions._ |
| import org.apache.flink.table.plan.util.AggregateInfo |
| |
| /** |
| * It is for code generate aggregation functions that are specified using expressions. |
| * The aggregate buffer is embedded inside of a larger shared aggregation buffer. |
| * |
| * @param ctx the code gen context |
| * @param aggInfo the aggregate information |
| * @param filterExpression filter argument access expression, none if no filter |
| * @param mergedAccOffset the mergedAcc may come from local aggregate, |
| * this is the first buffer offset in the row |
| * @param aggBufferOffset the offset in the buffers of this aggregate |
| * @param aggBufferSize the total size of aggregate buffers |
| * @param inputTypes the input field type infos |
| * @param constantExprs the constant expressions |
| * @param relBuilder the rel builder to translate expressions to calcite rex nodes |
| */ |
| class DeclarativeAggCodeGen( |
| ctx: CodeGeneratorContext, |
| aggInfo: AggregateInfo, |
| filterExpression: Option[Expression], |
| mergedAccOffset: Int, |
| aggBufferOffset: Int, |
| aggBufferSize: Int, |
| inputTypes: Seq[InternalType], |
| constantExprs: Seq[GeneratedExpression], |
| relBuilder: RelBuilder) |
| extends AggCodeGen { |
| |
| val function: DeclarativeAggregateFunction = |
| aggInfo.function.asInstanceOf[DeclarativeAggregateFunction] |
| |
| val bufferTypes: Array[DataType] = aggInfo.externalAccTypes |
| val bufferIndexes: Array[Int] = Array.range( |
| aggBufferOffset, aggBufferOffset + bufferTypes.length) |
| val bufferTerms: Array[String] = function.aggBufferAttributes |
| .map(a => s"agg${aggInfo.aggIndex}_${a.name}").toArray |
| val bufferNullTerms: Array[String] = bufferTerms.map(_ + "_isNull") |
| |
| val argIndexes: Array[Int] = aggInfo.argIndexes |
| val argTypes: Array[InternalType] = { |
| val types = inputTypes ++ constantExprs.map(_.resultType) |
| argIndexes.map(types(_)) |
| } |
| |
| def createAccumulator(generator: ExprCodeGenerator): Seq[GeneratedExpression] = { |
| function.initialValuesExpressions |
| .map(expr => generator.generateExpression(expr.toRexNode(relBuilder))) |
| } |
| |
| def setAccumulator(generator: ExprCodeGenerator): String = { |
| val aggBufferAccesses = function.aggBufferAttributes.zipWithIndex |
| .map { case (attr, index) => |
| ResolvedAggInputReference( |
| attr.name, bufferIndexes(index), bufferTypes(index).toInternalType) |
| } |
| .map(expr => generator.generateExpression(expr.toRexNode(relBuilder))) |
| |
| val setters = aggBufferAccesses.zipWithIndex.map { |
| case (access, index) => |
| val typeTerm = primitiveTypeTermForType(access.resultType) |
| val memberName = bufferTerms(index) |
| val memberNullTerm = bufferNullTerms(index) |
| ctx.addReusableMember(s"private $typeTerm $memberName;") |
| ctx.addReusableMember(s"private boolean $memberNullTerm;") |
| s""" |
| |${access.copyResultTermToTargetIfChanged(ctx, memberName)}; |
| |$memberNullTerm = ${access.nullTerm}; |
| """.stripMargin |
| } |
| |
| setters.mkString("\n") |
| } |
| |
| override def resetAccumulator(generator: ExprCodeGenerator): String = { |
| val initialExprs = function.initialValuesExpressions |
| .map(expr => generator.generateExpression(expr.toRexNode(relBuilder))) |
| val codes = initialExprs.zipWithIndex.map { |
| case (init, index) => |
| val memberName = bufferTerms(index) |
| val memberNullTerm = bufferNullTerms(index) |
| s""" |
| |${init.code} |
| |$memberName = ${init.resultTerm}; |
| |$memberNullTerm = ${init.nullTerm}; |
| """.stripMargin |
| } |
| codes.mkString("\n") |
| } |
| |
| def getAccumulator(generator: ExprCodeGenerator): Seq[GeneratedExpression] = { |
| bufferTypes.zipWithIndex.map { case (bufferType, index) => |
| GeneratedExpression( |
| bufferTerms(index), bufferNullTerms(index), "", bufferType.toInternalType) |
| } |
| } |
| |
| def accumulate(generator: ExprCodeGenerator): String = { |
| val resolvedExprs = if (generator.input1Term.startsWith(DISTINCT_KEY_TERM)) { |
| // called from distinct merge |
| function.accumulateExpressions |
| .map(_.postOrderTransform(resolveReference(isDistinctMerge = true))) |
| } else { |
| // called from accumulate |
| function.accumulateExpressions |
| .map(_.postOrderTransform(resolveReference())) |
| } |
| |
| val exprs = resolvedExprs |
| .map(_.toRexNode(relBuilder)) // rex nodes |
| .map(generator.generateExpression) // generated expressions |
| |
| val codes = exprs.zipWithIndex.map { case (expr, index) => |
| s""" |
| |${expr.code} |
| |${expr.copyResultTermToTargetIfChanged(ctx, bufferTerms(index))}; |
| |${bufferNullTerms(index)} = ${expr.nullTerm}; |
| """.stripMargin |
| } |
| |
| filterExpression match { |
| case Some(expr) => |
| val generated = generator.generateExpression(expr.toRexNode(relBuilder)) |
| s""" |
| |if (${generated.resultTerm}) { |
| | ${codes.mkString("\n")} |
| |} |
| """.stripMargin |
| case None => |
| codes.mkString("\n") |
| } |
| } |
| |
| def retract(generator: ExprCodeGenerator): String = { |
| val resolvedExprs = if (generator.input1Term.startsWith(DISTINCT_KEY_TERM)) { |
| // called from distinct merge |
| function.retractExpressions |
| .map(_.postOrderTransform(resolveReference(isDistinctMerge = true))) |
| } else { |
| // called from retract |
| function.retractExpressions |
| .map(_.postOrderTransform(resolveReference())) |
| } |
| |
| val exprs = resolvedExprs |
| .map(_.toRexNode(relBuilder)) // rex nodes |
| .map(generator.generateExpression) // generated expressions |
| |
| val codes = exprs.zipWithIndex.map { case (expr, index) => |
| s""" |
| |${expr.code} |
| |${expr.copyResultTermToTargetIfChanged(ctx, bufferTerms(index))}; |
| |${bufferNullTerms(index)} = ${expr.nullTerm}; |
| """.stripMargin |
| } |
| |
| filterExpression match { |
| case Some(expr) => |
| val generated = generator.generateExpression(expr.toRexNode(relBuilder)) |
| s""" |
| |if (${generated.resultTerm}) { |
| | ${codes.mkString("\n")} |
| |} |
| """.stripMargin |
| case None => |
| codes.mkString("\n") |
| } |
| } |
| |
| def merge(generator: ExprCodeGenerator): String = { |
| val exprs = function.mergeExpressions |
| .map(_.postOrderTransform(resolveReference(isMerge = true))) |
| .map(_.toRexNode(relBuilder)) // rex nodes |
| .map(generator.generateExpression) // generated expressions |
| |
| val codes = exprs.zipWithIndex.map { case (expr, index) => |
| s""" |
| |${expr.code} |
| |${expr.copyResultTermToTargetIfChanged(ctx, bufferTerms(index))}; |
| |${bufferNullTerms(index)} = ${expr.nullTerm}; |
| """.stripMargin |
| } |
| |
| codes.mkString("\n") |
| } |
| |
| def getValue(generator: ExprCodeGenerator): GeneratedExpression = { |
| val resolvedGetValueExpression = function.getValueExpression |
| .postOrderTransform(resolveReference()) |
| generator.generateExpression(resolvedGetValueExpression.toRexNode(relBuilder)) |
| } |
| |
| /** |
| * Resolves the given expression to a [[NamedExpression]]. |
| * |
| * @param isMerge this is called from merge() method |
| * @param isDistinctMerge this is called from distinct merge method |
| */ |
| private def resolveReference(isMerge: Boolean = false, isDistinctMerge: Boolean = false) |
| : PartialFunction[Expression, Expression] = { |
| case input: UnresolvedFieldReference => |
| // We always use UnresolvedFieldReference to represent reference of input field. |
| // In non-merge case, the input is the operand of the aggregate function. |
| // In merge case, the input is the aggregate buffers sent by local aggregate. |
| if (isMerge) { |
| val localIndex = function.inputAggBufferAttributes.indexOf(input) |
| // in merge case, the input1 is mergedAcc |
| ResolvedAggInputReference( |
| input.name, |
| mergedAccOffset + bufferIndexes(localIndex), |
| bufferTypes(localIndex).toInternalType) |
| } else { |
| val localIndex = function.operands.indexOf(input) |
| val inputIndex = argIndexes(localIndex) |
| // index to constant |
| if (inputIndex >= inputTypes.length) { |
| val constantIndex = inputIndex - inputTypes.length |
| val constantTerm = constantExprs(constantIndex).resultTerm |
| val nullTerm = constantExprs(constantIndex).nullTerm |
| val constantType = constantExprs(constantIndex).resultType |
| // constant is reused as member variable |
| ResolvedAggLocalReference(constantTerm, nullTerm, constantType) |
| } else { |
| if (isDistinctMerge) { // this is called from distinct merge |
| if (function.inputCount == 1) { |
| // the distinct key is a BoxedValue |
| ResolvedDistinctKeyReference(input.name, argTypes(0)) |
| } else { |
| // the distinct key is a BaseRow |
| ResolvedAggInputReference(input.name, localIndex, argTypes(localIndex)) |
| } |
| } else { |
| // the input is the inputRow |
| ResolvedAggInputReference(input.name, argIndexes(localIndex), argTypes(localIndex)) |
| } |
| } |
| } |
| |
| case buffer: UnresolvedAggBufferReference => |
| val localIndex = function.aggBufferAttributes.indexOf(buffer) |
| val name = bufferTerms(localIndex) |
| val nullTerm = bufferNullTerms(localIndex) |
| // buffer access is reused as member variable |
| ResolvedAggLocalReference(name, nullTerm, bufferTypes(localIndex).toInternalType) |
| } |
| |
| def checkNeededMethods( |
| needAccumulate: Boolean = false, |
| needRetract: Boolean = false, |
| needMerge: Boolean = false, |
| needReset: Boolean = false): Unit = { |
| // skip the check for DeclarativeAggregateFunction for now |
| } |
| } |
