maxframe 2.0.0b2__cp38-cp38-win32.whl → 2.2.0__cp38-cp38-win32.whl

maxframe_client/fetcher.py

@@ -13,13 +13,16 @@
 # limitations under the License.
 
 import sys
+import warnings
 from abc import ABC, abstractmethod
 from numbers import Integral
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
+import numpy as np
 import pandas as pd
 import pyarrow as pa
 from odps import ODPS
+from odps.errors import ODPSError
 from odps.models import ExternalVolume
 
 from maxframe import utils
@@ -175,6 +178,36 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
         table = self._odps_entry.get_table(table_name)
         return getattr(table, "comment", None)
 
+    def _get_table_record_num(self, table, part_specs):
+        if not part_specs or part_specs == [None]:
+            rec_num = getattr(table, "record_num", None)
+        else:
+            rec_num = 0
+            for part_spec in part_specs:
+                pt_rec_num = getattr(table.get_partition(part_spec), "record_num", None)
+                if pt_rec_num is None or pt_rec_num < 0:
+                    rec_num = None
+                    break
+                else:
+                    rec_num += pt_rec_num
+
+        if rec_num is not None and rec_num >= 0:
+            return rec_num
+
+        try:
+            part_sessions = TunnelTableIO.create_download_sessions(
+                self._odps_entry, table.full_table_name, part_specs
+            )
+            return sum(session.count for session in part_sessions.values())
+        except ODPSError as ex:
+            if "StatusConflict" not in str(ex):
+                raise
+            warnings.warn(
+                f"Failed to obtain record count of table {table.full_table_name}. "
+                f"The original error is:\n{ex}"
+            )
+            return None
+
     async def update_tileable_meta(
         self,
         tileable: TileableType,
@@ -204,10 +237,8 @@ class ODPSTableFetcher(ToThreadMixin, ResultFetcher):
                     dtypes = odps_schema_to_pandas_dtypes(table.table_schema)
                     tileable.refresh_from_dtypes(dtypes)
 
-                part_sessions = TunnelTableIO.create_download_sessions(
-                    self._odps_entry, info.full_table_name, part_specs
-                )
-                total_records = sum(session.count for session in part_sessions.values())
+                total_records = self._get_table_record_num(table, part_specs)
+                total_records = np.nan if total_records is None else total_records
 
             new_shape_list = list(tileable.shape)
             new_shape_list[0] = total_records

maxframe_client/session/odps.py

@@ -91,10 +91,15 @@ logger = logging.getLogger(__name__)
 
 class MaxFrameServiceCaller(metaclass=abc.ABCMeta):
     def get_settings_to_upload(self) -> Dict[str, Any]:
+        odps_entry = getattr(self, "_odps_entry", None)
+        entry_quota_name = getattr(odps_entry, "quota_name", None)
+
         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
+        quota_name = (
+            options.session.quota_name
+            or entry_quota_name
+            or getattr(odps_options, "quota_name", None)
         )
         quota_settings = {
             sql_settings.get("odps.task.wlm.quota", None),

maxframe_client/tests/test_fetcher.py

@@ -14,11 +14,13 @@
 
 import uuid
 
+import mock
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 import pytest
 from odps import ODPS
+from odps.errors import ODPSError
 
 import maxframe.dataframe as md
 from maxframe.config import options
@@ -40,9 +42,9 @@ def switch_table_io(request):
 
 
 @pytest.mark.parametrize("switch_table_io", [False, True], indirect=True)
-async def test_table_fetcher(switch_table_io):
+async def test_fetch_table_data(switch_table_io):
     odps_entry = ODPS.from_environments()
-    halo_table_io = ODPSTableIO(odps_entry)
+    table_io = ODPSTableIO(odps_entry)
     fetcher = ODPSTableFetcher(odps_entry)
 
     data = pd.DataFrame(
@@ -60,7 +62,7 @@ async def test_table_fetcher(switch_table_io):
     odps_entry.create_table(
         table_name, "_idx_0 bigint, a double, b bigint, c string", lifecycle=1
     )
-    with halo_table_io.open_writer(table_name) as writer:
+    with table_io.open_writer(table_name) as writer:
         writer.write(pa.Table.from_pandas(data, preserve_index=False))
 
     raw_data = data[list("abc")]
@@ -105,3 +107,74 @@ async def test_table_fetcher(switch_table_io):
     pd.testing.assert_frame_equal(raw_data.iloc[-1:-6:-1, :1], fetched)
 
     odps_entry.delete_table(table_name, if_exists=True)
+
+
+async def test_fetch_data_shape():
+    odps_entry = ODPS.from_environments()
+    table_io = ODPSTableIO(odps_entry)
+    fetcher = ODPSTableFetcher(odps_entry)
+
+    data = pd.DataFrame(
+        {
+            "_idx_0": np.arange(1000),
+            "a": np.random.rand(1000),
+            "b": np.random.randint(0, 10, 1000),
+            "c": np.random.choice(list("ABC"), 1000),
+        }
+    )
+
+    table_names = [
+        tn("mf_test_groupby_table_" + str(uuid.uuid4().hex)) for _ in range(3)
+    ]
+    pt_specs = [None, None, ["pt='2020-01-01'"]]
+    for tn_to_check in table_names:
+        odps_entry.delete_table(tn_to_check, if_exists=True)
+
+    odps_entry.create_table(
+        table_names[0], "_idx_0 bigint, a double, b bigint, c string", lifecycle=1
+    )
+    odps_entry.create_table(
+        table_names[1], "_idx_0 bigint, a double, b bigint, c string", lifecycle=1
+    )
+    odps_entry.create_table(
+        table_names[2],
+        ("_idx_0 bigint, a double, b bigint, c string", "pt string"),
+        lifecycle=1,
+    )
+    with table_io.open_writer(table_names[0]) as writer:
+        writer.write(pa.Table.from_pandas(data, preserve_index=False))
+
+    odps_entry.execute_sql(
+        f"insert overwrite table {table_names[1]} select * from {table_names[0]}"
+    )
+    odps_entry.execute_sql(
+        f"insert overwrite table {table_names[2]} partition ({pt_specs[-1][0]}) "
+        f"select * from {table_names[0]}"
+    )
+
+    for tn_to_check, pt_spec in zip(table_names, pt_specs):
+        tileable = md.read_pandas(data[list("abc")])
+        tileable._shape = (np.nan, 3)
+        result_info = ODPSTableResultInfo(
+            ResultType.ODPS_TABLE, full_table_name=tn_to_check, partition_specs=pt_spec
+        )
+        await fetcher.update_tileable_meta(tileable, result_info)
+        assert tileable.shape == (1000, 3)
+
+    def _create_session_with_error(*_, **__):
+        raise ODPSError("StatusConflict: mock")
+
+    with mock.patch(
+        "maxframe.io.odpsio.tableio.TunnelTableIO.create_download_sessions",
+        new=_create_session_with_error,
+    ):
+        tileable = md.read_pandas(data[list("abc")])
+        tileable._shape = (np.nan, 3)
+        result_info = ODPSTableResultInfo(
+            ResultType.ODPS_TABLE, full_table_name=table_names[0]
+        )
+        await fetcher.update_tileable_meta(tileable, result_info)
+        assert tileable.shape == (np.nan, 3)
+
+    for tn_to_check in table_names:
+        odps_entry.delete_table(tn_to_check, if_exists=True, wait=False)

maxframe_client/tests/test_session.py

@@ -96,7 +96,10 @@ def test_local_run_dataframe(start_mock_session, enable_local_execution):
     cat_df = md.concat([df, df2])
     executed = cat_df.execute()
     result = executed.fetch()
-    pd.testing.assert_frame_equal(pd.concat([pd_result, pd_df2]), result)
+    try:
+        pd.testing.assert_frame_equal(pd.concat([pd_result, pd_df2]), result)
+    except AssertionError:
+        pd.testing.assert_frame_equal(pd.concat([pd_df2, pd_result]), result)
 
 
 def test_simple_run_dataframe(start_mock_session):