解决Flink的KeyBy分布不均衡问题（映射新Key）

直接上一段代码：

java 复制代码

import org.apache.flink.runtime.state.KeyGroupRangeAssignment;

import java.util.HashMap;
import java.util.Map;

public class KeyBalanceUtil {

    private static Integer[] balanceKeys;
    private static volatile int parallelism;

    public static void init(int currentParallelism) {
        if (parallelism == 0) {
            synchronized (KeyBalanceUtil.class) {
                if (parallelism == 0) {
                    parallelism = currentParallelism;
                    balanceKeys = createBalanceKeys(currentParallelism);
                }
            }
        }
    }

    public static Integer[] createBalanceKeys(int parallelism) {
        int maxParallelism = KeyGroupRangeAssignment.computeDefaultMaxParallelism(parallelism);
        int maxRandomKey = parallelism * 12;
        Map<Integer, Integer> key_subIndex_map = new HashMap<>();
        for (int randomKey = 0; randomKey < maxRandomKey; randomKey++) {
            int subtaskIndex = KeyGroupRangeAssignment.assignKeyToParallelOperator(randomKey, maxParallelism, parallelism);
            if (key_subIndex_map.containsKey(subtaskIndex))
                continue;
            key_subIndex_map.put(subtaskIndex, randomKey);
        }
        return key_subIndex_map.values().toArray(new Integer[0]);
    }

    public static Integer mappingKey(int key) {
        if (parallelism == 0) {
            return 0;
        }
        return balanceKeys[key % parallelism];
    }
}

这段代码是用于解决 Flink 中 keyBy 可能导致的数据倾斜问题的工具类，其核心作用是通过预生成的平衡键（balance keys）实现更均匀的数据分布。以下从功能实现、设计思路和关键细节进行分点说明：

一、功能作用

解决数据倾斜问题
当 Flink 的 keyBy 操作中某些 Key 的哈希值分布不均匀时，会导致部分子任务（subtask）负载过高。此工具通过 预计算一组均衡分布的 Key，确保每个子任务都能分配到近似数量的 Key，从而缓解数据倾斜。
动态适配并行度
通过 init(currentParallelism) 方法初始化并行度，支持在任务并行度变化时重新生成平衡键（需重启作业），确保 Key 分布始终与当前并行度匹配。

二、设计原理

平衡键生成逻辑（createBalanceKeys）
- 计算最大并行度 ：基于 Flink 内置方法 computeDefaultMaxParallelism 确定 Key 组的最大范围，避免硬编码限制。
- 采样候选 Key ：生成 parallelism * 12 个随机 Key（经验值），遍历这些 Key 并记录每个子任务首次分配到的 Key。
- 构建映射表 ：最终为每个子任务保留一个唯一 Key（key_subIndex_map），保证每个子任务至少有一个 Key 被选中。
Key 映射逻辑（mappingKey）
- 输入任意原始 Key 时，通过 key % parallelism 计算索引，从 balanceKeys 中获取预生成的平衡键。
- 例如：若并行度为 3，原始 Key 为 100，则 100 % 3 = 1，返回 balanceKeys[1] 对应的平衡键。