如何在Apache Beam Dataflow中获取文件行号并填充至POJO？

现有方法的问题说明 #

• 自定义带锁计数器：分布式环境下每个worker进程会持有独立的static变量副本，锁仅能保证单进程内线程安全，跨worker完全无法共享计数，导致行号重复或不连续。
• Beam Metrics计数器：仅用于作业监控统计，无法在数据处理流程中实时获取当前值并赋值给POJO，不符合业务数据赋值需求。

方案1：获取每个文件内的独立行号（常用场景） #

如果只需每个文件内部的行号（数据行从1开始，表头跳过），可利用Beam的State API为每个文件维护独立计数器，适配分布式处理场景。

代码实现 #

1. 读取文件并保留文件元数据：

PCollection<ReadableFile> files = pipeline.apply(FileIO.match().filepattern("path/to/your/files/*.txt"))
                                         .apply(FileIO.readMatches());

2. 自定义DoFn处理每行并维护文件内行号：

public class AddFileLineNumberDoFn extends DoFn<ReadableFile, BatchData> {
    // 以文件名为键，维护每个文件的行号计数器State
    private final StateSpec<String, ValueState<Integer>> fileLineCounter =
            StateSpecs.value(VarIntCoder.of());

    @ProcessElement
    public void processElement(@Element ReadableFile file, ProcessContext context,
                               @StateId("fileLineCounter") ValueState<Integer> counterState) throws IOException {
        try (BufferedReader reader = file.openBufferedReader()) {
            String line;
            // 初始化或读取当前文件已计数的行号
            int lineNum = counterState.read() != null ? counterState.read() : 0;
            
            // 跳过表头行
            reader.readLine();
            
            while ((line = reader.readLine()) != null) {
                lineNum++;
                // 解析行数据到BatchData并设置行号
                BatchData data = parseLine(line);
                data.setLineNumber(lineNum);
                context.output(data);
            }
            // 更新State，保存当前文件的最后行号
            counterState.write(lineNum);
        }
    }

    private BatchData parseLine(String line) {
        String[] parts = line.split("\\|");
        BatchData data = new BatchData();
        data.setTransactionId(Integer.parseInt(parts[0]));
        data.setClientId(Integer.parseInt(parts[1]));
        data.setProductCode(Integer.parseInt(parts[2]));
        return data;
    }
}

3. 应用DoFn处理数据：

PCollection<BatchData> batchData = files.apply(ParDo.of(new AddFileLineNumberDoFn()));

方案2：获取全局唯一行号（跨所有文件） #

如果需要所有文件的所有数据行从1开始连续递增的全局行号，可采用GlobalWindow + State方案，避免提前统计数据量带来的性能瓶颈。

代码实现 #

1. 自定义DoFn维护全局计数器：

public class AddGlobalLineNumberDoFn extends DoFn<String, BatchData> {
    // 维护全局行号计数器State
    private final StateSpec<String, ValueState<Long>> globalCounter =
            StateSpecs.value(VarLongCoder.of());

    @ProcessElement
    public void processElement(@Element String line, ProcessContext context,
                               @StateId("globalCounter") ValueState<Long> counterState) {
        // 跳过表头行
        if (line.startsWith("TRANSACTION_ID")) {
            return;
        }
        
        // 递增全局计数器
        long currentNum = counterState.read() != null ? counterState.read() : 0;
        currentNum++;
        counterState.write(currentNum);
        
        // 解析并设置行号
        BatchData data = parseLine(line);
        data.setLineNumber((int) currentNum);
        context.output(data);
    }

    private BatchData parseLine(String line) {
        String[] parts = line.split("\\|");
        BatchData data = new BatchData();
        data.setTransactionId(Integer.parseInt(parts[0]));
        data.setClientId(Integer.parseInt(parts[1]));
        data.setProductCode(Integer.parseInt(parts[2]));
        return data;
    }
}

2. 应用DoFn并指定GlobalWindow：

PCollection<BatchData> batchData = pipeline.apply(TextIO.read().from("path/to/your/files/*.txt"))
                                       .apply(Window.into(new GlobalWindows())
                                                  .triggering(AfterWatermark.pastEndOfWindow())
                                                  .discardingFiredPanes())
                                       .apply(ParDo.of(new AddGlobalLineNumberDoFn()));

方案选择建议 #

• 若业务只需文件内独立行号，优先选择方案1，性能更优且适配分布式场景；
• 若必须全局连续行号，选择方案2，相比提前统计数据量的方法，避免全局聚合带来的性能瓶颈。

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