diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java index e048a996a8c7..045a74a8507e 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java @@ -655,6 +655,14 @@ public static WriteRecords writeRecords() { ///////////////////////// Read Support \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\ + /** + * Default number of keys to redistribute Kafka inputs into. + * + *

This value is used when {@link Read#withRedistribute()} is used without {@link + * Read#withRedistributeNumKeys(int redistributeNumKeys)}. + */ + 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 +1107,11 @@ public Read withTopicPartitions(List topicPartitions) { * @return an updated {@link Read} transform. */ public Read withRedistribute() { - return toBuilder().setRedistributed(true).build(); + Builder builder = toBuilder().setRedistributed(true); + if (getRedistributeNumKeys() == 0) { + builder = builder.setRedistributeNumKeys(DEFAULT_REDISTRIBUTE_NUM_KEYS); + } + return builder.build(); } /** @@ -1121,10 +1133,11 @@ public Read withAllowDuplicates(Boolean allowDuplicates) { * Redistributes Kafka messages into a distinct number of keys for processing in subsequent * steps. * - *

Specifying an explicit number of keys is generally recommended over redistributing into an - * unbounded key space. + *

If unset, defaults to {@link KafkaIO#DEFAULT_REDISTRIBUTE_NUM_KEYS}. * - *

Must be used with {@link KafkaIO#withRedistribute()}. + *

Use zero to disable bucketing into a distinct number of keys. + * + *

Must be used with {@link Read#withRedistribute()}. * * @param redistributeNumKeys specifies the total number of keys for redistributing inputs. * @return an updated {@link Read} transform. @@ -2667,13 +2680,30 @@ public ReadSourceDescriptors withProcessingTime() { /** Enable Redistribute. */ public ReadSourceDescriptors withRedistribute() { - return toBuilder().setRedistribute(true).build(); + Builder builder = toBuilder().setRedistribute(true); + if (getRedistributeNumKeys() == 0) { + builder = builder.setRedistributeNumKeys(DEFAULT_REDISTRIBUTE_NUM_KEYS); + } + return builder.build(); } public ReadSourceDescriptors withAllowDuplicates() { return toBuilder().setAllowDuplicates(true).build(); } + /** + * Redistributes Kafka messages into a distinct number of keys for processing in subsequent + * steps. + * + *

If unset, defaults to {@link KafkaIO#DEFAULT_REDISTRIBUTE_NUM_KEYS}. + * + *

Use zero to disable bucketing into a distinct number of keys. + * + *

Must be used with {@link ReadSourceDescriptors#withRedistribute()}. + * + * @param redistributeNumKeys specifies the total number of keys for redistributing inputs. + * @return an updated {@link Read} transform. + */ public ReadSourceDescriptors withRedistributeNumKeys(int redistributeNumKeys) { return toBuilder().setRedistributeNumKeys(redistributeNumKeys).build(); } diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java index 3d441f8dc521..83c2e1b38826 100644 --- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java +++ b/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 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);