Passing in key metadata when Reading From Kafka #34426
Conversation
|
Hi @ahmedabu98 , I think this PR is in a good place for some initial feedback. There are a couple of things/open notes I had: (1) Need to still implement this with the Any other comments on the general design/flow would be good as well! |
| beamSchema = Schema.builder().addField("payload", Schema.FieldType.BYTES).build(); | ||
| valueMapper = getRawBytesToRowFunction(beamSchema); | ||
| boolean withKeyMetadata = Boolean.TRUE.equals(configuration.getWithKeyMetadata()); | ||
| if (withKeyMetadata) { |
There was a problem hiding this comment.
This if block is really long. Can we put the logic here into its own function? Something like
if (withKeyMetadata) {
return someFnHere(beamSchema);
} else {
beamSchema = Schema.builder().addField("payload", Schema.FieldType.BYTES).build();
valueMapper = getRawBytesToRowFunction(beamSchema);
}
| } | ||
|
|
||
| @Test | ||
| public void testGetRawBytesKvToRowFunctionEmptyKey() { |
There was a problem hiding this comment.
Can we also add an integration test?
There was a problem hiding this comment.
I just saw above comment about ITs. You can add them to KafkaIOIT, where you can write a pipeline which triggers your code path as if it were a user's pipeline.
There was a problem hiding this comment.
+1 to keep using KafkaIOIT. It should already have some integration tests for KafkaIO as a ManagedIO (look for Managed.write/read), which uses the SchemaTransform under the hood
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
| @@ -139,6 +151,59 @@ Row getConfigurationRow() { | |||
| } | |||
| } | |||
|
|
|||
| private PCollectionRowTuple expandRawWithKeyMetadata( | |||
There was a problem hiding this 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?
| @@ -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; | |||
There was a problem hiding this comment.
Does this ever get initialized?
| .addField("payload", Schema.FieldType.BYTES) | ||
| .build(); | ||
|
|
||
| // 3. Read from Kafka using KafkaIO |
There was a problem hiding this comment.
we wanna test with the SchemaTransform right?
| 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(); | ||
| } | ||
| }; | ||
| } | ||
|
|
There was a problem hiding this comment.
This should probably be a fixed function with a fixed output Schema (i.e. not a method)
Fixes #34391
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.