[源码解析] Flink UDAF 背后做了什么

源码解析 Flink UDAF 背后做了什么

​

本文涉及到Flink SQL UDAF，Window 状态管理等部分，希望能起到抛砖引玉的作用，让大家可以借此深入了解这个领域。

大家知道，Flink的自定义聚合函数（UDAF）可以将多条记录聚合成1条记录，这功能是通过accumulate方法来完成的，官方参考指出：

但是实时计算还有一些特殊的场景，在此场景下，还需要提供merge方法才能完成。

之前因为没亲身操作，所以一直忽略merge的特殊性。最近无意中看到了一个UDAF的实现，突然觉得有一个地方很奇怪，即 accumulate 和 merge 这两个函数不应该定义在一个类中。因为这是两个完全不同的处理方法。应该定义在两个不同的类中。

比如用UDAF做word count，则：

然后又想出了一个问题：Flink是如何管理 UDAF的accumulator？其状态存在哪里？

看起来应该是Flink在背后做了一些黑魔法，把这两个函数从一个类中拆分了。为了验证我们的推测，让我们从源码入手来看看这些问题：

示例代码摘要如下 ：

public class CountUdaf extends AggregateFunction {    //定义存放count UDAF状态的accumulator的数据的结构。    public static class CountAccum {        public long total;    }      //初始化count UDAF的accumulator。    public CountAccum createAccumulator() {        CountAccum acc = new CountAccum();        acc.total = 0;        return acc;    }      //accumulate提供了，如何根据输入的数据，更新count UDAF存放状态的accumulator。    public void accumulate(CountAccum accumulator, Object iValue) {        accumulator.total++;    }    public void merge(CountAccum accumulator, Iterable its) {        for (CountAccum other : its) {            accumulator.total += other.total;        }    }}

批处理相对简单，因为数据是有边界的，其逻辑比较清晰。

首先给出测试代码

val input = env.fromElements(WC("hello", 1), WC("hello", 1), WC("ciao", 1))// register the DataSet as a view "WordCount"tEnv.createTemporaryView("WordCount", input, 'word, 'frequency)tEnv.registerFunction("countUdaf", new CountUdaf())// run a SQL query on the Table and retrieve the result as a new Tableval table = tEnv.sqlQuery("SELECT word, countUdaf(frequency), SUM(frequency) FROM WordCount GROUP BY word")case class WC(word: String, frequency: Long)

在

DataSetAggregate.translateToPlan

于是我们推断，这很有可能就是 combineGroup 调用accumulate，reduceGroup 调用 merge。

关于combineGroup，如果有兴趣，可以看看我之前文章 [源码解析] Flink的groupBy和reduce究竟做了什么 以及 源码解析 GroupReduce，GroupCombine 和 Flink SQL group by](https://cloud.tencent.com/developer/article/1693307)

override def translateToPlan(tableEnv: BatchTableEnvImpl,    queryConfig: BatchQueryConfig): DataSet[Row] = {    if (grouping.length > 0) {      // grouped aggregation      if (preAgg.isDefined) {        // 执行到这里        inputDS          // pre-aggregation          .groupBy(grouping: _*)          .combineGroup(preAgg.get) // 第一阶段          .returns(preAggType.get)          .name(aggOpName)                    // final aggregation          .groupBy(grouping.indices: _*)          .reduceGroup(finalAgg.right.get) // 第二阶段          .returns(rowTypeInfo)          .name(aggOpName)      }    }}

SQL语句对应的执行计划大致为：

在执行看，确实对应了两个阶段。

阶段 1 确实是 GroupReduceCombineDriver 调用到了 accumulate。

//堆栈如下accumulate:25, CountUdaf (mytest)accumulate:-1, DataSetAggregatePrepareMapHelper$5combine:71, DataSetPreAggFunction (org.apache.flink.table.runtime.aggregate)sortAndCombine:213, GroupReduceCombineDriver (org.apache.flink.runtime.operators)run:188, GroupReduceCombineDriver (org.apache.flink.runtime.operators)  //SQL UDAF生成的代码如下  function = {DataSetAggregatePrepareMapHelper$5@10085}  function_mytest$CountUdaf$5ae272a09e5f36214da5c4e5436c4c48 = {CountUdaf@10079} "CountUdaf" function_org$apache$flink$table$functions$aggfunctions$LongSumAggFunction$a5214701531789b3139223681d = {LongSumAggFunction@10087} "LongSumAggFunction"  

阶段 2 中 GroupReduceDriver 调用到了 merge

//堆栈如下merge:29, CountUdaf (mytest)mergeAccumulatorsPair:-1, DataSetAggregateFinalHelper$6reduce:71, DataSetFinalAggFunction (org.apache.flink.table.runtime.aggregate)run:131, GroupReduceDriver (org.apache.flink.runtime.operators)  //SQL UDAF生成的代码如下   function = {DataSetAggregateFinalHelper$6@10245}  function_mytest$CountUdaf$5ae272a09e5f36214da5c4e5436c4c48 = {CountUdaf@10238} "CountUdaf" function_org$apache$flink$table$functions$aggfunctions$LongSumAggFunction$a5214701531789b3139223681d = {LongSumAggFunction@10247} "LongSumAggFunction"  

Flink对用户定义的UDAF代码分别生成了两个不同的功能类：

UDAF有一个accumulator，这个会在程序运行过程中始终存在，Flink是如何管理这个accumulator呢？

GroupReduceCombineDriver类有一个成员变量 combiner，

public class GroupReduceCombineDriver implements Driver, OUT> {  private GroupCombineFunction combiner;}

而 combiner 被赋予了 DataSetPreAggFunction 类的一个实例。

class DataSetPreAggFunction(genAggregations: GeneratedAggregationsFunction)  extends AbstractRichFunction{  private var accumulators: Row = _ //这里存储历史状态  private var function: GeneratedAggregations = _}

Flink就是把 UDAF的accumulator 存储在

combiner.accumulators

combiner = {DataSetPreAggFunction@10063}  genAggregations = {GeneratedAggregationsFunction@10070}  accumulators = {Row@10117} "mytest.CountUdaf$CountAccum@1e343db7,(0,false)" function = {DataSetAggregatePrepareMapHelper$5@10066}  // function是包含用户代码的功能类。  function_mytest$CountUdaf$5ae272a09e5f36214da5c4e5436c4c48 = {CountUdaf@10076} "CountUdaf" 

让我们总结一下，批处理被分成两个阶段：

Flink在GroupReduceCombineDriver类的成员变量 combiner 中存储 accumulator历史状态。

流处理则是和批处理完全不同的世界，下面我们看看流处理背后有什么奥秘。

Yodayo

一个专为动漫迷和vTuber打造的AI艺术创作平台、交流社区

下载

在流计算场景中，数据没有边界源源不断的流入的，每条数据流入都可能会触发计算，比如在进行count或sum这些操作是如何计算的呢？

val query: Table = tableEnv.sqlQuery(  """    |SELECT    |countUdaf(num)    |FROM tb_num    |GROUP BY TUMBLE(proctime, INTERVAL '10' SECOND)   """.stripMargin)

DataStreamGroupWindowAggregateBase.translateToPlan

WindowedStream

WindowAssigner

WindowedStream

KeyedStream

AllWindowedStream

AllWindowedStream

DataStream

windowAll(...)

我们的示例代码是基于Key的，所以走

WindowedStream

// grouped / keyed aggregationif (grouping.length > 0) {      // 有key，所以是 WindowedStream，我们示例走这里      val windowFunction = AggregateUtil.createAggregationGroupWindowFunction(...)      val keySelector = new CRowKeySelector(grouping, inputSchema.projectedTypeInfo(grouping))      val keyedStream = timestampedInput.keyBy(keySelector)      val windowedStream =        createKeyedWindowedStream(queryConfig, window, keyedStream)          .asInstanceOf[WindowedStream[CRow, Row, DataStreamWindow]]      val (aggFunction, accumulatorRowType) =        AggregateUtil.createDataStreamGroupWindowAggregateFunction(...)      windowedStream        .aggregate(aggFunction, windowFunction, accumulatorRowType, outRowType)        .name(keyedAggOpName)}// global / non-keyed aggregationelse {      // 没有key，所以是AllWindowedStream       val windowFunction = AggregateUtil.createAggregationAllWindowFunction(...)      val windowedStream =        createNonKeyedWindowedStream(queryConfig, window, timestampedInput)          .asInstanceOf[AllWindowedStream[CRow, DataStreamWindow]]      val (aggFunction, accumulatorRowType) =        AggregateUtil.createDataStreamGroupWindowAggregateFunction(...)      windowedStream        .aggregate(aggFunction, windowFunction, accumulatorRowType, outRowType)        .name(nonKeyedAggOpName)}

SQL语句对应的执行计划大致如下，我们能看出来 accumulate & merge 都在 Window 中处理。

可以看到，流处理对UDAF的管理，就完全是进入了Window的地盘，而UDAF历史状态管理其实就是Flink Window状态管理的领域了。

我们以基于key的WindowedStream为例继续进行研究。

当Window接受到一个输入item时候，item会被分配到一个key，由KeySelector完成。WindowOperator 类首先使用用户选择的 windowAssigner 将流入的数据分配到响应的window中，有可能是1个，0个甚至多个window。这里就会做accumulate。

本例

windowAssigner = {TumblingProcessingTimeWindows}

windowState.add(element.getValue());

public abstract class StateTableimplements StateSnapshotRestore, Iterable> {   /**   * Map for holding the actual state objects. The outer array represents the key-groups.   * All array positions will be initialized with an empty state map.   */protected final StateMap[] keyGroupedStateMaps;}// 真实中变量摘录如下keyGroupedStateMaps = {StateMap[1]@9266}  0 = {CopyOnWriteStateMap@9262} // 这里就是将要保存用户accumulator的地方  stateSerializer = {RowSerializer@9254}   snapshotVersions = {TreeSet@9277}  size = 0  primaryTable = {CopyOnWriteStateMap$StateMapEntry[128]@9278}   incrementalRehashTable = {CopyOnWriteStateMap$StateMapEntry[2]@9280}   lastNamespace = {TimeWindow@9239} "TimeWindow{start=1593934200000, end=1593934210000}"

3.1.2)stateMap.transform(key, namespace, value, transformation);

@Overridepublic  void transform(   K key,   N namespace,   T value,   StateTransformationFunction transformation) throws Exception {   final StateMapEntry entry = putEntry(key, namespace);   // copy-on-write check for state   entry.state = transformation.apply(      (entry.stateVersion < highestRequiredSnapshotVersion) ?         getStateSerializer().copy(entry.state) : entry.state,         value);    // 当执行完用户代码之后，数据会存储在这里，这个就是CopyOnWriteStateMap的一个Entry   entry.stateVersion = stateMapVersion;