Flink 1.19中RichParallelSourceFunction迁移至Source接口并保留并行度咨询

你的现有实现存在核心问题：每个并行子任务的 SourceReader 都会拿到完整的输入数据迭代器，导致所有子任务重复处理全量数据，并行度没有起到正确的分片处理作用。要保留原 RichParallelSourceFunction 的并行能力，必须通过「Split 拆分-分发」机制让每个并行子任务处理独立的数据分片，以下是修正后的完整实现和关键说明：

核心修改思路 #

Flink 新的 Source 接口通过 SplitEnumerator 负责数据分片的拆分与分发，SourceReader 负责处理分配给自己的分片数据。要实现并行度，必须：

1. 自定义 Split 类承载分片数据
2. 在 Enumerator 中根据并行度拆分输入数据，将分片分发给不同子任务
3. Reader 只处理分配给自己的分片

完整修正代码 #

1. 自定义 Split 类 #

class ByteArraySplit(val data: List<ByteArray>, val splitId: String) : SourceSplit {
    override fun splitId(): String = splitId
}

2. 实现分片分发的 SplitEnumerator #

class SimpleSplitEnumerator(
    private val input: List<ByteArray>,
    private val enumeratorContext: SplitEnumeratorContext<ByteArraySplit>
) : SplitEnumerator<ByteArraySplit, Void> {
    private var splitsAssigned = false

    override fun start() {
        if (!splitsAssigned) {
            val parallelism = enumeratorContext.currentParallelism()
            // 按并行度拆分输入数据，每个分片对应一个子任务
            val splits = splitInputIntoParts(input, parallelism)
            // 将分片分配给对应子任务
            splits.forEachIndexed { index, split ->
                enumeratorContext.assignSplit(split, index)
            }
            // 通知所有子任务无更多分片
            enumeratorContext.signalNoMoreSplitsToAll()
            splitsAssigned = true
        }
    }

    // 按并行度均匀拆分输入数据
    private fun splitInputIntoParts(input: List<ByteArray>, numParts: Int): List<ByteArraySplit> {
        if (numParts <= 0) return emptyList()
        val splits = mutableListOf<ByteArraySplit>()
        val chunkSize = (input.size + numParts - 1) / numParts // 向上取整保证每个分片数据均匀
        for (i in 0 until numParts) {
            val start = i * chunkSize
            val end = minOf((i + 1) * chunkSize, input.size)
            val part = input.subList(start, end)
            if (part.isNotEmpty()) {
                splits.add(ByteArraySplit(part, "split-$i"))
            }
        }
        return splits
    }

    override fun handleSplitRequest(subtaskId: Int, requesterHostname: String?) {}
    override fun addSplitsBack(splits: List<ByteArraySplit>, subtaskId: Int) {}
    override fun addReader(readerId: Int) {}
    override fun snapshotState(checkpointId: Long): Void? = null
    override fun close() {}
}

3. 实现分片处理的 SourceReader #

class SourceWithSleepReader(
    private val sleep: Long
) : SourceReader<ByteArray, ByteArraySplit> {
    private var inputIterator: Iterator<ByteArray>? = null
    @Volatile
    private var stop = false

    override fun start() {}

    override fun pollNext(output: ReaderOutput<ByteArray>?): InputStatus {
        if (stop) return InputStatus.END_OF_INPUT

        val iterator = inputIterator ?: return InputStatus.NOTHING_AVAILABLE

        if (iterator.hasNext()) {
            output?.collect(iterator.next())
            Thread.sleep(sleep)
            return InputStatus.MORE_AVAILABLE
        }

        return InputStatus.END_OF_INPUT
    }

    override fun snapshotState(checkpointId: Long): List<ByteArraySplit> {
        // 容错场景可在此记录迭代位置，示例简化暂不实现
        return emptyList()
    }

    override fun isAvailable(): CompletableFuture<Void> {
        return CompletableFuture.completedFuture(null)
    }

    override fun notifyNoMoreSplits() {}

    // 接收分配给自己的分片并初始化迭代器
    override fun addSplits(splits: MutableList<ByteArraySplit>?) {
        splits?.firstOrNull()?.let { split ->
            inputIterator = split.data.iterator()
        }
    }

    override fun close() {
        stop = true
    }
}

4. 主 Source 类适配 #

class SourceWithSleep(
    private val input: List<ByteArray>,
    private val sleep: Long
) : Source<ByteArray, ByteArraySplit, Void> {

    override fun getBoundedness(): Boundedness = Boundedness.BOUNDED

    override fun createReader(
        readerContext: SourceReaderContext
    ): SourceReader<ByteArray, ByteArraySplit> {
        return SourceWithSleepReader(sleep)
    }

    override fun createEnumerator(
        enumeratorContext: SplitEnumeratorContext<ByteArraySplit>
    ): SplitEnumerator<ByteArraySplit, Void> {
        return SimpleSplitEnumerator(input, enumeratorContext)
    }

    override fun restoreEnumerator(
        enumeratorContext: SplitEnumeratorContext<ByteArraySplit>,
        checkpoint: Void?
    ): SplitEnumerator<ByteArraySplit, Void> {
        return SimpleSplitEnumerator(input, enumeratorContext)
    }

    override fun getSplitSerializer(): SimpleVersionedSerializer<ByteArraySplit> {
        return ByteArraySplitSerializer()
    }

    override fun getEnumeratorCheckpointSerializer(): SimpleVersionedSerializer<Void> {
        return VoidSerializer()
    }
}

// Split 序列化器（生产环境建议用Protobuf等高效序列化方式）
class ByteArraySplitSerializer : SimpleVersionedSerializer<ByteArraySplit> {
    override fun getVersion(): Int = 1

    override fun serialize(split: ByteArraySplit): ByteArray {
        val bos = ByteArrayOutputStream()
        val oos = ObjectOutputStream(bos)
        oos.writeObject(split.data)
        oos.writeUTF(split.splitId())
        oos.close()
        return bos.toByteArray()
    }

    @Suppress("UNCHECKED_CAST")
    override fun deserialize(version: Int, serialized: ByteArray): ByteArraySplit {
        val bis = ByteArrayInputStream(serialized)
        val ois = ObjectInputStream(bis)
        val data = ois.readObject() as List<ByteArray>
        val splitId = ois.readUTF()
        ois.close()
        return ByteArraySplit(data, splitId)
    }
}

// 空状态序列化器
class VoidSerializer : SimpleVersionedSerializer<Void> {
    override fun getVersion(): Int = 1
    override fun serialize(split: Void?): ByteArray {
        return ByteArray(0)
    }
    override fun deserialize(version: Int, serialized: ByteArray?): Void? {
        return null
    }
}

关键说明 #

1. 并行度生效逻辑：当你设置 Source 的并行度时，SplitEnumerator 会根据当前并行度拆分输入数据，每个并行子任务只会收到自己对应的分片，避免了原实现中重复处理全量数据的问题，和原 RichParallelSourceFunction 的并行行为完全一致。
2. 分片拆分策略：示例中采用了均匀拆分的方式，你也可以根据业务需求调整拆分逻辑（比如按数据大小、按自定义规则拆分）。
3. 容错扩展：如果需要支持 checkpoint 容错，可以在 snapshotState 方法中记录当前 Reader 的迭代位置，恢复时重新生成对应位置的分片。

内容的提问来源于stack exchange，提问作者João Carneiro