OffsetFetch v8 (KIP-709): Use batch interface when available (#3024)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: dpkp

kafka/admin/_groups.py

@@ -179,78 +179,126 @@ def list_groups(self, broker_ids=None, states_filter=None, types_filter=None):
 
     # -- List group offsets -------------------------------------------
 
-    def _list_group_offsets_request(self, group_id, partitions=None):
+    def _list_group_offsets_requests(self, group_specs):
         _Topic = OffsetFetchRequest.OffsetFetchRequestTopic
-        if partitions is None:
-            min_version = 1
-            topics = None
+        _Group = OffsetFetchRequest.OffsetFetchRequestGroup
+        _GroupTopic = _Group.OffsetFetchRequestTopics
+        max_version = 8
+
+        groups = []
+        for group_id, partitions in group_specs.items():
+            if partitions is None:
+                group_topics = None
+            else:
+                topics_partitions = defaultdict(set)
+                for topic, partition in partitions:
+                    topics_partitions[topic].add(partition)
+                group_topics = [
+                    _GroupTopic(name=name, partition_indexes=list(parts))
+                    for name, parts in topics_partitions.items()
+                ]
+            groups.append(_Group(group_id=group_id, topics=group_topics))
+
+        if len(groups) == 0:
+            raise ValueError('Empty group_specs!')
+        # Return multple requests when broker does not support v8+
+        if self._manager.broker_version_data.api_version(OffsetFetchRequest) < 8:
+            for group in groups:
+                min_version = 2 if group.topics is None else 0
+                yield (group.group_id, OffsetFetchRequest(group_id=group.group_id,
+                                                          topics=group.topics,
+                                                          min_version=min_version,
+                                                          max_version=max_version))
         else:
-            min_version = 0
-            topics_partitions_dict = defaultdict(set)
-            for topic, partition in partitions:
-                topics_partitions_dict[topic].add(partition)
-            topics = [
-                _Topic(name=name, partition_indexes=list(partitions))
-                for name, partitions in topics_partitions_dict.items()
-            ]
-        return OffsetFetchRequest(group_id=group_id, topics=topics,
-                                  min_version=min_version, max_version=6)
-
-    def _list_group_offsets_process_response(self, response):
-        """Process an OffsetFetchResponse."""
-        if response.API_VERSION > 1:
-            error_type = Errors.for_code(response.error_code)
-            if error_type is not Errors.NoError:
-                raise error_type(
-                    "OffsetFetchResponse failed with response '{}'."
-                    .format(response))
+            yield (None, OffsetFetchRequest(groups=groups,
+                                            min_version=8,
+                                            max_version=max_version))
+
+    @staticmethod
+    def _parse_group_offsets(group):
+        """Build {TopicPartition: OffsetAndMetadata} from an OffsetFetchResponse or OffsetFetchResponseGroup."""
+        error_type = Errors.for_code(group.error_code)
+        if error_type is not Errors.NoError:
+            raise error_type(
+                "OffsetFetchResponse failed for group '{}'.".format(group.group_id))
         results = {}
-        for topic in response.topics:
+        for topic in group.topics:
             for partition in topic.partitions:
                 tp = TopicPartition(topic.name, partition.partition_index)
-                error_type = Errors.for_code(partition.error_code)
-                if error_type is not Errors.NoError:
-                    raise error_type(
+                partition_error = Errors.for_code(partition.error_code)
+                if partition_error is not Errors.NoError:
+                    raise partition_error(
                         f"OffsetFetchResponse failed for partition {tp.partition}")
                 results[tp] = OffsetAndMetadata(
                     offset=partition.committed_offset,
                     metadata=partition.metadata,
-                    leader_epoch=partition.committed_leader_epoch
+                    leader_epoch=partition.committed_leader_epoch,
                 )
         return results
 
-    async def _async_list_group_offsets(self, group_id, group_coordinator_id=None, partitions=None):
-        if group_coordinator_id is None:
-            group_coordinator_id = await self._find_coordinator_id(group_id)
-        request = self._list_group_offsets_request(group_id, partitions)
-        response = await self._manager.send(request, node_id=group_coordinator_id)
-        return self._list_group_offsets_process_response(response)
+    def _list_group_offsets_process_response(self, response, group_id=None):
+        """Process an OffsetFetchResponse."""
+        error_type = Errors.for_code(response.error_code)
+        if error_type is not Errors.NoError:
+            raise error_type(
+                "OffsetFetchResponse failed with response '{}'."
+                .format(response))
+        if response.API_VERSION >= 8:
+            return {group.group_id: self._parse_group_offsets(group)
+                    for group in response.groups}
+        else:
+            return {group_id: self._parse_group_offsets(response)}
 
-    def list_group_offsets(self, group_id, group_coordinator_id=None, partitions=None):
-        """Fetch committed offsets for a single consumer group.
+    async def _async_list_group_offsets(self, group_specs):
+        # Bucket groups by coordinator. One OffsetFetch per coordinator.
+        coordinators_groups = defaultdict(list)
+        for group_id in group_specs:
+            coordinator_id = await self._find_coordinator_id(group_id)
+            coordinators_groups[coordinator_id].append(group_id)
 
-        Note:
-        This does not verify that the group_id or partitions actually exist
-        in the cluster.
+        results = {}
+        _Group = OffsetFetchRequest.OffsetFetchRequestGroup
+        _GroupTopic = _Group.OffsetFetchRequestTopics
+        for coordinator_id, group_ids in coordinators_groups.items():
+            for group_id, request in self._list_group_offsets_requests({group_id: group_specs[group_id]
+                                                                        for group_id in group_ids}):
+                response = await self._manager.send(request, node_id=coordinator_id)
+                results.update(self._list_group_offsets_process_response(response, group_id=group_id))
+        return results
 
-        As soon as any error is encountered, it is immediately raised.
+    def list_group_offsets(self, group_specs):
+        """Fetch committed offsets for one or more consumer groups.
 
-        Arguments:
-            group_id (str): The consumer group id name for which to fetch offsets.
+        On brokers supporting OffsetFetch v8+ (Apache Kafka 3.0+, KIP-709), this
+        issues a single OffsetFetch per coordinator covering all groups
+        hosted by that coordinator. On older brokers it currently only supports
+        one consumer group (per coordinator).
 
-        Keyword Arguments:
-            group_coordinator_id (int, optional): The node_id of the group's coordinator
-                broker. If set to None, will query the cluster to find the group
-                coordinator. Default: None.
-            partitions: A list of TopicPartitions for which to fetch
-                offsets. On brokers >= 0.10.2, this can be set to None to fetch all
-                known offsets for the consumer group. Default: None.
+        Arguments:
+            group_specs (dict): Mapping of group_id (str) to either a list of
+                :class:`~kafka.TopicPartition` to fetch, or None to fetch all
+                committed offsets for that group.
+                Or, one or more group_id (str or list[str]) to fetch all offsets
+                for each group.
 
         Returns:
-            A dict mapping :class:`~kafka.TopicPartition` to
+            A dict mapping group_id (str) to a dict mapping
+                :class:`~kafka.TopicPartition` to
                 :class:`~kafka.structs.OffsetAndMetadata`.
+
+        Raises:
+            UnsupportedVersionError: if multiple groups are requested against
+                a broker that does not support OffsetFetch v8+; or if group_spec
+                with value None against a broker that does not support
+                OffsetFetch v2+.
+            BrokerResponseError: as soon as any group- or partition-level error
+                is encountered.
         """
-        return self._manager.run(self._async_list_group_offsets, group_id, group_coordinator_id, partitions)
+        if isinstance(group_specs, list):
+            group_specs = {group_id: None for group_id in group_specs}
+        elif isinstance(group_specs, str):
+            group_specs = {group_specs: None}
+        return self._manager.run(self._async_list_group_offsets, group_specs)
 
     # -- Delete groups ------------------------------------------------
 
@@ -410,7 +458,8 @@ async def _async_reset_group_offsets(self, group_id, offset_specs, group_coordin
         if group_coordinator_id is None:
             group_coordinator_id = await self._find_coordinator_id(group_id)
 
-        current = await self._async_list_group_offsets(group_id, group_coordinator_id, all_tps)
+        #import pdb; pdb.set_trace()
+        current = (await self._async_list_group_offsets({group_id: list(all_tps)}))[group_id]
         earliest = await self._async_list_partition_offsets({tp: OffsetSpec.EARLIEST for tp in all_tps})
         latest = await self._async_list_partition_offsets({tp: OffsetSpec.LATEST for tp in all_tps})
 

kafka/cli/admin/groups/list_offsets.py

@@ -14,7 +14,7 @@ def add_arguments(cls, parser):
 
     @classmethod
     def command(cls, client, args):
-        offsets = client.list_group_offsets(args.group_id)
+        offsets = client.list_group_offsets(args.group_id)[args.group_id]
         latest = client.list_partition_offsets({tp: OffsetSpec.LATEST for tp in offsets})
         results = defaultdict(dict)
         for tp in latest:

kafka/coordinator/consumer.py

@@ -1097,24 +1097,32 @@ async def _send_offset_fetch_request(self, partitions):
         log.debug("Group %s fetching committed offsets for partitions: %s",
                   self.group_id, '(all)' if partitions is None else partitions)
         # construct the request
-        max_version = 7
+        _Topic = OffsetFetchRequest.OffsetFetchRequestTopic
+        _Group = OffsetFetchRequest.OffsetFetchRequestGroup
+        _GroupTopic = _Group.OffsetFetchRequestTopics
         if partitions is not None:
             topic_partitions = collections.defaultdict(set)
             for tp in partitions:
                 topic_partitions[tp.topic].add(tp.partition)
-            topic_partitions = list(topic_partitions.items())
+            topics = [_Topic(name=t, partition_indexes=list(p))
+                      for t, p in topic_partitions.items()]
+            group_topics = [_GroupTopic(name=t, partition_indexes=list(p))
+                            for t, p in topic_partitions.items()]
             min_version = 0
         else:
-            topic_partitions = None
+            topics = None
+            group_topics = None
             min_version = 2
 
+        groups = [_Group(group_id=self.group_id, topics=group_topics)]
         require_stable = self._isolation_level == IsolationLevel.READ_COMMITTED
         request = OffsetFetchRequest(
             group_id=self.group_id,
-            topics=topic_partitions,
+            topics=topics,
+            groups=groups,
             require_stable=require_stable,
             min_version=min_version,
-            max_version=max_version,
+            max_version=8,
         )
 
         try:
@@ -1126,8 +1134,16 @@ async def _send_offset_fetch_request(self, partitions):
 
     def _handle_offset_fetch_response(self, response):
         log.debug("Received OffsetFetchResponse: %s", response)
-        if response.API_VERSION >= 2 and response.error_code != Errors.NoError.errno:
-            error_type = Errors.for_code(response.error_code)
+        if response.API_VERSION >= 8:
+            group = response.groups[0]
+            top_level_error_code = group.error_code
+            topics = group.topics
+        else:
+            top_level_error_code = response.error_code if response.API_VERSION >= 2 else Errors.NoError.errno
+            topics = response.topics
+
+        if top_level_error_code != Errors.NoError.errno:
+            error_type = Errors.for_code(top_level_error_code)
             log.debug("Offset fetch failed: %s", error_type.__name__)
             error = error_type()
             if error_type is Errors.NotCoordinatorError:
@@ -1142,14 +1158,13 @@ def _handle_offset_fetch_response(self, response):
                 raise error
 
         offsets = {}
-        for topic, partitions in response.topics:
+        for topic, partitions in ((t.name, t.partitions) for t in topics):
             for partition_data in partitions:
-                partition, offset = partition_data[:2]
-                if response.API_VERSION >= 5:
-                    leader_epoch, metadata, error_code = partition_data[2:]
-                else:
-                    metadata, error_code = partition_data[2:]
-                    leader_epoch = -1
+                partition = partition_data.partition_index
+                offset = partition_data.committed_offset
+                leader_epoch = partition_data.committed_leader_epoch
+                metadata = partition_data.metadata
+                error_code = partition_data.error_code
                 tp = TopicPartition(topic, partition)
                 error_type = Errors.for_code(error_code)
                 if error_type is not Errors.NoError: