[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). |
1 participant
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.
(i of N)the WatermarkManager needs isstamped upstream, by the producing
ExecutableStagewhen its output feeds aFlatten (a translation pre-pass records which PCollections are Flatten inputs),
because Kafka Streams does not tell a processor which parent forwarded a record.
Discussed with je-ik on Slack -- stamping at the producer is free since Beam
fuses it and the runner (not user code) generates the watermark. The
(i, N)only matters to Flatten: in the shapes this PR covers (incl. PAssert's
GroupGlobally) a stamped PCollection feeds only the Flatten, so there is nofan-out. The general fan-out case -- a stamped output also feeding a single-input
consumer, which would then need to ignore the
(i, N)-- is deferred.FlattenTranslatorwires the N parents to one node;FLATTENregistered in thetranslator map.
Flattens) is rejected with a clear
UnsupportedOperationException-- its singleproducer cannot stamp two branch identities. Proper support is a follow-up.
Scope: only
ExecutableStageproducers stamp the branch identity so far, which covers PAssert'sGroupGlobally(its Flatten inputs are stages). ARead/Impulse/GBKoutput feeding a Flatten directly still reports(0 of 1)--that needs the same stamp wiring and is a follow-up when those tests are enabled.
Tests:
FlattenTestunions twoCreate -> ParDobranches and asserts a downstream ParDo sees all four elements (which only holds because the watermark is held until both branches drain), and asserts a self-flatten is rejected.