This is an automated email from the ASF dual-hosted git repository.
je-ik pushed a commit to branch feat/18479-kafka-streams-runner-skeleton
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to
refs/heads/feat/18479-kafka-streams-runner-skeleton by this push:
new fca14356614 [GSoC 2026] Kafka Streams runner: multi-output executable
stages (#39410)
fca14356614 is described below
commit fca14356614faa904e69633463d710bf860b63bd
Author: M Junaid Shaukat <[email protected]>
AuthorDate: Wed Jul 22 16:28:18 2026 +0500
[GSoC 2026] Kafka Streams runner: multi-output executable stages (#39410)
---
runners/kafka-streams/build.gradle | 8 +-
.../translation/ExecutableStageProcessor.java | 53 +++++++---
.../translation/ExecutableStageTranslator.java | 52 ++++++----
.../streams/translation/StageOutputProcessor.java | 93 ++++++++++++++++++
.../kafka/streams/MultiOutputStageTest.java | 87 +++++++++++++++++
.../ExecutableStageProcessorWatermarkTest.java | 5 +-
.../translation/StageOutputProcessorTest.java | 108 +++++++++++++++++++++
7 files changed, 373 insertions(+), 33 deletions(-)
diff --git a/runners/kafka-streams/build.gradle
b/runners/kafka-streams/build.gradle
index 81099560ba7..52535c6654e 100644
--- a/runners/kafka-streams/build.gradle
+++ b/runners/kafka-streams/build.gradle
@@ -87,8 +87,11 @@ dependencies {
// Known-failing @ValidatesRunner tests, excluded until the feature they need
lands.
def sickbayTests = [
- // Needs multi-output executable stages (output-tag dispatch in
ExecutableStageProcessor).
-
'org.apache.beam.sdk.transforms.FlattenTest.testFlattenMultiplePCollectionsHavingMultipleConsumers',
+ // Non-global windowing (FixedWindows, merging windows, timestamp combiners)
is not supported
+ // yet; these apply a window and assert on window-derived output, hitting a
GlobalWindow cast.
+ 'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests',
+
'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerLatest',
+
'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerEarliest',
]
tasks.register("validatesRunner", Test) {
@@ -139,6 +142,7 @@ tasks.register("validatesRunner", Test) {
// JUnit category to exclude.
includeTestsMatching 'org.apache.beam.sdk.transforms.CreateTest'
includeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest'
+ includeTestsMatching 'org.apache.beam.sdk.transforms.GroupByKeyTest*'
for (String test : sickbayTests) {
excludeTestsMatching test
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
index fb26e9e71f6..15265073dd2 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.runners.kafka.streams.translation;
+import java.util.Map;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
@@ -36,6 +37,7 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import org.apache.kafka.streams.processor.api.Processor;
import org.apache.kafka.streams.processor.api.ProcessorContext;
@@ -85,11 +87,14 @@ class ExecutableStageProcessor
// 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.
- // The element type is intentionally wildcarded: the runner does not need to
know the runtime
- // value type — the bundle factory handles all coder application at the
Fn-API boundary using
- // the PCollection coders from the ExecutableStagePayload. Pretending the
type was byte[] was
- // only safe because the Impulse output coder happens to be ByteArrayCoder.
- private final Queue<WindowedValue<?>> pendingOutputs = new
ConcurrentLinkedQueue<>();
+ // Each entry carries the output PCollection id so it can be routed to that
output's downstream on
+ // flush. The element type is intentionally wildcarded: the runner does not
need to know the
+ // runtime value type — the bundle factory handles all coder application at
the Fn-API boundary
+ // using the PCollection coders from the ExecutableStagePayload.
+ private final Queue<PendingOutput> pendingOutputs = new
ConcurrentLinkedQueue<>();
+ // Output PCollection id -> the child node (a StageOutputProcessor relay) to
forward that output
+ // to. Empty for a single-output stage, which forwards to its one downstream
directly.
+ private final Map<String, String> outputChildByPCollectionId;
// Computes this stage's input watermark from its upstream transform's
reports, holding until
// every partition of the upstream transform has reported (see
WatermarkAggregator).
@@ -114,12 +119,25 @@ class ExecutableStageProcessor
JobInfo jobInfo,
String transformId,
Set<String> upstreamTransformIds,
- MetricsContainerImpl metricsContainer) {
+ MetricsContainerImpl metricsContainer,
+ Map<String, String> outputChildByPCollectionId) {
this.stagePayload = stagePayload;
this.jobInfo = jobInfo;
this.transformId = transformId;
this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
this.metricsContainer = metricsContainer;
+ this.outputChildByPCollectionId =
ImmutableMap.copyOf(outputChildByPCollectionId);
+ }
+
+ /** A harness output element together with the id of the output PCollection
it belongs to. */
+ private static final class PendingOutput {
+ final String pCollectionId;
+ final WindowedValue<?> value;
+
+ PendingOutput(String pCollectionId, WindowedValue<?> value) {
+ this.pCollectionId = pCollectionId;
+ this.value = value;
+ }
}
@Override
@@ -182,11 +200,12 @@ class ExecutableStageProcessor
new OutputReceiverFactory() {
@Override
public <OutputT> FnDataReceiver<OutputT> create(String
pCollectionId) {
- // Outputs are queued here on harness threads and drained on the
processing thread
- // after the bundle closes.
+ // Outputs are queued here on harness threads, tagged with their
output PCollection id,
+ // and drained on the processing thread after the bundle closes.
return receivedElement -> {
if (receivedElement != null) {
- pendingOutputs.add((WindowedValue<?>) receivedElement);
+ pendingOutputs.add(
+ new PendingOutput(pCollectionId, (WindowedValue<?>)
receivedElement));
}
};
}
@@ -249,12 +268,20 @@ class ExecutableStageProcessor
}
ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
// The harness has finished the bundle (close() returned) so no further
enqueues happen.
- // Drain via poll() so each element is removed as it is forwarded.
- WindowedValue<?> output;
+ // Drain via poll() so each element is removed as it is forwarded. Each
output is routed to its
+ // own output's relay child for a multi-output stage; a single-output
stage forwards directly to
+ // its one downstream (empty routing map).
+ PendingOutput output;
while ((output = pendingOutputs.poll()) != null) {
- ctx.forward(
+ Record<byte[], KStreamsPayload<?>> outputRecord =
new Record<byte[], KStreamsPayload<?>>(
- record.key(), KStreamsPayload.data(output), record.timestamp()));
+ record.key(), KStreamsPayload.data(output.value),
record.timestamp());
+ String childNode = outputChildByPCollectionId.get(output.pCollectionId);
+ if (childNode == null) {
+ ctx.forward(outputRecord);
+ } else {
+ ctx.forward(outputRecord, childNode);
+ }
}
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
index 59c233a78bb..71e24f32c1f 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
@@ -18,9 +18,13 @@
package org.apache.beam.runners.kafka.streams.translation;
import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
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;
/**
@@ -66,23 +70,28 @@ class ExecutableStageTranslator implements
PTransformTranslator {
+ " uses user state or timers; stateful ParDo is not yet
supported by the Kafka"
+ " Streams runner.");
}
- if (transform.getOutputsMap().size() > 1) {
- // Multi-output stages (DoFns with side outputs, etc.) are a planned
follow-up — they need
- // an output-tag dispatch in the processor + per-output PCollection
routing. The current
- // rejection just fails loudly until that's wired in.
- throw new UnsupportedOperationException(
- "ExecutableStage "
- + transformId
- + " has "
- + transform.getOutputsMap().size()
- + " outputs; multi-output stages are not yet supported by the
Kafka Streams runner.");
- }
-
// The payload distinguishes the main input from side inputs, so reading
it from the payload
// is unambiguous even before we add side-input support.
String inputPCollectionId = stagePayload.getInput();
String parentProcessor =
context.getProcessorNameForPCollection(inputPCollectionId);
+ // A multi-output stage (a DoFn with side outputs, or a Read whose SDF
wrapper produces several
+ // outputs) needs each output routed to the right downstream. Since
downstream transforms are
+ // wired to a producer node by PCollection id, and that node must exist
when the stage is
+ // translated, give each output its own relay node (StageOutputProcessor)
and route to it by
+ // name. A single-output stage needs none of this — it forwards to its one
downstream directly
+ // and registers itself as that output's producer. Outputs are sorted so
the routing is
+ // deterministic across topology builds.
+ List<String> outputPCollectionIds = new
ArrayList<>(transform.getOutputsMap().values());
+ Collections.sort(outputPCollectionIds);
+ boolean multiOutput = outputPCollectionIds.size() > 1;
+ Map<String, String> outputChildByPCollectionId = new LinkedHashMap<>();
+ if (multiOutput) {
+ for (int i = 0; i < outputPCollectionIds.size(); i++) {
+ outputChildByPCollectionId.put(outputPCollectionIds.get(i),
transformId + "-output-" + i);
+ }
+ }
+
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
@@ -96,12 +105,21 @@ class ExecutableStageTranslator implements
PTransformTranslator {
context.getJobInfo(),
transformId,
ImmutableSet.of(parentProcessor),
-
context.getMetricsContainerStepMap().getContainer(transformId)),
+ context.getMetricsContainerStepMap().getContainer(transformId),
+ outputChildByPCollectionId),
parentProcessor);
- if (!transform.getOutputsMap().isEmpty()) {
- String outputPCollectionId =
Iterables.getOnlyElement(transform.getOutputsMap().values());
- context.registerPCollectionProducer(outputPCollectionId, transformId);
+ if (multiOutput) {
+ // One relay per output; downstream transforms wire to the relay, which
re-stamps the stage's
+ // watermark with the relay's own id so their watermark aggregation
stays consistent.
+ outputChildByPCollectionId.forEach(
+ (outputPCollectionId, relayName) -> {
+ topology.addProcessor(
+ relayName, () -> new StageOutputProcessor(relayName),
transformId);
+ context.registerPCollectionProducer(outputPCollectionId,
relayName);
+ });
+ } else if (!outputPCollectionIds.isEmpty()) {
+ context.registerPCollectionProducer(outputPCollectionIds.get(0),
transformId);
}
}
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java
new file mode 100644
index 00000000000..eec3f2bae08
--- /dev/null
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java
@@ -0,0 +1,93 @@
+/*
+ * 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.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.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * One output port of a multi-output {@link ExecutableStageProcessor}: a relay
node that stands in
+ * as the producer of one of the stage's output PCollections.
+ *
+ * <p>A Kafka Streams processor forwards a record either to all of its
children or to one named
+ * child, but downstream transforms are wired to a <em>producer node</em> by
PCollection id, and
+ * that node's identity has to be known when the stage is translated — before
the downstream
+ * transforms are. So each output PCollection of a multi-output stage gets its
own relay: the stage
+ * routes each harness output to the matching relay by name, and downstream
transforms wire to the
+ * relay. (A single-output stage needs none of this and forwards directly.)
+ *
+ * <p>The relay forwards data records unchanged, and on the watermark it
relabels <em>only</em> the
+ * transform id — keeping the stage instance's source partition and total
partition count intact.
+ * This is a 1:1 pass-through of one stage task's watermark (relay task {@code
i} sees stage task
+ * {@code i}), so preserving the partition identity is what lets a downstream
aggregator both know
+ * the report comes from this output (the relay's id) and still infer how many
parallel instances of
+ * the stage produced it. The relay does no aggregation of its own.
+ */
+class StageOutputProcessor
+ implements Processor<byte[], KStreamsPayload<?>, byte[],
KStreamsPayload<?>> {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(StageOutputProcessor.class);
+
+ private final String transformId;
+ private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
+
+ StageOutputProcessor(String transformId) {
+ this.transformId = transformId;
+ }
+
+ @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 = context;
+ if (ctx == null) {
+ throw new IllegalStateException("StageOutputProcessor used before
init()");
+ }
+ if (payload == null) {
+ LOG.warn(
+ "Stage output {} dropping record with null payload (external write
or tombstone)",
+ transformId);
+ return;
+ }
+ if (!payload.isWatermark()) {
+ // Data for this output: forward unchanged.
+ ctx.forward(record);
+ return;
+ }
+ // Relabel the transform id to this output port's, but keep the reporting
stage instance's
+ // partition and partition count so downstream can still infer the stage's
parallelism.
+ WatermarkPayload report = payload.asWatermark();
+ ctx.forward(
+ new Record<byte[], KStreamsPayload<?>>(
+ record.key(),
+ KStreamsPayload.watermark(
+ report.getWatermarkMillis(),
+ transformId,
+ report.getSourcePartition(),
+ report.getTotalSourcePartitions()),
+ record.timestamp()));
+ }
+}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/MultiOutputStageTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/MultiOutputStageTest.java
new file mode 100644
index 00000000000..7fe85d732f2
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/MultiOutputStageTest.java
@@ -0,0 +1,87 @@
+/*
+ * 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;
+
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Count;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.junit.Rule;
+import org.junit.Test;
+
+/**
+ * Verifies that a fused executable stage with more than one output routes
each output to its own
+ * downstream — including the case the mentor asked about, where each output
is followed by a
+ * separate GroupByKey (so each output ends up in its own repartition topic).
+ *
+ * <p>A DoFn splits its input to a main and an additional (side) output. One
output is grouped with
+ * {@code Count.perElement} (a GroupByKey under the hood, so that output lands
in its own
+ * repartition topic); the other is asserted directly (a plain in-process
output). A {@link PAssert}
+ * on each confirms the two outputs carried the right elements to the right
downstream.
+ */
+public class MultiOutputStageTest {
+
+ private static final TupleTag<Integer> EVENS = new TupleTag<Integer>() {};
+ private static final TupleTag<Integer> ODDS = new TupleTag<Integer>() {};
+
+ /** Routes even elements to the main output and odd elements to the side
output. */
+ private static class SplitByParityFn extends DoFn<Integer, Integer> {
+ @ProcessElement
+ public void processElement(@Element Integer input, MultiOutputReceiver
out) {
+ if (input % 2 == 0) {
+ out.get(EVENS).output(input);
+ } else {
+ out.get(ODDS).output(input);
+ }
+ }
+ }
+
+ private static PipelineOptions options() {
+ PipelineOptions options = PipelineOptionsFactory.create();
+ options.setRunner(TestKafkaStreamsRunner.class);
+ return options;
+ }
+
+ @Rule public final transient TestPipeline pipeline =
TestPipeline.fromOptions(options());
+
+ @Test
+ public void eachStageOutputFeedsItsOwnGroupByKey() {
+ PCollectionTuple split =
+ pipeline
+ .apply("create", Create.of(1, 2, 3, 4, 5, 6))
+ .apply(
+ "split",
+ ParDo.of(new SplitByParityFn()).withOutputTags(EVENS,
TupleTagList.of(ODDS)));
+
+ // The evens branch goes through a GroupByKey (Count.perElement), so that
output lands in its
+ // own repartition topic; the odds branch is asserted directly, exercising
a plain output.
+ PAssert.that(split.get(EVENS).apply("countEvens", Count.perElement()))
+ .containsInAnyOrder(KV.of(2, 1L), KV.of(4, 1L), KV.of(6, 1L));
+ PAssert.that(split.get(ODDS)).containsInAnyOrder(1, 3, 5);
+
+ pipeline.run();
+ }
+}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
index de3a23911ad..010d98d52e8 100644
---
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
@@ -26,6 +26,7 @@ import
org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
import org.apache.kafka.streams.processor.api.MockProcessorContext;
import org.apache.kafka.streams.processor.api.Record;
@@ -59,7 +60,9 @@ public class ExecutableStageProcessorWatermarkTest {
jobInfo,
STAGE_ID,
ImmutableSet.of(UPSTREAM_ID),
- new MetricsContainerImpl(STAGE_ID));
+ new MetricsContainerImpl(STAGE_ID),
+ // Single-output: no per-output routing (this test drives the
watermark path directly).
+ ImmutableMap.of());
}
/** A report from the upstream transform's given partition. */
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java
new file mode 100644
index 00000000000..83653f13c89
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java
@@ -0,0 +1,108 @@
+/*
+ * 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 static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import java.util.List;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.kafka.streams.processor.api.MockProcessorContext;
+import
org.apache.kafka.streams.processor.api.MockProcessorContext.CapturedForward;
+import org.apache.kafka.streams.processor.api.Record;
+import org.junit.Test;
+
+/**
+ * Unit tests for {@link StageOutputProcessor}, the per-output relay of a
multi-output stage: it
+ * forwards data unchanged and relabels the watermark's transform id while
preserving the reporting
+ * stage instance's partition and partition count.
+ *
+ * <p>The partition preservation is what makes the watermark propagate
correctly once the multi-
+ * output stage runs in several parallel instances (per je-ik's review): each
stage task's report
+ * must reach downstream carrying its own {@code (partition,
totalPartitions)}, so the downstream
+ * aggregator can tell how many parallel instances produced the output. A
{@link
+ * MockProcessorContext} lets the reports from different stage partitions be
fed directly, which a
+ * single-instance {@code TopologyTestDriver} cannot do.
+ */
+public class StageOutputProcessorTest {
+
+ private static final String RELAY_ID = "stage-output-0";
+ private static final String STAGE_ID = "stage";
+
+ private static Record<byte[], KStreamsPayload<?>> watermark(
+ long millis, int sourcePartition, int totalPartitions) {
+ return new Record<>(
+ new byte[0],
+ KStreamsPayload.watermark(millis, STAGE_ID, sourcePartition,
totalPartitions),
+ 0L);
+ }
+
+ @Test
+ public void watermarkKeepsPartitionIdentityAndRelabelsTransformId() {
+ MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
+ StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID);
+ processor.init(ctx);
+
+ // A report from partition 1 of a 3-instance stage.
+ processor.process(watermark(500L, 1, 3));
+
+ List<CapturedForward<? extends byte[], ? extends KStreamsPayload<?>>>
forwarded =
+ ctx.forwarded();
+ assertThat(forwarded.size(), is(1));
+ WatermarkPayload out = forwarded.get(0).record().value().asWatermark();
+ assertThat(out.getWatermarkMillis(), is(500L));
+ // Relabeled to the relay's id, but the stage instance's partition and
count are preserved.
+ assertThat(out.getTransformId(), is(RELAY_ID));
+ assertThat(out.getSourcePartition(), is(1));
+ assertThat(out.getTotalSourcePartitions(), is(3));
+ }
+
+ @Test
+ public void distinctStagePartitionsStayDistinctDownstream() {
+ MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
+ StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID);
+ processor.init(ctx);
+
+ processor.process(watermark(100L, 0, 3));
+ processor.process(watermark(200L, 2, 3));
+
+ List<CapturedForward<? extends byte[], ? extends KStreamsPayload<?>>>
forwarded =
+ ctx.forwarded();
+ assertThat(forwarded.size(), is(2));
+
assertThat(forwarded.get(0).record().value().asWatermark().getSourcePartition(),
is(0));
+
assertThat(forwarded.get(1).record().value().asWatermark().getSourcePartition(),
is(2));
+
assertThat(forwarded.get(0).record().value().asWatermark().getTotalSourcePartitions(),
is(3));
+ }
+
+ @Test
+ public void dataIsForwardedUnchanged() {
+ MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
+ StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID);
+ processor.init(ctx);
+
+ WindowedValue<byte[]> element = WindowedValues.valueInGlobalWindow(new
byte[] {7});
+ KStreamsPayload<byte[]> data = KStreamsPayload.data(element);
+ processor.process(new Record<>(new byte[0], data, 0L));
+
+ List<CapturedForward<? extends byte[], ? extends KStreamsPayload<?>>>
forwarded =
+ ctx.forwarded();
+ assertThat(forwarded.size(), is(1));
+ assertThat(forwarded.get(0).record().value(), is(data));
+ }
+}
