Java Beam按Key分目录写入Avro文件实现方案咨询

要实现按每个Key输出独立Avro文件，核心是利用Beam的Dynamic Destinations功能动态路由输出路径，同时需要先处理GroupByKey后的Iterable数据结构。以下是具体实现步骤：

步骤1：展开GroupByKey后的Iterable元素 #

GroupByKey输出的是KV<String, Iterable<String>>，每个Key对应一组值。需要先将这组值拆分为单个带Key的记录，方便后续按Key路由：

// 将每个Key对应的Iterable<String>展开为单个KV<String, String>
PCollection<KV<String, String>> flattenedRecords = sucessResponsepCollection.apply(
    "Flatten Iterable Values",
    ParDo.of(new DoFn<KV<String, Iterable<String>>, KV<String, String>>() {
        @ProcessElement
        public void processElement(ProcessContext ctx) {
            String key = ctx.element().getKey();
            for (String value : ctx.element().getValue()) {
                ctx.output(KV.of(key, value));
            }
        }
    })
);

步骤2：实现DynamicAvroDestinations类 #

这个类负责完成三个核心操作：提取路由Key、指定输出路径、将数据转换为Avro GenericRecord：

import org.apache.beam.sdk.io.avro.DynamicAvroDestinations;
import org.apache.beam.sdk.io.avro.FileNaming;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.ValueInSingleWindow;
import org.apache.avro.Schema;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericRecord;

public class KeyAvroDestinations extends DynamicAvroDestinations<KV<String, String>, String, GenericRecord> {
    private final Schema avroSchema;

    public KeyAvroDestinations(Schema schema) {
        this.avroSchema = schema;
    }

    // 提取用于路由的Key值
    @Override
    public String getDestination(ValueInSingleWindow<KV<String, String>> element) {
        return element.getValue().getKey();
    }

    // 定义每个Key对应的输出路径前缀
    @Override
    public String getOutputFilePrefix(String destinationKey, FileNaming fileNaming) {
        return destinationKey + "/part";
    }

    // 将KV转换为符合Schema的Avro GenericRecord
    @Override
    public GenericRecord formatRecord(KV<String, String> element) {
        GenericRecord record = new GenericData.Record(avroSchema);
        record.put("key", element.getKey()); // 需与你的Avro Schema字段名对应
        record.put("value", element.getValue()); // 需与你的Avro Schema字段名对应
        return record;
    }

    // 返回统一的Avro Schema（若不同Key需不同Schema可在此分支处理）
    @Override
    public Schema getSchema(String destinationKey) {
        return avroSchema;
    }
}

步骤3：配置AvroIO使用动态目标 #

替换原固定路径的写入逻辑，改用DynamicDestinations实现按Key输出：

import org.apache.avro.Schema;
import java.io.File;

// 加载你的Avro Schema文件
Schema avroSchema = new Schema.Parser().parse(new File(schemaFile));

// 按Key写入独立Avro文件
flattenedRecords.apply(
    "Write Avro Files by Key",
    AvroIO.write(KV.class)
           .to(new KeyAvroDestinations(avroSchema))
           .withWindowedWrites()
           .withNumShards(options.getNumberShards() == null ? 1 : options.getNumberShards())
           .withSuffix(AVRO_FILENAME_SUFFIX)
);

关键注意事项 #

• 若需每个Key仅生成一个文件，建议设置withNumShards(1)，避免同一Key生成多个分片文件。
• 确保formatRecord方法中字段名与你的Avro Schema完全匹配，否则会抛出Schema不匹配异常。
• DynamicDestinations会自动创建Key对应的目录，例如Key为"model_001"时，文件会写入model_001/part-xxxx.avro。

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