maxframe 1.0.0rc2__cp38-cp38-macosx_10_9_universal2.whl → 1.0.0rc4__cp38-cp38-macosx_10_9_universal2.whl

maxframe_client/fetcher.py

@@ -19,14 +19,18 @@ 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 odps.models import ExternalVolume
 from odps.tunnel import TableTunnel
-from tornado import httpclient
 
 from maxframe.core import OBJECT_TYPE
 from maxframe.dataframe.core import DATAFRAME_TYPE
-from maxframe.lib import wrapped_pickle as pickle
-from maxframe.odpsio import ODPSTableIO, arrow_to_pandas, build_dataframe_table_meta
+from maxframe.io.objects import get_object_io_handler
+from maxframe.io.odpsio import (
+    ODPSTableIO,
+    ODPSVolumeReader,
+    arrow_to_pandas,
+    build_dataframe_table_meta,
+)
 from maxframe.protocol import (
     DataFrameTableMeta,
     ODPSTableResultInfo,
@@ -36,7 +40,7 @@ from maxframe.protocol import (
 )
 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()
 
@@ -116,13 +120,15 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
 
         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
-            tunnel = TableTunnel(self._odps_entry)
-            total_records = 0
-            for part_spec in part_specs:
-                session = tunnel.create_download_session(
-                    info.full_table_name, part_spec
-                )
-                total_records += session.count
+
+            with sync_pyodps_options():
+                table = self._odps_entry.get_table(info.full_table_name)
+                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)}
@@ -222,47 +228,24 @@ class ODPSVolumeFetcher(ToThreadMixin, ResultFetcher):
     ) -> None:
         return
 
-    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
-
-                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,
@@ -271,5 +254,5 @@ class ODPSVolumeFetcher(ToThreadMixin, ResultFetcher):
         indexes: List[Union[Integral, slice]],
     ) -> Any:
         if isinstance(tileable, (OBJECT_TYPE, TENSOR_TYPE)):
-            return await self._fetch_object(info)
+            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__)
 

maxframe_client/session/odps.py

@@ -18,15 +18,17 @@ import logging
 import time
 import weakref
 from numbers import Integral
-from typing import Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
 from urllib.parse import urlparse
 
 import numpy as np
 import pandas as pd
 from odps import ODPS
+from odps import options as odps_options
 
 from maxframe.config import options
-from maxframe.core import Entity, TileableGraph, enter_mode
+from maxframe.core import Entity, TileableGraph, build_fetch, enter_mode
+from maxframe.core.operator import Fetch
 from maxframe.dataframe import read_odps_table
 from maxframe.dataframe.core import DATAFRAME_TYPE, SERIES_TYPE
 from maxframe.dataframe.datasource import PandasDataSourceOperator
@@ -36,11 +38,18 @@ from maxframe.errors import (
     NoTaskServerResponseError,
     SessionAlreadyClosedError,
 )
-from maxframe.odpsio import ODPSTableIO, pandas_to_arrow, pandas_to_odps_schema
+from maxframe.io.objects import get_object_io_handler
+from maxframe.io.odpsio import (
+    ODPSTableIO,
+    ODPSVolumeWriter,
+    pandas_to_arrow,
+    pandas_to_odps_schema,
+)
 from maxframe.protocol import (
     DagInfo,
     DagStatus,
     ODPSTableResultInfo,
+    ODPSVolumeResultInfo,
     ResultInfo,
     SessionInfo,
 )
@@ -51,8 +60,15 @@ from maxframe.session import (
     Profiling,
     Progress,
 )
+from maxframe.tensor.datasource import ArrayDataSource
 from maxframe.typing_ import TileableType
-from maxframe.utils import ToThreadMixin, build_temp_table_name
+from maxframe.utils import (
+    ToThreadMixin,
+    build_session_volume_name,
+    build_temp_table_name,
+    str_to_bool,
+    sync_pyodps_options,
+)
 
 from ..clients.framedriver import FrameDriverClient
 from ..fetcher import get_fetcher_cls
@@ -63,6 +79,43 @@ logger = logging.getLogger(__name__)
 
 
 class MaxFrameServiceCaller(metaclass=abc.ABCMeta):
+    def get_settings_to_upload(self) -> Dict[str, Any]:
+        sql_settings = (odps_options.sql.settings or {}).copy()
+        sql_settings.update(options.sql.settings or {})
+
+        quota_name = options.session.quota_name or getattr(
+            odps_options, "quota_name", None
+        )
+        lifecycle = options.session.table_lifecycle or odps_options.lifecycle
+        temp_lifecycle = (
+            options.session.temp_table_lifecycle or odps_options.temp_lifecycle
+        )
+
+        enable_schema = options.session.enable_schema
+        default_schema = options.session.default_schema
+        if hasattr(self, "_odps_entry"):
+            default_schema = default_schema or self._odps_entry.schema
+
+        # use flags in sql settings
+        if sql_settings.get("odps.default.schema"):
+            default_schema = sql_settings["odps.default.schema"]
+        if str_to_bool(
+            sql_settings.get("odps.namespace.schema") or "false"
+        ) or str_to_bool(
+            sql_settings.get("odps.sql.allow.namespace.schema") or "false"
+        ):
+            enable_schema = True
+
+        mf_settings = dict(options.to_dict(remote_only=True).items())
+        mf_settings["sql.settings"] = sql_settings
+        mf_settings["session.table_lifecycle"] = lifecycle
+        mf_settings["session.temp_table_lifecycle"] = temp_lifecycle
+        mf_settings["session.quota_name"] = quota_name
+        if enable_schema is not None:
+            mf_settings["session.enable_schema"] = enable_schema
+        mf_settings["session.default_schema"] = default_schema or "default"
+        return mf_settings
+
     @abc.abstractmethod
     def create_session(self) -> SessionInfo:
         raise NotImplementedError
@@ -73,7 +126,10 @@ class MaxFrameServiceCaller(metaclass=abc.ABCMeta):
 
     @abc.abstractmethod
     def submit_dag(
-        self, dag: TileableGraph, managed_input_infos: Dict[str, ResultInfo]
+        self,
+        dag: TileableGraph,
+        managed_input_infos: Dict[str, ResultInfo],
+        new_settings: Dict[str, Any] = None,
     ) -> DagInfo:
         raise NotImplementedError
 
@@ -127,6 +183,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         self._tileable_to_infos = weakref.WeakKeyDictionary()
 
         self._caller = self._create_caller(odps_entry, address, **kwargs)
+        self._last_settings = None
 
     @classmethod
     def _create_caller(
@@ -136,18 +193,14 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
 
     async def _init(self, _address: str):
         session_info = await self.ensure_async_call(self._caller.create_session)
+        self._last_settings = self._caller.get_settings_to_upload()
         self._session_id = session_info.session_id
         await self._show_logview_address()
 
-    def _upload_and_get_read_tileable(self, t: TileableType) -> Optional[TileableType]:
-        if (
-            not isinstance(t.op, PandasDataSourceOperator)
-            or t.op.get_data() is None
-            or t.inputs
-        ):
-            return None
-
-        schema, table_meta = pandas_to_odps_schema(t, unknown_as_string=True)
+    def _upload_and_get_table_read_tileable(
+        self, t: TileableType
+    ) -> Optional[TileableType]:
+        table_schema, table_meta = pandas_to_odps_schema(t, unknown_as_string=True)
         if self._odps_entry.exist_table(table_meta.table_name):
             self._odps_entry.delete_table(
                 table_meta.table_name, hints=options.sql.settings
@@ -155,7 +208,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         table_name = build_temp_table_name(self.session_id, t.key)
         table_obj = self._odps_entry.create_table(
             table_name,
-            schema,
+            table_schema,
             lifecycle=options.session.temp_table_lifecycle,
             hints=options.sql.settings,
         )
@@ -193,8 +246,30 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         read_tileable.params = t.params
         return read_tileable.data
 
+    def _upload_and_get_vol_read_tileable(
+        self, t: TileableType
+    ) -> Optional[TileableType]:
+        vol_name = build_session_volume_name(self.session_id)
+        writer = ODPSVolumeWriter(self._odps_entry, vol_name, t.key)
+        io_handler = get_object_io_handler(t)
+        io_handler().write_object(writer, t, t.op.data)
+        return build_fetch(t).data
+
+    def _upload_and_get_read_tileable(self, t: TileableType) -> Optional[TileableType]:
+        if (
+            not isinstance(t.op, (ArrayDataSource, PandasDataSourceOperator))
+            or t.op.get_data() is None
+            or t.inputs
+        ):
+            return None
+        with sync_pyodps_options():
+            if isinstance(t.op, PandasDataSourceOperator):
+                return self._upload_and_get_table_read_tileable(t)
+            else:
+                return self._upload_and_get_vol_read_tileable(t)
+
     @enter_mode(kernel=True, build=True)
-    def _scan_and_replace_pandas_sources(
+    def _scan_and_replace_local_sources(
         self, graph: TileableGraph
     ) -> Dict[TileableType, TileableType]:
         """Replaces Pandas data sources with temp table sources in the graph"""
@@ -215,7 +290,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
 
             for succ in successors:
                 graph.add_edge(replaced, succ)
-                succ.inputs = [replacements.get(t, t) for t in succ.inputs]
+                succ.op._set_inputs([replacements.get(t, t) for t in succ.inputs])
 
         graph.results = [replacements.get(t, t) for t in graph.results]
         return replacements
@@ -223,16 +298,41 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
     @enter_mode(kernel=True, build=True)
     def _get_input_infos(self, tileables: List[TileableType]) -> Dict[str, ResultInfo]:
         """Generate ResultInfo structs from generated temp tables"""
+        vol_name = build_session_volume_name(self.session_id)
+
         infos = dict()
         for t in tileables:
             key = t.key
-            if not isinstance(t.op, DataFrameReadODPSTable):
-                if not isinstance(t.inputs[0].op, DataFrameReadODPSTable):
-                    continue
-                t = t.inputs[0]
-            infos[key] = ODPSTableResultInfo(full_table_name=t.op.table_name)
+            if isinstance(t.op, DataFrameReadODPSTable):
+                infos[key] = ODPSTableResultInfo(full_table_name=t.op.table_name)
+            else:
+                if isinstance(t.op, Fetch):
+                    infos[key] = ODPSVolumeResultInfo(
+                        volume_name=vol_name, volume_path=t.key
+                    )
+                elif t.inputs and isinstance(t.inputs[0].op, DataFrameReadODPSTable):
+                    t = t.inputs[0]
+                    infos[key] = ODPSTableResultInfo(full_table_name=t.op.table_name)
         return infos
 
+    def _get_diff_settings(self) -> Dict[str, Any]:
+        new_settings = self._caller.get_settings_to_upload()
+        if not self._last_settings:  # pragma: no cover
+            self._last_settings = new_settings
+            return new_settings
+
+        update = dict()
+        for k in new_settings.keys():
+            old_item = self._last_settings.get(k)
+            new_item = new_settings.get(k)
+            try:
+                if old_item != new_item:
+                    update[k] = new_item
+            except:  # noqa: E722  # nosec  # pylint: disable=bare-except
+                update[k] = new_item
+        self._last_settings = new_settings
+        return update
+
     async def execute(self, *tileables, **kwargs) -> ExecutionInfo:
         tileables = [
             tileable.data if isinstance(tileable, Entity) else tileable
@@ -242,7 +342,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         tileable_graph, to_execute_tileables = gen_submit_tileable_graph(
             self, tileables, tileable_to_copied
         )
-        source_replacements = self._scan_and_replace_pandas_sources(tileable_graph)
+        source_replacements = self._scan_and_replace_local_sources(tileable_graph)
 
         # we need to manage uploaded data sources with refcounting mechanism
         # as nodes in tileable_graph are copied, we need to use original nodes
@@ -252,7 +352,10 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
 
         replaced_infos = self._get_input_infos(list(source_replacements.values()))
         dag_info = await self.ensure_async_call(
-            self._caller.submit_dag, tileable_graph, replaced_infos
+            self._caller.submit_dag,
+            tileable_graph,
+            replaced_infos,
+            self._get_diff_settings(),
         )
 
         await self._show_logview_address(dag_info.dag_id)
@@ -462,7 +565,8 @@ class MaxFrameRestCaller(MaxFrameServiceCaller):
     _client: FrameDriverClient
     _session_id: Optional[str]
 
-    def __init__(self, client: FrameDriverClient):
+    def __init__(self, odps_entry: ODPS, client: FrameDriverClient):
+        self._odps_entry = odps_entry
         self._client = client
         self._session_id = None
 
@@ -475,7 +579,10 @@ class MaxFrameRestCaller(MaxFrameServiceCaller):
         await self._client.delete_session(self._session_id)
 
     async def submit_dag(
-        self, dag: TileableGraph, managed_input_infos: Dict[str, ResultInfo]
+        self,
+        dag: TileableGraph,
+        managed_input_infos: Dict[str, ResultInfo] = None,
+        new_settings: Dict[str, Any] = None,
     ) -> DagInfo:
         return await self._client.submit_dag(self._session_id, dag, managed_input_infos)
 
@@ -515,7 +622,7 @@ class MaxFrameRestSession(MaxFrameSession):
 
     @classmethod
     def _create_caller(cls, odps_entry: ODPS, address: str, **kwargs):
-        return MaxFrameRestCaller(FrameDriverClient(address))
+        return MaxFrameRestCaller(odps_entry, FrameDriverClient(address))
 
 
 def register_session_schemes(overwrite: bool = False):

maxframe_client/session/task.py

@@ -16,7 +16,7 @@ import base64
 import json
 import logging
 import time
-from typing import Dict, List, Optional, Type, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 import msgpack
 from odps import ODPS
@@ -24,6 +24,12 @@ from odps import options as odps_options
 from odps.errors import parse_instance_error
 from odps.models import Instance, MaxFrameTask
 
+try:
+    from odps.errors import EmptyTaskInfoError
+except ImportError:  # pragma: no cover
+    # todo remove when pyodps>=0.12.0 is enforced
+    EmptyTaskInfoError = type("EmptyTaskInfoError", (Exception,), {})
+
 from maxframe.config import options
 from maxframe.core import TileableGraph
 from maxframe.errors import NoTaskServerResponseError, SessionAlreadyClosedError
@@ -36,6 +42,7 @@ except ImportError:
     mf_version = None
 
 from .consts import (
+    EMPTY_RESPONSE_RETRY_COUNT,
     MAXFRAME_DEFAULT_PROTOCOL,
     MAXFRAME_OUTPUT_JSON_FORMAT,
     MAXFRAME_OUTPUT_MAXFRAME_FORMAT,
@@ -92,6 +99,10 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             self._nested = True
             self._instance = odps_entry.get_instance(nested_instance_id)
 
+    @property
+    def instance(self):
+        return self._instance
+
     def _deserial_task_info_result(
         self, content: Union[bytes, str, dict], target_cls: Type[JsonSerializable]
     ):
@@ -125,16 +136,8 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             major_version=self._major_version,
             service_endpoint=self._odps_entry.endpoint,
         )
-
-        # merge sql options
-        sql_settings = (odps_options.sql.settings or {}).copy()
-        sql_settings.update(options.sql.settings or {})
-
-        mf_settings = dict(options.to_dict(remote_only=True).items())
-        mf_settings["sql.settings"] = sql_settings
-
         mf_opts = {
-            "odps.maxframe.settings": json.dumps(mf_settings),
+            "odps.maxframe.settings": json.dumps(self.get_settings_to_upload()),
             "odps.maxframe.output_format": self._output_format,
         }
         if mf_version:
@@ -189,18 +192,39 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             interval = min(max_interval, interval * 2)
 
     def _put_task_info(self, method_name: str, json_data: dict):
-        resp_data = self._instance.put_task_info(
-            self._task_name, method_name, json.dumps(json_data)
-        )
-        if not resp_data:
-            raise NoTaskServerResponseError(f"No response for request {method_name}")
-        return resp_data
+        for trial in range(EMPTY_RESPONSE_RETRY_COUNT):
+            try:
+                return self._instance.put_task_info(
+                    self._task_name,
+                    method_name,
+                    json.dumps(json_data),
+                    raise_empty=True,
+                )
+            except TypeError:  # pragma: no cover
+                # todo remove when pyodps>=0.12.0 is enforced
+                resp_data = self._instance.put_task_info(
+                    self._task_name, method_name, json.dumps(json_data)
+                )
+                if resp_data:
+                    return resp_data
+                else:
+                    raise NoTaskServerResponseError(
+                        f"No response for request {method_name}. "
+                        f"Instance ID: {self._instance.id}"
+                    )
+            except EmptyTaskInfoError as ex:
+                # retry when server returns HTTP 204, which is designed for retry
+                if ex.code != 204 or trial >= EMPTY_RESPONSE_RETRY_COUNT - 1:
+                    raise NoTaskServerResponseError(
+                        f"No response for request {method_name}. "
+                        f"Instance ID: {self._instance.id}. "
+                        f"Request ID: {ex.request_id}"
+                    ) from None
+                time.sleep(0.5)
 
     def get_session(self) -> SessionInfo:
         req_data = {"output_format": self._output_format}
-        serialized = self._instance.put_task_info(
-            self._task_name, MAXFRAME_TASK_GET_SESSION_METHOD, json.dumps(req_data)
-        )
+        serialized = self._put_task_info(MAXFRAME_TASK_GET_SESSION_METHOD, req_data)
         info: SessionInfo = self._deserial_task_info_result(serialized, SessionInfo)
         info.session_id = self._instance.id
         return info
@@ -217,13 +241,18 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
         self,
         dag: TileableGraph,
         managed_input_infos: Optional[Dict[str, ResultInfo]] = None,
+        new_settings: Dict[str, Any] = None,
     ) -> DagInfo:
+        new_settings_value = {
+            "odps.maxframe.settings": json.dumps(new_settings),
+        }
         req_data = {
             "protocol": MAXFRAME_DEFAULT_PROTOCOL,
             "dag": base64.b64encode(serialize_serializable(dag)).decode(),
             "managed_input_infos": base64.b64encode(
                 serialize_serializable(managed_input_infos)
             ).decode(),
+            "new_settings": json.dumps(new_settings_value),
             "output_format": self._output_format,
         }
         res = self._put_task_info(MAXFRAME_TASK_SUBMIT_DAG_METHOD, req_data)
@@ -276,7 +305,7 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
 class MaxFrameTaskSession(MaxFrameSession):
     schemes = [ODPS_SESSION_INSECURE_SCHEME, ODPS_SESSION_SECURE_SCHEME]
 
-    _instance: Instance
+    _caller: MaxFrameInstanceCaller
 
     @classmethod
     def _create_caller(
@@ -296,6 +325,15 @@ class MaxFrameTaskSession(MaxFrameSession):
             **kwargs,
         )
 
+    @property
+    def closed(self) -> bool:
+        if super().closed:
+            return True
+        if not self._caller or not self._caller.instance:
+            # session not initialized yet
+            return False
+        return self._caller.instance.is_terminated()
+
 
 def register_session_schemes(overwrite: bool = False):
     MaxFrameTaskSession.register_schemes(overwrite=overwrite)

maxframe_client/tests/test_fetcher.py

@@ -22,7 +22,7 @@ from odps import ODPS
 
 import maxframe.dataframe as md
 from maxframe.config import options
-from maxframe.odpsio import ODPSTableIO
+from maxframe.io.odpsio import ODPSTableIO
 from maxframe.protocol import ODPSTableResultInfo, ResultType
 from maxframe.tests.utils import tn
 

maxframe_client/tests/test_session.py

@@ -137,6 +137,15 @@ def test_simple_run_dataframe(start_mock_session):
     assert not odps_entry.exist_table(build_temp_table_name(start_mock_session, key))
 
 
+def test_run_and_fetch_slice(start_mock_session):
+    pd_df = pd.DataFrame(np.random.rand(1000, 5), columns=list("ABCDE"))
+    df = md.DataFrame(pd_df)
+    result = df.execute()
+
+    sliced = result.head(10).fetch()
+    assert len(sliced) == 10
+
+
 def test_run_empty_table(start_mock_session):
     odps_entry = ODPS.from_environments()
 
@@ -189,7 +198,7 @@ def test_run_dataframe_from_to_odps_table(start_mock_session):
     table_name = build_temp_table_name(start_mock_session, "tmp_save")
     table_obj = odps_entry.get_table(table_name)
     try:
-        md.to_odps_table(md.DataFrame(pd_df), table_obj).execute().fetch()
+        md.to_odps_table(md.DataFrame(pd_df), table_obj, lifecycle=1).execute().fetch()
         with table_obj.open_reader() as reader:
             result_df = reader.to_pandas()
         assert len(result_df) == 10
@@ -247,7 +256,22 @@ def test_run_and_fetch_series(start_mock_session):
         )
 
 
-def test_run_remote_success(start_mock_session):
+def test_execute_with_tensor(oss_config, start_mock_session):
+    pd_df = pd.DataFrame(
+        {"angles": [0, 3, 4], "degrees": [360, 180, 360]},
+        index=["circle", "triangle", "rectangle"],
+    )
+    df = md.DataFrame(pd_df)
+
+    result = (df - [1, 2]).execute().fetch()
+    expected = pd_df - [1, 2]
+    # TODO: currently the record order in tensor reading from table is the index
+    # sorting order
+    expected.sort_index(axis=0, inplace=True)
+    pd.testing.assert_frame_equal(result, expected, check_like=True)
+
+
+def test_run_remote_success(oss_config, start_mock_session):
     def func(a, b):
         return a + b
 
@@ -258,7 +282,7 @@ def test_run_remote_success(start_mock_session):
     assert result == 21
 
 
-def test_run_remote_error(start_mock_session):
+def test_run_remote_error(oss_config, start_mock_session):
     def func():
         raise ValueError