maxframe 0.1.0b5__cp310-cp310-macosx_10_9_universal2.whl → 1.0.0__cp310-cp310-macosx_10_9_universal2.whl

maxframe_client/tests/test_fetcher.py

@@ -17,19 +17,32 @@ import uuid
 import numpy as np
 import pandas as pd
 import pyarrow as pa
+import pytest
 from odps import ODPS
 
 import maxframe.dataframe as md
-from maxframe.odpsio import HaloTableIO
+from maxframe.config import options
+from maxframe.io.odpsio import ODPSTableIO
 from maxframe.protocol import ODPSTableResultInfo, ResultType
 from maxframe.tests.utils import tn
 
 from ..fetcher import ODPSTableFetcher
 
 
-async def test_table_fetcher():
+@pytest.fixture
+def switch_table_io(request):
+    old_use_common_table = options.use_common_table
+    try:
+        options.use_common_table = request.param
+        yield
+    finally:
+        options.use_common_table = old_use_common_table
+
+
+@pytest.mark.parametrize("switch_table_io", [False, True], indirect=True)
+async def test_table_fetcher(switch_table_io):
     odps_entry = ODPS.from_environments()
-    halo_table_io = HaloTableIO(odps_entry)
+    halo_table_io = ODPSTableIO(odps_entry)
     fetcher = ODPSTableFetcher(odps_entry)
 
     data = pd.DataFrame(
@@ -58,6 +71,11 @@ async def test_table_fetcher():
     assert len(fetched) == 1000
     pd.testing.assert_frame_equal(raw_data, fetched)
 
+    result_info = ODPSTableResultInfo(ResultType.ODPS_TABLE, full_table_name=table_name)
+    fetched = await fetcher.fetch(tileable, result_info, [slice(None, 2000), None])
+    assert len(fetched) == 1000
+    pd.testing.assert_frame_equal(raw_data, fetched)
+
     result_info = ODPSTableResultInfo(ResultType.ODPS_TABLE, full_table_name=table_name)
     fetched = await fetcher.fetch(tileable, result_info, [2, None])
     assert len(fetched) == 1

maxframe_client/tests/test_session.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import time
-from typing import Dict
+from typing import Any, Dict
 
 import mock
 import numpy as np
@@ -23,7 +23,10 @@ from odps import ODPS
 
 import maxframe.dataframe as md
 import maxframe.remote as mr
+from maxframe.config import options
+from maxframe.config.config import option_context
 from maxframe.core import ExecutableTuple, TileableGraph
+from maxframe.errors import NoTaskServerResponseError
 from maxframe.lib.aio import stop_isolation
 from maxframe.protocol import ResultInfo
 from maxframe.serialization import RemoteException
@@ -35,6 +38,7 @@ from maxframe_framedriver.app.tests.test_framedriver_webapp import (  # noqa: F4
 )
 
 from ..clients.framedriver import FrameDriverClient
+from ..session.odps import MaxFrameRestCaller
 
 pytestmark = pytest.mark.maxframe_engine(["MCSQL", "SPE"])
 
@@ -82,15 +86,32 @@ def test_simple_run_dataframe(start_mock_session):
         session_id: str,
         dag: TileableGraph,
         managed_input_infos: Dict[str, ResultInfo] = None,
+        new_settings: Dict[str, Any] = None,
     ):
         assert len(dag) == 2
-        return await original_submit_dag(self, session_id, dag, managed_input_infos)
+        return await original_submit_dag(
+            self, session_id, dag, managed_input_infos, new_settings
+        )
+
+    no_task_server_raised = False
+    original_get_dag_info = MaxFrameRestCaller.get_dag_info
+
+    async def patched_get_dag_info(self, dag_id: str):
+        nonlocal no_task_server_raised
+
+        if not no_task_server_raised:
+            no_task_server_raised = True
+            raise NoTaskServerResponseError
+        return await original_get_dag_info(self, dag_id)
 
     df["H"] = "extra_content"
 
     with mock.patch(
         "maxframe_client.clients.framedriver.FrameDriverClient.submit_dag",
         new=patched_submit_dag,
+    ), mock.patch(
+        "maxframe_client.session.odps.MaxFrameRestCaller.get_dag_info",
+        new=patched_get_dag_info,
     ):
         result = df.execute().fetch()
         assert len(result) == 1000
@@ -112,13 +133,30 @@ def test_simple_run_dataframe(start_mock_session):
     )
     assert odps_entry.exist_table(build_temp_table_name(start_mock_session, key))
     del df
-    time.sleep(5)
+    retry_times = 10
+    while (
+        odps_entry.exist_table(
+            build_temp_table_name(start_mock_session, intermediate_key)
+        )
+        and retry_times > 0
+    ):
+        time.sleep(1)
+        retry_times -= 1
     assert not odps_entry.exist_table(
         build_temp_table_name(start_mock_session, intermediate_key)
     )
     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()
 
@@ -139,6 +177,25 @@ def test_run_empty_table(start_mock_session):
     empty_table.drop()
 
 
+def test_run_odps_query_without_schema(start_mock_session):
+    odps_entry = ODPS.from_environments()
+
+    table_name = tn("test_session_empty_table")
+    odps_entry.delete_table(table_name, if_exists=True)
+    test_table = odps_entry.create_table(table_name, "a double, b double", lifecycle=1)
+
+    with test_table.open_writer() as writer:
+        writer.write([123, 456])
+
+    df = md.read_odps_query(
+        f"select a, b, a + b as `special: name` from {table_name}", skip_schema=True
+    )
+    executed = df.execute().fetch()
+    assert len(executed.dtypes) == 3
+
+    test_table.drop()
+
+
 def test_run_dataframe_with_pd_source(start_mock_session):
     odps_entry = ODPS.from_environments()
 
@@ -171,19 +228,38 @@ 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
         assert len(result_df.columns) == 6
 
-        df = md.read_odps_table(table_obj, index_col="index").head(10).execute().fetch()
+        df = md.read_odps_table(table_obj, index_col="index").head(10).execute()
+        assert df.shape == (10, 5)
         assert len(df) == 10
         assert len(df.columns) == 5
     finally:
         odps_entry.delete_table(table_name, if_exists=True)
 
 
+def test_create_session_with_options(framedriver_app):  # noqa: F811
+    odps_entry = ODPS.from_environments()
+    framedriver_addr = f"mf://localhost:{framedriver_app.port}"
+    old_value = options.session.max_alive_seconds
+    session = None
+    try:
+        options.session.max_alive_seconds = 10
+        session = new_session(framedriver_addr, odps_entry=odps_entry)
+        session_id = session.session_id
+        session_conf = framedriver_app.session_manager.get_session_settings(session_id)
+        with option_context(session_conf) as session_options:
+            assert session_options.session.max_alive_seconds == 10
+    finally:
+        options.session.max_alive_seconds = old_value
+        if session is not None:
+            session.destroy()
+
+
 def test_run_and_fetch_series(start_mock_session):
     odps_entry = ODPS.from_environments()
 
@@ -210,7 +286,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
 
@@ -221,7 +312,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
 
@@ -244,7 +335,7 @@ def test_pivot_dataframe(start_mock_session):
     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)
+    assert pivot.shape == (4, 2)
     pd.testing.assert_index_equal(
         pivot.dtypes.index, pd.Index(["large", "small"], name="C")
     )
@@ -253,3 +344,13 @@ def test_pivot_dataframe(start_mock_session):
         values="D", index=["A", "B"], columns=["C"], aggfunc="sum"
     )
     pd.testing.assert_frame_equal(executed.to_pandas(), expected)
+
+
+def test_index_drop_duplicates(start_mock_session):
+    pd_idx = pd.Index(["lame", "cow", "lame", "beetle", "lame", "hippo"])
+    idx = md.Index(pd_idx)
+    executed = idx.drop_duplicates(keep="first").execute()
+    expected = pd_idx.drop_duplicates(keep="first")
+    pd.testing.assert_index_equal(
+        executed.to_pandas().sort_values(), expected.sort_values()
+    )

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)