[GSoC 2026] Kafka Streams runner #39273: Kafka Streams runner: Flatten support

Author: junaiddshaukat

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.

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.
  *
- * <p>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.
+ * <p>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.
  *
  * <p>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,6 +71,9 @@ class ExecutableStageProcessor
 
   private final RunnerApi.ExecutableStagePayload stagePayload;
   private final JobInfo jobInfo;
+  // 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.
@@ -79,9 +83,9 @@ class ExecutableStageProcessor
   // only safe because the Impulse output coder happens to be ByteArrayCoder.
   private final Queue<WindowedValue<?>> 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;
 
@@ -90,9 +94,20 @@ class ExecutableStageProcessor
   private @Nullable StageBundleFactory stageBundleFactory;
   private @Nullable RemoteBundle currentBundle;
 
-  ExecutableStageProcessor(RunnerApi.ExecutableStagePayload stagePayload, JobInfo jobInfo) {
+  /**
+   * @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,
+      String transformId,
+      Set<String> upstreamTransformIds) {
     this.stagePayload = stagePayload;
     this.jobInfo = jobInfo;
+    this.transformId = transformId;
+    this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
   }
 
   @Override
@@ -116,18 +131,21 @@ private void ensureStageBundleFactory() {
   @Override
   public void process(Record<byte[], KStreamsPayload<?>> 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.
       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());
@@ -203,14 +221,16 @@ private void closeBundleAndFlush(Record<byte[], KStreamsPayload<?>> record) {
   }
 
   private void forwardWatermark(Record<byte[], KStreamsPayload<?>> 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).
+    // 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<byte[], KStreamsPayload<?>> ctx = checkInitialized(context);
     ctx.forward(
         new Record<byte[], KStreamsPayload<?>>(
-            record.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1), record.timestamp()));
+            record.key(),
+            KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1),
+            record.timestamp()));
   }
 
   @Override

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;
 
@@ -83,9 +84,14 @@ public void translate(
     String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId);
 
     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()),
+        () ->
+            new ExecutableStageProcessor(
+                stagePayload, context.getJobInfo(), transformId, ImmutableSet.of(parentProcessor)),
         parentProcessor);
 
     if (!transform.getOutputsMap().isEmpty()) {

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java

@@ -0,0 +1,119 @@
+/*
+ * 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.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;
+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
+ * 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 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
+ * <em>every</em> input branch has reported, so a downstream GroupByKey does not fire before all
+ * flattened branches are drained.
+ *
+ * <p>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 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<?>, 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
+  // 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<byte[], KStreamsPayload<?>> context;
+
+  /**
+   * @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<String> upstreamTransformIds) {
+    this.transformId = transformId;
+    this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
+  }
+
+  @Override
+  public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
+    this.context = context;
+  }
+
+  @Override
+  public void process(Record<byte[], KStreamsPayload<?>> 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<byte[], KStreamsPayload<?>> ctx = checkInitialized(context);
+    if (!payload.isWatermark()) {
+      // Data: the union of the parents' data streams is the flatten — forward unchanged.
+      ctx.forward(record);
+      return;
+    }
+    watermarkAggregator.observe(payload.asWatermark());
+    Instant advanced = watermarkAggregator.advance();
+    if (advanced.isAfter(lastForwardedWatermark)) {
+      lastForwardedWatermark = advanced;
+      // 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<byte[], KStreamsPayload<?>>(
+              record.key(),
+              KStreamsPayload.watermark(advanced.getMillis(), transformId, 0, 1),
+              record.timestamp()));
+    }
+  }
+
+  private static ProcessorContext<byte[], KStreamsPayload<?>> checkInitialized(
+      @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context) {
+    if (context == null) {
+      throw new IllegalStateException("FlattenProcessor used before init()");
+    }
+    return context;
+  }
+}

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java

@@ -0,0 +1,82 @@
+/*
+ * 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.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;
+
+/**
+ * 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 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.
+ *
+ * <p>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 {
+
+  @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());
+
+    Set<String> seenInputs = new HashSet<>();
+    List<String> parentProcessors = new ArrayList<>();
+    Set<String> upstreamTransformIds = new HashSet<>();
+    for (String inputPCollectionId : transform.getInputsMap().values()) {
+      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);
+    }
+
+    Topology topology = context.getTopology();
+    topology.addProcessor(
+        transformId,
+        () -> new FlattenProcessor(transformId, upstreamTransformIds),
+        parentProcessors.toArray(new String[0]));
+
+    context.registerPCollectionProducer(outputPCollectionId, transformId);
+  }
+}