[FLINK-36209] Remove redundant operations in the initialization of KafkaSourceEnumState (#116)

In certain methods, such as the DynamicKafkaSourceEnumerator#onHandleSubscribedStreamsFetch() method, partitions are divided into assignedPartitions and unassignedInitialPartitions before being passed as parameters to the KafkaSourceEnumState constructor. However, within the constructor, these assignedPartitions and unassignedInitialPartitions are recombined into partitions, leading to unnecessary operations and reduced performance. By optimizing the code to pass partitions directly as a parameter when initializing KafkaSourceEnumState, we can eliminate redundant operations and enhance performance.

---------

Co-authored-by: ClownXC <[email protected]>

Author: xiaochen-zhou

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumerator.java

@@ -35,6 +35,7 @@
 import org.apache.flink.connector.kafka.source.KafkaPropertiesUtil;
 import org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumState;
 import org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumerator;
+import org.apache.flink.connector.kafka.source.enumerator.TopicPartitionAndAssignmentStatus;
 import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
 import org.apache.flink.connector.kafka.source.enumerator.subscriber.KafkaSubscriber;
 import org.apache.flink.connector.kafka.source.split.KafkaPartitionSplit;
@@ -298,24 +299,16 @@ private void onHandleSubscribedStreamsFetch(Set<KafkaStream> fetchedKafkaStreams
                 final Set<String> activeTopics = activeClusterTopics.getValue();
 
                 // filter out removed topics
-                Set<TopicPartition> activeAssignedPartitions =
-                        kafkaSourceEnumState.assignedPartitions().stream()
-                                .filter(tp -> activeTopics.contains(tp.topic()))
-                                .collect(Collectors.toSet());
-                Set<TopicPartition> activeUnassignedInitialPartitions =
-                        kafkaSourceEnumState.unassignedInitialPartitions().stream()
-                                .filter(tp -> activeTopics.contains(tp.topic()))
+                Set<TopicPartitionAndAssignmentStatus> partitions =
+                        kafkaSourceEnumState.partitions().stream()
+                                .filter(tp -> activeTopics.contains(tp.topicPartition().topic()))
                                 .collect(Collectors.toSet());
 
                 newKafkaSourceEnumState =
                         new KafkaSourceEnumState(
-                                activeAssignedPartitions,
-                                activeUnassignedInitialPartitions,
-                                kafkaSourceEnumState.initialDiscoveryFinished());
+                                partitions, kafkaSourceEnumState.initialDiscoveryFinished());
             } else {
-                newKafkaSourceEnumState =
-                        new KafkaSourceEnumState(
-                                Collections.emptySet(), Collections.emptySet(), false);
+                newKafkaSourceEnumState = new KafkaSourceEnumState(Collections.emptySet(), false);
             }
 
             // restarts enumerator from state using only the active topic partitions, to avoid

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumState.java

@@ -32,7 +32,8 @@ public class KafkaSourceEnumState {
     /** Partitions with status: ASSIGNED or UNASSIGNED_INITIAL. */
     private final Set<TopicPartitionAndAssignmentStatus> partitions;
     /**
-     * this flag will be marked as true if inital partitions are discovered after enumerator starts.
+     * this flag will be marked as true if initial partitions are discovered after enumerator
+     * starts.
      */
     private final boolean initialDiscoveryFinished;
 

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumStateSerializer.java

@@ -89,9 +89,7 @@ public KafkaSourceEnumState deserialize(int version, byte[] serialized) throws I
             case CURRENT_VERSION:
                 return deserializeTopicPartitionAndAssignmentStatus(serialized);
             case VERSION_1:
-                final Set<TopicPartition> assignedPartitions =
-                        deserializeTopicPartitions(serialized);
-                return new KafkaSourceEnumState(assignedPartitions, new HashSet<>(), true);
+                return deserializeAssignedTopicPartitions(serialized);
             case VERSION_0:
                 Map<Integer, Set<KafkaPartitionSplit>> currentPartitionAssignment =
                         SerdeUtils.deserializeSplitAssignments(
@@ -113,23 +111,24 @@ public KafkaSourceEnumState deserialize(int version, byte[] serialized) throws I
         }
     }
 
-    private static Set<TopicPartition> deserializeTopicPartitions(byte[] serializedTopicPartitions)
-            throws IOException {
+    private static KafkaSourceEnumState deserializeAssignedTopicPartitions(
+            byte[] serializedTopicPartitions) throws IOException {
         try (ByteArrayInputStream bais = new ByteArrayInputStream(serializedTopicPartitions);
                 DataInputStream in = new DataInputStream(bais)) {
 
             final int numPartitions = in.readInt();
-            Set<TopicPartition> topicPartitions = new HashSet<>(numPartitions);
+            Set<TopicPartitionAndAssignmentStatus> partitions = new HashSet<>(numPartitions);
             for (int i = 0; i < numPartitions; i++) {
                 final String topic = in.readUTF();
                 final int partition = in.readInt();
-                topicPartitions.add(new TopicPartition(topic, partition));
+                partitions.add(
+                        new TopicPartitionAndAssignmentStatus(
+                                new TopicPartition(topic, partition), AssignmentStatus.ASSIGNED));
             }
             if (in.available() > 0) {
                 throw new IOException("Unexpected trailing bytes in serialized topic partitions");
             }
-
-            return topicPartitions;
+            return new KafkaSourceEnumState(partitions, true);
         }
     }
 

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumerator.java

@@ -112,7 +112,7 @@ public KafkaSourceEnumerator(
                 properties,
                 context,
                 boundedness,
-                new KafkaSourceEnumState(Collections.emptySet(), Collections.emptySet(), false));
+                new KafkaSourceEnumState(Collections.emptySet(), false));
     }
 
     public KafkaSourceEnumerator(