[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.
…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.
| 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.
| transformId, | ||
| () -> new ExecutableStageProcessor(stagePayload, context.getJobInfo()), | ||
| () -> | ||
| new ExecutableStageProcessor(stagePayload, context.getJobInfo(), watermarkSourceId, 1), |
There was a problem hiding this comment.
We are definitely confusing something here. We cannot pass "watermarkSourceId" to sourcePartition. These are distinct things.
There was a problem hiding this comment.
You're right, I was mixing two different things. Reworked it: the watermark payload now carries all three fields, the producing transform's id (new transform_id proto field), the source partition, and the partition count. Every producer stamps its own transform id and its own partition (0 of 1 while single-instance), without knowing who consumes it. On the consuming side there's a new WatermarkAggregator used by ExecutableStage, GBK and Flatten (CombinePerKey later): it gets the expected upstream transform ids from the pipeline graph at translation time, tracks each upstream's partitions with its own WatermarkManager, and advances to the min across upstreams once all have reported. A report from an unexpected transform now throws, so a wiring mistake like this one fails loudly instead of silently working.
On why tests didn't catch it: everything is single-instance (0 of 1), so the two versions behaved identically in every topology we can currently build, it would only have become observable with multi-instance transforms. The stage watermark test now also asserts the forwarded report carries the stage's own transform id.
One side finding: a self-flatten (Flatten.of(pc, pc)) is handled by the fuser, it folds the Flatten into the consuming harness stage, which does the duplication. The previous revision falsely rejected it; there's now a test asserting each element appears twice.
…load 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.
je-ik
left a comment
je-ik
left a comment
There was a problem hiding this comment.
This is nearly perfect with few points:
a) we should resolve the open comments (either commenting on why we think it needs to action, or take an action and resolve or comment on them)
b) the WatermarkAggregator definitely deserves its own test suite so that we can explicitly test all corner cases on this separate class
Thanks for the good work!
…covery 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.
Thanks! Added the WatermarkAggregator test suite (holds across upstreams and partitions, min aggregation, monotonicity, duplicate broadcasts, terminal MAX, repartition re-opening the hold, unexpected-transform fail-fast) and the null-payload handling, warn and drop in all four payload-consuming processors. Also went through the open threads and resolved them with a note each. |
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.