Fix dataframe splitting numpy issue (#91)

Resolves #87

Author: amitschang

dplutils/pipeline/ray.py

@@ -10,6 +10,7 @@
 from dplutils import observer
 from dplutils.pipeline import OutputBatch, PipelineExecutor, PipelineTask
 from dplutils.pipeline.stream import StreamBatch, StreamingGraphExecutor
+from dplutils.pipeline.utils import split_dataframe
 
 
 def set_run_id(inputs, run_id):
@@ -31,7 +32,7 @@ def wrapper(indf):
         if task.batch_size is None:
             return funcwrapper(indf, **kwargs)
 
-        splits = np.array_split(indf, np.ceil(len(indf) / task.batch_size))
+        splits = split_dataframe(indf, max_rows=task.batch_size)
         refs = [
             ray.remote(funcwrapper)
             .options(
@@ -123,7 +124,7 @@ def wrapper(*df_list):
 
 
 def stream_split_func(df, splits):
-    splits = np.array_split(df, splits)
+    splits = split_dataframe(df, num_splits=splits)
     return [len(i) for i in splits] + splits
 
 

dplutils/pipeline/stream.py

@@ -9,7 +9,7 @@
 import pandas as pd
 
 from dplutils.pipeline import OutputBatch, PipelineExecutor, PipelineTask
-from dplutils.pipeline.utils import deque_extract
+from dplutils.pipeline.utils import deque_extract, split_dataframe
 
 
 @dataclass
@@ -310,7 +310,7 @@ def task_submit(self, pt, df_list):
 
     def split_batch_submit(self, stream_batch, max_rows):
         df = stream_batch.data
-        return np.array_split(df, np.ceil(len(df) / max_rows))
+        return split_dataframe(df, max_rows=max_rows)
 
     def task_resolve_output(self, to):
         if isinstance(to, list):

dplutils/pipeline/utils.py

@@ -1,3 +1,6 @@
+import numpy as np
+
+
 def dict_from_coord(coord, value):
     d = {}
     if "." in coord:
@@ -16,3 +19,19 @@ def deque_extract(queue, condition):
             yield queue.pop()
         else:
             queue.rotate()
+
+
+def split_dataframe(df, max_rows=None, num_splits=None):
+    """Split dataframe by max number of rows or given splits.
+
+    If num_splits is provided, ignore max_rows and produce a set of num_splits
+    splits with each dataframe having approximately equal number of rows (one may be
+    different in size).
+
+    If max_rows provided, produce the smallest number of splits that ensure no
+    single split has more than max_rows
+    """
+    if num_splits is None:
+        num_splits = np.ceil(len(df) / max_rows)
+    chunks = np.linspace(0, num_splits, num=len(df), endpoint=False, dtype=np.int32)
+    return [chunk for _, chunk in df.groupby(chunks)]

tests/pipeline/test_pipeline_utils.py

@@ -1,5 +1,43 @@
-from dplutils.pipeline.utils import dict_from_coord
+from collections import defaultdict
+
+import pandas as pd
+import pytest
+
+from dplutils.pipeline.utils import dict_from_coord, split_dataframe
 
 
 def test_dict_from_coord():
     assert dict_from_coord("a.b.c", "value") == {"a": {"b": {"c": "value"}}}
+
+
+@pytest.mark.parametrize(
+    "df_len, max_rows, num_splits, expected_splits",
+    [
+        (100, 20, None, 5),
+        (101, 20, None, 6),
+        (101, 21, None, 5),
+        (10, 1, None, 10),
+        (10, 5, None, 2),
+        (10, 10, None, 1),
+        (100, None, 5, 5),
+        (101, None, 5, 5),
+        (101, 50, 5, 5),  # this should prefer num_splits, not max
+    ],
+)
+def test_split_dataframe(df_len, max_rows, num_splits, expected_splits):
+    df = pd.DataFrame({"a": range(df_len), "b": range(df_len)})
+    splits = split_dataframe(df, max_rows=max_rows, num_splits=num_splits)
+    assert len(splits) == expected_splits
+    assert sum(len(i) for i in splits) == df_len
+    if max_rows and not num_splits:
+        assert max(len(i) for i in splits) <= max_rows
+    for s in splits:
+        assert isinstance(s, pd.DataFrame)
+        assert list(s.columns) == ["a", "b"]
+    # all but one should be same length
+    len_map = defaultdict(int)
+    for s in splits:
+        len_map[len(s)] += 1
+    assert len(len_map.keys()) <= 2
+    if len(len_map) > 1:
+        assert 1 in len_map.values()