diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java index d5a0b7b152ba..60cbe87d43d3 100644 --- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java +++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java @@ -373,7 +373,12 @@ public void testBuildTransformWithManaged() { + "bootstrap_servers: some bootstrap\n" + "format: RAW\n" + "with_gcp_adc: true\n" - + "num_partitions: 3"); + + "num_partitions: 3\n" + + "redistributed: true\n" + + "redistribute_num_keys: 10\n" + + "allow_duplicates: false\n" + + "offset_deduplication: true\n" + + "redistribute_by_record_key: true"); for (String config : configs) { // Kafka Read SchemaTransform gets built in ManagedSchemaTransformProvider's expand @@ -419,6 +424,44 @@ public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) { || readWithPartitions.get().getTopics().isEmpty()); } + @Test + public void testBuildTransformWithRedistribute() { + ServiceLoader serviceLoader = + ServiceLoader.load(SchemaTransformProvider.class); + List providers = + StreamSupport.stream(serviceLoader.spliterator(), false) + .filter(provider -> provider.getClass() == KafkaReadSchemaTransformProvider.class) + .collect(Collectors.toList()); + KafkaReadSchemaTransformProvider kafkaProvider = + (KafkaReadSchemaTransformProvider) providers.get(0); + + SchemaTransform validTransform = + kafkaProvider.from( + KafkaReadSchemaTransformConfiguration.builder() + .setTopic("anytopic") + .setBootstrapServers("anybootstrap") + .setFormat("RAW") + .setRedistributed(true) + .setRedistributeNumKeys(10) + .setAllowDuplicates(false) + .setOffsetDeduplication(true) + .setRedistributeByRecordKey(true) + .build()); + validTransform.expand(PCollectionRowTuple.empty(Pipeline.create())); + + SchemaTransform invalidTransform = + kafkaProvider.from( + KafkaReadSchemaTransformConfiguration.builder() + .setTopic("anytopic") + .setBootstrapServers("anybootstrap") + .setFormat("RAW") + .setRedistributeNumKeys(10) + .build()); + assertThrows( + IllegalStateException.class, + () -> invalidTransform.expand(PCollectionRowTuple.empty(Pipeline.create()))); + } + // This test verifies that the schema for KafkaReadSchemaTransformConfiguration is correctly // generated. This schema is used when KafkaReadSchemaTransformConfiguration are // serialized/deserialized with diff --git a/sdks/python/apache_beam/yaml/extended_tests/messaging/kafka.yaml b/sdks/python/apache_beam/yaml/extended_tests/messaging/kafka.yaml index 1c0e3fc6adaa..3ef6ac1d9f5a 100644 --- a/sdks/python/apache_beam/yaml/extended_tests/messaging/kafka.yaml +++ b/sdks/python/apache_beam/yaml/extended_tests/messaging/kafka.yaml @@ -62,6 +62,11 @@ pipelines: auto.offset.reset: "earliest" group.id: "yaml-kafka-test-group" max_read_time_seconds: 10 # will read forever if not set + redistributed: true + redistribute_num_keys: 2 + allow_duplicates: false + offset_deduplication: true + redistribute_by_record_key: true - type: MapToFields config: language: python diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml index 15aa9186c0fc..13dd9c16818c 100644 --- a/sdks/python/apache_beam/yaml/standard_io.yaml +++ b/sdks/python/apache_beam/yaml/standard_io.yaml @@ -112,6 +112,11 @@ 'max_read_time_seconds': 'max_read_time_seconds' 'with_gcp_adc': 'with_gcp_adc' 'num_partitions': 'num_partitions' + 'redistributed': 'redistributed' + 'allow_duplicates': 'allow_duplicates' + 'redistribute_num_keys': 'redistribute_num_keys' + 'offset_deduplication': 'offset_deduplication' + 'redistribute_by_record_key': 'redistribute_by_record_key' 'WriteToKafka': 'format': 'format' 'topic': 'topic'