diff --git a/CHANGES.md b/CHANGES.md index f98790bc2e19..4c723c792521 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -31,6 +31,7 @@ ## New Features / Improvements +* (Java) Flink runner: `Reshuffle` in batch distributes elements round-robin from deterministic start channels instead of random ones, so inputs with few elements per sender stay balanced ([#40607](https://github.com/apache/beam/issues/40607)). * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Breaking Changes diff --git a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java index abeb9daaf044..797e78487381 100644 --- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java @@ -34,6 +34,7 @@ import org.apache.beam.runners.core.SystemReduceFn; import org.apache.beam.runners.core.construction.SerializablePipelineOptions; import org.apache.beam.runners.flink.adapter.FlinkKey; +import org.apache.beam.runners.flink.translation.functions.DeterministicRebalance; import org.apache.beam.runners.flink.translation.functions.FlinkAssignWindows; import org.apache.beam.runners.flink.translation.functions.ImpulseSourceFunction; import org.apache.beam.runners.flink.translation.types.CoderTypeInformation; @@ -936,7 +937,16 @@ public void translateNode( DataStream>> inputDataSet = context.getInputDataStream(context.getInput(transform)); - context.setOutputDataStream(context.getOutput(transform), inputDataSet.rebalance()); + // In batch, a reshuffle often redistributes only a few elements per sender, such as file + // names or splittable DoFn restrictions, where rebalance()'s random start channels leave + // some readers idle and overload others. Start channels derived from the sender index keep + // such inputs balanced. + DataStream>> reshuffled = + context.isStreaming() + ? inputDataSet.rebalance() + : DeterministicRebalance.apply(inputDataSet, context.getOutput(transform).getName()); + + context.setOutputDataStream(context.getOutput(transform), reshuffled); } } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java index b3b40d2874a7..96eba2e097aa 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java @@ -34,6 +34,7 @@ import org.apache.beam.runners.core.SystemReduceFn; import org.apache.beam.runners.core.construction.SerializablePipelineOptions; import org.apache.beam.runners.flink.adapter.FlinkKey; +import org.apache.beam.runners.flink.translation.functions.DeterministicRebalance; import org.apache.beam.runners.flink.translation.functions.FlinkAssignWindows; import org.apache.beam.runners.flink.translation.functions.ImpulseSourceFunction; import org.apache.beam.runners.flink.translation.types.CoderTypeInformation; @@ -936,7 +937,16 @@ public void translateNode( DataStream>> inputDataSet = context.getInputDataStream(context.getInput(transform)); - context.setOutputDataStream(context.getOutput(transform), inputDataSet.rebalance()); + // In batch, a reshuffle often redistributes only a few elements per sender, such as file + // names or splittable DoFn restrictions, where rebalance()'s random start channels leave + // some readers idle and overload others. Start channels derived from the sender index keep + // such inputs balanced. + DataStream>> reshuffled = + context.isStreaming() + ? inputDataSet.rebalance() + : DeterministicRebalance.apply(inputDataSet, context.getOutput(transform).getName()); + + context.setOutputDataStream(context.getOutput(transform), reshuffled); } } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/DeterministicRebalance.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/DeterministicRebalance.java new file mode 100644 index 000000000000..caf332199a5a --- /dev/null +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/DeterministicRebalance.java @@ -0,0 +1,135 @@ +/* + * 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.flink.translation.functions; + +import org.apache.flink.api.common.TaskInfo; +import org.apache.flink.api.common.functions.Partitioner; +import org.apache.flink.api.common.functions.RichMapFunction; +import org.apache.flink.api.common.typeinfo.BasicTypeInfo; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.java.typeutils.TupleTypeInfo; +import org.apache.flink.streaming.api.datastream.DataStream; + +/** + * Round-robin redistribution where each sender starts at a channel derived from its own subtask + * index rather than at a random one. + * + *

Flink's {@link DataStream#rebalance()} starts each sender's round-robin at a random channel. + * That is even when senders emit many elements, but when they emit only a few each, as when a + * {@code Reshuffle} redistributes file names or the restrictions of a splittable DoFn, the + * assignment is close to random: with {@code k} elements per sender on average, the number of + * elements per receiver follows a Poisson distribution of mean {@code k}, so some receivers get + * none and others several times the mean. + * + *

Here sender {@code i} of {@code M} sends its {@code j}-th element to channel {@code (floor(i * + * N / M) + j) mod N} for {@code N} receivers. The senders' start channels are spread evenly, so an + * input that is already balanced stays balanced, and a single sender's elements still go + * round-robin over all receivers. + * + *

A {@link Partitioner} is not told which sender it runs in, so each element is first tagged + * with its sender's index and parallelism. The tagging and untagging maps are chained to their + * neighbours, so the only cost is the tag in the serialized elements. + */ +public final class DeterministicRebalance { + + private DeterministicRebalance() {} + + /** Redistributes {@code input} round-robin from deterministic per-sender start channels. */ + public static DataStream apply(DataStream input, String name) { + TypeInformation typeInfo = input.getType(); + TypeInformation> taggedTypeInfo = + new TupleTypeInfo<>(BasicTypeInfo.LONG_TYPE_INFO, typeInfo); + return input + .map(new TagWithSender()) + .name(name + "/TagWithSender") + .returns(taggedTypeInfo) + .partitionCustom(new DeterministicRoundRobinPartitioner(), tagged -> tagged.f0) + .map(tagged -> tagged.f1) + .name(name + "/StripSender") + .returns(typeInfo); + } + + /** + * Encodes a sender's subtask index and the senders' parallelism in one key: the sender index in + * the upper 32 bits, the number of senders in the lower 32 bits. + */ + static final class SenderKey { + + private SenderKey() {} + + static long of(int senderIndex, int numSenders) { + return ((long) senderIndex << 32) | (numSenders & 0xFFFFFFFFL); + } + + static int senderIndex(long senderKey) { + return (int) (senderKey >>> 32); + } + + static int numSenders(long senderKey) { + return (int) senderKey; + } + } + + /** Tags each element with its {@link SenderKey}, which is constant for a sender subtask. */ + static class TagWithSender extends RichMapFunction> { + + // Transient fields are reset on deserialization, so the key is computed on first use. + private transient boolean keyed; + private transient long key; + + @Override + public Tuple2 map(T value) { + if (!keyed) { + TaskInfo taskInfo = getRuntimeContext().getTaskInfo(); + key = + SenderKey.of(taskInfo.getIndexOfThisSubtask(), taskInfo.getNumberOfParallelSubtasks()); + keyed = true; + } + return Tuple2.of(key, value); + } + } + + /** + * Sends the {@code j}-th element of sender {@code i} of {@code M} to channel {@code (floor(i * N + * / M) + j) mod N}. + * + *

Stateful: Flink gives each sender subtask its own deserialized partitioner, which counts + * that sender's elements. + */ + static class DeterministicRoundRobinPartitioner implements Partitioner { + + // Transient fields are reset on deserialization, so the start channel is set on first use. + private transient boolean started; + private transient int nextChannel; + + @Override + public int partition(Long key, int numPartitions) { + if (!started) { + int senderIndex = SenderKey.senderIndex(key); + int numSenders = SenderKey.numSenders(key); + nextChannel = + (int) ((long) senderIndex * numPartitions / Math.max(1, numSenders)) % numPartitions; + started = true; + } + int channel = nextChannel; + nextChannel = (nextChannel + 1) % numPartitions; + return channel; + } + } +} diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/DeterministicRebalanceTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/DeterministicRebalanceTest.java new file mode 100644 index 000000000000..a48911d8fc19 --- /dev/null +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/DeterministicRebalanceTest.java @@ -0,0 +1,124 @@ +/* + * 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.flink.translation.functions; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.beam.runners.flink.FlinkPipelineOptions; +import org.apache.beam.runners.flink.FlinkTestPipeline; +import org.apache.beam.runners.flink.translation.functions.DeterministicRebalance.DeterministicRoundRobinPartitioner; +import org.apache.beam.runners.flink.translation.functions.DeterministicRebalance.SenderKey; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.Reshuffle; +import org.apache.beam.sdk.values.PCollection; +import org.junit.Test; + +public class DeterministicRebalanceTest { + + /** + * Returns the number of elements each receiver gets when the {@code k}-th of {@code numElements} + * elements starts on sender {@code k mod numSenders}, as after a perfect round-robin. + */ + private static int[] receiverCounts(int numElements, int numSenders, int numReceivers) { + int[] counts = new int[numReceivers]; + for (int sender = 0; sender < numSenders; sender++) { + // Each sender subtask has its own partitioner instance. + DeterministicRoundRobinPartitioner partitioner = new DeterministicRoundRobinPartitioner(); + long key = SenderKey.of(sender, numSenders); + for (int k = sender; k < numElements; k += numSenders) { + counts[partitioner.partition(key, numReceivers)]++; + } + } + return counts; + } + + private static void assertBalanced(int[] counts, int numElements) { + assertEquals(numElements, Arrays.stream(counts).sum()); + int min = Arrays.stream(counts).min().getAsInt(); + int max = Arrays.stream(counts).max().getAsInt(); + assertTrue( + "receiver counts should differ by at most one, got min " + min + " and max " + max, + max - min <= 1); + } + + @Test + public void testFewElementsPerSenderStayBalanced() { + // About 2.6 elements per sender, as when 17,776 files are reshuffled at parallelism 6,900: + // random start channels leave about 7.6% of the receivers empty. + assertBalanced(receiverCounts(17_776, 6_900, 6_900), 17_776); + } + + @Test + public void testSingleSenderRoundRobinsOverAllReceivers() { + int[] counts = receiverCounts(70, 1, 7); + int[] expected = new int[7]; + Arrays.fill(expected, 10); + assertArrayEquals(expected, counts); + } + + @Test + public void testFewerSendersThanReceivers() { + assertBalanced(receiverCounts(25, 5, 20), 25); + assertBalanced(receiverCounts(10_003, 3, 10), 10_003); + } + + @Test + public void testMoreSendersThanReceiversStayCloseToBalanced() { + // With more senders than receivers, groups of about M / N senders share a start channel, so + // the counts can differ by about one element per sender in a group. Random start channels + // would give a Poisson spread instead: here a standard deviation of about 2.5 around 6.2. + int numSenders = 20; + int numReceivers = 6; + int[] counts = receiverCounts(37, numSenders, numReceivers); + assertEquals(37, Arrays.stream(counts).sum()); + int min = Arrays.stream(counts).min().getAsInt(); + int max = Arrays.stream(counts).max().getAsInt(); + int bound = (numSenders + numReceivers - 1) / numReceivers; + assertTrue("got min " + min + " and max " + max, min > 0 && max - min <= bound); + } + + @Test + public void testReshuffleInDataStreamBatchKeepsAllElements() { + FlinkTestPipeline p = FlinkTestPipeline.createForBatch(); + FlinkPipelineOptions options = p.getOptions().as(FlinkPipelineOptions.class); + options.setParallelism(4); + options.setUseDataStreamForBatch(true); + List input = new ArrayList<>(); + for (int i = 0; i < 1_000; i++) { + input.add(i); + } + + PCollection output = p.apply(Create.of(input)).apply(Reshuffle.viaRandomKey()); + + PAssert.that(output).containsInAnyOrder(input); + p.run(); + } + + @Test + public void testSenderKeyRoundTrips() { + long key = SenderKey.of(6_899, 13_800); + assertEquals(6_899, SenderKey.senderIndex(key)); + assertEquals(13_800, SenderKey.numSenders(key)); + } +}