FIX-modin-project#2322: fix count of return values of ray remote funcs

Signed-off-by: Anatoly Myachev <anatoly.myachev@intel.com>
anmyachev · Nov 6, 2020 · 6a00aef · 6a00aef
1 parent d0fb8ef
commit 6a00aef
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/modin/engines/ray/pandas_on_ray/frame/axis_partition.py b/modin/engines/ray/pandas_on_ray/frame/axis_partition.py
@@ -33,6 +33,7 @@ def __init__(self, list_of_blocks):
     def deploy_axis_func(
         cls, axis, func, num_splits, kwargs, maintain_partitioning, *partitions
     ):
+        lengths = kwargs.get("_lengths", None)
         return deploy_ray_func._remote(
             args=(
                 PandasFrameAxisPartition.deploy_axis_func,
@@ -43,7 +44,7 @@ def deploy_axis_func(
                 maintain_partitioning,
             )
             + tuple(partitions),
-            num_returns=num_splits * 3,
+            num_returns=num_splits * 3 if lengths is None else len(lengths) * 3,
         )
 
     @classmethod