Repository: apache/kafka
Pull request #19611: KAFKA-17747: [3/N] Get rid of TopicMetadata in SubscribedTopicDescriberImpl
Author: login: FrankYang0529
Labels: name: performance
Created: 2025-05-01T12:49:12Z | Merged: 2025-05-19T12:46:25Z

Replace `TopicMetadata` with `MetadataImage` in
`SubscribedTopicDescriberImpl` and  `TargetAssignmentBuilder`.

Reviewers: Chia-Ping Tsai <chia7712@gmail.com>, David Jacot
 <djacot@confluent.io>



Diff:
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
index a2e682722d78a..5412b2ab87eca 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
@@ -276,16 +276,13 @@ public class GroupMetadataManager {
 
     private static class UpdateSubscriptionMetadataResult {
         private final int groupEpoch;
-        private final Map<String, TopicMetadata> subscriptionMetadata;
         private final SubscriptionType subscriptionType;
 
         UpdateSubscriptionMetadataResult(
             int groupEpoch,
-            Map<String, TopicMetadata> subscriptionMetadata,
             SubscriptionType subscriptionType
         ) {
             this.groupEpoch = groupEpoch;
-            this.subscriptionMetadata = Objects.requireNonNull(subscriptionMetadata);
             this.subscriptionType = Objects.requireNonNull(subscriptionType);
         }
     }
@@ -2231,7 +2228,6 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, CoordinatorRecord>
         );
 
         int groupEpoch = group.groupEpoch();
-        Map<String, TopicMetadata> subscriptionMetadata = group.subscriptionMetadata();
         SubscriptionType subscriptionType = group.subscriptionType();
 
         if (bumpGroupEpoch || group.hasMetadataExpired(currentTimeMs)) {
@@ -2247,7 +2243,6 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, CoordinatorRecord>
             );
 
             groupEpoch = result.groupEpoch;
-            subscriptionMetadata = result.subscriptionMetadata;
             subscriptionType = result.subscriptionType;
         }
 
@@ -2262,7 +2257,6 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, CoordinatorRecord>
                 groupEpoch,
                 member,
                 updatedMember,
-                subscriptionMetadata,
                 subscriptionType,
                 records
             );
@@ -2373,7 +2367,6 @@ private CoordinatorResult<Void, CoordinatorRecord> classicGroupJoinToConsumerGro
         }
 
         int groupEpoch = group.groupEpoch();
-        Map<String, TopicMetadata> subscriptionMetadata = group.subscriptionMetadata();
         SubscriptionType subscriptionType = group.subscriptionType();
         final ConsumerProtocolSubscription subscription = deserializeSubscription(protocols);
 
@@ -2416,7 +2409,6 @@ private CoordinatorResult<Void, CoordinatorRecord> classicGroupJoinToConsumerGro
             );
 
             groupEpoch = result.groupEpoch;
-            subscriptionMetadata = result.subscriptionMetadata;
             subscriptionType = result.subscriptionType;
         }
 
@@ -2431,7 +2423,6 @@ private CoordinatorResult<Void, CoordinatorRecord> classicGroupJoinToConsumerGro
                 groupEpoch,
                 member,
                 updatedMember,
-                subscriptionMetadata,
                 subscriptionType,
                 records
             );
@@ -2605,7 +2596,6 @@ private CoordinatorResult<Map.Entry<ShareGroupHeartbeatResponseData, Optional<In
                 group,
                 groupEpoch,
                 updatedMember,
-                subscriptionMetadata,
                 subscriptionType,
                 records
             );
@@ -3624,7 +3614,6 @@ private UpdateSubscriptionMetadataResult updateSubscriptionMetadata(
 
         return new UpdateSubscriptionMetadataResult(
             groupEpoch,
-            subscriptionMetadata,
             subscriptionType
         );
     }
@@ -3632,13 +3621,12 @@ private UpdateSubscriptionMetadataResult updateSubscriptionMetadata(
     /**
      * Updates the target assignment according to the updated member and subscription metadata.
      *
-     * @param group                 The ConsumerGroup.
-     * @param groupEpoch            The group epoch.
-     * @param member                The existing member.
-     * @param updatedMember         The updated member.
-     * @param subscriptionMetadata  The subscription metadata.
-     * @param subscriptionType      The group subscription type.
-     * @param records               The list to accumulate any new records.
+     * @param group            The ConsumerGroup.
+     * @param groupEpoch       The group epoch.
+     * @param member           The existing member.
+     * @param updatedMember    The updated member.
+     * @param subscriptionType The group subscription type.
+     * @param records          The list to accumulate any new records.
      * @return The new target assignment.
      */
     private Assignment updateTargetAssignment(
@@ -3646,7 +3634,6 @@ private Assignment updateTargetAssignment(
         int groupEpoch,
         ConsumerGroupMember member,
         ConsumerGroupMember updatedMember,
-        Map<String, TopicMetadata> subscriptionMetadata,
         SubscriptionType subscriptionType,
         List<CoordinatorRecord> records
     ) {
@@ -3659,11 +3646,10 @@ private Assignment updateTargetAssignment(
                 new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(group.groupId(), groupEpoch, consumerGroupAssignors.get(preferredServerAssignor))
                     .withMembers(group.members())
                     .withStaticMembers(group.staticMembers())
-                    .withSubscriptionMetadata(subscriptionMetadata)
                     .withSubscriptionType(subscriptionType)
                     .withTargetAssignment(group.targetAssignment())
                     .withInvertedTargetAssignment(group.invertedTargetAssignment())
-                    .withTopicsImage(metadataImage.topics())
+                    .withMetadataImage(metadataImage)
                     .withResolvedRegularExpressions(group.resolvedRegularExpressions())
                     .addOrUpdateMember(updatedMember.memberId(), updatedMember);
 
@@ -3706,19 +3692,17 @@ private Assignment updateTargetAssignment(
     /**
      * Updates the target assignment according to the updated member and subscription metadata.
      *
-     * @param group                 The ShareGroup.
-     * @param groupEpoch            The group epoch.
-     * @param updatedMember         The updated member.
-     * @param subscriptionMetadata  The subscription metadata.
-     * @param subscriptionType      The group subscription type.
-     * @param records               The list to accumulate any new records.
+     * @param group            The ShareGroup.
+     * @param groupEpoch       The group epoch.
+     * @param updatedMember    The updated member.
+     * @param subscriptionType The group subscription type.
+     * @param records          The list to accumulate any new records.
      * @return The new target assignment.
      */
     private Assignment updateTargetAssignment(
         ShareGroup group,
         int groupEpoch,
         ShareGroupMember updatedMember,
-        Map<String, TopicMetadata> subscriptionMetadata,
         SubscriptionType subscriptionType,
         List<CoordinatorRecord> records
     ) {
@@ -3730,12 +3714,11 @@ private Assignment updateTargetAssignment(
             TargetAssignmentBuilder.ShareTargetAssignmentBuilder assignmentResultBuilder =
                 new TargetAssignmentBuilder.ShareTargetAssignmentBuilder(group.groupId(), groupEpoch, shareGroupAssignor)
                     .withMembers(group.members())
-                    .withSubscriptionMetadata(subscriptionMetadata)
                     .withSubscriptionType(subscriptionType)
                     .withTargetAssignment(group.targetAssignment())
                     .withTopicAssignablePartitionsMap(initializedTopicPartitions)
                     .withInvertedTargetAssignment(group.invertedTargetAssignment())
-                    .withTopicsImage(metadataImage.topics())
+                    .withMetadataImage(metadataImage)
                     .addOrUpdateMember(updatedMember.memberId(), updatedMember);
 
             long startTimeMs = time.milliseconds();
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
index 5a5f2336d14ff..d8237724b748b 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
@@ -19,12 +19,17 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.metadata.BrokerRegistration;
+import org.apache.kafka.metadata.PartitionRegistration;
 
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain
@@ -32,28 +37,26 @@
  */
 public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber {
     /**
-     * The topic Ids mapped to their corresponding {@link TopicMetadata}
-     * object, which contains topic and partition metadata.
+     * The map of topic Ids to the set of allowed partitions for each topic.
+     * If this is empty, all partitions are allowed.
      */
-    private final Map<Uuid, TopicMetadata> topicMetadata;
-    private final Map<Uuid, Set<Integer>> topicPartitionAllowedMap;
+    private final Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap;
 
-    public SubscribedTopicDescriberImpl(Map<Uuid, TopicMetadata> topicMetadata) {
-        this(topicMetadata, null);
-    }
+    /**
+     * The metadata image that contains the latest metadata information.
+     */
+    private final MetadataImage metadataImage;
 
-    public SubscribedTopicDescriberImpl(Map<Uuid, TopicMetadata> topicMetadata, Map<Uuid, Set<Integer>> topicPartitionAllowedMap) {
-        this.topicMetadata = Objects.requireNonNull(topicMetadata);
-        this.topicPartitionAllowedMap = topicPartitionAllowedMap;
+    public SubscribedTopicDescriberImpl(MetadataImage metadataImage) {
+        this(metadataImage, Optional.empty());
     }
 
-    /**
-     * Map of topic Ids to topic metadata.
-     *
-     * @return The map of topic Ids to topic metadata.
-     */
-    public Map<Uuid, TopicMetadata> topicMetadata() {
-        return this.topicMetadata;
+    public SubscribedTopicDescriberImpl(
+        MetadataImage metadataImage,
+        Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap
+    ) {
+        this.metadataImage = Objects.requireNonNull(metadataImage);
+        this.topicPartitionAllowedMap = Objects.requireNonNull(topicPartitionAllowedMap);
     }
 
     /**
@@ -65,8 +68,8 @@ public Map<Uuid, TopicMetadata> topicMetadata() {
      */
     @Override
     public int numPartitions(Uuid topicId) {
-        TopicMetadata topic = this.topicMetadata.get(topicId);
-        return topic == null ? -1 : topic.numPartitions();
+        TopicImage topicImage = this.metadataImage.topics().getTopic(topicId);
+        return topicImage == null ? -1 : topicImage.partitions().size();
     }
 
     /**
@@ -79,13 +82,28 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    BrokerRegistration brokerRegistration = metadataImage.cluster().broker(replica);
+                    if (brokerRegistration != null) {
+                        brokerRegistration.rack().ifPresent(racks::add);
+                    }
+                }
+                return Collections.unmodifiableSet(racks);
+            }
+        }
         return Set.of();
     }
 
     /**
-     * Returns a set of assignable partitions from the topic metadata.
-     * If the allowed partition map is null, all the partitions in the corresponding
-     * topic metadata are returned for the argument topic id. If allowed map is empty,
+     * Returns a set of assignable partitions from the metadata image.
+     * If the allowed partition map is Optional.empty(), all the partitions in the corresponding
+     * topic image are returned for the argument topic id. If allowed map is empty,
      * empty set is returned.
      *
      * @param topicId The uuid of the topic
@@ -93,16 +111,16 @@ public Set<String> racksForPartition(Uuid topicId, int partition) {
      */
     @Override
     public Set<Integer> assignablePartitions(Uuid topicId) {
-        TopicMetadata topic = this.topicMetadata.get(topicId);
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
         if (topic == null) {
             return Set.of();
         }
 
-        if (topicPartitionAllowedMap == null) {
-            return IntStream.range(0, topic.numPartitions()).boxed().collect(Collectors.toUnmodifiableSet());
+        if (topicPartitionAllowedMap.isEmpty()) {
+            return Collections.unmodifiableSet(topic.partitions().keySet());
         }
 
-        return topicPartitionAllowedMap.getOrDefault(topicId, Set.of());
+        return topicPartitionAllowedMap.get().getOrDefault(topicId, Set.of());
     }
 
     @Override
@@ -110,18 +128,22 @@ public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         SubscribedTopicDescriberImpl that = (SubscribedTopicDescriberImpl) o;
-        return topicMetadata.equals(that.topicMetadata);
+        if (!topicPartitionAllowedMap.equals(that.topicPartitionAllowedMap)) return false;
+        return metadataImage.equals(that.metadataImage);
     }
 
     @Override
     public int hashCode() {
-        return topicMetadata.hashCode();
+        int result = metadataImage.hashCode();
+        result = 31 * result + topicPartitionAllowedMap.hashCode();
+        return result;
     }
 
     @Override
     public String toString() {
         return "SubscribedTopicMetadata(" +
-            "topicMetadata=" + topicMetadata +
+            "metadataImage=" + metadataImage +
+            ", topicPartitionAllowedMap=" + topicPartitionAllowedMap +
             ')';
     }
 }
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java
index 63197eaba5ad5..ed1755e483e18 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java
@@ -27,7 +27,7 @@
 import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember;
 import org.apache.kafka.coordinator.group.modern.consumer.ResolvedRegularExpression;
 import org.apache.kafka.coordinator.group.modern.share.ShareGroupMember;
-import org.apache.kafka.image.TopicsImage;
+import org.apache.kafka.image.MetadataImage;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -250,11 +250,6 @@ protected MemberSubscriptionAndAssignmentImpl newMemberSubscriptionAndAssignment
      */
     private Map<String, T> members = Map.of();
 
-    /**
-     * The subscription metadata.
-     */
-    private Map<String, TopicMetadata> subscriptionMetadata = Map.of();
-
     /**
      * The subscription type of the consumer group.
      */
@@ -272,9 +267,9 @@ protected MemberSubscriptionAndAssignmentImpl newMemberSubscriptionAndAssignment
     private Map<Uuid, Map<Integer, String>> invertedTargetAssignment = Map.of();
 
     /**
-     * The topics image.
+     * The metadata image.
      */
-    private TopicsImage topicsImage = TopicsImage.EMPTY;
+    private MetadataImage metadataImage = MetadataImage.EMPTY;
 
     /**
      * The members which have been updated or deleted. Deleted members
@@ -290,7 +285,7 @@ protected MemberSubscriptionAndAssignmentImpl newMemberSubscriptionAndAssignment
     /**
      * Topic partition assignable map.
      */
-    private Map<Uuid, Set<Integer>> topicAssignablePartitionsMap = new HashMap<>();
+    private Optional<Map<Uuid, Set<Integer>>> topicAssignablePartitionsMap = Optional.empty();
 
     /**
      * Constructs the object.
@@ -335,19 +330,6 @@ public U withStaticMembers(
         return self();
     }
 
-    /**
-     * Adds the subscription metadata to use.
-     *
-     * @param subscriptionMetadata  The subscription metadata.
-     * @return This object.
-     */
-    public U withSubscriptionMetadata(
-        Map<String, TopicMetadata> subscriptionMetadata
-    ) {
-        this.subscriptionMetadata = subscriptionMetadata;
-        return self();
-    }
-
     /**
      * Adds the subscription type in use.
      *
@@ -388,22 +370,22 @@ public U withInvertedTargetAssignment(
     }
 
     /**
-     * Adds the topics image.
+     * Adds the metadata image.
      *
-     * @param topicsImage    The topics image.
+     * @param metadataImage    The metadata image.
      * @return This object.
      */
-    public U withTopicsImage(
-        TopicsImage topicsImage
+    public U withMetadataImage(
+        MetadataImage metadataImage
     ) {
-        this.topicsImage = topicsImage;
+        this.metadataImage = metadataImage;
         return self();
     }
 
     public U withTopicAssignablePartitionsMap(
         Map<Uuid, Set<Integer>> topicAssignablePartitionsMap
     ) {
-        this.topicAssignablePartitionsMap = topicAssignablePartitionsMap;
+        this.topicAssignablePartitionsMap = Optional.of(topicAssignablePartitionsMap);
         return self();
     }
 
@@ -445,7 +427,7 @@ public U removeMember(
      */
     public TargetAssignmentResult build() throws PartitionAssignorException {
         Map<String, MemberSubscriptionAndAssignmentImpl> memberSpecs = new HashMap<>();
-        TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(topicsImage);
+        TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics());
 
         // Prepare the member spec for all members.
         members.forEach((memberId, member) ->
@@ -479,15 +461,6 @@ public TargetAssignmentResult build() throws PartitionAssignorException {
             }
         });
 
-        // Prepare the topic metadata.
-        Map<Uuid, TopicMetadata> topicMetadataMap = new HashMap<>();
-        subscriptionMetadata.forEach((topicName, topicMetadata) ->
-            topicMetadataMap.put(
-                topicMetadata.id(),
-                topicMetadata
-            )
-        );
-
         // Compute the assignment.
         GroupAssignment newGroupAssignment = assignor.assign(
             new GroupSpecImpl(
@@ -495,7 +468,7 @@ public TargetAssignmentResult build() throws PartitionAssignorException {
                 subscriptionType,
                 invertedTargetAssignment
             ),
-            new SubscribedTopicDescriberImpl(topicMetadataMap, topicAssignablePartitionsMap)
+            new SubscribedTopicDescriberImpl(metadataImage, topicAssignablePartitionsMap)
         );
 
         // Compute delta from previous to new target assignment and create the relevant records.
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java
index d37b75ce9fee9..1038c3833d710 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.coordinator.group.assignor;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.coordinator.group.MetadataImageBuilder;
 import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
 import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignorException;
@@ -24,7 +25,7 @@
 import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
 import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
 import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
-import org.apache.kafka.coordinator.group.modern.TopicMetadata;
+import org.apache.kafka.image.MetadataImage;
 
 import org.junit.jupiter.api.Test;
 
@@ -61,15 +62,11 @@ public class OptimizedUniformAssignmentBuilderTest {
 
     @Test
     public void testOneMemberNoTopicSubscription() {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .build();
         SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                topic1Uuid,
-                new TopicMetadata(
-                    topic1Uuid,
-                    topic1Name,
-                    3
-                )
-            )
+            metadataImage
         );
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -98,15 +95,11 @@ public void testOneMemberNoTopicSubscription() {
 
     @Test
     public void testOneMemberSubscribedToNonexistentTopic() {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .build();
         SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                topic1Uuid,
-                new TopicMetadata(
-                    topic1Uuid,
-                    topic1Name,
-                    3
-                )
-            )
+            metadataImage
         );
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -131,17 +124,10 @@ public void testOneMemberSubscribedToNonexistentTopic() {
 
     @Test
     public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic3Uuid, topic3Name, 2)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -172,7 +158,9 @@ public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() {
             HOMOGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -185,12 +173,9 @@ public void testFirstAssignmentTwoMembersTwoTopicsNoMemberRacks() {
 
     @Test
     public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic3Uuid, topic3Name, 2)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -232,7 +217,9 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
             HOMOGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -245,22 +232,21 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
 
     @Test
     public void testValidityAndBalanceForLargeSampleSet() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        for (int i = 1; i < 100; i++) {
+        MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder();
+        Set<Uuid> topicIds = new HashSet<>();
+        for (int i = 1; i < 101; i++) {
             Uuid topicId = Uuid.randomUuid();
-            topicMetadata.put(topicId, new TopicMetadata(
-                topicId,
-                "topic-" + i,
-                3
-            ));
+            metadataImageBuilder.addTopic(topicId, "topic-" + i, 3);
+            topicIds.add(topicId);
         }
+        MetadataImage metadataImage = metadataImageBuilder.build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
         for (int i = 1; i < 50; i++) {
             members.put("member" + i, new MemberSubscriptionAndAssignmentImpl(
                 Optional.empty(),
                 Optional.empty(),
-                topicMetadata.keySet(),
+                topicIds,
                 Assignment.EMPTY
             ));
         }
@@ -270,7 +256,9 @@ public void testValidityAndBalanceForLargeSampleSet() {
             HOMOGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -282,17 +270,10 @@ public void testValidityAndBalanceForLargeSampleSet() {
 
     @Test
     public void testReassignmentForTwoMembersTwoTopicsGivenUnbalancedPrevAssignment() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            3
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic2Uuid, topic2Name, 3)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -331,7 +312,9 @@ public void testReassignmentForTwoMembersTwoTopicsGivenUnbalancedPrevAssignment(
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -345,17 +328,10 @@ public void testReassignmentForTwoMembersTwoTopicsGivenUnbalancedPrevAssignment(
     @Test
     public void testReassignmentWhenPartitionsAreAddedForTwoMembersTwoTopics() {
         // Simulating adding partition to T1 and T2 - originally T1 -> 3 Partitions and T2 -> 3 Partitions
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            6
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            5
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 6)
+            .addTopic(topic2Uuid, topic2Name, 5)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -394,7 +370,9 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembersTwoTopics() {
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -407,17 +385,10 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembersTwoTopics() {
 
     @Test
     public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembersTwoTopics() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            3
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic2Uuid, topic2Name, 3)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -466,7 +437,9 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -479,17 +452,10 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe
 
     @Test
     public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeMembersTwoTopics() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            3
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic2Uuid, topic2Name, 3)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -530,7 +496,9 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -543,17 +511,10 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM
 
     @Test
     public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWithTwoMembersTwoTopics() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            2
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 2)
+            .addTopic(topic2Uuid, topic2Name, 2)
+            .build();
 
         // Initial subscriptions were [T1, T2]
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -591,7 +552,9 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -604,12 +567,9 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith
 
     @Test
     public void testReassignmentStickinessWhenAlreadyBalanced() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            5
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 5)
+            .build();
 
         // A TreeMap ensures that memberA is first in the iteration order.
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -660,7 +620,7 @@ public void testReassignmentStickinessWhenAlreadyBalanced() {
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage);
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java
index 5ab032e47ca9f..f83655a3196c7 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/RangeAssignorTest.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.coordinator.group.assignor;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.coordinator.group.MetadataImageBuilder;
 import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
 import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
 import org.apache.kafka.coordinator.group.api.assignor.MemberAssignment;
@@ -28,7 +29,7 @@
 import org.apache.kafka.coordinator.group.modern.MemberAssignmentImpl;
 import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
 import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
-import org.apache.kafka.coordinator.group.modern.TopicMetadata;
+import org.apache.kafka.image.MetadataImage;
 
 import org.junit.jupiter.api.Test;
 
@@ -61,14 +62,7 @@ public class RangeAssignorTest {
     @Test
     public void testOneMemberNoTopic() {
         SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                topic1Uuid,
-                new TopicMetadata(
-                    topic1Uuid,
-                    topic1Name,
-                    3
-                )
-            )
+            MetadataImage.EMPTY
         );
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -102,15 +96,11 @@ public void testOneMemberNoTopic() {
 
     @Test
     public void testOneMemberSubscribedToNonExistentTopic() {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .build();
         SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                topic1Uuid,
-                new TopicMetadata(
-                    topic1Uuid,
-                    topic1Name,
-                    3
-                )
-            )
+            metadataImage
         );
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -135,17 +125,10 @@ public void testOneMemberSubscribedToNonExistentTopic() {
 
     @Test
     public void testFirstAssignmentTwoMembersTwoTopicsSameSubscriptions() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic3Uuid, topic3Name, 2)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -168,7 +151,9 @@ public void testFirstAssignmentTwoMembersTwoTopicsSameSubscriptions() {
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -190,22 +175,11 @@ public void testFirstAssignmentTwoMembersTwoTopicsSameSubscriptions() {
 
     @Test
     public void testFirstAssignmentThreeMembersThreeTopicsDifferentSubscriptions() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            3
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic2Uuid, topic2Name, 3)
+            .addTopic(topic3Uuid, topic3Name, 2)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -235,7 +209,9 @@ public void testFirstAssignmentThreeMembersThreeTopicsDifferentSubscriptions() {
             HETEROGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -260,17 +236,10 @@ public void testFirstAssignmentThreeMembersThreeTopicsDifferentSubscriptions() {
 
     @Test
     public void testFirstAssignmentNumMembersGreaterThanNumPartitions() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic3Uuid, topic3Name, 2)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -300,7 +269,9 @@ public void testFirstAssignmentNumMembersGreaterThanNumPartitions() {
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -326,15 +297,11 @@ public void testFirstAssignmentNumMembersGreaterThanNumPartitions() {
 
     @Test
     public void testStaticMembership() throws PartitionAssignorException {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .build();
         SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                topic1Uuid,
-                new TopicMetadata(
-                    topic1Uuid,
-                    topic1Name,
-                    3
-                )
-            )
+            metadataImage
         );
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -396,15 +363,11 @@ public void testStaticMembership() throws PartitionAssignorException {
 
     @Test
     public void testMixedStaticMembership() throws PartitionAssignorException {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 5)
+            .build();
         SubscribedTopicDescriber subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                topic1Uuid,
-                new TopicMetadata(
-                    topic1Uuid,
-                    topic1Name,
-                    5
-                )
-            )
+            metadataImage
         );
 
         // Initialize members with instance Ids.
@@ -480,17 +443,10 @@ public void testMixedStaticMembership() throws PartitionAssignorException {
 
     @Test
     public void testReassignmentNumMembersGreaterThanNumPartitionsWhenOneMemberAdded() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            2
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 2)
+            .addTopic(topic2Uuid, topic2Name, 2)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -527,7 +483,9 @@ public void testReassignmentNumMembersGreaterThanNumPartitionsWhenOneMemberAdded
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -552,17 +510,10 @@ public void testReassignmentNumMembersGreaterThanNumPartitionsWhenOneMemberAdded
     @Test
     public void testReassignmentWhenOnePartitionAddedForTwoMembersTwoTopics() {
         // Simulating adding a partition - originally T1 -> 3 Partitions and T2 -> 3 Partitions
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            4
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            4
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 4)
+            .addTopic(topic2Uuid, topic2Name, 4)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -591,7 +542,9 @@ public void testReassignmentWhenOnePartitionAddedForTwoMembersTwoTopics() {
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -613,17 +566,10 @@ public void testReassignmentWhenOnePartitionAddedForTwoMembersTwoTopics() {
 
     @Test
     public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembersTwoTopics() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            3
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic2Uuid, topic2Name, 3)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -660,7 +606,9 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -687,17 +635,10 @@ public void testReassignmentWhenOneMemberAddedAfterInitialAssignmentWithTwoMembe
     @Test
     public void testReassignmentWhenOneMemberAddedAndOnePartitionAfterInitialAssignmentWithTwoMembersTwoTopics() {
         // Add a new partition to topic 1, initially T1 -> 3 partitions
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            4
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            3
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 4)
+            .addTopic(topic2Uuid, topic2Name, 3)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -734,7 +675,9 @@ public void testReassignmentWhenOneMemberAddedAndOnePartitionAfterInitialAssignm
             HETEROGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -759,17 +702,10 @@ public void testReassignmentWhenOneMemberAddedAndOnePartitionAfterInitialAssignm
 
     @Test
     public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithTwoMembersTwoTopics() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            3
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic2Uuid, topic2Name, 3)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -790,7 +726,9 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithTwoMem
             HOMOGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -808,22 +746,11 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithTwoMem
 
     @Test
     public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignmentWithThreeMembersTwoTopics() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            3
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic2Uuid, topic2Name, 3)
+            .addTopic(topic3Uuid, topic3Name, 2)
+            .build();
 
         // Let initial subscriptions be A -> T1, T2 // B -> T2 // C -> T2, T3
         // Change the subscriptions to A -> T1 // B -> T1, T2, T3 // C -> T2
@@ -863,7 +790,9 @@ public void testReassignmentWhenMultipleSubscriptionsRemovedAfterInitialAssignme
             HETEROGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java
index 3f33868e99955..01ad2f17d05b7 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/SimpleAssignorTest.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.coordinator.group.assignor;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.coordinator.group.MetadataImageBuilder;
 import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
 import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
 import org.apache.kafka.coordinator.group.api.assignor.MemberAssignment;
@@ -26,7 +27,7 @@
 import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
 import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
 import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
-import org.apache.kafka.coordinator.group.modern.TopicMetadata;
+import org.apache.kafka.image.MetadataImage;
 import org.apache.kafka.server.common.TopicIdPartition;
 
 import org.junit.jupiter.api.Test;
@@ -74,7 +75,7 @@ public void testName() {
     @Test
     public void testAssignWithEmptyMembers() {
         SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of()
+            MetadataImage.EMPTY
         );
 
         GroupSpec groupSpec = new GroupSpecImpl(
@@ -104,15 +105,11 @@ public void testAssignWithEmptyMembers() {
 
     @Test
     public void testAssignWithNoSubscribedTopic() {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
+            .build();
         SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                TOPIC_1_UUID,
-                new TopicMetadata(
-                    TOPIC_1_UUID,
-                    TOPIC_1_NAME,
-                    3
-                )
-            )
+            metadataImage
         );
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -141,15 +138,11 @@ public void testAssignWithNoSubscribedTopic() {
 
     @Test
     public void testAssignWithSubscribedToNonExistentTopic() {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
+            .build();
         SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                TOPIC_1_UUID,
-                new TopicMetadata(
-                    TOPIC_1_UUID,
-                    TOPIC_1_NAME,
-                    3
-                )
-            )
+            metadataImage
         );
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = Map.of(
@@ -174,17 +167,10 @@ public void testAssignWithSubscribedToNonExistentTopic() {
 
     @Test
     public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(TOPIC_1_UUID, new TopicMetadata(
-            TOPIC_1_UUID,
-            TOPIC_1_NAME,
-            3
-        ));
-        topicMetadata.put(TOPIC_3_UUID, new TopicMetadata(
-            TOPIC_3_UUID,
-            TOPIC_3_NAME,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
+            .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new HashMap<>();
 
@@ -211,7 +197,9 @@ public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() {
             HOMOGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -239,23 +227,11 @@ public void testAssignWithTwoMembersAndTwoTopicsHomogeneous() {
 
     @Test
     public void testAssignWithThreeMembersThreeTopicsHeterogeneous() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(TOPIC_1_UUID, new TopicMetadata(
-            TOPIC_1_UUID,
-            TOPIC_1_NAME,
-            3
-        ));
-
-        topicMetadata.put(TOPIC_2_UUID, new TopicMetadata(
-            TOPIC_2_UUID,
-            "topic2",
-            3
-        ));
-        topicMetadata.put(TOPIC_3_UUID, new TopicMetadata(
-            TOPIC_3_UUID,
-            TOPIC_3_NAME,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
+            .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3)
+            .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
+            .build();
 
         Set<Uuid> memberATopicsSubscription = new LinkedHashSet<>();
         memberATopicsSubscription.add(TOPIC_1_UUID);
@@ -291,7 +267,9 @@ public void testAssignWithThreeMembersThreeTopicsHeterogeneous() {
             HETEROGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -321,18 +299,10 @@ public void testAssignWithThreeMembersThreeTopicsHeterogeneous() {
 
     @Test
     public void testAssignWithOneMemberNoAssignedTopicHeterogeneous() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(TOPIC_1_UUID, new TopicMetadata(
-            TOPIC_1_UUID,
-            TOPIC_1_NAME,
-            3
-        ));
-
-        topicMetadata.put(TOPIC_2_UUID, new TopicMetadata(
-            TOPIC_2_UUID,
-            "topic2",
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
+            .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2)
+            .build();
 
         Set<Uuid> memberATopicsSubscription = new LinkedHashSet<>();
         memberATopicsSubscription.add(TOPIC_1_UUID);
@@ -357,7 +327,9 @@ public void testAssignWithOneMemberNoAssignedTopicHeterogeneous() {
             HETEROGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -495,17 +467,10 @@ public void testRoundRobinAssignmentWithCountTooManyPartitions() {
     @Test
     public void testAssignWithCurrentAssignmentHomogeneous() {
         // Current assignment setup - Two members A, B subscribing to T1 and T2.
-        Map<Uuid, TopicMetadata> topicMetadata1 = new HashMap<>();
-        topicMetadata1.put(TOPIC_1_UUID, new TopicMetadata(
-            TOPIC_1_UUID,
-            TOPIC_1_NAME,
-            3
-        ));
-        topicMetadata1.put(TOPIC_2_UUID, new TopicMetadata(
-            TOPIC_2_UUID,
-            TOPIC_2_NAME,
-            2
-        ));
+        MetadataImage metadataImage1 = new MetadataImageBuilder()
+            .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
+            .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members1 = new HashMap<>();
 
@@ -532,7 +497,9 @@ public void testAssignWithCurrentAssignmentHomogeneous() {
             HOMOGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl(topicMetadata1);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl(
+            metadataImage1
+        );
 
         GroupAssignment computedAssignment1 = assignor.assign(
             groupSpec1,
@@ -558,17 +525,10 @@ public void testAssignWithCurrentAssignmentHomogeneous() {
         assertAssignment(expectedAssignment1, computedAssignment1);
 
         // New assignment setup - Three members A, B, C subscribing to T2 and T3.
-        Map<Uuid, TopicMetadata> topicMetadata2 = new HashMap<>();
-        topicMetadata2.put(TOPIC_2_UUID, new TopicMetadata(
-            TOPIC_2_UUID,
-            TOPIC_2_NAME,
-            2
-        ));
-        topicMetadata2.put(TOPIC_3_UUID, new TopicMetadata(
-            TOPIC_3_UUID,
-            TOPIC_3_NAME,
-            3
-        ));
+        MetadataImage metadataImage2 = new MetadataImageBuilder()
+            .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 2)
+            .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 3)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members2 = new HashMap<>();
 
@@ -607,7 +567,9 @@ public void testAssignWithCurrentAssignmentHomogeneous() {
             HOMOGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl(topicMetadata2);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl(
+            metadataImage2
+        );
 
         GroupAssignment computedAssignment2 = assignor.assign(
             groupSpec2,
@@ -639,23 +601,11 @@ public void testAssignWithCurrentAssignmentHomogeneous() {
     @Test
     public void testAssignWithCurrentAssignmentHeterogeneous() {
         // Current assignment setup - 3 members A - {T1, T2}, B - {T3}, C - {T2, T3}.
-        Map<Uuid, TopicMetadata> topicMetadata1 = new HashMap<>();
-        topicMetadata1.put(TOPIC_1_UUID, new TopicMetadata(
-            TOPIC_1_UUID,
-            TOPIC_1_NAME,
-            3
-        ));
-
-        topicMetadata1.put(TOPIC_2_UUID, new TopicMetadata(
-            TOPIC_2_UUID,
-            TOPIC_2_NAME,
-            3
-        ));
-        topicMetadata1.put(TOPIC_3_UUID, new TopicMetadata(
-            TOPIC_3_UUID,
-            TOPIC_3_NAME,
-            2
-        ));
+        MetadataImage metadataImage1 = new MetadataImageBuilder()
+            .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
+            .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3)
+            .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
+            .build();
 
         Set<Uuid> memberATopicsSubscription1 = new LinkedHashSet<>();
         memberATopicsSubscription1.add(TOPIC_1_UUID);
@@ -691,7 +641,9 @@ public void testAssignWithCurrentAssignmentHeterogeneous() {
             HETEROGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl(topicMetadata1);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata1 = new SubscribedTopicDescriberImpl(
+            metadataImage1
+        );
 
         GroupAssignment computedAssignment1 = assignor.assign(
             groupSpec1,
@@ -720,27 +672,12 @@ public void testAssignWithCurrentAssignmentHeterogeneous() {
 
         // New assignment setup - 2 members A - {T1, T2, T3}, B - {T3, T4}.
 
-        Map<Uuid, TopicMetadata> topicMetadata2 = new HashMap<>();
-        topicMetadata2.put(TOPIC_1_UUID, new TopicMetadata(
-            TOPIC_1_UUID,
-            TOPIC_1_NAME,
-            3
-        ));
-        topicMetadata2.put(TOPIC_2_UUID, new TopicMetadata(
-            TOPIC_2_UUID,
-            TOPIC_2_NAME,
-            3
-        ));
-        topicMetadata2.put(TOPIC_3_UUID, new TopicMetadata(
-            TOPIC_3_UUID,
-            TOPIC_3_NAME,
-            2
-        ));
-        topicMetadata2.put(TOPIC_4_UUID, new TopicMetadata(
-            TOPIC_4_UUID,
-            TOPIC_4_NAME,
-            1
-        ));
+        MetadataImage metadataImage2 = new MetadataImageBuilder()
+            .addTopic(TOPIC_1_UUID, TOPIC_1_NAME, 3)
+            .addTopic(TOPIC_2_UUID, TOPIC_2_NAME, 3)
+            .addTopic(TOPIC_3_UUID, TOPIC_3_NAME, 2)
+            .addTopic(TOPIC_4_UUID, TOPIC_4_NAME, 1)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members2 = new HashMap<>();
 
@@ -776,7 +713,9 @@ public void testAssignWithCurrentAssignmentHeterogeneous() {
             Map.of()
         );
 
-        SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl(topicMetadata2);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata2 = new SubscribedTopicDescriberImpl(
+            metadataImage2
+        );
 
         GroupAssignment computedAssignment2 = assignor.assign(
             groupSpec2,
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java
index decffabb6f492..2c12ed467f17c 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/UniformHeterogeneousAssignmentBuilderTest.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.coordinator.group.assignor;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.coordinator.group.MetadataImageBuilder;
 import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment;
 import org.apache.kafka.coordinator.group.api.assignor.GroupSpec;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignorException;
@@ -25,7 +26,7 @@
 import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
 import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
 import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
-import org.apache.kafka.coordinator.group.modern.TopicMetadata;
+import org.apache.kafka.image.MetadataImage;
 
 import org.junit.jupiter.api.Test;
 
@@ -89,15 +90,11 @@ public Collection<String> memberIds() {
 
     @Test
     public void testTwoMembersNoTopicSubscription() {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .build();
         SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                topic1Uuid,
-                new TopicMetadata(
-                    topic1Uuid,
-                    topic1Name,
-                    3
-                )
-            )
+            metadataImage
         );
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -130,15 +127,11 @@ public void testTwoMembersNoTopicSubscription() {
 
     @Test
     public void testTwoMembersSubscribedToNonexistentTopics() {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .build();
         SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            Map.of(
-                topic1Uuid,
-                new TopicMetadata(
-                    topic1Uuid,
-                    topic1Name,
-                    3
-                )
-            )
+            metadataImage
         );
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -168,17 +161,10 @@ public void testTwoMembersSubscribedToNonexistentTopics() {
 
     @Test
     public void testFirstAssignmentTwoMembersTwoTopics() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            6
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic3Uuid, topic3Name, 6)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -201,7 +187,9 @@ public void testFirstAssignmentTwoMembersTwoTopics() {
             HETEROGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -222,17 +210,10 @@ public void testFirstAssignmentTwoMembersTwoTopics() {
 
     @Test
     public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            1
-        ));
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 2)
+            .addTopic(topic3Uuid, topic3Name, 1)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -262,7 +243,9 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
             HETEROGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -286,22 +269,11 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
 
     @Test
     public void testReassignmentForTwoMembersThreeTopicsGivenUnbalancedPrevAssignment() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            6
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            4
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            4
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 6)
+            .addTopic(topic2Uuid, topic2Name, 4)
+            .addTopic(topic3Uuid, topic3Name, 4)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -340,7 +312,9 @@ public void testReassignmentForTwoMembersThreeTopicsGivenUnbalancedPrevAssignmen
             HETEROGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -366,27 +340,14 @@ public void testReassignmentForTwoMembersThreeTopicsGivenUnbalancedPrevAssignmen
     @Test
     public void testReassignmentWhenPartitionsAreAddedForTwoMembers() {
         // Simulating adding partitions to T1, T2, T3 - originally T1 -> 4, T2 -> 3, T3 -> 2, T4 -> 3
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            6
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            5
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            3
-        ));
-        topicMetadata.put(topic4Uuid, new TopicMetadata(
-            topic4Uuid,
-            topic4Name,
-            3
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 6)
+            .addTopic(topic2Uuid, topic2Name, 5)
+            .addTopic(topic3Uuid, topic3Name, 3)
+            .addTopic(topic4Uuid, topic4Name, 3)
+            .build();
+
+
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -415,7 +376,9 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembers() {
             HETEROGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -438,17 +401,10 @@ public void testReassignmentWhenPartitionsAreAddedForTwoMembers() {
     @Test
     public void testReassignmentWhenOneMemberAddedAndPartitionsAddedTwoMembersTwoTopics() {
         // Initially T1 -> 3, T2 -> 3 partitions.
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            6
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            7
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 6)
+            .addTopic(topic2Uuid, topic2Name, 7)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -485,7 +441,9 @@ public void testReassignmentWhenOneMemberAddedAndPartitionsAddedTwoMembersTwoTop
             HETEROGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -509,22 +467,11 @@ public void testReassignmentWhenOneMemberAddedAndPartitionsAddedTwoMembersTwoTop
 
     @Test
     public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeMembersThreeTopics() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            8
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            3
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic2Uuid, topic2Name, 8)
+            .addTopic(topic3Uuid, topic3Name, 3)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -554,7 +501,9 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM
             HETEROGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -575,17 +524,10 @@ public void testReassignmentWhenOneMemberRemovedAfterInitialAssignmentWithThreeM
 
     @Test
     public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWithTwoMembersTwoTopics() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            5
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addTopic(topic2Uuid, topic2Name, 5)
+            .build();
 
         // Initial subscriptions were [T1, T2]
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
@@ -615,7 +557,9 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith
             HETEROGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -641,22 +585,11 @@ public void testReassignmentWhenOneSubscriptionRemovedAfterInitialAssignmentWith
      */
     @Test
     public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            2
-        ));
-        topicMetadata.put(topic2Uuid, new TopicMetadata(
-            topic2Uuid,
-            topic2Name,
-            2
-        ));
-        topicMetadata.put(topic3Uuid, new TopicMetadata(
-            topic3Uuid,
-            topic3Name,
-            2
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 2)
+            .addTopic(topic2Uuid, topic2Name, 2)
+            .addTopic(topic3Uuid, topic3Name, 2)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -683,7 +616,9 @@ public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions(
             HETEROGENEOUS,
             invertedTargetAssignment(members)
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
@@ -705,12 +640,9 @@ public void testReassignmentWhenTopicPartitionsRunOutAndMembersHaveNoPartitions(
 
     @Test
     public void testFirstAssignmentWithTwoMembersIncludingOneWithoutSubscriptions() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        topicMetadata.put(topic1Uuid, new TopicMetadata(
-            topic1Uuid,
-            topic1Name,
-            3
-        ));
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .build();
 
         Map<String, MemberSubscriptionAndAssignmentImpl> members = new TreeMap<>();
 
@@ -733,7 +665,9 @@ public void testFirstAssignmentWithTwoMembersIncludingOneWithoutSubscriptions()
             HETEROGENEOUS,
             Map.of()
         );
-        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadata);
+        SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
+            metadataImage
+        );
 
         GroupAssignment computedAssignment = assignor.assign(
             groupSpec,
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java
index c0fd4a5a400c8..b2ecfcbce78d4 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java
@@ -17,12 +17,14 @@
 package org.apache.kafka.coordinator.group.modern;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.coordinator.group.MetadataImageBuilder;
+import org.apache.kafka.image.MetadataImage;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -31,80 +33,98 @@
 
 public class SubscribedTopicMetadataTest {
 
-    private Map<Uuid, TopicMetadata> topicMetadataMap;
     private SubscribedTopicDescriberImpl subscribedTopicMetadata;
+    private MetadataImage metadataImage;
+    private final int numPartitions = 5;
 
     @BeforeEach
     public void setUp() {
-        topicMetadataMap = new HashMap<>();
+        MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder();
         for (int i = 0; i < 5; i++) {
             Uuid topicId = Uuid.randomUuid();
             String topicName = "topic" + i;
-            topicMetadataMap.put(
-                topicId,
-                new TopicMetadata(topicId, topicName, 5)
-            );
+            metadataImageBuilder.addTopic(topicId, topicName, numPartitions);
         }
-        subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadataMap);
+        metadataImage = metadataImageBuilder.addRacks().build();
+
+        subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage);
     }
 
     @Test
-    public void testAttribute() {
-        assertEquals(topicMetadataMap, subscribedTopicMetadata.topicMetadata());
+    public void testMetadataImageCannotBeNull() {
+        assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(null));
     }
 
     @Test
-    public void testTopicMetadataCannotBeNull() {
-        assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(null));
+    public void testTopicPartitionAllowedMapCannotBeNull() {
+        assertThrows(NullPointerException.class, () -> new SubscribedTopicDescriberImpl(metadataImage, null));
     }
 
     @Test
     public void testNumberOfPartitions() {
         Uuid topicId = Uuid.randomUuid();
 
-        // Test -1 is returned when the topic Id doesn't exist.
+        // Test -1 is returned when the topic ID doesn't exist.
         assertEquals(-1, subscribedTopicMetadata.numPartitions(topicId));
 
-        topicMetadataMap.put(topicId, new TopicMetadata(topicId, "topic6", 3));
+        // Test that the correct number of partitions are returned for a given topic ID.
+        metadataImage.topics().topicsById().forEach((id, name) ->
+            // Test that the correct number of partitions are returned for a given topic ID.
+            assertEquals(numPartitions, subscribedTopicMetadata.numPartitions(id))
+        );
+    }
+
+    @Test
+    public void testRacksForPartition() {
+        Uuid topicId = Uuid.randomUuid();
+
+        // Test empty set is returned when the topic ID doesn't exist.
+        assertEquals(Set.of(), subscribedTopicMetadata.racksForPartition(topicId, 0));
+        metadataImage.topics().topicsById().forEach((id, name) -> {
+            // Test empty set is returned when the partition ID doesn't exist.
+            assertEquals(Set.of(), subscribedTopicMetadata.racksForPartition(id, 10));
 
-        // Test that the correct number of partitions are returned for a given topic Id.
-        assertEquals(3, subscribedTopicMetadata.numPartitions(topicId));
+            // Test that the correct racks of partition are returned for a given topic ID.
+            assertEquals(Set.of("rack0", "rack1"), subscribedTopicMetadata.racksForPartition(id, 0));
+        });
     }
 
     @Test
     public void testEquals() {
-        assertEquals(new SubscribedTopicDescriberImpl(topicMetadataMap), subscribedTopicMetadata);
+        assertEquals(new SubscribedTopicDescriberImpl(metadataImage), subscribedTopicMetadata);
 
-        Map<Uuid, TopicMetadata> topicMetadataMap2 = new HashMap<>();
         Uuid topicId = Uuid.randomUuid();
-        topicMetadataMap2.put(topicId, new TopicMetadata(topicId, "newTopic", 5));
-        assertNotEquals(new SubscribedTopicDescriberImpl(topicMetadataMap2), subscribedTopicMetadata);
+        MetadataImage metadataImage2 = new MetadataImageBuilder()
+            .addTopic(topicId, "newTopic", 5)
+            .addRacks()
+            .build();
+        assertNotEquals(new SubscribedTopicDescriberImpl(metadataImage2), subscribedTopicMetadata);
     }
 
     @Test
     public void testAssignablePartitions() {
-        // null allow map (all partitions assignable)
-        subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadataMap, null);
         String t1Name = "t1";
         Uuid t1Id = Uuid.randomUuid();
-        topicMetadataMap.put(t1Id, new TopicMetadata(t1Id, t1Name, 5));
+        metadataImage = new MetadataImageBuilder().addTopic(t1Id, t1Name, numPartitions).build();
+        // Optional.empty() allow map (all partitions assignable)
+        subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage, Optional.empty());
         assertEquals(Set.of(0, 1, 2, 3, 4), subscribedTopicMetadata.assignablePartitions(t1Id));
 
         // empty allow map (nothing assignable)
-        subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadataMap, Map.of());
+        subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage, Optional.of(Map.of()));
         assertEquals(Set.of(), subscribedTopicMetadata.assignablePartitions(t1Id));
 
         // few assignable partitions
         subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            topicMetadataMap,
-            Map.of(t1Id, Set.of(0, 5))
+            metadataImage,
+            Optional.of(Map.of(t1Id, Set.of(0, 5)))
         );
         assertEquals(Set.of(0, 5), subscribedTopicMetadata.assignablePartitions(t1Id));
 
         // all assignable partitions
         subscribedTopicMetadata = new SubscribedTopicDescriberImpl(
-            topicMetadataMap,
-            Map.of(t1Id, Set.of(0, 1, 2, 3, 4))
+            metadataImage,
+            Optional.of(Map.of(t1Id, Set.of(0, 1, 2, 3, 4)))
         );
         assertEquals(Set.of(0, 1, 2, 3, 4), subscribedTopicMetadata.assignablePartitions(t1Id));
     }
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java
index f33eed569180a..59541ebad6a23 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java
@@ -25,7 +25,7 @@
 import org.apache.kafka.coordinator.group.api.assignor.SubscriptionType;
 import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember;
 import org.apache.kafka.coordinator.group.modern.consumer.ResolvedRegularExpression;
-import org.apache.kafka.image.TopicsImage;
+import org.apache.kafka.image.MetadataImage;
 
 import org.junit.jupiter.api.Test;
 
@@ -57,13 +57,12 @@ public static class TargetAssignmentBuilderTestContext {
         private final int groupEpoch;
         private final PartitionAssignor assignor = mock(PartitionAssignor.class);
         private final Map<String, ConsumerGroupMember> members = new HashMap<>();
-        private final Map<String, TopicMetadata> subscriptionMetadata = new HashMap<>();
         private final Map<String, ConsumerGroupMember> updatedMembers = new HashMap<>();
         private final Map<String, Assignment> targetAssignment = new HashMap<>();
         private final Map<String, MemberAssignment> memberAssignments = new HashMap<>();
         private final Map<String, String> staticMembers = new HashMap<>();
         private final Map<String, ResolvedRegularExpression> resolvedRegularExpressions = new HashMap<>();
-        private MetadataImageBuilder topicsImageBuilder = new MetadataImageBuilder();
+        private MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder();
 
         public TargetAssignmentBuilderTestContext(
             String groupId,
@@ -123,12 +122,7 @@ public Uuid addTopicMetadata(
             int numPartitions
         ) {
             Uuid topicId = Uuid.randomUuid();
-            subscriptionMetadata.put(topicName, new TopicMetadata(
-                topicId,
-                topicName,
-                numPartitions
-            ));
-            topicsImageBuilder = topicsImageBuilder.addTopic(topicId, topicName, numPartitions);
+            metadataImageBuilder = metadataImageBuilder.addTopic(topicId, topicName, numPartitions);
 
             return topicId;
         }
@@ -218,8 +212,8 @@ private MemberSubscriptionAndAssignmentImpl newMemberSubscriptionAndAssignment(
         }
 
         public TargetAssignmentBuilder.TargetAssignmentResult build() {
-            TopicsImage topicsImage = topicsImageBuilder.build().topics();
-            TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(topicsImage);
+            MetadataImage metadataImage = metadataImageBuilder.build();
+            TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics());
             // Prepare expected member specs.
             Map<String, MemberSubscriptionAndAssignmentImpl> memberSubscriptions = new HashMap<>();
 
@@ -256,13 +250,8 @@ public TargetAssignmentBuilder.TargetAssignmentResult build() {
                 }
             });
 
-            // Prepare the expected topic metadata.
-            Map<Uuid, TopicMetadata> topicMetadataMap = new HashMap<>();
-            subscriptionMetadata.forEach((topicName, topicMetadata) ->
-                topicMetadataMap.put(topicMetadata.id(), topicMetadata));
-
             // Prepare the expected subscription topic metadata.
-            SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(topicMetadataMap);
+            SubscribedTopicDescriberImpl subscribedTopicMetadata = new SubscribedTopicDescriberImpl(metadataImage);
             SubscriptionType subscriptionType = HOMOGENEOUS;
 
             // Prepare the member assignments per topic partition.
@@ -286,11 +275,10 @@ public TargetAssignmentBuilder.TargetAssignmentResult build() {
                 new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(groupId, groupEpoch, assignor)
                     .withMembers(members)
                     .withStaticMembers(staticMembers)
-                    .withSubscriptionMetadata(subscriptionMetadata)
                     .withSubscriptionType(subscriptionType)
                     .withTargetAssignment(targetAssignment)
                     .withInvertedTargetAssignment(invertedTargetAssignment)
-                    .withTopicsImage(topicsImage)
+                    .withMetadataImage(metadataImage)
                     .withResolvedRegularExpressions(resolvedRegularExpressions);
 
             // Add the updated members or delete the deleted members.
diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java
index 9402728e1eba2..d6515a5f56536 100644
--- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java
+++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java
@@ -27,12 +27,10 @@
 import org.apache.kafka.coordinator.group.modern.GroupSpecImpl;
 import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
 import org.apache.kafka.coordinator.group.modern.TopicIds;
-import org.apache.kafka.coordinator.group.modern.TopicMetadata;
 import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember;
 import org.apache.kafka.image.MetadataDelta;
 import org.apache.kafka.image.MetadataImage;
 import org.apache.kafka.image.MetadataProvenance;
-import org.apache.kafka.image.TopicsImage;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -89,69 +87,29 @@ public static List<String> createTopicNames(int topicCount) {
     }
 
     /**
-     * Creates a subscription metadata map for the given topics.
+     * Creates a TopicsImage from the given topic names.
      *
-     * @param topicNames                The names of the topics.
-     * @param partitionsPerTopic        The number of partitions per topic.
-     * @return The subscription metadata map.
-     */
-    public static Map<String, TopicMetadata> createSubscriptionMetadata(
-        List<String> topicNames,
-        int partitionsPerTopic
-    ) {
-        Map<String, TopicMetadata> subscriptionMetadata = new HashMap<>();
-
-        for (String topicName : topicNames) {
-            Uuid topicId = Uuid.randomUuid();
-
-            TopicMetadata metadata = new TopicMetadata(
-                topicId,
-                topicName,
-                partitionsPerTopic
-            );
-            subscriptionMetadata.put(topicName, metadata);
-        }
-
-        return subscriptionMetadata;
-    }
-
-    /**
-     * Creates a topic metadata map from the given subscription metadata.
-     *
-     * @param subscriptionMetadata  The subscription metadata.
-     * @return The topic metadata map.
-     */
-    public static Map<Uuid, TopicMetadata> createTopicMetadata(
-        Map<String, TopicMetadata> subscriptionMetadata
-    ) {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>((int) (subscriptionMetadata.size() / 0.75f + 1));
-        for (Map.Entry<String, TopicMetadata> entry : subscriptionMetadata.entrySet()) {
-            topicMetadata.put(entry.getValue().id(), entry.getValue());
-        }
-        return topicMetadata;
-    }
-
-    /**
-     * Creates a TopicsImage from the given subscription metadata.
-     *
-     * @param subscriptionMetadata  The subscription metadata.
+     * @param allTopicNames      The topic names.
+     * @param partitionsPerTopic Number of partitions per topic.
      * @return A TopicsImage containing the topic ids, names and partition counts from the
      *         subscription metadata.
      */
-    public static TopicsImage createTopicsImage(Map<String, TopicMetadata> subscriptionMetadata) {
+    public static MetadataImage createMetadataImage(
+        List<String> allTopicNames,
+        int partitionsPerTopic
+    ) {
         MetadataDelta delta = new MetadataDelta(MetadataImage.EMPTY);
 
-        for (Map.Entry<String, TopicMetadata> entry : subscriptionMetadata.entrySet()) {
-            TopicMetadata topicMetadata = entry.getValue();
+        for (String topicName : allTopicNames) {
             AssignorBenchmarkUtils.addTopic(
                 delta,
-                topicMetadata.id(),
-                topicMetadata.name(),
-                topicMetadata.numPartitions()
+                Uuid.randomUuid(),
+                topicName,
+                partitionsPerTopic
             );
         }
 
-        return delta.apply(MetadataProvenance.EMPTY).topics();
+        return delta.apply(MetadataProvenance.EMPTY);
     }
 
     /**
diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java
index 971828c1c4726..fa232c9427fc7 100644
--- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java
+++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/ServerSideAssignorBenchmark.java
@@ -31,9 +31,8 @@
 import org.apache.kafka.coordinator.group.modern.MemberSubscriptionAndAssignmentImpl;
 import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
 import org.apache.kafka.coordinator.group.modern.TopicIds;
-import org.apache.kafka.coordinator.group.modern.TopicMetadata;
 import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember;
-import org.apache.kafka.image.TopicsImage;
+import org.apache.kafka.image.MetadataImage;
 
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -125,9 +124,7 @@ public enum AssignmentType {
 
     private List<String> allTopicNames = Collections.emptyList();
 
-    private Map<String, TopicMetadata> subscriptionMetadata = Collections.emptyMap();
-
-    private TopicsImage topicsImage = TopicsImage.EMPTY;
+    private MetadataImage metadataImage = MetadataImage.EMPTY;
 
     private TopicIds.TopicResolver topicResolver;
 
@@ -151,16 +148,11 @@ private void setupTopics() {
         allTopicNames = AssignorBenchmarkUtils.createTopicNames(topicCount);
 
         int partitionsPerTopic = (memberCount * partitionsToMemberRatio) / topicCount;
-        subscriptionMetadata = AssignorBenchmarkUtils.createSubscriptionMetadata(
-            allTopicNames,
-            partitionsPerTopic
-        );
 
-        topicsImage = AssignorBenchmarkUtils.createTopicsImage(subscriptionMetadata);
-        topicResolver = new TopicIds.CachedTopicResolver(topicsImage);
+        metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionsPerTopic);
+        topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics());
 
-        Map<Uuid, TopicMetadata> topicMetadata = AssignorBenchmarkUtils.createTopicMetadata(subscriptionMetadata);
-        subscribedTopicDescriber = new SubscribedTopicDescriberImpl(topicMetadata);
+        subscribedTopicDescriber = new SubscribedTopicDescriberImpl(metadataImage);
     }
 
     private Map<String, ConsumerGroupMember> createMembers() {
diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java
index 6fbb7908622b0..c150d94551a89 100644
--- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java
+++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/TargetAssignmentBuilderBenchmark.java
@@ -28,9 +28,8 @@
 import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
 import org.apache.kafka.coordinator.group.modern.TargetAssignmentBuilder;
 import org.apache.kafka.coordinator.group.modern.TopicIds;
-import org.apache.kafka.coordinator.group.modern.TopicMetadata;
 import org.apache.kafka.coordinator.group.modern.consumer.ConsumerGroupMember;
-import org.apache.kafka.image.TopicsImage;
+import org.apache.kafka.image.MetadataImage;
 
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -93,9 +92,7 @@ public class TargetAssignmentBuilderBenchmark {
 
     private List<String> allTopicNames = Collections.emptyList();
 
-    private Map<String, TopicMetadata> subscriptionMetadata = Collections.emptyMap();
-
-    private TopicsImage topicsImage;
+    private MetadataImage metadataImage;
 
     private TopicIds.TopicResolver topicResolver;
 
@@ -118,11 +115,10 @@ public void setup() {
 
         targetAssignmentBuilder = new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(GROUP_ID, GROUP_EPOCH, partitionAssignor)
             .withMembers(members)
-            .withSubscriptionMetadata(subscriptionMetadata)
             .withSubscriptionType(subscriptionType)
             .withTargetAssignment(existingTargetAssignment)
             .withInvertedTargetAssignment(invertedTargetAssignment)
-            .withTopicsImage(topicsImage)
+            .withMetadataImage(metadataImage)
             .addOrUpdateMember(newMember.memberId(), newMember);
     }
 
@@ -130,16 +126,11 @@ private void setupTopics() {
         allTopicNames = AssignorBenchmarkUtils.createTopicNames(topicCount);
 
         int partitionsPerTopic = (memberCount * partitionsToMemberRatio) / topicCount;
-        subscriptionMetadata = AssignorBenchmarkUtils.createSubscriptionMetadata(
-            allTopicNames,
-            partitionsPerTopic
-        );
 
-        topicsImage = AssignorBenchmarkUtils.createTopicsImage(subscriptionMetadata);
-        topicResolver = new TopicIds.CachedTopicResolver(topicsImage);
+        metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionsPerTopic);
+        topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics());
 
-        Map<Uuid, TopicMetadata> topicMetadata = AssignorBenchmarkUtils.createTopicMetadata(subscriptionMetadata);
-        subscribedTopicDescriber = new SubscribedTopicDescriberImpl(topicMetadata);
+        subscribedTopicDescriber = new SubscribedTopicDescriberImpl(metadataImage);
     }
 
     private Map<String, Assignment> generateMockInitialTargetAssignmentAndUpdateInvertedTargetAssignment(


Reviews:
user: login: dajac
body: @FrankYang0529 Thanks for the patch! I have started to look into it and I left a few high level comments to start with. Out of curiosity, have you tried to run the benchmarks before and after the change to see the impact?
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T15:19:20Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T06:19:40Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T06:20:54Z
user: login: AndrewJSchofield
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T17:26:51Z
user: login: dajac
body: @FrankYang0529 Thanks for the update. I just made a pass and I left some minor comments for consideration.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-17T16:58:47Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-18T05:31:38Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-18T05:32:34Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-18T05:33:27Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-18T05:33:33Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-18T05:38:30Z
user: login: dajac
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-18T06:54:01Z
user: login: chia7712
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-18T16:30:18Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T03:43:48Z
user: login: dajac
body: @FrankYang0529 Thanks for the update. One more nit and we should be good to do.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T07:02:48Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T07:08:11Z
user: login: dajac
body: lgtm, thanks
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-19T07:24:32Z
user: login: chia7712
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T07:31:22Z
user: login: chia7712
body: @FrankYang0529 thanks for this patch.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T08:33:24Z
user: login: dajac
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T09:19:00Z
user: login: dajac
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T09:20:53Z
user: login: dajac
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T09:21:11Z
user: login: dajac
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T09:25:44Z
user: login: chia7712
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T09:32:25Z
user: login: FrankYang0529
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T09:35:54Z
user: login: chia7712
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-19T12:45:32Z

Review comments:
diff_hunk: @@ -3652,11 +3652,11 @@ private Assignment updateTargetAssignment(
                 new TargetAssignmentBuilder.ConsumerTargetAssignmentBuilder(group.groupId(), groupEpoch, consumerGroupAssignors.get(preferredServerAssignor))
                     .withMembers(group.members())
                     .withStaticMembers(group.staticMembers())
-                    .withSubscriptionMetadata(subscriptionMetadata)
+                    .withSubscriptionTopicIdSet(subscriptionMetadata.values().stream().map(TopicMetadata::id).collect(Collectors.toSet()))
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
user: login: dajac
body: I think that our goal is to get rid of `subscriptionMetadata` so it is a bit surprising to still rely on it here. Could we avoid it?
created_at: 2025-05-09T15:14:43Z
author_association: MEMBER
diff_hunk: @@ -19,41 +19,58 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.metadata.PartitionRegistration;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain
  * topic and partition metadata for the topics that the modern group is subscribed to.
  */
 public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber {
     /**
-     * The topic Ids mapped to their corresponding {@link TopicMetadata}
-     * object, which contains topic and partition metadata.
+     * The topic Ids set.
      */
-    private final Map<Uuid, TopicMetadata> topicMetadata;
+    private final Set<Uuid> subscriptionTopicIdSet;
     private final Map<Uuid, Set<Integer>> topicPartitionAllowedMap;
 
-    public SubscribedTopicDescriberImpl(Map<Uuid, TopicMetadata> topicMetadata) {
-        this(topicMetadata, null);
+    /**
+     * The metadata image that contains the latest metadata information.
+     */
+    private final MetadataImage metadataImage;
+
+    public SubscribedTopicDescriberImpl(Set<Uuid> subscriptionTopicIdSet, MetadataImage metadataImage) {
+        this(subscriptionTopicIdSet, metadataImage, null);
     }
 
-    public SubscribedTopicDescriberImpl(Map<Uuid, TopicMetadata> topicMetadata, Map<Uuid, Set<Integer>> topicPartitionAllowedMap) {
-        this.topicMetadata = Objects.requireNonNull(topicMetadata);
+    public SubscribedTopicDescriberImpl(Set<Uuid> subscriptionTopicIdSet, MetadataImage metadataImage, Map<Uuid, Set<Integer>> topicPartitionAllowedMap) {
+        this.subscriptionTopicIdSet = Objects.requireNonNull(subscriptionTopicIdSet);
+        this.metadataImage = Objects.requireNonNull(metadataImage);
         this.topicPartitionAllowedMap = topicPartitionAllowedMap;
     }
 
     /**
-     * Map of topic Ids to topic metadata.
+     * Set of topic Ids.
+     *
+     * @return The set of topic Ids.
+     */
+    public Set<Uuid> subscriptionTopicIdSet() {
+        return this.subscriptionTopicIdSet;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: dajac
body: Given how this is used (mainly in tests), I think that we should rather try to remove it. Would it be possible?
created_at: 2025-05-09T15:17:50Z
author_association: MEMBER
diff_hunk: @@ -19,41 +19,42 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.metadata.PartitionRegistration;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain
  * topic and partition metadata for the topics that the modern group is subscribed to.
  */
 public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber {
     /**
-     * The topic Ids mapped to their corresponding {@link TopicMetadata}
-     * object, which contains topic and partition metadata.
+     * The map of topic Ids to the set of allowed partitions for each topic.
+     * If this is empty, all partitions are allowed.
      */
-    private final Map<Uuid, TopicMetadata> topicMetadata;
-    private final Map<Uuid, Set<Integer>> topicPartitionAllowedMap;
+    private final Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: FrankYang0529
body: @AndrewJSchofield @smjn I changed this to `Optional`, because `topicPartitionAllowedMap` cannot be `null` in production code. The `TargetAssignmentBuilder` set default value as empty map [0] and always input the value to `SubscribedTopicDescriberImpl` [1]. We still need a value to indicate all partitions are assignable. I think it's better to use `Optional.empty`, so I change the default value in `TargetAssignmentBuilder` to `Optional.empty`. Developers still can use `withTopicAssignablePartitionsMap` to set a map value, but the value cannot be null, because it will be used in `Optional.of`. WDYT?

[0]
https://github.com/apache/kafka/blob/8fc41f9ca246616a75c888696a7dff4fa4db4b7f/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java#L290-L293

[1]
https://github.com/apache/kafka/blob/8fc41f9ca246616a75c888696a7dff4fa4db4b7f/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java#L491-L499
created_at: 2025-05-16T06:19:40Z
author_association: MEMBER
diff_hunk: @@ -19,41 +19,42 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.metadata.PartitionRegistration;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain
  * topic and partition metadata for the topics that the modern group is subscribed to.
  */
 public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber {
     /**
-     * The topic Ids mapped to their corresponding {@link TopicMetadata}
-     * object, which contains topic and partition metadata.
+     * The map of topic Ids to the set of allowed partitions for each topic.
+     * If this is empty, all partitions are allowed.
      */
-    private final Map<Uuid, TopicMetadata> topicMetadata;
-    private final Map<Uuid, Set<Integer>> topicPartitionAllowedMap;
+    private final Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: FrankYang0529
body: BTW, with `Optional` value, it's easier to include `topicAssignablePartitionsMap` in `equals`, `hashCode`, and `toString` functions.
created_at: 2025-05-16T06:20:54Z
author_association: MEMBER
diff_hunk: @@ -19,41 +19,42 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.metadata.PartitionRegistration;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain
  * topic and partition metadata for the topics that the modern group is subscribed to.
  */
 public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber {
     /**
-     * The topic Ids mapped to their corresponding {@link TopicMetadata}
-     * object, which contains topic and partition metadata.
+     * The map of topic Ids to the set of allowed partitions for each topic.
+     * If this is empty, all partitions are allowed.
      */
-    private final Map<Uuid, TopicMetadata> topicMetadata;
-    private final Map<Uuid, Set<Integer>> topicPartitionAllowedMap;
+    private final Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: AndrewJSchofield
body: There is a related PR in this area https://github.com/apache/kafka/pull/19739. It's correcting a mistake.
created_at: 2025-05-16T17:26:51Z
author_association: MEMBER
diff_hunk: @@ -19,41 +19,42 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.metadata.PartitionRegistration;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain
  * topic and partition metadata for the topics that the modern group is subscribed to.
  */
 public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber {
     /**
-     * The topic Ids mapped to their corresponding {@link TopicMetadata}
-     * object, which contains topic and partition metadata.
+     * The map of topic Ids to the set of allowed partitions for each topic.
+     * If this is empty, all partitions are allowed.
      */
-    private final Map<Uuid, TopicMetadata> topicMetadata;
-    private final Map<Uuid, Set<Integer>> topicPartitionAllowedMap;
+    private final Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap;
 
-    public SubscribedTopicDescriberImpl(Map<Uuid, TopicMetadata> topicMetadata) {
-        this(topicMetadata, null);
-    }
+    /**
+     * The metadata image that contains the latest metadata information.
+     */
+    private final MetadataImage metadataImage;
 
-    public SubscribedTopicDescriberImpl(Map<Uuid, TopicMetadata> topicMetadata, Map<Uuid, Set<Integer>> topicPartitionAllowedMap) {
-        this.topicMetadata = Objects.requireNonNull(topicMetadata);
-        this.topicPartitionAllowedMap = topicPartitionAllowedMap;
+    public SubscribedTopicDescriberImpl(MetadataImage metadataImage) {
+        this(metadataImage, Optional.empty());
     }
 
-    /**
-     * Map of topic Ids to topic metadata.
-     *
-     * @return The map of topic Ids to topic metadata.
-     */
-    public Map<Uuid, TopicMetadata> topicMetadata() {
-        return this.topicMetadata;
+    public SubscribedTopicDescriberImpl(
+        MetadataImage metadataImage,
+        Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap
+    ) {
+        this.metadataImage = Objects.requireNonNull(metadataImage);
+        this.topicPartitionAllowedMap = topicPartitionAllowedMap;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: dajac
body: nit: Should we requireNonNull too?
created_at: 2025-05-17T16:41:22Z
author_association: MEMBER
diff_hunk: @@ -79,49 +80,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
+                }
+                return racks;
+            }
+        }
         return Set.of();
     }
 
     /**
-     * Returns a set of assignable partitions from the topic metadata.
-     * If the allowed partition map is null, all the partitions in the corresponding
-     * topic metadata are returned for the argument topic id. If allowed map is empty,
+     * Returns a set of assignable partitions from the metadata image.
+     * If the allowed partition map is Optional.empty(), all the partitions in the corresponding
+     * topic image are returned for the argument topic id. If allowed map is empty,
      * empty set is returned.
      *
      * @param topicId The uuid of the topic
      * @return Set of integers if assignable partitions available, empty otherwise.
      */
     @Override
     public Set<Integer> assignablePartitions(Uuid topicId) {
-        TopicMetadata topic = this.topicMetadata.get(topicId);
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
         if (topic == null) {
             return Set.of();
         }
 
-        if (topicPartitionAllowedMap == null) {
-            return IntStream.range(0, topic.numPartitions()).boxed().collect(Collectors.toUnmodifiableSet());
+        if (topicPartitionAllowedMap.isEmpty()) {
+            return Set.copyOf(topic.partitions().keySet());
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: dajac
body: Do we really need to make a copy of the Set? The image is immutable so we could just return a non modifiable version. What do you think?
created_at: 2025-05-17T16:42:47Z
author_association: MEMBER
diff_hunk: @@ -61,15 +63,12 @@ public class OptimizedUniformAssignmentBuilderTest {
 
     @Test
     public void testOneMemberNoTopicSubscription() {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addRacks()
path: group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java
user: login: dajac
body: nit: Is `addRacks` necessary here and in all the other places?
created_at: 2025-05-17T16:44:18Z
author_association: MEMBER
diff_hunk: @@ -245,22 +237,21 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
 
     @Test
     public void testValidityAndBalanceForLargeSampleSet() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        for (int i = 1; i < 100; i++) {
+        MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder();
+        Set<Uuid> topicIds = new HashSet<>();
+        IntStream.range(1, 101).forEach(i -> {
path: group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java
user: login: dajac
body: Is 1 to 101 correct? It may be easier to keep the for loop.
created_at: 2025-05-17T16:45:33Z
author_association: MEMBER
diff_hunk: @@ -17,94 +17,111 @@
 package org.apache.kafka.coordinator.group.modern;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.coordinator.group.MetadataImageBuilder;
+import org.apache.kafka.image.MetadataImage;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.stream.IntStream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class SubscribedTopicMetadataTest {
 
-    private Map<Uuid, TopicMetadata> topicMetadataMap;
     private SubscribedTopicDescriberImpl subscribedTopicMetadata;
+    private MetadataImage metadataImage;
+    private final int numPartitions = 5;
 
     @BeforeEach
     public void setUp() {
-        topicMetadataMap = new HashMap<>();
-        for (int i = 0; i < 5; i++) {
+        MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder();
+        IntStream.range(0, 5).forEach(i -> {
path: group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java
user: login: dajac
body: nit: Let's keep the for loop here as well.
created_at: 2025-05-17T16:48:14Z
author_association: MEMBER
diff_hunk: @@ -19,41 +19,42 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.metadata.PartitionRegistration;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain
  * topic and partition metadata for the topics that the modern group is subscribed to.
  */
 public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber {
     /**
-     * The topic Ids mapped to their corresponding {@link TopicMetadata}
-     * object, which contains topic and partition metadata.
+     * The map of topic Ids to the set of allowed partitions for each topic.
+     * If this is empty, all partitions are allowed.
      */
-    private final Map<Uuid, TopicMetadata> topicMetadata;
-    private final Map<Uuid, Set<Integer>> topicPartitionAllowedMap;
+    private final Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap;
 
-    public SubscribedTopicDescriberImpl(Map<Uuid, TopicMetadata> topicMetadata) {
-        this(topicMetadata, null);
-    }
+    /**
+     * The metadata image that contains the latest metadata information.
+     */
+    private final MetadataImage metadataImage;
 
-    public SubscribedTopicDescriberImpl(Map<Uuid, TopicMetadata> topicMetadata, Map<Uuid, Set<Integer>> topicPartitionAllowedMap) {
-        this.topicMetadata = Objects.requireNonNull(topicMetadata);
-        this.topicPartitionAllowedMap = topicPartitionAllowedMap;
+    public SubscribedTopicDescriberImpl(MetadataImage metadataImage) {
+        this(metadataImage, Optional.empty());
     }
 
-    /**
-     * Map of topic Ids to topic metadata.
-     *
-     * @return The map of topic Ids to topic metadata.
-     */
-    public Map<Uuid, TopicMetadata> topicMetadata() {
-        return this.topicMetadata;
+    public SubscribedTopicDescriberImpl(
+        MetadataImage metadataImage,
+        Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap
+    ) {
+        this.metadataImage = Objects.requireNonNull(metadataImage);
+        this.topicPartitionAllowedMap = topicPartitionAllowedMap;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: FrankYang0529
body: Yes, updated it. Thanks.
created_at: 2025-05-18T05:31:38Z
author_association: MEMBER
diff_hunk: @@ -61,15 +63,12 @@ public class OptimizedUniformAssignmentBuilderTest {
 
     @Test
     public void testOneMemberNoTopicSubscription() {
+        MetadataImage metadataImage = new MetadataImageBuilder()
+            .addTopic(topic1Uuid, topic1Name, 3)
+            .addRacks()
path: group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java
user: login: FrankYang0529
body: Currently, assignor doesn't use `racksForPartition`, so it's not necessary. Remove it from assignor test cases.
created_at: 2025-05-18T05:32:34Z
author_association: MEMBER
diff_hunk: @@ -245,22 +237,21 @@ public void testFirstAssignmentNumMembersGreaterThanTotalNumPartitions() {
 
     @Test
     public void testValidityAndBalanceForLargeSampleSet() {
-        Map<Uuid, TopicMetadata> topicMetadata = new HashMap<>();
-        for (int i = 1; i < 100; i++) {
+        MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder();
+        Set<Uuid> topicIds = new HashSet<>();
+        IntStream.range(1, 101).forEach(i -> {
path: group-coordinator/src/test/java/org/apache/kafka/coordinator/group/assignor/OptimizedUniformAssignmentBuilderTest.java
user: login: FrankYang0529
body: Updated it.
created_at: 2025-05-18T05:33:26Z
author_association: MEMBER
diff_hunk: @@ -17,94 +17,111 @@
 package org.apache.kafka.coordinator.group.modern;
 
 import org.apache.kafka.common.Uuid;
+import org.apache.kafka.coordinator.group.MetadataImageBuilder;
+import org.apache.kafka.image.MetadataImage;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.stream.IntStream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class SubscribedTopicMetadataTest {
 
-    private Map<Uuid, TopicMetadata> topicMetadataMap;
     private SubscribedTopicDescriberImpl subscribedTopicMetadata;
+    private MetadataImage metadataImage;
+    private final int numPartitions = 5;
 
     @BeforeEach
     public void setUp() {
-        topicMetadataMap = new HashMap<>();
-        for (int i = 0; i < 5; i++) {
+        MetadataImageBuilder metadataImageBuilder = new MetadataImageBuilder();
+        IntStream.range(0, 5).forEach(i -> {
path: group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicMetadataTest.java
user: login: FrankYang0529
body: Updated it.
created_at: 2025-05-18T05:33:33Z
author_association: MEMBER
diff_hunk: @@ -79,49 +80,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
+                }
+                return racks;
+            }
+        }
         return Set.of();
     }
 
     /**
-     * Returns a set of assignable partitions from the topic metadata.
-     * If the allowed partition map is null, all the partitions in the corresponding
-     * topic metadata are returned for the argument topic id. If allowed map is empty,
+     * Returns a set of assignable partitions from the metadata image.
+     * If the allowed partition map is Optional.empty(), all the partitions in the corresponding
+     * topic image are returned for the argument topic id. If allowed map is empty,
      * empty set is returned.
      *
      * @param topicId The uuid of the topic
      * @return Set of integers if assignable partitions available, empty otherwise.
      */
     @Override
     public Set<Integer> assignablePartitions(Uuid topicId) {
-        TopicMetadata topic = this.topicMetadata.get(topicId);
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
         if (topic == null) {
             return Set.of();
         }
 
-        if (topicPartitionAllowedMap == null) {
-            return IntStream.range(0, topic.numPartitions()).boxed().collect(Collectors.toUnmodifiableSet());
+        if (topicPartitionAllowedMap.isEmpty()) {
+            return Set.copyOf(topic.partitions().keySet());
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: FrankYang0529
body: IIUC, the production code uses `HashMap` to build partitions in `TopicImage`. From [Map#keySet Java doc](https://devdocs.io/openjdk~17/java.base/java/util/map#keySet()), the change to the set can be reflected in the map. Is it safer to use `Set.copyOf` here?

https://github.com/apache/kafka/blob/f26974b16db7903bd350f06b5796acac52604d4c/metadata/src/main/java/org/apache/kafka/image/TopicDelta.java#L142-L157
created_at: 2025-05-18T05:38:30Z
author_association: MEMBER
diff_hunk: @@ -79,49 +80,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
+                }
+                return racks;
+            }
+        }
         return Set.of();
     }
 
     /**
-     * Returns a set of assignable partitions from the topic metadata.
-     * If the allowed partition map is null, all the partitions in the corresponding
-     * topic metadata are returned for the argument topic id. If allowed map is empty,
+     * Returns a set of assignable partitions from the metadata image.
+     * If the allowed partition map is Optional.empty(), all the partitions in the corresponding
+     * topic image are returned for the argument topic id. If allowed map is empty,
      * empty set is returned.
      *
      * @param topicId The uuid of the topic
      * @return Set of integers if assignable partitions available, empty otherwise.
      */
     @Override
     public Set<Integer> assignablePartitions(Uuid topicId) {
-        TopicMetadata topic = this.topicMetadata.get(topicId);
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
         if (topic == null) {
             return Set.of();
         }
 
-        if (topicPartitionAllowedMap == null) {
-            return IntStream.range(0, topic.numPartitions()).boxed().collect(Collectors.toUnmodifiableSet());
+        if (topicPartitionAllowedMap.isEmpty()) {
+            return Set.copyOf(topic.partitions().keySet());
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: dajac
body: Yeah, this is correct. However, the TopicImage is never mutated once built. The metadata image is immutable.
created_at: 2025-05-18T06:54:01Z
author_association: MEMBER
diff_hunk: @@ -79,49 +80,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
+                }
+                return racks;
+            }
+        }
         return Set.of();
     }
 
     /**
-     * Returns a set of assignable partitions from the topic metadata.
-     * If the allowed partition map is null, all the partitions in the corresponding
-     * topic metadata are returned for the argument topic id. If allowed map is empty,
+     * Returns a set of assignable partitions from the metadata image.
+     * If the allowed partition map is Optional.empty(), all the partitions in the corresponding
+     * topic image are returned for the argument topic id. If allowed map is empty,
      * empty set is returned.
      *
      * @param topicId The uuid of the topic
      * @return Set of integers if assignable partitions available, empty otherwise.
      */
     @Override
     public Set<Integer> assignablePartitions(Uuid topicId) {
-        TopicMetadata topic = this.topicMetadata.get(topicId);
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
         if (topic == null) {
             return Set.of();
         }
 
-        if (topicPartitionAllowedMap == null) {
-            return IntStream.range(0, topic.numPartitions()).boxed().collect(Collectors.toUnmodifiableSet());
+        if (topicPartitionAllowedMap.isEmpty()) {
+            return Set.copyOf(topic.partitions().keySet());
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: chia7712
body: As @dajac mentioned, maybe we can make `TopicImage` create immutable collection as it is never modified once built.
```java
    public TopicImage(String name,
                      Uuid id,
                      Map<Integer, PartitionRegistration> partitions) {
        this.name = name;
        this.id = id;
        this.partitions = Collections.unmodifiableMap(partitions);
    }
```

with this change, we can return `topic.partitions().keySet()` directly as it is immutable now.

BTW, I have opened https://issues.apache.org/jira/browse/KAFKA-19305 to fix other image classes too
created_at: 2025-05-18T16:30:18Z
author_association: MEMBER
diff_hunk: @@ -79,49 +80,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
+                }
+                return racks;
+            }
+        }
         return Set.of();
     }
 
     /**
-     * Returns a set of assignable partitions from the topic metadata.
-     * If the allowed partition map is null, all the partitions in the corresponding
-     * topic metadata are returned for the argument topic id. If allowed map is empty,
+     * Returns a set of assignable partitions from the metadata image.
+     * If the allowed partition map is Optional.empty(), all the partitions in the corresponding
+     * topic image are returned for the argument topic id. If allowed map is empty,
      * empty set is returned.
      *
      * @param topicId The uuid of the topic
      * @return Set of integers if assignable partitions available, empty otherwise.
      */
     @Override
     public Set<Integer> assignablePartitions(Uuid topicId) {
-        TopicMetadata topic = this.topicMetadata.get(topicId);
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
         if (topic == null) {
             return Set.of();
         }
 
-        if (topicPartitionAllowedMap == null) {
-            return IntStream.range(0, topic.numPartitions()).boxed().collect(Collectors.toUnmodifiableSet());
+        if (topicPartitionAllowedMap.isEmpty()) {
+            return Set.copyOf(topic.partitions().keySet());
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: FrankYang0529
body: Thanks for the review and suggestion. Remove `Set.copyOf` here.
created_at: 2025-05-19T03:43:48Z
author_association: MEMBER
diff_hunk: @@ -112,7 +112,7 @@ public Set<Integer> assignablePartitions(Uuid topicId) {
         }
 
         if (topicPartitionAllowedMap.isEmpty()) {
-            return Set.copyOf(topic.partitions().keySet());
+            return topic.partitions().keySet();
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: dajac
body: nit: I would use Collections.unmodifiableSet to ensure that it cannot be mutated by the caller.
created_at: 2025-05-19T06:59:40Z
author_association: MEMBER
diff_hunk: @@ -112,7 +112,7 @@ public Set<Integer> assignablePartitions(Uuid topicId) {
         }
 
         if (topicPartitionAllowedMap.isEmpty()) {
-            return Set.copyOf(topic.partitions().keySet());
+            return topic.partitions().keySet();
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: FrankYang0529
body: Updated it. Thank you.
created_at: 2025-05-19T07:08:11Z
author_association: MEMBER
diff_hunk: @@ -112,7 +112,7 @@ public Set<Integer> assignablePartitions(Uuid topicId) {
         }
 
         if (topicPartitionAllowedMap.isEmpty()) {
-            return Set.copyOf(topic.partitions().keySet());
+            return topic.partitions().keySet();
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: chia7712
body: @dajac is right, we can remove the wrap in addressing KAFKA-19305 
created_at: 2025-05-19T07:31:22Z
author_association: MEMBER
diff_hunk: @@ -79,49 +81,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: chia7712
body: Should we check the offline replica?
created_at: 2025-05-19T08:29:17Z
author_association: MEMBER
diff_hunk: @@ -79,49 +81,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
+                }
+                return racks;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: chia7712
body: Should we return immutable set?
created_at: 2025-05-19T08:29:31Z
author_association: MEMBER
diff_hunk: @@ -19,41 +19,42 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.metadata.PartitionRegistration;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain
  * topic and partition metadata for the topics that the modern group is subscribed to.
  */
 public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber {
     /**
-     * The topic Ids mapped to their corresponding {@link TopicMetadata}
-     * object, which contains topic and partition metadata.
+     * The map of topic Ids to the set of allowed partitions for each topic.
+     * If this is empty, all partitions are allowed.
      */
-    private final Map<Uuid, TopicMetadata> topicMetadata;
-    private final Map<Uuid, Set<Integer>> topicPartitionAllowedMap;
+    private final Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: chia7712
body: #19739 uses null, in contrast, this PR uses `Optional`. @AndrewJSchofield WDYT?
created_at: 2025-05-19T08:31:29Z
author_association: MEMBER
diff_hunk: @@ -19,41 +19,42 @@
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.coordinator.group.api.assignor.PartitionAssignor;
 import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
+import org.apache.kafka.image.MetadataImage;
+import org.apache.kafka.image.TopicImage;
+import org.apache.kafka.metadata.PartitionRegistration;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * The subscribed topic metadata class is used by the {@link PartitionAssignor} to obtain
  * topic and partition metadata for the topics that the modern group is subscribed to.
  */
 public class SubscribedTopicDescriberImpl implements SubscribedTopicDescriber {
     /**
-     * The topic Ids mapped to their corresponding {@link TopicMetadata}
-     * object, which contains topic and partition metadata.
+     * The map of topic Ids to the set of allowed partitions for each topic.
+     * If this is empty, all partitions are allowed.
      */
-    private final Map<Uuid, TopicMetadata> topicMetadata;
-    private final Map<Uuid, Set<Integer>> topicPartitionAllowedMap;
+    private final Optional<Map<Uuid, Set<Integer>>> topicPartitionAllowedMap;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: dajac
body: Using an Optional is better in my opinion.
created_at: 2025-05-19T09:19:00Z
author_association: MEMBER
diff_hunk: @@ -79,49 +81,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: dajac
body: I was assuming that we should not. If I recall correctly, this is what we did in KIP-881. I need to double check to be sure.
created_at: 2025-05-19T09:20:53Z
author_association: MEMBER
diff_hunk: @@ -79,49 +81,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
+                }
+                return racks;
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: dajac
body: Good call.
created_at: 2025-05-19T09:21:10Z
author_association: MEMBER
diff_hunk: @@ -79,49 +81,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: dajac
body: In KIP-881, we don't check whether replicas are online or offline. We basically assume that they will eventually come back. I would follow the same approach here.

Side note, this needs to be aligned with the computation of the hash. There, we include all replicas regardless of whether they are online or not.
created_at: 2025-05-19T09:25:44Z
author_association: MEMBER
diff_hunk: @@ -79,49 +81,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: chia7712
body: If the offline replicas could be existent in this path, we need to add null check for `broker`, right?
created_at: 2025-05-19T09:32:24Z
author_association: MEMBER
diff_hunk: @@ -79,49 +81,65 @@ public int numPartitions(Uuid topicId) {
      */
     @Override
     public Set<String> racksForPartition(Uuid topicId, int partition) {
+        TopicImage topic = metadataImage.topics().getTopic(topicId);
+        if (topic != null) {
+            PartitionRegistration partitionRegistration = topic.partitions().get(partition);
+            if (partitionRegistration != null) {
+                Set<String> racks = new HashSet<>();
+                for (int replica : partitionRegistration.replicas) {
+                    // Only add the rack if it is available for the broker/replica.
+                    metadataImage.cluster().broker(replica).rack().ifPresent(racks::add);
path: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/SubscribedTopicDescriberImpl.java
user: login: FrankYang0529
body: We do null check in hash function. It's better to add the check here.

https://github.com/apache/kafka/blob/ce4940f9891a96819e54f8db097ce3824876e8e5/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/Utils.java#L408-L413
created_at: 2025-05-19T09:35:54Z
author_association: MEMBER

PR comments:
user: login: FrankYang0529
created_at: 2025-05-13T11:42:34Z
body: > Out of curiosity, have you tried to run the benchmarks before and after the change to see the impact?

Since previous `SubscribedTopicDescriberImpl` only supports `numOfPartitions`, I only test this function. The result is 10^(-5) vs 10^(-6).

### This PR

```
Benchmark                                                  (partitionsPerTopic)  (topicCount)  Mode  Cnt   Score    Error  Units
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                     5            10  avgt    5  ≈ 10⁻⁵           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                     5           100  avgt    5  ≈ 10⁻⁵           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                     5          1000  avgt    5  ≈ 10⁻⁵           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    10            10  avgt    5  ≈ 10⁻⁵           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    10           100  avgt    5  ≈ 10⁻⁵           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    10          1000  avgt    5  ≈ 10⁻⁵           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    50            10  avgt    5  ≈ 10⁻⁵           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    50           100  avgt    5  ≈ 10⁻⁵           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    50          1000  avgt    5  ≈ 10⁻⁵           ms/op
JMH benchmarks done
```

<details>
<summary> SubscribedTopicDescriberImplBenchmark.java
</summary>

```java
/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.jmh.assignor;

import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
import org.apache.kafka.coordinator.group.modern.TopicIds;
import org.apache.kafka.image.MetadataImage;

import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.Level;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;

import java.util.List;
import java.util.Set;
import java.util.concurrent.TimeUnit;

@State(Scope.Benchmark)
@Fork(value = 1)
@Warmup(iterations = 3)
@Measurement(iterations = 5)
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
public class SubscribedTopicDescriberImplBenchmark {
    @Param({"5", "10", "50"})
    private int partitionsPerTopic;

    @Param({"10", "100", "1000"})
    private int topicCount;

    private List<String> allTopicNames = List.of();

    private Set<Uuid> subscriptionTopicIdSet;

    private MetadataImage metadataImage;

    private TopicIds.TopicResolver topicResolver;

    private SubscribedTopicDescriber subscribedTopicDescriber;

    private String topicName;
    private Uuid topicId;

    private void setupTopics() {
        allTopicNames = AssignorBenchmarkUtils.createTopicNames(topicCount);
        topicName = allTopicNames.get(0);

        metadataImage = AssignorBenchmarkUtils.createMetadataImage(allTopicNames, partitionsPerTopic);
        topicResolver = new TopicIds.CachedTopicResolver(metadataImage.topics());
        topicId = topicResolver.id(topicName);

        subscriptionTopicIdSet = AssignorBenchmarkUtils.subscriptionTopicIdSet(metadataImage);
        subscribedTopicDescriber = new SubscribedTopicDescriberImpl(subscriptionTopicIdSet, metadataImage);
    }

    @Setup(Level.Trial)
    public void setup() {
        setupTopics();
    }

    @Benchmark
    public void testNumOfPartitions() {
        subscribedTopicDescriber.numPartitions(topicId);
    }
}
```
</details>

### trunk

```
Benchmark                                                  (partitionsPerTopic)  (topicCount)  Mode  Cnt   Score    Error  Units
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                     5            10  avgt    5  ≈ 10⁻⁶           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                     5           100  avgt    5  ≈ 10⁻⁶           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                     5          1000  avgt    5  ≈ 10⁻⁶           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    10            10  avgt    5  ≈ 10⁻⁶           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    10           100  avgt    5  ≈ 10⁻⁶           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    10          1000  avgt    5  ≈ 10⁻⁶           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    50            10  avgt    5  ≈ 10⁻⁵           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    50           100  avgt    5  ≈ 10⁻⁶           ms/op
SubscribedTopicDescriberImplBenchmark.testNumOfPartitions                    50          1000  avgt    5  ≈ 10⁻⁶           ms/op
JMH benchmarks done
```

<details>
<summary> SubscribedTopicDescriberImplBenchmark.java
</summary>

```java
/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.jmh.assignor;

import org.apache.kafka.common.Uuid;
import org.apache.kafka.coordinator.group.api.assignor.SubscribedTopicDescriber;
import org.apache.kafka.coordinator.group.modern.SubscribedTopicDescriberImpl;
import org.apache.kafka.coordinator.group.modern.TopicIds;
import org.apache.kafka.coordinator.group.modern.TopicMetadata;
import org.apache.kafka.image.MetadataImage;

import org.apache.kafka.image.TopicsImage;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.Level;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;

import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

@State(Scope.Benchmark)
@Fork(value = 1)
@Warmup(iterations = 3)
@Measurement(iterations = 5)
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
public class SubscribedTopicDescriberImplBenchmark {
    @Param({"5", "10", "50"})
    private int partitionsPerTopic;

    @Param({"10", "100", "1000"})
    private int topicCount;

    private List<String> allTopicNames = List.of();

    private Map<String, TopicMetadata> subscriptionMetadata;

    private TopicsImage topicsImage;

    private TopicIds.TopicResolver topicResolver;

    private SubscribedTopicDescriber subscribedTopicDescriber;

    private String topicName;
    private Uuid topicId;

    private void setupTopics() {
        allTopicNames = AssignorBenchmarkUtils.createTopicNames(topicCount);
        topicName = allTopicNames.get(0);

        subscriptionMetadata = AssignorBenchmarkUtils.createSubscriptionMetadata(allTopicNames, partitionsPerTopic);
        topicsImage = AssignorBenchmarkUtils.createTopicsImage(subscriptionMetadata);
        topicResolver = new TopicIds.CachedTopicResolver(topicsImage);
        topicId = topicResolver.id(topicName);

        Map<Uuid, TopicMetadata> topicMetadata = AssignorBenchmarkUtils.createTopicMetadata(subscriptionMetadata);
        subscribedTopicDescriber = new SubscribedTopicDescriberImpl(topicMetadata);
    }

    @Setup(Level.Trial)
    public void setup() {
        setupTopics();
    }

    @Benchmark
    public void testNumOfPartitions() {
        subscribedTopicDescriber.numPartitions(topicId);
    }
}
```
</details>
author_association: MEMBER
user: login: dajac
created_at: 2025-05-13T12:38:26Z
body: @FrankYang0529 Thanks for the results. Have you also tried to run `ServerSideAssignorBenchmark`?
author_association: MEMBER
user: login: FrankYang0529
created_at: 2025-05-14T13:13:14Z
body: > Have you also tried to run `ServerSideAssignorBenchmark`?

Run this benchmark takes 4 hours. I will paste result tomorrow.
author_association: MEMBER
user: login: FrankYang0529
created_at: 2025-05-15T04:28:37Z
body: The benchmark results are similar.

```
./jmh-benchmarks/jmh.sh org.apache.kafka.jmh.assignor.ServerSideAssignorBenchmark -f 1 -i 1 -wi 1
```

[pr-result.txt](https://github.com/user-attachments/files/20219606/pr-result.txt)
[trunk-result.txt](https://github.com/user-attachments/files/20219608/trunk-result.txt)

author_association: MEMBER
user: login: dajac
created_at: 2025-05-15T08:44:24Z
body: @FrankYang0529 I had another look at your PR and I am not fully satisfied by the end state. I won't like the fact that we introduce `subscriptionTopicIdSet`. It does not seem necessary to me. I wonder whether we could also try to actually remove all the usages of `TopicMetadata` and replace them by `ImageMetadata`. Ideally, `TopicMetadata` should only be used to trigger a rebalance after this PR but it should not be used for anything else. Would it be possible?
author_association: MEMBER
user: login: FrankYang0529
created_at: 2025-05-15T15:35:11Z
body: @dajac Thanks for the suggestion. I agree we can remove `subscriptionTopicIdSet` because `PartitionAssignor` can return subscribed topic UUID in `MemberSubscription`. However, if we remove `subscriptionTopicIdSet`, the implementation of `SubscribedTopicDescriber` doesn't match the naming, because the `SubscribedTopicDescriber` can return all data in `MetadataImage`. How about we also change `SubscribedTopicDescriber` to `TopicDescriber`?
author_association: MEMBER
user: login: dajac
created_at: 2025-05-15T18:30:23Z
body: > @dajac Thanks for the suggestion. I agree we can remove `subscriptionTopicIdSet` because `PartitionAssignor` can return subscribed topic UUID in `MemberSubscription`. However, if we remove `subscriptionTopicIdSet`, the implementation of `SubscribedTopicDescriber` doesn't match the naming, because the `SubscribedTopicDescriber` can return all data in `MetadataImage`. How about we also change `SubscribedTopicDescriber` to `TopicDescriber`?

Yeah, we are indeed diverging a bit from the name but I don't think that it is an issue. At the end of the day, it will be used based on the subscribed topics anyway. Regarding changing the name, `SubscribedTopicDescriber` is actually part of our public API so we cannot change it.
author_association: MEMBER
user: login: FrankYang0529
created_at: 2025-05-16T09:34:17Z
body: @dajac Thanks for the review. I remove `subscriptionTopicIdSet` and `TopicMetadata` in `SubscribedTopicDescriberImpl`.
author_association: MEMBER
user: login: AndrewJSchofield
created_at: 2025-05-22T12:46:07Z
body: @FrankYang0529 I wonder whether this PR could be the cause of the recent flakiness in `PlaintextAdminIntegrationTest.testDeleteConsumerGroupOffsets`. The test checks that the correct error code is returned if an attempt is made to delete a consumer group offset when there is a subscription to the topic whose offsets are being deleted.

https://develocity.apache.org/scans/tests?search.relativeStartTime=P28D&search.rootProjectNames=kafka&search.timeZoneId=Europe%2FLondon&tests.container=kafka.api.PlaintextAdminIntegrationTest&tests.test=testDeleteConsumerGroupOffsets(String)%5B1%5D
author_association: MEMBER
user: login: chia7712
created_at: 2025-05-22T13:10:40Z
body: @AndrewJSchofield thanks for this report. we will check it ASAP
author_association: MEMBER
user: login: chia7712
created_at: 2025-05-22T16:26:06Z
body: @AndrewJSchofield the error is caused by #19380 which changed the partition order in request. I will file a PR to fix it
author_association: MEMBER