gemini-code-assist[bot] commented on code in PR #39273:
URL: https://github.com/apache/beam/pull/39273#discussion_r3554609854
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java:
##########
@@ -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++;
+ }
Review Comment:
There are two issues here:
1. **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 different
`sourcePartition` indices to the same PCollections.
2. **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 the `SourceStamp`
mapping, resulting in only one partition being registered and stamped.
Consequently, the `WatermarkManager` in `FlattenProcessor` will wait forever
for the other partition, causing the watermark to get stuck.
Deduplicating and sorting the input PCollection IDs resolves both issues.
```java
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++;
}
```
##########
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.
+ *
+ * <p><b>Data</b> records are forwarded straight through unchanged — the merge
of the N parents'
+ * data streams <em>is</em> the flatten.
+ *
+ * <p><b>Watermark</b> 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
+ * <em>every</em> input branch has reported, so a downstream GroupByKey does
not fire before all
+ * flattened branches are drained.
+ *
+ * <p>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<?>, 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<byte[], KStreamsPayload<?>> context;
+
+ @Override
+ public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
+ this.context = context;
+ }
+
+ @Override
+ public void process(Record<byte[], KStreamsPayload<?>> record) {
+ KStreamsPayload<?> payload = record.value();
+ ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
+ if (!payload.isWatermark()) {
+ // Data: the union of the parents' data streams is the flatten — forward
unchanged.
Review 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.
```java
KStreamsPayload<?> payload = record.value();
ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
if (payload == null) {
ctx.forward(record);
return;
}
if (!payload.isWatermark()) {
```
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java:
##########
@@ -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));
Review 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.
```java
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.");
}
```
##########
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.
+ *
+ * <p>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}.
+ *
+ * <p>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<String> parentProcessors = new ArrayList<>();
+ for (String inputPCollectionId : transform.getInputsMap().values()) {
+
parentProcessors.add(context.getProcessorNameForPCollection(inputPCollectionId));
+ }
Review 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`.
```suggestion
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));
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
