From a5f946f4e960f53595b0d56bff1fce57632f77aa Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Fri, 10 Jul 2026 01:24:00 +0500 Subject: [PATCH 1/7] [GSoC 2026] Kafka Streams runner: Flatten support Add a translator for Beam's Flatten primitive (beam:transform:flatten:v1), the union of N input PCollections into one. - FlattenProcessor forwards data straight through and owns its output watermark the way GroupByKey does: it runs a WatermarkManager over its input branches and emits its own single-source (0 of 1) watermark only when the min() across them advances, holding until every branch has drained so a downstream GroupByKey does not fire early. - The branch identity (i of N) the WatermarkManager needs is stamped upstream by the producing ExecutableStage when its output feeds a Flatten -- a translation pre-pass records which PCollections are Flatten inputs -- because Kafka Streams does not tell a processor which parent forwarded a record. Producers whose output does not feed a Flatten keep reporting as the single source (0 of 1). Only ExecutableStage producers stamp the branch identity so far, which covers the PAssert GroupGlobally shape the ValidatesRunner tests use. A Read/Impulse/GBK output feeding a Flatten directly is a follow-up. Test: FlattenTest unions two Create -> ParDo branches and asserts a downstream ParDo sees every element from both. --- .../translation/ExecutableStageProcessor.java | 26 +++-- .../ExecutableStageTranslator.java | 14 ++- .../streams/translation/FlattenProcessor.java | 97 +++++++++++++++++++ .../translation/FlattenTranslator.java | 60 ++++++++++++ .../KafkaStreamsPipelineTranslator.java | 24 +++++ .../KafkaStreamsTranslationContext.java | 49 ++++++++++ ...ExecutableStageProcessorWatermarkTest.java | 3 +- .../streams/translation/FlattenTest.java | 90 +++++++++++++++++ 8 files changed, 355 insertions(+), 8 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java index 00f85032d041..f2e355638a0a 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java @@ -70,6 +70,11 @@ class ExecutableStageProcessor private final RunnerApi.ExecutableStagePayload stagePayload; private final JobInfo jobInfo; + // The source identity this stage stamps on its forwarded watermark. Normally (0 of 1), a single + // source to the next stage; but when this stage's output feeds a Flatten it is that branch's + // (i of N), so the Flatten's WatermarkManager can tell the branches apart. See FlattenProcessor. + private final int sourcePartition; + private final int totalPartitions; // pendingOutputs is enqueued by SDK harness threads (inside the OutputReceiverFactory callback) // and drained by the Kafka Streams processing thread on bundle close; needs to be thread-safe. @@ -90,9 +95,15 @@ class ExecutableStageProcessor private @Nullable StageBundleFactory stageBundleFactory; private @Nullable RemoteBundle currentBundle; - ExecutableStageProcessor(RunnerApi.ExecutableStagePayload stagePayload, JobInfo jobInfo) { + ExecutableStageProcessor( + RunnerApi.ExecutableStagePayload stagePayload, + JobInfo jobInfo, + int sourcePartition, + int totalPartitions) { this.stagePayload = stagePayload; this.jobInfo = jobInfo; + this.sourcePartition = sourcePartition; + this.totalPartitions = totalPartitions; } @Override @@ -203,14 +214,17 @@ private void closeBundleAndFlush(Record> record) { } private void forwardWatermark(Record> record, long watermarkMillis) { - // This stage is a single instance for now, so it forwards its watermark as the only source - // partition (0 of 1). Fanning the watermark out to every downstream partition — and producing - // it atomically with the offset commit so it is durable — lands with the topic-based shuffle - // work, when there are real source partitions to track (#18479). + // This stage is a single instance for now, so downstream it is one source — reported as + // (sourcePartition of totalPartitions), which is (0 of 1) unless its output feeds a Flatten, in + // which case it is that branch's identity so the Flatten can tell its branches apart. Fanning + // the watermark out to every downstream partition — and producing it atomically with the offset + // commit so it is durable — lands with the topic-based shuffle work (#18479). ProcessorContext> ctx = checkInitialized(context); ctx.forward( new Record>( - record.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1), record.timestamp())); + record.key(), + KStreamsPayload.watermark(watermarkMillis, sourcePartition, totalPartitions), + record.timestamp())); } @Override diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java index dc56d57f57c2..89ef345dbc61 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java @@ -82,10 +82,22 @@ public void translate( String inputPCollectionId = stagePayload.getInput(); String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); + // A stage has at most one output (multi-output rejected above). Stamp its watermark with that + // output's Flatten branch identity if it feeds one, else the default single source (0 of 1). + KafkaStreamsTranslationContext.SourceStamp stamp = + transform.getOutputsMap().isEmpty() + ? context.getSourceStamp("") + : context.getSourceStamp(Iterables.getOnlyElement(transform.getOutputsMap().values())); + Topology topology = context.getTopology(); topology.addProcessor( transformId, - () -> new ExecutableStageProcessor(stagePayload, context.getJobInfo()), + () -> + new ExecutableStageProcessor( + stagePayload, + context.getJobInfo(), + stamp.getSourcePartition(), + stamp.getTotalPartitions()), parentProcessor); if (!transform.getOutputsMap().isEmpty()) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java new file mode 100644 index 000000000000..0b3fa5031297 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java @@ -0,0 +1,97 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.kafka.streams.translation; + +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * Kafka Streams {@link Processor} implementing Beam's {@code Flatten} primitive ({@code + * beam:transform:flatten:v1}): the union of N input PCollections into one output PCollection. + * + *

Data records are forwarded straight through unchanged — the merge of the N parents' + * data streams is the flatten. + * + *

Watermark reports are where Flatten does real work, and it owns its output watermark + * the same way GroupByKey does: it runs a {@link WatermarkManager} over its inputs, forwards its + * own watermark only when the {@code min()} across them advances, and stamps that as a single + * source ({@code 0 of 1}) to its downstream. This holds the output watermark back until + * every input branch has reported, so a downstream GroupByKey does not fire before all + * flattened branches are drained. + * + *

The {@link WatermarkManager} tells the N input branches apart by the {@code (sourcePartition, + * totalSourcePartitions)} carried in each watermark payload. Because Kafka Streams does not tell a + * processor which parent forwarded a record, the branch identity {@code (i of N)} is stamped + * upstream — by the parent transform that produces the branch — so the reports arriving here + * already carry distinct partitions. (Every parent stamping the same {@code 0 of 1} would collide + * and release the watermark too early.) + */ +class FlattenProcessor + implements Processor, byte[], KStreamsPayload> { + + // Computes the output watermark as min() over the input branches, holding until every branch has + // reported (see WatermarkManager). Flatten is a single instance for now. + private final WatermarkManager watermarkManager = new WatermarkManager(); + // The last watermark actually forwarded downstream, so we only forward when it advances. + private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + + private @Nullable ProcessorContext> context; + + @Override + public void init(ProcessorContext> context) { + this.context = context; + } + + @Override + public void process(Record> record) { + KStreamsPayload payload = record.value(); + ProcessorContext> ctx = checkInitialized(context); + if (!payload.isWatermark()) { + // Data: the union of the parents' data streams is the flatten — forward unchanged. + ctx.forward(record); + return; + } + WatermarkPayload report = payload.asWatermark(); + watermarkManager.observe( + report.getSourcePartition(), + new Instant(report.getWatermarkMillis()), + report.getTotalSourcePartitions()); + Instant advanced = watermarkManager.advance(); + if (advanced.isAfter(lastForwardedWatermark)) { + lastForwardedWatermark = advanced; + // Flatten is one logical source to the next stage; report it as partition 0 of 1. + ctx.forward( + new Record>( + record.key(), + KStreamsPayload.watermark(advanced.getMillis(), 0, 1), + record.timestamp())); + } + } + + private static ProcessorContext> checkInitialized( + @Nullable ProcessorContext> context) { + if (context == null) { + throw new IllegalStateException("FlattenProcessor used before init()"); + } + return context; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java new file mode 100644 index 000000000000..864fb3adc558 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.kafka.streams.translation; + +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.kafka.streams.Topology; + +/** + * Translates Beam's {@code Flatten} primitive ({@code beam:transform:flatten:v1}): the union of N + * input PCollections into one output PCollection. + * + *

Wires a single {@link FlattenProcessor} node to the producer of every input PCollection (Kafka + * Streams lets a processor have many parents), so the parents' data streams merge into it, and + * registers it as the producer of the flattened output so downstream translators wire to it. The + * processor forwards data through and owns its output watermark via a {@link WatermarkManager}, + * exactly like GroupByKey — see {@link FlattenProcessor}. + * + *

The {@code (i of N)} branch identity the {@link WatermarkManager} needs is stamped upstream, + * by the parent transform that produces each branch, because Kafka Streams does not tell a + * processor which parent forwarded a record. + */ +class FlattenTranslator implements PTransformTranslator { + + @Override + public void translate( + String transformId, RunnerApi.Pipeline pipeline, KafkaStreamsTranslationContext context) { + RunnerApi.PTransform transform = pipeline.getComponents().getTransformsOrThrow(transformId); + // Flatten produces exactly one output PCollection, fed by all of its input PCollections. + String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + + List parentProcessors = new ArrayList<>(); + for (String inputPCollectionId : transform.getInputsMap().values()) { + parentProcessors.add(context.getProcessorNameForPCollection(inputPCollectionId)); + } + + Topology topology = context.getTopology(); + topology.addProcessor( + transformId, FlattenProcessor::new, parentProcessors.toArray(new String[0])); + + context.registerPCollectionProducer(outputPCollectionId, transformId); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index 943673dfe98c..a02c9620cc99 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -52,6 +52,7 @@ public KafkaStreamsPipelineTranslator() { .put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator()) .put(PTransformTranslation.READ_TRANSFORM_URN, new ReadTranslator()) .put(PTransformTranslation.REDISTRIBUTE_ARBITRARILY_URN, new RedistributeTranslator()) + .put(PTransformTranslation.FLATTEN_TRANSFORM_URN, new FlattenTranslator()) .put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new GroupByKeyTranslator()) .put(ExecutableStage.URN, new ExecutableStageTranslator()) .build()); @@ -106,6 +107,7 @@ public RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) { * Throws {@link UnsupportedOperationException} on the first unsupported URN. */ public void translate(KafkaStreamsTranslationContext context, RunnerApi.Pipeline pipeline) { + registerFlattenSourceStamps(context, pipeline); QueryablePipeline queryable = QueryablePipeline.forTransforms( pipeline.getRootTransformIdsList(), pipeline.getComponents()); @@ -126,6 +128,28 @@ public void translate(KafkaStreamsTranslationContext context, RunnerApi.Pipeline } } + /** + * Records each Flatten input PCollection's branch identity {@code (i of N)} with the context, so + * the producer of that PCollection stamps its watermark with a distinct source partition. Without + * this every branch would report as the single source {@code (0 of 1)} and the Flatten's {@link + * WatermarkManager} could not tell them apart, releasing its watermark before every branch + * drains. + */ + private static void registerFlattenSourceStamps( + KafkaStreamsTranslationContext context, RunnerApi.Pipeline pipeline) { + for (RunnerApi.PTransform transform : pipeline.getComponents().getTransformsMap().values()) { + if (!PTransformTranslation.FLATTEN_TRANSFORM_URN.equals(transform.getSpec().getUrn())) { + continue; + } + int totalPartitions = transform.getInputsMap().size(); + int sourcePartition = 0; + for (String inputPCollectionId : transform.getInputsMap().values()) { + context.registerFlattenSourceStamp(inputPCollectionId, sourcePartition, totalPartitions); + sourcePartition++; + } + } + } + /** * Tells the SDK that URNs handled directly by the Kafka Streams runner should be treated as * primitives by {@link QueryablePipeline}. Mirrors Flink's {@code IsFlinkNativeTransform} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java index 0a045ff7395b..0c943319620c 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java @@ -41,10 +41,18 @@ public class KafkaStreamsTranslationContext { /** Characters not legal in a Kafka topic name; a topic's legal set is {@code [a-zA-Z0-9._-]}. */ private static final Pattern ILLEGAL_TOPIC_CHARS = Pattern.compile("[^a-zA-Z0-9._-]"); + /** The watermark source stamp a producer reports for an output not consumed by a Flatten. */ + private static final SourceStamp SINGLE_SOURCE = new SourceStamp(0, 1); + private final JobInfo jobInfo; private final KafkaStreamsPipelineOptions pipelineOptions; private final Topology topology; private final Map pCollectionIdToProcessorName; + // PCollection id -> the (sourcePartition, totalPartitions) its producer stamps its watermark + // with. + // Only PCollections consumed by a Flatten get a non-default entry (their branch index of the + // Flatten's fan-in); everything else reports as the single source (0 of 1). See getSourceStamp. + private final Map pCollectionIdToSourceStamp; public static KafkaStreamsTranslationContext create( JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { @@ -62,6 +70,7 @@ private KafkaStreamsTranslationContext( this.pipelineOptions = pipelineOptions; this.topology = topology; this.pCollectionIdToProcessorName = new HashMap<>(); + this.pCollectionIdToSourceStamp = new HashMap<>(); } public JobInfo getJobInfo() { @@ -120,4 +129,44 @@ public String getReadBootstrapTopic(String transformId) { + "_" + sanitizedTransformId; } + + /** + * Records that the given PCollection is the {@code sourcePartition}-th of {@code totalPartitions} + * input branches of a Flatten, so its producer stamps its watermark with that identity instead of + * the default single source. Lets the Flatten's {@link WatermarkManager} tell its branches apart + * (Kafka Streams does not tell a processor which parent forwarded a record). + */ + public void registerFlattenSourceStamp( + String pCollectionId, int sourcePartition, int totalPartitions) { + pCollectionIdToSourceStamp.put( + pCollectionId, new SourceStamp(sourcePartition, totalPartitions)); + } + + /** + * Returns the watermark source stamp a producer should report for the given output PCollection: + * its Flatten branch identity if it feeds a Flatten (see {@link #registerFlattenSourceStamp}), or + * the single source {@code (0 of 1)} otherwise. + */ + public SourceStamp getSourceStamp(String pCollectionId) { + return pCollectionIdToSourceStamp.getOrDefault(pCollectionId, SINGLE_SOURCE); + } + + /** A watermark report's source identity: partition {@code sourcePartition} of {@code total}. */ + public static final class SourceStamp { + private final int sourcePartition; + private final int totalPartitions; + + SourceStamp(int sourcePartition, int totalPartitions) { + this.sourcePartition = sourcePartition; + this.totalPartitions = totalPartitions; + } + + public int getSourcePartition() { + return sourcePartition; + } + + public int getTotalPartitions() { + return totalPartitions; + } + } } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java index fc5797a12c26..c1425d94dec9 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java @@ -46,8 +46,9 @@ private static ExecutableStageProcessor newProcessor() { "job-name", "", PipelineOptionsTranslation.toProto(PipelineOptionsFactory.create())); + // Forwards its watermark as the single source (0 of 1); this test exercises the consume side. return new ExecutableStageProcessor( - RunnerApi.ExecutableStagePayload.getDefaultInstance(), jobInfo); + RunnerApi.ExecutableStagePayload.getDefaultInstance(), jobInfo, 0, 1); } private static Record> watermark( diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java new file mode 100644 index 000000000000..8e9dbbb634c2 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.kafka.streams.translation; + +import static org.hamcrest.CoreMatchers.hasItems; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.util.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.junit.Test; + +/** + * End-to-end test for {@link FlattenTranslator}: two branches are flattened into one PCollection + * and a recording ParDo sees every element from both. + * + *

Each branch is a {@code Create -> identity ParDo}, so its producer feeding the Flatten is an + * {@link ExecutableStageProcessor} — the same shape PAssert's {@code GroupGlobally} produces. This + * is the case the {@code (i of N)} watermark stamping exists for: without it both branches would + * report as the single source {@code (0 of 1)}, the Flatten would release its watermark after the + * first branch drained, and the downstream stage's bundle would close early and drop the second + * branch's elements. + */ +public class FlattenTest { + + private static class IdentityFn extends DoFn { + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver out) { + out.output(input); + } + } + + /** Records every element the harness feeds it so the test can assert the flatten's union. */ + private static class RecordingFn extends DoFn { + private final SharedTestCollector collector; + + RecordingFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver out) { + collector.record(input); + out.output(input); + } + } + + @Test + public void flattenUnionsEveryBranch() { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection a = + pipeline.apply("createA", Create.of(1, 2)).apply("idA", ParDo.of(new IdentityFn())); + PCollection b = + pipeline.apply("createB", Create.of(3, 4)).apply("idB", ParDo.of(new IdentityFn())); + PCollectionList.of(a) + .and(b) + .apply("flatten", Flatten.pCollections()) + .apply("record", ParDo.of(new RecordingFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List recorded = collector.recorded(); + assertThat(recorded.size(), is(4)); + assertThat(recorded, hasItems(1, 2, 3, 4)); + } + } +} From 1784261561fa5c0b57d2624503f16154cf8449a6 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Fri, 10 Jul 2026 01:44:21 +0500 Subject: [PATCH 2/7] Address review: fail fast when a PCollection feeds a Flatten more than once A PCollection flattened with itself (Flatten.of(pc, pc)) or consumed by two Flattens would need a distinct watermark source per branch, but its single producer can only stamp one identity, so the branch watermark would get stuck. registerFlattenSourceStamp now throws UnsupportedOperationException on the second registration rather than silently overwriting. Deduping is not an option: a self-flatten must emit its input twice (bag union), so dropping the copy would lose data. Also sort the Flatten input PCollection ids so each branch index is assigned deterministically. Adds a test that a self-flatten is rejected. --- .../KafkaStreamsPipelineTranslator.java | 13 +++++++++++-- .../KafkaStreamsTranslationContext.java | 16 ++++++++++++++-- .../kafka/streams/translation/FlattenTest.java | 18 ++++++++++++++++++ 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index a02c9620cc99..2c3d6861573f 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -18,6 +18,9 @@ package org.apache.beam.runners.kafka.streams.translation; import com.google.auto.service.AutoService; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.Map; import java.util.Set; import org.apache.beam.model.pipeline.v1.RunnerApi; @@ -141,9 +144,15 @@ private static void registerFlattenSourceStamps( if (!PTransformTranslation.FLATTEN_TRANSFORM_URN.equals(transform.getSpec().getUrn())) { continue; } - int totalPartitions = transform.getInputsMap().size(); + // Sort the input PCollection ids so each branch's index is assigned deterministically. Not + // required for correctness — a Flatten and its producers sit in one task, so the indices only + // have to agree within a single topology build — but a stable assignment is easier to reason + // about and reproduce. registerFlattenSourceStamp fails fast on a duplicate (self-flatten). + List inputPCollectionIds = new ArrayList<>(transform.getInputsMap().values()); + Collections.sort(inputPCollectionIds); + int totalPartitions = inputPCollectionIds.size(); int sourcePartition = 0; - for (String inputPCollectionId : transform.getInputsMap().values()) { + for (String inputPCollectionId : inputPCollectionIds) { context.registerFlattenSourceStamp(inputPCollectionId, sourcePartition, totalPartitions); sourcePartition++; } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java index 0c943319620c..c0fa914e3803 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java @@ -138,8 +138,20 @@ public String getReadBootstrapTopic(String transformId) { */ public void registerFlattenSourceStamp( String pCollectionId, int sourcePartition, int totalPartitions) { - pCollectionIdToSourceStamp.put( - pCollectionId, new SourceStamp(sourcePartition, totalPartitions)); + SourceStamp existing = + pCollectionIdToSourceStamp.putIfAbsent( + pCollectionId, new SourceStamp(sourcePartition, totalPartitions)); + if (existing != null) { + // A PCollection feeding a Flatten more than once — the same Flatten twice, or two different + // Flattens — would need a distinct watermark source per branch, but its single producer can + // only stamp one identity. Fail fast rather than get the watermark stuck or silently drop the + // duplicate branch. + throw new UnsupportedOperationException( + "PCollection " + + pCollectionId + + " feeds a Flatten more than once (the same Flatten twice, or multiple Flattens);" + + " this is not yet supported by the Kafka Streams runner."); + } } /** diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java index 8e9dbbb634c2..c595cbfdb5d2 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java @@ -20,6 +20,7 @@ import static org.hamcrest.CoreMatchers.hasItems; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertThrows; import java.util.List; import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; @@ -87,4 +88,21 @@ public void flattenUnionsEveryBranch() { assertThat(recorded, hasItems(1, 2, 3, 4)); } } + + @Test + public void flattenOfOneBranchTwiceIsRejected() { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection branch = + pipeline.apply("create", Create.of(1, 2)).apply("id", ParDo.of(new IdentityFn())); + // Flattening a branch with itself needs a distinct watermark source per branch, but the branch + // has a single producer that can only stamp one identity, so the runner rejects it (rather than + // getting the watermark stuck or silently dropping the duplicate copy). + PCollectionList.of(branch) + .and(branch) + .apply("flatten", Flatten.pCollections()) + .apply("sink", ParDo.of(new IdentityFn())); + + assertThrows( + UnsupportedOperationException.class, () -> KafkaStreamsTestRunner.translate(pipeline)); + } } From 445b236a3eec4e935a70b1c3874c25add55b63c4 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Sun, 12 Jul 2026 16:20:17 +0500 Subject: [PATCH 3/7] Address review: number producers globally so a PCollection can feed two Flattens je-ik found that per-Flatten (i of N) branch numbering breaks when a PCollection feeds two Flattens (input2 is "1 of 2" for one and "0 of 2" for the other) -- its single producer cannot stamp two identities, and the previous guard wrongly rejected this valid pipeline. Number producers instead: each Flatten-input PCollection gets one stable global id, its producer stamps that id (always as a single source, 1 of 1), and each Flatten holds its watermark using its own input count. A shared input reports one id and every Flatten still waits only for its own branches. This also removes the fan-out hazard (a single-input consumer always sees "1 of 1"). Self-flatten (Flatten.of(pc, pc)) stays rejected -- one producer cannot be two branches. Adds a test that a PCollection feeding two Flattens produces both unions. --- .../ExecutableStageTranslator.java | 18 ++--- .../streams/translation/FlattenProcessor.java | 28 ++++--- .../translation/FlattenTranslator.java | 16 +++- .../KafkaStreamsPipelineTranslator.java | 37 ++++++---- .../KafkaStreamsTranslationContext.java | 74 ++++++------------- .../streams/translation/FlattenTest.java | 33 +++++++++ 6 files changed, 117 insertions(+), 89 deletions(-) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java index 89ef345dbc61..d6e7df99537c 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java @@ -83,21 +83,21 @@ public void translate( String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); // A stage has at most one output (multi-output rejected above). Stamp its watermark with that - // output's Flatten branch identity if it feeds one, else the default single source (0 of 1). - KafkaStreamsTranslationContext.SourceStamp stamp = + // output's global producer id (0 if no Flatten consumes it), always as a single source (1 of + // 1): + // a downstream Flatten tells its branches apart by this id and holds using its own input count, + // while a single-input consumer just sees one source. + int watermarkSourceId = transform.getOutputsMap().isEmpty() - ? context.getSourceStamp("") - : context.getSourceStamp(Iterables.getOnlyElement(transform.getOutputsMap().values())); + ? 0 + : context.getProducerWatermarkId( + Iterables.getOnlyElement(transform.getOutputsMap().values())); Topology topology = context.getTopology(); topology.addProcessor( transformId, () -> - new ExecutableStageProcessor( - stagePayload, - context.getJobInfo(), - stamp.getSourcePartition(), - stamp.getTotalPartitions()), + new ExecutableStageProcessor(stagePayload, context.getJobInfo(), watermarkSourceId, 1), parentProcessor); if (!transform.getOutputsMap().isEmpty()) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java index 0b3fa5031297..84ea47eb0726 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java @@ -38,12 +38,13 @@ * every input branch has reported, so a downstream GroupByKey does not fire before all * flattened branches are drained. * - *

The {@link WatermarkManager} tells the N input branches apart by the {@code (sourcePartition, - * totalSourcePartitions)} carried in each watermark payload. Because Kafka Streams does not tell a - * processor which parent forwarded a record, the branch identity {@code (i of N)} is stamped - * upstream — by the parent transform that produces the branch — so the reports arriving here - * already carry distinct partitions. (Every parent stamping the same {@code 0 of 1} would collide - * and release the watermark too early.) + *

The {@link WatermarkManager} tells the input branches apart by the global producer id each + * branch's producer stamps on its watermark (Kafka Streams does not tell a processor which parent + * forwarded a record). Each producer stamps its own stable id regardless of who consumes it, so a + * PCollection feeding several Flattens reports one id and every Flatten still waits only for its + * own inputs. The Flatten holds its watermark using its own input count — passed in at construction + * — not the reported total (producers always report the single source {@code 1 of 1}); that is what + * lets an input shared by two Flattens work. */ class FlattenProcessor implements Processor, byte[], KStreamsPayload> { @@ -51,11 +52,19 @@ class FlattenProcessor // Computes the output watermark as min() over the input branches, holding until every branch has // reported (see WatermarkManager). Flatten is a single instance for now. private final WatermarkManager watermarkManager = new WatermarkManager(); + // How many input branches feed this Flatten. The watermark is held until every branch's producer + // has reported, so this is the count the WatermarkManager waits for — not the reported total, + // which is always 1 because a producer stamps its watermark without regard to who consumes it. + private final int expectedInputCount; // The last watermark actually forwarded downstream, so we only forward when it advances. private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; private @Nullable ProcessorContext> context; + FlattenProcessor(int expectedInputCount) { + this.expectedInputCount = expectedInputCount; + } + @Override public void init(ProcessorContext> context) { this.context = context; @@ -71,10 +80,11 @@ public void process(Record> record) { return; } WatermarkPayload report = payload.asWatermark(); + // Key on the producer id the branch stamped, but wait for this Flatten's own input count — not + // the report's total (always 1) — so a producer shared with another Flatten reports one id yet + // each Flatten still holds until all of its own branches arrive. watermarkManager.observe( - report.getSourcePartition(), - new Instant(report.getWatermarkMillis()), - report.getTotalSourcePartitions()); + report.getSourcePartition(), new Instant(report.getWatermarkMillis()), expectedInputCount); Instant advanced = watermarkManager.advance(); if (advanced.isAfter(lastForwardedWatermark)) { lastForwardedWatermark = advanced; diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java index 864fb3adc558..18a6cbae4a46 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java @@ -33,9 +33,11 @@ * processor forwards data through and owns its output watermark via a {@link WatermarkManager}, * exactly like GroupByKey — see {@link FlattenProcessor}. * - *

The {@code (i of N)} branch identity the {@link WatermarkManager} needs is stamped upstream, - * by the parent transform that produces each branch, because Kafka Streams does not tell a - * processor which parent forwarded a record. + *

The producer id the {@link WatermarkManager} keys each branch on is stamped upstream, by the + * branch's producing transform, because Kafka Streams does not tell a processor which parent + * forwarded a record. Ids are assigned once per Flatten-input PCollection (see {@link + * KafkaStreamsPipelineTranslator}), so an input shared by two Flattens reports one id and each + * Flatten still waits only for its own branches. */ class FlattenTranslator implements PTransformTranslator { @@ -51,9 +53,15 @@ public void translate( parentProcessors.add(context.getProcessorNameForPCollection(inputPCollectionId)); } + // The Flatten holds its watermark until all of its input branches report. Inputs are distinct + // (a self-flatten is rejected during producer-id assignment), so the input count is the number + // of producer ids to wait for. + int inputCount = transform.getInputsMap().size(); Topology topology = context.getTopology(); topology.addProcessor( - transformId, FlattenProcessor::new, parentProcessors.toArray(new String[0])); + transformId, + () -> new FlattenProcessor(inputCount), + parentProcessors.toArray(new String[0])); context.registerPCollectionProducer(outputPCollectionId, transformId); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index 2c3d6861573f..bd63a04e1b45 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -20,6 +20,7 @@ import com.google.auto.service.AutoService; import java.util.ArrayList; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; @@ -110,7 +111,7 @@ public RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) { * Throws {@link UnsupportedOperationException} on the first unsupported URN. */ public void translate(KafkaStreamsTranslationContext context, RunnerApi.Pipeline pipeline) { - registerFlattenSourceStamps(context, pipeline); + assignFlattenProducerIds(context, pipeline); QueryablePipeline queryable = QueryablePipeline.forTransforms( pipeline.getRootTransformIdsList(), pipeline.getComponents()); @@ -132,29 +133,35 @@ public void translate(KafkaStreamsTranslationContext context, RunnerApi.Pipeline } /** - * Records each Flatten input PCollection's branch identity {@code (i of N)} with the context, so - * the producer of that PCollection stamps its watermark with a distinct source partition. Without - * this every branch would report as the single source {@code (0 of 1)} and the Flatten's {@link - * WatermarkManager} could not tell them apart, releasing its watermark before every branch - * drains. + * Assigns each Flatten input PCollection a stable global producer id, so the producer of that + * PCollection stamps its watermark with that id and each downstream Flatten can tell its input + * branches apart. A PCollection shared by several Flattens keeps one id, so every Flatten still + * waits only for its own inputs' ids — this is what makes a PCollection feeding two Flattens work + * (je-ik's review of #39273). A self-flatten (the same PCollection listed as an input twice) is + * rejected: its single producer cannot be two branches, and Kafka Streams cannot wire the same + * parent to a child twice, so the duplicate copy would be silently dropped. */ - private static void registerFlattenSourceStamps( + private static void assignFlattenProducerIds( KafkaStreamsTranslationContext context, RunnerApi.Pipeline pipeline) { for (RunnerApi.PTransform transform : pipeline.getComponents().getTransformsMap().values()) { if (!PTransformTranslation.FLATTEN_TRANSFORM_URN.equals(transform.getSpec().getUrn())) { continue; } - // Sort the input PCollection ids so each branch's index is assigned deterministically. Not - // required for correctness — a Flatten and its producers sit in one task, so the indices only - // have to agree within a single topology build — but a stable assignment is easier to reason - // about and reproduce. registerFlattenSourceStamp fails fast on a duplicate (self-flatten). + // Sorted so the numbering is deterministic across topology builds. List inputPCollectionIds = new ArrayList<>(transform.getInputsMap().values()); Collections.sort(inputPCollectionIds); - int totalPartitions = inputPCollectionIds.size(); - int sourcePartition = 0; + Set seenInputs = new HashSet<>(); for (String inputPCollectionId : inputPCollectionIds) { - context.registerFlattenSourceStamp(inputPCollectionId, sourcePartition, totalPartitions); - sourcePartition++; + if (!seenInputs.add(inputPCollectionId)) { + throw new UnsupportedOperationException( + "Flatten " + + transform.getUniqueName() + + " has PCollection " + + inputPCollectionId + + " as an input more than once; a self-flatten is not yet supported by the Kafka" + + " Streams runner."); + } + context.assignFlattenProducerId(inputPCollectionId); } } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java index c0fa914e3803..506bc7342b1f 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java @@ -41,18 +41,21 @@ public class KafkaStreamsTranslationContext { /** Characters not legal in a Kafka topic name; a topic's legal set is {@code [a-zA-Z0-9._-]}. */ private static final Pattern ILLEGAL_TOPIC_CHARS = Pattern.compile("[^a-zA-Z0-9._-]"); - /** The watermark source stamp a producer reports for an output not consumed by a Flatten. */ - private static final SourceStamp SINGLE_SOURCE = new SourceStamp(0, 1); + /** The watermark source id a producer reports for an output no Flatten consumes. */ + private static final int SINGLE_SOURCE_ID = 0; private final JobInfo jobInfo; private final KafkaStreamsPipelineOptions pipelineOptions; private final Topology topology; private final Map pCollectionIdToProcessorName; - // PCollection id -> the (sourcePartition, totalPartitions) its producer stamps its watermark - // with. - // Only PCollections consumed by a Flatten get a non-default entry (their branch index of the - // Flatten's fan-in); everything else reports as the single source (0 of 1). See getSourceStamp. - private final Map pCollectionIdToSourceStamp; + // PCollection id -> a stable global id its producer stamps on the watermark it emits (always as + // "1 source of 1"), so a downstream Flatten can tell its input branches apart by producer. Only + // PCollections feeding a Flatten are numbered; everything else reports id 0. A PCollection + // feeding + // several Flattens keeps ONE id, so each Flatten just waits for its own inputs' ids to report + // (je-ik's steer — the watermark is generated without regard to who consumes it). + private final Map pCollectionIdToProducerId; + private int nextProducerId = SINGLE_SOURCE_ID + 1; public static KafkaStreamsTranslationContext create( JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { @@ -70,7 +73,7 @@ private KafkaStreamsTranslationContext( this.pipelineOptions = pipelineOptions; this.topology = topology; this.pCollectionIdToProcessorName = new HashMap<>(); - this.pCollectionIdToSourceStamp = new HashMap<>(); + this.pCollectionIdToProducerId = new HashMap<>(); } public JobInfo getJobInfo() { @@ -131,54 +134,21 @@ public String getReadBootstrapTopic(String transformId) { } /** - * Records that the given PCollection is the {@code sourcePartition}-th of {@code totalPartitions} - * input branches of a Flatten, so its producer stamps its watermark with that identity instead of - * the default single source. Lets the Flatten's {@link WatermarkManager} tell its branches apart - * (Kafka Streams does not tell a processor which parent forwarded a record). + * Assigns the given Flatten-input PCollection a stable global producer id (or returns the one it + * already has, so a PCollection feeding several Flattens keeps a single id). Its producer stamps + * this id on the watermark it emits so each downstream Flatten can tell its input branches apart. */ - public void registerFlattenSourceStamp( - String pCollectionId, int sourcePartition, int totalPartitions) { - SourceStamp existing = - pCollectionIdToSourceStamp.putIfAbsent( - pCollectionId, new SourceStamp(sourcePartition, totalPartitions)); - if (existing != null) { - // A PCollection feeding a Flatten more than once — the same Flatten twice, or two different - // Flattens — would need a distinct watermark source per branch, but its single producer can - // only stamp one identity. Fail fast rather than get the watermark stuck or silently drop the - // duplicate branch. - throw new UnsupportedOperationException( - "PCollection " - + pCollectionId - + " feeds a Flatten more than once (the same Flatten twice, or multiple Flattens);" - + " this is not yet supported by the Kafka Streams runner."); - } + public int assignFlattenProducerId(String pCollectionId) { + return pCollectionIdToProducerId.computeIfAbsent(pCollectionId, id -> nextProducerId++); } /** - * Returns the watermark source stamp a producer should report for the given output PCollection: - * its Flatten branch identity if it feeds a Flatten (see {@link #registerFlattenSourceStamp}), or - * the single source {@code (0 of 1)} otherwise. + * Returns the watermark source id a producer should stamp for the given output PCollection: its + * global producer id if it feeds a Flatten (see {@link #assignFlattenProducerId}), or {@link + * #SINGLE_SOURCE_ID} otherwise. The producer always reports "1 source of 1"; a Flatten holds its + * watermark using its own input count, not the reported total. */ - public SourceStamp getSourceStamp(String pCollectionId) { - return pCollectionIdToSourceStamp.getOrDefault(pCollectionId, SINGLE_SOURCE); - } - - /** A watermark report's source identity: partition {@code sourcePartition} of {@code total}. */ - public static final class SourceStamp { - private final int sourcePartition; - private final int totalPartitions; - - SourceStamp(int sourcePartition, int totalPartitions) { - this.sourcePartition = sourcePartition; - this.totalPartitions = totalPartitions; - } - - public int getSourcePartition() { - return sourcePartition; - } - - public int getTotalPartitions() { - return totalPartitions; - } + public int getProducerWatermarkId(String pCollectionId) { + return pCollectionIdToProducerId.getOrDefault(pCollectionId, SINGLE_SOURCE_ID); } } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java index c595cbfdb5d2..d2b57b514e30 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java @@ -89,6 +89,39 @@ public void flattenUnionsEveryBranch() { } } + @Test + public void pCollectionFeedingTwoFlattensIsSupported() { + // je-ik's #39273 case: input2 feeds both flattens. Per-Flatten branch numbering could not + // express this (input2 would need two identities); a global producer id per PCollection can — + // input2's producer stamps one id, and each flatten still holds until its own two branches + // drain. Verify both flattens produce the right union. + try (SharedTestCollector left = SharedTestCollector.create(); + SharedTestCollector right = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection input1 = + pipeline.apply("c1", Create.of(1, 2)).apply("id1", ParDo.of(new IdentityFn())); + PCollection input2 = + pipeline.apply("c2", Create.of(3, 4)).apply("id2", ParDo.of(new IdentityFn())); + PCollection input3 = + pipeline.apply("c3", Create.of(5, 6)).apply("id3", ParDo.of(new IdentityFn())); + PCollectionList.of(input1) + .and(input2) + .apply("l1", Flatten.pCollections()) + .apply("recordL1", ParDo.of(new RecordingFn(left))); + PCollectionList.of(input2) + .and(input3) + .apply("l2", Flatten.pCollections()) + .apply("recordL2", ParDo.of(new RecordingFn(right))); + + KafkaStreamsTestRunner.run(pipeline); + + assertThat(left.recorded().size(), is(4)); + assertThat(left.recorded(), hasItems(1, 2, 3, 4)); + assertThat(right.recorded().size(), is(4)); + assertThat(right.recorded(), hasItems(3, 4, 5, 6)); + } + } + @Test public void flattenOfOneBranchTwiceIsRejected() { Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); From aa78709bd4ddea9e9e6ba5bc4aee7891c2bdb2a6 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Mon, 13 Jul 2026 18:20:45 +0500 Subject: [PATCH 4/7] Address review: carry the producing transform id in the watermark payload A watermark report now carries three separate fields: which transform produced it (new transform_id proto field), which of that transform's partitions it is for, and how many partitions that transform has. Every producer (Impulse, Read, ExecutableStage, GroupByKey, Flatten) stamps its own transform id and its own physical partition (0 of 1, single-instance for now), without regard to who consumes the report. The consuming side is a new WatermarkAggregator, used by every transform that aggregates a watermark (ExecutableStage, GroupByKey, Flatten; CombinePerKey later): it is constructed with the upstream transform ids the consumer expects, known from the pipeline graph at translation time, and tracks each upstream transform's partitions with a dedicated WatermarkManager, advancing to the min across upstream transforms once all of them have reported. A report from an unexpected transform now fails fast. The producer-id numbering from the previous revision is deleted. A self-flatten (Flatten.of(pc, pc)) turns out to be handled by the fuser, which folds the Flatten into the consuming SDK-harness stage; the harness performs the duplication. The previous revision falsely rejected it. The test now asserts every element appears twice. --- .../main/proto/kafka_streams_payload.proto | 14 ++- .../translation/ExecutableStageProcessor.java | 63 +++++------ .../ExecutableStageTranslator.java | 18 +-- .../streams/translation/FlattenProcessor.java | 48 ++++---- .../translation/FlattenTranslator.java | 40 ++++--- .../translation/GroupByKeyProcessor.java | 44 ++++--- .../translation/GroupByKeyTranslator.java | 13 ++- .../streams/translation/ImpulseProcessor.java | 13 +-- .../streams/translation/KStreamsPayload.java | 27 ++++- .../translation/KStreamsPayloadSerde.java | 2 + .../KafkaStreamsPipelineTranslator.java | 39 ------- .../KafkaStreamsTranslationContext.java | 31 ----- .../streams/translation/ReadProcessor.java | 13 +-- .../translation/WatermarkAggregator.java | 107 ++++++++++++++++++ .../streams/translation/WatermarkPayload.java | 18 ++- ...ExecutableStageProcessorWatermarkTest.java | 22 +++- .../streams/translation/FlattenTest.java | 45 ++++---- .../translation/KStreamsPayloadSerdeTest.java | 7 +- 18 files changed, 341 insertions(+), 223 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java diff --git a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto index 87e3de44397b..4dadf2d14b70 100644 --- a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto +++ b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto @@ -27,16 +27,22 @@ option java_outer_classname = "KafkaStreamsPayloadProtos"; // topic boundary (e.g. the GroupByKey repartition topic and the watermark fan-out). Protobuf is // used for compatible schema evolution and compact varint encoding. message KafkaStreamsPayload { - // A watermark report: the watermark plus the in-band coordination fields the downstream - // WatermarkManager needs. + // A watermark report: the watermark plus the in-band coordination fields a downstream + // watermark aggregator needs to reconstruct its input watermark. message WatermarkPayload { // Event-time watermark in milliseconds. Signed (sint64, zigzag-encoded) because Beam event // times can be negative, e.g. BoundedWindow.TIMESTAMP_MIN_VALUE. sint64 millis = 1; - // The source partition this report is for. + // Which partition (physical instance) of the producing transform this report is for, in + // [0, total_partitions). uint32 source_partition = 2; - // The total number of source partitions feeding the downstream stage. + // How many partitions (physical instances) the producing transform has in total. uint32 total_partitions = 3; + // Globally unique id of the transform that produced this report. A producer stamps its own id + // without regard to who consumes the report; a consumer with several upstream transforms + // (e.g. Flatten) aggregates per producing transform, holding its output watermark until every + // partition of every upstream transform it expects has reported. + string transform_id = 4; } // A data element: the Beam WindowedValue encoded with the PCollection's windowed-value coder. diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java index f2e355638a0a..36adf8f2c002 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.kafka.streams.translation; import java.util.Queue; +import java.util.Set; import java.util.concurrent.ConcurrentLinkedQueue; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.fnexecution.control.BundleProgressHandler; @@ -51,12 +52,12 @@ * ProcessorContext#forward} must only be called from the processing thread, so outputs are never * forwarded directly from a harness callback. * - *

A {@link KStreamsPayload#isWatermark() watermark} payload is a per-source-partition report and - * marks a bundle boundary: the open bundle (if any) is closed (flushing outputs), the report is fed - * to the {@link WatermarkManager}, and the stage's output watermark is forwarded downstream only - * when the {@code min()} across its source partitions actually advances. Until every source - * partition has reported, the watermark is held and nothing is forwarded — but data is still - * processed in the meantime. + *

A {@link KStreamsPayload#isWatermark() watermark} payload is a report from one partition of + * one upstream transform and marks a bundle boundary: the open bundle (if any) is closed (flushing + * outputs), the report is fed to the {@link WatermarkAggregator}, and the stage's output watermark + * is forwarded downstream — stamped with this stage's own transform id — only when the aggregate + * across the upstream transform's partitions actually advances. Until every partition has reported, + * the watermark is held and nothing is forwarded — but data is still processed in the meantime. * *

This is the Kafka Streams analogue of Flink's {@code ExecutableStageDoFnOperator} and Spark's * {@code SparkExecutableStageFunction}. State, timers, and side inputs are out of scope for this @@ -70,11 +71,9 @@ class ExecutableStageProcessor private final RunnerApi.ExecutableStagePayload stagePayload; private final JobInfo jobInfo; - // The source identity this stage stamps on its forwarded watermark. Normally (0 of 1), a single - // source to the next stage; but when this stage's output feeds a Flatten it is that branch's - // (i of N), so the Flatten's WatermarkManager can tell the branches apart. See FlattenProcessor. - private final int sourcePartition; - private final int totalPartitions; + // This stage's own transform id, stamped on every watermark it forwards so downstream watermark + // aggregators know which transform the report came from — regardless of who consumes it. + private final String transformId; // pendingOutputs is enqueued by SDK harness threads (inside the OutputReceiverFactory callback) // and drained by the Kafka Streams processing thread on bundle close; needs to be thread-safe. @@ -84,9 +83,9 @@ class ExecutableStageProcessor // only safe because the Impulse output coder happens to be ByteArrayCoder. private final Queue> pendingOutputs = new ConcurrentLinkedQueue<>(); - // Computes this stage's output watermark as min() over its source partitions' reported - // watermarks, holding until every source partition has reported (see WatermarkManager). - private final WatermarkManager watermarkManager = new WatermarkManager(); + // Computes this stage's input watermark from its upstream transform's reports, holding until + // every partition of the upstream transform has reported (see WatermarkAggregator). + private final WatermarkAggregator watermarkAggregator; // The last watermark actually forwarded downstream, so we only forward when it advances. private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; @@ -95,15 +94,20 @@ class ExecutableStageProcessor private @Nullable StageBundleFactory stageBundleFactory; private @Nullable RemoteBundle currentBundle; + /** + * @param transformId this stage's own transform id, stamped on the watermarks it emits + * @param upstreamTransformIds the transform ids feeding this stage (known from the pipeline + * graph), whose reports the {@link WatermarkAggregator} waits for + */ ExecutableStageProcessor( RunnerApi.ExecutableStagePayload stagePayload, JobInfo jobInfo, - int sourcePartition, - int totalPartitions) { + String transformId, + Set upstreamTransformIds) { this.stagePayload = stagePayload; this.jobInfo = jobInfo; - this.sourcePartition = sourcePartition; - this.totalPartitions = totalPartitions; + this.transformId = transformId; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); } @Override @@ -131,14 +135,10 @@ public void process(Record> record) { // Emit any buffered outputs before the watermark. Data is processed regardless of watermark // readiness; only the watermark itself is held until every source partition has reported. closeBundleAndFlush(record); - // Feed the report into the WatermarkManager and forward the stage's output watermark only - // when min() across the source partitions actually advances, not on every received watermark. - WatermarkPayload report = payload.asWatermark(); - watermarkManager.observe( - report.getSourcePartition(), - new Instant(report.getWatermarkMillis()), - report.getTotalSourcePartitions()); - Instant advanced = watermarkManager.advance(); + // Feed the report into the aggregator and forward the stage's output watermark only when the + // aggregate across the upstream transform's partitions actually advances. + watermarkAggregator.observe(payload.asWatermark()); + Instant advanced = watermarkAggregator.advance(); if (advanced.isAfter(lastForwardedWatermark)) { lastForwardedWatermark = advanced; forwardWatermark(record, advanced.getMillis()); @@ -214,16 +214,15 @@ private void closeBundleAndFlush(Record> record) { } private void forwardWatermark(Record> record, long watermarkMillis) { - // This stage is a single instance for now, so downstream it is one source — reported as - // (sourcePartition of totalPartitions), which is (0 of 1) unless its output feeds a Flatten, in - // which case it is that branch's identity so the Flatten can tell its branches apart. Fanning - // the watermark out to every downstream partition — and producing it atomically with the offset - // commit so it is durable — lands with the topic-based shuffle work (#18479). + // Stamped with this stage's own transform id; this stage is a single instance for now, so the + // report is for its only partition (0 of 1). Fanning the watermark out to every downstream + // partition — and producing it atomically with the offset commit so it is durable — lands with + // the topic-based shuffle work (#18479). ProcessorContext> ctx = checkInitialized(context); ctx.forward( new Record>( record.key(), - KStreamsPayload.watermark(watermarkMillis, sourcePartition, totalPartitions), + KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), record.timestamp())); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java index d6e7df99537c..a2e6ed837c7d 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java @@ -19,6 +19,7 @@ import java.io.IOException; import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.kafka.streams.Topology; @@ -82,22 +83,15 @@ public void translate( String inputPCollectionId = stagePayload.getInput(); String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); - // A stage has at most one output (multi-output rejected above). Stamp its watermark with that - // output's global producer id (0 if no Flatten consumes it), always as a single source (1 of - // 1): - // a downstream Flatten tells its branches apart by this id and holds using its own input count, - // while a single-input consumer just sees one source. - int watermarkSourceId = - transform.getOutputsMap().isEmpty() - ? 0 - : context.getProducerWatermarkId( - Iterables.getOnlyElement(transform.getOutputsMap().values())); - Topology topology = context.getTopology(); + // The stage stamps its own transform id on the watermarks it emits, and aggregates its input + // watermark from the reports of its single upstream transform (the producer of its input + // PCollection, whose node name is the upstream transform id). topology.addProcessor( transformId, () -> - new ExecutableStageProcessor(stagePayload, context.getJobInfo(), watermarkSourceId, 1), + new ExecutableStageProcessor( + stagePayload, context.getJobInfo(), transformId, ImmutableSet.of(parentProcessor)), parentProcessor); if (!transform.getOutputsMap().isEmpty()) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java index 84ea47eb0726..9da483bfad4a 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.kafka.streams.translation; +import java.util.Set; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.kafka.streams.processor.api.Processor; import org.apache.kafka.streams.processor.api.ProcessorContext; @@ -38,31 +39,34 @@ * every input branch has reported, so a downstream GroupByKey does not fire before all * flattened branches are drained. * - *

The {@link WatermarkManager} tells the input branches apart by the global producer id each + *

The {@link WatermarkAggregator} tells the input branches apart by the transform id each * branch's producer stamps on its watermark (Kafka Streams does not tell a processor which parent - * forwarded a record). Each producer stamps its own stable id regardless of who consumes it, so a - * PCollection feeding several Flattens reports one id and every Flatten still waits only for its - * own inputs. The Flatten holds its watermark using its own input count — passed in at construction - * — not the reported total (producers always report the single source {@code 1 of 1}); that is what - * lets an input shared by two Flattens work. + * forwarded a record). Each producer stamps its own identity regardless of who consumes it, so a + * PCollection feeding several Flattens reports one identity and every Flatten still waits only for + * the upstream transforms it expects — the set handed to it at construction from the pipeline + * graph. */ class FlattenProcessor implements Processor, byte[], KStreamsPayload> { - // Computes the output watermark as min() over the input branches, holding until every branch has - // reported (see WatermarkManager). Flatten is a single instance for now. - private final WatermarkManager watermarkManager = new WatermarkManager(); - // How many input branches feed this Flatten. The watermark is held until every branch's producer - // has reported, so this is the count the WatermarkManager waits for — not the reported total, - // which is always 1 because a producer stamps its watermark without regard to who consumes it. - private final int expectedInputCount; + // This transform's own id, stamped on every watermark it forwards downstream. + private final String transformId; + // Computes the output watermark as min() over the upstream transforms' reports, holding until + // every partition of every expected upstream transform has reported (see WatermarkAggregator). + private final WatermarkAggregator watermarkAggregator; // The last watermark actually forwarded downstream, so we only forward when it advances. private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; private @Nullable ProcessorContext> context; - FlattenProcessor(int expectedInputCount) { - this.expectedInputCount = expectedInputCount; + /** + * @param transformId this Flatten's own transform id, stamped on the watermarks it emits + * @param upstreamTransformIds the producers of this Flatten's input PCollections (known from the + * pipeline graph), whose reports the {@link WatermarkAggregator} waits for + */ + FlattenProcessor(String transformId, Set upstreamTransformIds) { + this.transformId = transformId; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); } @Override @@ -79,20 +83,16 @@ public void process(Record> record) { ctx.forward(record); return; } - WatermarkPayload report = payload.asWatermark(); - // Key on the producer id the branch stamped, but wait for this Flatten's own input count — not - // the report's total (always 1) — so a producer shared with another Flatten reports one id yet - // each Flatten still holds until all of its own branches arrive. - watermarkManager.observe( - report.getSourcePartition(), new Instant(report.getWatermarkMillis()), expectedInputCount); - Instant advanced = watermarkManager.advance(); + watermarkAggregator.observe(payload.asWatermark()); + Instant advanced = watermarkAggregator.advance(); if (advanced.isAfter(lastForwardedWatermark)) { lastForwardedWatermark = advanced; - // Flatten is one logical source to the next stage; report it as partition 0 of 1. + // Stamped with this Flatten's own transform id; Flatten is a single instance for now, so the + // report is for its only partition (0 of 1). ctx.forward( new Record>( record.key(), - KStreamsPayload.watermark(advanced.getMillis(), 0, 1), + KStreamsPayload.watermark(advanced.getMillis(), transformId, 0, 1), record.timestamp())); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java index 18a6cbae4a46..8f6ce2b546dd 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java @@ -18,7 +18,9 @@ package org.apache.beam.runners.kafka.streams.translation; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; +import java.util.Set; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.kafka.streams.Topology; @@ -30,14 +32,17 @@ *

Wires a single {@link FlattenProcessor} node to the producer of every input PCollection (Kafka * Streams lets a processor have many parents), so the parents' data streams merge into it, and * registers it as the producer of the flattened output so downstream translators wire to it. The - * processor forwards data through and owns its output watermark via a {@link WatermarkManager}, - * exactly like GroupByKey — see {@link FlattenProcessor}. + * processor forwards data through and owns its output watermark via a {@link WatermarkAggregator}, + * which is handed the producers of the input PCollections — the upstream transform ids whose + * watermark reports the Flatten must hear from. Producers stamp their own transform id on the + * reports they emit, without regard to who consumes them, so an input shared with another Flatten + * needs no special handling. * - *

The producer id the {@link WatermarkManager} keys each branch on is stamped upstream, by the - * branch's producing transform, because Kafka Streams does not tell a processor which parent - * forwarded a record. Ids are assigned once per Flatten-input PCollection (see {@link - * KafkaStreamsPipelineTranslator}), so an input shared by two Flattens reports one id and each - * Flatten still waits only for its own branches. + *

A user-written self-flatten never reaches this translator: the fuser folds the Flatten into + * the consuming SDK-harness stage, which performs the duplication itself. The duplicate-input check + * below is defensive — if a runner-executed Flatten ever did receive the same PCollection twice, + * Kafka Streams could not wire the same parent to a child twice and the duplicate copy would be + * silently dropped, so failing fast is safer. */ class FlattenTranslator implements PTransformTranslator { @@ -48,19 +53,28 @@ public void translate( // Flatten produces exactly one output PCollection, fed by all of its input PCollections. String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + Set seenInputs = new HashSet<>(); List parentProcessors = new ArrayList<>(); + Set upstreamTransformIds = new HashSet<>(); for (String inputPCollectionId : transform.getInputsMap().values()) { - parentProcessors.add(context.getProcessorNameForPCollection(inputPCollectionId)); + if (!seenInputs.add(inputPCollectionId)) { + throw new UnsupportedOperationException( + "Flatten " + + transform.getUniqueName() + + " has PCollection " + + inputPCollectionId + + " as an input more than once; a self-flatten is not yet supported by the Kafka" + + " Streams runner."); + } + String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); + parentProcessors.add(parentProcessor); + upstreamTransformIds.add(parentProcessor); } - // The Flatten holds its watermark until all of its input branches report. Inputs are distinct - // (a self-flatten is rejected during producer-id assignment), so the input count is the number - // of producer ids to wait for. - int inputCount = transform.getInputsMap().size(); Topology topology = context.getTopology(); topology.addProcessor( transformId, - () -> new FlattenProcessor(inputCount), + () -> new FlattenProcessor(transformId, upstreamTransformIds), parentProcessors.toArray(new String[0])); context.registerPCollectionProducer(outputPCollectionId, transformId); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java index a6c0cb8c4061..208dde54204e 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java @@ -19,6 +19,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Set; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.coders.IterableCoder; @@ -41,10 +42,10 @@ * *

Records arrive on the repartition topic keyed by the encoded Beam key, so every value of a key * is co-located here. Each value is appended to a per-key buffer in a Kafka Streams state store. - * Watermark reports are fed to a {@link WatermarkManager}; when the input watermark reaches {@link - * BoundedWindow#TIMESTAMP_MAX_VALUE} (the end of the global window) every buffered key is emitted - * once as {@code KV>} and the buffer cleared, then the watermark is forwarded - * downstream. + * Watermark reports are fed to a {@link WatermarkAggregator}; when the input watermark reaches + * {@link BoundedWindow#TIMESTAMP_MAX_VALUE} (the end of the global window) every buffered key is + * emitted once as {@code KV>} and the buffer cleared, then the watermark is + * forwarded downstream. * *

Buffering whole value lists and re-encoding on each append is O(n^2) per key; fine for this * first GroupByKey, and replaced when this moves to runner-core {@code GroupAlsoByWindow}. @@ -53,10 +54,15 @@ class GroupByKeyProcessor implements Processor, byte[], KStreamsPayload> { private final String stateStoreName; + // This transform's own id, stamped on every watermark it forwards downstream. + private final String transformId; private final Coder keyCoder; private final IterableCoder<@Nullable Object> bufferCoder; - private final WatermarkManager watermarkManager = new WatermarkManager(); + // Aggregates the input watermark from the upstream transform's reports, which arrive through the + // repartition topic with the upstream producer's transform id intact (the shuffle forwards + // watermark payloads unchanged). + private final WatermarkAggregator watermarkAggregator; private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; // The global window fires exactly once, when the watermark first reaches its end. Later watermark // reports (e.g. the same terminal watermark broadcast across repartition partitions) must not @@ -69,9 +75,20 @@ class GroupByKeyProcessor private @Nullable ProcessorContext> context; private @Nullable KeyValueStore store; + /** + * @param transformId this transform's own id, stamped on the watermarks it emits + * @param upstreamTransformIds the transform ids feeding this GroupByKey (known from the pipeline + * graph), whose reports the {@link WatermarkAggregator} waits for + */ GroupByKeyProcessor( - String stateStoreName, Coder keyCoder, Coder<@Nullable Object> valueCoder) { + String stateStoreName, + String transformId, + Set upstreamTransformIds, + Coder keyCoder, + Coder<@Nullable Object> valueCoder) { this.stateStoreName = stateStoreName; + this.transformId = transformId; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); this.keyCoder = keyCoder; this.bufferCoder = IterableCoder.of(valueCoder); } @@ -94,12 +111,8 @@ public void process(Record> record) { appendValue(encodedKey, element); return; } - WatermarkPayload report = payload.asWatermark(); - watermarkManager.observe( - report.getSourcePartition(), - new Instant(report.getWatermarkMillis()), - report.getTotalSourcePartitions()); - Instant advanced = watermarkManager.advance(); + watermarkAggregator.observe(payload.asWatermark()); + Instant advanced = watermarkAggregator.advance(); if (!fired && !advanced.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) { fireAll(record); fired = true; @@ -153,10 +166,13 @@ private void fireAll(Record> trigger) { private void forwardWatermark(Record> trigger, long watermarkMillis) { ProcessorContext> ctx = checkInitialized(context); - // GroupByKey is a single logical source for the next stage; report it as partition 0 of 1. + // Stamped with this transform's own id; GroupByKey is a single instance for now, so the report + // is for its only partition (0 of 1). ctx.forward( new Record>( - trigger.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1), trigger.timestamp())); + trigger.key(), + KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), + trigger.timestamp())); } private byte[] encodeBuffer(List<@Nullable Object> values) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index d7c4a309d9c3..9e23dbb5cfb0 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.streams.Topology; @@ -110,10 +111,18 @@ public void translate( payloadSerde.deserializer(), repartitionTopic); - // Buffer values per key and fire KV> at the terminal watermark. + // Buffer values per key and fire KV> at the terminal watermark. Watermark + // reports cross the repartition topic unchanged, so they still carry the id of the transform + // that produced this GroupByKey's input — the parent the shuffle is attached to. topology.addProcessor( transformId, - () -> new GroupByKeyProcessor(stateStoreName, keyCoder, valueCoder), + () -> + new GroupByKeyProcessor( + stateStoreName, + transformId, + ImmutableSet.of(parentProcessor), + keyCoder, + valueCoder), sourceName); topology.addStateStore( Stores.keyValueStoreBuilder( diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java index 675bee4d8591..bac91978a298 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java @@ -122,17 +122,16 @@ private void maybeFire() { } /** - * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. - * - *

Impulse is a single-instance source, so the report is stamped as the only source partition: - * {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real per-partition identities - * arrive once the topology gains topic-based shuffle. + * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors, + * stamped with this transform's id. Impulse is a single-instance source, so the report is for its + * only partition: {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real + * per-partition identities arrive once the topology gains topic-based shuffle. */ - private static void forwardWatermarkMax(ProcessorContext> ctx) { + private void forwardWatermarkMax(ProcessorContext> ctx) { long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); ctx.forward( new Record>( - new byte[0], KStreamsPayload.watermark(maxMillis, 0, 1), 0L)); + new byte[0], KStreamsPayload.watermark(maxMillis, transformId, 0, 1), 0L)); } /** Cancels the wall-clock punctuator after the impulse has fired to stop periodic wakeups. */ diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java index 93b40346761a..c165f0e875d4 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java @@ -57,6 +57,7 @@ private enum Kind { private final Kind kind; private final @Nullable WindowedValue data; private final long watermarkMillis; + private final String transformId; private final int sourcePartition; private final int totalSourcePartitions; @@ -64,27 +65,33 @@ private KStreamsPayload( Kind kind, @Nullable WindowedValue data, long watermarkMillis, + String transformId, int sourcePartition, int totalSourcePartitions) { this.kind = kind; this.data = data; this.watermarkMillis = watermarkMillis; + this.transformId = transformId; this.sourcePartition = sourcePartition; this.totalSourcePartitions = totalSourcePartitions; } /** Returns a data payload wrapping the given {@link WindowedValue}. */ public static KStreamsPayload data(WindowedValue value) { - return new KStreamsPayload<>(Kind.DATA, value, 0L, 0, 0); + return new KStreamsPayload<>(Kind.DATA, value, 0L, "", 0, 0); } /** * Returns a watermark report payload: the event-time milliseconds together with the in-band - * coordination fields the downstream stage's {@link WatermarkManager} needs — which source - * partition this report is for and how many source partitions feed the stage in total. + * coordination fields a downstream watermark aggregator needs — which transform produced the + * report ({@code transformId}, stamped by the producer without regard to who consumes it), which + * of that transform's partitions this report is for, and how many partitions that transform has + * in total. */ public static KStreamsPayload watermark( - long watermarkMillis, int sourcePartition, int totalSourcePartitions) { + long watermarkMillis, String transformId, int sourcePartition, int totalSourcePartitions) { + Preconditions.checkArgument( + transformId != null && !transformId.isEmpty(), "transformId must be non-empty"); Preconditions.checkArgument( totalSourcePartitions > 0, "totalSourcePartitions must be positive: %s", @@ -95,7 +102,7 @@ public static KStreamsPayload watermark( sourcePartition, totalSourcePartitions); return new KStreamsPayload<>( - Kind.WATERMARK, null, watermarkMillis, sourcePartition, totalSourcePartitions); + Kind.WATERMARK, null, watermarkMillis, transformId, sourcePartition, totalSourcePartitions); } public boolean isData() { @@ -134,6 +141,11 @@ public long getWatermarkMillis() { return watermarkMillis; } + @Override + public String getTransformId() { + return transformId; + } + @Override public int getSourcePartition() { return sourcePartition; @@ -156,6 +168,7 @@ public boolean equals(@Nullable Object o) { KStreamsPayload that = (KStreamsPayload) o; return kind == that.kind && watermarkMillis == that.watermarkMillis + && transformId.equals(that.transformId) && sourcePartition == that.sourcePartition && totalSourcePartitions == that.totalSourcePartitions && Objects.equals(data, that.data); @@ -163,7 +176,8 @@ public boolean equals(@Nullable Object o) { @Override public int hashCode() { - return Objects.hash(kind, data, watermarkMillis, sourcePartition, totalSourcePartitions); + return Objects.hash( + kind, data, watermarkMillis, transformId, sourcePartition, totalSourcePartitions); } @Override @@ -174,6 +188,7 @@ public String toString() { } else { helper .add("watermarkMillis", watermarkMillis) + .add("transformId", transformId) .add("sourcePartition", sourcePartition) .add("totalSourcePartitions", totalSourcePartitions); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java index 912eb5fb6047..1363740d58bd 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java @@ -82,6 +82,7 @@ public byte[] serialize(String topic, KStreamsPayload payload) { proto.setWatermark( KafkaStreamsPayload.WatermarkPayload.newBuilder() .setMillis(watermark.getWatermarkMillis()) + .setTransformId(watermark.getTransformId()) .setSourcePartition(watermark.getSourcePartition()) .setTotalPartitions(watermark.getTotalSourcePartitions())); } @@ -109,6 +110,7 @@ public KStreamsPayload deserialize(String topic, byte[] bytes) { KafkaStreamsPayload.WatermarkPayload watermark = proto.getWatermark(); return KStreamsPayload.watermark( watermark.getMillis(), + watermark.getTransformId(), watermark.getSourcePartition(), watermark.getTotalPartitions()); case PAYLOAD_NOT_SET: diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index bd63a04e1b45..a9e26ecd6412 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -18,10 +18,6 @@ package org.apache.beam.runners.kafka.streams.translation; import com.google.auto.service.AutoService; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashSet; -import java.util.List; import java.util.Map; import java.util.Set; import org.apache.beam.model.pipeline.v1.RunnerApi; @@ -111,7 +107,6 @@ public RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) { * Throws {@link UnsupportedOperationException} on the first unsupported URN. */ public void translate(KafkaStreamsTranslationContext context, RunnerApi.Pipeline pipeline) { - assignFlattenProducerIds(context, pipeline); QueryablePipeline queryable = QueryablePipeline.forTransforms( pipeline.getRootTransformIdsList(), pipeline.getComponents()); @@ -132,40 +127,6 @@ public void translate(KafkaStreamsTranslationContext context, RunnerApi.Pipeline } } - /** - * Assigns each Flatten input PCollection a stable global producer id, so the producer of that - * PCollection stamps its watermark with that id and each downstream Flatten can tell its input - * branches apart. A PCollection shared by several Flattens keeps one id, so every Flatten still - * waits only for its own inputs' ids — this is what makes a PCollection feeding two Flattens work - * (je-ik's review of #39273). A self-flatten (the same PCollection listed as an input twice) is - * rejected: its single producer cannot be two branches, and Kafka Streams cannot wire the same - * parent to a child twice, so the duplicate copy would be silently dropped. - */ - private static void assignFlattenProducerIds( - KafkaStreamsTranslationContext context, RunnerApi.Pipeline pipeline) { - for (RunnerApi.PTransform transform : pipeline.getComponents().getTransformsMap().values()) { - if (!PTransformTranslation.FLATTEN_TRANSFORM_URN.equals(transform.getSpec().getUrn())) { - continue; - } - // Sorted so the numbering is deterministic across topology builds. - List inputPCollectionIds = new ArrayList<>(transform.getInputsMap().values()); - Collections.sort(inputPCollectionIds); - Set seenInputs = new HashSet<>(); - for (String inputPCollectionId : inputPCollectionIds) { - if (!seenInputs.add(inputPCollectionId)) { - throw new UnsupportedOperationException( - "Flatten " - + transform.getUniqueName() - + " has PCollection " - + inputPCollectionId - + " as an input more than once; a self-flatten is not yet supported by the Kafka" - + " Streams runner."); - } - context.assignFlattenProducerId(inputPCollectionId); - } - } - } - /** * Tells the SDK that URNs handled directly by the Kafka Streams runner should be treated as * primitives by {@link QueryablePipeline}. Mirrors Flink's {@code IsFlinkNativeTransform} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java index 506bc7342b1f..0a045ff7395b 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java @@ -41,21 +41,10 @@ public class KafkaStreamsTranslationContext { /** Characters not legal in a Kafka topic name; a topic's legal set is {@code [a-zA-Z0-9._-]}. */ private static final Pattern ILLEGAL_TOPIC_CHARS = Pattern.compile("[^a-zA-Z0-9._-]"); - /** The watermark source id a producer reports for an output no Flatten consumes. */ - private static final int SINGLE_SOURCE_ID = 0; - private final JobInfo jobInfo; private final KafkaStreamsPipelineOptions pipelineOptions; private final Topology topology; private final Map pCollectionIdToProcessorName; - // PCollection id -> a stable global id its producer stamps on the watermark it emits (always as - // "1 source of 1"), so a downstream Flatten can tell its input branches apart by producer. Only - // PCollections feeding a Flatten are numbered; everything else reports id 0. A PCollection - // feeding - // several Flattens keeps ONE id, so each Flatten just waits for its own inputs' ids to report - // (je-ik's steer — the watermark is generated without regard to who consumes it). - private final Map pCollectionIdToProducerId; - private int nextProducerId = SINGLE_SOURCE_ID + 1; public static KafkaStreamsTranslationContext create( JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { @@ -73,7 +62,6 @@ private KafkaStreamsTranslationContext( this.pipelineOptions = pipelineOptions; this.topology = topology; this.pCollectionIdToProcessorName = new HashMap<>(); - this.pCollectionIdToProducerId = new HashMap<>(); } public JobInfo getJobInfo() { @@ -132,23 +120,4 @@ public String getReadBootstrapTopic(String transformId) { + "_" + sanitizedTransformId; } - - /** - * Assigns the given Flatten-input PCollection a stable global producer id (or returns the one it - * already has, so a PCollection feeding several Flattens keeps a single id). Its producer stamps - * this id on the watermark it emits so each downstream Flatten can tell its input branches apart. - */ - public int assignFlattenProducerId(String pCollectionId) { - return pCollectionIdToProducerId.computeIfAbsent(pCollectionId, id -> nextProducerId++); - } - - /** - * Returns the watermark source id a producer should stamp for the given output PCollection: its - * global producer id if it feeds a Flatten (see {@link #assignFlattenProducerId}), or {@link - * #SINGLE_SOURCE_ID} otherwise. The producer always reports "1 source of 1"; a Flatten holds its - * watermark using its own input count, not the reported total. - */ - public int getProducerWatermarkId(String pCollectionId) { - return pCollectionIdToProducerId.getOrDefault(pCollectionId, SINGLE_SOURCE_ID); - } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java index 604eca54014a..eb20a8b83588 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java @@ -183,17 +183,16 @@ private WindowedValue toRunnerWire(WindowedValue element) { } /** - * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. - * - *

Read is a single-instance source, so the report is stamped as the only source partition: - * {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real per-partition identities - * arrive once the topology gains topic-based shuffle. + * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors, + * stamped with this transform's id. Read is a single-instance source, so the report is for its + * only partition: {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real + * per-partition identities arrive once the topology gains topic-based shuffle. */ - private static void forwardWatermarkMax(ProcessorContext> ctx) { + private void forwardWatermarkMax(ProcessorContext> ctx) { long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); ctx.forward( new Record>( - new byte[0], KStreamsPayload.watermark(maxMillis, 0, 1), 0L)); + new byte[0], KStreamsPayload.watermark(maxMillis, transformId, 0, 1), 0L)); } /** Cancels the wall-clock punctuator after the read has fired to stop periodic wakeups. */ diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java new file mode 100644 index 000000000000..c9af03df26a9 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.kafka.streams.translation; + +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.joda.time.Instant; + +/** + * Computes a transform's input watermark from the watermark reports of its upstream transforms. + * + *

A watermark report carries three orthogonal pieces of information (see {@link + * WatermarkPayload}): which transform produced it, which partition (physical + * instance) of that transform it is for, and how many partitions that transform has. A + * producer stamps its own identity without regard to who consumes the report. This aggregator is + * the consuming side, used by every transform that aggregates a watermark — ExecutableStage, + * GroupByKey, Flatten (and CombinePerKey later): + * + *

    + *
  • It is constructed with the set of upstream transform ids the consumer expects, known from + * the pipeline graph at translation time (a single-input transform passes its one parent; a + * Flatten passes the producers of all of its input PCollections). + *
  • Per upstream transform it tracks partitions with a dedicated {@link WatermarkManager}, + * which holds until every partition of that transform has reported and keeps each partition + * monotonic. + *
  • The aggregate input watermark is the {@code min()} across the upstream transforms' + * watermarks, defined only once every expected upstream transform is ready; until + * then {@link #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} and the caller + * emits nothing. + *
+ * + *

Not thread-safe; the caller (a single Kafka Streams processor thread) serializes access. + */ +final class WatermarkAggregator { + + /** Upstream transform ids this consumer must hear from, fixed by the pipeline graph. */ + private final Set expectedUpstreamTransformIds; + + /** Per-upstream-transform partition tracking. */ + private final Map managerByTransformId = new HashMap<>(); + + WatermarkAggregator(Set expectedUpstreamTransformIds) { + Preconditions.checkArgument( + !expectedUpstreamTransformIds.isEmpty(), "expectedUpstreamTransformIds must not be empty"); + this.expectedUpstreamTransformIds = ImmutableSet.copyOf(expectedUpstreamTransformIds); + } + + /** + * Records one upstream watermark report. A report from a transform this consumer does not expect + * indicates a translation wiring bug and fails fast. + */ + void observe(WatermarkPayload report) { + String transformId = report.getTransformId(); + if (!expectedUpstreamTransformIds.contains(transformId)) { + throw new IllegalStateException( + "Received a watermark report from unexpected transform " + + transformId + + "; expected one of " + + expectedUpstreamTransformIds); + } + managerByTransformId + .computeIfAbsent(transformId, id -> new WatermarkManager()) + .observe( + report.getSourcePartition(), + new Instant(report.getWatermarkMillis()), + report.getTotalSourcePartitions()); + } + + /** + * Returns the aggregate input watermark: {@code min()} across all expected upstream transforms, + * or {@link BoundedWindow#TIMESTAMP_MIN_VALUE} while any upstream transform has not yet fully + * reported (the hold). + */ + Instant advance() { + if (managerByTransformId.size() < expectedUpstreamTransformIds.size()) { + return BoundedWindow.TIMESTAMP_MIN_VALUE; + } + Instant min = BoundedWindow.TIMESTAMP_MAX_VALUE; + for (WatermarkManager manager : managerByTransformId.values()) { + // A not-yet-ready manager advances to TIMESTAMP_MIN_VALUE, which correctly holds the min. + Instant watermark = manager.advance(); + if (watermark.isBefore(min)) { + min = watermark; + } + } + return min; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java index 194be701ecec..bac5314a6da7 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java @@ -24,18 +24,26 @@ * {@link KStreamsPayload#isWatermark()} and narrowed the payload, so there is no kind check to do * on each accessor. * - *

A watermark report is the in-band coordination message a downstream stage's {@link - * WatermarkManager} consumes: the watermark value plus which source partition reported it and how - * many source partitions feed the stage in total. + *

A watermark report is the in-band coordination message a downstream watermark aggregator + * consumes: the watermark value, which transform produced it, which of that transform's partitions + * reported it, and how many partitions that transform has in total. The producer stamps its own + * identity without regard to who consumes the report; a consumer with several upstream transforms + * (e.g. Flatten) aggregates per producing transform. */ public interface WatermarkPayload { /** The reported watermark, in event-time milliseconds. */ long getWatermarkMillis(); - /** The source partition this report is for, in {@code [0, getTotalSourcePartitions())}. */ + /** Globally unique id of the transform that produced this report. */ + String getTransformId(); + + /** + * Which partition (physical instance) of the producing transform this report is for, in {@code + * [0, getTotalSourcePartitions())}. + */ int getSourcePartition(); - /** The total number of source partitions feeding the downstream stage. */ + /** How many partitions (physical instances) the producing transform has in total. */ int getTotalSourcePartitions(); } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java index c1425d94dec9..74169acb81cb 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java @@ -25,13 +25,14 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.kafka.streams.processor.api.MockProcessorContext; import org.apache.kafka.streams.processor.api.Record; import org.junit.Test; /** * Tests the watermark wiring of {@link ExecutableStageProcessor}: how it feeds incoming watermark - * reports to the {@link WatermarkManager} and forwards the stage's output watermark. + * reports to the {@link WatermarkAggregator} and forwards the stage's output watermark. * *

Only the watermark path is exercised, so the SDK harness is never started (it is created * lazily on the first data element). A {@link MockProcessorContext} captures what the processor @@ -39,6 +40,12 @@ */ public class ExecutableStageProcessorWatermarkTest { + /** The stage's own transform id, expected on every watermark it forwards. */ + private static final String STAGE_ID = "stage"; + + /** The single upstream transform whose reports the stage aggregates. */ + private static final String UPSTREAM_ID = "upstream"; + private static ExecutableStageProcessor newProcessor() { JobInfo jobInfo = JobInfo.create( @@ -46,15 +53,18 @@ private static ExecutableStageProcessor newProcessor() { "job-name", "", PipelineOptionsTranslation.toProto(PipelineOptionsFactory.create())); - // Forwards its watermark as the single source (0 of 1); this test exercises the consume side. return new ExecutableStageProcessor( - RunnerApi.ExecutableStagePayload.getDefaultInstance(), jobInfo, 0, 1); + RunnerApi.ExecutableStagePayload.getDefaultInstance(), + jobInfo, + STAGE_ID, + ImmutableSet.of(UPSTREAM_ID)); } + /** A report from the upstream transform's given partition. */ private static Record> watermark( long millis, int sourcePartition, int totalSourcePartitions) { KStreamsPayload payload = - KStreamsPayload.watermark(millis, sourcePartition, totalSourcePartitions); + KStreamsPayload.watermark(millis, UPSTREAM_ID, sourcePartition, totalSourcePartitions); return new Record<>(new byte[0], payload, 0L); } @@ -76,7 +86,9 @@ public void singleSourcePartitionForwardsImmediatelyStampedAsItsOwnSource() { assertThat(out.isWatermark(), is(true)); WatermarkPayload report = out.asWatermark(); assertThat(report.getWatermarkMillis(), is(100L)); - // The stage forwards as its own single source (0 of 1), not the upstream's identity. + // The stage forwards under its own identity — its transform id and its own single partition + // (0 of 1) — not the upstream's. + assertThat(report.getTransformId(), is(STAGE_ID)); assertThat(report.getSourcePartition(), is(0)); assertThat(report.getTotalSourcePartitions(), is(1)); } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java index d2b57b514e30..24a82cab374d 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java @@ -20,7 +20,6 @@ import static org.hamcrest.CoreMatchers.hasItems; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.Assert.assertThrows; import java.util.List; import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; @@ -39,10 +38,11 @@ * *

Each branch is a {@code Create -> identity ParDo}, so its producer feeding the Flatten is an * {@link ExecutableStageProcessor} — the same shape PAssert's {@code GroupGlobally} produces. This - * is the case the {@code (i of N)} watermark stamping exists for: without it both branches would - * report as the single source {@code (0 of 1)}, the Flatten would release its watermark after the - * first branch drained, and the downstream stage's bundle would close early and drop the second - * branch's elements. + * exercises the per-producing-transform watermark aggregation: each branch's producer stamps its + * own transform id on its watermark, and the Flatten holds its output watermark until every + * upstream transform it expects has reported. Without that, the Flatten would release its watermark + * after the first branch drained and the downstream stage's bundle would close early, dropping the + * second branch's elements. */ public class FlattenTest { @@ -123,19 +123,26 @@ public void pCollectionFeedingTwoFlattensIsSupported() { } @Test - public void flattenOfOneBranchTwiceIsRejected() { - Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); - PCollection branch = - pipeline.apply("create", Create.of(1, 2)).apply("id", ParDo.of(new IdentityFn())); - // Flattening a branch with itself needs a distinct watermark source per branch, but the branch - // has a single producer that can only stamp one identity, so the runner rejects it (rather than - // getting the watermark stuck or silently dropping the duplicate copy). - PCollectionList.of(branch) - .and(branch) - .apply("flatten", Flatten.pCollections()) - .apply("sink", ParDo.of(new IdentityFn())); - - assertThrows( - UnsupportedOperationException.class, () -> KafkaStreamsTestRunner.translate(pipeline)); + public void flattenOfOneBranchTwiceDuplicatesEveryElement() { + // A self-flatten is a bag union with itself: every element must appear twice. The fuser folds + // the Flatten into the SDK-harness stage (it never reaches the runner's Flatten translator as a + // duplicate-input node), so the duplication happens in the harness. + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection branch = + pipeline.apply("create", Create.of(1, 2)).apply("id", ParDo.of(new IdentityFn())); + PCollectionList.of(branch) + .and(branch) + .apply("flatten", Flatten.pCollections()) + .apply("record", ParDo.of(new RecordingFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List recorded = collector.recorded(); + assertThat(recorded.size(), is(4)); + assertThat(recorded, hasItems(1, 2)); + assertThat(recorded.stream().filter(v -> v == 1).count(), is(2L)); + assertThat(recorded.stream().filter(v -> v == 2).count(), is(2L)); + } } } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java index 19346f54763e..95ce70b88578 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java @@ -58,10 +58,11 @@ public void roundTripsDataPayload() { @Test public void roundTripsWatermarkPayload() { - KStreamsPayload payload = KStreamsPayload.watermark(12345L, 2, 4); + KStreamsPayload payload = KStreamsPayload.watermark(12345L, "transform-a", 2, 4); KStreamsPayload out = roundTrip(payload); assertThat(out.isWatermark(), is(true)); assertThat(out.asWatermark().getWatermarkMillis(), is(12345L)); + assertThat(out.asWatermark().getTransformId(), is("transform-a")); assertThat(out.asWatermark().getSourcePartition(), is(2)); assertThat(out.asWatermark().getTotalSourcePartitions(), is(4)); assertThat(out, is(payload)); @@ -70,7 +71,7 @@ public void roundTripsWatermarkPayload() { @Test public void roundTripsTerminalMaxWatermark() { KStreamsPayload payload = - KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(), 0, 1); + KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(), "t", 0, 1); assertThat( roundTrip(payload).asWatermark().getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); @@ -80,7 +81,7 @@ public void roundTripsTerminalMaxWatermark() { public void roundTripsNegativeWatermark() { // Beam event times can be negative; sint64 must round-trip them losslessly. KStreamsPayload payload = - KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(), 0, 1); + KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(), "t", 0, 1); assertThat( roundTrip(payload).asWatermark().getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis())); From 769579b4ffe7f6223ec3fa43d6071551a2f51b95 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Tue, 14 Jul 2026 07:48:48 +0500 Subject: [PATCH 5/7] Retrigger CI (Actions infra failure: Bad Gateway during action download) From e6be254aaa5b1d04008e3399a8d978515dc32cba Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Tue, 14 Jul 2026 20:51:11 +0500 Subject: [PATCH 6/7] Address review: WatermarkAggregator test suite + drop-null-payload recovery Adds a dedicated WatermarkAggregatorTest covering the aggregation corner cases: holding until every expected upstream transform and every partition of each has reported, min across upstreams, per-upstream monotonicity, duplicate broadcast reports, terminal MAX aggregation, a repartition of one upstream re-opening the hold, and fail-fast on a report from an unexpected transform. A record with a null payload (an external write to a topic the runner reads, or a tombstone) is now logged and dropped by the payload-consuming processors (Flatten, GroupByKey, ShuffleByKey, ExecutableStage) instead of failing the task. --- .../translation/ExecutableStageProcessor.java | 7 + .../streams/translation/FlattenProcessor.java | 14 +- .../translation/GroupByKeyProcessor.java | 12 ++ .../translation/ShuffleByKeyProcessor.java | 10 ++ .../translation/WatermarkAggregatorTest.java | 160 ++++++++++++++++++ 5 files changed, 202 insertions(+), 1 deletion(-) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java index 36adf8f2c002..38d3601e814f 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java @@ -131,6 +131,13 @@ private void ensureStageBundleFactory() { @Override public void process(Record> record) { KStreamsPayload payload = record.value(); + if (payload == null) { + // A topic feeding the runner can always be written to from outside (or carry a tombstone), + // so recover from the obvious error instead of crashing the task: warn and drop. + LOG.warn( + "Stage {} dropping record with null payload (external write or tombstone)", transformId); + return; + } if (payload.isWatermark()) { // Emit any buffered outputs before the watermark. Data is processed regardless of watermark // readiness; only the watermark itself is held until every source partition has reported. diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java index 9da483bfad4a..d09fed8185a1 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java @@ -24,6 +24,8 @@ import org.apache.kafka.streams.processor.api.Record; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Kafka Streams {@link Processor} implementing Beam's {@code Flatten} primitive ({@code @@ -33,7 +35,7 @@ * data streams is the flatten. * *

Watermark reports are where Flatten does real work, and it owns its output watermark - * the same way GroupByKey does: it runs a {@link WatermarkManager} over its inputs, forwards its + * the same way GroupByKey does: it runs a {@link WatermarkAggregator} over its inputs, forwards its * own watermark only when the {@code min()} across them advances, and stamps that as a single * source ({@code 0 of 1}) to its downstream. This holds the output watermark back until * every input branch has reported, so a downstream GroupByKey does not fire before all @@ -49,6 +51,8 @@ class FlattenProcessor implements Processor, byte[], KStreamsPayload> { + private static final Logger LOG = LoggerFactory.getLogger(FlattenProcessor.class); + // This transform's own id, stamped on every watermark it forwards downstream. private final String transformId; // Computes the output watermark as min() over the upstream transforms' reports, holding until @@ -77,6 +81,14 @@ public void init(ProcessorContext> context) { @Override public void process(Record> record) { KStreamsPayload payload = record.value(); + if (payload == null) { + // A topic feeding the runner can always be written to from outside (or carry a tombstone), + // so recover from the obvious error instead of crashing the task: warn and drop. + LOG.warn( + "Flatten {} dropping record with null payload (external write or tombstone)", + transformId); + return; + } ProcessorContext> ctx = checkInitialized(context); if (!payload.isWatermark()) { // Data: the union of the parents' data streams is the flatten — forward unchanged. diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java index 208dde54204e..3e82935b807d 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java @@ -36,6 +36,8 @@ import org.apache.kafka.streams.state.KeyValueStore; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Executes a {@code GroupByKey} (GlobalWindow, default trigger, no allowed lateness). @@ -53,6 +55,8 @@ class GroupByKeyProcessor implements Processor, byte[], KStreamsPayload> { + private static final Logger LOG = LoggerFactory.getLogger(GroupByKeyProcessor.class); + private final String stateStoreName; // This transform's own id, stamped on every watermark it forwards downstream. private final String transformId; @@ -102,6 +106,14 @@ public void init(ProcessorContext> context) { @Override public void process(Record> record) { KStreamsPayload payload = record.value(); + if (payload == null) { + // The repartition topic can be written to from outside the runner (or carry a tombstone), + // so recover from the obvious error instead of crashing the task: warn and drop. + LOG.warn( + "GroupByKey {} dropping record with null payload (external write or tombstone)", + transformId); + return; + } if (payload.isData()) { byte[] encodedKey = record.key(); Object element = payload.getData().getValue(); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java index 3184d838e09d..774d36db0f2e 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java @@ -25,6 +25,8 @@ import org.apache.kafka.streams.processor.api.ProcessorContext; import org.apache.kafka.streams.processor.api.Record; import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Re-keys a {@code KV}-valued stream by the Beam key so Kafka Streams shuffles by it. @@ -39,6 +41,8 @@ class ShuffleByKeyProcessor implements Processor, byte[], KStreamsPayload> { + private static final Logger LOG = LoggerFactory.getLogger(ShuffleByKeyProcessor.class); + private final Coder keyCoder; private @Nullable ProcessorContext> context; @@ -55,6 +59,12 @@ public void init(ProcessorContext> context) { public void process(Record> record) { ProcessorContext> ctx = checkInitialized(context); KStreamsPayload payload = record.value(); + if (payload == null) { + // A topic feeding the runner can always be written to from outside (or carry a tombstone), + // so recover from the obvious error instead of crashing the task: warn and drop. + LOG.warn("Shuffle dropping record with null payload (external write or tombstone)"); + return; + } if (payload.isData()) { Object element = payload.getData().getValue(); if (element == null) { diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java new file mode 100644 index 000000000000..88d541bdc0c2 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java @@ -0,0 +1,160 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.kafka.streams.translation; + +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertThrows; + +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.joda.time.Instant; +import org.junit.Test; + +/** + * Unit tests for {@link WatermarkAggregator}: aggregation of watermark reports across multiple + * upstream transforms, each with its own partition set. Per-partition behaviour within one upstream + * transform (monotonicity, repartition reset) is covered in depth by {@code WatermarkManagerTest}; + * here it is exercised through the aggregate. + */ +public class WatermarkAggregatorTest { + + private static Instant ts(long millis) { + return new Instant(millis); + } + + /** A report from the given transform's partition {@code sourcePartition} of {@code total}. */ + private static WatermarkPayload report( + String transformId, long millis, int sourcePartition, int total) { + return KStreamsPayload.watermark(millis, transformId, sourcePartition, total).asWatermark(); + } + + @Test + public void holdsBeforeAnyReport() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + } + + @Test + public void singleUpstreamSinglePartitionAdvances() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + aggregator.observe(report("a", 100L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void holdsUntilEveryExpectedUpstreamReports() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b")); + aggregator.observe(report("a", 100L, 0, 1)); + // Only one of the two expected upstream transforms has reported — still holding. + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + aggregator.observe(report("b", 200L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void holdsUntilEveryPartitionOfEachUpstreamReports() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b")); + aggregator.observe(report("a", 100L, 0, 1)); + aggregator.observe(report("b", 200L, 0, 2)); + // Upstream "b" has reported only one of its two partitions — still holding. + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + aggregator.observe(report("b", 300L, 1, 2)); + // Ready: min(100, min(200, 300)) = 100. + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void aggregateIsMinAcrossUpstreams() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b", "c")); + aggregator.observe(report("a", 300L, 0, 1)); + aggregator.observe(report("b", 100L, 0, 1)); + aggregator.observe(report("c", 500L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + + // The slowest upstream advances; the aggregate follows the new min. + aggregator.observe(report("b", 400L, 0, 1)); + assertThat(aggregator.advance(), is(ts(300L))); + } + + @Test + public void perUpstreamWatermarkIsMonotonic() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + aggregator.observe(report("a", 200L, 0, 1)); + assertThat(aggregator.advance(), is(ts(200L))); + + // A late lower report from the same upstream partition is ignored. + aggregator.observe(report("a", 100L, 0, 1)); + assertThat(aggregator.advance(), is(ts(200L))); + } + + @Test + public void duplicateReportDoesNotAdvance() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + aggregator.observe(report("a", 100L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + + // The same report again (e.g. a broadcast duplicate) leaves the aggregate unchanged. + aggregator.observe(report("a", 100L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void terminalMaxWatermarkAggregates() { + long max = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b")); + aggregator.observe(report("a", max, 0, 1)); + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + aggregator.observe(report("b", max, 0, 1)); + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MAX_VALUE)); + } + + @Test + public void repartitionOfOneUpstreamReopensTheHold() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b")); + aggregator.observe(report("a", 100L, 0, 1)); + aggregator.observe(report("b", 200L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + + // Upstream "b" changes its partition count (repartition): its per-partition state resets and + // the aggregate holds again until b's new full partition set has reported. + aggregator.observe(report("b", 250L, 0, 2)); + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + aggregator.observe(report("b", 300L, 1, 2)); + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void reportFromUnexpectedTransformFailsFast() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, () -> aggregator.observe(report("intruder", 100L, 0, 1))); + assertThat(thrown.getMessage(), containsString("intruder")); + } + + @Test + public void emptyExpectedUpstreamsIsRejected() { + assertThrows(IllegalArgumentException.class, () -> new WatermarkAggregator(ImmutableSet.of())); + } +} From e4d2a77436ff1b9bd8116bb8e84644ac8d0f29a0 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Wed, 15 Jul 2026 15:08:07 +0500 Subject: [PATCH 7/7] Address review: verify watermark propagation through Flatten Adds a FlattenTest that puts a GroupByKey downstream of the Flatten with both branches sharing one key. GroupByKey fires exactly once, at the end-of-window watermark, so a single group containing the elements of both branches proves the Flatten forwarded its watermark only after every branch drained -- a premature release would fire a partial group. Also rewords a test comment to drop the reviewer attribution. --- .../streams/translation/FlattenTest.java | 72 +++++++++++++++++-- 1 file changed, 68 insertions(+), 4 deletions(-) diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java index 24a82cab374d..a5af0cdc9223 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java @@ -21,13 +21,20 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; +import java.util.ArrayList; +import java.util.Collections; import java.util.List; import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.GroupByKey; 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.PCollectionList; import org.junit.Test; @@ -91,10 +98,9 @@ public void flattenUnionsEveryBranch() { @Test public void pCollectionFeedingTwoFlattensIsSupported() { - // je-ik's #39273 case: input2 feeds both flattens. Per-Flatten branch numbering could not - // express this (input2 would need two identities); a global producer id per PCollection can — - // input2's producer stamps one id, and each flatten still holds until its own two branches - // drain. Verify both flattens produce the right union. + // input2 feeds both flattens, so its producer's watermark report is consumed by two different + // aggregators. The producer stamps its own transform id once, and each flatten holds until its + // own two branches drain. Verify both flattens produce the right union. try (SharedTestCollector left = SharedTestCollector.create(); SharedTestCollector right = SharedTestCollector.create()) { Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); @@ -122,6 +128,64 @@ public void pCollectionFeedingTwoFlattensIsSupported() { } } + /** Maps each int to {@code KV("k", int)} so a downstream GroupByKey groups all branches. */ + private static class ToKvFn extends DoFn> { + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver> out) { + out.output(KV.of("k", input)); + } + } + + /** Records each grouped result as {@code "key=[sorted values]"}. */ + private static class RecordGroupFn extends DoFn>, Void> { + private final SharedTestCollector collector; + + RecordGroupFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element KV> group) { + List values = new ArrayList<>(); + group.getValue().forEach(values::add); + Collections.sort(values); + collector.record(group.getKey() + "=" + values); + } + } + + @Test + public void watermarkPropagatesThroughFlattenAndFiresDownstreamGroupByKey() { + // GroupByKey fires exactly once, when its input watermark reaches the end of the global + // window, and the Flatten forwards its watermark only after every branch has drained. Both + // branches share the key, so a single group holding the elements of both branches proves the + // watermark propagated through the Flatten at the right time — a premature release would fire + // a partial group instead. + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection> a = + pipeline + .apply("createA", Create.of(1, 2)) + .apply("kvA", ParDo.of(new ToKvFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); + PCollection> b = + pipeline + .apply("createB", Create.of(3, 4)) + .apply("kvB", ParDo.of(new ToKvFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); + PCollectionList.of(a) + .and(b) + .apply("flatten", Flatten.pCollections()) + .apply("gbk", GroupByKey.create()) + .apply("record", ParDo.of(new RecordGroupFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List groups = collector.recorded(); + assertThat(groups.size(), is(1)); + assertThat(groups, hasItems("k=[1, 2, 3, 4]")); + } + } + @Test public void flattenOfOneBranchTwiceDuplicatesEveryElement() { // A self-flatten is a bag union with itself: every element must appear twice. The fuser folds