[GSoC 2026] Kafka Streams runner: Flatten support#39273
Conversation
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.
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request adds support for the Beam Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements the Flatten primitive for the Kafka Streams runner by introducing FlattenProcessor and FlattenTranslator, and updating watermark tracking to support multi-branch fan-in via SourceStamp. The reviewer feedback highlights several critical issues: potential non-deterministic behavior and watermark stalls due to duplicate input PCollections in registerFlattenSourceStamps, consistency issues in parent processor resolution, unsupported multi-Flatten consumption of a single PCollection, and a potential NullPointerException when processing tombstone records with null payloads.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| int totalPartitions = transform.getInputsMap().size(); | ||
| int sourcePartition = 0; | ||
| for (String inputPCollectionId : transform.getInputsMap().values()) { | ||
| context.registerFlattenSourceStamp(inputPCollectionId, sourcePartition, totalPartitions); | ||
| sourcePartition++; | ||
| } |
There was a problem hiding this comment.
There are two issues here:
- Non-deterministic iteration order: The iteration order of
transform.getInputsMap().values()is not guaranteed to be deterministic across different JVMs or runs. Since the topology is built independently on each worker, this can lead to different workers assigning differentsourcePartitionindices to the same PCollections. - Duplicate input PCollections: If a PCollection is flattened with itself (e.g.,
PCollectionList.of(pc).and(pc)), the same PCollection ID will appear multiple times. The current loop will overwrite theSourceStampmapping, resulting in only one partition being registered and stamped. Consequently, theWatermarkManagerinFlattenProcessorwill wait forever for the other partition, causing the watermark to get stuck.
Deduplicating and sorting the input PCollection IDs resolves both issues.
java.util.List<String> uniqueInputs =
new java.util.ArrayList<>(new java.util.HashSet<>(transform.getInputsMap().values()));
java.util.Collections.sort(uniqueInputs);
int totalPartitions = uniqueInputs.size();
int sourcePartition = 0;
for (String inputPCollectionId : uniqueInputs) {
context.registerFlattenSourceStamp(inputPCollectionId, sourcePartition, totalPartitions);
sourcePartition++;
}| List<String> parentProcessors = new ArrayList<>(); | ||
| for (String inputPCollectionId : transform.getInputsMap().values()) { | ||
| parentProcessors.add(context.getProcessorNameForPCollection(inputPCollectionId)); | ||
| } |
There was a problem hiding this comment.
To match the deduplicated and sorted input PCollections used for watermark tracking, we should also deduplicate and sort the inputs when building the parent processors list. This ensures consistency and avoids passing duplicate parent names to Topology.addProcessor.
| List<String> parentProcessors = new ArrayList<>(); | |
| for (String inputPCollectionId : transform.getInputsMap().values()) { | |
| parentProcessors.add(context.getProcessorNameForPCollection(inputPCollectionId)); | |
| } | |
| java.util.List<String> uniqueInputs = | |
| new java.util.ArrayList<>(new java.util.HashSet<>(transform.getInputsMap().values())); | |
| java.util.Collections.sort(uniqueInputs); | |
| List<String> parentProcessors = new ArrayList<>(); | |
| for (String inputPCollectionId : uniqueInputs) { | |
| parentProcessors.add(context.getProcessorNameForPCollection(inputPCollectionId)); | |
| } |
| pCollectionIdToSourceStamp.put( | ||
| pCollectionId, new SourceStamp(sourcePartition, totalPartitions)); |
There was a problem hiding this comment.
If a PCollection is consumed by multiple different Flatten transforms, the SourceStamp mapping will be overwritten, leading to incorrect watermark tracking. We should explicitly detect this conflict and throw an UnsupportedOperationException to prevent silent correctness bugs.
SourceStamp existing =
pCollectionIdToSourceStamp.put(
pCollectionId, new SourceStamp(sourcePartition, totalPartitions));
if (existing != null) {
throw new UnsupportedOperationException(
"PCollection "
+ pCollectionId
+ " is consumed by multiple Flattens, which is not supported yet due to watermark tracking limitations.");
}| ProcessorContext<byte[], KStreamsPayload<?>> ctx = checkInitialized(context); | ||
| if (!payload.isWatermark()) { | ||
| // Data: the union of the parents' data streams is the flatten — forward unchanged. |
There was a problem hiding this comment.
In Kafka Streams, tombstone records can have a null value. Accessing payload.isWatermark() without a null check will throw a NullPointerException. Adding a defensive null check to forward tombstone records downstream is recommended.
KStreamsPayload<?> payload = record.value();
ProcessorContext<byte[], KStreamsPayload<?>> ctx = checkInitialized(context);
if (payload == null) {
ctx.forward(record);
return;
}
if (!payload.isWatermark()) {…n 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.
|
Assigning reviewers: R: @tvalentyn added as fallback since no labels match configuration Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
je-ik
left a comment
je-ik
left a comment
There was a problem hiding this comment.
Reading through the code I realized we overlooked something.
Suppose the following case:
Pipeline p = Pipeline.create();
PCollection<..> input1 = p.apply(...)
PCollection<...> input2 = p.apply(...)
PCollection<...> input3 = p.apply(...)
PCollectionList l1 = PCollectionList.of(input1).and(input2);
l1.apply(Flatten.pCollections());
PCollectionList l2 = PCollectionList.of(input2).and(input3);
l2.apply(Flatten.pCollection());input2 is "1 of 2" for l1, but "0 of 2" for l2.
So the static numbering does not work here. What we actually need is a static numbering of PTransforms:
input1 = 1
input2 = 2
input3 = 3
and information that flatten l1 should wait for watermark from 1 and 2, while flatten l2 for watemark from 2 and 3. This way, the watermark can be generated without worrying about WHO and HOW consumes it.
| // 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; |
There was a problem hiding this comment.
We should use a distinct naming convetion for there "partitions", because we already have Kafka partitions and this creates a confusion.
| record.key(), | ||
| KStreamsPayload.watermark(advanced.getMillis(), 0, 1), | ||
| record.timestamp())); | ||
| } |
There was a problem hiding this comment.
This logic seems like it is a general watermark propagation (will be the same for GBK, Combine, ExecutableStage), we could make this part of the WatermarkManager.
| // about and reproduce. registerFlattenSourceStamp fails fast on a duplicate (self-flatten). | ||
| List<String> inputPCollectionIds = new ArrayList<>(transform.getInputsMap().values()); | ||
| Collections.sort(inputPCollectionIds); | ||
| int totalPartitions = inputPCollectionIds.size(); |
There was a problem hiding this comment.
Ah! Here is the confusion! There are Kafka partitions and this "partitions" here is the number of inputs.
There was a problem hiding this comment.
Good catch, thanks, you're right, per-Flatten (i of N) can't express a PCollection that feeds two Flattens. Reworked it the way you described: each Flatten-input PCollection gets one stable global producer 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. So input2 stamps one id, l1 waits for {1,2}, l2 for {2,3}, and the shared input just works. Nice side effect it also removes the fan-out hazard, since any single-input consumer always sees 1 of 1. Added a test for your exact case (input2 feeding both flattens produces both unions). Self-flatten (Flatten.of(pc,pc)) is still rejected, since one producer can't be two branches.
…wo 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.
2 participants
Summary
Adds a translator for Beam's Flatten primitive (
beam:transform:flatten:v1) -- the union of N input PCollections into one. Part of #18479; the last primitive needed before the first PAssert-based@ValidatesRunnertest (PAssert'sGroupGloballyusesGBK + Flatten, no side inputs).What's here:
FlattenProcessor-- forwards data through, and owns its output watermark theway GroupByKey does: a
WatermarkManagerover the input branches, emitting itsown single-source
(0 of 1)watermark only when themin()advances. Thisholds the watermark back until every branch has drained.
gets one stable global producer id; its producing transform stamps that id on
its watermark (always as a single source,
1 of 1), because Kafka Streams doesnot tell a processor which parent forwarded a record. Each Flatten then holds
using its own input count. A PCollection that feeds two Flattens reports one id
and each Flatten still waits only for its own branches -- so
input2inl1 = Flatten(input1, input2)/l2 = Flatten(input2, input3)works (thanks@je-ik for catching that per-Flatten numbering could not express this). It also
means a single-input consumer always sees
1 of 1, so there is no fan-outhazard.
FlattenTranslatorwires the N parents to one node;FLATTENregistered in thetranslator map.
Flatten.of(pc, pc)) is rejected with a clearUnsupportedOperationException-- its single producer cannot be two branches,and Kafka Streams cannot wire the same parent to a child twice. Proper support
is a follow-up.
Scope: only
ExecutableStageproducers stamp the producer id so far, which covers PAssert'sGroupGlobally(its Flatten inputs are stages). ARead/Impulse/GBKoutput feeding a Flatten directly still reports id 0 -- thatneeds the same stamp wiring and is a follow-up when those tests are enabled.
Tests:
FlattenTestunions twoCreate -> ParDobranches and asserts all four elements arrive (which only holds because the watermark is held until both branches drain), verifies a PCollection feeding two Flattens produces bothunions, and asserts a self-flatten is rejected.