maxframe 0.1.0b4__cp311-cp311-macosx_10_9_universal2.whl → 1.0.0__cp311-cp311-macosx_10_9_universal2.whl

{maxframe-0.1.0b4.dist-info → maxframe-1.0.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: bdist_wheel (0.44.0)
 Root-Is-Purelib: false
 Tag: cp311-cp311-macosx_10_9_universal2
 

maxframe_client/__init__.py

@@ -13,5 +13,4 @@
 # limitations under the License.
 
 from .clients.framedriver import FrameDriverClient
-from .clients.spe import SPEClient
 from .session import MaxFrameRestSession

maxframe_client/clients/framedriver.py

@@ -63,9 +63,12 @@ class FrameDriverClient:
         session_id: str,
         dag: TileableGraph,
         managed_input_infos: Dict[str, ResultInfo] = None,
+        new_settings: Dict[str, Any] = None,
     ) -> DagInfo:
         req_url = f"{self._endpoint}/api/sessions/{session_id}/dags"
-        req_body = ExecuteDagRequest(session_id, dag, managed_input_infos)
+        req_body = ExecuteDagRequest(
+            session_id, dag, managed_input_infos, new_settings=new_settings
+        )
         resp = await httpclient.AsyncHTTPClient().fetch(
             req_url,
             method="POST",

maxframe_client/fetcher.py

@@ -14,16 +14,24 @@
 
 from abc import ABC, abstractmethod
 from numbers import Integral
-from typing import Any, Dict, List, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
+import pandas as pd
 import pyarrow as pa
 from odps import ODPS
-from odps.models import ExternalVolume, PartedVolume
-from tornado import httpclient
+from odps.models import ExternalVolume
+from odps.tunnel import TableTunnel
 
 from maxframe.core import OBJECT_TYPE
-from maxframe.lib import wrapped_pickle as pickle
-from maxframe.odpsio import HaloTableIO, arrow_to_pandas, build_dataframe_table_meta
+from maxframe.dataframe.core import DATAFRAME_TYPE
+from maxframe.io.objects import get_object_io_handler
+from maxframe.io.odpsio import (
+    ODPSTableIO,
+    ODPSVolumeReader,
+    arrow_to_pandas,
+    build_dataframe_table_meta,
+    odps_schema_to_pandas_dtypes,
+)
 from maxframe.protocol import (
     DataFrameTableMeta,
     ODPSTableResultInfo,
@@ -31,8 +39,9 @@ from maxframe.protocol import (
     ResultInfo,
     ResultType,
 )
+from maxframe.tensor.core import TENSOR_TYPE
 from maxframe.typing_ import PandasObjectTypes, TileableType
-from maxframe.utils import ToThreadMixin
+from maxframe.utils import ToThreadMixin, sync_pyodps_options
 
 _result_fetchers: Dict[ResultType, Type["ResultFetcher"]] = dict()
 
@@ -52,6 +61,14 @@ class ResultFetcher(ABC):
     def __init__(self, odps_entry: ODPS):
         self._odps_entry = odps_entry
 
+    @abstractmethod
+    async def update_tileable_meta(
+        self,
+        tileable: TileableType,
+        info: ResultInfo,
+    ) -> None:
+        raise NotImplementedError
+
     @abstractmethod
     async def fetch(
         self,
@@ -66,6 +83,13 @@ class ResultFetcher(ABC):
 class NullFetcher(ResultFetcher):
     result_type = ResultType.NULL
 
+    async def update_tileable_meta(
+        self,
+        tileable: TileableType,
+        info: ResultInfo,
+    ) -> None:
+        return
+
     async def fetch(
         self,
         tileable: TileableType,
@@ -79,13 +103,79 @@ class NullFetcher(ResultFetcher):
 class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
     result_type = ResultType.ODPS_TABLE
 
+    def _get_table_comment(self, table_name: str) -> Optional[str]:
+        table = self._odps_entry.get_table(table_name)
+        return getattr(table, "comment", None)
+
+    async def update_tileable_meta(
+        self,
+        tileable: TileableType,
+        info: ODPSTableResultInfo,
+    ) -> None:
+        if (
+            isinstance(tileable, DATAFRAME_TYPE)
+            and tileable.dtypes is None
+            and info.table_meta is not None
+        ):
+            if info.table_meta.pd_column_dtypes is not None:
+                tileable.refresh_from_table_meta(info.table_meta)
+            else:
+                # need to get meta directly from table
+                table = self._odps_entry.get_table(info.full_table_name)
+                pd_dtypes = odps_schema_to_pandas_dtypes(table.table_schema).drop(
+                    info.table_meta.table_index_column_names
+                )
+                tileable.refresh_from_dtypes(pd_dtypes)
+
+        if tileable.shape and any(pd.isna(x) for x in tileable.shape):
+            part_specs = [None] if not info.partition_specs else info.partition_specs
+
+            with sync_pyodps_options():
+                table = self._odps_entry.get_table(info.full_table_name)
+                if isinstance(tileable, DATAFRAME_TYPE) and tileable.dtypes is None:
+                    dtypes = odps_schema_to_pandas_dtypes(table.table_schema)
+                    tileable.refresh_from_dtypes(dtypes)
+
+                tunnel = TableTunnel(self._odps_entry)
+                total_records = 0
+                for part_spec in part_specs:
+                    session = tunnel.create_download_session(table, part_spec)
+                    total_records += session.count
+
+            new_shape_list = list(tileable.shape)
+            new_shape_list[0] = total_records
+            tileable.params = {"shape": tuple(new_shape_list)}
+
+    @staticmethod
+    def _align_selection_with_shape(
+        row_sel: slice, shape: Tuple[Optional[int], ...]
+    ) -> dict:
+        size = shape[0]
+        if not row_sel.start and not row_sel.stop:
+            return {}
+        is_reversed = row_sel.step is not None and row_sel.step < 0
+        read_kw = {
+            "start": row_sel.start,
+            "stop": row_sel.stop,
+            "reverse_range": is_reversed,
+        }
+        if pd.isna(size):
+            return read_kw
+
+        if is_reversed and row_sel.start is not None:
+            read_kw["start"] = min(size - 1, row_sel.start)
+        if not is_reversed and row_sel.stop is not None:
+            read_kw["stop"] = min(size, row_sel.stop)
+        return read_kw
+
     def _read_single_source(
         self,
         table_meta: DataFrameTableMeta,
         info: ODPSTableResultInfo,
         indexes: List[Union[None, Integral, slice]],
+        shape: Tuple[Optional[int], ...],
     ):
-        table_io = HaloTableIO(self._odps_entry)
+        table_io = ODPSTableIO(self._odps_entry)
         read_kw = {}
         row_step = None
         if indexes:
@@ -93,13 +183,8 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
                 indexes += [None]
             row_sel, col_sel = indexes
             if isinstance(row_sel, slice):
-                if row_sel.start or row_sel.stop:
-                    read_kw["start"] = row_sel.start
-                    read_kw["stop"] = row_sel.stop
-                    read_kw["reverse_range"] = (
-                        row_sel.step is not None and row_sel.step < 0
-                    )
-                    row_step = row_sel.step
+                row_step = row_sel.step
+                read_kw = self._align_selection_with_shape(row_sel, shape)
             elif isinstance(row_sel, int):
                 read_kw["start"] = row_sel
                 read_kw["stop"] = row_sel + 1
@@ -118,8 +203,8 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
         with table_io.open_reader(
             info.full_table_name, info.partition_specs, **read_kw
         ) as reader:
-            reader_count = reader.count
             result = reader.read_all()
+            reader_count = result.num_rows
 
         if not row_step:
             return result
@@ -140,7 +225,7 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
     ) -> PandasObjectTypes:
         table_meta = build_dataframe_table_meta(tileable)
         arrow_table: pa.Table = await self.to_thread(
-            self._read_single_source, table_meta, info, indexes
+            self._read_single_source, table_meta, info, indexes, tileable.shape
         )
         return arrow_to_pandas(arrow_table, table_meta)
 
@@ -149,47 +234,31 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
 class ODPSVolumeFetcher(ToThreadMixin, ResultFetcher):
     result_type = ResultType.ODPS_VOLUME
 
-    async def _read_parted_volume_data(
-        self, volume: PartedVolume, partition: str, file_name: str
-    ) -> bytes:
-        def sync_read():
-            with volume.open_reader(partition, file_name) as reader:
-                return reader.read()
-
-        return await self.to_thread(sync_read)
-
-    async def _read_external_volume_data(
-        self, volume: ExternalVolume, path: str, file_name: str
-    ) -> bytes:
-        signed_url = await self.to_thread(
-            volume.get_sign_url, path + "/" + file_name, "GET"
-        )
-        http_client = httpclient.AsyncHTTPClient()
-
-        resp = await http_client.fetch(signed_url)
-        if hasattr(resp, "status_code") and resp.code >= 400:
-            try:
-                import oss2.exceptions
+    async def update_tileable_meta(
+        self,
+        tileable: TileableType,
+        info: ODPSVolumeResultInfo,
+    ) -> None:
+        return
 
-                oss_exc = oss2.exceptions.make_exception(resp.body)
-                raise oss_exc
-            except ImportError:
-                raise SystemError(resp.body)
-        return resp.body
+    async def _fetch_object(
+        self,
+        tileable: TileableType,
+        info: ODPSVolumeResultInfo,
+        indexes: List[Union[Integral, slice]],
+    ) -> Any:
+        def volume_fetch_func():
+            reader = ODPSVolumeReader(
+                self._odps_entry, info.volume_name, info.volume_path
+            )
+            io_handler = get_object_io_handler(tileable)()
+            return io_handler.read_object(reader, tileable, indexes)
 
-    async def _fetch_object(self, info: ODPSVolumeResultInfo) -> Any:
         volume = await self.to_thread(self._odps_entry.get_volume, info.volume_name)
-        if isinstance(volume, PartedVolume):
-            byte_data = await self._read_parted_volume_data(
-                volume, info.volume_path, "data"
-            )
-        elif isinstance(volume, ExternalVolume):
-            byte_data = await self._read_external_volume_data(
-                volume, info.volume_path, "data"
-            )
+        if isinstance(volume, ExternalVolume):
+            return await self.to_thread(volume_fetch_func)
         else:
             raise NotImplementedError(f"Volume type {type(volume)} not supported")
-        return pickle.loads(byte_data)
 
     async def fetch(
         self,
@@ -197,6 +266,6 @@ class ODPSVolumeFetcher(ToThreadMixin, ResultFetcher):
         info: ODPSVolumeResultInfo,
         indexes: List[Union[Integral, slice]],
     ) -> Any:
-        if isinstance(tileable, OBJECT_TYPE):
-            return await self._fetch_object(info)
+        if isinstance(tileable, (OBJECT_TYPE, TENSOR_TYPE)):
+            return await self._fetch_object(tileable, info, indexes)
         raise NotImplementedError(f"Fetching {type(tileable)} not implemented")

maxframe_client/session/consts.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# retry consts
+EMPTY_RESPONSE_RETRY_COUNT = 5
+
 # Restful Service
 RESTFUL_SESSION_INSECURE_SCHEME = "mf"
 RESTFUL_SESSION_SECURE_SCHEME = "mfs"

maxframe_client/session/graph.py

@@ -19,10 +19,16 @@ from dataclasses import dataclass
 from typing import Any, Dict, List, Tuple, Union
 from weakref import WeakSet
 
-from maxframe.core import ChunkType, TileableGraph, TileableType, enter_mode
+from maxframe.core import (
+    ChunkType,
+    TileableGraph,
+    TileableType,
+    build_fetch,
+    enter_mode,
+)
 from maxframe.core.operator import Fetch
 from maxframe.session import AbstractSession
-from maxframe.utils import build_fetch, copy_tileables
+from maxframe.utils import copy_tileables
 
 logger = logging.getLogger(__name__)