You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by au...@apache.org on 2023/10/27 04:46:46 UTC
[beam] branch master updated: [YAML] Kafka Read Provider (#28865)
This is an automated email from the ASF dual-hosted git repository.
austin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new e98e37f15b8 [YAML] Kafka Read Provider (#28865)
e98e37f15b8 is described below
commit e98e37f15b898c55b63af680aff33f528d2ede38
Author: Ferran Fernández Garrido <ff...@gmail.com>
AuthorDate: Fri Oct 27 06:46:38 2023 +0200
[YAML] Kafka Read Provider (#28865)
* [YAML] Kafka provider
---
.../KafkaReadSchemaTransformConfiguration.java | 2 +-
.../io/kafka/KafkaReadSchemaTransformProvider.java | 54 +++++++++++++++++++++-
.../KafkaReadSchemaTransformProviderTest.java | 18 ++++++++
sdks/python/apache_beam/yaml/standard_io.yaml | 21 +++++++++
4 files changed, 92 insertions(+), 3 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 f7e09158068..299a4f222e3 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
@@ -39,7 +39,7 @@ public abstract class KafkaReadSchemaTransformConfiguration {
public static final Set<String> VALID_START_OFFSET_VALUES = Sets.newHashSet("earliest", "latest");
- public static final String VALID_FORMATS_STR = "AVRO,JSON";
+ public static final String VALID_FORMATS_STR = "RAW,AVRO,JSON";
public static final Set<String> VALID_DATA_FORMATS =
Sets.newHashSet(VALID_FORMATS_STR.split(","));
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 0c091bf9ba8..fcba5c355df 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
@@ -49,6 +49,7 @@ import org.apache.beam.sdk.transforms.DoFn.FinishBundle;
import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionRowTuple;
@@ -112,16 +113,56 @@ public class KafkaReadSchemaTransformProvider
consumerConfigs.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, 100);
consumerConfigs.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, autoOffsetReset);
+ String format = configuration.getFormat();
+
+ if (format != null && format.equals("RAW")) {
+ if (inputSchema != null) {
+ throw new IllegalArgumentException(
+ "To read from Kafka in RAW format, you can't provide a schema.");
+ }
+ Schema rawSchema = Schema.builder().addField("payload", Schema.FieldType.BYTES).build();
+ SerializableFunction<byte[], Row> valueMapper = getRawBytesToRowFunction(rawSchema);
+ return new SchemaTransform() {
+ @Override
+ public PCollectionRowTuple expand(PCollectionRowTuple input) {
+ KafkaIO.Read<byte[], byte[]> kafkaRead =
+ KafkaIO.readBytes()
+ .withConsumerConfigUpdates(consumerConfigs)
+ .withConsumerFactoryFn(new ConsumerFactoryWithGcsTrustStores())
+ .withTopic(configuration.getTopic())
+ .withBootstrapServers(configuration.getBootstrapServers());
+ if (isTest) {
+ kafkaRead = kafkaRead.withMaxReadTime(Duration.standardSeconds(testTimeoutSecs));
+ }
+
+ PCollection<byte[]> kafkaValues =
+ input.getPipeline().apply(kafkaRead.withoutMetadata()).apply(Values.create());
+
+ PCollectionTuple outputTuple =
+ kafkaValues.apply(
+ ParDo.of(new ErrorFn("Kafka-read-error-counter", valueMapper))
+ .withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
+
+ return PCollectionRowTuple.of(
+ "output",
+ outputTuple.get(OUTPUT_TAG).setRowSchema(rawSchema),
+ "errors",
+ outputTuple.get(ERROR_TAG).setRowSchema(ERROR_SCHEMA));
+ }
+ };
+ }
+
if (inputSchema != null && !inputSchema.isEmpty()) {
assert Strings.isNullOrEmpty(configuration.getConfluentSchemaRegistryUrl())
: "To read from Kafka, a schema must be provided directly or though Confluent "
+ "Schema Registry, but not both.";
+
final Schema beamSchema =
- Objects.equals(configuration.getFormat(), "JSON")
+ Objects.equals(format, "JSON")
? JsonUtils.beamSchemaFromJsonSchema(inputSchema)
: AvroUtils.toBeamSchema(new org.apache.avro.Schema.Parser().parse(inputSchema));
SerializableFunction<byte[], Row> valueMapper =
- Objects.equals(configuration.getFormat(), "JSON")
+ Objects.equals(format, "JSON")
? JsonUtils.getJsonBytesToRowFunction(beamSchema)
: AvroUtils.getAvroBytesToRowFunction(beamSchema);
return new SchemaTransform() {
@@ -193,6 +234,15 @@ public class KafkaReadSchemaTransformProvider
}
}
+ public static SerializableFunction<byte[], Row> getRawBytesToRowFunction(Schema rawSchema) {
+ return new SimpleFunction<byte[], Row>() {
+ @Override
+ public Row apply(byte[] input) {
+ return Row.withSchema(rawSchema).addValue(input).build();
+ }
+ };
+ }
+
@Override
public String identifier() {
return "beam:schematransform:org.apache.beam:kafka_read:v1";
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 6b9dde4dc95..367f0f7a535 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
@@ -147,4 +147,22 @@ public class KafkaReadSchemaTransformProviderTest {
StandardCharsets.UTF_8))
.build());
}
+
+ @Test
+ public void testBuildTransformWithRawFormat() throws IOException {
+ ServiceLoader<SchemaTransformProvider> serviceLoader =
+ ServiceLoader.load(SchemaTransformProvider.class);
+ List<SchemaTransformProvider> providers =
+ StreamSupport.stream(serviceLoader.spliterator(), false)
+ .filter(provider -> provider.getClass() == KafkaReadSchemaTransformProvider.class)
+ .collect(Collectors.toList());
+ KafkaReadSchemaTransformProvider kafkaProvider =
+ (KafkaReadSchemaTransformProvider) providers.get(0);
+ kafkaProvider.from(
+ KafkaReadSchemaTransformConfiguration.builder()
+ .setTopic("anytopic")
+ .setBootstrapServers("anybootstrap")
+ .setFormat("RAW")
+ .build());
+ }
}
diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml
index d796ecc28bc..6499ffc7e55 100644
--- a/sdks/python/apache_beam/yaml/standard_io.yaml
+++ b/sdks/python/apache_beam/yaml/standard_io.yaml
@@ -46,6 +46,27 @@
config:
gradle_target: 'sdks:java:extensions:sql:expansion-service:shadowJar'
+- type: renaming
+ transforms:
+ 'ReadFromKafka': 'ReadFromKafka'
+ config:
+ mappings:
+ 'ReadFromKafka':
+ 'schema': 'schema'
+ 'consumer_config': 'consumerConfigUpdates'
+ 'format': 'format'
+ 'topic': 'topic'
+ 'bootstrap_servers': 'bootstrapServers'
+ 'confluent_schema_registry_url': 'confluentSchemaRegistryUrl'
+ 'confluent_schema_registry_subject': 'confluentSchemaRegistrySubject'
+ 'auto_offset_reset_config': 'autoOffsetResetConfig'
+ underlying_provider:
+ type: beamJar
+ transforms:
+ 'ReadFromKafka': 'beam:schematransform:org.apache.beam:kafka_read:v1'
+ config:
+ gradle_target: 'sdks:java:io:expansion-service:shadowJar'
+
- type: python
transforms:
'ReadFromBigQuery': 'apache_beam.yaml.yaml_io.read_from_bigquery'