能否向Apache Beam PTransform传递侧输入？TensorFlow预处理咨询

当然可以给Apache Beam的PTransform传递侧输入！这正是处理你这种依赖全局计算值（比如样本总数）来动态配置处理逻辑的常用方案，刚好能解决你根据样本数量设置TFRecord分片数的需求。

核心结论 #

PTransform不仅支持传递侧输入，而且对于你这种场景，推荐将计算好的样本总数封装成PCollectionView（单值视图），通过自定义PTransform的构造参数传入，在expand方法中使用这个视图来动态计算分片数。

针对你的场景的具体实现步骤 #

我结合你的需求拆解成几个关键步骤，附带代码示例：

1. 先计算数据集的样本总数并转成单值视图

首先你需要先统计数据集中的样本总量，然后把这个全局值转换成PCollectionView——因为Beam是延迟执行的，只有通过视图才能在后续的Transform中安全获取这个计算结果：

// 假设你的原始样本数据是PCollection<YourSampleType> rawSamples
PCollection<Long> totalSamples = rawSamples.apply(Count.globally());
// 转成单值视图，方便后续侧输入使用
PCollectionView<Long> totalSamplesView = totalSamples.apply(View.asSingleton());

2. 自定义支持侧输入的PTransform来处理写入逻辑

接下来你可以写一个自定义的PTransform，把刚才的样本数视图作为构造参数传入，在expand方法里根据总数计算分片数，再完成TFRecord的写入：

public class DynamicTFRecordWriter extends PTransform<PCollection<YourSampleType>, PDone> {
    private final PCollectionView<Long> totalSamplesView;
    private final long targetRecordsPerShard; // 你期望每个分片的样本数

    // 通过构造函数传入侧输入视图和分片阈值
    public DynamicTFRecordWriter(PCollectionView<Long> totalSamplesView, long targetRecordsPerShard) {
        this.totalSamplesView = totalSamplesView;
        this.targetRecordsPerShard = targetRecordsPerShard;
    }

    @Override
    public PDone expand(PCollection<YourSampleType> input) {
        // 第一步：把样本转换成TFRecord字节
        PCollection<byte[]> tfRecordBytes = input.apply(ParDo.of(new DoFn<YourSampleType, byte[]>() {
            @ProcessElement
            public void process(ProcessContext ctx) {
                // 这里替换成你自己的样本转TFRecord逻辑
                byte[] tfBytes = convertSampleToTFRecord(ctx.element());
                ctx.output(tfBytes);
            }
        }));

        // 第二步：通过侧输入获取样本总数，计算目标分片数
        PCollectionView<Integer> targetShardsView = input.getPipeline()
                .apply(Create.of(null))
                .apply(ParDo.of(new DoFn<Void, Integer>() {
                    @SideInput
                    private final PCollectionView<Long> samplesView = totalSamplesView;

                    @ProcessElement
                    public void calculateShards(ProcessContext ctx) {
                        long total = ctx.sideInput(samplesView);
                        // 向上取整计算分片数，避免最后一个分片数据过少
                        int numShards = (int) Math.ceil((double) total / targetRecordsPerShard);
                        ctx.output(numShards);
                    }
                }))
                .apply(View.asSingleton());

        // 第三步：用动态计算的分片数写入TFRecord
        return tfRecordBytes.apply(FileIO.write()
                .to("gs://your-bucket/output/path") // 替换成你的输出路径
                .withSuffix(".tfrecord")
                .via(FileIO.sinkForByteArray())
                .withNumShards(targetShardsView)); // 这里传入动态计算的分片数视图
    }

    // 你的样本转TFRecord方法
    private byte[] convertSampleToTFRecord(YourSampleType sample) {
        // 实现你的转换逻辑，比如用TensorFlow的Example类构建
        return new byte[0];
    }
}

3. 在主流程中使用这个自定义Transform

最后在你的主Pipeline里调用这个Transform就行：

Pipeline pipeline = Pipeline.create(options);
PCollection<YourSampleType> rawSamples = pipeline.apply(ReadYourDataset());

// 调用自定义Writer，传入侧输入视图和每个分片的目标样本数
rawSamples.apply(new DynamicTFRecordWriter(totalSamplesView, 10000L)); // 每10000条一个分片

pipeline.run().waitUntilFinish();

关于你之前代码执行失败的可能原因 #

如果之前的代码抛出异常，大概率是这几个问题：

• 没有正确将全局计算值封装成PCollectionView，直接在PTransform构造阶段尝试获取未计算的值（Beam延迟执行，此时值还不存在）
• 在DoFn中使用侧输入时，没有用@SideInput注解声明
• 动态设置分片数时，没有用withNumShards(PCollectionView<Integer>)这个重载方法，而是直接传入了一个静态值

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