Flink SQL 之 TOPN

// org.apache.flink.table.runtime.operators.rank.AppendOnlyTopNFunction#processElement
public void processElement(RowData input, Context context, Collector<RowData> out)
        throws Exception {
    long currentTime = context.timerService().currentProcessingTime();
    // 注册清除状态的定时器，OnTimer 函数中清状态
    registerProcessingCleanupTimer(context, currentTime);
    // Map<PartitionKey, TopNBuffer>
    // TopNBuffer(TreeMap<sortKey, RowData>)：这是一个好的TopN类，自己项目使用到也可以用
    // 初始化PartitionKey，对应的TopNBuffer
    initHeapStates();
    // 确定 rankend，本例中是常量2
    initRankEnd(input);

    RowData sortKey = sortKeySelector.getKey(input);
    // 检测input 数据是否在 topn中，代码在org.apache.flink.table.runtime.operators.rank.TopNBuffer#checkSortKeyInBufferRange
    // 具体逻辑，查看input 与TreeMap值比较，是否在 topn个数内
    if (checkSortKeyInBufferRange(sortKey, buffer)) {
        // input 插入 TreeMap
        buffer.put(sortKey, inputRowSer.copy(input));
        Collection<RowData> inputs = buffer.get(sortKey);
        // update data state
        // copy a new collection to avoid mutating state values, see CopyOnWriteStateMap,
        // otherwise, the result might be corrupt.
        // don't need to perform a deep copy, because RowData elements will not be updated
        // 保存TreeMap 数据到状态
        dataState.put(sortKey, new ArrayList<>(inputs));
        if (outputRankNumber || hasOffset()) {
            // the without-number-algorithm can't handle topN with offset,
            // so use the with-number-algorithm to handle offset
            processElementWithRowNumber(sortKey, input, out);
        } else {
            processElementWithoutRowNumber(input, out);
        }
    }
    // 若不满足无输出也不会保存状态
}

// LRU，treeMap 初始化
private void initHeapStates() throws Exception {
    requestCount += 1;
    RowData currentKey = (RowData) keyContext.getCurrentKey();
    buffer = kvSortedMap.get(currentKey);
    if (buffer == null) {
        buffer = new TopNBuffer(sortKeyComparator, ArrayList::new);
        kvSortedMap.put(currentKey, buffer);
        // restore buffer
        // 恢复 buffer，为什么这里会有恢复操作？
        // kvSortedMap 是在内存中的LRU，size=默认1000，很久不用的partitionkey会被删除
        // 在这里从状态 dataState重新组装成 buffer(treeMap)
        Iterator<Map.Entry<RowData, List<RowData>>> iter = dataState.iterator();
        if (iter != null) {
            while (iter.hasNext()) {
                Map.Entry<RowData, List<RowData>> entry = iter.next();
                RowData sortKey = entry.getKey();
                List<RowData> values = entry.getValue();
                // the order is preserved
                buffer.putAll(sortKey, values);
            }
        }
    } else {
        hitCount += 1;
    }
}

// 定时清理数据
public void onTimer(long timestamp, OnTimerContext ctx, Collector<RowData> out)
        throws Exception {
    if (stateCleaningEnabled) {
        // cleanup cache
        kvSortedMap.remove(keyContext.getCurrentKey());
        cleanupState(dataState);
    }
}

public void processElement(RowData input, Context ctx, Collector<RowData> out)
        throws Exception {
    long currentTime = ctx.timerService().currentProcessingTime();
    // register state-cleanup timer
    registerProcessingCleanupTimer(ctx, currentTime);
    initRankEnd(input);
    SortedMap<RowData, Long> sortedMap = treeMap.value();
    if (sortedMap == null) {
        sortedMap = new TreeMap<>(sortKeyComparator);
    }
    RowData sortKey = sortKeySelector.getKey(input);
    boolean isAccumulate = RowDataUtil.isAccumulateMsg(input);
    input.setRowKind(RowKind.INSERT); // erase row kind for further state accessing
    if (isAccumulate) {
        // update sortedMap
        if (sortedMap.containsKey(sortKey)) {
            sortedMap.put(sortKey, sortedMap.get(sortKey) + 1);
        } else {
            sortedMap.put(sortKey, 1L);
        }

        // emit
        if (outputRankNumber || hasOffset()) {
            // the without-number-algorithm can't handle topN with offset,
            // so use the with-number-algorithm to handle offset
            // rank 也输出
            emitRecordsWithRowNumber(sortedMap, sortKey, input, out);
        } else {
            emitRecordsWithoutRowNumber(sortedMap, sortKey, input, out);
        }
        // update data state
        List<RowData> inputs = dataState.get(sortKey);
        if (inputs == null) {
            // the sort key is never seen
            inputs = new ArrayList<>();
        }
        inputs.add(input);
        dataState.put(sortKey, inputs);
    } else {
        final boolean stateRemoved;
        // emit updates first
        if (outputRankNumber || hasOffset()) {
            // the without-number-algorithm can't handle topN with offset,
            // so use the with-number-algorithm to handle offset
            // rank 也输出
            stateRemoved = retractRecordWithRowNumber(sortedMap, sortKey, input, out);
        } else {
            stateRemoved = retractRecordWithoutRowNumber(sortedMap, sortKey, input, out);
        }

        // and then update sortedMap
        if (sortedMap.containsKey(sortKey)) {
            long count = sortedMap.get(sortKey) - 1;
            if (count == 0) {
                sortedMap.remove(sortKey);
            } else {
                sortedMap.put(sortKey, count);
            }
        } else {
            if (sortedMap.isEmpty()) {
                if (lenient) {
                    LOG.warn(STATE_CLEARED_WARN_MSG);
                } else {
                    throw new RuntimeException(STATE_CLEARED_WARN_MSG);
                }
            } else {
                throw new RuntimeException(
                        "Can not retract a non-existent record. This should never happen.");
            }
        }

        if (!stateRemoved) {
            // the input record has not been removed from state
            // should update the data state
            List<RowData> inputs = dataState.get(sortKey);
            if (inputs != null) {
                // comparing record by equaliser
                Iterator<RowData> inputsIter = inputs.iterator();
                while (inputsIter.hasNext()) {
                    if (equaliser.equals(inputsIter.next(), input)) {
                        inputsIter.remove();
                        break;
                    }
                }
                if (inputs.isEmpty()) {
                    dataState.remove(sortKey);
                } else {
                    dataState.put(sortKey, inputs);
                }
            }
        }
    }
    treeMap.update(sortedMap);
}

//org.apache.flink.table.planner.plan.rules.physical.stream.StreamPhysicalDeduplicateRule#canConvertToDeduplicate
def canConvertToDeduplicate(rank: FlinkLogicalRank): Boolean = {
  val sortCollation = rank.orderKey
  val rankRange = rank.rankRange
  val isRowNumberType = rank.rankType == RankType.ROW_NUMBER
  val isLimit1 = rankRange match {
    case rankRange: ConstantRankRange =>
      rankRange.getRankStart == 1 && rankRange.getRankEnd == 1
    case _ => false
  }
  val inputRowType = rank.getInput.getRowType
  val isSortOnTimeAttribute = sortOnTimeAttribute(sortCollation, inputRowType)
  !rank.outputRankNumber && isLimit1 && isSortOnTimeAttribute && isRowNumberType
}