Test2

modin/core/execution/ray/common/deferred_execution.py

@@ -12,8 +12,19 @@
 # governing permissions and limitations under the License.
 from builtins import NotImplementedError
 from enum import Enum
+from itertools import islice
 from types import GeneratorType
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union, Generator
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    Generator,
+)
 
 import pandas
 import ray
@@ -101,12 +112,7 @@ def _ref(obj):
 
     def exec(
         self,
-    ) -> Tuple[
-        ObjectRefOrListType,
-        ObjectRefOrListType,
-        ObjectRefOrListType,
-        ObjectRefOrListType,
-    ]:
+    ) -> Tuple[ObjectRefOrListType, Union["MetaList", List], Union[int, List[int]]]:
         """
         Execute this task, if required.
 
@@ -116,63 +122,58 @@ def exec(
             The execution result, length, width and the worker's ip address.
         """
         if not self.has_result:
-            if self.num_returns == 1 and not isinstance(self.data, DeferredExecution):
+            data = self.data
+            if self.num_returns == 1 and not isinstance(data, DeferredExecution):
                 task = self.task
                 if task[3] and task[4]:
-                    result = remote_exec_func.remote(
-                        task[0], self.data, *task[1], **task[2]
+                    if isinstance(data, ray.ObjectRef):
+                        ray.wait([data])
+                    result, length, width, ip = remote_exec_func.remote(
+                        task[0], data, *task[1], **task[2]
                     )
-                    self._set_result(*result)
-                    return result
+                    meta = MetaList([length, width, ip])
+                    self._set_result(result, meta, 0)
+                    return result, meta, 0
 
             consumers, output = deconstruct(self)
-            num_returns = sum(c.num_returns for c in consumers)
+            num_returns = sum(c.num_returns for c in consumers) + 1
+            if isinstance(output[0], ray.ObjectRef):
+                ray.wait(output[0:1])
             results = _remote_exec_chain(num_returns, *output)
-            if not isinstance(results, ListOrTuple):
-                results = [results]
+            meta = MetaList(results.pop())
+            meta_off = 0
             results = iter(results)
             for de in consumers:
                 if de.num_returns == 1:
-                    de._set_result(
-                        next(results),
-                        None,
-                        None,
-                        None,
-                    )
+                    de._set_result(next(results), meta, meta_off)
+                    meta_off += 2
                 else:
-                    res = []
-                    lengths = []
-                    widths = []
-                    for _ in range(de.num_returns):
-                        res.append(next(results))
-                        lengths.append(None)
-                        widths.append(None)
-                    de._set_result(res, lengths, widths, None)
+                    res = list(islice(results, num_returns))
+                    offsets = list(range(0, 2 * num_returns, 2))
+                    de._set_result(res, res, offsets)
+                    meta_off += 2 * num_returns
 
-        return self.data, self.length, self.width, self.ip
+        return self.data, self.meta, self.meta_off
 
     def _set_result(
         self,
         result: ObjectRefOrListType,
-        length: ObjectRefOrListType,
-        width: ObjectRefOrListType,
-        ip: ObjectRefOrListType,
+        meta: Union["MetaList", List],
+        meta_off: Union[int, List[int]],
     ):
         """
         Set the execution result.
 
         Parameters
         ----------
         result : ObjectRefOrListType
-        length : ObjectRefOrListType
-        width : ObjectRefOrListType
-        ip : ObjectRefOrListType
+        meta : MetaList or list
+        meta_off : int or list of int
         """
         del self.task
         self.data = result
-        self.length = length
-        self.width = width
-        self.ip = ip
+        self.meta = meta
+        self.meta_off = meta_off
 
     @property
     def has_result(self):
@@ -200,11 +201,25 @@ def __reduce__(self):
         raise NotImplementedError()
 
 
-ListOrDe = (DeferredExecution, list, tuple)
+class MetaList:
+    def __init__(self, obj: Union[ray.ObjectRef, ClientObjectRef, List]):
+        self._obj = obj
 
+    def __getitem__(self, item):
+        obj = self._obj
+        if not isinstance(obj, list):
+            from modin.core.execution.ray.common import RayWrapper
 
-def has_list_or_de(it: Iterable):
-    return any(isinstance(i, ListOrDe) for i in it)
+            self._obj = obj = RayWrapper.materialize(obj)
+        return obj[item]
+
+    def __setitem__(self, key, value):
+        obj = self._obj
+        if not isinstance(obj, list):
+            from modin.core.execution.ray.common import RayWrapper
+
+            self._obj = obj = RayWrapper.materialize(obj)
+        obj[key] = value
 
 
 class DeferredExecutionException(Exception):
@@ -223,6 +238,13 @@ class _Tag(Enum):
     END = 3
 
 
+ListOrDe = (DeferredExecution, list, tuple)
+
+
+def has_list_or_de(it: Iterable):
+    return any(isinstance(i, ListOrDe) for i in it)
+
+
 def deconstruct(de: DeferredExecution) -> Tuple[List[DeferredExecution], List[Any]]:
     stack = []
     result_consumers = []
@@ -399,8 +421,9 @@ def _remote_exec_multi_chain(num_returns: int, *args: Tuple) -> List:
 
 def construct(num_returns: int, args: Tuple) -> List:
     chain = list(reversed(args))
+    meta = []
     try:
-        stack = [construct_chain(chain, {}, None)]
+        stack = [construct_chain(chain, {}, meta, None)]
         while stack:
             try:
                 gen = stack.pop()
@@ -411,7 +434,8 @@ def construct(num_returns: int, args: Tuple) -> List:
                 else:
                     yield obj
             except StopIteration:
-                ...
+                meta.append(get_node_ip_address())
+                yield meta
     except DeferredExecutionException as err:
         for _ in range(num_returns):
             yield err
@@ -426,7 +450,8 @@ def construct(num_returns: int, args: Tuple) -> List:
 def construct_chain(
     chain: List,
     refs: Dict[int, Any],
-    lst: List,
+    meta: List,
+    lst: Optional[List],
 ):
     pop = chain.pop
     tg_e = _Tag.END
@@ -438,7 +463,7 @@ def construct_chain(
         obj = refs[pop()]
     elif obj is _Tag.LIST:
         obj = []
-        yield construct_list(obj, chain, refs)
+        yield construct_list(obj, chain, refs, meta)
         if isinstance(obj[0], DeferredExecutionException):
             raise obj[0]
 
@@ -457,38 +482,35 @@ def construct_chain(
                 args.reverse()
         else:
             args = []
-            yield construct_list(args, chain, refs)
+            yield construct_list(args, chain, refs, meta)
         if (args_len := pop()) >= 0:
             kwargs = {pop(): pop() for _ in range(args_len)}
         else:
             values = []
-            yield construct_list(values, chain, refs)
+            yield construct_list(values, chain, refs, meta)
             kwargs = {pop(): v for v in values}
 
         obj = _exec_func(fn, obj, args, kwargs)
 
         if ref := pop():
             refs[ref] = obj
         if pop():
-            yield obj
-
-
-def _add_result(obj: Any, output: List):
-    out_append = output.append
-    if isinstance(obj, pandas.DataFrame):
-        out_append(ray.put(obj) if obj.memory_usage(deep=True).sum() > 100000 else obj)
-        out_append(len(obj))
-        out_append(len(obj.columns))
-    else:
-        out_append(obj)
-        out_append(0)
-        out_append(0)
+            if isinstance(obj, ListOrTuple):
+                for o in obj:
+                    meta.append(len(o) if hasattr(o, "__len__") else 0)
+                    meta.append(len(o.columns) if hasattr(o, "columns") else 0)
+                    yield o
+            else:
+                meta.append(len(obj) if hasattr(obj, "__len__") else 0)
+                meta.append(len(obj.columns) if hasattr(obj, "columns") else 0)
+                yield obj
 
 
 def construct_list(
     lst: List,
     chain: List,
     refs: Dict[int, Any],
+    meta: List,
 ):
     pop = chain.pop
     lst_append = lst.append
@@ -498,10 +520,10 @@ def construct_list(
             if obj == _Tag.END:
                 break
             elif obj == _Tag.CHAIN:
-                yield construct_chain(chain, refs, lst)
+                yield construct_chain(chain, refs, meta, lst)
             elif obj == _Tag.LIST:
                 lst_append([])
-                yield construct_list(lst[-1], chain, refs)
+                yield construct_list(lst[-1], chain, refs, meta)
             elif obj is _Tag.REF:
                 lst_append(refs[pop()])
             else:

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

@@ -25,6 +25,7 @@
     DeferredExecution,
     has_list_or_de,
     remote_exec_func,
+    MetaList,
 )
 from modin.core.execution.ray.common.utils import ObjectIDType
 from modin.logging import get_logger
@@ -64,9 +65,8 @@ def __init__(
         if isinstance(data, DeferredExecution):
             data.ref_count(1)
         self._data_ref = data
-        self._length_cache = length
-        self._width_cache = width
-        self._ip_cache = ip
+        self.meta = MetaList([length, width, ip])
+        self.meta_off = 0
 
     def __del__(self):
         """Decrement the reference counter."""
@@ -116,9 +116,8 @@ def drain_call_queue(self):
         if isinstance(data, DeferredExecution):
             (
                 self._data_ref,
-                self._length_cache,
-                self._width_cache,
-                self._ip_cache,
+                self._meta,
+                self._meta_off,
             ) = data.exec()
             data.ref_count(-1)
 
@@ -302,6 +301,30 @@ def _data(self):
         self.drain_call_queue()
         return self._data_ref
 
+    @property
+    def _length_cache(self):
+        return self.meta[self.meta_off]
+
+    @_length_cache.setter
+    def _length_cache(self, value):
+        self.meta[self.meta_off] = value
+
+    @property
+    def _width_cache(self):
+        return self.meta[self.meta_off + 1]
+
+    @_width_cache.setter
+    def _width_cache(self, value):
+        self.meta[self.meta_off + 1] = value
+
+    @property
+    def _ip_cache(self):
+        return self.meta[-1]
+
+    @_ip_cache.setter
+    def _ip_cache(self, value):
+        self.meta[-1] = value
+
 
 @ray.remote(num_returns=2)
 def _get_index_and_columns(df):  # pragma: no cover