flink - 分配数据至指定keygroup及 subTask

本博客解决了：发送数据到指定并行度的问题，如我们指定某条数据的 key = 0，那么将该数据发送到 subTaskIndex = 0 的并行度，类推。

这是项目中遇到的问题，flink 数据在落库的时候，需要将数据拼成大的 String，而且想让几个并行度上拼出的 String 更均衡一些。于是想着 keyby random 的随机值来分配数据。当并行度为 2 时，我们 random 出来的数据是 0 和 1，keyby 这两个值会发现，数据都到了 subIndex = 0 的并行度上，显然这不是想要的，预想的是，keyBy=1，数据就到 subIndex=1的并行度上。

由于想让数据放入状态中，所以不使用在 subTask 上创建 List 收集数据的方式。

解决这个问题，需要理解：

1、flink 的 keygroup 包含哪些；

2、一个 key 该分配到哪个 keyGroup 中；

3、一个 subTask 处理哪些 keyGroup。

1、查看官网发现了以下，翻译过来，key group 的个数等于 maxParallelism，即索引位于 [0,maxParallelism - 1] 内。

Keyed State is further organized into so-called Key Groups. Key Groups are the atomic unit by which Flink can redistribute Keyed State; there are exactly as many Key Groups as the defined maximum parallelism.

2、在 Flink 1.11 源码的 KeyGroupRangeAssignment 类中可以看到以下，即 key 的 hashCode 经过 murmurhash 之后对 manParallelism 取余，所得到的 int 类型的值即为 keygroup Index。

public static int assignToKeyGroup(Object key, int maxParallelism) {
    Preconditions.checkNotNull(key, "Assigned key must not be null!");
    return computeKeyGroupForKeyHash(key.hashCode(), maxParallelism);
}

public static int computeKeyGroupForKeyHash(int keyHash, int maxParallelism) {
    return MathUtils.murmurHash(keyHash) % maxParallelism;
}

3、同样是 KeyGroupRangeAssignment 中，即 keygroupId * 当前并行度后，然后对 maxParallelism 取整。如果 maxParallelism = 8，curTaskParallellism = 3 时，keygroupId 分配如下：subTaskIndex = 0 [0,1,2]，subTaskIndex=1 [3,4,5]，subTaskIndex=2 [6,7]。

public static int computeOperatorIndexForKeyGroup(int maxParallelism, int parallelism, int keyGroupId) {
    return keyGroupId * parallelism / maxParallelism;
}

回到上边的问题，我们会发现，random.nextInt(2) 生成的 key 只包含 0和1，而这两个值都被分配到了 subTaskIndex = 0 上。

如果我们为每条数据通过 random.nextInt(curTaskParallism)生成后续 keyBy 的 key，得到的 0 为 key 的数据 keyby 后分配到 subTaskIndex = 0 的并行度上，得到的 1 为 key 的数据 keyby 后分配到 subTaskIndex = 1 的并行度上，那么我们需要通过以下步骤来解决问题：

1、首先我们需要根据作业的最大并行度，当前 operator 并行度，计算出当前算子的每个并行度上会分配哪些 keygroupId；

2、我们既然知道了某个 keygroupId 位于哪个并行度上了，那么每个并行度上选出一个 keygroupId 作为代表，然后根据 random 的整数生成一个“映射数字”到该 keygroupId 即可。我们知道 keygroupId 分配如下：subTaskIndex = 0 [0,1,2] ，即 random.nextInt(3) = 0，通过 MathUtils.murmurHash(keyHash) % maxParallelism = keygroupId = 0，计算出 keyHash 的值 即为 random = 0 的映射数字。

keygroupId 分配如下：subTaskIndex = 0 [0,1,2]，subTaskIndex=1 [3,4,5]，subTaskIndex=2 [6,7]

根据上边的分配，如果我们 random.nextInt(3)=1时，我们想让这条数据到 subTaskIndex=1的并行度上，那么我们只需要根据上一步计算出 MathUtils.murmurHash(keyHash) % maxParallelism = keygroupId = 3 即可，那么这个 keyHash 就是映射数字，数据分配到了对应的 keygroup，自然而然就到了对应的 subTask 上。

好了，将上边的逻辑生成一个 keyselector 的代码如下，很容易就可以看懂了。

对了，用的时候要注意两点：

1、如果直接random 生成随机数来使用 selector，库可能出现乱序覆盖问题；

2、可以根据 key 的 hashCode 来作为分配数据的依据，这样相同 key 的数据就会到同一个并行度上。如：selector.subTaskIndex(MathUtils.murmurHash(x.getField(3).toString.hashCode) % param.getProperty("curParallel").toInt)

class FlinkSubtaskIndexSelector extends Serializable {

private var subIndexMap: mutable.HashMap[Int,ListBuffer[Int]] = _
private var murmurKeyMap: mutable.HashMap[Int,Int] = _
private var maxParal: Int = _
private var subTaskParal: Int = _


def this(maxParallel: Int,subTaskParallel: Int){
    this()
    this.maxParal = maxParallel
    this.subTaskParal = subTaskParallel
    subIndexMap = new mutable.HashMap[Int,ListBuffer[Int]]()
    murmurKeyMap = new mutable.HashMap[Int,Int]()

    // 分配 keyGroup 至 subTask
    for(i <- 0 to maxParallel - 1){
        val curIndex = indexForKeyGroup(i)
        if(subIndexMap.get(curIndex).isEmpty){
            val list = new ListBuffer[Int]()
            list.append(i)
            subIndexMap.put(curIndex,list)
        }else {
            subIndexMap.get(curIndex).get.append(i)
        }
    }

    // 为 subIndex 查找 murmurKey
    var temp = 0
    val subIndexMapKeys = subIndexMap.keySet
    for(key <- subIndexMapKeys){
        val value = subIndexMap.get(key).get(0)
        while((MathUtils.murmurHash(temp) % maxParallel != value)){
            temp = temp + 1
        }
        murmurKeyMap.put(key,temp)
    }
}


private def indexForKeyGroup(kgroup: Int): Int = {
    kgroup * subTaskParal / maxParal
}


// 对 murmurHash 后的 key，会根据对当前 task 并行度的余数，返回对应 index
def subTaskIndex(key: Int): Int = {
    if(subIndexMap == null || murmurKeyMap == null){
        throw new RuntimeException("xxx")
    }
    murmurKeyMap.get(key).get
}

}