ahmedabu98 commented on code in PR #34426:
URL: https://github.com/apache/beam/pull/34426#discussion_r2025281767
##########
sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java:
##########
@@ -103,6 +105,16 @@ public Row apply(byte[] input) {
};
}
+ public static SerializableFunction<KV<byte[], byte[]>, Row>
getRawBytesKvToRowFunction(
+ Schema rawSchema) {
+ return new SimpleFunction<KV<byte[], byte[]>, Row>() {
+ @Override
+ public Row apply(KV<byte[], byte[]> input) {
+ return Row.withSchema(rawSchema).addValues(input.getKey(),
input.getValue()).build();
+ }
+ };
+ }
+
Review Comment:
This should probably be a fixed function with a fixed output Schema (i.e.
not a method)
##########
sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProvider.java:
##########
@@ -139,6 +151,59 @@ Row getConfigurationRow() {
}
}
+ private PCollectionRowTuple expandRawWithKeyMetadata(
Review Comment:
There's a lot of duplicate code from `expand()` here :(
I'm a fan of making small code changes if possible. Does it make sense to
consolidate this with the existing `expand()` code block?
##########
sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOIT.java:
##########
@@ -155,6 +160,8 @@ public class KafkaIOIT {
private static final int BUFFER_MEMORY_CONFIG = 100554432;
private static final int RETRY_BACKOFF_MS_CONFIG = 5000;
+ private static String bootstrapServers;
Review Comment:
Does this ever get initialized?
##########
sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOIT.java:
##########
@@ -274,6 +283,77 @@ public void testKafkaIOFailsFastWithInvalidTopics() throws
IOException {
cancelIfTimeouted(readResult, readState);
}
+ @Test
+ public void testKafkaReadWithKeyMetadata() {
+ // 1. Produce test data to Kafka
+ Properties producerProps = new Properties();
+ producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG,
bootstrapServers);
+ producerProps.put(
+ ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG,
ByteArraySerializer.class.getName());
+ producerProps.put(
+ ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG,
ByteArraySerializer.class.getName());
+
+ try (KafkaProducer<byte[], byte[]> producer = new
KafkaProducer<>(producerProps)) {
+ producer.send(
+ new ProducerRecord<>(
+ topicName,
+ "key1".getBytes(StandardCharsets.UTF_8),
+ "value1".getBytes(StandardCharsets.UTF_8)));
+ producer.send(
+ new ProducerRecord<>(
+ topicName,
+ "key2".getBytes(Charset.defaultCharset()),
+ "value2".getBytes(Charset.defaultCharset())));
+ }
+
+ // 2. Define the expected schema for the output Rows
+ Schema expectedSchema =
+ Schema.builder()
+ .addField("key", Schema.FieldType.BYTES)
+ .addField("payload", Schema.FieldType.BYTES)
+ .build();
+
+ // 3. Read from Kafka using KafkaIO
Review Comment:
we wanna test with the SchemaTransform right?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
