Add support for listing Kafka offsets in bulk #26168
Conversation
|
Thank you for your pull request and welcome to the Trino community. We require contributors to sign our Contributor License Agreement, and we don't seem to have you on file. Continue to work with us on the review and improvements in this PR, and submit the signed CLA to [email protected]. Photos, scans, or digitally-signed PDF files are all suitable. Processing may take a few days. The CLA needs to be on file before we merge your changes. For more information, see https://github.com/trinodb/cla |
There was a problem hiding this comment.
Change commit comment to
pull all partition offsets in a single call to Kafka. -> Retrieve in bulk partition offsets
| @@ -37,6 +37,7 @@ | |||
| import org.apache.kafka.common.config.ConfigResource; | |||
There was a problem hiding this comment.
Squash the two commits into one
| @@ -37,6 +37,7 @@ | |||
| import org.apache.kafka.common.config.ConfigResource; | |||
|
|
|||
| import java.util.Collections; | |||
| import java.util.HashMap; | |||
| import java.util.List; | |||
| import java.util.Map; | |||
| import java.util.Optional; | |||
There was a problem hiding this comment.
In the description
"By changing the call to a bulk operation, the number of API calls can be significantly reduced, improving query startup time."
please add some specific numbers to add the reviewers understand the impact of this change.
|
https://github.com/trinodb/trino/actions/runs/16201941001/job/45742893962?pr=26168 |
| Map<TopicPartition, Long> topicPartitionOffsets = new HashMap<>(); | ||
| topicPartitionOffsetAndTimestamps.forEach((topicPartition, offsetAndTimestamp) -> { | ||
| if (offsetAndTimestamp != null) { | ||
| topicPartitionOffsets.put(topicPartition, offsetAndTimestamp.offset()); | ||
| } | ||
| }); | ||
| return topicPartitionOffsets; |
There was a problem hiding this comment.
return topicPartitionOffsetAndTimestamps.entrySet().stream()
.filter(entry -> entry.getValue() != null)
.collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().offset()));
| Map<TopicPartition, Long> partitionBeginTimestamps = new HashMap<>(); | ||
| partitionBeginOffsets.forEach((partition, partitionIndex) -> { | ||
| partitionBeginTimestamps.put(partition, offsetTimestampRanged.get().begin()); | ||
| }); |
There was a problem hiding this comment.
long partitionBeginTimestamp = floorDiv(offsetTimestampRanged.get().begin(), MICROSECONDS_PER_MILLISECOND);
Map<TopicPartition, Long> partitionBeginTimestamps = partitionBeginOffsets.entrySet().stream()
.collect(Collectors.toMap(Map.Entry::getKey, _ -> partitionBeginTimestamp));
There was a problem hiding this comment.
No need to mutate the map anymore
timestamps.replaceAll((k, v) -> floorDiv(v, MICROSECONDS_PER_MILLISECOND));
in findOffsetsForTimestampGreaterOrEqual method.
| @@ -172,11 +182,17 @@ private boolean isTimestampUpperBoundPushdownEnabled(ConnectorSession session, S | |||
| return KafkaSessionProperties.isTimestampUpperBoundPushdownEnabled(session); | |||
| } | |||
|
|
|||
| private static Optional<Long> findOffsetsForTimestampGreaterOrEqual(KafkaConsumer<byte[], byte[]> kafkaConsumer, TopicPartition topicPartition, long timestamp) | |||
| private static Map<TopicPartition, Long> findOffsetsForTimestampGreaterOrEqual(KafkaConsumer<byte[], byte[]> kafkaConsumer, Map<TopicPartition, Long> timestamps) | |||
There was a problem hiding this comment.
optional: Maybe we could think rather returning Map<TopicPartition, Optional<Long> instead
It is better to avoid having null values.
Description
This PR modifies how the Trino Kafka integration performs translation of timestamps into offsets.
The current implementation makes a Kafka API call per partition to translate the timestamp, however the API can accept a list of partitions as part of the call, allowing for a bulk translation.
By changing the call to a bulk operation, the number of API calls can be significantly reduced, improving query startup time.
Release notes
(X) This is not user-visible or is docs only, and no release notes are required.
Since the only impact for end users is increased query performance, release notes are probably optional.