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..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,6 +56,14 @@ public Optional> partitions( } 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 000000000000..601fa40e679c --- /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. + * + *

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 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> { + @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; + } + + /** 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; + } + + 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()); + } +}