Set default redistribute key limit for KafkaIO read.

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

@@ -655,6 +655,8 @@ public static <K, V> WriteRecords<K, V> writeRecords() {
 
   ///////////////////////// Read Support \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
 
+  private static final int DEFAULT_REDISTRIBUTE_NUM_KEYS = 32768;
+
   /**
    * A {@link PTransform} to read from Kafka topics. See {@link KafkaIO} for more information on
    * usage and configuration.
@@ -1099,7 +1101,11 @@ public Read<K, V> withTopicPartitions(List<TopicPartition> topicPartitions) {
      * @return an updated {@link Read} transform.
      */
     public Read<K, V> withRedistribute() {
-      return toBuilder().setRedistributed(true).build();
+      Builder<K, V> builder = toBuilder().setRedistributed(true);
+      if (getRedistributeNumKeys() == 0) {
+        builder = builder.setRedistributeNumKeys(DEFAULT_REDISTRIBUTE_NUM_KEYS);
+      }
+      return builder.build();
     }
 
     /**
@@ -2667,7 +2673,11 @@ public ReadSourceDescriptors<K, V> withProcessingTime() {
 
     /** Enable Redistribute. */
     public ReadSourceDescriptors<K, V> withRedistribute() {
-      return toBuilder().setRedistribute(true).build();
+      Builder<K, V> builder = toBuilder().setRedistribute(true);
+      if (getRedistributeNumKeys() == 0) {
+        builder = builder.setRedistributeNumKeys(DEFAULT_REDISTRIBUTE_NUM_KEYS);
+      }
+      return builder.build();
     }
 
     public ReadSourceDescriptors<K, V> withAllowDuplicates() {

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

@@ -30,6 +30,7 @@
 import static org.hamcrest.Matchers.matchesPattern;
 import static org.hamcrest.Matchers.not;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
@@ -792,6 +793,53 @@ public void testNumKeysIgnoredWithRedistributeNotEnabled() {
     p.run();
   }
 
+  @Test
+  public void testDefaultRedistributeNumKeys() {
+    int numElements = 1000;
+    // Redistribute is not used and does not modify the read transform further.
+    KafkaIO.Read<Integer, Long> read =
+        mkKafkaReadTransform(
+            numElements,
+            numElements,
+            new ValueAsTimestampFn(),
+            false, /*redistribute*/
+            false, /*allowDuplicates*/
+            null, /*numKeys*/
+            null, /*offsetDeduplication*/
+            null /*topics*/);
+    assertFalse(read.isRedistributed());
+    assertEquals(0, read.getRedistributeNumKeys());
+
+    // Redistribute is used and defaulted the number of keys due to no user setting.
+    read =
+        mkKafkaReadTransform(
+            numElements,
+            numElements,
+            new ValueAsTimestampFn(),
+            true, /*redistribute*/
+            false, /*allowDuplicates*/
+            null, /*numKeys*/
+            null, /*offsetDeduplication*/
+            null /*topics*/);
+    assertTrue(read.isRedistributed());
+    // Default is defined by DEFAULT_REDISTRIBUTE_NUM_KEYS in KafkaIO.
+    assertEquals(32768, read.getRedistributeNumKeys());
+
+    // Redistribute is set with user-specified the number of keys.
+    read =
+        mkKafkaReadTransform(
+            numElements,
+            numElements,
+            new ValueAsTimestampFn(),
+            true, /*redistribute*/
+            false, /*allowDuplicates*/
+            10, /*numKeys*/
+            null, /*offsetDeduplication*/
+            null /*topics*/);
+    assertTrue(read.isRedistributed());
+    assertEquals(10, read.getRedistributeNumKeys());
+  }
+
   @Test
   public void testDisableRedistributeKafkaOffsetLegacy() {
     thrown.expect(Exception.class);