PR comments

Signed-off-by: arunjose696 <[email protected]>

Author: arunjose696

modin/core/storage_formats/pandas/utils.py

@@ -84,34 +84,11 @@ def split_result_of_axis_func_pandas(
     list of pandas.DataFrames
         Splitted dataframe represented by list of frames.
     """
-    if num_splits == 1:
-        return [result]
-
-    if length_list is None:
-        length_list = get_length_list(result.shape[axis], num_splits, min_block_size)
-    # Inserting the first "zero" to properly compute cumsum indexing slices
-    length_list = np.insert(length_list, obj=0, values=[0])
-
-    sums = np.cumsum(length_list)
-    axis = 0 if isinstance(result, pandas.Series) else axis
-    # We do this to restore block partitioning
-    if axis == 0:
-        chunked = [result.iloc[sums[i] : sums[i + 1]] for i in range(len(sums) - 1)]
-    else:
-        chunked = [result.iloc[:, sums[i] : sums[i + 1]] for i in range(len(sums) - 1)]
-
-    return [
-        # Sliced MultiIndex still stores all encoded values of the original index, explicitly
-        # asking it to drop unused values in order to save memory.
-        (
-            chunk.set_axis(
-                chunk.axes[axis].remove_unused_levels(), axis=axis, copy=False
-            )
-            if isinstance(chunk.axes[axis], pandas.MultiIndex)
-            else chunk
+    return list(
+        generate_result_of_axis_func_pandas(
+            axis, num_splits, result, length_list, min_block_size
         )
-        for chunk in chunked
-    ]
+    )
 
 
 def generate_result_of_axis_func_pandas(