From 564e4642b3ea1cfc64017e9e5cafa433540cd2da Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Fri, 18 Sep 2026 14:10:09 -0400 Subject: [PATCH 1/2] [KafkaIO] Add `with_gcp_adc` and `num_partitions` options to Kafka SchemaTransforms and Beam YAML - Expose `with_gcp_adc` (`withGcpAdc`) in KafkaIO read/write SchemaTransform, Python xlang IO and YAML - Add `num_partitions` in KafkaIO read SchemaTransform, Python xlang IO and YAML. This avoids the need to Kafka instance connection at pipeline submission time for Dataflow runner v1 (Streaming runner) if using SchemaTransform --- ...KafkaReadSchemaTransformConfiguration.java | 19 ++++++ .../KafkaReadSchemaTransformProvider.java | 28 ++++++-- .../KafkaWriteSchemaTransformProvider.java | 52 +++++++++------ .../KafkaReadSchemaTransformProviderTest.java | 64 ++++++++++++++++++- ...KafkaWriteSchemaTransformProviderTest.java | 13 +++- sdks/python/apache_beam/yaml/standard_io.yaml | 3 + 6 files changed, 152 insertions(+), 27 deletions(-) diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java index 0cf40f9b7eba..59157d8ff506 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java @@ -198,6 +198,21 @@ public static Builder builder() { @Nullable public abstract Boolean getRedistributeByRecordKey(); + @SchemaFieldDescription( + "Whether to use Google Cloud Platform Application Default Credentials (ADC) for" + + " authenticating with a Google Managed Kafka cluster.") + @SchemaFieldNumber("17") + @Nullable + public abstract Boolean getWithGcpAdc(); + + @SchemaFieldDescription( + "The number of partitions to read from the Kafka topic. If specified, partitions" + + " 0 to numPartitions-1 will be assigned statically without querying Kafka at pipeline" + + " construction time.") + @SchemaFieldNumber("18") + @Nullable + public abstract Integer getNumPartitions(); + /** Builder for the {@link KafkaReadSchemaTransformConfiguration}. */ @AutoValue.Builder public abstract static class Builder { @@ -238,6 +253,10 @@ public abstract static class Builder { public abstract Builder setRedistributeByRecordKey(Boolean redistributeByRecordKey); + public abstract Builder setWithGcpAdc(Boolean withGcpAdc); + + public abstract Builder setNumPartitions(Integer numPartitions); + /** Builds a {@link KafkaReadSchemaTransformConfiguration} instance. */ public abstract KafkaReadSchemaTransformConfiguration build(); } diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java index c5764b39bc69..aa8b627c63e0 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java @@ -31,6 +31,7 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; +import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; @@ -69,6 +70,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.kafka.clients.consumer.Consumer; import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.serialization.ByteArrayDeserializer; import org.joda.time.Duration; import org.slf4j.Logger; @@ -166,6 +168,20 @@ private SchemaRegistryProvider getSchemaRegistryProvider(String confluentSchemaR return SchemaRegistryProvider.UNSPECIFIED; } + private static KafkaIO.Read applyTopicOrPartitions( + KafkaIO.Read kafkaRead, KafkaReadSchemaTransformConfiguration configuration) { + Integer numPartitions = configuration.getNumPartitions(); + String topic = configuration.getTopic(); + if (numPartitions != null && numPartitions > 0) { + List topicPartitions = new ArrayList<>(numPartitions); + for (int i = 0; i < numPartitions; i++) { + topicPartitions.add(new TopicPartition(topic, i)); + } + return kafkaRead.withTopicPartitions(topicPartitions); + } + return kafkaRead.withTopic(topic); + } + private static KafkaIO.Read applyRedistributeSettings( KafkaIO.Read kafkaRead, KafkaReadSchemaTransformConfiguration configuration) { Boolean redistribute = configuration.getRedistributed(); @@ -221,12 +237,14 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { KafkaIO.Read kafkaRead; kafkaRead = - KafkaIO.read() - .withTopic(configuration.getTopic()) + applyTopicOrPartitions(KafkaIO.read(), configuration) .withConsumerFactoryFn(new ConsumerFactoryWithGcsTrustStores()) .withBootstrapServers(configuration.getBootstrapServers()) .withConsumerConfigUpdates(consumerConfigs) .withKeyDeserializer(ByteArrayDeserializer.class); + if (Boolean.TRUE.equals(configuration.getWithGcpAdc())) { + kafkaRead = kafkaRead.withGCPApplicationDefaultCredentials(); + } SchemaRegistryProvider provider = getSchemaRegistryProvider(confluentSchemaRegUrl); switch (provider) { @@ -300,11 +318,13 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { } KafkaIO.Read kafkaRead = - KafkaIO.readBytes() + applyTopicOrPartitions(KafkaIO.readBytes(), configuration) .withConsumerConfigUpdates(consumerConfigs) .withConsumerFactoryFn(new ConsumerFactoryWithGcsTrustStores()) - .withTopic(configuration.getTopic()) .withBootstrapServers(configuration.getBootstrapServers()); + if (Boolean.TRUE.equals(configuration.getWithGcpAdc())) { + kafkaRead = kafkaRead.withGCPApplicationDefaultCredentials(); + } Integer maxReadTimeSeconds = configuration.getMaxReadTimeSeconds(); if (maxReadTimeSeconds != null) { kafkaRead = kafkaRead.withMaxReadTime(Duration.standardSeconds(maxReadTimeSeconds)); diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java index d8d0fe478e35..213a8543e9d4 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java @@ -256,17 +256,20 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { handleErrors)) .withOutputTags(RECORD_OUTPUT_TAG, TupleTagList.of(ERROR_TAG))); HashMap producerConfig = new HashMap<>(configOverrides); + KafkaIO.Write kafkaWrite = + KafkaIO.write() + .withTopic(configuration.getTopic()) + .withBootstrapServers(configuration.getBootstrapServers()) + .withProducerConfigUpdates(producerConfig) + .withKeySerializer(ByteArraySerializer.class) + .withValueSerializer((Class) KafkaAvroSerializer.class); + if (Boolean.TRUE.equals(configuration.getWithGcpAdc())) { + kafkaWrite = kafkaWrite.withGCPApplicationDefaultCredentials(); + } outputTuple .get(RECORD_OUTPUT_TAG) .setCoder(KvCoder.of(NullableCoder.of(ByteArrayCoder.of()), AvroCoder.of(avroSchema))) - .apply( - "Map Rows to GenericRecords", - KafkaIO.write() - .withTopic(configuration.getTopic()) - .withBootstrapServers(configuration.getBootstrapServers()) - .withProducerConfigUpdates(producerConfig) - .withKeySerializer(ByteArraySerializer.class) - .withValueSerializer((Class) KafkaAvroSerializer.class)); + .apply("Map Rows to GenericRecords", kafkaWrite); } else { outputTuple = input @@ -278,19 +281,23 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { "Kafka-write-error-counter", toBytesFn, errorSchema, handleErrors)) .withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG))); + KafkaIO.Write kafkaWrite = + KafkaIO.write() + .withTopic(configuration.getTopic()) + .withBootstrapServers(configuration.getBootstrapServers()) + .withProducerConfigUpdates( + configOverrides == null + ? new HashMap<>() + : new HashMap(configOverrides)) + .withKeySerializer(ByteArraySerializer.class) + .withValueSerializer(ByteArraySerializer.class); + if (Boolean.TRUE.equals(configuration.getWithGcpAdc())) { + kafkaWrite = kafkaWrite.withGCPApplicationDefaultCredentials(); + } outputTuple .get(OUTPUT_TAG) .setCoder(KvCoder.of(NullableCoder.of(ByteArrayCoder.of()), ByteArrayCoder.of())) - .apply( - KafkaIO.write() - .withTopic(configuration.getTopic()) - .withBootstrapServers(configuration.getBootstrapServers()) - .withProducerConfigUpdates( - configOverrides == null - ? new HashMap<>() - : new HashMap(configOverrides)) - .withKeySerializer(ByteArraySerializer.class) - .withValueSerializer(ByteArraySerializer.class)); + .apply(kafkaWrite); } // TODO: include output from KafkaIO Write once updated from PDone @@ -381,6 +388,13 @@ public abstract static class KafkaWriteSchemaTransformConfiguration implements S @Nullable public abstract String getSchema(); + @SchemaFieldDescription( + "Whether to use Google Cloud Platform Application Default Credentials (ADC) for" + + " authenticating with a Google Managed Kafka cluster.") + @SchemaFieldNumber("8") + @Nullable + public abstract Boolean getWithGcpAdc(); + public static Builder builder() { return new AutoValue_KafkaWriteSchemaTransformProvider_KafkaWriteSchemaTransformConfiguration .Builder(); @@ -404,6 +418,8 @@ public abstract static class Builder { public abstract Builder setSchema(String schema); + public abstract Builder setWithGcpAdc(Boolean withGcpAdc); + public abstract KafkaWriteSchemaTransformConfiguration build(); } } 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 9d276fa0e55e..d5a0b7b152ba 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 @@ -18,7 +18,9 @@ package org.apache.beam.sdk.io.kafka; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -26,10 +28,12 @@ import java.util.List; import java.util.Objects; import java.util.ServiceLoader; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.managed.Managed; +import org.apache.beam.sdk.runners.TransformHierarchy; import org.apache.beam.sdk.schemas.NoSuchSchemaException; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.SchemaRegistry; @@ -41,6 +45,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams; +import org.apache.kafka.common.TopicPartition; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -138,7 +143,9 @@ public void testFindTransformAndMakeItWork() { "allow_duplicates", "offset_deduplication", "redistribute_num_keys", - "redistribute_by_record_key"), + "redistribute_by_record_key", + "with_gcp_adc", + "num_partitions"), kafkaProvider.configurationSchema().getFields().stream() .map(field -> field.getName()) .collect(Collectors.toSet())); @@ -361,7 +368,12 @@ public void testBuildTransformWithManaged() { + "schema: '" + PROTO_SCHEMA + "'\n" - + "message_name: MyMessage"); + + "message_name: MyMessage", + "topic: topic_6\n" + + "bootstrap_servers: some bootstrap\n" + + "format: RAW\n" + + "with_gcp_adc: true\n" + + "num_partitions: 3"); for (String config : configs) { // Kafka Read SchemaTransform gets built in ManagedSchemaTransformProvider's expand @@ -371,6 +383,42 @@ public void testBuildTransformWithManaged() { } } + @Test + public void testBuildTransformWithNumPartitions() { + KafkaReadSchemaTransformProvider kafkaProvider = new KafkaReadSchemaTransformProvider(); + SchemaTransform transformWithPartitions = + kafkaProvider.from( + KafkaReadSchemaTransformConfiguration.builder() + .setTopic("anytopic") + .setBootstrapServers("anybootstrap") + .setFormat("RAW") + .setNumPartitions(3) + .build()); + Pipeline pipelineWithPartitions = Pipeline.create(); + transformWithPartitions.expand(PCollectionRowTuple.empty(pipelineWithPartitions)); + AtomicReference> readWithPartitions = new AtomicReference<>(); + pipelineWithPartitions.traverseTopologically( + new Pipeline.PipelineVisitor.Defaults() { + @Override + public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) { + if (node.getTransform() instanceof KafkaIO.Read) { + readWithPartitions.set((KafkaIO.Read) node.getTransform()); + } + return CompositeBehavior.ENTER_TRANSFORM; + } + }); + assertNotNull(readWithPartitions.get()); + assertEquals( + Arrays.asList( + new TopicPartition("anytopic", 0), + new TopicPartition("anytopic", 1), + new TopicPartition("anytopic", 2)), + readWithPartitions.get().getTopicPartitions()); + assertTrue( + readWithPartitions.get().getTopics() == null + || readWithPartitions.get().getTopics().isEmpty()); + } + // This test verifies that the schema for KafkaReadSchemaTransformConfiguration is correctly // generated. This schema is used when KafkaReadSchemaTransformConfiguration are // serialized/deserialized with @@ -380,7 +428,7 @@ public void testKafkaReadSchemaTransformConfigurationSchema() throws NoSuchSchem Schema schema = SchemaRegistry.createDefault().getSchema(KafkaReadSchemaTransformConfiguration.class); - assertEquals(17, schema.getFieldCount()); + assertEquals(19, schema.getFieldCount()); // Check field name, type, and nullability. Descriptions are not checked as they are not // critical for serialization. @@ -478,5 +526,15 @@ public void testKafkaReadSchemaTransformConfigurationSchema() throws NoSuchSchem Schema.Field.nullable("redistributeByRecordKey", Schema.FieldType.BOOLEAN) .withDescription(schema.getField(16).getDescription()), schema.getField(16)); + + assertEquals( + Schema.Field.nullable("withGcpAdc", Schema.FieldType.BOOLEAN) + .withDescription(schema.getField(17).getDescription()), + schema.getField(17)); + + assertEquals( + Schema.Field.nullable("numPartitions", Schema.FieldType.INT32) + .withDescription(schema.getField(18).getDescription()), + schema.getField(18)); } } diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java index cc5bd82a5a88..d7215eec99b0 100644 --- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java +++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java @@ -268,7 +268,11 @@ public void testBuildTransformWithManaged() { + "schema: '" + PROTO_SCHEMA + "'\n" - + "message_name: MyMessage"); + + "message_name: MyMessage", + "topic: topic_4\n" + + "bootstrap_servers: some bootstrap\n" + + "format: RAW\n" + + "with_gcp_adc: true"); for (String config : configs) { // Kafka Write SchemaTransform gets built in ManagedSchemaTransformProvider's expand @@ -289,7 +293,7 @@ public void testKafkaWriteSchemaTransformConfigurationSchema() throws NoSuchSche System.out.println("schema = " + schema); - assertEquals(8, schema.getFieldCount()); + assertEquals(9, schema.getFieldCount()); // Check field name, type, and nullability. Descriptions are not checked as they are not // critical for serialization. @@ -341,5 +345,10 @@ public void testKafkaWriteSchemaTransformConfigurationSchema() throws NoSuchSche Schema.Field.nullable("schema", Schema.FieldType.STRING) .withDescription(schema.getField(7).getDescription()), schema.getField(7)); + + assertEquals( + Schema.Field.nullable("withGcpAdc", Schema.FieldType.BOOLEAN) + .withDescription(schema.getField(8).getDescription()), + schema.getField(8)); } } diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml index c41dd6f2e07f..2f46bd7cc41a 100644 --- a/sdks/python/apache_beam/yaml/standard_io.yaml +++ b/sdks/python/apache_beam/yaml/standard_io.yaml @@ -75,6 +75,8 @@ 'file_descriptor_path': 'file_descriptor_path' 'message_name': 'message_name' 'max_read_time_seconds': 'max_read_time_seconds' + 'with_gcp_adc': 'with_gcp_adc' + 'num_partitions': 'num_partitions' 'WriteToKafka': 'format': 'format' 'topic': 'topic' @@ -84,6 +86,7 @@ 'file_descriptor_path': 'file_descriptor_path' 'message_name': 'message_name' 'schema': 'schema' + 'with_gcp_adc': 'with_gcp_adc' underlying_provider: type: beamJar transforms: From 7d00dc19da79e02cc016dbb7b59fda5997e9dfd8 Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Tue, 22 Sep 2026 11:49:56 -0400 Subject: [PATCH 2/2] fix comments --- .../sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java | 4 ++-- .../beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java index 59157d8ff506..aad7d742769d 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java @@ -199,8 +199,8 @@ public static Builder builder() { public abstract Boolean getRedistributeByRecordKey(); @SchemaFieldDescription( - "Whether to use Google Cloud Platform Application Default Credentials (ADC) for" - + " authenticating with a Google Managed Kafka cluster.") + "Whether to use Google Application Default Credentials (ADC) for authenticating with a " + + "Google Managed Kafka cluster.") @SchemaFieldNumber("17") @Nullable public abstract Boolean getWithGcpAdc(); diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java index 213a8543e9d4..55202764c689 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java @@ -389,8 +389,8 @@ public abstract static class KafkaWriteSchemaTransformConfiguration implements S public abstract String getSchema(); @SchemaFieldDescription( - "Whether to use Google Cloud Platform Application Default Credentials (ADC) for" - + " authenticating with a Google Managed Kafka cluster.") + "Whether to use Google Application Default Credentials (ADC) for authenticating with a " + + "Google Managed Kafka cluster.") @SchemaFieldNumber("8") @Nullable public abstract Boolean getWithGcpAdc();