Flink 源码之数据分区

Flink源码分析系列文档目录

请点击：Flink 源码分析系列文档目录

前言

Flink流处理作业支持并行操作。每一个并行度可以理解为一个数据管道。上游和下游的并行度也可能是不同的。为了解决数据从上游到下游的分发策略问题，Flink提供了一系列分区（partition）算子。下面为大家介绍分区算子以及他们对应的Partitioner（分区器）。

5种分区算子

Flink共有5种类型的分区算子：

• 随机分区：随机挑选一个下游管道，将数据发送给它。代表算子为shuffle。
• 轮询分区：轮流将数据发送给各个下游管道。代表算子为rebalance（上游和所有下游建立连接，轮询发送）和 rescale（下游和上游对应关系是确定的，无论扩容还是缩容）。
• 广播：发送数据到所有下游管道。代表算子为broadcast。
• 按键值分区：需要用户指定key抽取逻辑。根据数据流中的元素，提取出这条数据的key，然后将key做hash运算，决定发送给下游哪个管道，确保key相同的数据一定被发送到同一个下游管道。代表算子为keyBy。
• 发送到下游算子第一个并行任务：上游所有管道的数据合在一起。代表算子为global。
• 自定义规则：自定义数据分发策略。代表算子为partitionCustom。

这些算子在底层对应了不同的Partitioner（分区器）。分区器为分区策略的具体实现方式。下面为大家列出这些算子创建分区器的逻辑。分区器的具体实现即将在下一个章节详细介绍。

shuffle

shuffle算子对应的分区器为ShufflePartitioner。

@PublicEvolving
public DataStream shuffle() {
    return setConnectionType(new ShufflePartitioner());
}

rebalance

rebalance算子对应的分区器为RebalancePartitioner。

public DataStream rebalance() {
    return setConnectionType(new RebalancePartitioner());
}

rescale

rescale算子对应的分区器为RescalePartitioner。

@PublicEvolving
public DataStream rescale() {
    return setConnectionType(new RescalePartitioner());
}

broadcast

broadcast算子对应的分区器为BroadcastPartitioner。

public DataStream broadcast() {
    return setConnectionType(new BroadcastPartitioner());
}

keyBy

keyBy算子对应的分区器为KeyGroupStreamPartitioner。

DataStream的keyBy方法代码如下：

public  KeyedStream keyBy(KeySelector key) {
    Preconditions.checkNotNull(key);
    return new KeyedStream<>(this, clean(key));
}

发现这里面没有使用到Partitioner，而是创建了一个KeyedStream。我们继续跟踪它的构造函数：

KeyedStream.java

public KeyedStream(DataStream dataStream, KeySelector keySelector) {
    this(
            dataStream,
            keySelector,
            TypeExtractor.getKeySelectorTypes(keySelector, dataStream.getType()));
}

public KeyedStream(
        DataStream dataStream,
        KeySelector keySelector,
        TypeInformation keyType) {
    this(
            dataStream,
            // 在这里创建出了一个`PartitionTransformation`，专用于分区操作
            // 它内部封装了`KeyGroupStreamPartitioner`
            new PartitionTransformation<>(
                    dataStream.getTransformation(),
                    new KeyGroupStreamPartitioner<>(
                            keySelector,
                            StreamGraphGenerator.DEFAULT_LOWER_BOUND_MAX_PARALLELISM)),
            keySelector,
            keyType);
}

global

global算子对应的分区器为GlobalPartitioner。

@PublicEvolving
public DataStream global() {
    return setConnectionType(new GlobalPartitioner());
}

partitionCustom

partitionCustom算子对应的分区器为CustomPartitionerWrapper。

public  DataStream partitionCustom(
        Partitioner partitioner, KeySelector keySelector) {
    return setConnectionType(
            new CustomPartitionerWrapper<>(clean(partitioner), clean(keySelector)));
}

8种分区器

前面章节讲解了Flink的5种分区算子，接下来为大家带来分区器的讲解。

Flink共有8种类型的分区器：

• BroadcastPartitioner
• CustomPartitionerWrapper
• ForwardPartitioner
• GlobalPartitioner
• KeyGroupStreamPartitioner
• RebalancePartitioner
• RescalePartitioner
• ShufflePartitioner

其中除了ForwardPartitioner，其余的在上面都已经提到过。

这些分区器的父类为StreamPartitioner。它本质上是一种ChannelSelector，用来根据规则决定数据被发往下游哪一个channel。

接下来我们父类StreamPartitioner开始，逐个介绍下分区器的实现。

StreamPartitioner和ChannelSelector

将StreamPartitioner和ChannelSelector放在第一位介绍的原因是我们需要提前了解它内部各个方法的作用。它的8个子类仅仅是覆盖或重写了这些方法而已。

ChannelSelector接口的各个方法定义如下所示：

public interface ChannelSelector {

    // 初始化ChannelSelector，传入的参数为下游channel（output channel）的数量
    void setup(int numberOfChannels);

    // 返回选择的channel索引编号，这个方法决定的上游的数据需要写入到哪个channel中
    // 这个方法的Partitioner子类重点需要实现的方法
    // 对于broadcast广播类型算子，不需要实现这个方法
    // 传入的参数为记录数据流中的元素，该方法需要根据元素来推断出需要发送到的下游channel
    int selectChannel(T record);

    // 返回是否为广播类型
    // 广播类型指的是上游数据发送给所有下游channel
    boolean isBroadcast();
}

StreamPartitioner抽象类实现了StreamPartitioner接口，它的代码如下所示：

@Internal
public abstract class StreamPartitioner
        implements ChannelSelector>>, Serializable {
    private static final long serialVersionUID = 1L;

    // 持有output channel数量
    protected int numberOfChannels;

    // 初始化的时候设置numberOfChannels的值
    @Override
    public void setup(int numberOfChannels) {
        this.numberOfChannels = numberOfChannels;
    }

    @Override
    public boolean isBroadcast() {
        return false;
    }

    public abstract StreamPartitioner copy();

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }
        final StreamPartitioner that = (StreamPartitioner) o;
        return numberOfChannels == that.numberOfChannels;
    }

    @Override
    public int hashCode() {
        return Objects.hash(numberOfChannels);
    }

    // 决定了作业恢复时候上游遇到扩缩容的话，需要处理哪些上游状态保存的数据
    public SubtaskStateMapper getUpstreamSubtaskStateMapper() {
        return SubtaskStateMapper.ARBITRARY;
    }

    /**
     * Defines the behavior of this partitioner, when downstream rescaled during recovery of
     * in-flight data.
     */
    // 同上，但关注的是下游扩缩容
    public abstract SubtaskStateMapper getDownstreamSubtaskStateMapper();

    // 重要
    // isPointwise方法决定了上游和下游的对应关系。
    // false表示没有指向性，上游和下游没有明确的对应关系
    // true表示上游和下游存在对应关系
    public abstract boolean isPointwise();
}

接下来我们开始介绍它的子类。

ForwardPartitioner

ForwardPartitioner直接将数据发送给本地下游（一对一）。要求上游的并行度和下游必须相同。

@Override
public int selectChannel(SerializationDelegate> record) {
    return 0;
}

@Override
public boolean isPointwise() {
    return true;
}

ForwardPartitioner在StreamGraph的addEdgeInternal方法中自动创建（生成StreamGraph的过程），代码片段如下所示：

// ...
if (partitioner == null
        && upstreamNode.getParallelism() == downstreamNode.getParallelism()) {
    // 只有在上游和下游的并行度相同的时候，才会使用ForwardPartitioner
    partitioner = new ForwardPartitioner