Apache Beam KafkaIO写入多主题实现方案咨询及代码示例请求

Apache Beam的KafkaIO.write（Java）或WriteToKafka（Python）本身仅支持绑定单个Topic，因此最优实现思路是将处理后的数据集按对象类型拆分，为每个类型单独创建Kafka写入分支。以下提供两种常用实现方式及代码示例：

核心思路 #

先从多输入Topic读取数据，统一处理生成带类型标识的对象，再通过拆分操作将数据分流到不同的PCollection，最后每个PCollection独立写入对应输出Topic。

Java 实现示例 #

假设处理后的数据是自定义对象ProcessedData，包含getType()方法返回类型标识（如TYPE_A/TYPE_B/TYPE_C）。

1. 读取多输入Topic #

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.kafka.KafkaIO;
import org.apache.beam.sdk.values.KV;
import java.util.Arrays;
import org.apache.kafka.common.serialization.StringDeserializer;

Pipeline pipeline = Pipeline.create();

PCollection<KV<String, String>> input = pipeline.apply(KafkaIO.<String, String>read()
        .withBootstrapServers("kafka-broker:9092")
        .withTopics(Arrays.asList("input-topic-1", "input-topic-2", "input-topic-3"))
        .withKeyDeserializer(StringDeserializer.class)
        .withValueDeserializer(StringDeserializer.class)
        .withoutMetadata());

2. 处理数据生成带类型的对象 #

import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollection;

// 自定义处理后的对象
class ProcessedData {
    private String type;
    private String content;
    // 构造方法、getter/setter省略
    public String getType() { return type; }
}

PCollection<ProcessedData> processedData = input.apply(ParDo.of(new DoFn<KV<String, String>, ProcessedData>() {
    @ProcessElement
    public void processElement(ProcessContext c) {
        String key = c.element().getKey();
        String value = c.element().getValue();
        // 自定义逻辑：将原始KV转换为ProcessedData
        ProcessedData data = new ProcessedData();
        data.setContent(value);
        // 根据业务逻辑设置type，示例仅做演示
        if (value.contains("typeA")) data.setType("TYPE_A");
        else if (value.contains("typeB")) data.setType("TYPE_B");
        else data.setType("TYPE_C");
        c.output(data);
    }
}));

方式一：使用Partition拆分后写入 #

适合固定数量的分流场景，代码更紧凑：

import org.apache.beam.sdk.transforms.Partition;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.kafka.common.serialization.StringSerializer;

// 按类型拆分PCollection为3个分区
PCollectionList<ProcessedData> partitionedData = processedData.apply(Partition.of(3, (data, numPartitions) -> {
    switch (data.getType()) {
        case "TYPE_A": return 0;
        case "TYPE_B": return 1;
        case "TYPE_C": return 2;
        default: throw new IllegalArgumentException("Unknown data type: " + data.getType());
    }
}));

// 分别写入对应Topic
partitionedData.get(0).apply("Write to Topic A", KafkaIO.<String, ProcessedData>write()
        .withBootstrapServers("kafka-broker:9092")
        .withTopic("output-topic-a")
        .withKeySerializer(StringSerializer.class)
        .withValueSerializer(ProcessedDataSerializer.class) // 自定义序列化器，需实现Kafka Serializer接口
        .values());

partitionedData.get(1).apply("Write to Topic B", KafkaIO.<String, ProcessedData>write()
        .withBootstrapServers("kafka-broker:9092")
        .withTopic("output-topic-b")
        .withKeySerializer(StringSerializer.class)
        .withValueSerializer(ProcessedDataSerializer.class)
        .values());

partitionedData.get(2).apply("Write to Topic C", KafkaIO.<String, ProcessedData>write()
        .withBootstrapServers("kafka-broker:9092")
        .withTopic("output-topic-c")
        .withKeySerializer(StringSerializer.class)
        .withValueSerializer(ProcessedDataSerializer.class)
        .values());

方式二：使用Filter过滤后写入 #

逻辑更直观，便于单独维护每个分支：

import org.apache.beam.sdk.transforms.Filter;

// 过滤TYPE_A并写入
processedData.apply("Filter TYPE_A", Filter.by(data -> "TYPE_A".equals(data.getType())))
        .apply("Write to Topic A", KafkaIO.<String, ProcessedData>write()
                .withBootstrapServers("kafka-broker:9092")
                .withTopic("output-topic-a")
                .withKeySerializer(StringSerializer.class)
                .withValueSerializer(ProcessedDataSerializer.class)
                .values());

// 过滤TYPE_B并写入
processedData.apply("Filter TYPE_B", Filter.by(data -> "TYPE_B".equals(data.getType())))
        .apply("Write to Topic B", KafkaIO.<String, ProcessedData>write()
                .withBootstrapServers("kafka-broker:9092")
                .withTopic("output-topic-b")
                .withKeySerializer(StringSerializer.class)
                .withValueSerializer(ProcessedDataSerializer.class)
                .values());

// 过滤TYPE_C并写入
processedData.apply("Filter TYPE_C", Filter.by(data -> "TYPE_C".equals(data.getType())))
        .apply("Write to Topic C", KafkaIO.<String, ProcessedData>write()
                .withBootstrapServers("kafka-broker:9092")
                .withTopic("output-topic-c")
                .withKeySerializer(StringSerializer.class)
                .withValueSerializer(ProcessedDataSerializer.class)
                .values());

Python 实现示例 #

1. 读取多输入Topic #

import apache_beam as beam
import json

pipeline = beam.Pipeline()

input_data = pipeline | "Read from Kafka" >> beam.io.ReadFromKafka(
    consumer_config={'bootstrap.servers': 'kafka-broker:9092'},
    topics=['input-topic-1', 'input-topic-2', 'input-topic-3']
)

2. 处理数据生成带类型的对象 #

def process_raw_data(element):
    key, value = element
    # 自定义处理逻辑，返回带type标识的字典
    processed = {
        'key': key,
        'content': value,
        'type': 'TYPE_A' if 'typeA' in value else 'TYPE_B' if 'typeB' in value else 'TYPE_C'
    }
    return processed

processed_data = input_data | "Process Data" >> beam.Map(process_raw_data)

过滤后写入多Topic #

# 写入Topic A
processed_data | "Filter TYPE_A" >> beam.Filter(lambda x: x['type'] == 'TYPE_A') \
               | "Write to Topic A" >> beam.io.WriteToKafka(
                   producer_config={'bootstrap.servers': 'kafka-broker:9092'},
                   topic='output-topic-a',
                   key_serializer=lambda x: x['key'].encode('utf-8'),
                   value_serializer=lambda x: json.dumps(x).encode('utf-8')
               )

# 写入Topic B
processed_data | "Filter TYPE_B" >> beam.Filter(lambda x: x['type'] == 'TYPE_B') \
               | "Write to Topic B" >> beam.io.WriteToKafka(
                   producer_config={'bootstrap.servers': 'kafka-broker:9092'},
                   topic='output-topic-b',
                   key_serializer=lambda x: x['key'].encode('utf-8'),
                   value_serializer=lambda x: json.dumps(x).encode('utf-8')
               )

# 写入Topic C
processed_data | "Filter TYPE_C" >> beam.Filter(lambda x: x['type'] == 'TYPE_C') \
               | "Write to Topic C" >> beam.io.WriteToKafka(
                   producer_config={'bootstrap.servers': 'kafka-broker:9092'},
                   topic='output-topic-c',
                   key_serializer=lambda x: x['key'].encode('utf-8'),
                   value_serializer=lambda x: json.dumps(x).encode('utf-8')
               )

注意事项 #

• 自定义对象需实现Kafka对应的序列化器（Java），或在写入前转换为可序列化的格式（如Python中转为JSON字符串）。
• 若分流逻辑复杂，可结合ParDo输出到不同的PCollection（通过SideOutput），再分别写入Topic，这种方式适合更多类型的分流场景。

内容的提问来源于stack exchange，提问作者Prasad