Repository: modin-project/modin

Issue:
Dask engine ignore lengths parameter



Discussion:

Resolving patch:
diff --git a/modin/data_management/utils.py b/modin/data_management/utils.py
--- a/modin/data_management/utils.py
+++ b/modin/data_management/utils.py
@@ -82,18 +82,19 @@ def split_result_of_axis_func_pandas(axis, num_splits, result, length_list=None)
     list
         A list of Pandas DataFrames.
     """
-    if num_splits == 1:
-        return result
     if length_list is not None:
         length_list.insert(0, 0)
         sums = np.cumsum(length_list)
-        if axis == 0:
+        if axis == 0 or isinstance(result, pandas.Series):
             return [result.iloc[sums[i] : sums[i + 1]] for i in range(len(sums) - 1)]
         else:
             return [result.iloc[:, sums[i] : sums[i + 1]] for i in range(len(sums) - 1)]
+
+    if num_splits == 1:
+        return [result]
     # We do this to restore block partitioning
     chunksize = compute_chunksize(result, num_splits, axis=axis)
-    if axis == 0:
+    if axis == 0 or isinstance(result, pandas.Series):
         return [
             result.iloc[chunksize * i : chunksize * (i + 1)] for i in range(num_splits)
         ]
diff --git a/modin/engines/base/frame/axis_partition.py b/modin/engines/base/frame/axis_partition.py
--- a/modin/engines/base/frame/axis_partition.py
+++ b/modin/engines/base/frame/axis_partition.py
@@ -91,10 +91,7 @@ def shuffle(self, func, lengths, **kwargs):
     partition_type = None
 
     def _wrap_partitions(self, partitions):
-        if isinstance(partitions, self.instance_type):
-            return [self.partition_type(partitions)]
-        else:
-            return [self.partition_type(obj) for obj in partitions]
+        return [self.partition_type(obj) for obj in partitions]
 
 
 class PandasFrameAxisPartition(BaseFrameAxisPartition):
@@ -216,10 +213,6 @@ def deploy_axis_func(
 
         dataframe = pandas.concat(list(partitions), axis=axis, copy=False)
         result = func(dataframe, **kwargs)
-        if isinstance(result, pandas.Series):
-            if num_splits == 1:
-                return result
-            return [result] + [pandas.Series([]) for _ in range(num_splits - 1)]
 
         if manual_partition:
             # The split function is expecting a list
diff --git a/modin/engines/dask/pandas_on_dask/frame/axis_partition.py b/modin/engines/dask/pandas_on_dask/frame/axis_partition.py
--- a/modin/engines/dask/pandas_on_dask/frame/axis_partition.py
+++ b/modin/engines/dask/pandas_on_dask/frame/axis_partition.py
@@ -43,13 +43,15 @@ def deploy_axis_func(
             *partitions,
             pure=False,
         )
-        if num_splits == 1:
-            return axis_result
+
+        lengths = kwargs.get("_lengths", None)
+        result_num_splits = len(lengths) if lengths else num_splits
+
         # We have to do this to split it back up. It is already split, but we need to
         # get futures for each.
         return [
             client.submit(lambda l: l[i], axis_result, pure=False)
-            for i in range(num_splits)
+            for i in range(result_num_splits)
         ]
 
     @classmethod
@@ -68,8 +70,6 @@ def deploy_func_between_two_axis_partitions(
             *partitions,
             pure=False,
         )
-        if num_splits == 1:
-            return axis_result
         # We have to do this to split it back up. It is already split, but we need to
         # get futures for each.
         return [


Test patch: