Add byRecordKey property to Kafka read compatibility.

Author: tomstepp

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/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,7 +108,7 @@ private PipelineResult testReadTransformCreationWithImplementationBoundPropertie
       Function<KafkaIO.Read<Integer, Long>, KafkaIO.Read<Integer, Long>> kafkaReadDecorator) {
     p.apply(
         kafkaReadDecorator.apply(
-            mkKafkaReadTransform(
+            KafkaIOTest.mkKafkaReadTransformBase(
                 1000,
                 null,
                 new ValueAsTimestampFn(),

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

@@ -386,7 +386,7 @@ public Consumer<byte[], byte[]> apply(Map<String, Object> config) {
 
   static KafkaIO.Read<Integer, Long> mkKafkaReadTransform(
       int numElements, @Nullable SerializableFunction<KV<Integer, Long>, Instant> timestampFn) {
-    return mkKafkaReadTransform(
+    return mkKafkaReadTransformBase(
         numElements,
         numElements,
         timestampFn,
@@ -400,7 +400,7 @@ static KafkaIO.Read<Integer, Long> mkKafkaReadTransform(
 
   static KafkaIO.Read<Integer, Long> mkKafkaReadTransformWithOffsetDedup(
       int numElements, @Nullable SerializableFunction<KV<Integer, Long>, Instant> timestampFn) {
-    return mkKafkaReadTransform(
+    return mkKafkaReadTransformBase(
         numElements,
         numElements,
         timestampFn,
@@ -416,7 +416,7 @@ static KafkaIO.Read<Integer, Long> mkKafkaReadTransformWithRedistributeByRecordK
       int numElements,
       @Nullable SerializableFunction<KV<Integer, Long>, Instant> timestampFn,
       boolean byRecordKey) {
-    return mkKafkaReadTransform(
+    return mkKafkaReadTransformBase(
         numElements,
         numElements,
         timestampFn,
@@ -432,7 +432,7 @@ static KafkaIO.Read<Integer, Long> mkKafkaReadTransformWithTopics(
       int numElements,
       @Nullable SerializableFunction<KV<Integer, Long>, Instant> timestampFn,
       List<String> topics) {
-    return mkKafkaReadTransform(
+    return mkKafkaReadTransformBase(
         numElements,
         numElements,
         timestampFn,
@@ -448,7 +448,7 @@ static KafkaIO.Read<Integer, Long> mkKafkaReadTransformWithTopics(
    * Creates a consumer with two topics, with 10 partitions each. numElements are (round-robin)
    * assigned all the 20 partitions.
    */
-  static KafkaIO.Read<Integer, Long> mkKafkaReadTransform(
+  static KafkaIO.Read<Integer, Long> mkKafkaReadTransformBase(
       int numElements,
       @Nullable Integer maxNumRecords,
       @Nullable SerializableFunction<KV<Integer, Long>, Instant> timestampFn,
@@ -737,7 +737,7 @@ public void warningsWithAllowDuplicatesEnabledAndCommitOffsets() {
 
     PCollection<Long> input =
         p.apply(
-                mkKafkaReadTransform(
+                mkKafkaReadTransformBase(
                         numElements,
                         numElements,
                         new ValueAsTimestampFn(),
@@ -766,7 +766,7 @@ public void noWarningsWithNoAllowDuplicatesAndCommitOffsets() {
 
     PCollection<Long> input =
         p.apply(
-                mkKafkaReadTransform(
+                mkKafkaReadTransformBase(
                         numElements,
                         numElements,
                         new ValueAsTimestampFn(),
@@ -796,7 +796,7 @@ public void testNumKeysIgnoredWithRedistributeNotEnabled() {
 
     PCollection<Long> input =
         p.apply(
-                mkKafkaReadTransform(
+                mkKafkaReadTransformBase(
                         numElements,
                         numElements,
                         new ValueAsTimestampFn(),
@@ -2170,7 +2170,7 @@ public void testUnboundedSourceStartReadTime() {
 
     PCollection<Long> input =
         p.apply(
-                mkKafkaReadTransform(
+                mkKafkaReadTransformBase(
                         numElements,
                         maxNumRecords,
                         new ValueAsTimestampFn(),
@@ -2247,7 +2247,7 @@ public void testUnboundedSourceStartReadTimeException() {
     int startTime = numElements / 20;
 
     p.apply(
-            mkKafkaReadTransform(
+            mkKafkaReadTransformBase(
                     numElements,
                     numElements,
                     new ValueAsTimestampFn(),