diff --git a/src/mapply/mapply.py b/src/mapply/mapply.py index 63d6428..9604e47 100644 --- 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 index e9f3299..ce6861e 100644 --- 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)