maxframe 0.1.0b4__cp310-cp310-macosx_10_9_universal2.whl → 0.1.0b5__cp310-cp310-macosx_10_9_universal2.whl

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

@@ -99,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(
@@ -226,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)