如何在Kafka Streams中获取输出记录的对应分区信息

有两种可行的实现方式，不需要依赖外部工具，全部逻辑可以在KStream内部完成：

方案1：自定义StreamPartitioner（通用性最高，兼容所有版本） #

你可以在调用to()方法写入输出主题时，传入自定义的分区器实现，在分区计算逻辑中直接获取并记录目标分区：

• 自定义分区器可以复用Kafka默认的哈希/轮询分区逻辑，也可以替换为你自己的业务分区规则
• 计算得到的目标分区可以写入消息头，供后续处理逻辑读取，也可以直接在分区器中完成你需要的相关处理逻辑

示例代码如下：

import org.apache.kafka.streams.processor.StreamPartitioner;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.ThreadLocalRandom;

public class OutputPartitionCapturePartitioner<K, V> implements StreamPartitioner<K, V> {
    @Override
    public Integer partition(String topic, K key, V value, int numPartitions) {
        // 此处为默认分区逻辑，可按需替换为自定义分区规则
        int targetPartition = key == null ? 
            ThreadLocalRandom.current().nextInt(numPartitions) :
            Math.abs(key.hashCode() % numPartitions);
        
        // 将输出分区写入消息头，后续处理节点可直接从header中读取
        context.headers().add("output-target-partition", 
            String.valueOf(targetPartition).getBytes(StandardCharsets.UTF_8));
        
        return targetPartition;
    }
}

调用时传入分区器即可：

yourKStream.to(
    "your-output-topic",
    Produced.with(keySerde, valueSerde, new OutputPartitionCapturePartitioner<>())
);

方案2：使用低阶Processor API（适合需要细粒度控制的场景） #

如果你的分支逻辑本身就用了Processor API实现，可以在process()方法中调用context.forward()时直接指定目标分区，同时记录该分区值即可，该方案仅适用于Kafka Streams 2.7及以上版本。

注意事项 #

• 不要尝试通过输入消息的分区推断输出分区，输出主题分区数和输入不一致时两者没有固定对应关系
• 如果你不打算自定义分区逻辑，也可以在业务代码中复现Kafka默认分区器的计算逻辑获取分区，注意要和Kafka官方逻辑保持一致，避免计算结果和实际写入分区不一致的问题

内容的提问来源于stack exchange，提问作者B.Gen.Jack.O.Neill