Add dynamic read support to KafkaIO in Python SDK (#35292)

* Add dynamic read support to KafkaIO in Python SDK

* removed the comment

Author: liferoad

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java

@@ -875,9 +875,13 @@ static <K, V> void setupExternalBuilder(
           builder.setStopReadTime(Instant.ofEpochMilli(config.stopReadTime));
         }
 
-        // We can expose dynamic read to external build when ReadFromKafkaDoFn is the default
-        // implementation.
-        builder.setDynamicRead(false);
+        if (config.dynamicReadPollIntervalSeconds != null) {
+          builder.setDynamicRead(true);
+          builder.setWatchTopicPartitionDuration(
+              Duration.standardSeconds(config.dynamicReadPollIntervalSeconds));
+        } else {
+          builder.setDynamicRead(false);
+        }
 
         if (config.consumerPollingTimeout != null) {
           if (config.consumerPollingTimeout <= 0) {
@@ -975,6 +979,7 @@ public static class Configuration {
         private Boolean redistribute;
         private Boolean allowDuplicates;
         private Boolean offsetDeduplication;
+        private Long dynamicReadPollIntervalSeconds;
 
         public void setConsumerConfig(Map<String, String> consumerConfig) {
           this.consumerConfig = consumerConfig;
@@ -1035,6 +1040,10 @@ public void setAllowDuplicates(Boolean allowDuplicates) {
         public void setOffsetDeduplication(Boolean offsetDeduplication) {
           this.offsetDeduplication = offsetDeduplication;
         }
+
+        public void setDynamicReadPollIntervalSeconds(Long dynamicReadPollIntervalSeconds) {
+          this.dynamicReadPollIntervalSeconds = dynamicReadPollIntervalSeconds;
+        }
       }
     }
 

sdks/python/apache_beam/io/kafka.py

@@ -124,6 +124,7 @@
         ('redistribute', typing.Optional[bool]),
         ('redistribute_num_keys', typing.Optional[np.int32]),
         ('allow_duplicates', typing.Optional[bool]),
+        ('dynamic_read_poll_interval_seconds', typing.Optional[int]),
     ])
 
 
@@ -171,6 +172,7 @@ def __init__(
       redistribute=False,
       redistribute_num_keys=np.int32(0),
       allow_duplicates=False,
+      dynamic_read_poll_interval_seconds: typing.Optional[int] = None,
   ):
     """
     Initializes a read operation from Kafka.
@@ -195,7 +197,7 @@ def __init__(
     :param timestamp_policy: The built-in timestamp policy which is used for
         extracting timestamp from KafkaRecord.
     :param consumer_polling_timeout: Kafka client polling request
-        timeout time in seconds. A lower timeout optimizes for latency. Increase                                   
+        timeout time in seconds. A lower timeout optimizes for latency. Increase
         the timeout if the consumer is not fetching any records. Default is 2
         seconds.
     :param with_metadata: whether the returned PCollection should contain
@@ -205,12 +207,15 @@ def __init__(
         this only works when using default key and value deserializers where
         Java Kafka Reader reads keys and values as 'byte[]'.
     :param expansion_service: The address (host:port) of the ExpansionService.
-    :param redistribute: whether a Redistribute transform should be applied 
+    :param redistribute: whether a Redistribute transform should be applied
         immediately after the read.
-    :param redistribute_num_keys: Configures how many keys the Redistribute 
+    :param redistribute_num_keys: Configures how many keys the Redistribute
         spreads the data across.
-    :param allow_duplicates: whether the Redistribute transform allows for 
+    :param allow_duplicates: whether the Redistribute transform allows for
         duplicates (this serves solely as a hint to the underlying runner).
+    :param dynamic_read_poll_interval_seconds: The interval in seconds at which
+        to check for new partitions. If not None, dynamic partition discovery
+        is enabled.
     """
     if timestamp_policy not in [ReadFromKafka.processing_time_policy,
                                 ReadFromKafka.create_time_policy,
@@ -235,7 +240,9 @@ def __init__(
                 consumer_polling_timeout=consumer_polling_timeout,
                 redistribute=redistribute,
                 redistribute_num_keys=redistribute_num_keys,
-                allow_duplicates=allow_duplicates)),
+                allow_duplicates=allow_duplicates,
+                dynamic_read_poll_interval_seconds=
+                dynamic_read_poll_interval_seconds)),
         expansion_service or default_io_expansion_service())