maxframe 1.0.0rc1__cp38-cp38-macosx_10_9_universal2.whl → 1.0.0rc3__cp38-cp38-macosx_10_9_universal2.whl

maxframe/io/objects/core.py

@@ -0,0 +1,140 @@
+# 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 abc import ABCMeta, abstractmethod
+from typing import Any, Dict, Type, Union
+
+import msgpack
+
+from ...core import Entity, EntityData
+from ...core.entity import ObjectData, TileableData
+from ...lib import wrapped_pickle as pickle
+from ...typing_ import SlicesType, TileableType
+from ...utils import TypeDispatcher
+from ..odpsio.volumeio import ODPSVolumeReader, ODPSVolumeWriter
+
+_MetaType = Dict[str, Any]
+
+_META_FILE_NAME = ".meta"
+_META_PICKLED_KEYS_KEY = ".pickled_keys"
+
+
+_io_handler_dispatcher = TypeDispatcher()
+
+
+def register_object_io_handler(tileable_data_type: Type[TileableData]):
+    def wrapper(handler_cls):
+        _io_handler_dispatcher.register(tileable_data_type, handler_cls)
+        return handler_cls
+
+    return wrapper
+
+
+def get_object_io_handler(
+    tileable_data_type: Union[Entity, EntityData, Type[EntityData]]
+) -> 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)
+
+
+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 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)
+
+    @abstractmethod
+    def _read_object_body(
+        self,
+        reader: ODPSVolumeReader,
+        tileable: TileableType,
+        meta: Dict[str, Any],
+        slices: SlicesType = None,
+    ) -> Any:
+        raise NotImplementedError
+
+    def read_object(
+        self,
+        reader: ODPSVolumeReader,
+        tileable: TileableType,
+        slices: SlicesType = None,
+    ) -> Any:
+        meta = self.read_object_meta(reader, tileable)
+        return self._read_object_body(reader, tileable, meta, slices)
+
+    @abstractmethod
+    def _write_object_body(
+        self, writer: ODPSVolumeWriter, tileable: TileableType, value: Any
+    ):
+        raise NotImplementedError
+
+    def write_object_meta(
+        self,
+        writer: ODPSVolumeWriter,
+        tileable: TileableType,
+        extra_meta: Dict[str, Any] = None,
+    ):
+        meta_obj = tileable.params.copy()
+        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 write_object(
+        self, writer: ODPSVolumeWriter, tileable: TileableType, value: Any
+    ):
+        self.write_object_meta(writer, tileable)
+        self._write_object_body(writer, tileable, value)
+
+
+@register_object_io_handler(ObjectData)
+class ObjectIOHandler(AbstractObjectIOHandler):
+    def _read_object_body(
+        self,
+        reader: ODPSVolumeReader,
+        tileable: TileableType,
+        meta: Dict[str, Any],
+        slices: SlicesType = None,
+    ) -> Any:
+        return pickle.loads(reader.read_file("data"))
+
+    def _write_object_body(
+        self, writer: ODPSVolumeWriter, tileable: TileableType, value: Any
+    ):
+        writer.write_file("data", pickle.dumps(value))

maxframe/io/objects/tensor.py

@@ -0,0 +1,76 @@
+# 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 struct
+from io import BytesIO
+from typing import Any, Dict
+
+import msgpack
+import numpy as np
+
+from ...lib import wrapped_pickle as pickle
+from ...tensor.core import TensorData
+from ...typing_ import SlicesType, TileableType
+from ..odpsio import ODPSVolumeReader, ODPSVolumeWriter
+from .core import AbstractObjectIOHandler, register_object_io_handler
+
+
+@register_object_io_handler(TensorData)
+class TensorIOHandler(AbstractObjectIOHandler):
+    def write_object_meta(
+        self,
+        writer: ODPSVolumeWriter,
+        tileable: TileableType,
+        extra_meta: Dict[str, Any] = None,
+    ):
+        # fixme upload in real slices when tensors are supported in DPE
+        extra_meta = extra_meta or dict()
+        extra_meta["nsplits"] = ((np.nan,),)
+
+        super().write_object_meta(writer, tileable, extra_meta=extra_meta)
+
+    def _read_object_body(
+        self,
+        reader: ODPSVolumeReader,
+        tileable: TileableType,
+        meta: Dict[str, Any],
+        slices: SlicesType = None,
+    ) -> Any:
+        # fixme read data with slices when tensors are supported in DPE
+        body = reader.read_file("0,0.dat")
+        bio = BytesIO(body)
+        (header_len,) = struct.unpack("<I", bio.read(4))
+        header_data = msgpack.loads(bio.read(header_len))
+
+        pickled = bio.read(header_data[0])
+        bufs = [bio.read(size) for size in header_data[1:]]
+        return pickle.loads(pickled, buffers=bufs)
+
+    def _write_object_body(
+        self, writer: ODPSVolumeWriter, tileable: TileableType, value: Any
+    ):
+        # fixme upload in real slices when tensors are supported in DPE
+        def data_gen():
+            bufs = []
+            pickled = pickle.dumps(value, buffer_callback=bufs.append)
+            header_data = msgpack.dumps(
+                [len(pickled)] + [len(buf.raw()) for buf in bufs]
+            )
+            yield struct.pack("<I", len(header_data))
+            yield header_data
+            yield pickled
+            for buf in bufs:
+                yield buf
+
+        writer.write_file("0,0.dat", data_gen())

maxframe/io/objects/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/io/objects/tests/test_object_io.py

@@ -0,0 +1,97 @@
+# 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 numpy as np
+import pytest
+from odps import ODPS
+
+from ....core import OutputType
+from ....core.operator import ObjectOperatorMixin, Operator
+from ....tensor.datasource import ArrayDataSource
+from ....tests.utils import tn
+from ...odpsio import ODPSVolumeReader, ODPSVolumeWriter
+from ..core import get_object_io_handler
+
+
+class TestObjectOp(Operator, ObjectOperatorMixin):
+    def __call__(self):
+        self._output_types = [OutputType.object]
+        return self.new_tileable([])
+
+
+@pytest.fixture(scope="module")
+def create_volume(request, oss_config):
+    test_vol_name = tn("test_object_io_volume")
+    odps_entry = ODPS.from_environments()
+
+    try:
+        odps_entry.delete_volume(test_vol_name, auto_remove_dir=True, recursive=True)
+    except:
+        pass
+
+    oss_test_dir_name = tn("test_oss_directory")
+    if oss_config is None:
+        pytest.skip("Need oss and its config to run this test")
+    (
+        oss_access_id,
+        oss_secret_access_key,
+        oss_bucket_name,
+        oss_endpoint,
+    ) = oss_config.oss_config
+    test_location = "oss://%s:%s@%s/%s/%s" % (
+        oss_access_id,
+        oss_secret_access_key,
+        oss_endpoint,
+        oss_bucket_name,
+        oss_test_dir_name,
+    )
+    oss_config.oss_bucket.put_object(oss_test_dir_name + "/", b"")
+    odps_entry.create_external_volume(test_vol_name, location=test_location)
+
+    try:
+        yield test_vol_name
+    finally:
+        try:
+            odps_entry.delete_volume(
+                test_vol_name, auto_remove_dir=True, recursive=True
+            )
+        except:
+            pass
+
+
+def test_simple_object_io(create_volume):
+    obj = TestObjectOp()()
+    data = "abcdefg"
+
+    odps_entry = ODPS.from_environments()
+
+    reader = ODPSVolumeReader(odps_entry, create_volume, obj.key)
+    writer = ODPSVolumeWriter(odps_entry, create_volume, obj.key)
+
+    handler = get_object_io_handler(obj)()
+    handler.write_object(writer, obj, data)
+    assert data == handler.read_object(reader, obj)
+
+
+def test_tensor_object_io(create_volume):
+    data = np.array([[4, 9, 2], [3, 5, 7], [8, 1, 6]])
+    obj = ArrayDataSource(data, dtype=data.dtype)(data.shape)
+
+    odps_entry = ODPS.from_environments()
+
+    reader = ODPSVolumeReader(odps_entry, create_volume, obj.key)
+    writer = ODPSVolumeWriter(odps_entry, create_volume, obj.key)
+
+    handler = get_object_io_handler(obj)()
+    handler.write_object(writer, obj, data)
+    np.testing.assert_equal(data, handler.read_object(reader, obj))

maxframe/{odpsio → io/odpsio}/__init__.py

@@ -14,8 +14,10 @@
 
 from .arrow import arrow_to_pandas, pandas_to_arrow
 from .schema import (
+    arrow_schema_to_odps_schema,
     build_dataframe_table_meta,
     odps_schema_to_pandas_dtypes,
     pandas_to_odps_schema,
 )
-from .tableio import HaloTableIO
+from .tableio import HaloTableIO, ODPSTableIO
+from .volumeio import ODPSVolumeReader, ODPSVolumeWriter

maxframe/{odpsio → io/odpsio}/arrow.py

@@ -17,10 +17,10 @@ from typing import Any, Tuple, Union
 import pandas as pd
 import pyarrow as pa
 
-from ..core import OutputType
-from ..protocol import DataFrameTableMeta
-from ..tensor.core import TENSOR_TYPE
-from ..typing_ import ArrowTableType, PandasObjectTypes
+from ...core import OutputType
+from ...protocol import DataFrameTableMeta
+from ...tensor.core import TENSOR_TYPE
+from ...typing_ import ArrowTableType, PandasObjectTypes
 from .schema import build_dataframe_table_meta
 
 
@@ -45,9 +45,13 @@ def _rebuild_dataframe(
 
 def _rebuild_index(df: pd.DataFrame, table_meta: DataFrameTableMeta) -> pd.Index:
     if df.shape[1] > 1:
-        df.columns = pd.Index(table_meta.pd_index_level_names)
-        return pd.MultiIndex.from_frame(df)
-    return pd.Index(df.iloc[:, 0], name=table_meta.pd_index_level_names[0])
+        idx = pd.MultiIndex.from_frame(df)
+        idx.names = table_meta.pd_index_level_names
+    else:
+        # make sure even if None names are updated properly
+        idx = pd.Index(df.iloc[:, 0])
+        idx.name = table_meta.pd_index_level_names[0]
+    return idx
 
 
 def arrow_to_pandas(
@@ -75,7 +79,7 @@ def pandas_to_arrow(
         df.columns = pd.Index(table_meta.table_column_names)
         if not ignore_index:
             df = df.rename_axis(table_meta.table_index_column_names).reset_index()
-    elif ignore_index:
+    elif ignore_index and table_meta.type != OutputType.index:
         df = pd.DataFrame([], columns=[])
     elif table_meta.type == OutputType.index:
         names = [f"_idx_{idx}" for idx in range(len(df.names))]

maxframe/{odpsio → io/odpsio}/schema.py

@@ -21,9 +21,9 @@ import pyarrow as pa
 from odps import types as odps_types
 from pandas.api import types as pd_types
 
-from ..core import TILEABLE_TYPE, OutputType
-from ..protocol import DataFrameTableMeta
-from ..tensor.core import TENSOR_TYPE
+from ...core import TILEABLE_TYPE, OutputType
+from ...protocol import DataFrameTableMeta
+from ...tensor.core import TENSOR_TYPE
 
 _TEMP_TABLE_PREFIX = "tmp_mf_"
 
@@ -126,10 +126,15 @@ def odps_type_to_arrow_type(
             ]
             col_type = pa.struct(fields)
         elif isinstance(odps_type, odps_types.Decimal):
-            col_type = pa.decimal128(
-                odps_type.precision or odps_types.Decimal._max_precision,
-                odps_type.scale or odps_types.Decimal._max_scale,
-            )
+            if odps_type.name == "decimal":
+                # legacy decimal data without precision or scale
+                # precision data from internal compat mode
+                col_type = pa.decimal128(38, 18)
+            else:
+                col_type = pa.decimal128(
+                    odps_type.precision or odps_types.Decimal._max_precision,
+                    odps_type.scale or odps_types.Decimal._max_scale,
+                )
         elif isinstance(odps_type, (odps_types.Varchar, odps_types.Char)):
             col_type = pa.string()
         else:
@@ -179,7 +184,7 @@ def pandas_to_odps_schema(
     unknown_as_string: bool = False,
     ignore_index=False,
 ) -> Tuple[odps_types.OdpsSchema, DataFrameTableMeta]:
-    from .. import dataframe as md
+    from ... import dataframe as md
     from .arrow import pandas_to_arrow
 
     if _is_scalar_object(df_obj):
@@ -273,7 +278,7 @@ def build_table_column_name(
 def build_dataframe_table_meta(
     df_obj: Any, ignore_index: bool = False
 ) -> DataFrameTableMeta:
-    from .. import dataframe as md
+    from ... import dataframe as md
 
     col_to_count = defaultdict(lambda: 0)
     col_to_idx = defaultdict(lambda: 0)
@@ -289,8 +294,6 @@ def build_dataframe_table_meta(
     else:  # pragma: no cover
         raise TypeError(f"Cannot accept type {type(df_obj)}")
 
-    assert not ignore_index or obj_type in (OutputType.dataframe, OutputType.series)
-
     if obj_type == OutputType.scalar:
         pd_dtypes = pd.Series([])
         column_index_names = []
@@ -346,7 +349,7 @@ def build_dataframe_table_meta(
     else:
         index_dtypes = pd.Series([pd_index_val.dtype], index=pd_index_val.names)
 
-    if ignore_index:
+    if ignore_index and obj_type != OutputType.index:
         table_index_column_names = []
         pd_index_dtypes = pd.Series([], index=[])
     else: