diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index 8438e6af2a97..2057c52f561a 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -47,8 +47,8 @@ * *

The Beam key becomes the Kafka record key so Kafka Streams shuffles by it. The topology is a * {@link ShuffleByKeyProcessor} that sets that key and passes watermark reports through, a sink to - * an internal repartition topic using a {@link GroupByKeyBroadcastPartitioner} that hashes data by - * key and fans watermarks out to every partition, a source reading that topic back, and the {@link + * an internal repartition topic using a {@link KStreamsPayloadPartitioner} that hashes data by key + * and fans watermarks out to every partition, a source reading that topic back, and the {@link * WindowedGroupByKeyProcessor} with its state and timer stores. */ class GroupByKeyTranslator implements PTransformTranslator { @@ -114,16 +114,20 @@ public void translate( shuffleName, () -> new ShuffleByKeyProcessor( - keyCoder, upstreamPartitionCount, shuffleName, context.getTerminationTracker()), + keyCoder, + upstreamPartitionCount, + partitionCount, + shuffleName, + context.getTerminationTracker()), parentProcessor); - // Shuffle through the repartition topic: data partitioned by key, watermark broadcast. + // Shuffle through the repartition topic: data by key, watermarks broadcast, flushes targeted. topology.addSink( sinkName, repartitionTopic, Serdes.ByteArray().serializer(), payloadSerde.serializer(), - new GroupByKeyBroadcastPartitioner<>(), + new KStreamsPayloadPartitioner<>(partitionCount), shuffleName); topology.addSource( sourceName, 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/KStreamsPayloadPartitioner.java similarity index 74% rename from runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java rename to runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadPartitioner.java index 3c775c86f14c..8a892300225b 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/KStreamsPayloadPartitioner.java @@ -25,19 +25,27 @@ import org.apache.kafka.streams.processor.StreamPartitioner; /** - * Partitions records on the GroupByKey repartition topic. + * Partitions records on a repartition topic. * *

* * @param the data element type carried by data payloads */ -class GroupByKeyBroadcastPartitioner implements StreamPartitioner> { +class KStreamsPayloadPartitioner implements StreamPartitioner> { + + /** The partition count the flush targets were computed for. */ + private final int expectedPartitions; + + KStreamsPayloadPartitioner(int expectedPartitions) { + this.expectedPartitions = expectedPartitions; + } @Override public Integer partition(String topic, byte[] key, KStreamsPayload value, int numPartitions) { @@ -56,6 +64,17 @@ public Optional> partitions( } return Optional.of(all); } + if (value.isFlush()) { + // Targets assume internalParallelism partitions; a topic from an earlier run may differ. + if (numPartitions != expectedPartitions) { + throw new IllegalStateException( + String.format( + "Repartition topic %s has %d partitions but the pipeline expects %d. It was" + + " probably created by an earlier run with a different --internalParallelism.", + topic, numPartitions, expectedPartitions)); + } + return Optional.of(value.asFlush().getTargetPartitions()); + } 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 diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java index 638239406ed9..ce4bc9bb48f4 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java @@ -17,10 +17,12 @@ */ package org.apache.beam.runners.kafka.streams.translation; +import java.util.Set; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.values.KV; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.kafka.streams.processor.api.Processor; import org.apache.kafka.streams.processor.api.ProcessorContext; import org.apache.kafka.streams.processor.api.Record; @@ -39,11 +41,13 @@ *

Watermark reports are relabelled here with the reporting instance's real partition identity. * This is the point at which a report stops being delivered in-process and starts crossing a topic: * upstream of it a transform forwards to its fused children, which see exactly one instance of it, - * so the report names a single source. The {@link GroupByKeyBroadcastPartitioner} on the sink below + * so the report names a single source. The {@link KStreamsPayloadPartitioner} on the sink below * fans each report out to every partition, so a downstream task instead sees a report from * every instance of the upstream transform, and has to be able to tell them apart to know when it * has heard from all of them. The transform id is left alone, so the report still names the * transform that produced it. + * + *

Flush markers are readdressed here too; see {@link #flushTargets}. */ class ShuffleByKeyProcessor implements Processor, byte[], KStreamsPayload> { @@ -57,6 +61,9 @@ class ShuffleByKeyProcessor private int upstreamPartition; + /** Partitions of the repartition topic. */ + private final int downstreamPartitionCount; + // Reports this shuffle as finished once it has written the terminal watermark to the repartition // topic. The downstream side reading that topic reports separately, which is why the pipeline // waits for every processor rather than the first. @@ -67,10 +74,12 @@ class ShuffleByKeyProcessor ShuffleByKeyProcessor( Coder keyCoder, int upstreamPartitionCount, + int downstreamPartitionCount, String nodeName, TerminationTracker terminationTracker) { this.keyCoder = keyCoder; this.upstreamPartitionCount = upstreamPartitionCount; + this.downstreamPartitionCount = downstreamPartitionCount; this.terminationReporter = new TerminationReporter(terminationTracker, nodeName); } @@ -114,6 +123,16 @@ public void process(Record> record) { throw new RuntimeException("Failed to encode shuffle key", e); } ctx.forward(record.withKey(encodedKey)); + } else if (payload.isFlush()) { + // Forwarded once: the sink's partitioner makes Kafka Streams write a copy to each target. + // Fanning in, an instance may have nothing to address. + Set targets = + flushTargets(upstreamPartition, upstreamPartitionCount, downstreamPartitionCount); + if (!targets.isEmpty()) { + ctx.forward( + new Record>( + record.key(), KStreamsPayload.flush(targets), record.timestamp())); + } } else { WatermarkPayload report = payload.asWatermark(); ctx.forward( @@ -129,6 +148,20 @@ public void process(Record> record) { } } + /** + * Downstream partitions {@code [i*D/U, (i+1)*D/U)} for upstream partition {@code i}. Over all + * {@code i} these cover each downstream partition once, and do not change on a rebalance. + */ + static Set flushTargets(int upstreamPartition, int upstreamCount, int downstreamCount) { + long from = (long) upstreamPartition * downstreamCount / upstreamCount; + long to = ((long) upstreamPartition + 1) * downstreamCount / upstreamCount; + ImmutableSet.Builder targets = ImmutableSet.builder(); + for (long partition = from; partition < to; partition++) { + targets.add((int) partition); + } + return targets.build(); + } + private static T checkInitialized(@Nullable T value) { if (value == null) { throw new IllegalStateException("ShuffleByKeyProcessor used before init()"); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java index d0a04e38126a..b035791ea8a3 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java @@ -33,8 +33,8 @@ * exhausted. This collects those reports and fires a callback when nothing is left to do. * *

No coordination between instances is needed: a watermark crossing a repartition topic is - * broadcast to every partition (see {@link GroupByKeyBroadcastPartitioner}), so every task observes - * the terminal watermark itself and all instances reach the same conclusion independently. + * broadcast to every partition (see {@link KStreamsPayloadPartitioner}), so every task observes the + * terminal watermark itself and all instances reach the same conclusion independently. * *

Every local processor is counted, not just the first. One instance can own tasks from both * sides of a repartition topic, and the upstream side goes terminal as soon as it has written to diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadPartitionerBrokerIT.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadPartitionerBrokerIT.java new file mode 100644 index 000000000000..7645d8fc6bf2 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadPartitionerBrokerIT.java @@ -0,0 +1,176 @@ +/* + * 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.nio.file.Files; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.TreeMap; +import java.util.UUID; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.values.WindowedValues; +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.clients.admin.Admin; +import org.apache.kafka.clients.admin.AdminClientConfig; +import org.apache.kafka.clients.admin.NewTopic; +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.clients.consumer.ConsumerRecord; +import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.clients.producer.ProducerConfig; +import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.serialization.ByteArrayDeserializer; +import org.apache.kafka.common.serialization.ByteArraySerializer; +import org.apache.kafka.streams.KafkaStreams; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.Topology; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.testcontainers.kafka.KafkaContainer; +import org.testcontainers.utility.DockerImageName; + +/** + * Checks on a real broker that a flush is written once to each partition it names. The {@code + * TopologyTestDriver} cannot show this, since it reports every topic as having one partition. + */ +@RunWith(JUnit4.class) +public class KStreamsPayloadPartitionerBrokerIT { + + private static final int PARTITIONS = 4; + + private static KafkaContainer kafka; + + @BeforeClass + public static void startBroker() { + kafka = new KafkaContainer(DockerImageName.parse("apache/kafka:4.0.0")); + kafka.start(); + } + + @AfterClass + public static void stopBroker() { + if (kafka != null) { + kafka.stop(); + } + } + + @Test + public void aFlushIsWrittenOnceToEachPartitionItNames() throws Exception { + String id = UUID.randomUUID().toString(); + String input = "flush-in-" + id; + String output = "flush-out-" + id; + try (Admin admin = + Admin.create( + ImmutableMap.of( + AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, kafka.getBootstrapServers()))) { + admin + .createTopics( + Arrays.asList( + new NewTopic(input, 1, (short) 1), new NewTopic(output, PARTITIONS, (short) 1))) + .all() + .get(); + } + + KStreamsPayloadSerde serde = + new KStreamsPayloadSerde<>( + WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE)); + Topology topology = new Topology(); + topology.addSource("in", new ByteArrayDeserializer(), serde.deserializer(), input); + topology.addSink( + "out", + output, + new ByteArraySerializer(), + serde.serializer(), + new KStreamsPayloadPartitioner<>(PARTITIONS), + "in"); + + Properties config = new Properties(); + config.put(StreamsConfig.APPLICATION_ID_CONFIG, "flush-it-" + id); + config.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, kafka.getBootstrapServers()); + config.put(StreamsConfig.STATE_DIR_CONFIG, Files.createTempDirectory("flush-it").toString()); + KafkaStreams streams = new KafkaStreams(topology, config); + streams.start(); + try { + try (KafkaProducer> producer = + new KafkaProducer<>( + ImmutableMap.of( + ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka.getBootstrapServers()), + new ByteArraySerializer(), + serde.serializer())) { + producer + .send( + new ProducerRecord<>( + input, new byte[0], KStreamsPayload.flush(ImmutableSet.of(1, 2)))) + .get(); + } + + Map expected = ImmutableMap.of(1, 1, 2, 1); + assertThat(copiesPerPartition(output, serde, 2), is(expected)); + } finally { + streams.close(Duration.ofSeconds(30)); + streams.cleanUp(); + } + } + + /** Reads the topic until {@code expected} records arrive, then a little longer for extras. */ + private static Map copiesPerPartition( + String topic, KStreamsPayloadSerde serde, int expected) { + List partitions = new ArrayList<>(); + for (int partition = 0; partition < PARTITIONS; partition++) { + partitions.add(new TopicPartition(topic, partition)); + } + Map copies = new TreeMap<>(); + try (KafkaConsumer> consumer = + new KafkaConsumer<>( + ImmutableMap.of( + ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka.getBootstrapServers()), + new ByteArrayDeserializer(), + serde.deserializer())) { + consumer.assign(partitions); + consumer.seekToBeginning(partitions); + int seen = 0; + boolean draining = false; + long deadline = System.currentTimeMillis() + 60_000; + while (System.currentTimeMillis() < deadline) { + for (ConsumerRecord> record : + consumer.poll(Duration.ofMillis(500))) { + assertThat(record.value().isFlush(), is(true)); + copies.merge(record.partition(), 1, Integer::sum); + seen++; + } + if (!draining && seen >= expected) { + draining = true; + deadline = System.currentTimeMillis() + 3_000; + } + } + } + return copies; + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadPartitionerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadPartitionerTest.java new file mode 100644 index 000000000000..821847f53781 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadPartitionerTest.java @@ -0,0 +1,74 @@ +/* + * 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 static org.junit.Assert.assertThrows; + +import java.util.Optional; +import java.util.Set; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.junit.Test; + +/** Tests for {@link KStreamsPayloadPartitioner}. */ +public class KStreamsPayloadPartitionerTest { + + private static final String TOPIC = "repartition"; + private static final byte[] KEY = new byte[] {1, 2, 3}; + + private final KStreamsPayloadPartitioner partitioner = + new KStreamsPayloadPartitioner<>(8); + + @Test + public void aWatermarkGoesToEveryPartition() { + Optional> targets = + partitioner.partitions(TOPIC, KEY, KStreamsPayload.watermark(5L, "t", 0, 1), 8); + assertThat(targets, is(Optional.of(ImmutableSet.of(0, 1, 2, 3, 4, 5, 6, 7)))); + } + + @Test + public void aFlushGoesToExactlyThePartitionsItNames() { + Optional> targets = + partitioner.partitions(TOPIC, KEY, KStreamsPayload.flush(ImmutableSet.of(4, 5)), 8); + assertThat(targets, is(Optional.of(ImmutableSet.of(4, 5)))); + } + + @Test + public void dataWithTheSameKeyGoesToTheSamePartition() { + KStreamsPayload data = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(1)); + Optional> first = partitioner.partitions(TOPIC, KEY, data, 8); + assertThat(first.get().size(), is(1)); + assertThat(partitioner.partitions(TOPIC, KEY.clone(), data, 8), is(first)); + } + + @Test + public void keylessDataIsLeftToTheProducer() { + KStreamsPayload data = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(1)); + assertThat(partitioner.partitions(TOPIC, null, data, 8), is(Optional.empty())); + } + + @Test + public void aFlushFailsOnATopicWithADifferentPartitionCount() { + // The targets were chosen for 8 partitions; this topic was left with 4 by an earlier run. + assertThrows( + IllegalStateException.class, + () -> partitioner.partitions(TOPIC, KEY, KStreamsPayload.flush(ImmutableSet.of(0)), 4)); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java index b6716764da7f..a002c6bc7299 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java @@ -21,16 +21,18 @@ import static org.hamcrest.MatcherAssert.assertThat; import java.util.Properties; +import java.util.Set; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.kafka.streams.processor.TaskId; import org.apache.kafka.streams.processor.api.MockProcessorContext; import org.apache.kafka.streams.processor.api.Record; import org.junit.Test; /** - * Tests how {@link ShuffleByKeyProcessor} restamps a watermark report as it is about to cross a - * repartition topic. + * Tests how {@link ShuffleByKeyProcessor} restamps watermark reports and addresses flush markers as + * they are about to cross a repartition topic. * *

Upstream of the shuffle a transform forwards its watermark in process, to its fused children, * which see exactly one instance of it — so the report names a single source. The sink below the @@ -42,13 +44,19 @@ public class ShuffleByKeyProcessorTest { private static final String UPSTREAM_ID = "upstream"; - @SuppressWarnings("unchecked") private static ShuffleByKeyProcessor processorFor(int taskPartition, int upstreamPartitions) { + return processorFor(taskPartition, upstreamPartitions, upstreamPartitions); + } + + @SuppressWarnings("unchecked") + private static ShuffleByKeyProcessor processorFor( + int taskPartition, int upstreamPartitions, int downstreamPartitions) { ShuffleByKeyProcessor processor = new ShuffleByKeyProcessor( (org.apache.beam.sdk.coders.Coder) (org.apache.beam.sdk.coders.Coder) StringUtf8Coder.of(), upstreamPartitions, + downstreamPartitions, "shuffle-node", new TerminationTracker()); MockProcessorContext> ctx = @@ -122,4 +130,65 @@ public void dataIsRekeyedByTheBeamKeyAndNotRestamped() { assertThat(lastContext.forwarded().size(), is(1)); assertThat(lastContext.forwarded().get(0).record().value().isData(), is(true)); } + + private static Record> flush() { + // The shuffle ignores these targets and picks its own. + return new Record<>(new byte[0], KStreamsPayload.flush(ImmutableSet.of(0)), 0L); + } + + @Test + public void aFlushIsAddressedToThisInstancesShareOfTheRepartitionTopic() { + // Partition 2 of 10 upstream, 8 downstream. + processorFor(2, 10, 8).process(flush()); + + assertThat(lastContext.forwarded().size(), is(1)); + FlushPayload out = lastContext.forwarded().get(0).record().value().asFlush(); + assertThat(out.getTargetPartitions(), is(ImmutableSet.of(1))); + } + + @Test + public void anInstanceWithNothingToAddressForwardsNoFlush() { + // 10 upstream into 8 downstream: partition 0 addresses nothing. + processorFor(0, 10, 8).process(flush()); + + assertThat(lastContext.forwarded().isEmpty(), is(true)); + } + + @Test + public void flushTargetsMatchTheWorkedExample() { + // 10 upstream partitions owned in pairs, 8 downstream. + assertThat(pairTargets(0, 10, 8), is(ImmutableSet.of(0))); + assertThat(pairTargets(2, 10, 8), is(ImmutableSet.of(1, 2))); + assertThat(pairTargets(4, 10, 8), is(ImmutableSet.of(3))); + assertThat(pairTargets(6, 10, 8), is(ImmutableSet.of(4, 5))); + assertThat(pairTargets(8, 10, 8), is(ImmutableSet.of(6, 7))); + } + + private static Set pairTargets(int first, int upstream, int downstream) { + return ImmutableSet.builder() + .addAll(ShuffleByKeyProcessor.flushTargets(first, upstream, downstream)) + .addAll(ShuffleByKeyProcessor.flushTargets(first + 1, upstream, downstream)) + .build(); + } + + @Test + public void everyDownstreamPartitionIsTargetedExactlyOnce() { + int[][] shapes = {{10, 8}, {2, 8}, {8, 2}, {4, 4}, {1, 1}, {1, 16}, {16, 1}, {3, 7}, {7, 3}}; + for (int[] shape : shapes) { + int upstream = shape[0]; + int downstream = shape[1]; + int[] hits = new int[downstream]; + for (int partition = 0; partition < upstream; partition++) { + for (int target : ShuffleByKeyProcessor.flushTargets(partition, upstream, downstream)) { + hits[target]++; + } + } + for (int target = 0; target < downstream; target++) { + assertThat( + "upstream=" + upstream + " downstream=" + downstream + " target=" + target, + hits[target], + is(1)); + } + } + } }