Fix streaming input for batch and multiple sources (#99)

This fixes the behavior of multiple source tasks to match that
of a branched task (broadcast outputs to children). It also adds
support for generated input of length > 1 to be handled as expected
included possibly splitting the batch.

Author: amitschang

dplutils/pipeline/stream.py

@@ -65,8 +65,10 @@ class StreamingGraphExecutor(PipelineExecutor, ABC):
           to the input task(s). Default is None, which means either exhaust the
           source generator or run indefinitely.
         generator: A callable that when called returns a generator which yields
-          dataframes. The yielded dataframes are assumed to be a single row, in
-          which case input task batching will be honored.
+          dataframes. The driver will call ``len()`` on the yielded dataframes to
+          obtain the number of rows and will split and batch according to the input
+          task settings. Each generated dataframe, regardless of size, counts as
+          a single source batch with respect to ``max_batches``.
 
 
     Implementations must override abstract methods for (remote) task submission
@@ -136,39 +138,41 @@ def resolve_completed(self):
                 task.data_in.extendleft(self.task_resolve_output(ready))
         return None
 
-    def process_source(self, source):
-        source_batch = []
-        for _ in range(source.task.batch_size or 1):
+    def _feed_source(self, source):
+        if self.source_exhausted:
+            return
+        total_length = sum(i.length for i in source.data_in)
+        while total_length < (source.task.batch_size or 1):
             try:
-                source_batch.append(next(self.source_generator))
+                next_df = next(self.source_generator)
             except StopIteration:
                 self.source_exhausted = True
-                if len(source_batch) == 0:
-                    return
                 break
+            # We feed any generated source to all source tasks similar the way
+            # upstream forked outputs broadcast. We add to data_in so that any
+            # necessary batching and splitting can be handled by normal procedure.
+            for task in self.stream_graph.source_tasks:
+                task.data_in.append(StreamBatch(data=next_df, length=len(next_df)))
             self.n_sourced += 1
             if self.n_sourced == self.max_batches:
                 self.source_exhausted = True
                 break
-        source.pending.appendleft(self.task_submit(source.task, source_batch))
-        source.counter += 1
-        return
+            total_length += len(next_df)
 
     def enqueue_tasks(self):
         # Work through the graph in reverse order, submitting any tasks as
         # needed. Reverse order ensures we prefer to send tasks that are closer
         # to the end of the pipeline and only feed as necessary.
-        rank = 0
-        for task in self.stream_graph.walk_back(sort_key=lambda x: x.counter):
-            if task in self.stream_graph.source_tasks or len(task.data_in) == 0:
-                continue
+        def _handle_one_task(task, rank):
+            eligible = submitted = False
+            if len(task.data_in) == 0:
+                return (eligible, submitted)
 
             batch_size = task.task.batch_size
             if batch_size is not None:
                 for batch in deque_extract(task.data_in, lambda b: b.length > batch_size):
                     task.split_pending.appendleft(self.split_batch_submit(batch, batch_size))
 
-            eligible = False
             while len(task.data_in) > 0:
                 num_to_merge = deque_num_merge(task.data_in, batch_size)
                 if num_to_merge == 0:
@@ -184,18 +188,29 @@ def enqueue_tasks(self):
                 merged = [task.data_in.pop().data for i in range(num_to_merge)]
                 task.pending.appendleft(self.task_submit(task.task, merged))
                 task.counter += 1
+                submitted = True
+            return (eligible, submitted)
 
+        # proceed through all non-source tasks, which will be handled separately
+        # below due to the need to feed from generator.
+        rank = 0
+        for task in self.stream_graph.walk_back(sort_key=lambda x: x.counter):
+            if task in self.stream_graph.source_tasks:
+                continue
+            eligible, _ = _handle_one_task(task, rank)
             if eligible:  # update rank of this task if it _could_ be done, whether or not it was
                 rank += 1
 
-        # in least-run order try to enqueue as many of the source tasks as can fit
+        # Source as many inputs as can fit on source tasks. We prioritize flushing the
+        # input queue and secondarily on number of invocations in case batch sizes differ.
         while True:
             task_scheduled = False
-            for source in sorted(self.stream_graph.source_tasks, key=lambda x: x.counter):
-                if self.source_exhausted or not self.task_submittable(source.task, rank):
-                    continue
-                self.process_source(source)
-                task_scheduled = True
+            for task in sorted(self.stream_graph.source_tasks, key=lambda x: (-len(x.data_in), x.counter)):
+                if self.task_submittable(task.task, rank):
+                    self._feed_source(task)
+                    _, task_scheduled = _handle_one_task(task, rank)
+                    if task_scheduled:  # we want to re-evalute the sort order
+                        break
             if not task_scheduled:
                 break
 

tests/conftest.py

@@ -68,6 +68,14 @@ def multi_output_graph():
     return PipelineGraph([(t1, t2A), (t1, t2B)])
 
 
+@pytest.fixture
+def multi_source_graph():
+    t1A = PipelineTask("srcA", lambda x: x.assign(a="A"))
+    t1B = PipelineTask("srcB", lambda x: x.assign(b="B"))
+    t2 = PipelineTask("sink", lambda x: x.assign(sink="sink"))
+    return PipelineGraph([(t1A, t2), (t1B, t2)])
+
+
 @pytest.fixture
 def graph_suite(dummy_steps, dummy_pipeline_graph, multi_output_graph):
     return {

tests/pipeline/test_stream_executor.py

@@ -62,3 +62,18 @@ def generator():
     else:
         # do not submit empty source batches
         assert len(res) == 0
+
+
+def test_stream_executor_input_batch_size_splits(dummy_steps):
+    def generator():
+        for i in range(4):
+            yield pd.DataFrame({"col": range(2)})
+
+    # sanity check to ensure below test is actually inspect the split action
+    pl = LocalSerialExecutor(dummy_steps, generator=generator)
+    res = [i.data for i in pl.run()]
+    assert len(res) == 4
+    # explicitly set batch size so we should call split on each input
+    pl = LocalSerialExecutor(dummy_steps, generator=generator).set_config("task1.batch_size", 1)
+    res = [i.data for i in pl.run()]
+    assert len(res) == 8

tests/pipeline/test_suite.py

@@ -1,3 +1,5 @@
+from math import ceil
+
 import pandas as pd
 import pytest
 
@@ -105,3 +107,14 @@ def test_with_merge_batch(self, dummy_steps, max_batches):
         assert all([len(i.data) == expected_len for i in res])
         final = pd.concat([i.data for i in res])
         assert final["id"].nunique() == max_batches
+
+    def test_all_sources_get_each_input_and_batch(self, multi_source_graph, max_batches):
+        pl = self.executor(multi_source_graph, max_batches=max_batches).set_config("srcA.batch_size", 4)
+        res = list(pl.run())
+        expected_len = max_batches + ceil(max_batches / 4)
+        assert len(res) == expected_len
+        # now ensure that each source task got same set of batches
+        final = pd.concat([i.data for i in res])
+        idset_a = set(final[final.a == "A"].id)
+        idset_b = set(final[final.b == "B"].id)
+        assert idset_a == idset_b == set(range(max_batches))