maxframe 0.1.0b3__cp37-cp37m-win_amd64.whl → 0.1.0b5__cp37-cp37m-win_amd64.whl

maxframe_client/fetcher.py

@@ -12,16 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import base64
+import json
 from abc import ABC, abstractmethod
 from numbers import Integral
-from typing import Any, Dict, List, Type, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
+import pandas as pd
 import pyarrow as pa
 from odps import ODPS
 from odps.models import ExternalVolume, PartedVolume
+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 HaloTableIO, arrow_to_pandas, build_dataframe_table_meta
 from maxframe.protocol import (
@@ -31,8 +36,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, deserialize_serializable
 
 _result_fetchers: Dict[ResultType, Type["ResultFetcher"]] = dict()
 
@@ -52,6 +58,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 +80,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,6 +100,40 @@ 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:
+            tb_comment = await self.to_thread(
+                self._get_table_comment, info.full_table_name
+            )
+            if tb_comment:  # pragma: no branch
+                comment_data = json.loads(tb_comment)
+
+                table_meta: DataFrameTableMeta = deserialize_serializable(
+                    base64.b64decode(comment_data["table_meta"])
+                )
+                tileable.refresh_from_table_meta(table_meta)
+
+        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
+            new_shape_list = list(tileable.shape)
+            new_shape_list[-1] = total_records
+            tileable.params = {"shape": tuple(new_shape_list)}
+
     def _read_single_source(
         self,
         table_meta: DataFrameTableMeta,
@@ -149,6 +204,13 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
 class ODPSVolumeFetcher(ToThreadMixin, ResultFetcher):
     result_type = ResultType.ODPS_VOLUME
 
+    async def update_tileable_meta(
+        self,
+        tileable: TileableType,
+        info: ODPSVolumeResultInfo,
+    ) -> None:
+        return
+
     async def _read_parted_volume_data(
         self, volume: PartedVolume, partition: str, file_name: str
     ) -> bytes:
@@ -197,6 +259,6 @@ class ODPSVolumeFetcher(ToThreadMixin, ResultFetcher):
         info: ODPSVolumeResultInfo,
         indexes: List[Union[Integral, slice]],
     ) -> Any:
-        if isinstance(tileable, OBJECT_TYPE):
+        if isinstance(tileable, (OBJECT_TYPE, TENSOR_TYPE)):
             return await self._fetch_object(info)
         raise NotImplementedError(f"Fetching {type(tileable)} not implemented")

maxframe_client/session/odps.py

@@ -84,6 +84,9 @@ class MaxFrameServiceCaller(metaclass=abc.ABCMeta):
     def decref(self, tileable_keys: List[str]) -> None:
         raise NotImplementedError
 
+    def get_logview_address(self, dag_id=None, hours=None) -> Optional[str]:
+        return None
+
 
 class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
     _odps_entry: Optional[ODPS]
@@ -115,7 +118,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
     ):
         super().__init__(address, session_id)
         self.timeout = timeout
-        self._odps_entry = odps_entry or ODPS.from_environments()
+        self._odps_entry = odps_entry or ODPS.from_global() or ODPS.from_environments()
         self._tileable_to_infos = weakref.WeakKeyDictionary()
 
         self._caller = self._create_caller(odps_entry, address, **kwargs)
@@ -129,6 +132,7 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
     async def _init(self, _address: str):
         session_info = await self.ensure_async_call(self._caller.create_session)
         self._session_id = session_info.session_id
+        await self._show_logview_address()
 
     def _upload_and_get_read_tileable(self, t: TileableType) -> Optional[TileableType]:
         if (
@@ -142,20 +146,23 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         if self._odps_entry.exist_table(table_meta.table_name):
             self._odps_entry.delete_table(table_meta.table_name)
         table_name = build_temp_table_name(self.session_id, t.key)
-        table_obj = self._odps_entry.create_table(table_name, schema)
+        table_obj = self._odps_entry.create_table(
+            table_name, schema, lifecycle=options.session.temp_table_lifecycle
+        )
 
         data = t.op.get_data()
         batch_size = options.session.upload_batch_size
 
-        halo_client = HaloTableIO(self._odps_entry)
-        with halo_client.open_writer(table_obj.full_table_name) as writer:
-            for batch_start in range(0, len(data), batch_size):
-                if isinstance(data, pd.Index):
-                    batch = data[batch_start : batch_start + batch_size]
-                else:
-                    batch = data.iloc[batch_start : batch_start + batch_size]
-                arrow_batch, _ = pandas_to_arrow(batch)
-                writer.write(arrow_batch)
+        if len(data):
+            halo_client = HaloTableIO(self._odps_entry)
+            with halo_client.open_writer(table_obj.full_table_name) as writer:
+                for batch_start in range(0, len(data), batch_size):
+                    if isinstance(data, pd.Index):
+                        batch = data[batch_start : batch_start + batch_size]
+                    else:
+                        batch = data.iloc[batch_start : batch_start + batch_size]
+                    arrow_batch, _ = pandas_to_arrow(batch)
+                    writer.write(arrow_batch)
 
         read_tileable = read_odps_table(
             table_obj.full_table_name,
@@ -238,6 +245,8 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
             self._caller.submit_dag, tileable_graph, replaced_infos
         )
 
+        await self._show_logview_address(dag_info.dag_id)
+
         progress = Progress()
         profiling = Profiling()
         aio_task = asyncio.create_task(
@@ -293,6 +302,8 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
 
             for key, result_info in dag_info.tileable_to_result_infos.items():
                 t = key_to_tileables[key]
+                fetcher = get_fetcher_cls(result_info.result_type)(self._odps_entry)
+                await fetcher.update_tileable_meta(t, result_info)
                 self._tileable_to_infos[t] = result_info
 
     def _get_data_tileable_and_indexes(
@@ -387,6 +398,25 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
     async def get_mutable_tensor(self, name: str):
         raise NotImplementedError
 
+    async def get_logview_address(self, hours=None) -> Optional[str]:
+        return await self.get_dag_logview_address(None, hours)
+
+    async def get_dag_logview_address(self, dag_id=None, hours=None) -> Optional[str]:
+        return await self.ensure_async_call(
+            self._caller.get_logview_address, dag_id, hours
+        )
+
+    async def _show_logview_address(self, dag_id=None, hours=None):
+        identity = f"Session ID: {self._session_id}"
+        if dag_id:
+            identity += f", DAG ID: {dag_id}"
+
+        logview_addr = await self.get_dag_logview_address(dag_id, hours)
+        if logview_addr:
+            logger.info("%s, Logview: %s", identity, logview_addr)
+        else:
+            logger.info("%s, Logview address does not exist", identity)
+
 
 class MaxFrameRestCaller(MaxFrameServiceCaller):
     _client: FrameDriverClient

maxframe_client/session/task.py

@@ -21,9 +21,8 @@ from typing import Dict, List, Optional, Type, Union
 import msgpack
 from odps import ODPS
 from odps import options as odps_options
-from odps import serializers
 from odps.errors import parse_instance_error
-from odps.models import Instance, Task
+from odps.models import Instance, MaxFrameTask
 
 from maxframe.config import options
 from maxframe.core import TileableGraph
@@ -55,55 +54,6 @@ from .odps import MaxFrameServiceCaller, MaxFrameSession
 logger = logging.getLogger(__name__)
 
 
-class MaxFrameTask(Task):
-    __slots__ = ("_output_format", "_major_version", "_service_endpoint")
-    _root = "MaxFrame"
-    _anonymous_task_name = "AnonymousMaxFrameTask"
-
-    command = serializers.XMLNodeField("Command", default="CREATE_SESSION")
-
-    def __init__(self, **kwargs):
-        kwargs["name"] = kwargs.get("name") or self._anonymous_task_name
-        self._output_format = kwargs.pop(
-            "output_format", MAXFRAME_OUTPUT_MSGPACK_FORMAT
-        )
-        self._major_version = kwargs.pop("major_version", None)
-        self._service_endpoint = kwargs.pop("service_endpoint", None)
-        super().__init__(**kwargs)
-
-    def serial(self):
-        if self.properties is None:
-            self.properties = dict()
-
-        if odps_options.default_task_settings:
-            settings = odps_options.default_task_settings
-        else:
-            settings = dict()
-
-        if self._major_version is not None:
-            settings["odps.task.major.version"] = self._major_version
-
-        if "settings" in self.properties:
-            settings.update(json.loads(self.properties["settings"]))
-
-        # 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.output_format": self._output_format,
-            "odps.service.endpoint": self._service_endpoint,
-        }
-        if mf_version:
-            mf_opts["odps.maxframe.client_version"] = mf_version
-        settings.update(mf_opts)
-        self.properties["settings"] = json.dumps(settings)
-        return super().serial()
-
-
 class MaxFrameInstanceCaller(MaxFrameServiceCaller):
     _instance: Optional[Instance]
 
@@ -159,13 +109,31 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
                 f"Serialization format {self._output_format} not supported"
             )
 
-    def create_session(self) -> SessionInfo:
+    def _create_maxframe_task(self) -> MaxFrameTask:
         task = MaxFrameTask(
             name=self._task_name,
             major_version=self._major_version,
-            output_format=self._output_format,
             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.output_format": self._output_format,
+        }
+        if mf_version:
+            mf_opts["odps.maxframe.client_version"] = mf_version
+        task.update_settings(mf_opts)
+        return task
+
+    def create_session(self) -> SessionInfo:
+        task = self._create_maxframe_task()
         if not self._nested:
             self._task_name = task.name
             project = self._odps_entry.get_project(self._project)
@@ -278,6 +246,11 @@ class MaxFrameInstanceCaller(MaxFrameServiceCaller):
             self._task_name, MAXFRAME_TASK_DECREF_METHOD, json.dumps(req_data)
         )
 
+    def get_logview_address(self, dag_id=None, hours=None) -> Optional[str]:
+        hours = hours or options.session.logview_hours
+        subquery_suffix = f"&subQuery={dag_id}" if dag_id else ""
+        return self._instance.get_logview_address(hours) + subquery_suffix
+
 
 class MaxFrameTaskSession(MaxFrameSession):
     schemes = [ODPS_SESSION_INSECURE_SCHEME, ODPS_SESSION_SECURE_SCHEME]

maxframe_client/tests/test_session.py

@@ -28,6 +28,7 @@ from maxframe.lib.aio import stop_isolation
 from maxframe.protocol import ResultInfo
 from maxframe.serialization import RemoteException
 from maxframe.session import new_session
+from maxframe.tests.utils import tn
 from maxframe.utils import build_temp_table_name
 from maxframe_framedriver.app.tests.test_framedriver_webapp import (  # noqa: F401
     framedriver_app,
@@ -98,9 +99,12 @@ def test_simple_run_dataframe(start_mock_session):
     corner_top, corner_bottom = ExecutableTuple([df.iloc[:10], df.iloc[-10:]]).fetch()
     assert len(corner_top) == len(corner_bottom) == 10
 
-    # check ellipsis mark in DataFrame errors
+    # check ellipsis mark in DataFrame reprs
     df_str_repr = str(df)
     assert ".." in df_str_repr
+    # check ellipsis mark in Series reprs
+    series_str_repr = str(df.A.execute())
+    assert ".." in series_str_repr
 
     key = df.key
     assert odps_entry.exist_table(
@@ -115,6 +119,26 @@ 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_empty_table(start_mock_session):
+    odps_entry = ODPS.from_environments()
+
+    table_name = tn("test_session_empty_table")
+    odps_entry.delete_table(table_name, if_exists=True)
+    empty_table = odps_entry.create_table(
+        table_name, "_idx_0 bigint, a double, b double", lifecycle=1
+    )
+    df = md.read_odps_table(table_name, index_col="_idx_0")
+    df["d"] = df["a"] + df["b"]
+
+    executed = df.execute()
+    assert "Index: []" in str(executed)
+
+    fetched = executed.fetch()
+    assert 0 == len(fetched)
+
+    empty_table.drop()
+
+
 def test_run_dataframe_with_pd_source(start_mock_session):
     odps_entry = ODPS.from_environments()
 
@@ -205,3 +229,27 @@ def test_run_remote_error(start_mock_session):
 
     with pytest.raises((ValueError, RemoteException)):
         v.execute()
+
+
+def test_pivot_dataframe(start_mock_session):
+    pd_df = pd.DataFrame(
+        {
+            "A": "foo foo foo foo foo bar bar bar bar".split(),
+            "B": "one one one two two one one two two".split(),
+            "C": "small large large small small large small small large".split(),
+            "D": [1, 2, 2, 3, 3, 4, 5, 6, 7],
+            "E": [2, 4, 5, 5, 6, 6, 8, 9, 9],
+        }
+    )
+    df = md.DataFrame(pd_df)
+    pivot = df.pivot_table(values="D", index=["A", "B"], columns=["C"], aggfunc="sum")
+    executed = pivot.execute()
+    assert pivot.shape == (2, 4)
+    pd.testing.assert_index_equal(
+        pivot.dtypes.index, pd.Index(["large", "small"], name="C")
+    )
+
+    expected = pd_df.pivot_table(
+        values="D", index=["A", "B"], columns=["C"], aggfunc="sum"
+    )
+    pd.testing.assert_frame_equal(executed.to_pandas(), expected)