Apache Beam中PCollection<KV<Long,Double>>聚合及窗口触发异常排查

问题描述 #

尝试在Apache Beam中对PCollection<KV<Long, Double>>进行聚合，需求是对所有不同的Key求和、对应的Value求和。代码设置了5秒固定窗口，触发条件为每个Pane至少2个元素，原数据共10个元素，期望输出5条聚合结果，但实际输出了10条原始KV数据。

相关代码 #

public class StreamPipelineBuilder {
    public void execute() {
        final List<UserTxn> txn = Utils.getUserTxnList().subList(0, 10);
        // create Pipeline
        final Pipeline pipeline = Pipeline.create();
        TestStream.Builder<KV<Long, UserTxn>>
                streamBuilder = TestStream.create(UserTxnKVCoder.of());
        // add all lines with timestamps to the TestStream
        final List<TimestampedValue<KV<Long, UserTxn>>> timestamped =
                txn.stream().map(i -> {
                    final KV<Long, UserTxn> kv = KV.of(i.getId(), i);
                    final LocalDateTime time = i.getTime();
                    final long millis = time.toInstant(ZoneOffset.UTC).toEpochMilli();
                    final Instant instant = new Instant(millis);
                    return TimestampedValue.of(kv, instant);
                }).collect(Collectors.toList());

        for (TimestampedValue<KV<Long, UserTxn>> value : timestamped) {
            streamBuilder = streamBuilder.addElements(value);
        }

        // create the unbounded PCollection from TestStream
        PCollection<KV<Long, UserTxn>> input = pipeline.apply(streamBuilder.advanceWatermarkToInfinity());
        PCollection<KV<Long, UserTxn>> windowed =
                input.apply(Window.<KV<Long, UserTxn>>into(FixedWindows.of(Duration.standardSeconds(5)))
                        .discardingFiredPanes()
                        .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(2)))
                        .withAllowedLateness(Duration.ZERO));

        PCollection<KV<Long, Double>> added = windowed.apply("aggregate", new PTransform<>() {
            @Override
            public PCollection<KV<Long, Double>> expand(PCollection<KV<Long, UserTxn>> input) {
                return input.apply(
                        MapElements.into(
                                TypeDescriptors.kvs(TypeDescriptors.longs(), TypeDescriptors.doubles())
                        ).via((record) -> KV.of(record.getKey(), record.getValue().getAmount()))
                ).apply(Combine.globally((SerializableFunction<Iterable<KV<Long, Double>>, KV<Long, Double>>) input1 -> {
                    AtomicLong keys = new AtomicLong();
                    AtomicDouble amounts = new AtomicDouble();
                    input1.forEach(e -> {
                        keys.addAndGet(e.getKey());
                        amounts.addAndGet(e.getValue());
                    });
                    return KV.of(keys.get(), amounts.get());
                }).withoutDefaults());

            }
        });

        added.apply(PrintPCollection.with());

        pipeline.run().waitUntilFinish();
    }
}

实际输出 #

[INFO] 2022-11-03 00:12:21.010 PrintPCollection - KV{1, 821.21}
[INFO] 2022-11-03 00:12:21.014 PrintPCollection - KV{6, 973.31}
[INFO] 2022-11-03 00:12:21.014 PrintPCollection - KV{8, 980.26}
[INFO] 2022-11-03 00:12:21.014 PrintPCollection - KV{4, 37.53}
[INFO] 2022-11-03 00:12:21.014 PrintPCollection - KV{2, 541.95}
[INFO] 2022-11-03 00:12:21.014 PrintPCollection - KV{7, 705.49}
[INFO] 2022-11-03 00:12:21.014 PrintPCollection - KV{3, 384.09}
[INFO] 2022-11-03 00:12:21.015 PrintPCollection - KV{9, 106.96}
[INFO] 2022-11-03 00:12:21.015 PrintPCollection - KV{5, 207.3}
[INFO] 2022-11-03 00:12:21.015 PrintPCollection - KV{10, 675.48}

问题原因与修复方案 #

核心问题1：全局聚合的类型处理错误 #

你使用Combine.globally()时，通过强制类型转换将lambda转为SerializableFunction，这会因Java类型擦除导致Beam无法正确识别聚合逻辑，最终跳过聚合步骤，直接输出MapElements转换后的原始元素。

核心问题2：窗口触发与TestStream的元素注入逻辑不匹配 #

你通过streamBuilder.advanceWatermarkToInfinity()一次性注入所有10个元素，导致Watermark直接推进到无穷大，窗口被立即关闭。此时所有元素会被放入同一个Pane中，即使设置了AfterPane.elementCountAtLeast(2)的触发条件，也只会触发一次聚合，而非期望的5次。

修复方案 #

方案1：修复全局聚合的类型问题（针对“所有Key求和、所有Value求和”的需求）

使用显式的CombineFn替代lambda强制转换，避免类型擦除问题：

PCollection<KV<Long, Double>> added = windowed.apply("aggregate", new PTransform<>() {
    @Override
    public PCollection<KV<Long, Double>> expand(PCollection<KV<Long, UserTxn>> input) {
        return input.apply(
                MapElements.into(
                        TypeDescriptors.kvs(TypeDescriptors.longs(), TypeDescriptors.doubles())
                ).via((record) -> KV.of(record.getKey(), record.getValue().getAmount()))
        ).apply(Combine.globally(new CombineFn<KV<Long, Double>, CombineFn.Accumulator<AtomicLong, AtomicDouble>, KV<Long, Double>>() {
            @Override
            public Accumulator createAccumulator() {
                return new Accumulator(new AtomicLong(), new AtomicDouble());
            }

            @Override
            public Accumulator addInput(Accumulator accumulator, KV<Long, Double> input) {
                accumulator.keys.addAndGet(input.getKey());
                accumulator.amounts.addAndGet(input.getValue());
                return accumulator;
            }

            @Override
            public Accumulator mergeAccumulators(Iterable<Accumulator> accumulators) {
                Accumulator result = createAccumulator();
                for (Accumulator acc : accumulators) {
                    result.keys.addAndGet(acc.keys.get());
                    result.amounts.addAndGet(acc.amounts.get());
                }
                return result;
            }

            @Override
            public KV<Long, Double> extractOutput(Accumulator accumulator) {
                return KV.of(accumulator.keys.get(), accumulator.amounts.get());
            }

            static class Accumulator {
                AtomicLong keys;
                AtomicDouble amounts;

                Accumulator(AtomicLong keys, AtomicDouble amounts) {
                    this.keys = keys;
                    this.amounts = amounts;
                }
            }
        }).withoutDefaults());
    }
});

方案2：调整TestStream的元素注入逻辑，实现分批触发

将元素分批注入TestStream，并逐步推进Watermark，确保每个触发条件（2个元素）被满足：

// 替换原有的TestStream构建逻辑
TestStream.Builder<KV<Long, UserTxn>> streamBuilder = TestStream.create(UserTxnKVCoder.of());
// 分批添加元素，每2个元素后推进一次Watermark（确保窗口触发）
for (int i = 0; i < timestamped.size(); i += 2) {
    // 添加2个元素
    streamBuilder = streamBuilder.addElements(timestamped.get(i), timestamped.get(i+1));
    // 推进Watermark到当前最后一个元素的时间戳+1ms，触发窗口检查
    Instant nextWatermark = timestamped.get(i+1).getTimestamp().plus(Duration.millis(1));
    streamBuilder = streamBuilder.advanceWatermarkTo(nextWatermark);
}
// 最后推进Watermark到无穷大，关闭窗口
streamBuilder = streamBuilder.advanceWatermarkToInfinity();

方案3：如果需求是按Key分组求和（而非全局聚合）

如果你的实际需求是按Key分组，对每个Key的Value求和，则应使用Combine.perKey()替代Combine.globally()：

PCollection<KV<Long, Double>> added = windowed.apply("aggregate", new PTransform<>() {
    @Override
    public PCollection<KV<Long, Double>> expand(PCollection<KV<Long, UserTxn>> input) {
        return input.apply(
                MapElements.into(
                        TypeDescriptors.kvs(TypeDescriptors.longs(), TypeDescriptors.doubles())
                ).via((record) -> KV.of(record.getKey(), record.getValue().getAmount()))
        ).apply(Combine.perKey(Sum.ofDoubles()));
    }
});

内容的提问来源于stack exchange，提问作者Frederick Álvarez