Repository: nalepae/pandarallel

Issue:
parallel_apply(..., axis=0) return duplicated rows.
this behavior is different from pandas.

Discussion:

Patch:
diff --git a/pandarallel/core.py b/pandarallel/core.py
index f96359c..8eaaf12 100644
--- a/pandarallel/core.py
+++ b/pandarallel/core.py
@@ -235,7 +235,7 @@ def parallelize_with_memory_file_system(
         progresses_length = [len(chunk_) * multiplicator_factor for chunk_ in chunks]
 
         work_extra = data_type.get_work_extra(data)
-        reduce_extra = data_type.get_reduce_extra(data)
+        reduce_extra = data_type.get_reduce_extra(data, user_defined_function_kwargs)
 
         show_progress_bars = progress_bars_type != ProgressBarsType.No
 
@@ -376,7 +376,7 @@ def parallelize_with_pipe(
         progresses_length = [len(chunk_) * multiplicator_factor for chunk_ in chunks]
 
         work_extra = data_type.get_work_extra(data)
-        reduce_extra = data_type.get_reduce_extra(data)
+        reduce_extra = data_type.get_reduce_extra(data, user_defined_function_kwargs)
 
         show_progress_bars = progress_bars_type != ProgressBarsType.No
 
diff --git a/pandarallel/data_types/dataframe.py b/pandarallel/data_types/dataframe.py
index 5e2e067..4454c6f 100644
--- a/pandarallel/data_types/dataframe.py
+++ b/pandarallel/data_types/dataframe.py
@@ -4,6 +4,7 @@ import pandas as pd
 
 from ..utils import chunk
 from .generic import DataType
+from ..utils import _get_axis_int, _opposite_axis_int
 
 
 class DataFrame:
@@ -13,13 +14,9 @@ class DataFrame:
             nb_workers: int, data: pd.DataFrame, **kwargs
         ) -> Iterator[pd.DataFrame]:
             user_defined_function_kwargs = kwargs["user_defined_function_kwargs"]
-            axis = user_defined_function_kwargs.get("axis", 0)
 
-            if axis not in {0, 1, "index", "columns"}:
-                raise ValueError(f"No axis named {axis} for object type DataFrame")
-
-            axis_int = {0: 0, 1: 1, "index": 0, "columns": 1}[axis]
-            opposite_axis_int = 1 - axis_int
+            axis_int = _get_axis_int(user_defined_function_kwargs)
+            opposite_axis_int = _opposite_axis_int(axis_int)
 
             for chunk_ in chunk(data.shape[opposite_axis_int], nb_workers):
                 yield data.iloc[chunk_] if axis_int == 1 else data.iloc[:, chunk_]
@@ -38,11 +35,19 @@ class DataFrame:
                 **user_defined_function_kwargs,
             )
 
+        @staticmethod
+        def get_reduce_extra(data: Any, user_defined_function_kwargs) -> Dict[str, Any]:
+            return {"axis": _get_axis_int(user_defined_function_kwargs)}
+
         @staticmethod
         def reduce(
             datas: Iterable[pd.DataFrame], extra: Dict[str, Any]
         ) -> pd.DataFrame:
-            return pd.concat(datas, copy=False)
+            if isinstance(datas[0], pd.Series):
+                axis = 0
+            else:
+                axis = _opposite_axis_int(extra["axis"])
+            return pd.concat(datas, copy=False, axis=axis)
 
     class ApplyMap(DataType):
         @staticmethod
diff --git a/pandarallel/data_types/dataframe_groupby.py b/pandarallel/data_types/dataframe_groupby.py
index f6d119f..184088b 100644
--- a/pandarallel/data_types/dataframe_groupby.py
+++ b/pandarallel/data_types/dataframe_groupby.py
@@ -40,7 +40,7 @@ class DataFrameGroupBy:
             return [compute_result(key, df) for key, df in data]
 
         @staticmethod
-        def get_reduce_extra(data: PandasDataFrameGroupBy) -> Dict[str, Any]:
+        def get_reduce_extra(data: PandasDataFrameGroupBy, user_defined_function_kwargs) -> Dict[str, Any]:
             return {"df_groupby": data}
 
         @staticmethod
diff --git a/pandarallel/data_types/generic.py b/pandarallel/data_types/generic.py
index 3b50383..34dfeb4 100644
--- a/pandarallel/data_types/generic.py
+++ b/pandarallel/data_types/generic.py
@@ -24,7 +24,7 @@ class DataType(ABC):
         ...
 
     @staticmethod
-    def get_reduce_extra(data: Any) -> Dict[str, Any]:
+    def get_reduce_extra(data: Any, user_defined_function_kwargs) -> Dict[str, Any]:
         return dict()
 
     @staticmethod
diff --git a/pandarallel/utils.py b/pandarallel/utils.py
index c83e93e..64c85b8 100644
--- a/pandarallel/utils.py
+++ b/pandarallel/utils.py
@@ -87,6 +87,18 @@ def get_pandas_version() -> Tuple[int, int]:
     return int(major_str), int(minor_str)
 
 
+def _get_axis_int(user_defined_function_kwargs):
+    axis = user_defined_function_kwargs.get("axis", 0)
+
+    if axis not in {0, 1, "index", "columns"}:
+        raise ValueError(f"No axis named {axis} for object type DataFrame")
+
+    axis_int = {0: 0, 1: 1, "index": 0, "columns": 1}[axis]
+    return axis_int
+
+def _opposite_axis_int(axis: int):
+    return 1 - axis
+
 class WorkerStatus(int, Enum):
     Running = 0
     Success = 1


Test patch:
diff --git a/tests/test_pandarallel.py b/tests/test_pandarallel.py
index 6cbdc8d..03d5cb4 100644
--- a/tests/test_pandarallel.py
+++ b/tests/test_pandarallel.py
@@ -135,6 +135,26 @@ def func_dataframe_groupby_expanding_apply(request):
     )[request.param]
 
 
+@pytest.fixture(params=("named", "anonymous"))
+def func_dataframe_apply_axis_0_no_reduce(request):
+    def func(x):
+        return x
+
+    return dict(
+        named=func, anonymous=lambda x: x
+    )[request.param]
+
+
+@pytest.fixture(params=("named", "anonymous"))
+def func_dataframe_apply_axis_1_no_reduce(request):
+    def func(x):
+        return x**2
+
+    return dict(
+        named=func, anonymous=lambda x: x**2
+    )[request.param]
+
+
 @pytest.fixture
 def pandarallel_init(progress_bar, use_memory_fs):
     pandarallel.initialize(
@@ -290,3 +310,29 @@ def test_dataframe_groupby_expanding_apply(
         .parallel_apply(func_dataframe_groupby_expanding_apply, raw=False)
     )
     res.equals(res_parallel)
+
+def test_dataframe_axis_0_no_reduction(
+    pandarallel_init, func_dataframe_apply_axis_0_no_reduce, df_size
+):
+    df = pd.DataFrame(
+        dict(a=np.random.randint(1, 10, df_size), b=np.random.randint(1, 10, df_size), c=np.random.randint(1, 10, df_size))
+    )
+    res = df.apply(func_dataframe_apply_axis_0_no_reduce)
+
+    res_parallel = df.parallel_apply(func_dataframe_apply_axis_0_no_reduce)
+
+    assert res.equals(res_parallel)
+
+def test_dataframe_axis_1_no_reduction(
+    pandarallel_init, func_dataframe_apply_axis_1_no_reduce, df_size
+):
+    df = pd.DataFrame(
+        dict(a=np.random.randint(1, 10, df_size), b=np.random.randint(1, 10, df_size), c=np.random.randint(1, 10, df_size))
+    )
+
+    res = df.apply(func_dataframe_apply_axis_1_no_reduce, axis=1)
+
+    res_parallel = df.parallel_apply(func_dataframe_apply_axis_1_no_reduce, axis=1)
+
+    assert res.equals(res_parallel)
+