Skip to content

[KafkaIO] Add with_gcp_adc and num_partitions options to Kafka SchemaTransforms and YAML #40181

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Abacn merged 2 commits into from
Oct 5, 2026
Merged

[KafkaIO] Add with_gcp_adc and num_partitions options to Kafka SchemaTransforms and YAML #40181

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
564e464
[KafkaIO] Add `with_gcp_adc` and `num_partitions` options to Kafka Sc…
Abacn Sep 18, 2026
7d00dc1
fix comments
Abacn Sep 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions ...fka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,21 @@ public static Builder builder() {
@Nullable
public abstract Boolean getRedistributeByRecordKey();

@SchemaFieldDescription(
"Whether to use Google Application Default Credentials (ADC) for authenticating with a "
+ "Google Managed Kafka cluster.")
@SchemaFieldNumber("17")
@Nullable
public abstract Boolean getWithGcpAdc();

@sjvanrossum sjvanrossum Sep 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: This doesn't match the name used to configure the transform. Application default credentials is annoyingly long to spell so I'm not opposed to ADC, but I'd consider keeping it the same for consistency and clarity to users who might be unfamiliar with the acronym.

@Abacn Abacn Sep 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Understand. Since this is mostly a ease-of-use change (users can current set GCP configs themselves by adding all needed configurations to config) a simple name is preferred

@sjvanrossum sjvanrossum Oct 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SGTM


@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();

@sjvanrossum sjvanrossum Sep 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In most cases the partition set for a topic will match the range of integers from 0 to N, but users may arbitrarily retire individual partitions.
Examples:

  • Create a partition per entity with a non-sequential integer id (e.g., a partition per player in a game).
  • Retire a partition after the last usable offset is used to produce a record.

Replacing the number of partitions property with a topic partition list property would spare adding it later with more boilerplate to validate the different combinations of configuration properties.

@Abacn Abacn Sep 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's true there is a risk number of partitions evolving won't get honored, but in fact it's the same limitation for a typical setup

  • on Dataflow streaming runner, it uses UnboundedSource, which resolves partitions at once, at initial split that happens at pipeline expansion time

  • on Dataflow portable runner, it uses SDF implementation. Initial split happens at pipeline execution time, but after that the partitions to consume is still fixed

There is an .withDynamicRead(...) option that support dynamic partitions, led to SDK implementation and requires Dataflow runner v2.

One of the motivation of exposing this option is to be able to use KafkaIO via SchemaTransform to connect to a GMK on Dataflow Streaming runner, therefore .withDynamicRead(...) isn't available here


/** Builder for the {@link KafkaReadSchemaTransformConfiguration}. */
@AutoValue.Builder
public abstract static class Builder {
Expand Down Expand Up @@ -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();
}
Expand Down
28 changes: 24 additions & 4 deletions ...io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -166,6 +168,20 @@ private SchemaRegistryProvider getSchemaRegistryProvider(String confluentSchemaR
return SchemaRegistryProvider.UNSPECIFIED;
}

private static <K, V> KafkaIO.Read<K, V> applyTopicOrPartitions(
KafkaIO.Read<K, V> kafkaRead, KafkaReadSchemaTransformConfiguration configuration) {
Integer numPartitions = configuration.getNumPartitions();
String topic = configuration.getTopic();
if (numPartitions != null && numPartitions > 0) {
List<TopicPartition> 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 <K, V> KafkaIO.Read<K, V> applyRedistributeSettings(
KafkaIO.Read<K, V> kafkaRead, KafkaReadSchemaTransformConfiguration configuration) {
Boolean redistribute = configuration.getRedistributed();
Expand Down Expand Up @@ -221,12 +237,14 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
KafkaIO.Read<byte[], GenericRecord> kafkaRead;

kafkaRead =
KafkaIO.<byte[], GenericRecord>read()
.withTopic(configuration.getTopic())
applyTopicOrPartitions(KafkaIO.<byte[], GenericRecord>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) {
Expand Down Expand Up @@ -300,11 +318,13 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
}

KafkaIO.Read<byte[], byte[]> 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));
Expand Down
52 changes: 34 additions & 18 deletions ...o/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,17 +256,20 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
handleErrors))
.withOutputTags(RECORD_OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
HashMap<String, Object> producerConfig = new HashMap<>(configOverrides);
KafkaIO.Write<byte[], GenericRecord> kafkaWrite =
KafkaIO.<byte[], GenericRecord>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.<byte[], GenericRecord>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
Expand All @@ -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<byte[], byte[]> kafkaWrite =
KafkaIO.<byte[], byte[]>write()
.withTopic(configuration.getTopic())
.withBootstrapServers(configuration.getBootstrapServers())
.withProducerConfigUpdates(
configOverrides == null
? new HashMap<>()
: new HashMap<String, Object>(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.<byte[], byte[]>write()
.withTopic(configuration.getTopic())
.withBootstrapServers(configuration.getBootstrapServers())
.withProducerConfigUpdates(
configOverrides == null
? new HashMap<>()
: new HashMap<String, Object>(configOverrides))
.withKeySerializer(ByteArraySerializer.class)
.withValueSerializer(ByteArraySerializer.class));
.apply(kafkaWrite);
}

// TODO: include output from KafkaIO Write once updated from PDone
Expand Down Expand Up @@ -381,6 +388,13 @@ public abstract static class KafkaWriteSchemaTransformConfiguration implements S
@Nullable
public abstract String getSchema();

@SchemaFieldDescription(
"Whether to use Google 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();
Expand All @@ -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();
}
}
Expand Down
64 changes: 61 additions & 3 deletions ...afka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,22 @@
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;
import java.util.Arrays;
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;
Expand All @@ -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;
Expand Down Expand Up @@ -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()));
Expand Down Expand Up @@ -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
Expand All @@ -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<KafkaIO.Read<?, ?>> 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
Expand All @@ -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.
Expand Down Expand Up @@ -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));
}
}
13 changes: 11 additions & 2 deletions ...fka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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));
}
}
3 changes: 3 additions & 0 deletions sdks/python/apache_beam/yaml/standard_io.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand All @@ -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:
Expand Down
Loading