maxframe 0.1.0b5__cp39-cp39-macosx_11_0_arm64.whl

maxframe/core/operator/core.py

@@ -0,0 +1,276 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from typing import Any, Callable, Dict, List, Type
+
+import numpy as np
+
+try:
+    from numpy.core._exceptions import UFuncTypeError
+except ImportError:  # pragma: no cover
+    UFuncTypeError = None
+
+from ...typing_ import ChunkType, OperatorType, TileableType
+from ..entity import (
+    ExecutableTuple,
+    OutputType,
+    get_fetch_class,
+    get_output_types,
+    get_tileable_types,
+)
+from ..mode import is_eager_mode
+
+_op_type_to_executor: Dict[Type[OperatorType], Callable] = dict()
+_op_type_to_size_estimator: Dict[Type[OperatorType], Callable] = dict()
+
+
+class TileableOperatorMixin:
+    __slots__ = ()
+
+    def check_inputs(self, inputs: List[TileableType]):
+        if not inputs:
+            return
+
+        for inp in inputs:
+            if inp is not None and inp._need_execution():
+                raise ValueError(
+                    f"{inp} has unknown dtypes, "
+                    f"it must be executed first before {str(type(self))}"
+                )
+
+    @classmethod
+    def _check_if_gpu(cls, inputs: List[TileableType]):
+        if not inputs:
+            return None
+        true_num = 0
+        for inp in inputs:
+            op = getattr(inp, "op", None)
+            if op is None or op.gpu is None:
+                return None
+            true_num += int(op.gpu)
+        if true_num == len(inputs):
+            return True
+        elif true_num == 0:
+            return False
+        return None
+
+    def _tokenize_output(self, output_idx: int, **kw):
+        return f"{self._key}_{output_idx}"
+
+    @staticmethod
+    def _fill_nan_shape(kw: dict):
+        nsplits = kw.get("nsplits")
+        shape = kw.get("shape")
+        if nsplits is not None and shape is not None:
+            nsplits = tuple(nsplits)
+            shape = list(shape)
+            for idx, (s, sp) in enumerate(zip(shape, nsplits)):
+                if not np.isnan(s):
+                    continue
+                s = sum(sp)
+                if not np.isnan(s):
+                    shape[idx] = s
+            kw["shape"] = tuple(shape)
+            kw["nsplits"] = nsplits
+        return kw
+
+    def _create_tileable(self, output_idx: int, **kw) -> TileableType:
+        output_type = kw.pop("output_type", self._get_output_type(output_idx))
+        if output_type is None:
+            raise ValueError("output_type should be specified")
+
+        if isinstance(output_type, (list, tuple)):
+            output_type = output_type[output_idx]
+
+        tileable_type, tileable_data_type = get_tileable_types(output_type)
+        kw["_i"] = output_idx
+        kw["op"] = self
+        if output_type == OutputType.scalar:
+            # tensor
+            kw["order"] = "C_ORDER"
+
+        kw = self._fill_nan_shape(kw)
+
+        # key of output chunks may only contain keys for its output ids
+        if "_key" not in kw:
+            kw["_key"] = self._tokenize_output(output_idx, **kw)
+
+        data = tileable_data_type(**kw)
+        return tileable_type(data)
+
+    def _new_tileables(
+        self, inputs: List[TileableType], kws: List[dict] = None, **kw
+    ) -> List[TileableType]:
+        output_limit = kw.pop("output_limit", None)
+        if output_limit is None:
+            output_limit = getattr(self, "output_limit")
+
+        self._set_inputs(inputs)
+        if self.gpu is None:
+            self.gpu = self._check_if_gpu(self._inputs)
+        if getattr(self, "_key", None) is None:
+            self._update_key()  # update key when inputs are set
+
+        tileables = []
+        for j in range(output_limit):
+            create_tensor_kw = kw.copy()
+            if kws:
+                create_tensor_kw.update(kws[j])
+            tileable = self._create_tileable(j, **create_tensor_kw)
+            tileables.append(tileable)
+
+        self.outputs = tileables
+        if len(tileables) > 1:
+            # for each output tileable, hold the reference to the other outputs
+            # so that either no one or everyone are gc collected
+            for j, t in enumerate(tileables):
+                t.data._siblings = [
+                    tileable.data for tileable in tileables[:j] + tileables[j + 1 :]
+                ]
+        return tileables
+
+    def new_tileables(
+        self, inputs: List[TileableType], kws: List[dict] = None, **kw
+    ) -> List[TileableType]:
+        """
+        Create tileable objects(Tensors or DataFrames).
+
+        This is a misc function for create tileable objects like tensors or dataframes,
+        it will be called inside the `new_tensors` and `new_dataframes`.
+        If eager mode is on, it will trigger the execution after tileable objects are created.
+
+        Parameters
+        ----------
+        inputs : list
+            Input tileables
+        kws : List[dict]
+            Kwargs for each output.
+        kw : dict
+            Common kwargs for all outputs.
+
+        Returns
+        -------
+        tileables : list
+            Output tileables.
+
+        .. note::
+            It's a final method, do not override.
+            Override the method `_new_tileables` if needed.
+        """
+        tileables = self._new_tileables(inputs, kws=kws, **kw)
+        if is_eager_mode():
+            ExecutableTuple(tileables).execute()
+        return tileables
+
+    def new_tileable(
+        self, inputs: List[TileableType], kws: List[Dict] = None, **kw
+    ) -> TileableType:
+        if getattr(self, "output_limit") != 1:
+            raise TypeError("cannot new chunk with more than 1 outputs")
+
+        return self.new_tileables(inputs, kws=kws, **kw)[0]
+
+    @classmethod
+    def concat_tileable_chunks(cls, tileable: TileableType):
+        raise NotImplementedError
+
+    @classmethod
+    def create_tileable_from_chunks(
+        cls, chunks: List[ChunkType], inputs: List[TileableType] = None, **kw
+    ) -> TileableType:
+        raise NotImplementedError
+
+    def get_fetch_op_cls(self, obj: ChunkType):
+        from .shuffle import ShuffleProxy
+
+        output_types = get_output_types(obj, unknown_as=OutputType.object)
+        fetch_cls, fetch_shuffle_cls = get_fetch_class(output_types[0])
+        if isinstance(self, ShuffleProxy):
+            cls = fetch_shuffle_cls
+        else:
+            cls = fetch_cls
+
+        def _inner(**kw):
+            return cls(output_types=output_types, **kw)
+
+        return _inner
+
+    @classmethod
+    def register_executor(cls, executor: Callable):
+        _op_type_to_executor[cls] = executor
+
+    @classmethod
+    def unregister_executor(cls):
+        del _op_type_to_executor[cls]
+
+    @classmethod
+    def register_size_estimator(cls, size_estimator: Callable):
+        _op_type_to_size_estimator[cls] = size_estimator
+
+    @classmethod
+    def unregister_size_estimator(cls):
+        del _op_type_to_size_estimator[cls]
+
+
+def execute(results: Dict[str, Any], op: OperatorType):
+    try:
+        executor = _op_type_to_executor[type(op)]
+    except KeyError:
+        executor = type(op).execute
+
+    # pre execute
+    op.pre_execute(results, op)
+    succeeded = False
+    try:
+        if UFuncTypeError is None:  # pragma: no cover
+            return executor(results, op)
+        else:
+            # Cast `UFuncTypeError` to `TypeError` since subclasses of the former is unpickleable.
+            # The `UFuncTypeError` was introduced by numpy#12593 since v1.17.0.
+            try:
+                result = executor(results, op)
+                succeeded = True
+                return result
+            except UFuncTypeError as e:  # pragma: no cover
+                raise TypeError(str(e)).with_traceback(sys.exc_info()[2]) from None
+    except NotImplementedError:
+        for op_cls in type(op).__mro__:
+            if op_cls in _op_type_to_executor:
+                executor = _op_type_to_executor[op_cls]
+                _op_type_to_executor[type(op)] = executor
+                result = executor(results, op)
+                succeeded = True
+                return result
+        raise KeyError(f"No handler found for op: {op}")
+    finally:
+        if succeeded:
+            op.post_execute(results, op)
+
+
+def estimate_size(results: Dict[str, Any], op: OperatorType):
+    try:
+        size_estimator = _op_type_to_size_estimator[type(op)]
+    except KeyError:
+        size_estimator = type(op).estimate_size
+
+    try:
+        return size_estimator(results, op)
+    except NotImplementedError:
+        for op_cls in type(op).__mro__:
+            if op_cls in _op_type_to_size_estimator:
+                size_estimator = _op_type_to_size_estimator[op_cls]
+                _op_type_to_size_estimator[type(op)] = size_estimator
+                return size_estimator(results, op)
+        raise KeyError(f"No handler found for op: {op} to estimate size")

maxframe/core/operator/fetch.py

@@ -0,0 +1,53 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import enum
+
+from ... import opcodes
+from ...serialization.serializables import (
+    FieldTypes,
+    Int32Field,
+    ListField,
+    ReferenceField,
+    StringField,
+)
+from .base import Operator
+from .core import TileableOperatorMixin
+
+
+class Fetch(Operator):
+    _op_type_ = opcodes.FETCH
+
+    source_key = StringField("source_key", default=None)
+
+
+class FetchMixin(TileableOperatorMixin):
+    def check_inputs(self, inputs):
+        # no inputs
+        if inputs and len(inputs) > 0:
+            raise ValueError(f"{type(self).__name__} has no inputs")
+
+
+class FetchShuffle(Operator):
+    _op_type_ = opcodes.FETCH_SHUFFLE
+
+    source_keys = ListField("source_keys", FieldTypes.string)
+    n_mappers = Int32Field("n_mappers")
+    n_reducers = Int32Field("n_reducers")
+    shuffle_fetch_type = ReferenceField("shuffle_fetch_type")
+
+
+class ShuffleFetchType(enum.Enum):
+    FETCH_BY_KEY = 0
+    FETCH_BY_INDEX = 1

maxframe/core/operator/fuse.py

@@ -0,0 +1,29 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ... import opcodes
+from ...serialization.serializables import ReferenceField
+from ..graph import ChunkGraph
+from .base import Operator
+
+
+class Fuse(Operator):
+    __slots__ = ("_fuse_graph",)
+    _op_type_ = opcodes.FUSE
+
+    fuse_graph = ReferenceField("fuse_graph", ChunkGraph)
+
+
+class FuseChunkMixin:
+    __slots__ = ()

maxframe/core/operator/objects.py

@@ -0,0 +1,72 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...serialization.serializables import BoolField
+from ..entity import OutputType, register_fetch_class
+from .base import Operator
+from .core import TileableOperatorMixin
+from .fetch import Fetch, FetchMixin
+from .fuse import Fuse, FuseChunkMixin
+
+
+class ObjectOperator(Operator):
+    pass
+
+
+class ObjectOperatorMixin(TileableOperatorMixin):
+    _output_type_ = OutputType.object
+
+
+class ObjectFuseChunkMixin(FuseChunkMixin, ObjectOperatorMixin):
+    __slots__ = ()
+
+
+class ObjectFuseChunk(ObjectFuseChunkMixin, Fuse):
+    pass
+
+
+class ObjectFetch(FetchMixin, ObjectOperatorMixin, Fetch):
+    _output_type_ = OutputType.object
+
+    def __init__(self, **kw):
+        kw.pop("output_types", None)
+        kw.pop("_output_types", None)
+        super().__init__(**kw)
+
+    def _new_tileables(self, inputs, kws=None, **kw):
+        if "_key" in kw and self.source_key is None:
+            self.source_key = kw["_key"]
+        return super()._new_tileables(inputs, kws=kws, **kw)
+
+
+register_fetch_class(OutputType.object, ObjectFetch, None)
+
+
+class MergeDictOperator(ObjectOperator, ObjectOperatorMixin):
+    _merge = BoolField("merge")
+
+    def __init__(self, merge=None, **kw):
+        super().__init__(_merge=merge, **kw)
+
+    @property
+    def merge(self):
+        return self._merge
+
+    @classmethod
+    def concat_tileable_chunks(cls, tileable):
+        assert not tileable.is_coarse()
+
+        op = cls(merge=True)
+        chunk = cls(merge=True).new_chunk(tileable.chunks)
+        return op.new_tileable([tileable], chunks=[chunk], nsplits=((1,),))

maxframe/core/operator/shuffle.py

@@ -0,0 +1,111 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ... import opcodes
+from ...serialization.serializables import (
+    FieldTypes,
+    Int32Field,
+    StringField,
+    TupleField,
+)
+from . import FetchShuffle, ShuffleFetchType
+from .base import Operator, OperatorStage, VirtualOperator
+
+
+class ShuffleProxy(VirtualOperator):
+    _op_type_ = opcodes.SHUFFLE_PROXY
+    n_mappers = Int32Field("n_mappers", default=0)
+    # `n_reducers` will be updated in `MapReduceOperator._new_chunks`
+    n_reducers = Int32Field("n_reducers", default=0)
+
+
+class MapReduceOperator(Operator):
+    """
+    An operator for shuffle execution which partitions data by the value in each record’s partition key, and
+    send the partitioned data from all mappers to all reducers.
+    """
+
+    # for reducer
+    reducer_index = TupleField("reducer_index", FieldTypes.uint64)
+    # Total reducer nums, which also be shuffle blocks for single mapper.
+    n_reducers = Int32Field("n_reducers")
+    # The reducer ordinal in all reducers. It's different from reducer_index,
+    # which might be a tuple.
+    # `reducer_ordinal` will be set in `_new_chunks`.
+    reducer_ordinal = Int32Field("reducer_ordinal")
+    reducer_phase = StringField("reducer_phase", default=None)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.stage == OperatorStage.reduce:
+            # for reducer, we assign worker at first
+            self.scheduling_hint.reassign_worker = True
+
+    def get_dependent_data_keys(self):
+        from .fetch import FetchShuffle
+
+        if self.stage == OperatorStage.reduce:
+            inputs = self.inputs or ()
+            deps = []
+            for inp in inputs:
+                if isinstance(inp.op, ShuffleProxy):
+                    deps.extend(
+                        [(chunk.key, self.reducer_index) for chunk in inp.inputs or ()]
+                    )
+                elif isinstance(inp.op, FetchShuffle):
+                    # fetch shuffle by index doesn't store data keys, so it won't run into this function.
+                    assert inp.op.shuffle_fetch_type == ShuffleFetchType.FETCH_BY_KEY
+                    deps.extend([(k, self.reducer_index) for k in inp.op.source_keys])
+                else:
+                    deps.append(inp.key)
+            return deps
+        return super().get_dependent_data_keys()
+
+    def iter_mapper_keys(self, input_id=0):
+        # key is mapper chunk key, index is mapper chunk index.
+        input_chunk = self.inputs[input_id]
+        if isinstance(input_chunk.op, ShuffleProxy):
+            keys = [inp.key for inp in input_chunk.inputs]
+        else:
+            assert isinstance(input_chunk.op, FetchShuffle), input_chunk.op
+            if input_chunk.op.shuffle_fetch_type == ShuffleFetchType.FETCH_BY_INDEX:
+                # For fetch shuffle by index, all shuffle block of same reducers are
+                # identified by their index. chunk key are not needed any more.
+                # so just mock key here.
+                # keep this in sync with ray executor `execute_subtask`.
+                return list(range(input_chunk.op.n_mappers))
+            keys = input_chunk.op.source_keys
+        return keys
+
+    def iter_mapper_data(self, ctx, input_id=0, pop=False, skip_none=False):
+        for key in self.iter_mapper_keys(input_id):
+            try:
+                if pop:
+                    yield ctx.pop((key, self.reducer_index))
+                else:
+                    yield ctx[key, self.reducer_index]
+            except KeyError:
+                if not skip_none:  # pragma: no cover
+                    raise
+                if not pop:
+                    ctx[key, self.reducer_index] = None
+
+    def execute(self, ctx, op):
+        """The mapper stage must ensure all mapper blocks are inserted into ctx
+        and no blocks for some reducers are missing. This is needed by shuffle
+        fetch by index, which shuffle block are identified by the  index instead
+        of data keys. For operators implementation simplicity, we can sort the
+        `ctx` by key which are (chunk key, reducer index) tuple and relax the
+        insert order requirements.
+        """

maxframe/core/operator/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

maxframe/core/operator/tests/test_core.py

@@ -0,0 +1,64 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from ....dataframe import core  # noqa: F401  # pylint: disable=unused-variable
+from ... import OutputType
+from .. import Operator, TileableOperatorMixin
+
+
+class MyOperator(Operator, TileableOperatorMixin):
+    @classmethod
+    def execute(cls, ctx, op):
+        return 1
+
+    @classmethod
+    def estimate_size(cls, ctx, op):
+        return 1
+
+
+class MyOperator2(MyOperator):
+    @classmethod
+    def execute(cls, ctx, op):
+        raise NotImplementedError
+
+    @classmethod
+    def estimate_size(cls, ctx, op):
+        raise NotImplementedError
+
+
+class _OperatorMixin(TileableOperatorMixin):
+    pass
+
+
+class MyOperator3(Operator, _OperatorMixin):
+    pass
+
+
+class MyOperator4(Operator, _OperatorMixin):
+    pass
+
+
+class MyOperator5(MyOperator4):
+    pass
+
+
+def test_unknown_dtypes():
+    op = MyOperator(_output_types=[OutputType.dataframe])
+    df = op.new_tileable(None, dtypes=None)
+    op2 = MyOperator(_output_types=[OutputType.scalar])
+    with pytest.raises(ValueError) as exc_info:
+        op2.new_tileable([df])
+    assert "executed first" in exc_info.value.args[0]

maxframe/core/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.