maxframe 1.0.0rc2__cp311-cp311-macosx_10_9_universal2.whl → 1.0.0rc3__cp311-cp311-macosx_10_9_universal2.whl

maxframe_client/session/odps.py

@@ -26,7 +26,8 @@ import pandas as pd
 from odps import ODPS
 
 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 +37,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 +59,13 @@ 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,
+)
 
 from ..clients.framedriver import FrameDriverClient
 from ..fetcher import get_fetcher_cls
@@ -139,14 +152,9 @@ class MaxFrameSession(ToThreadMixin, IsolatedAsyncSession):
         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
-
+    def _upload_and_get_table_read_tileable(
+        self, t: TileableType
+    ) -> Optional[TileableType]:
         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(
@@ -193,8 +201,29 @@ 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
+        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"""
@@ -223,14 +252,21 @@ 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
 
     async def execute(self, *tileables, **kwargs) -> ExecutionInfo:
@@ -242,7 +278,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

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

@@ -247,7 +247,19 @@ 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]
+    pd.testing.assert_frame_equal(result, expected)
+
+
+def test_run_remote_success(oss_config, start_mock_session):
     def func(a, b):
         return a + b
 
@@ -258,7 +270,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
 

maxframe/core/entity/chunks.py

@@ -1,68 +0,0 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from ...serialization.serializables import BoolField, FieldTypes, TupleField
-from ...utils import tokenize
-from .core import Entity, EntityData
-
-
-class ChunkData(EntityData):
-    __slots__ = ()
-
-    is_broadcaster = BoolField("is_broadcaster", default=False)
-    # If the operator is a shuffle mapper, this flag indicates whether the current chunk is mapper chunk when
-    # the operator produce multiple chunks such as TensorUnique.
-    is_mapper = BoolField("is_mapper", default=None)
-    # optional fields
-    _index = TupleField("index", FieldTypes.uint32)
-
-    def __repr__(self):
-        if self.op.stage is None:
-            return (
-                f"{type(self).__name__} <op={type(self.op).__name__}, "
-                f"key={self.key}>"
-            )
-        else:
-            return (
-                f"{type(self).__name__} <op={type(self.op).__name__}, "
-                f"stage={self.op.stage.name}, key={self.key}>"
-            )
-
-    @property
-    def index(self):
-        return getattr(self, "_index", None)
-
-    @property
-    def device(self):
-        return self.op.device
-
-    def _update_key(self):
-        object.__setattr__(
-            self,
-            "_key",
-            tokenize(
-                type(self).__name__,
-                *(getattr(self, k, None) for k in self._keys_ if k != "_index"),
-            ),
-        )
-
-
-class Chunk(Entity):
-    _allow_data_type_ = (ChunkData,)
-
-    def __repr__(self):
-        return f"{type(self).__name__}({self._data.__repr__()})"
-
-
-CHUNK_TYPE = (Chunk, ChunkData)

maxframe/core/entity/fuse.py

@@ -1,73 +0,0 @@
-# Copyright 1999-2024 Alibaba Group Holding Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import numpy as np
-
-from ...serialization.serializables import ReferenceField
-from .chunks import CHUNK_TYPE, Chunk, ChunkData
-
-
-class FuseChunkData(ChunkData):
-    __slots__ = ("_inited",)
-
-    _chunk = ReferenceField(
-        "chunk", CHUNK_TYPE, on_serialize=lambda x: x.data if hasattr(x, "data") else x
-    )
-
-    def __init__(self, *args, **kwargs):
-        self._inited = False
-        super().__init__(*args, **kwargs)
-        self._extra_params = {}
-        self._inited = True
-
-    @property
-    def chunk(self):
-        return self._chunk
-
-    @property
-    def composed(self):
-        # for compatibility, just return the topological ordering,
-        # once we apply optimization on the subgraph,
-        # `composed` is not needed any more and should be removed then.
-        assert getattr(self._op, "fuse_graph", None) is not None
-        fuse_graph = self._op.fuse_graph
-        return list(fuse_graph.topological_iter())
-
-    def __getattr__(self, attr):
-        if not self._inited:
-            return object.__getattribute__(self, attr)
-        if attr in self._extra_params:
-            return self._extra_params[attr]
-        try:
-            return getattr(self._chunk, attr)
-        except AttributeError:
-            return object.__getattribute__(self, attr)
-
-    def __setattr__(self, attr, value):
-        if attr == "params":
-            self._chunk.params = value
-        else:
-            super().__setattr__(attr, value)
-
-    @property
-    def nbytes(self):
-        return np.prod(self.shape) * self.dtype.itemsize
-
-
-class FuseChunk(Chunk):
-    __slots__ = ()
-    _allow_data_type_ = (FuseChunkData,)
-
-
-FUSE_CHUNK_TYPE = (FuseChunkData, FuseChunk)