Make npartitions optional & fix Kafka checkpointing with auto.offset.reset=latest which is now default (2). No support for adding partitions on the fly

Author: chinmaychandak

streamz/sources.py

@@ -459,7 +459,7 @@ def __init__(self, topic, consumer_params, poll_interval='1s',
         # Override the auto-commit config to enforce custom streamz checkpointing
         self.consumer_params['enable.auto.commit'] = 'false'
         if 'auto.offset.reset' not in self.consumer_params.keys():
-            consumer_params['auto.offset.reset'] = 'earliest'
+            consumer_params['auto.offset.reset'] = 'latest'
         self.topic = topic
         self.npartitions = npartitions
         self.poll_interval = convert_interval(poll_interval)
@@ -484,6 +484,11 @@ def checkpoint_emit(_part):
             ref = RefCounter(cb=lambda: commit(_part))
             yield self._emit(_part, metadata=[{'ref': ref}])
 
+        if self.npartitions is None:
+            kafka_cluster_metadata = self.consumer.list_topics(self.topic)
+            self.npartitions = len(kafka_cluster_metadata.topics[self.topic].partitions)
+        self.positions = [0] * self.npartitions
+
         tps = []
         for partition in range(self.npartitions):
             tps.append(ck.TopicPartition(self.topic, partition))
@@ -501,10 +506,6 @@ def checkpoint_emit(_part):
         try:
             while not self.stopped:
                 out = []
-                # kafka_cluster_metadata = self.consumer.list_topics(self.topic)
-                # self.npartitions = len(kafka_cluster_metadata.topics[self.topic].partitions)
-                # if self.npartitions > len(self.positions):
-                #     self.positions.extend([0] * (self.npartitions - len(self.positions)))
                 for partition in range(self.npartitions):
                     tp = ck.TopicPartition(self.topic, partition, 0)
                     try:
@@ -546,12 +547,6 @@ def start(self):
             else:
                 self.consumer = ck.Consumer(self.consumer_params)
             self.stopped = False
-
-            if self.npartitions is None:
-                kafka_cluster_metadata = self.consumer.list_topics(self.topic)
-                self.npartitions = len(kafka_cluster_metadata.topics[self.topic].partitions)
-            self.positions = [0] * self.npartitions
-
             tp = ck.TopicPartition(self.topic, 0, 0)
 
             # blocks for consumer thread to come up
@@ -597,6 +592,8 @@ def from_kafka_batched(topic, consumer_params, poll_interval='1s',
     npartitions: int (None)
         | Number of partitions in the topic.
         | If None, streamz will poll Kafka to get the number of partitions.
+        | As of now, streamz does not support changing number of partitions on the fly.
+        | It is recommended to restart the stream after changing the number of partitions.
     start: bool (False)
         Whether to start polling upon instantiation
     max_batch_size: int

streamz/tests/test_kafka.py

@@ -223,6 +223,7 @@ def test_kafka_batch():
         stream.upstream.stopped = True
 
 
+@flaky(max_runs=3, min_passes=1)
 @gen_cluster(client=True, timeout=60)
 def test_kafka_dask_batch(c, s, w1, w2):
     j = random.randint(0, 10000)
@@ -334,8 +335,7 @@ def test_kafka_batch_checkpointing_async_nodes_1():
     j = random.randint(0, 10000)
     ARGS = {'bootstrap.servers': 'localhost:9092',
             'group.id': 'streamz-test%i' % j,
-            'enable.auto.commit': False,
-            'auto.offset.reset': 'earliest'}
+            'enable.auto.commit': False}
     with kafka_service() as kafka:
         kafka, TOPIC = kafka
         stream1 = Stream.from_kafka_batched(TOPIC, ARGS)