🐛 Fix concatenation when trailing result chunk is smaller (#28)

ddelange · Sep 7, 2022 · 019ee26 · 019ee26
1 parent d997507
commit 019ee26
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/src/mapply/mapply.py b/src/mapply/mapply.py
@@ -116,11 +116,7 @@ def run_apply(func, df_or_series, args=(), **kwargs):
         )
     )
 
-    if (
-        isseries
-        or len(results) == 1
-        or len(results[0]) * len(results) in df_or_series.shape
-    ):
+    if isseries or len(results) == 1 or sum(map(len, results)) in df_or_series.shape:
         return concat(results, copy=False)
 
     return concat(results, axis=1, copy=False)
diff --git a/tests/test_mapply.py b/tests/test_mapply.py
@@ -57,6 +57,14 @@ def test_df_mapply():
         df.mapply(lambda x: x**2),
     )
 
+    # not all result chunks have equal size (trailing chunk)
+    mapply.init(progressbar=False, chunk_size=100, n_workers=2)
+    df = pd.DataFrame(np.random.randint(2, size=(5, 201)))
+    pd.testing.assert_series_equal(
+        df.apply(np.var),
+        df.mapply(np.var),
+    )
+
     # concat for only one result
     mapply.init(progressbar=False, chunk_size=100, n_workers=2)
     df = pd.DataFrame(list(range(1, 200)))  # (199, 1)