Fix bug in source generation for input batch size > 1 (#75)

Resolves #74

Author: amitschang

dplutils/pipeline/stream.py

@@ -135,7 +135,9 @@ def process_source(self, source):
                 source_batch.append(next(self.source_generator))
             except StopIteration:
                 self.source_exhausted = True
-                return
+                if len(source_batch) == 0:
+                    return
+                break
             self.n_sourced += 1
             if self.n_sourced == self.max_batches:
                 self.source_exhausted = True

tests/pipeline/test_stream_executor.py

@@ -40,3 +40,20 @@ def generator():
     expected_rows = max_batches if max_batches else 12
     assert len(res) == expected_rows
     assert pd.concat(res).customgen.to_list() == list(range(expected_rows))
+
+
+@pytest.mark.parametrize('gen_n', [0, 4])
+def test_stream_executor_exhausts_input_when_source_batchsize_larger_than_input(gen_n):
+    st = PipelineTask('task_name', lambda x: x, batch_size=10)
+    def generator():
+        n = gen_n
+        for i in range(n):
+            yield pd.DataFrame({'customgen': [i]})
+    pl = LocalSerialExecutor([st], generator=generator)
+    res = [i.data for i in pl.run()]
+    if gen_n > 0:
+        assert len(res) == 1
+        assert len(res[0]) == gen_n
+    else:
+        # do not submit empty source batches
+        assert len(res) == 0