diff --git a/modin/core/execution/ray/common/deferred_execution.py b/modin/core/execution/ray/common/deferred_execution.py index 0ace6c27b10..3e55ff546db 100644 --- a/modin/core/execution/ray/common/deferred_execution.py +++ b/modin/core/execution/ray/common/deferred_execution.py @@ -33,9 +33,11 @@ from ray.util.client.common import ClientObjectRef from modin.core.execution.ray.common import MaterializationHook, RayWrapper +from modin.core.execution.utils import remote_function from modin.logging import get_logger +from modin.utils import _inherit_docstrings -ObjectRefType = Union[ray.ObjectRef, ClientObjectRef, None] +ObjectRefType = Union[ray.ObjectRef, ClientObjectRef] ObjectRefOrListType = Union[ObjectRefType, List[ObjectRefType]] ListOrTuple = (list, tuple) @@ -68,16 +70,18 @@ class DeferredExecution: Attributes ---------- - data : ObjectRefType or DeferredExecution + data : object The execution input. func : callable or ObjectRefType A function to be executed. - args : list or tuple + args : list or tuple, optional Additional positional arguments to be passed in `func`. - kwargs : dict + kwargs : dict, optional Additional keyword arguments to be passed in `func`. - num_returns : int + num_returns : int, default: 1 The number of the return values. + flat_data : bool + True means that the data is neither DeferredExecution nor list. flat_args : bool True means that there are no lists or DeferredExecution objects in `args`. In this case, no arguments processing is performed and `args` is passed @@ -88,26 +92,29 @@ class DeferredExecution: def __init__( self, - data: Union[ - ObjectRefType, - "DeferredExecution", - List[Union[ObjectRefType, "DeferredExecution"]], - ], + data: Any, func: Union[Callable, ObjectRefType], - args: Union[List[Any], Tuple[Any]], - kwargs: Dict[str, Any], + args: Union[List[Any], Tuple[Any]] = None, + kwargs: Dict[str, Any] = None, num_returns=1, ): - if isinstance(data, DeferredExecution): - data.subscribe() + self.flat_data = self._flat_args((data,)) self.data = data self.func = func - self.args = args - self.kwargs = kwargs self.num_returns = num_returns - self.flat_args = self._flat_args(args) - self.flat_kwargs = self._flat_args(kwargs.values()) self.subscribers = 0 + if args is not None: + self.args = args + self.flat_args = self._flat_args(args) + else: + self.args = () + self.flat_args = True + if kwargs is not None: + self.kwargs = kwargs + self.flat_kwargs = self._flat_args(kwargs.values()) + else: + self.kwargs = {} + self.flat_kwargs = True @classmethod def _flat_args(cls, args: Iterable): @@ -134,7 +141,7 @@ def _flat_args(cls, args: Iterable): def exec( self, - ) -> Tuple[ObjectRefOrListType, Union["MetaList", List], Union[int, List[int]]]: + ) -> Tuple[ObjectRefOrListType, "MetaList", Union[int, List[int]]]: """ Execute this task, if required. @@ -150,7 +157,7 @@ def exec( return self.data, self.meta, self.meta_offset if ( - not isinstance(self.data, DeferredExecution) + self.flat_data and self.flat_args and self.flat_kwargs and self.num_returns == 1 @@ -166,6 +173,7 @@ def exec( # it back. After the execution, the result is saved and the counter has no effect. self.subscribers += 2 consumers, output = self._deconstruct() + # The last result is the MetaList, so adding +1 here. num_returns = sum(c.num_returns for c in consumers) + 1 results = self._remote_exec_chain(num_returns, *output) @@ -173,12 +181,13 @@ def exec( meta_offset = 0 results = iter(results) for de in consumers: - if de.num_returns == 1: + num_returns = de.num_returns + if num_returns == 1: de._set_result(next(results), meta, meta_offset) meta_offset += 2 else: res = list(islice(results, num_returns)) - offsets = list(range(0, 2 * num_returns, 2)) + offsets = list(range(meta_offset, meta_offset + 2 * num_returns, 2)) de._set_result(res, meta, offsets) meta_offset += 2 * num_returns return self.data, self.meta, self.meta_offset @@ -303,7 +312,9 @@ def _deconstruct_chain( out_extend = output.extend while True: de.unsubscribe() - if (out_pos := getattr(de, "out_pos", None)) and not de.has_result: + if not (has_result := de.has_result) and ( + out_pos := getattr(de, "out_pos", None) + ): out_append(_Tag.REF) out_append(out_pos) output[out_pos] = out_pos @@ -318,12 +329,13 @@ def _deconstruct_chain( break elif not isinstance(data := de.data, DeferredExecution): if isinstance(data, ListOrTuple): + out_append(_Tag.LIST) yield cls._deconstruct_list( data, output, stack, result_consumers, out_append ) else: out_append(data) - if not de.has_result: + if not has_result: stack.append(de) break else: @@ -391,22 +403,24 @@ def _deconstruct_list( """ for obj in lst: if isinstance(obj, DeferredExecution): - if out_pos := getattr(obj, "out_pos", None): + if obj.has_result: + obj = obj.data + elif out_pos := getattr(obj, "out_pos", None): obj.unsubscribe() - if obj.has_result: - out_append(obj.data) - else: - out_append(_Tag.REF) - out_append(out_pos) - output[out_pos] = out_pos - if obj.subscribers == 0: - output[out_pos + 1] = 0 - result_consumers.remove(obj) + out_append(_Tag.REF) + out_append(out_pos) + output[out_pos] = out_pos + if obj.subscribers == 0: + output[out_pos + 1] = 0 + result_consumers.remove(obj) + continue else: out_append(_Tag.CHAIN) yield cls._deconstruct_chain(obj, output, stack, result_consumers) out_append(_Tag.END) - elif isinstance(obj, ListOrTuple): + continue + + if isinstance(obj, ListOrTuple): out_append(_Tag.LIST) yield cls._deconstruct_list( obj, output, stack, result_consumers, out_append @@ -432,13 +446,13 @@ def _remote_exec_chain(num_returns: int, *args: Tuple) -> List[Any]: list The execution results. The last element of this list is the ``MetaList``. """ - # Prefer _remote_exec_single_chain(). It has fewer arguments and - # does not require the num_returns to be specified in options. + # Prefer _remote_exec_single_chain(). It does not require the num_returns + # to be specified in options. if num_returns == 2: return _remote_exec_single_chain.remote(*args) else: return _remote_exec_multi_chain.options(num_returns=num_returns).remote( - num_returns, *args + *args ) def _set_result( @@ -456,7 +470,7 @@ def _set_result( meta : MetaList meta_offset : int or list of int """ - del self.func, self.args, self.kwargs, self.flat_args, self.flat_kwargs + del self.func, self.args, self.kwargs self.data = result self.meta = meta self.meta_offset = meta_offset @@ -466,6 +480,64 @@ def __reduce__(self): raise NotImplementedError("DeferredExecution is not serializable!") +ObjectRefOrDeType = Union[ObjectRefType, DeferredExecution] + + +class DeferredGetItem(DeferredExecution): + """ + Deferred execution task that returns an item at the specified index. + + Parameters + ---------- + data : ObjectRefOrDeType + The object to get the item from. + index : int + The item index. + """ + + def __init__(self, data: ObjectRefOrDeType, index: int): + super().__init__(data, self._remote_fn, [index]) + self.index = index + + @property + @_inherit_docstrings(DeferredExecution.has_result) + def has_result(self): + if super().has_result: + return True + + if ( + isinstance(self.data, DeferredExecution) + and self.data.has_result + and self.data.num_returns != 1 + ): + # If `data` is a `DeferredExecution`, that returns multiple results, we + # don't need to execute `_remote_fn`, but can get the result by index instead. + self._set_result( + self.data.data[self.index], + self.data.meta, + self.data.meta_offset[self.index], + ) + return True + + return False + + @remote_function + def _remote_fn(obj, index): # pragma: no cover + """ + Return the item by index. + + Parameters + ---------- + obj : collection + index : int + + Returns + ------- + object + """ + return obj[index] + + class MetaList: """ Meta information, containing the result lengths and the worker address. @@ -478,6 +550,11 @@ class MetaList: def __init__(self, obj: Union[ray.ObjectID, ClientObjectRef, List]): self._obj = obj + def materialize(self): + """Materialized the list, if required.""" + if not isinstance(self._obj, list): + self._obj = RayWrapper.materialize(self._obj) + def __getitem__(self, index): """ Get item at the specified index. @@ -508,7 +585,7 @@ def __setitem__(self, index, value): obj[index] = value -class MetaListHook(MaterializationHook): +class MetaListHook(MaterializationHook, DeferredGetItem): """ Used by MetaList.__getitem__() for lazy materialization and getting a single value from the list. @@ -516,13 +593,13 @@ class MetaListHook(MaterializationHook): ---------- meta : MetaList Non-materialized list to get the value from. - idx : int + index : int The value index in the list. """ - def __init__(self, meta: MetaList, idx: int): + def __init__(self, meta: MetaList, index: int): + super().__init__(meta._obj, index) self.meta = meta - self.idx = idx def pre_materialize(self): """ @@ -533,7 +610,7 @@ def pre_materialize(self): object """ obj = self.meta._obj - return obj[self.idx] if isinstance(obj, list) else obj + return obj[self.index] if isinstance(obj, list) else obj def post_materialize(self, materialized): """ @@ -548,7 +625,7 @@ def post_materialize(self, materialized): object """ self.meta._obj = materialized - return materialized[self.idx] + return materialized[self.index] class _Tag(Enum): # noqa: PR01 @@ -605,7 +682,7 @@ def exec_func(fn: Callable, obj: Any, args: Tuple, kwargs: Dict) -> Any: raise err @classmethod - def construct(cls, num_returns: int, args: Tuple): # pragma: no cover + def construct(cls, args: Tuple): # pragma: no cover """ Construct and execute the specified chain. @@ -615,7 +692,6 @@ def construct(cls, num_returns: int, args: Tuple): # pragma: no cover Parameters ---------- - num_returns : int args : tuple Yields @@ -687,7 +763,7 @@ def construct_chain( while chain: fn = pop() - if fn == tg_e: + if fn is tg_e: lst.append(obj) break @@ -717,10 +793,10 @@ def construct_chain( itr = iter([obj] if num_returns == 1 else obj) for _ in range(num_returns): - obj = next(itr) - meta.append(len(obj) if hasattr(obj, "__len__") else 0) - meta.append(len(obj.columns) if hasattr(obj, "columns") else 0) - yield obj + o = next(itr) + meta.append(len(o) if hasattr(o, "__len__") else 0) + meta.append(len(o.columns) if hasattr(o, "columns") else 0) + yield o @classmethod def construct_list( @@ -834,20 +910,18 @@ def _remote_exec_single_chain( ------- Generator """ - return remote_executor.construct(num_returns=2, args=args) + return remote_executor.construct(args=args) @ray.remote def _remote_exec_multi_chain( - num_returns: int, *args: Tuple, remote_executor=_REMOTE_EXEC + *args: Tuple, remote_executor=_REMOTE_EXEC ) -> Generator: # pragma: no cover """ Execute the deconstructed chain with a multiple return values in a worker process. Parameters ---------- - num_returns : int - The number of return values. *args : tuple A deconstructed chain to be executed. remote_executor : _RemoteExecutor, default: _REMOTE_EXEC @@ -857,4 +931,4 @@ def _remote_exec_multi_chain( ------- Generator """ - return remote_executor.construct(num_returns, args) + return remote_executor.construct(args) diff --git a/modin/core/execution/ray/common/engine_wrapper.py b/modin/core/execution/ray/common/engine_wrapper.py index fc5f8a643d2..178b25e1d03 100644 --- a/modin/core/execution/ray/common/engine_wrapper.py +++ b/modin/core/execution/ray/common/engine_wrapper.py @@ -20,7 +20,7 @@ import asyncio import os from types import FunctionType -from typing import Sequence +from typing import Iterable, Sequence import ray from ray.util.client.common import ClientObjectRef @@ -106,7 +106,7 @@ def materialize(cls, obj_id): Parameters ---------- - obj_id : ray.ObjectID + obj_id : ObjectRefTypes Ray object identifier to get the value by. Returns @@ -214,7 +214,7 @@ def wait(cls, obj_ids, num_returns=None): num_returns : int, optional """ if not isinstance(obj_ids, Sequence): - obj_ids = list(obj_ids) + obj_ids = list(obj_ids) if isinstance(obj_ids, Iterable) else [obj_ids] ids = set() for obj in obj_ids: diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py index a4c35bf7e95..9c5e3a43c8f 100644 --- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py +++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py @@ -148,7 +148,7 @@ def add_to_apply_calls( def drain_call_queue(self): data = self._data_ref if not isinstance(data, DeferredExecution): - return data + return log = get_logger() self._is_debug(log) and log.debug( @@ -419,7 +419,7 @@ def eager_exec(self, func, *args, length=None, width=None, **kwargs): LazyExecution.subscribe(_configure_lazy_exec) -class SlicerHook(MaterializationHook): +class SlicerHook(MaterializationHook, DeferredExecution): """ Used by mask() for the slilced length computation. @@ -432,6 +432,7 @@ class SlicerHook(MaterializationHook): """ def __init__(self, ref: ObjectIDType, slc: slice): + super().__init__(slc, compute_sliced_len, [ref]) self.ref = ref self.slc = slc diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py index 0748efb219a..dfe782dc50e 100644 --- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py +++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py @@ -12,16 +12,20 @@ # governing permissions and limitations under the License. """Module houses class that implements ``GenericRayDataframePartitionManager`` using Ray.""" +import math import numpy as np +import pandas from pandas.core.dtypes.common import is_numeric_dtype -from modin.config import AsyncReadMode +from modin.config import AsyncReadMode, MinPartitionSize from modin.core.execution.modin_aqp import progress_bar_wrapper from modin.core.execution.ray.common import RayWrapper +from modin.core.execution.ray.common.deferred_execution import DeferredExecution from modin.core.execution.ray.generic.partitioning import ( GenericRayDataframePartitionManager, ) +from modin.core.execution.utils import remote_function from modin.logging import get_logger from modin.utils import _inherit_docstrings @@ -29,6 +33,7 @@ from .virtual_partition import ( PandasOnRayDataframeColumnPartition, PandasOnRayDataframeRowPartition, + PandasOnRayDataframeVirtualPartition, ) @@ -42,6 +47,74 @@ class PandasOnRayDataframePartitionManager(GenericRayDataframePartitionManager): _execution_wrapper = RayWrapper materialize_futures = RayWrapper.materialize + @remote_function + def _remote_get_index(*dfs): # pragma: no cover # noqa: GL08 + return [df.index for df in dfs] + + @remote_function + def _remote_get_columns(*dfs): # pragma: no cover # noqa: GL08 + return [df.columns for df in dfs] + + @classmethod + @_inherit_docstrings(GenericRayDataframePartitionManager.get_indices) + def get_indices(cls, axis, partitions, index_func=None): + partitions = partitions.T if axis == 0 else partitions + if len(partitions): + partitions = [part for part in partitions[0]] + num_partitions = len(partitions) + if num_partitions == 0: + return pandas.Index([]), [] + + non_split, lengths, _ = ( + PandasOnRayDataframeVirtualPartition.find_non_split_block(partitions) + ) + if non_split is not None: + partitions = [non_split] + if lengths is None: + lengths = [] + else: + partitions = [part._data for part in partitions] + + if index_func is None: + remote_fn = cls._remote_get_index if axis == 0 else cls._remote_get_columns + data, args = partitions[0], partitions[1:] + else: + + @remote_function + def remote_fn(fn, *dfs): # pragma: no cover + return [fn(df) for df in dfs] + + data, args = index_func, partitions + + de = DeferredExecution(data, remote_fn, args, num_returns=len(partitions)) + part_indices = de.exec()[0] + + if non_split is not None: + materialized = RayWrapper.materialize([part_indices] + lengths) + idx = materialized[0][0] + lengths = materialized[1:] + idx_len = len(idx) + + if idx_len != sum(lengths): + chunk_len = max( + math.ceil(idx_len / num_partitions), MinPartitionSize.get() + ) + lengths = [chunk_len] * num_partitions + + part_indices = [] + start = 0 + for length in lengths: + end = start + length + part_indices.append(idx[start:end]) + start = end + return idx, part_indices + + part_indices = RayWrapper.materialize(part_indices) + indices = [idx for idx in part_indices if len(idx)] + if len(indices) == 0: + return part_indices[0], part_indices + return indices[0].append(indices[1:]), part_indices + @classmethod def wait_partitions(cls, partitions): """ diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py index f30b474305a..bb263b51831 100644 --- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py +++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py @@ -12,326 +12,471 @@ # governing permissions and limitations under the License. """Module houses classes responsible for storing a virtual partition and applying a function to it.""" +import math +from typing import Collection, Iterable, List, Optional, Set, Tuple, Union import pandas import ray -from ray.util import get_node_ip_address +from modin.config import MinPartitionSize +from modin.core.dataframe.base.partitioning.axis_partition import ( + BaseDataframeAxisPartition, +) from modin.core.dataframe.pandas.partitioning.axis_partition import ( PandasDataframeAxisPartition, ) from modin.core.execution.ray.common import RayWrapper +from modin.core.execution.ray.common.deferred_execution import ( + DeferredExecution, + DeferredGetItem, + MetaList, + ObjectRefOrDeType, + ObjectRefType, +) +from modin.core.execution.utils import remote_function from modin.utils import _inherit_docstrings from .partition import PandasOnRayDataframePartition -class PandasOnRayDataframeVirtualPartition(PandasDataframeAxisPartition): +class PandasOnRayDataframeVirtualPartition(BaseDataframeAxisPartition): """ The class implements the interface in ``PandasDataframeAxisPartition``. Parameters ---------- - list_of_partitions : Union[list, PandasOnRayDataframePartition] - List of ``PandasOnRayDataframePartition`` and - ``PandasOnRayDataframeVirtualPartition`` objects, or a single - ``PandasOnRayDataframePartition``. - get_ip : bool, default: False - Whether to get node IP addresses to conforming partitions or not. + data : DeferredExecution or list of PandasOnRayDataframePartition full_axis : bool, default: True Whether or not the virtual partition encompasses the whole axis. - call_queue : list, optional - A list of tuples (callable, args, kwargs) that contains deferred calls. length : ray.ObjectRef or int, optional Length, or reference to length, of wrapped ``pandas.DataFrame``. width : ray.ObjectRef or int, optional Width, or reference to width, of wrapped ``pandas.DataFrame``. + num_splits : int, optional + The number of chunks to split the results on. + chunk_lengths : list of ints, optional + The chunk lengths. """ - _PARTITIONS_METADATA_LEN = 3 # (length, width, ip) partition_type = PandasOnRayDataframePartition instance_type = ray.ObjectRef axis = None - # these variables are intentionally initialized at runtime (see #6023) - _DEPLOY_AXIS_FUNC = None - _DEPLOY_SPLIT_FUNC = None - _DRAIN_FUNC = None + def __init__( + self, + data: Union[ + DeferredExecution, + PandasOnRayDataframePartition, + List[PandasOnRayDataframePartition], + ], + full_axis: bool = True, + length: Union[int, ObjectRefType] = None, + width: Union[int, ObjectRefType] = None, + num_splits=None, + chunk_lengths=None, + ): + self.full_axis = full_axis + self._meta = MetaList([length, width, None]) + self._meta_offset = 0 + self._chunk_lengths_cache = chunk_lengths + + if isinstance(data, DeferredExecution): + self._set_data_ref(data) + self._num_splits = num_splits + self._list_of_block_partitions = None + return + + if not isinstance(data, Collection) or len(data) == 1: + if not isinstance(data, Collection): + data = [data] + self._set_data_ref(data[0]._data) + self._num_splits = 1 + self._list_of_block_partitions = data + return + + self._num_splits = len(data) + self._list_of_block_partitions = data + + non_split, lengths, full_concat = self.find_non_split_block(data) + if non_split is not None: + if lengths is not None: + self._chunk_lengths_cache = lengths + if full_concat: + self._set_data_ref(non_split) + return + # TODO: We have a subset of the same frame here and can just get a single chunk + # from the original frame instead of concatenating all these chunks. + + data = DeferredExecution([part._data for part in data], self._remote_concat) + self._set_data_ref(data) + + @staticmethod + def find_non_split_block( + partitions: Iterable[PandasOnRayDataframePartition], + ) -> Tuple[ + Union[ObjectRefOrDeType, None], + Union[List[Union[ObjectRefType, int]], None], + bool, + ]: + """ + Find a non-split block if there is one. + + If all the partitions are the sequential chunks of the same DataFrame, + the concatenation of all these chunks will give an identical DF. Thus, we + don't need to concatenate but can get the original one instead. - @classmethod - def _get_deploy_axis_func(cls): # noqa: GL08 - if cls._DEPLOY_AXIS_FUNC is None: - cls._DEPLOY_AXIS_FUNC = RayWrapper.put( - PandasDataframeAxisPartition.deploy_axis_func + Parameters + ---------- + partitions : list of PandasOnRayDataframePartition + + Returns + ------- + ObjectRefOrDeType or None + The non-split block or None if not found. + list of lengths or None + Estimated chunk lengths, that could be different form the real ones. + bool + Whether the specified partitions represent the full block or just the + first part of this block. + """ + refs = [part._data_ref for part in partitions] + if ( + isinstance(refs[0], _DeferredGetChunk) + and isinstance(split := refs[0].data, _DeferredSplit) + and (refs[0].index == 0) + and all(prev.is_next_chunk(next) for prev, next in zip(refs[:-1], refs[1:])) + ): + lengths = [ref.length for ref in refs if ref.length is not None] + return ( + split.data, + lengths if len(lengths) == len(refs) else None, + split.num_splits == refs[-1].index + 1, ) - return cls._DEPLOY_AXIS_FUNC + return None, None, False + + def _set_data_ref( + self, data: Union[DeferredExecution, ObjectRefType] + ): # noqa: GL08 + if isinstance(data, DeferredExecution): + data.subscribe() + self._data_ref = data + + def __del__(self): + """Unsubscribe from DeferredExecution.""" + if isinstance(self._data_ref, DeferredExecution): + self._data_ref.unsubscribe() + + @_inherit_docstrings(BaseDataframeAxisPartition.apply) + def apply( + self, + func, + *args, + num_splits=None, + other_axis_partition=None, + maintain_partitioning=True, + lengths=None, + manual_partition=False, + **kwargs, + ) -> Union[List[PandasOnRayDataframePartition], PandasOnRayDataframePartition]: + if not manual_partition: + if not self.full_axis: + # If this is not a full axis partition, it already contains a subset of + # the full axis, so we shouldn't split the result further. + num_splits = 1 + elif num_splits is None: + num_splits = self._num_splits + + if ( + num_splits == 1 + or not maintain_partitioning + or num_splits != self._num_splits + ): + lengths = None + elif lengths is None: + lengths = self._chunk_lengths + + if other_axis_partition is not None: + if isinstance(other_axis_partition, Collection): + if len(other_axis_partition) == 1: + other_part = other_axis_partition[0]._data + else: + concat_fn = ( + PandasOnRayDataframeColumnPartition + if self.axis + else PandasOnRayDataframeRowPartition + )._concat + other_part = concat_fn([p._data for p in other_axis_partition]) + else: + other_part = other_axis_partition._data + args = [other_part] + list(args) + + de = self._apply(func, args, kwargs) + if num_splits > 1: + de = _DeferredSplit(de, self.axis, num_splits, lengths) + if lengths is not None and len(lengths) != num_splits: + lengths = None + result = [ + PandasOnRayDataframePartition( + _DeferredGetChunk( + de, i, lengths[i] if lengths is not None else None + ) + ) + for i in range(num_splits) + ] + else: + result = [PandasOnRayDataframePartition(de)] + if self.full_axis or other_axis_partition is not None: + return result + else: + # If this is not a full axis partition, just take out the single split in the result. + return result[0] + + @_inherit_docstrings(PandasDataframeAxisPartition.add_to_apply_calls) + def add_to_apply_calls(self, func, *args, length=None, width=None, **kwargs): + de = self._apply(func, args, kwargs) + return type(self)( + de, self.full_axis, length, width, self._num_splits, self._chunk_lengths + ) + + @_inherit_docstrings(PandasDataframeAxisPartition.split) + def split( + self, split_func, num_splits, f_args=None, f_kwargs=None, extract_metadata=False + ) -> List[PandasOnRayDataframePartition]: + chunks, meta, offsets = DeferredExecution( + self._data_ref, + split_func, + args=f_args, + kwargs=f_kwargs, + num_returns=num_splits, + ).exec() + return [ + PandasOnRayDataframePartition(chunks[i], meta=meta, meta_offset=offsets[i]) + for i in range(num_splits) + ] - @classmethod - def _get_deploy_split_func(cls): # noqa: GL08 - if cls._DEPLOY_SPLIT_FUNC is None: - cls._DEPLOY_SPLIT_FUNC = RayWrapper.put( - PandasDataframeAxisPartition.deploy_splitting_func + @property + def _length_cache(self): # noqa: GL08 + return self._meta[self._meta_offset] + + def length(self, materialize=True): # noqa: GL08 + if self._length_cache is None: + self._calculate_lengths(materialize) + elif materialize: + self._meta.materialize() + return self._length_cache + + @property + def _width_cache(self): # noqa: GL08 + return self._meta[self._meta_offset + 1] + + def width(self, materialize=True): # noqa: GL08 + if self._width_cache is None: + self._calculate_lengths(materialize) + elif materialize: + self._meta.materialize() + return self._width_cache + + def _calculate_lengths(self, materialize=True): # noqa: GL08 + if self._list_of_block_partitions is not None: + from . import PandasOnRayDataframePartitionManager + + lengths = [part.length(False) for part in self._list_of_block_partitions] + widths = [part.width(False) for part in self._list_of_block_partitions] + materialized = PandasOnRayDataframePartitionManager.materialize_futures( + lengths + widths ) - return cls._DEPLOY_SPLIT_FUNC + self._meta[self._meta_offset] = sum(materialized[: len(lengths)]) + self._meta[self._meta_offset + 1] = sum(materialized[len(lengths) :]) + else: + self.force_materialization() + if materialize: + self._meta.materialize() + + @_inherit_docstrings(PandasDataframeAxisPartition.drain_call_queue) + def drain_call_queue(self, num_splits=None): + if num_splits: + self._num_splits = num_splits + + @_inherit_docstrings(PandasDataframeAxisPartition.force_materialization) + def force_materialization(self, get_ip=False): + self._data # Trigger execution + self._num_splits = 1 + self._chunk_lengths_cache = None + self._list_of_block_partitions = None + return self + + @_inherit_docstrings(PandasDataframeAxisPartition.wait) + def wait(self): + """Wait completing computations on the object wrapped by the partition.""" + RayWrapper.wait(self._data) + + @_inherit_docstrings(PandasDataframeAxisPartition.to_pandas) + def to_pandas(self): + return RayWrapper.materialize(self._data) - @classmethod - def _get_drain_func(cls): # noqa: GL08 - if cls._DRAIN_FUNC is None: - cls._DRAIN_FUNC = RayWrapper.put(PandasDataframeAxisPartition.drain) - return cls._DRAIN_FUNC + @_inherit_docstrings(PandasDataframeAxisPartition.to_numpy) + def to_numpy(self): + return self.to_pandas().to_numpy() + + @_inherit_docstrings(PandasDataframeAxisPartition.mask) + def mask(self, row_indices, col_indices): + part = PandasOnRayDataframePartition(self._data_ref).mask( + row_indices, col_indices + ) + return type(self)(part, False) + + @property + @_inherit_docstrings(BaseDataframeAxisPartition.list_of_blocks) + def list_of_blocks(self): + return [part._data for part in self.list_of_block_partitions] + + @property + @_inherit_docstrings(PandasDataframeAxisPartition.list_of_block_partitions) + def list_of_block_partitions(self) -> list: + if self._list_of_block_partitions is not None: + return self._list_of_block_partitions + + data = self._data_ref + num_splits = self._num_splits + if num_splits > 1: + lengths = self._chunk_lengths + data = _DeferredSplit(data, self.axis, num_splits, lengths) + if lengths is not None and len(lengths) != num_splits: + lengths = None + self._list_of_block_partitions = [ + PandasOnRayDataframePartition( + _DeferredGetChunk( + data, i, lengths[i] if lengths is not None else None + ) + ) + for i in range(num_splits) + ] + else: + self._list_of_block_partitions = [PandasOnRayDataframePartition(data)] + return self._list_of_block_partitions @property def list_of_ips(self): """ - Get the IPs holding the physical objects composing this partition. + Return the list of IP worker addresses. Returns ------- - List - A list of IPs as ``ray.ObjectRef`` or str. + list of str """ - # Defer draining call queue until we get the ip address - result = [None] * len(self.list_of_block_partitions) - for idx, partition in enumerate(self.list_of_block_partitions): - partition.drain_call_queue() - result[idx] = partition.ip(materialize=False) - return result + if (ip := self._meta[self._meta_offset + 2]) is not None: + return [ip] + if self._list_of_block_partitions is not None: + return [part.ip() for part in self._list_of_block_partitions] + return [] - @classmethod - @_inherit_docstrings(PandasDataframeAxisPartition.deploy_splitting_func) - def deploy_splitting_func( - cls, - axis, - func, - f_args, - f_kwargs, - num_splits, - *partitions, - extract_metadata=False, - ): - return _deploy_ray_func.options( - num_returns=( - num_splits * (1 + cls._PARTITIONS_METADATA_LEN) - if extract_metadata - else num_splits - ), - ).remote( - cls._get_deploy_split_func(), - *f_args, - num_splits, - *partitions, - axis=axis, - f_to_deploy=func, - f_len_args=len(f_args), - f_kwargs=f_kwargs, - extract_metadata=extract_metadata, - ) + @property + def _data(self): # noqa: GL08 + data = self._data_ref + if isinstance(data, DeferredExecution): + data, self._meta, self._meta_offset = data.exec() + self._data_ref = data + return data - @classmethod - def deploy_axis_func( - cls, - axis, - func, - f_args, - f_kwargs, - num_splits, - maintain_partitioning, - *partitions, - lengths=None, - manual_partition=False, - max_retries=None, + @property + def _chunk_lengths(self): # noqa: GL08 + if ( + self._chunk_lengths_cache is None + and self._list_of_block_partitions is not None + ): + attr = "length" if self.axis == 0 else "width" + self._chunk_lengths_cache = [ + getattr(p, attr)(materialize=False) + for p in self._list_of_block_partitions + ] + return self._chunk_lengths_cache + + def _apply(self, apply_fn, args, kwargs) -> DeferredExecution: # noqa: GL08 + return DeferredExecution(self._data_ref, apply_fn, args, kwargs) + + +class _DeferredSplit(DeferredExecution): # noqa: GL08 + def __init__( + self, + obj: ObjectRefOrDeType, + axis: int, + num_splits: int, + lengths: Union[List[Union[ObjectRefType, int]], None], ): - """ - Deploy a function along a full axis. + self.num_splits = num_splits + self.skip_chunks = set() + args = [axis, num_splits, MinPartitionSize.get(), self.skip_chunks] + if lengths and (len(lengths) == num_splits): + args.extend(lengths) + super().__init__(obj, self._split, args, num_returns=num_splits) + + @remote_function + def _split( + df: pandas.DataFrame, + axis: int, + num_splits: int, + min_chunk_len: int, + skip_chunks: Set[int], + *lengths: Optional[List[int]], + ): # pragma: no cover # noqa: GL08 + if not lengths or (sum(lengths) != df.shape[axis]): + length = df.shape[axis] + chunk_len = max(math.ceil(length / num_splits), min_chunk_len) + lengths = [chunk_len] * num_splits + + result = [] + start = 0 + for i in range(num_splits): + if i in skip_chunks: + result.append(None) + start += lengths[i] + continue + + end = start + lengths[i] + chunk = df.iloc[start:end] if axis == 0 else df.iloc[:, start:end] + start = end + result.append(chunk) + if isinstance(chunk.axes[axis], pandas.MultiIndex): + chunk.set_axis( + chunk.axes[axis].remove_unused_levels(), + axis=axis, + copy=False, + ) - Parameters - ---------- - axis : {0, 1} - The axis to perform the function along. - func : callable - The function to perform. - f_args : list or tuple - Positional arguments to pass to ``func``. - f_kwargs : dict - Keyword arguments to pass to ``func``. - num_splits : int - The number of splits to return (see ``split_result_of_axis_func_pandas``). - maintain_partitioning : bool - If True, keep the old partitioning if possible. - If False, create a new partition layout. - *partitions : iterable - All partitions that make up the full axis (row or column). - lengths : list, optional - The list of lengths to shuffle the object. - manual_partition : bool, default: False - If True, partition the result with `lengths`. - max_retries : int, default: None - The max number of times to retry the func. + return result - Returns - ------- - list - A list of ``ray.ObjectRef``-s. - """ - return _deploy_ray_func.options( - num_returns=(num_splits if lengths is None else len(lengths)) - * (1 + cls._PARTITIONS_METADATA_LEN), - **({"max_retries": max_retries} if max_retries is not None else {}), - ).remote( - cls._get_deploy_axis_func(), - *f_args, - num_splits, - maintain_partitioning, - *partitions, - axis=axis, - f_to_deploy=func, - f_len_args=len(f_args), - f_kwargs=f_kwargs, - manual_partition=manual_partition, - lengths=lengths, - return_generator=True, - ) - @classmethod - def deploy_func_between_two_axis_partitions( - cls, - axis, - func, - f_args, - f_kwargs, - num_splits, - len_of_left, - other_shape, - *partitions, - ): - """ - Deploy a function along a full axis between two data sets. +class _DeferredGetChunk(DeferredGetItem): # noqa: GL08 + def __init__(self, split: _DeferredSplit, index: int, length: Optional[int] = None): + super().__init__(split, index) + self.length = length - Parameters - ---------- - axis : {0, 1} - The axis to perform the function along. - func : callable - The function to perform. - f_args : list or tuple - Positional arguments to pass to ``func``. - f_kwargs : dict - Keyword arguments to pass to ``func``. - num_splits : int - The number of splits to return (see ``split_result_of_axis_func_pandas``). - len_of_left : int - The number of values in `partitions` that belong to the left data set. - other_shape : np.ndarray - The shape of right frame in terms of partitions, i.e. - (other_shape[i-1], other_shape[i]) will indicate slice to restore i-1 axis partition. - *partitions : iterable - All partitions that make up the full axis (row or column) for both data sets. + def __del__(self): + """Remove this chunk from _DeferredSplit if it's not executed yet.""" + if isinstance(self.data, _DeferredSplit): + self.data.skip_chunks.add(self.index) - Returns - ------- - list - A list of ``ray.ObjectRef``-s. - """ - return _deploy_ray_func.options( - num_returns=num_splits * (1 + cls._PARTITIONS_METADATA_LEN) - ).remote( - PandasDataframeAxisPartition.deploy_func_between_two_axis_partitions, - *f_args, - num_splits, - len_of_left, - other_shape, - *partitions, - axis=axis, - f_to_deploy=func, - f_len_args=len(f_args), - f_kwargs=f_kwargs, - return_generator=True, + def is_next_chunk(self, other): # noqa: GL08 + return ( + isinstance(other, _DeferredGetChunk) + and (self.data is other.data) + and (other.index == self.index + 1) ) - def wait(self): - """Wait completing computations on the object wrapped by the partition.""" - self.drain_call_queue() - futures = self.list_of_blocks - RayWrapper.wait(futures) - @_inherit_docstrings(PandasOnRayDataframeVirtualPartition.__init__) class PandasOnRayDataframeColumnPartition(PandasOnRayDataframeVirtualPartition): axis = 0 + @remote_function + def _remote_concat(dfs): # pragma: no cover # noqa: GL08 + return pandas.concat(dfs, axis=0, copy=False) + @_inherit_docstrings(PandasOnRayDataframeVirtualPartition.__init__) class PandasOnRayDataframeRowPartition(PandasOnRayDataframeVirtualPartition): axis = 1 - -@ray.remote -def _deploy_ray_func( - deployer, - *positional_args, - axis, - f_to_deploy, - f_len_args, - f_kwargs, - extract_metadata=True, - **kwargs, -): # pragma: no cover - """ - Execute a function on an axis partition in a worker process. - - This is ALWAYS called on either ``PandasDataframeAxisPartition.deploy_axis_func`` - or ``PandasDataframeAxisPartition.deploy_func_between_two_axis_partitions``, which both - serve to deploy another dataframe function on a Ray worker process. The provided `positional_args` - contains positional arguments for both: `deployer` and for `f_to_deploy`, the parameters can be separated - using the `f_len_args` value. The parameters are combined so they will be deserialized by Ray before the - kernel is executed (`f_kwargs` will never contain more Ray objects, and thus does not require deserialization). - - Parameters - ---------- - deployer : callable - A `PandasDataFrameAxisPartition.deploy_*` method that will call ``f_to_deploy``. - *positional_args : list - The first `f_len_args` elements in this list represent positional arguments - to pass to the `f_to_deploy`. The rest are positional arguments that will be - passed to `deployer`. - axis : {0, 1} - The axis to perform the function along. This argument is keyword only. - f_to_deploy : callable or RayObjectID - The function to deploy. This argument is keyword only. - f_len_args : int - Number of positional arguments to pass to ``f_to_deploy``. This argument is keyword only. - f_kwargs : dict - Keyword arguments to pass to ``f_to_deploy``. This argument is keyword only. - extract_metadata : bool, default: True - Whether to return metadata (length, width, ip) of the result. Passing `False` may relax - the load on object storage as the remote function would return 4 times fewer futures. - Passing `False` makes sense for temporary results where you know for sure that the - metadata will never be requested. This argument is keyword only. - **kwargs : dict - Keyword arguments to pass to ``deployer``. - - Returns - ------- - list : Union[tuple, list] - The result of the function call, and metadata for it. - - Notes - ----- - Ray functions are not detected by codecov (thus pragma: no cover). - """ - 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: - for item in result: - yield item - else: - ip = get_node_ip_address() - for r in result: - if isinstance(r, pandas.DataFrame): - for item in [r, len(r), len(r.columns), ip]: - yield item - else: - for item in [r, None, None, ip]: - yield item + @remote_function + def _remote_concat(dfs): # pragma: no cover # noqa: GL08 + return pandas.concat(dfs, axis=1, copy=False) diff --git a/modin/core/execution/utils.py b/modin/core/execution/utils.py index a26525234f4..269621b6cb3 100644 --- a/modin/core/execution/utils.py +++ b/modin/core/execution/utils.py @@ -33,43 +33,60 @@ def set_env(**environ): os.environ.update(old_environ) +if "_MODIN_DOC_CHECKER_" in os.environ: + + # The doc checker should get the non-processed functions + def remote_function(func, ignore_defaults=False): + return func + + # Check if the function already exists to avoid circular imports -if "remote_function" not in dir(): +elif "remote_function" not in dir(): from modin.config import Engine if Engine.get() == "Ray": from modin.core.execution.ray.common import RayWrapper - _remote_function_wrapper = RayWrapper.put + _preprocess_func = RayWrapper.put elif Engine.get() == "Unidist": from modin.core.execution.unidist.common import UnidistWrapper - _remote_function_wrapper = UnidistWrapper.put + _preprocess_func = UnidistWrapper.put elif Engine.get() == "Dask": from modin.core.execution.dask.common import DaskWrapper # The function cache is not supported for Dask - def remote_function(func): + def remote_function(func, ignore_defaults=False): return DaskWrapper.put(func) else: - def remote_function(func): + def remote_function(func, ignore_defaults=False): return func if "remote_function" not in dir(): _remote_function_cache = {} - def remote_function(func): # noqa: F811 - if func.__closure__: - ErrorMessage.warn( - f"The remote function {func} can not be cached, because " - + "it captures objects from the outer scope." - ) - return func - func_id = id(func.__code__) - ref = _remote_function_cache.get(func_id, None) + def remote_function(func, ignore_defaults=False): # noqa: F811 + if "" in func.__qualname__: # Nested function + if func.__closure__: + ErrorMessage.single_warning( + f"The nested function {func} can not be cached, because " + + "it captures objects from the outer scope." + ) + return func + if not ignore_defaults and func.__defaults__: + ErrorMessage.single_warning( + f"The nested function {func} can not be cached, because it has " + + "default values. Use `ignore_defaults` to forcibly enable caching." + ) + return func + # For the nested functions, use __code__ as the key + key = id(func.__code__) + else: + key = func + ref = _remote_function_cache.get(key, None) if ref is None: - ref = _remote_function_wrapper(func) - _remote_function_cache[func_id] = ref + ref = _preprocess_func(func) + _remote_function_cache[key] = ref return ref diff --git a/modin/experimental/batch/pipeline.py b/modin/experimental/batch/pipeline.py index d7cd60cae47..b17fb1d7fd1 100644 --- a/modin/experimental/batch/pipeline.py +++ b/modin/experimental/batch/pipeline.py @@ -242,13 +242,14 @@ def _complete_nodes(self, list_of_nodes, partitions): ) new_dfs[-1].drain_call_queue(num_splits=1) - def reducer(df): + def reducer(df, reduce_fn, *parts): df_inputs = [df] - for df in new_dfs: - df_inputs.append(df.to_pandas()) - return node.reduce_fn(df_inputs) + df_inputs.extend(parts) + return reduce_fn(df_inputs) - partitions = [partitions[0].add_to_apply_calls(reducer)] + args = [node.reduce_fn] + args.extend([dfs._data_ref for dfs in new_dfs]) + partitions = [partitions[0].add_to_apply_calls(reducer, *args)] elif node.repartition_after: if len(partitions) > 1: ErrorMessage.not_implemented( diff --git a/modin/pandas/test/dataframe/test_binary.py b/modin/pandas/test/dataframe/test_binary.py index d31840b0446..241ba8a4501 100644 --- a/modin/pandas/test/dataframe/test_binary.py +++ b/modin/pandas/test/dataframe/test_binary.py @@ -18,8 +18,8 @@ import modin.pandas as pd from modin.config import Engine, NPartitions, StorageFormat -from modin.core.dataframe.pandas.partitioning.axis_partition import ( - PandasDataframeAxisPartition, +from modin.core.dataframe.base.partitioning.axis_partition import ( + BaseDataframeAxisPartition, ) from modin.pandas.test.utils import ( CustomIntegerForAddition, @@ -228,7 +228,7 @@ def modin_df(is_virtual): # Modin should rebalance the partitions after the concat, producing virtual partitions. assert isinstance( result._query_compiler._modin_frame._partitions[0][0], - PandasDataframeAxisPartition, + BaseDataframeAxisPartition, ) return result diff --git a/modin/pandas/test/test_groupby.py b/modin/pandas/test/test_groupby.py index a381eb3cb7c..e0395aa77d1 100644 --- a/modin/pandas/test/test_groupby.py +++ b/modin/pandas/test/test_groupby.py @@ -28,8 +28,8 @@ StorageFormat, ) from modin.core.dataframe.algebra.default2pandas.groupby import GroupBy -from modin.core.dataframe.pandas.partitioning.axis_partition import ( - PandasDataframeAxisPartition, +from modin.core.dataframe.base.partitioning.axis_partition import ( + BaseDataframeAxisPartition, ) from modin.pandas.io import from_pandas from modin.pandas.utils import is_scalar @@ -264,6 +264,9 @@ def test_mixed_dtypes_groupby(as_index): ), ) # FIXME: https://github.com/modin-project/modin/issues/7032 + # Triger execution of deferred operations. If not executed, eval_shift() below fails with + # `could not convert string to float: '\x94'`. Probably, this is also related to #7032. + modin_groupby.shift() eval_general( modin_groupby, pandas_groupby, @@ -2626,7 +2629,7 @@ def test_groupby_with_virtual_partitions(): # Check that the constructed Modin DataFrame has virtual partitions when assert issubclass( type(big_modin_df._query_compiler._modin_frame._partitions[0][0]), - PandasDataframeAxisPartition, + BaseDataframeAxisPartition, ) eval_general( big_modin_df, big_pandas_df, lambda df: df.groupby(df.columns[0]).count() diff --git a/scripts/doc_checker.py b/scripts/doc_checker.py index c85f11a9ed3..82d6491afea 100644 --- a/scripts/doc_checker.py +++ b/scripts/doc_checker.py @@ -35,6 +35,9 @@ from numpydoc.docscrape import NumpyDocString from numpydoc.validate import Docstring +# Let the other modules to know that the doc checker is running. +os.environ["_MODIN_DOC_CHECKER_"] = "1" + # fake cuDF-related modules if they're missing for mod_name in ("cudf", "cupy"): try: