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

public class CountUdaf extends AggregateFunction<Long, CountUdaf.CountAccum> {
    //定义存放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<CountAccum> 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 Table
val table = tEnv.sqlQuery("SELECT word, countUdaf(frequency), SUM(frequency) FROM WordCount GROUP BY word")

case class WC(word: String, frequency: Long)

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)
      }
    }
}

//堆栈如下
accumulate:25, CountUdaf (mytest)
accumulate:-1, DataSetAggregatePrepareMapHelper$5
combine: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"  

//堆栈如下
merge:29, CountUdaf (mytest)
mergeAccumulatorsPair:-1, DataSetAggregateFinalHelper$6
reduce: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"  

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

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

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" 

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

// grouped / keyed aggregation
if (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 aggregation
else {
      // 没有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)
}

public abstract class StateTable<K, N, S>
	implements StateSnapshotRestore, Iterable<StateEntry<K, N, S>> {
   /**
   * 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<K, N, S>[] 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}"

@Override
public <T> void transform(
   K key,
   N namespace,
   T value,
   StateTransformationFunction<S, T> transformation) throws Exception {

   final StateMapEntry<K, N, S> 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;