Skip zero-length batches in stream executor (#90)

Zero length data batches risk starting tasks and serializing data
unnecessarily, consuming resources for batches that ultimately would
not be expected to yield any data.

Resolves #62

Author: amitschang

dplutils/pipeline/stream.py

@@ -124,6 +124,8 @@ def resolve_completed(self):
         for task in self.stream_graph.walk_fwd():
             for ready in deque_extract(task.pending, self.is_task_ready):
                 block_info = self.task_resolve_output(ready)
+                if block_info.length == 0:
+                    continue
                 if task in self.stream_graph.sink_tasks:
                     return OutputBatch(block_info.data, task=task.name)
                 else:

tests/conftest.py

@@ -46,6 +46,11 @@ def dummy_steps():
     ]
 
 
+@pytest.fixture
+def blackhole_step():
+    return [PipelineTask("blackhole", func=lambda x: pd.DataFrame([]))]
+
+
 @pytest.fixture
 def dummy_pipeline_graph():
     t1 = PipelineTask("task1", lambda x: x.assign(t1="1"))

tests/pipeline/test_suite.py

@@ -2,6 +2,7 @@
 import pytest
 
 from dplutils.pipeline import OutputBatch
+from dplutils.pipeline.task import PipelineTask
 
 
 @pytest.mark.parametrize("max_batches", (1, 4, 10))
@@ -22,6 +23,29 @@ def test_run_simple_pipeline_iterator(self, dummy_steps, max_batches):
                 assert isinstance(res.data, pd.DataFrame)
                 assert set(res.data.columns).issuperset({"id", "step1", "step2"})
 
+    def test_pipeline_throws_away_empty_batches(self, blackhole_step, max_batches, tmp_path):
+        # Flag to ensure that subsequent tasks are not called on empty batches. This is useful
+        # since the check of yielded results only checks output of pipeline, we want to ensure tasks
+        # don't needlessly get called as well. Flag implemented with file so it is more portable
+        # for subprocess calls (e.g. used in ray)
+        flag_file = tmp_path / "called.flag"
+
+        def set_counter(x):
+            flag_file.write_text("")
+            return x
+
+        # first ensure operation of our counter instrument
+        pl = self.executor([PipelineTask("nocalls", set_counter)], max_batches=max_batches)
+        res = list(pl.run())
+        assert len(res) == max_batches
+        assert flag_file.exists()
+        # now ensure we toss empty dataframes
+        flag_file.unlink()
+        pl = self.executor([*blackhole_step, PipelineTask("nocalls", set_counter)], max_batches=max_batches)
+        res = list(pl.run())
+        assert len(res) == 0
+        assert not flag_file.exists()
+
     def test_run_dag_pipeline(self, dummy_pipeline_graph, max_batches):
         pl = self.executor(dummy_pipeline_graph, max_batches=max_batches)
         it = pl.run()