maxframe 1.3.1__cp311-cp311-win_amd64.whl → 2.0.0__cp311-cp311-win_amd64.whl

maxframe/dataframe/utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import dataclasses
 import functools
 import inspect
 import itertools
@@ -20,26 +21,30 @@ import operator
 import sys
 from contextlib import contextmanager
 from numbers import Integral
-from typing import TYPE_CHECKING, Any, Callable, List
+from typing import TYPE_CHECKING, Any, Callable, List, Optional
 
 import numpy as np
 import pandas as pd
-from pandas.api.extensions import ExtensionDtype
 from pandas.api.types import is_string_dtype
 from pandas.core.dtypes.inference import is_dict_like, is_list_like
 
-from ..core import Entity, ExecutableTuple
+from ..core import Entity, ExecutableTuple, OutputType, get_output_types
 from ..lib.mmh3 import hash as mmh_hash
 from ..udf import MarkedFunction
 from ..utils import (
     ModulePlaceholder,
     is_full_slice,
     lazy_import,
-    parse_version,
+    make_dtype,
+    make_dtypes,
+    quiet_stdio,
     sbytes,
     tokenize,
 )
 
+if TYPE_CHECKING:
+    from .core import IndexValue
+
 try:
     import pyarrow as pa
 except ImportError:  # pragma: no cover
@@ -49,15 +54,6 @@ if TYPE_CHECKING:
     from .operators import DataFrameOperator
 
 cudf = lazy_import("cudf", rename="cudf")
-vineyard = lazy_import("vineyard")
-try:
-    import ray
-
-    ray_release_version = parse_version(ray.__version__).release
-    ray_deprecate_ml_dataset = ray_release_version[:2] >= (2, 0)
-except ImportError:
-    ray_release_version = None
-    ray_deprecate_ml_dataset = None
 logger = logging.getLogger(__name__)
 
 try:
@@ -603,6 +599,7 @@ def build_series(
     else:
         series_index = index[:0] if index is not None else None
 
+    name = name or getattr(series_obj, "name", None)
     for size, fill_value in zip(sizes, fill_values):
         empty_series = build_empty_series(dtype, name=name, index=series_index)
         record = _generate_value(dtype, fill_value)
@@ -1039,20 +1036,6 @@ def to_arrow_dtypes(dtypes, test_df=None):
     return new_dtypes
 
 
-def make_dtype(dtype):
-    if isinstance(dtype, (np.dtype, ExtensionDtype)):
-        return dtype
-    return np.dtype(dtype) if dtype is not None else None
-
-
-def make_dtypes(dtypes):
-    if dtypes is None:
-        return None
-    if not isinstance(dtypes, pd.Series):
-        dtypes = pd.Series(dtypes)
-    return dtypes.apply(make_dtype)
-
-
 def is_dataframe(x):
     if cudf is not None:  # pragma: no cover
         if isinstance(x, cudf.DataFrame):
@@ -1314,6 +1297,8 @@ def pack_func_args(df, funcs, *args, args_bind_position=1, **kwargs) -> Any:
     AttributeError :
         If there's a string but no corresponding function is found.
     """
+    from ..udf import MarkedFunction
+
     if not args and not kwargs:
         return funcs
 
@@ -1325,8 +1310,6 @@ def pack_func_args(df, funcs, *args, args_bind_position=1, **kwargs) -> Any:
 
     f = get_callable_by_name(df, funcs) if isinstance(funcs, str) else funcs
 
-    from ..udf import MarkedFunction
-
     if isinstance(f, MarkedFunction):
         # for marked function, pack the inner function, and reset as mark function
         packed_func = f.copy()
@@ -1351,7 +1334,7 @@ def get_callable_by_name(df: Any, func_name: str) -> Callable:
 
     Parameters
     ----------
-    df: padnas.Series or pandas.Dataframe
+    df: pandas.Series or pandas.Dataframe
         The receiver of the func name.
     func_name : str
         The func name.
@@ -1381,10 +1364,215 @@ def get_callable_by_name(df: Any, func_name: str) -> Callable:
     )
 
 
+@dataclasses.dataclass
+class InferredDataFrameMeta:
+    output_type: OutputType
+    dtypes: Optional[pd.Series] = None
+    dtype: Optional[Any] = None
+    name: Optional[str] = None
+    index_value: Optional["IndexValue"] = None
+    maybe_agg: bool = False
+    elementwise: bool = False
+
+    def check_absence(self, *args: str) -> None:
+        args_set = set(args)
+        if self.output_type == OutputType.dataframe:
+            args_set.difference_update(["dtype", "name"])
+        else:
+            args_set.difference_update(["dtypes"])
+        absent_args = [arg for arg in sorted(args_set) if getattr(self, arg) is None]
+        if absent_args:
+            raise TypeError(
+                f"Cannot determine {', '.join(absent_args)} by calculating "
+                "with mock data, please specify it as arguments"
+            )
+
+
+def _get_groupby_input_df(groupby):
+    in_df = groupby
+    while in_df.op.output_types[0] not in (OutputType.dataframe, OutputType.series):
+        in_df = in_df.inputs[0]
+    return in_df
+
+
+def infer_dataframe_return_value(
+    df_obj,
+    func,
+    output_type=None,
+    dtypes=None,
+    dtype=None,
+    name=None,
+    index=None,
+    inherit_index=False,
+    build_kw=None,
+    elementwise=None,
+) -> InferredDataFrameMeta:
+    from .core import GROUPBY_TYPE
+
+    if elementwise is None:
+        unwrapped_func = func
+        if isinstance(unwrapped_func, MarkedFunction):
+            unwrapped_func = unwrapped_func.func
+        while True:
+            if isinstance(unwrapped_func, functools.partial):
+                unwrapped_func = unwrapped_func.func
+            elif hasattr(unwrapped_func, "__wrapped__"):
+                unwrapped_func = unwrapped_func.__wrapped__
+            else:
+                break
+        elementwise = isinstance(unwrapped_func, np.ufunc)
+
+    ret_index_value = None
+    if output_type is not None and (dtypes is not None or dtype is not None):
+        if inherit_index:
+            ret_index_value = df_obj.index_value
+        elif index is not None:
+            ret_index_value = parse_index(index)
+
+        if ret_index_value is not None:
+            return InferredDataFrameMeta(
+                output_type,
+                dtypes,
+                dtype,
+                name,
+                ret_index_value,
+                elementwise=elementwise or False,
+            )
+
+    ret_output_type = ret_dtypes = None
+    maybe_agg = False
+    build_kw = build_kw or {}
+    obj_key = df_obj.key
+
+    if elementwise:
+        inherit_index = True
+        (ret_output_type,) = get_output_types(df_obj)
+    if index is not None:
+        ret_index_value = parse_index(index)
+
+    if isinstance(df_obj, GROUPBY_TYPE):
+        is_groupby = True
+        empty_df_obj = df_obj.op.build_mock_groupby(**build_kw)
+    else:
+        is_groupby = False
+        empty_df_obj = (
+            build_df(df_obj, **build_kw)
+            if df_obj.ndim == 2
+            else build_series(df_obj, **build_kw)
+        )
+    try:
+        with np.errstate(all="ignore"), quiet_stdio():
+            infer_df_obj = func(empty_df_obj)
+
+        if ret_index_value is None:
+            if (
+                infer_df_obj is None
+                or not hasattr(infer_df_obj, "index")
+                or infer_df_obj.index is None
+            ):
+                ret_index_value = parse_index(pd.RangeIndex(-1))
+            elif (
+                infer_df_obj.index is getattr(empty_df_obj, "index", None)
+                or inherit_index
+            ):
+                ret_index_value = df_obj.index_value
+            else:
+                ret_index_value = parse_index(infer_df_obj.index, obj_key, func)
+
+        if isinstance(infer_df_obj, pd.DataFrame):
+            if output_type is not None and output_type != OutputType.dataframe:
+                raise TypeError(
+                    f'Cannot infer output_type as "series", '
+                    f'please specify `output_type` as "dataframe"'
+                )
+            ret_output_type = ret_output_type or OutputType.dataframe
+            ret_dtypes = ret_dtypes or infer_df_obj.dtypes
+        else:
+            if output_type is not None and output_type == OutputType.dataframe:
+                raise TypeError(
+                    f'Cannot infer output_type as "dataframe", '
+                    f'please specify `output_type` as "series"'
+                )
+            ret_output_type = ret_output_type or OutputType.series
+            name = name or getattr(infer_df_obj, "name", None)
+            dtype = dtype or infer_df_obj.dtype
+
+        if is_groupby and len(infer_df_obj) <= 2:
+            # we create mock df with 4 rows, 2 groups
+            # if return df has 2 rows, we assume that
+            # it's an aggregation operation
+            maybe_agg = True
+
+        return InferredDataFrameMeta(
+            ret_output_type,
+            make_dtypes(ret_dtypes),
+            make_dtype(dtype),
+            name,
+            ret_index_value,
+            maybe_agg,
+            elementwise=elementwise,
+        )
+    except:  # noqa: E722  # nosec
+        logger.info(
+            "Exception raised while inferring meta of function result", exc_info=True
+        )
+        return InferredDataFrameMeta(
+            output_type,
+            make_dtypes(dtypes),
+            make_dtype(dtype),
+            name,
+            ret_index_value,
+            maybe_agg,
+            elementwise=elementwise,
+        )
+
+
 def copy_func_scheduling_hints(func, op: "DataFrameOperator") -> None:
+    from ..config import options
+
     if not isinstance(func, MarkedFunction):
         return
     if func.expect_engine:
         op.expect_engine = func.expect_engine
+
+    expect_resources = func.expect_resources or {}
+    default_function_running_options = options.function.default_running_options or {}
+
+    for key, value in default_function_running_options.items():
+        if key not in expect_resources or expect_resources.get(key) is None:
+            expect_resources[key] = value
+
     if func.expect_resources:
-        op.expect_resources = func.expect_resources
+        op.expect_resources = expect_resources
+
+
+def make_column_list(col, dtypes_or_columns, level=None):
+    """Returns [col] if col is a column in dtypes"""
+    try:
+        if isinstance(dtypes_or_columns, pd.Series):
+            idx = dtypes_or_columns.index
+        else:
+            idx = dtypes_or_columns
+
+        if level is None:
+            if col in idx:
+                return [col]
+            elif isinstance(col, int):
+                col = [col]
+            if all(c in idx for c in col):
+                return col
+            if all(isinstance(c, int) for c in col):
+                return [idx[c] for c in col]
+            return col
+        else:
+            level_idx = idx.get_level_values(level)
+            if isinstance(col, list):
+                cols = col
+            else:
+                cols = [col]
+            mask = level_idx.isin(cols)
+            if not mask.any():
+                mask = col
+            return idx[mask]
+    except (IndexError, TypeError, ValueError):
+        return col

maxframe/dataframe/window/rolling.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 from collections import OrderedDict
+from typing import List
 
 import pandas as pd
 
 from ... import opcodes
+from ...core import EntityData
 from ...serialization.serializables import (
     AnyField,
     BoolField,
@@ -55,10 +57,11 @@ class DataFrameRollingAgg(DataFrameOperator, DataFrameOperatorMixin):
     def __init__(self, output_types=None, **kw):
         super().__init__(_output_types=output_types, **kw)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        input_iter = iter(self._inputs)
-        self.input = next(input_iter)
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameRollingAgg", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        input_iter = iter(op._inputs)
+        op.input = next(input_iter)
 
     def __call__(self, rolling):
         inp = rolling.input

maxframe/env.py

@@ -16,6 +16,7 @@
 MAXFRAME_NAMESPACE = "MAXFRAME_NAMESPACE"
 
 # Maxframe Service common envs
+MAXFRAME_HTTP_PORT_FILE = "MAXFRAME_PROXY_PORT_FILE"
 MAXFRAME_SERVICE_PORT = "MAXFRAME_SERVICE_PORT"
 MAXFRAME_SERVICE_PORT_RETRIES = "MAXFRAME_SERVICE_PORT_RETRIES"
 MAXFRAME_SERVICE_LISTEN_ADDRESS = "MAXFRAME_SERVICE_LISTEN_ADDRESS"

maxframe/errors.py

@@ -29,6 +29,10 @@ class MaxFrameDeprecationError(MaxFrameUserError):
     pass
 
 
+class TileableNotExecutedError(MaxFrameError):
+    pass
+
+
 class NoTaskServerResponseError(MaxFrameError):
     pass
 
@@ -36,3 +40,8 @@ class NoTaskServerResponseError(MaxFrameError):
 class SessionAlreadyClosedError(MaxFrameError):
     def __init__(self, session_id: str):
         super().__init__(f"Session {session_id} is already closed")
+
+
+class EngineUnavailableError(MaxFrameIntentionalError):
+    def __init__(self, engine_type: str):
+        super().__init__(f"Engine {engine_type} is not ready")

maxframe/extension.py

@@ -17,7 +17,7 @@ import itertools
 from typing import TYPE_CHECKING, Iterable, Tuple, Type
 
 if TYPE_CHECKING:
-    from maxframe.codegen import BigDagCodeGenerator
+    from .codegen import DAGCodeGenerator
 
     try:
         import maxframe_framedriver.runners
@@ -28,7 +28,7 @@ if TYPE_CHECKING:
 class MaxFrameExtension(metaclass=abc.ABCMeta):
     @classmethod
     @abc.abstractmethod
-    def get_codegen(cls) -> Type["BigDagCodeGenerator"]:
+    def get_codegen(cls) -> Type["DAGCodeGenerator"]:
         raise NotImplementedError
 
     @classmethod
@@ -48,6 +48,17 @@ class MaxFrameExtension(metaclass=abc.ABCMeta):
         """
         pass
 
+    @classmethod
+    async def create_session(cls, session_id: str) -> None:
+        """
+        create the session
+        Parameters
+        ----------
+        session_id : str
+            The session id.
+        """
+        pass
+
     @classmethod
     async def reload_session(cls, session_id: str) -> None:
         """

maxframe/io/objects/core.py

@@ -16,74 +16,68 @@ from abc import ABCMeta, abstractmethod
 from typing import Any, Dict, Type, Union
 
 import msgpack
+import numpy as np
 
-from ...core import Entity, EntityData
-from ...core.entity import ObjectData, TileableData
+from ...core import Entity, EntityData, OutputType, get_output_types
 from ...lib import wrapped_pickle as pickle
 from ...typing_ import SlicesType, TileableType
-from ...utils import TypeDispatcher
+from ...utils import deserialize_serializable, serialize_serializable
 from ..odpsio.volumeio import ODPSVolumeReader, ODPSVolumeWriter
 
 _MetaType = Dict[str, Any]
 
 _META_FILE_NAME = ".meta"
 _META_PICKLED_KEYS_KEY = ".pickled_keys"
+_META_SERIALIZED_VALUE_KEY = ".serialized_values"
 
 
-_io_handler_dispatcher = TypeDispatcher()
+_io_handler_dispatcher = dict()
 
 
-def register_object_io_handler(tileable_data_type: Type[TileableData]):
+def register_object_io_handler(out_type: OutputType):
     def wrapper(handler_cls):
-        _io_handler_dispatcher.register(tileable_data_type, handler_cls)
+        _io_handler_dispatcher[out_type] = handler_cls
         return handler_cls
 
     return wrapper
 
 
 def get_object_io_handler(
-    tileable_data_type: Union[Entity, EntityData, Type[EntityData]]
+    tileable_data_type: Union[Entity, EntityData, Type[EntityData], OutputType]
 ) -> Type["AbstractObjectIOHandler"]:
-    if not isinstance(tileable_data_type, type):
-        if isinstance(tileable_data_type, Entity):
-            tileable_data_type = tileable_data_type.data
-        tileable_data_type = type(tileable_data_type)
-    return _io_handler_dispatcher.get_handler(tileable_data_type)
+    output_type = tileable_data_type
+    if not isinstance(tileable_data_type, OutputType):
+        if hasattr(tileable_data_type, "output_type"):
+            output_type = tileable_data_type.output_type
+        else:
+            [output_type] = get_output_types(tileable_data_type)
+    return _io_handler_dispatcher[output_type]
 
 
 class AbstractObjectIOHandler(metaclass=ABCMeta):
-    def _prepare_meta_for_serial(
-        self, tileable: TileableType, meta: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        to_pack = meta.copy()
-        pickled_keys = []
-        for k, v in meta.items():
-            if not isinstance(v, (str, bytes, int, float, bool)):
-                to_pack[k] = pickle.dumps(v)
-                pickled_keys.append(k)
-        to_pack[".pickled_keys"] = pickled_keys
-        return to_pack
-
-    def _prepare_meta_for_deserial(
-        self, tileable: TileableType, meta: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        pickled_keys = meta.pop(".pickled_keys", None) or []
-        for k in pickled_keys:
-            meta[k] = pickle.loads(meta[k])
-        return meta
+    def _prepare_meta_for_serial(self, tileable: TileableType) -> Dict[str, Any]:
+        return tileable.params.copy()
 
     def read_object_meta(
         self, reader: ODPSVolumeReader, tileable: TileableType
     ) -> Dict[str, Any]:
-        meta_obj = msgpack.loads(reader.read_file(_META_FILE_NAME))
-        return self._prepare_meta_for_deserial(tileable, meta_obj)
+        file_bin = reader.read_file(_META_FILE_NAME)
+        meta_obj = msgpack.loads(file_bin)
+        pickled_keys = meta_obj.pop(_META_PICKLED_KEYS_KEY, None) or []
+        serialized_values = meta_obj.pop(_META_SERIALIZED_VALUE_KEY, None) or {}
+        for k in pickled_keys:
+            try:
+                meta_obj[k] = deserialize_serializable(serialized_values[k])
+            except:
+                meta_obj[k] = pickle.loads(meta_obj[k])
+        return meta_obj
 
     @abstractmethod
-    def _read_object_body(
+    def read_object_body(
         self,
         reader: ODPSVolumeReader,
-        tileable: TileableType,
-        meta: Dict[str, Any],
+        params: Dict[str, Any],
+        extra_params: Dict[str, Any] = None,
         slices: SlicesType = None,
     ) -> Any:
         raise NotImplementedError
@@ -92,49 +86,71 @@ class AbstractObjectIOHandler(metaclass=ABCMeta):
         self,
         reader: ODPSVolumeReader,
         tileable: TileableType,
+        extra_params: Dict[str, Any] = None,
         slices: SlicesType = None,
     ) -> Any:
-        meta = self.read_object_meta(reader, tileable)
-        return self._read_object_body(reader, tileable, meta, slices)
+        obj_meta = self.read_object_meta(reader, tileable)
+        return self.read_object_body(reader, obj_meta, extra_params, slices)
 
     @abstractmethod
-    def _write_object_body(
-        self, writer: ODPSVolumeWriter, tileable: TileableType, value: Any
+    def write_object_body(
+        self, writer: ODPSVolumeWriter, params: Dict[str, Any], value: Any
     ):
         raise NotImplementedError
 
+    @classmethod
+    def _transform_serializable_meta(cls, meta: Dict[str, Any]) -> Dict[str, Any]:
+        to_pack = meta.copy()
+        pickled_keys = []
+        serialized_values = {}
+        for k, v in meta.items():
+            if not isinstance(v, (str, bytes, int, float, bool, type(None))):
+                to_pack[k] = pickle.dumps(v)
+                serialized_values[k] = serialize_serializable(v)
+                pickled_keys.append(k)
+        to_pack[_META_PICKLED_KEYS_KEY] = pickled_keys
+        to_pack[_META_SERIALIZED_VALUE_KEY] = serialized_values
+        return to_pack
+
     def write_object_meta(
         self,
         writer: ODPSVolumeWriter,
         tileable: TileableType,
         extra_meta: Dict[str, Any] = None,
     ):
-        meta_obj = tileable.params.copy()
+        meta_obj = self._prepare_meta_for_serial(tileable)
         if extra_meta:
             meta_obj.update(extra_meta)
-        meta_obj = self._prepare_meta_for_serial(tileable, meta_obj)
-        packed = msgpack.dumps(meta_obj)
-        writer.write_file(_META_FILE_NAME, packed)
+
+        def encode_np_num(obj):
+            if isinstance(obj, np.generic) and obj.shape == ():
+                return obj.item()
+            return obj
+
+        meta_obj = self._transform_serializable_meta(meta_obj)
+        writer.write_file(
+            _META_FILE_NAME, msgpack.dumps(meta_obj, default=encode_np_num)
+        )
 
     def write_object(
         self, writer: ODPSVolumeWriter, tileable: TileableType, value: Any
     ):
         self.write_object_meta(writer, tileable)
-        self._write_object_body(writer, tileable, value)
+        self.write_object_body(writer, tileable.params, value)
 
 
-@register_object_io_handler(ObjectData)
+@register_object_io_handler(OutputType.object)
 class ObjectIOHandler(AbstractObjectIOHandler):
-    def _read_object_body(
+    def read_object_body(
         self,
         reader: ODPSVolumeReader,
-        tileable: TileableType,
-        meta: Dict[str, Any],
+        params: Dict[str, Any],
+        extra_params: Dict[str, Any] = None,
         slices: SlicesType = None,
     ) -> Any:
         return pickle.loads(reader.read_file("data"))
 
-    def _write_object_body(
-        self, writer: ODPSVolumeWriter, tileable: TileableType, value: Any
+    def write_object_body(
+        self, writer: ODPSVolumeWriter, params: Dict[str, Any], value: Any
     ):
         writer.write_file("data", pickle.dumps(value))