maxframe 1.3.1__cp310-cp310-win_amd64.whl → 2.0.0b1__cp310-cp310-win_amd64.whl

maxframe/io/objects/tensor.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
 import struct
 from io import BytesIO
 from typing import Any, Dict
@@ -19,36 +20,65 @@ from typing import Any, Dict
 import msgpack
 import numpy as np
 
+from ...config import options
+from ...core import OutputType
 from ...lib import wrapped_pickle as pickle
-from ...tensor.core import TensorData
+from ...lib.version import parse as parse_version
 from ...typing_ import SlicesType, TileableType
 from ..odpsio import ODPSVolumeReader, ODPSVolumeWriter
 from .core import AbstractObjectIOHandler, register_object_io_handler
 
+try:
+    from maxframe import __version__ as mf_version
+except ImportError:
+    mf_version = None
 
-@register_object_io_handler(TensorData)
+
+@register_object_io_handler(OutputType.tensor)
 class TensorIOHandler(AbstractObjectIOHandler):
+    def _prepare_meta_for_serial(self, tileable: TileableType) -> Dict[str, Any]:
+        meta = super()._prepare_meta_for_serial(tileable)
+        meta["nsplits"] = (
+            getattr(tileable, "nsplits", None) or (np.nan,) * tileable.ndim
+        )
+        return meta
+
     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,),)
-
+        extra_meta["nsplits"] = (
+            getattr(tileable, "nsplits", None) or ((np.nan,),) * tileable.ndim
+        )
         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,
+    @staticmethod
+    def _get_chunk_file_name(params: Dict[str, Any]) -> str:
+        # fixme remove this when all v1.0 clients removed
+        mf_ver = options.session.client_version or mf_version
+        if mf_ver:
+            client_major = parse_version(
+                options.session.client_version or mf_version
+            ).major
+        else:
+            client_major = None
+
+        if client_major == 1:
+            # returns v1.0 file name as we do not support tiled tensors
+            return "0,0.dat"
+
+        ndim = len(params.get("shape") or (0,))
+        data_index = params.get("index") or (0,) * ndim
+        return ",".join(str(ix) for ix in data_index) + ".dat"
+
+    def _read_single_chunk(
+        self, reader: ODPSVolumeReader, params: Dict[str, Any]
     ) -> Any:
-        # fixme read data with slices when tensors are supported in DPE
-        body = reader.read_file("0,0.dat")
+        file_name = self._get_chunk_file_name(params)
+        body = reader.read_file(file_name)
         bio = BytesIO(body)
         (header_len,) = struct.unpack("<I", bio.read(4))
         header_data = msgpack.loads(bio.read(header_len))
@@ -57,10 +87,35 @@ class TensorIOHandler(AbstractObjectIOHandler):
         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
+    def read_object_body(
+        self,
+        reader: ODPSVolumeReader,
+        params: Dict[str, Any],
+        extra_params: Dict[str, Any] = None,
+        slices: SlicesType = None,
+    ) -> Any:
+        if "index" in params:
+            return self._read_single_chunk(reader, params)
+
+        tileable_params = params.copy()
+        tileable_params.update(extra_params or {})
+
+        # todo implements slices argument for head and tail fetching
+        assert slices is None
+
+        chunk_shape = tuple(len(x) for x in params["nsplits"])
+        block_array = np.empty(shape=chunk_shape, dtype="O")
+        for idx in itertools.product(*(range(sp) for sp in chunk_shape)):
+            chunk_params = tileable_params.copy()
+            chunk_params.pop("nsplits")
+            chunk_params["index"] = idx
+            block_array[idx] = self._read_single_chunk(reader, chunk_params)
+
+        return np.block(block_array.tolist())
+
+    def write_object_body(
+        self, writer: ODPSVolumeWriter, params: Dict[str, Any], 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)
@@ -73,4 +128,5 @@ class TensorIOHandler(AbstractObjectIOHandler):
             for buf in bufs:
                 yield buf
 
-        writer.write_file("0,0.dat", data_gen())
+        file_name = self._get_chunk_file_name(params)
+        writer.write_file(file_name, data_gen())

maxframe/io/objects/tests/test_object_io.py

@@ -19,7 +19,7 @@ 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 ....tests.utils import create_test_volume, tn
 from ...odpsio import ODPSVolumeReader, ODPSVolumeWriter
 from ..core import get_object_io_handler
 
@@ -31,61 +31,9 @@ class TestObjectOp(Operator, ObjectOperatorMixin):
 
 
 @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
-
-    if "test" in oss_endpoint:
-        # offline 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,
-        )
-        rolearn = None
-    else:
-        # online config
-        endpoint_parts = oss_endpoint.split(".", 1)
-        if "-internal" not in endpoint_parts[0]:
-            endpoint_parts[0] += "-internal"
-        test_location = "oss://%s/%s/%s" % (
-            ".".join(endpoint_parts),
-            oss_bucket_name,
-            oss_test_dir_name,
-        )
-        rolearn = oss_config.oss_rolearn
-
-    oss_config.oss_bucket.put_object(oss_test_dir_name + "/", b"")
-    odps_entry.create_external_volume(
-        test_vol_name, location=test_location, rolearn=rolearn
-    )
-
-    try:
+def create_volume(oss_config):
+    with create_test_volume(tn("test_object_io_vol"), oss_config) as test_vol_name:
         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):
@@ -119,6 +67,11 @@ def test_tensor_object_io(create_volume):
         odps_entry, create_volume, obj.key, replace_internal_host=True
     )
 
+    # test write and read full object
     handler = get_object_io_handler(obj)()
     handler.write_object(writer, obj, data)
     np.testing.assert_equal(data, handler.read_object(reader, obj))
+
+    # test read single chunk
+    params = {"index": (0, 0)}
+    np.testing.assert_equal(data, handler.read_object_body(reader, params))

maxframe/io/odpsio/arrow.py

@@ -68,6 +68,8 @@ def arrow_to_pandas(
         return _rebuild_dataframe(df, table_meta)
     elif table_meta.type == OutputType.index:
         return _rebuild_index(df, table_meta)
+    elif table_meta.type == OutputType.tensor:
+        return _rebuild_index(df, table_meta).to_numpy()
     elif table_meta.type == OutputType.scalar:
         return _rebuild_index(df, table_meta)[0]
     else:  # this could never happen  # pragma: no cover
@@ -107,9 +109,9 @@ def pandas_to_arrow(
             else:
                 table_datetime_cols = {"_idx_0"}
         df = df.to_frame(name=names[0] if len(names) == 1 else names)
-    elif table_meta.type == OutputType.scalar:
+    elif table_meta.type in (OutputType.scalar, OutputType.tensor):
         names = ["_idx_0"]
-        if isinstance(df, TENSOR_TYPE):
+        if isinstance(df, (TENSOR_TYPE, np.ndarray)):
             df = pd.DataFrame([], columns=names).astype({names[0]: df.dtype})
         else:
             df = pd.DataFrame([[df]], columns=names)
@@ -145,4 +147,15 @@ def pandas_to_arrow(
             col_data = pa_table.column(idx).cast(pa.timestamp("ms"))
             col_datas.append(col_data)
         pa_table = pa.Table.from_arrays(col_datas, names=col_names)
+
+        new_names, new_dtypes = [], []
+        for table_col, (pd_col, pd_dtype) in zip(
+            table_meta.table_column_names, table_meta.pd_column_dtypes.items()
+        ):
+            new_names.append(pd_col)
+            if table_col not in table_datetime_cols:
+                new_dtypes.append(pd_dtype)
+            else:
+                new_dtypes.append(np.dtype("datetime64[ms]"))
+        table_meta.pd_column_dtypes = pd.Series(new_dtypes, index=new_names)
     return pa_table, table_meta

maxframe/io/odpsio/schema.py

@@ -14,7 +14,7 @@
 
 import string
 from collections import defaultdict
-from typing import Any, Dict, Tuple, Union
+from typing import Any, Dict, Tuple
 
 import numpy as np
 import pandas as pd
@@ -23,9 +23,11 @@ from odps import types as odps_types
 from pandas.api import types as pd_types
 
 from ...core import TILEABLE_TYPE, OutputType
+from ...dataframe.core import DATAFRAME_TYPE, INDEX_TYPE, SERIES_TYPE
 from ...lib.dtypes_extension import ArrowDtype
 from ...protocol import DataFrameTableMeta
 from ...tensor.core import TENSOR_TYPE
+from ...utils import build_temp_table_name
 
 _TEMP_TABLE_PREFIX = "tmp_mf_"
 DEFAULT_SINGLE_INDEX_NAME = "_idx_0"
@@ -63,15 +65,15 @@ _odps_type_to_arrow = {
     odps_types.timestamp_ntz: pa.timestamp("ns"),
 }
 
-_based_for_pandas_pa_dtypes = Union[pa.MapType]
+_based_for_pandas_pa_types = (pa.ListType, pa.MapType)
 
 
-def is_based_for_pandas_dtype(dtype: pa.DataType) -> bool:
+def is_based_for_pandas_dtype(arrow_type: pa.DataType) -> bool:
     """
     Check whether the arrow type is based for one pandas data type.
     If true, we should make sure the environment support ArrowDtype.
     """
-    if not isinstance(dtype, _based_for_pandas_pa_dtypes):
+    if not isinstance(arrow_type, _based_for_pandas_pa_types):
         return False
 
     if ArrowDtype is None:
@@ -243,12 +245,37 @@ def pandas_dataframe_to_arrow_table(df: pd.DataFrame, nthreads=1) -> pa.Table:
     )
 
 
+def pandas_dtypes_to_arrow_schema(dtypes, unknown_as_string: bool = False) -> pa.Schema:
+    if unknown_as_string:
+        dt_list = [dt if dt != np.dtype("O") else pd.StringDtype() for dt in dtypes]
+        dtypes = pd.Series(dt_list, index=dtypes.index)
+    schema = pandas_types_to_arrow_schema(
+        pd.DataFrame([], columns=dtypes.index).astype(dtypes)
+    )
+    return schema
+
+
+def pandas_dtype_to_arrow_type(dtype, unknown_as_string: bool = False) -> pa.DataType:
+    if unknown_as_string and dtype == np.dtype("O"):
+        dtype = pd.StringDtype()
+    schema = pandas_types_to_arrow_schema(
+        pd.DataFrame([], columns=["a"]).astype({"a": dtype})
+    )
+    return schema.types[0]
+
+
 def is_scalar_object(df_obj: Any) -> bool:
     return (
         isinstance(df_obj, TENSOR_TYPE) and df_obj.shape == ()
     ) or pd_types.is_scalar(df_obj)
 
 
+def is_tensor_object(df_obj: Any) -> bool:
+    return (
+        isinstance(df_obj, TENSOR_TYPE) or isinstance(df_obj, np.ndarray)
+    ) and df_obj.ndim <= 1
+
+
 def _scalar_as_index(df_obj: Any) -> pd.Index:
     if isinstance(df_obj, TILEABLE_TYPE):
         return pd.Index([], dtype=df_obj.dtype)
@@ -264,7 +291,7 @@ def pandas_to_odps_schema(
     from ... import dataframe as md
     from .arrow import pandas_to_arrow
 
-    if is_scalar_object(df_obj):
+    if is_scalar_object(df_obj) or is_tensor_object(df_obj):
         empty_index = None
     elif hasattr(df_obj, "index_value"):
         empty_index = df_obj.index_value.to_pandas()[:0]
@@ -368,25 +395,25 @@ def build_table_column_name(
 
 
 def build_dataframe_table_meta(
-    df_obj: Any, ignore_index: bool = False
+    df_obj: Any, ignore_index: bool = False, session_id: str = None
 ) -> DataFrameTableMeta:
-    from ... import dataframe as md
-
     col_to_count = defaultdict(lambda: 0)
     col_to_idx = defaultdict(lambda: 0)
     pd_col_to_col_name = dict()
-    if isinstance(df_obj, (md.DataFrame, pd.DataFrame)):
+    if isinstance(df_obj, (DATAFRAME_TYPE, pd.DataFrame)):
         obj_type = OutputType.dataframe
-    elif isinstance(df_obj, (md.Series, pd.Series)):
+    elif isinstance(df_obj, (SERIES_TYPE, pd.Series)):
         obj_type = OutputType.series
-    elif isinstance(df_obj, (md.Index, pd.Index)):
+    elif isinstance(df_obj, (INDEX_TYPE, pd.Index)):
         obj_type = OutputType.index
+    elif is_tensor_object(df_obj) and df_obj.ndim == 1:
+        obj_type = OutputType.tensor
     elif is_scalar_object(df_obj):
         obj_type = OutputType.scalar
     else:  # pragma: no cover
         raise TypeError(f"Cannot accept type {type(df_obj)}")
 
-    if obj_type == OutputType.scalar:
+    if obj_type in (OutputType.scalar, OutputType.tensor):
         pd_dtypes = pd.Series([])
         column_index_names = []
         index_obj = _scalar_as_index(df_obj)
@@ -404,7 +431,10 @@ def build_dataframe_table_meta(
         index_obj = df_obj.index
 
     if isinstance(df_obj, TILEABLE_TYPE):
-        table_name = _TEMP_TABLE_PREFIX + str(df_obj.key)
+        if not session_id:
+            table_name = _TEMP_TABLE_PREFIX + str(df_obj.key)
+        else:
+            table_name = build_temp_table_name(session_id, df_obj.key)
     else:
         table_name = None
 

maxframe/io/odpsio/tableio.py

@@ -26,9 +26,9 @@ from odps.apis.storage_api import (
     TableBatchScanResponse,
     TableBatchWriteResponse,
 )
+from odps.errors import TableModified
 from odps.tunnel import TableDownloadSession, TableDownloadStatus, TableTunnel
 from odps.types import OdpsSchema, PartitionSpec, timestamp_ntz
-from odps.utils import call_with_retry
 
 try:
     import pyarrow.compute as pac
@@ -37,7 +37,7 @@ except ImportError:
 
 from ...config import options
 from ...env import ODPS_STORAGE_API_ENDPOINT
-from ...utils import is_empty, sync_pyodps_options
+from ...utils import call_with_retry, is_empty, sync_pyodps_options
 from .schema import odps_schema_to_arrow_schema
 
 PartitionsType = Union[List[str], str, None]
@@ -154,6 +154,32 @@ class TunnelMultiPartitionReader:
             return None
         return self._count
 
+    def _open_table_reader(self, partition: Optional[str], columns: List[str]):
+        attempts = 2
+        for trial in range(attempts):
+            try:
+                return self._table.open_reader(
+                    partition,
+                    columns=columns,
+                    arrow=True,
+                    download_id=self._partition_to_download_ids.get(partition),
+                    append_partitions=True,
+                )
+            except TableModified:
+                if trial == attempts - 1:
+                    raise
+                pt_to_session = TunnelTableIO.create_download_sessions(
+                    self._odps_entry,
+                    self._table.full_table_name,
+                    partition,
+                    reopen=True,
+                )
+                assert partition in pt_to_session
+                self._partition_to_download_ids[partition] = pt_to_session[partition].id
+        raise RuntimeError(
+            "Unexpected condition: all trial of open reader done and not raised"
+        )
+
     def _open_next_reader(self):
         if self._cur_reader is not None:
             self._reader_start_pos += self._cur_reader.count
@@ -170,12 +196,8 @@ class TunnelMultiPartitionReader:
             part_str = self._partitions[self._cur_partition_id]
             req_columns = self._schema.names
             with sync_pyodps_options():
-                self._cur_reader = self._table.open_reader(
-                    part_str,
-                    columns=req_columns,
-                    arrow=True,
-                    download_id=self._partition_to_download_ids.get(part_str),
-                    append_partitions=True,
+                self._cur_reader = self._open_table_reader(
+                    part_str, columns=req_columns
                 )
             if self._cur_reader.count + self._reader_start_pos > self._start:
                 start = self._start - self._reader_start_pos
@@ -193,13 +215,27 @@ class TunnelMultiPartitionReader:
 
     def read(self):
         with sync_pyodps_options():
+            is_first_batch = False
             if self._cur_reader is None:
+                is_first_batch = True
                 self._open_next_reader()
                 if self._cur_reader is None:
                     return None
             while self._cur_reader is not None:
                 try:
-                    batch = next(self._reader_iter)
+                    try:
+                        batch = next(self._reader_iter)
+                    except TableModified:
+                        if not is_first_batch:
+                            raise
+                        # clear download id cache to create new sessions
+                        self._partition_to_download_ids = dict()
+                        self._cur_reader = None
+                        self._open_next_reader()
+                        if self._cur_reader is None:
+                            return None
+                        batch = next(self._reader_iter)
+
                     if batch is not None:
                         if self._row_left is not None:
                             self._row_left -= batch.num_rows
@@ -222,6 +258,14 @@ class TunnelMultiPartitionReader:
 
 class TunnelTableIO(ODPSTableIO):
     _down_session_ids = OrderedDict()
+    _down_modified_time = dict()
+
+    @classmethod
+    def _get_modified_time(cls, odps_entry: ODPS, full_table_name, partition):
+        data_src = odps_entry.get_table(full_table_name)
+        if partition is not None:
+            data_src = data_src.partitions[partition]
+        return data_src.last_data_modified_time
 
     @classmethod
     def create_download_sessions(
@@ -229,6 +273,7 @@ class TunnelTableIO(ODPSTableIO):
         odps_entry: ODPS,
         full_table_name: str,
         partitions: List[Optional[str]] = None,
+        reopen: bool = False,
     ) -> Dict[Optional[str], TableDownloadSession]:
         table = odps_entry.get_table(full_table_name)
         tunnel = TableTunnel(odps_entry, quota_name=options.tunnel_quota_name)
@@ -240,9 +285,14 @@ class TunnelTableIO(ODPSTableIO):
         part_to_session = dict()
         for part in parts:
             part_key = (full_table_name, part)
+            modified_time = cls._get_modified_time(odps_entry, full_table_name, part)
             down_session = None
 
-            if part_key in cls._down_session_ids:
+            if (
+                not reopen
+                and part_key in cls._down_session_ids
+                and cls._down_modified_time.get(part_key) == modified_time
+            ):
                 down_id = cls._down_session_ids[part_key]
                 down_session = tunnel.create_download_session(
                     table, async_mode=True, partition_spec=part, download_id=down_id
@@ -256,8 +306,10 @@ class TunnelTableIO(ODPSTableIO):
                 )
 
             while len(cls._down_session_ids) >= _DOWNLOAD_ID_CACHE_SIZE:
-                cls._down_session_ids.popitem(False)
+                k, _ = cls._down_session_ids.popitem(False)
+                cls._down_modified_time.pop(k)
             cls._down_session_ids[part_key] = down_session.id
+            cls._down_modified_time[part_key] = modified_time
             part_to_session[part] = down_session
         return part_to_session
 

maxframe/io/odpsio/tests/test_arrow.py

@@ -17,8 +17,7 @@ import pandas as pd
 import pyarrow as pa
 import pytest
 
-from maxframe.lib.dtypes_extension import dict_
-
+from ....lib.dtypes_extension import dict_
 from ..arrow import arrow_to_pandas, pandas_to_arrow