From e1d47c649cd7be3f0b07b3771f3a80ad9f788f80 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Tue, 4 Aug 2026 14:27:01 +0500 Subject: [PATCH 1/2] [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 +- .../translation/FlattenTranslator.java | 10 +- .../GroupByKeyBroadcastPartitioner.java | 5 +- .../translation/FlattenParallelismTest.java | 110 ++++++++++++++++++ 4 files changed, 125 insertions(+), 4 deletions(-) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenParallelismTest.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index d1326c079e5d..616257ddd788 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 a5c8ce05baeb..793c1e1dc530 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 @@ public void translate( Set seenInputs = new HashSet<>(); List parentProcessors = new ArrayList<>(); Set 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 b5ddcf2536a2..030df30c6dc3 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,7 +56,10 @@ public Optional> partitions( } return Optional.of(all); } - int partition = Utils.toPositive(Utils.murmur2(key)) % numPartitions; + // A keyless record — a stateless stage carries no key — has nowhere in particular to go, so + // send it to partition 0 rather than hashing a null. This is the method Kafka Streams calls, + // so the guard has to be here and not only on partition() above. + int partition = key == null ? 0 : 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 000000000000..1e6cab197c00 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenParallelismTest.java @@ -0,0 +1,110 @@ +/* + * 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 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.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. + * + *

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. + * + *

That does not arise, and this test is here to record why: the fuser folds such a Flatten into + * the SDK harness stage rather than leaving a node for the runner to translate. 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, this test starts + * failing and the partition-count handling there needs revisiting. + */ +public class FlattenParallelismTest { + + private static class ToKvFn extends DoFn> { + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver> out) { + out.output(KV.of("k", input)); + } + } + + private static class UngroupFn extends DoFn>, Integer> { + @ProcessElement + public void processElement( + @Element KV> group, OutputReceiver 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 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 direct = pipeline.apply("createDirect", Create.of(4, 5, 6)); + + PCollectionList.of(shuffled).and(direct).apply("merge", Flatten.pCollections()); + return pipeline; + } + + @Test + public void branchesAtDifferentParallelismsAreFusedRatherThanLeftToTheRunner() { + // Translating is the assertion: the mismatched shape does not reach FlattenTranslator, because + // the fuser absorbs this Flatten into the harness stage. Were it to arrive there, the runner + // would build one Flatten node over branches of differing parallelism and stall. + KafkaStreamsTranslationContext context = + KafkaStreamsTestRunner.translate(mixedParallelismFlatten(4)); + + assertThat(context.getTopology().describe().subtopologies().isEmpty(), is(false)); + } + + @Test + public void theSameShapeTranslatesAtASingleParallelism() { + KafkaStreamsTranslationContext context = + KafkaStreamsTestRunner.translate(mixedParallelismFlatten(1)); + + assertThat(context.getTopology().describe().subtopologies().isEmpty(), is(false)); + } +} From 76e214468feffbccc4137709ca7f355e78b3facc Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Tue, 4 Aug 2026 14:49:46 +0500 Subject: [PATCH 2/2] Address review: spread keyless records, assert the topology shape A null record key was being sent to partition 0, which is a fixed partition rather than no choice at all: keyless records would have piled onto one partition of the topic. Returning an empty Optional instead tells Kafka Streams that no explicit partition was chosen, and the producer's default partitioner spreads them, which is what a record with no key should get. The Flatten tests asserted only that the topology was non-empty, which would have passed whatever the fuser did with the Flatten. They now assert the property they exist for: that no processor node stands for the Flatten, and that the branches remain in three separate subtopologies rather than being merged into one. Were a Flatten node to appear over branches of differing parallelism, the branch with fewer partitions could not reach all of its instances and the rest would wait forever for a watermark report. Both the mismatched and the single-parallelism case assert the same shape, since whether the Flatten is fused is a property of the fused graph and not of the parallelism. --- .../GroupByKeyBroadcastPartitioner.java | 13 ++-- .../translation/FlattenParallelismTest.java | 61 ++++++++++++++----- 2 files changed, 55 insertions(+), 19 deletions(-) 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 030df30c6dc3..3c775c86f14c 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,10 +56,15 @@ public Optional> partitions( } return Optional.of(all); } - // A keyless record — a stateless stage carries no key — has nowhere in particular to go, so - // send it to partition 0 rather than hashing a null. This is the method Kafka Streams calls, - // so the guard has to be here and not only on partition() above. - int partition = key == null ? 0 : Utils.toPositive(Utils.murmur2(key)) % numPartitions; + 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 index 1e6cab197c00..601fa40e679c 100644 --- 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 @@ -20,6 +20,8 @@ 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; @@ -31,6 +33,7 @@ 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; /** @@ -43,14 +46,18 @@ * 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. * - *

That does not arise, and this test is here to record why: the fuser folds such a Flatten into - * the SDK harness stage rather than leaving a node for the runner to translate. The Flattens that + *

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, this test starts + * 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> { @ProcessElement public void processElement(@Element Integer input, OutputReceiver> out) { @@ -89,22 +96,46 @@ private static Pipeline mixedParallelismFlatten(int internalParallelism) { return pipeline; } - @Test - public void branchesAtDifferentParallelismsAreFusedRatherThanLeftToTheRunner() { - // Translating is the assertion: the mismatched shape does not reach FlattenTranslator, because - // the fuser absorbs this Flatten into the harness stage. Were it to arrive there, the runner - // would build one Flatten node over branches of differing parallelism and stall. - KafkaStreamsTranslationContext context = - KafkaStreamsTestRunner.translate(mixedParallelismFlatten(4)); + /** Every processor node in the topology, across all subtopologies. */ + private static List processorNames(TopologyDescription description) { + List 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; + } - assertThat(context.getTopology().describe().subtopologies().isEmpty(), is(false)); + 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 theSameShapeTranslatesAtASingleParallelism() { - KafkaStreamsTranslationContext context = - KafkaStreamsTestRunner.translate(mixedParallelismFlatten(1)); + public void branchesAtDifferentParallelismsAreFusedRatherThanLeftToTheRunner() { + assertFlattenWasFusedAway( + KafkaStreamsTestRunner.translate(mixedParallelismFlatten(4)).getTopology().describe()); + } - assertThat(context.getTopology().describe().subtopologies().isEmpty(), is(false)); + @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()); } }