Add deterministic redistribute sharding for KafkaIO read.

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java

@@ -131,7 +131,7 @@ public void processElement(
   public static class RedistributeArbitrarily<T>
       extends PTransform<PCollection<T>, PCollection<T>> {
     // The number of buckets to shard into.
-    // A runner is free to ignore this (a runner may ignore the transorm
+    // A runner is free to ignore this (a runner may ignore the transform
     // entirely!) This is a performance optimization to prevent having
     // unit sized bundles on the output. If unset, uses a random integer key.
     private @Nullable Integer numBuckets = null;

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java

@@ -79,6 +79,7 @@
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.Redistribute;
+import org.apache.beam.sdk.transforms.Redistribute.RedistributeArbitrarily;
 import org.apache.beam.sdk.transforms.Reshuffle;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.SimpleFunction;
@@ -730,6 +731,9 @@ public abstract static class Read<K, V>
     @Pure
     public abstract @Nullable Boolean getOffsetDeduplication();
 
+    @Pure
+    public abstract @Nullable Boolean getRedistributeByRecordKey();
+
     @Pure
     public abstract @Nullable Duration getWatchTopicPartitionDuration();
 
@@ -800,6 +804,8 @@ abstract Builder<K, V> setConsumerFactoryFn(
 
       abstract Builder<K, V> setOffsetDeduplication(Boolean offsetDeduplication);
 
+      abstract Builder<K, V> setRedistributeByRecordKey(Boolean redistributeByRecordKey);
+
       abstract Builder<K, V> setTimestampPolicyFactory(
           TimestampPolicyFactory<K, V> timestampPolicyFactory);
 
@@ -915,11 +921,15 @@ static <K, V> void setupExternalBuilder(
               && config.offsetDeduplication != null) {
             builder.setOffsetDeduplication(config.offsetDeduplication);
           }
+          if (config.redistribute && config.redistributeByRecordKey != null) {
+            builder.setRedistributeByRecordKey(config.redistributeByRecordKey);
+          }
         } else {
           builder.setRedistributed(false);
           builder.setRedistributeNumKeys(0);
           builder.setAllowDuplicates(false);
           builder.setOffsetDeduplication(false);
+          builder.setRedistributeByRecordKey(false);
         }
       }
 
@@ -989,6 +999,7 @@ public static class Configuration {
         private Boolean redistribute;
         private Boolean allowDuplicates;
         private Boolean offsetDeduplication;
+        private Boolean redistributeByRecordKey;
         private Long dynamicReadPollIntervalSeconds;
 
         public void setConsumerConfig(Map<String, String> consumerConfig) {
@@ -1051,6 +1062,10 @@ public void setOffsetDeduplication(Boolean offsetDeduplication) {
           this.offsetDeduplication = offsetDeduplication;
         }
 
+        public void setRedistributeByRecordKey(Boolean redistributeByRecordKey) {
+          this.redistributeByRecordKey = redistributeByRecordKey;
+        }
+
         public void setDynamicReadPollIntervalSeconds(Long dynamicReadPollIntervalSeconds) {
           this.dynamicReadPollIntervalSeconds = dynamicReadPollIntervalSeconds;
         }
@@ -1161,6 +1176,10 @@ public Read<K, V> withOffsetDeduplication(Boolean offsetDeduplication) {
       return toBuilder().setOffsetDeduplication(offsetDeduplication).build();
     }
 
+    public Read<K, V> withRedistributeByRecordKey(Boolean redistributeByRecordKey) {
+      return toBuilder().setRedistributeByRecordKey(redistributeByRecordKey).build();
+    }
+
     /**
      * Internally sets a {@link java.util.regex.Pattern} of topics to read from. All the partitions
      * from each of the matching topics are read.
@@ -1679,6 +1698,11 @@ private void checkRedistributeConfiguration() {
         LOG.warn(
             "Offsets used for deduplication are available in WindowedValue's metadata. Combining, aggregating, mutating them may risk with data loss.");
       }
+      if (getRedistributeByRecordKey() != null && getRedistributeByRecordKey()) {
+        checkState(
+            isRedistributed(),
+            "withRedistributeByRecordKey can only be used when withRedistribute is set.");
+      }
     }
 
     private void warnAboutUnsafeConfigurations(PBegin input) {
@@ -1858,18 +1882,25 @@ public PCollection<KafkaRecord<K, V>> expand(PBegin input) {
                 "Offsets committed due to usage of commitOffsetsInFinalize() and may not capture all work processed due to use of withRedistribute() with duplicates enabled");
           }
 
-          if (kafkaRead.getRedistributeNumKeys() == 0) {
-            return output.apply(
-                "Insert Redistribute",
-                Redistribute.<KafkaRecord<K, V>>arbitrarily()
-                    .withAllowDuplicates(kafkaRead.isAllowDuplicates()));
-          } else {
-            return output.apply(
-                "Insert Redistribute with Shards",
-                Redistribute.<KafkaRecord<K, V>>arbitrarily()
-                    .withAllowDuplicates(kafkaRead.isAllowDuplicates())
-                    .withNumBuckets((int) kafkaRead.getRedistributeNumKeys()));
+          if (kafkaRead.getOffsetDeduplication() != null && kafkaRead.getOffsetDeduplication()) {
+            if (kafkaRead.getRedistributeByRecordKey() != null
+                && kafkaRead.getRedistributeByRecordKey()) {
+              return output.apply(
+                  KafkaReadRedistribute.<K, V>byRecordKey(kafkaRead.getRedistributeNumKeys()));
+            } else {
+              return output.apply(
+                  KafkaReadRedistribute.<K, V>byOffsetShard(kafkaRead.getRedistributeNumKeys()));
+            }
+          }
+          RedistributeArbitrarily<KafkaRecord<K, V>> redistribute =
+              Redistribute.<KafkaRecord<K, V>>arbitrarily()
+                  .withAllowDuplicates(kafkaRead.isAllowDuplicates());
+          String redistributeName = "Insert Redistribute";
+          if (kafkaRead.getRedistributeNumKeys() != 0) {
+            redistribute = redistribute.withNumBuckets((int) kafkaRead.getRedistributeNumKeys());
+            redistributeName = "Insert Redistribute with Shards";
           }
+          return output.apply(redistributeName, redistribute);
         }
         return output;
       }

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java

@@ -139,6 +139,12 @@ Object getDefaultValue() {
     },
     OFFSET_DEDUPLICATION(LEGACY),
     LOG_TOPIC_VERIFICATION,
+    REDISTRIBUTE_BY_RECORD_KEY {
+      @Override
+      Object getDefaultValue() {
+        return false;
+      }
+    },
     ;
 
     private final @NonNull ImmutableSet<KafkaIOReadImplementation> supportedImplementations;

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadRedistribute.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.sdk.io.kafka;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Redistribute;
+import org.apache.beam.sdk.transforms.Values;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedInteger;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+public class KafkaReadRedistribute<K, V>
+    extends PTransform<PCollection<KafkaRecord<K, V>>, PCollection<KafkaRecord<K, V>>> {
+  public static <K, V> KafkaReadRedistribute<K, V> byOffsetShard(@Nullable Integer numBuckets) {
+    return new KafkaReadRedistribute<>(numBuckets, false);
+  }
+
+  public static <K, V> KafkaReadRedistribute<K, V> byRecordKey(@Nullable Integer numBuckets) {
+    return new KafkaReadRedistribute<>(numBuckets, true);
+  }
+
+  // The number of buckets to shard into.
+  private @Nullable Integer numBuckets = null;
+  // When redistributing, group records by the Kafka record's key instead of by offset hash.
+  private boolean byRecordKey = false;
+
+  private KafkaReadRedistribute(@Nullable Integer numBuckets, boolean byRecordKey) {
+    this.numBuckets = numBuckets;
+    this.byRecordKey = byRecordKey;
+  }
+
+  @Override
+  public PCollection<KafkaRecord<K, V>> expand(PCollection<KafkaRecord<K, V>> input) {
+
+    if (byRecordKey) {
+      return input
+          .apply("Pair with shard from key", ParDo.of(new AssignRecordKeyFn<K, V>(numBuckets)))
+          .apply(Redistribute.<Integer, KafkaRecord<K, V>>byKey().withAllowDuplicates(false))
+          .apply(Values.create());
+    }
+
+    return input
+        .apply("Pair with shard from offset", ParDo.of(new AssignOffsetShardFn<K, V>(numBuckets)))
+        .apply(Redistribute.<Integer, KafkaRecord<K, V>>byKey().withAllowDuplicates(false))
+        .apply(Values.create());
+  }
+
+  static class AssignOffsetShardFn<K, V>
+      extends DoFn<KafkaRecord<K, V>, KV<Integer, KafkaRecord<K, V>>> {
+    private @NonNull UnsignedInteger numBuckets;
+
+    public AssignOffsetShardFn(@Nullable Integer numBuckets) {
+      if (numBuckets != null && numBuckets > 0) {
+        this.numBuckets = UnsignedInteger.fromIntBits(numBuckets);
+      } else {
+        this.numBuckets = UnsignedInteger.valueOf(0);
+      }
+    }
+
+    @ProcessElement
+    public void processElement(
+        @Element KafkaRecord<K, V> element,
+        OutputReceiver<KV<Integer, KafkaRecord<K, V>>> receiver) {
+      int hash = Hashing.farmHashFingerprint64().hashLong(element.getOffset()).asInt();
+
+      if (numBuckets != null) {
+        hash = UnsignedInteger.fromIntBits(hash).mod(numBuckets).intValue();
+      }
+
+      receiver.output(KV.of(hash, element));
+    }
+  }
+
+  static class AssignRecordKeyFn<K, V>
+      extends DoFn<KafkaRecord<K, V>, KV<Integer, KafkaRecord<K, V>>> {
+
+    private @NonNull UnsignedInteger numBuckets;
+
+    public AssignRecordKeyFn(@Nullable Integer numBuckets) {
+      if (numBuckets != null && numBuckets > 0) {
+        this.numBuckets = UnsignedInteger.fromIntBits(numBuckets);
+      } else {
+        this.numBuckets = UnsignedInteger.valueOf(0);
+      }
+    }
+
+    @ProcessElement
+    public void processElement(
+        @Element KafkaRecord<K, V> element,
+        OutputReceiver<KV<Integer, KafkaRecord<K, V>>> receiver) {
+      K key = element.getKV().getKey();
+      String keyString = key == null ? "" : key.toString();
+      int hash = Hashing.farmHashFingerprint64().hashBytes(keyString.getBytes(UTF_8)).asInt();
+
+      if (numBuckets != null) {
+        hash = UnsignedInteger.fromIntBits(hash).mod(numBuckets).intValue();
+      }
+
+      receiver.output(KV.of(hash, element));
+    }
+  }
+}

sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibilityTest.java

@@ -17,7 +17,6 @@
  */
 package org.apache.beam.sdk.io.kafka;
 
-import static org.apache.beam.sdk.io.kafka.KafkaIOTest.mkKafkaReadTransform;
 import static org.apache.beam.sdk.io.kafka.KafkaIOTest.mkKafkaReadTransformWithOffsetDedup;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsInAnyOrder;
@@ -109,15 +108,16 @@ private PipelineResult testReadTransformCreationWithImplementationBoundPropertie
       Function<KafkaIO.Read<Integer, Long>, KafkaIO.Read<Integer, Long>> kafkaReadDecorator) {
     p.apply(
         kafkaReadDecorator.apply(
-            mkKafkaReadTransform(
+            KafkaIOTest.mkKafkaReadTransform(
                 1000,
                 null,
                 new ValueAsTimestampFn(),
                 false, /*redistribute*/
                 false, /*allowDuplicates*/
                 0, /*numKeys*/
                 null, /*offsetDeduplication*/
-                null /*topics*/)));
+                null, /*topics*/
+                null /*redistributeByRecordKey*/)));
     return p.run();
   }