Apache Beam2.9使用writeDynamic将Avro文件写入到GCS上的多个目录

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.AvroIO;
import org.apache.beam.sdk.io.FileIO;
import org.apache.beam.sdk.io.fs.ResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;

public class AvroWriter {
  public static void main(String[] args) {
    // 创建Pipeline
    Pipeline pipeline = Pipeline.create(PipelineOptionsFactory.fromArgs(args).create());

    // 从输入源读取数据
    PCollection<MyData> input = pipeline.apply(AvroIO.read(MyData.class).from("input.avro"));

    // 定义用于将数据写入GCS的DoFn
    DoFn<MyData, KV<String, MyData>> writeToGcsFn = new DoFn<MyData, KV<String, MyData>>() {
      @ProcessElement
      public void processElement(ProcessContext c) {
        MyData data = c.element();
        String directory = determineDirectory(data); // 根据数据的某个属性值确定目录
        c.output(KV.of(directory, data));
      }

      private String determineDirectory(MyData data) {
        // 根据数据的某个属性值确定目录，这里只是示例，具体实现需要根据实际需求来定
        return "directory/" + data.getProperty();
      }
    };

    // 将数据按目录写入GCS
    TupleTag<KV<String, MyData>> mainTag = new TupleTag<>();
    TupleTagList additionalTags = TupleTagList.empty();
    PCollection<KV<String, MyData>> output = input.apply(ParDo.of(writeToGcsFn).withOutputTags(mainTag, additionalTags));
    output.apply(FileIO.<String, KV<String, MyData>>writeDynamic()
        .by(KV::getKey)
        .via(Contextful.fn(KV::getValue), AvroIO.sink(MyData.class))
        .to(new GcsDynamicDestination())
        .withDestinationCoder(StringUtf8Coder.of()));

    // 运行Pipeline
    pipeline.run().waitUntilFinish();
  }

  // 自定义GCS目标路径
  public static class GcsDynamicDestination extends FileIO.DynamicDestinations<String, KV<String, MyData>> {
    @Override
    public String formatRecord(KV<String, MyData> element) {
      return element.getValue().toString();
    }

    @Override
    public String getDestination(String element) {
      return element;
    }

    @Override
    public FileIO.Write<String, KV<String, MyData>> getWriter(String destination) {
      ResourceId resourceId = FileSystems.matchNewResource(destination, true);
      return FileIO.<String, KV<String, MyData>>write()
          .via(Contextful.fn(KV::getValue), AvroIO.sink(MyData.class))
          .to(resourceId)
          .withDestinationCoder(StringUtf8Coder.of())
          .withNumShards(1)
          .withSuffix(".avro");
    }
  }

  // 自定义数据类型
  public static class MyData {
    // 定义数据属性
    // ...
  }
}