sjvanrossum commented on code in PR #34258:
URL: https://github.com/apache/beam/pull/34258#discussion_r2005922333
##########
sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedSource.java:
##########
@@ -100,8 +102,31 @@ public List<KafkaUnboundedSource<K, V>> split(int
desiredNumSplits, PipelineOpti
}
}
} else {
- for (TopicPartition p : partitions) {
- Lineage.getSources().add("kafka", ImmutableList.of(bootStrapServers,
p.topic()));
+ try (Consumer<?, ?> consumer =
spec.getConsumerFactoryFn().apply(spec.getConsumerConfig())) {
+ Map<String, List<Integer>> topicsAndPartitions = new HashMap<>();
+ try {
+ for (TopicPartition p : partitions) {
+ if (!topicsAndPartitions.containsKey(p.topic())) {
+ List<Integer> partitionsList = new ArrayList<>();
+ List<PartitionInfo> partitionInfoList =
consumer.partitionsFor(p.topic());
+ for (PartitionInfo info : partitionInfoList) {
+ partitionsList.add(info.partition());
+ }
+ topicsAndPartitions.put(p.topic(), partitionsList);
+ }
+ checkState(
+ topicsAndPartitions.containsKey(p.topic())
+ &&
topicsAndPartitions.get(p.topic()).contains(p.partition()),
+ "Partition "
+ + p.partition()
+ + " does not exist for topic "
+ + p.topic()
+ + ". Please check Kafka configuration.");
+ Lineage.getSources().add("kafka",
ImmutableList.of(bootStrapServers, p.topic()));
+ }
+ } catch (KafkaException e) {
+ LOG.warn("Unable to access cluster. Skipping fail fast checks.");
+ }
Review Comment:
This can be simplified by first constructing a `Map<String, Set<Integer>>`
of provided `TopicPartition`s and then iterating over
`topicsAndPartitions.entrySet()`.
--
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: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
