adding generators to split_result_of_axis_func_pandas

Author: arunjose696

modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py

@@ -320,16 +320,17 @@ def _deploy_ray_func(
     f_args = positional_args[:f_len_args]
     deploy_args = positional_args[f_len_args:]
     result = deployer(axis, f_to_deploy, f_args, f_kwargs, *deploy_args, **kwargs)
+
     if not extract_metadata:
-        return result
+        for item in result:
+            yield item
+        return
+
     ip = get_node_ip_address()
-    if isinstance(result, pandas.DataFrame):
-        return result, len(result), len(result.columns), ip
-    elif all(isinstance(r, pandas.DataFrame) for r in result):
-        for r in result:
+    for r in result:
+        if isinstance(r, pandas.DataFrame):
             for item in [r, len(r), len(r.columns), ip]:
                 yield item
-    else:
-        for r in result:
+        else:
             for item in [r, None, None, ip]:
                 yield item

modin/core/storage_formats/pandas/utils.py

@@ -80,7 +80,8 @@ def split_result_of_axis_func_pandas(axis, num_splits, result, length_list=None)
         Splitted dataframe represented by list of frames.
     """
     if num_splits == 1:
-        return [result]
+        yield result
+        return
 
     if length_list is None:
         length_list = get_length_list(result.shape[axis], num_splits)
@@ -89,24 +90,21 @@ def split_result_of_axis_func_pandas(axis, num_splits, result, length_list=None)
 
     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 [
+    for i in range(len(sums) - 1):
+        # We do this to restore block partitioning
+        if axis == 0:
+            chunk = result.iloc[sums[i] : sums[i + 1]]
+        else:
+            chunk = result.iloc[:, sums[i] : sums[i + 1]]
+
         # 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(
+        if isinstance(chunk.axes[axis], pandas.MultiIndex):
+            chunk = chunk.set_axis(
                 chunk.axes[axis].remove_unused_levels(), axis=axis, copy=False
             )
-            if isinstance(chunk.axes[axis], pandas.MultiIndex)
-            else chunk
-        )
-        for chunk in chunked
-    ]
+        yield chunk
 
 
 def get_length_list(axis_len: int, num_splits: int) -> list: