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 fc36301e6c2 [GSoC 2026] Kafka Streams runner: CombineTest coverage and
two review follow-ups (#39610)
fc36301e6c2 is described below
commit fc36301e6c2e6e98a9758eea815e7f4f8a03f966
Author: M Junaid Shaukat <[email protected]>
AuthorDate: Tue Aug 4 14:54:53 2026 +0500
[GSoC 2026] Kafka Streams runner: CombineTest coverage and two review
follow-ups (#39610)
* [GSoC 2026] Kafka Streams runner: CombineTest coverage and two review
follow-ups
Enables CombineTest in the ValidatesRunner suite, taking it from 49 to 59
tests. Combine was expected to work without a translator of its own, since
the fuser expands Combine.perKey into a GroupByKey with the combining logic
running as ordinary ParDos in the SDK harness, but nothing exercised that.
BasicTests passes in full, including hot-key fanout and the
accumulation-mode
variant, and WindowingTests contributes the fixed-window and empty-window
cases. The remainder falls out on category excludes the task already
declares. testSessionsCombine is sickbayed alongside the existing merging
windows entry, and it is the only Combine failure.
Corrects the Flatten partition-count comment, which asserted that the inputs
are co-partitioned and so implied the Math.max over them was redundant.
Neither half held. The max is not a no-op in principle: Kafka Streams merges
the subtopologies of every parent a processor is wired to and gives the
result as many tasks as its largest source topic has partitions. But the
mismatched shape does not reach this translator, because the fuser folds
such a Flatten into the harness stage, and the runner Flattens that do
arrive come from the fuser deduplicating partial outputs of one PCollection.
FlattenParallelismTest records that, so a change letting the mismatched
shape
through starts failing there rather than producing a pipeline that stalls
waiting for a watermark report that never comes.
Guards the null record key in GroupByKeyBroadcastPartitioner.partitions().
partition() already guarded it, but partitions() is the method Kafka Streams
calls and it hashed the key unguarded.
---
runners/kafka-streams/build.gradle | 4 +-
.../streams/translation/FlattenTranslator.java | 10 +-
.../GroupByKeyBroadcastPartitioner.java | 8 ++
.../translation/FlattenParallelismTest.java | 141 +++++++++++++++++++++
4 files changed, 160 insertions(+), 3 deletions(-)
diff --git a/runners/kafka-streams/build.gradle
b/runners/kafka-streams/build.gradle
index d1326c079e5..616257ddd78 100644
--- a/runners/kafka-streams/build.gradle
+++ b/runners/kafka-streams/build.gradle
@@ -113,8 +113,9 @@ def sickbayTests = [
// Merging (session) windows are not supported yet: ReduceFnRunner drives
them through a merging
// window set that moves per-window state as windows merge, which this first
windowing pass does
// not implement. Non-merging windows (fixed, sliding), the default trigger
and timestamp
- // combiners do work. Lands with the follow-up windowing PR.
+ // combiners do work, for both GroupByKey and Combine. Lands with the
follow-up windowing PR.
'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests.testGroupByKeyMergingWindows',
+
'org.apache.beam.sdk.transforms.CombineTest$WindowingTests.testSessionsCombine',
// A DoFn whose @StartBundle throws never gets to report its error:
SdkHarnessClient.newBundle
// sends the ProcessBundleRequest and then blocks in
GrpcDataService.createOutboundAggregator
// waiting for the SDK harness to open its data stream, which a bundle that
failed during setup
@@ -177,6 +178,7 @@ tasks.register("validatesRunner", Test) {
includeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest'
includeTestsMatching 'org.apache.beam.sdk.transforms.GroupByKeyTest*'
includeTestsMatching 'org.apache.beam.sdk.transforms.ParDoTest*'
+ includeTestsMatching 'org.apache.beam.sdk.transforms.CombineTest*'
for (String test : sickbayTests) {
excludeTestsMatching test
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
index a5c8ce05bae..793c1e1dc53 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
@@ -57,8 +57,14 @@ class FlattenTranslator implements PTransformTranslator {
Set<String> seenInputs = new HashSet<>();
List<String> parentProcessors = new ArrayList<>();
Set<String> upstreamTransformIds = new HashSet<>();
- // Kafka Streams puts a processor and the parents it is wired to in one
subtopology, so the
- // inputs are co-partitioned and this Flatten runs at their partition
count.
+ // How many instances this Flatten runs as. Kafka Streams merges the
subtopologies of every
+ // parent a processor is wired to and gives the merged subtopology as many
tasks as its largest
+ // source topic has partitions, so the max is what that comes to. In
practice the inputs agree:
+ // a Flatten whose branches could disagree — one through a GroupByKey, one
straight from a
+ // source — is fused into the harness stage instead of becoming a node
here, and the runner
+ // Flattens that do reach this translator come from the fuser
deduplicating partial outputs of
+ // one PCollection. The max is kept as the cheap conservative choice
rather than asserting that
+ // agreement, which is not enforced anywhere.
int partitionCount = 1;
for (String inputPCollectionId : transform.getInputsMap().values()) {
if (!seenInputs.add(inputPCollectionId)) {
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java
index b5ddcf2536a..3c775c86f14 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java
@@ -56,6 +56,14 @@ class GroupByKeyBroadcastPartitioner<T> implements
StreamPartitioner<byte[], KSt
}
return Optional.of(all);
}
+ if (key == null) {
+ // A keyless record has no partition it must go to, so leave the choice
to Kafka rather than
+ // hashing a null or pinning one partition: an empty Optional tells
Kafka Streams no explicit
+ // partition was chosen, and the producer's default partitioner spreads
keyless records over
+ // the topic instead of piling them onto one. This is the method Kafka
Streams calls, so the
+ // null has to be handled here and not only in partition() above.
+ return Optional.empty();
+ }
int partition = Utils.toPositive(Utils.murmur2(key)) % numPartitions;
return Optional.of(Collections.singleton(partition));
}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenParallelismTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenParallelismTest.java
new file mode 100644
index 00000000000..601fa40e679
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenParallelismTest.java
@@ -0,0 +1,141 @@
+/*
+ * 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.ArrayList;
+import java.util.List;
+import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions;
+import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.Flatten;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionList;
+import org.apache.kafka.streams.TopologyDescription;
+import org.junit.Test;
+
+/**
+ * Pins down what happens to a Flatten whose branches would run at different
parallelisms — one
+ * through a GroupByKey and so at the shuffle's parallelism, one straight from
a source and so a
+ * single instance.
+ *
+ * <p>This matters because a Flatten runs as one set of tasks over all of its
inputs. Kafka Streams
+ * merges the subtopologies of every parent a processor is wired to and gives
the result as many
+ * tasks as its largest source topic has partitions, so a parent with fewer
partitions would only
+ * produce on some of those tasks and the rest would wait forever for a
watermark report from it.
+ *
+ * <p>That does not arise, and these tests record why: the fuser folds such a
Flatten into the SDK
+ * harness stages rather than leaving a node for the runner to translate, so
the branches never
+ * share a subtopology and no Flatten node exists to run at a single
parallelism. The Flattens that
+ * do reach {@link FlattenTranslator} come from the fuser deduplicating
partial outputs of a single
+ * PCollection. If a change ever makes the mismatched shape reach the
translator, these tests start
+ * failing and the partition-count handling there needs revisiting.
+ */
+public class FlattenParallelismTest {
+
+ /** The name given to the Flatten below, which no topology node should be
derived from. */
+ private static final String FLATTEN_NAME = "merge";
+
+ private static class ToKvFn extends DoFn<Integer, KV<String, Integer>> {
+ @ProcessElement
+ public void processElement(@Element Integer input,
OutputReceiver<KV<String, Integer>> out) {
+ out.output(KV.of("k", input));
+ }
+ }
+
+ private static class UngroupFn extends DoFn<KV<String, Iterable<Integer>>,
Integer> {
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, Iterable<Integer>> group, OutputReceiver<Integer>
out) {
+ for (int value : group.getValue()) {
+ out.output(value);
+ }
+ }
+ }
+
+ private static Pipeline mixedParallelismFlatten(int internalParallelism) {
+ KafkaStreamsPipelineOptions options =
+
KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class);
+ options.setInternalParallelism(internalParallelism);
+ Pipeline pipeline = Pipeline.create(options);
+
+ // Through a GroupByKey, so this branch runs at the shuffle's parallelism.
+ PCollection<Integer> shuffled =
+ pipeline
+ .apply("createGrouped", Create.of(1, 2, 3))
+ .apply("toKv", ParDo.of(new ToKvFn()))
+ .apply("group", GroupByKey.create())
+ .apply("ungroup", ParDo.of(new UngroupFn()));
+
+ // Straight from a source, so this branch is a single instance.
+ PCollection<Integer> direct = pipeline.apply("createDirect", Create.of(4,
5, 6));
+
+ PCollectionList.of(shuffled).and(direct).apply("merge",
Flatten.pCollections());
+ return pipeline;
+ }
+
+ /** Every processor node in the topology, across all subtopologies. */
+ private static List<String> processorNames(TopologyDescription description) {
+ List<String> names = new ArrayList<>();
+ for (TopologyDescription.Subtopology subtopology :
description.subtopologies()) {
+ for (TopologyDescription.Node node : subtopology.nodes()) {
+ if (node instanceof TopologyDescription.Processor) {
+ names.add(node.name());
+ }
+ }
+ }
+ return names;
+ }
+
+ private static void assertFlattenWasFusedAway(TopologyDescription
description) {
+ // No node stands for the Flatten. If one did, it would be wired to both
branches and so would
+ // run over a merged subtopology whose smaller-parallelism parent could
not reach all of its
+ // instances.
+ for (String name : processorNames(description)) {
+ assertThat(
+ "no processor node should stand for the Flatten, but found " + name,
+ name.contains(FLATTEN_NAME),
+ is(false));
+ }
+ // The branches stay in separate subtopologies for the same reason: the
source-fed branch, the
+ // one behind the shuffle, and the second source-fed branch.
+ assertThat(description.subtopologies().size(), is(3));
+ }
+
+ @Test
+ public void
branchesAtDifferentParallelismsAreFusedRatherThanLeftToTheRunner() {
+ assertFlattenWasFusedAway(
+
KafkaStreamsTestRunner.translate(mixedParallelismFlatten(4)).getTopology().describe());
+ }
+
+ @Test
+ public void theSameHoldsAtASingleParallelism() {
+ // Whether the Flatten is fused is a property of the fused graph, not of
the parallelism, so
+ // the shape is the same either way — which is why raising the parallelism
cannot introduce a
+ // Flatten node over mismatched branches.
+ assertFlattenWasFusedAway(
+
KafkaStreamsTestRunner.translate(mixedParallelismFlatten(1)).getTopology().describe());
+ }
+}
