[Flink] Reshuffle batch inputs round-robin from deterministic start channels

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

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<WindowedValue<KV<K, InputT>>> 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<WindowedValue<KV<K, InputT>>> reshuffled =
+          context.isStreaming()
+              ? inputDataSet.rebalance()
+              : DeterministicRebalance.apply(inputDataSet, context.getOutput(transform).getName());
+
+      context.setOutputDataStream(context.getOutput(transform), reshuffled);
     }
   }
 

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<WindowedValue<KV<K, InputT>>> 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<WindowedValue<KV<K, InputT>>> reshuffled =
+          context.isStreaming()
+              ? inputDataSet.rebalance()
+              : DeterministicRebalance.apply(inputDataSet, context.getOutput(transform).getName());
+
+      context.setOutputDataStream(context.getOutput(transform), reshuffled);
     }
   }
 

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.
+ *
+ * <p>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.
+ *
+ * <p>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.
+ *
+ * <p>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 <T> DataStream<T> apply(DataStream<T> input, String name) {
+    TypeInformation<T> typeInfo = input.getType();
+    TypeInformation<Tuple2<Long, T>> taggedTypeInfo =
+        new TupleTypeInfo<>(BasicTypeInfo.LONG_TYPE_INFO, typeInfo);
+    return input
+        .map(new TagWithSender<T>())
+        .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<T> extends RichMapFunction<T, Tuple2<Long, T>> {
+
+    // 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<Long, T> 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}.
+   *
+   * <p>Stateful: Flink gives each sender subtask its own deserialized partitioner, which counts
+   * that sender's elements.
+   */
+  static class DeterministicRoundRobinPartitioner implements Partitioner<Long> {
+
+    // 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;
+    }
+  }
+}

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<Integer> input = new ArrayList<>();
+    for (int i = 0; i < 1_000; i++) {
+      input.add(i);
+    }
+
+    PCollection<Integer> 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));
+  }
+}