maxframe 0.1.0b5__cp311-cp311-macosx_10_9_universal2.whl → 1.0.0rc2__cp311-cp311-macosx_10_9_universal2.whl

maxframe/odpsio/tests/test_tableio.py

@@ -12,22 +12,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
+
 import numpy as np
 import pandas as pd
 import pyarrow as pa
+import pytest
 from odps import ODPS
 
+from ...config import options
 from ...tests.utils import flaky, tn
 from ...utils import config_odps_default_options
-from ..tableio import HaloTableIO
+from ..tableio import ODPSTableIO
+
+
+@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
 
 
 @flaky(max_runs=3)
-def test_empty_table_io():
+@pytest.mark.parametrize("switch_table_io", [False, True], indirect=True)
+def test_empty_table_io(switch_table_io):
     config_odps_default_options()
 
     o = ODPS.from_environments()
-    halo_table_io = HaloTableIO(o)
+    table_io = ODPSTableIO(o)
 
     # test read from empty table
     empty_table_name = tn("test_empty_table_halo_read")
@@ -35,42 +50,53 @@ def test_empty_table_io():
     tb = o.create_table(empty_table_name, "col1 string", lifecycle=1)
 
     try:
-        with halo_table_io.open_reader(empty_table_name) as reader:
+        with table_io.open_reader(empty_table_name) as reader:
             assert len(reader.read_all()) == 0
     finally:
         tb.drop()
 
 
 @flaky(max_runs=3)
-def test_table_io_without_parts():
+@pytest.mark.parametrize("switch_table_io", [False, True], indirect=True)
+def test_table_io_without_parts(switch_table_io):
     config_odps_default_options()
 
     o = ODPS.from_environments()
-    halo_table_io = HaloTableIO(o)
+    table_io = ODPSTableIO(o)
 
     # test read and write tables without partition
     no_part_table_name = tn("test_no_part_halo_write")
     o.delete_table(no_part_table_name, if_exists=True)
-    tb = o.create_table(
-        no_part_table_name, ",".join(f"{c} double" for c in "abcde"), lifecycle=1
-    )
+    col_desc = ",".join(f"{c} double" for c in "abcde") + ", f datetime"
+    tb = o.create_table(no_part_table_name, col_desc, lifecycle=1)
 
     try:
         pd_data = pd.DataFrame(np.random.rand(100, 5), columns=list("abcde"))
-        with halo_table_io.open_writer(no_part_table_name) as writer:
+        date_val = [
+            (
+                datetime.datetime.now().replace(microsecond=0)
+                + datetime.timedelta(seconds=i)
+            )
+            for i in range(100)
+        ]
+        pd_data["f"] = pd.Series(date_val, dtype="datetime64[ms]").dt.tz_localize(
+            options.local_timezone
+        )
+        with table_io.open_writer(no_part_table_name) as writer:
             writer.write(pa.Table.from_pandas(pd_data, preserve_index=False))
-        with halo_table_io.open_reader(no_part_table_name) as reader:
+        with table_io.open_reader(no_part_table_name) as reader:
             pd.testing.assert_frame_equal(reader.read_all().to_pandas(), pd_data)
     finally:
         tb.drop()
 
 
 @flaky(max_runs=3)
-def test_table_io_with_range_reader():
+@pytest.mark.parametrize("switch_table_io", [False, True], indirect=True)
+def test_table_io_with_range_reader(switch_table_io):
     config_odps_default_options()
 
     o = ODPS.from_environments()
-    halo_table_io = HaloTableIO(o)
+    table_io = ODPSTableIO(o)
 
     # test read and write tables without partition
     no_part_table_name = tn("test_no_part_halo_write")
@@ -81,15 +107,15 @@ def test_table_io_with_range_reader():
 
     try:
         pd_data = pd.DataFrame(np.random.rand(100, 5), columns=list("abcde"))
-        with halo_table_io.open_writer(no_part_table_name) as writer:
+        with table_io.open_writer(no_part_table_name) as writer:
             writer.write(pa.Table.from_pandas(pd_data, preserve_index=False))
 
-        with halo_table_io.open_reader(
+        with table_io.open_reader(
             no_part_table_name, start=None, stop=100, row_batch_size=10
         ) as reader:
             pd.testing.assert_frame_equal(reader.read_all().to_pandas(), pd_data)
 
-        with halo_table_io.open_reader(
+        with table_io.open_reader(
             no_part_table_name,
             start=-2,
             stop=-52,
@@ -105,11 +131,12 @@ def test_table_io_with_range_reader():
 
 
 @flaky(max_runs=3)
-def test_table_io_with_parts():
+@pytest.mark.parametrize("switch_table_io", [False, True], indirect=True)
+def test_table_io_with_parts(switch_table_io):
     config_odps_default_options()
 
     o = ODPS.from_environments()
-    halo_table_io = HaloTableIO(o)
+    table_io = ODPSTableIO(o)
 
     # test read and write tables with partition
     parted_table_name = tn("test_parted_halo_write")
@@ -122,11 +149,11 @@ def test_table_io_with_parts():
 
     try:
         pd_data = pd.DataFrame(np.random.rand(100, 5), columns=list("abcde"))
-        with halo_table_io.open_writer(parted_table_name, "pt=test") as writer:
+        with table_io.open_writer(parted_table_name, "pt=test") as writer:
             writer.write(pa.Table.from_pandas(pd_data, preserve_index=False))
-        with halo_table_io.open_reader(parted_table_name, "pt=test") as reader:
+        with table_io.open_reader(parted_table_name, "pt=test") as reader:
             pd.testing.assert_frame_equal(reader.read_all().to_pandas(), pd_data)
-        with halo_table_io.open_reader(
+        with table_io.open_reader(
             parted_table_name, "pt=test", partition_columns=True
         ) as reader:
             expected_data = pd_data.copy()

maxframe/protocol.py

@@ -32,6 +32,7 @@ from .serialization.serializables import (
     EnumField,
     FieldTypes,
     Float64Field,
+    Int32Field,
     ListField,
     ReferenceField,
     Serializable,
@@ -71,6 +72,9 @@ class DagStatus(enum.Enum):
     CANCELLING = 4
     CANCELLED = 5
 
+    def is_terminated(self):
+        return self in (DagStatus.CANCELLED, DagStatus.SUCCEEDED, DagStatus.FAILED)
+
 
 class DimensionIndex(Serializable):
     is_slice: bool = BoolField("is_slice", default=None)
@@ -150,6 +154,9 @@ class ODPSTableResultInfo(ResultInfo):
     partition_specs: Optional[List[str]] = ListField(
         "partition_specs", FieldTypes.string, default=None
     )
+    table_meta: Optional["DataFrameTableMeta"] = ReferenceField(
+        "table_meta", default=None
+    )
 
     def __init__(self, result_type: ResultType = None, **kw):
         result_type = result_type or ResultType.ODPS_TABLE
@@ -160,8 +167,17 @@ class ODPSTableResultInfo(ResultInfo):
         ret["full_table_name"] = self.full_table_name
         if self.partition_specs:
             ret["partition_specs"] = self.partition_specs
+        if self.table_meta:
+            ret["table_meta"] = self.table_meta.to_json()
         return ret
 
+    @classmethod
+    def _json_to_kwargs(cls, serialized: dict) -> dict:
+        kw = super()._json_to_kwargs(serialized)
+        if "table_meta" in kw:
+            kw["table_meta"] = DataFrameTableMeta.from_json(kw["table_meta"])
+        return kw
+
 
 class ODPSVolumeResultInfo(ResultInfo):
     _result_type = ResultType.ODPS_VOLUME
@@ -190,9 +206,9 @@ class ErrorInfo(JsonSerializable):
         "error_tracebacks", FieldTypes.list
     )
     raw_error_source: ErrorSource = EnumField(
-        "raw_error_source", ErrorSource, FieldTypes.int8
+        "raw_error_source", ErrorSource, FieldTypes.int8, default=None
     )
-    raw_error_data: Optional[Exception] = AnyField("raw_error_data")
+    raw_error_data: Optional[Exception] = AnyField("raw_error_data", default=None)
 
     @classmethod
     def from_exception(cls, exc: Exception):
@@ -201,20 +217,29 @@ class ErrorInfo(JsonSerializable):
         return cls(messages, tracebacks, ErrorSource.PYTHON, exc)
 
     def reraise(self):
-        if self.raw_error_source == ErrorSource.PYTHON:
+        if (
+            self.raw_error_source == ErrorSource.PYTHON
+            and self.raw_error_data is not None
+        ):
             raise self.raw_error_data
         raise RemoteException(self.error_messages, self.error_tracebacks, [])
 
     @classmethod
     def from_json(cls, serialized: dict) -> "ErrorInfo":
         kw = serialized.copy()
-        kw["raw_error_source"] = ErrorSource(serialized["raw_error_source"])
+        if kw.get("raw_error_source") is not None:
+            kw["raw_error_source"] = ErrorSource(serialized["raw_error_source"])
+        else:
+            kw["raw_error_source"] = None
+
         if kw.get("raw_error_data"):
             bufs = [base64.b64decode(s) for s in kw["raw_error_data"]]
             try:
                 kw["raw_error_data"] = pickle.loads(bufs[0], buffers=bufs[1:])
             except:
-                kw["raw_error_data"] = None
+                # both error source and data shall be None to make sure
+                # RemoteException is raised.
+                kw["raw_error_source"] = kw["raw_error_data"] = None
         return cls(**kw)
 
     def to_json(self) -> dict:
@@ -227,7 +252,12 @@ class ErrorInfo(JsonSerializable):
         if isinstance(self.raw_error_data, (PickleContainer, RemoteException)):
             err_data_bufs = self.raw_error_data.get_buffers()
         elif isinstance(self.raw_error_data, BaseException):
-            err_data_bufs = pickle_buffers(self.raw_error_data)
+            try:
+                err_data_bufs = pickle_buffers(self.raw_error_data)
+            except:
+                err_data_bufs = None
+                ret["raw_error_source"] = None
+
         if err_data_bufs:
             ret["raw_error_data"] = [
                 base64.b64encode(s).decode() for s in err_data_bufs
@@ -249,9 +279,17 @@ class DagInfo(JsonSerializable):
     error_info: Optional[ErrorInfo] = ReferenceField("error_info", default=None)
     start_timestamp: Optional[float] = Float64Field("start_timestamp", default=None)
     end_timestamp: Optional[float] = Float64Field("end_timestamp", default=None)
+    subdag_infos: Dict[str, "SubDagInfo"] = DictField(
+        "subdag_infos",
+        key_type=FieldTypes.string,
+        value_type=FieldTypes.reference,
+        default_factory=dict,
+    )
 
     @classmethod
-    def from_json(cls, serialized: dict) -> "DagInfo":
+    def from_json(cls, serialized: dict) -> Optional["DagInfo"]:
+        if serialized is None:
+            return None
         kw = serialized.copy()
         kw["status"] = DagStatus(kw["status"])
         if kw.get("tileable_to_result_infos"):
@@ -261,6 +299,10 @@ class DagInfo(JsonSerializable):
             }
         if kw.get("error_info"):
             kw["error_info"] = ErrorInfo.from_json(kw["error_info"])
+        if kw.get("subdag_infos"):
+            kw["subdag_infos"] = {
+                k: SubDagInfo.from_json(v) for k, v in kw["subdag_infos"].items()
+            }
         return DagInfo(**kw)
 
     def to_json(self) -> dict:
@@ -279,6 +321,8 @@ class DagInfo(JsonSerializable):
             }
         if self.error_info:
             ret["error_info"] = self.error_info.to_json()
+        if self.subdag_infos:
+            ret["subdag_infos"] = {k: v.to_json() for k, v in self.subdag_infos.items()}
         return ret
 
 
@@ -302,7 +346,9 @@ class SessionInfo(JsonSerializable):
     error_info: Optional[ErrorInfo] = ReferenceField("error_info", default=None)
 
     @classmethod
-    def from_json(cls, serialized: dict) -> "SessionInfo":
+    def from_json(cls, serialized: dict) -> Optional["SessionInfo"]:
+        if serialized is None:
+            return None
         kw = serialized.copy()
         if kw.get("dag_infos"):
             kw["dag_infos"] = {
@@ -320,7 +366,10 @@ class SessionInfo(JsonSerializable):
             "idle_timestamp": self.idle_timestamp,
         }
         if self.dag_infos:
-            ret["dag_infos"] = {k: v.to_json() for k, v in self.dag_infos.items()}
+            ret["dag_infos"] = {
+                k: v.to_json() if v is not None else None
+                for k, v in self.dag_infos.items()
+            }
         if self.error_info:
             ret["error_info"] = self.error_info.to_json()
         return ret
@@ -342,7 +391,25 @@ class ExecuteDagRequest(Serializable):
     )
 
 
-class SubDagInfo(Serializable):
+class SubDagSubmitInstanceInfo(JsonSerializable):
+    submit_reason: str = StringField("submit_reason")
+    instance_id: str = StringField("instance_id")
+    subquery_id: Optional[int] = Int32Field("subquery_id", default=None)
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "SubDagSubmitInstanceInfo":
+        return SubDagSubmitInstanceInfo(**serialized)
+
+    def to_json(self) -> dict:
+        ret = {
+            "submit_reason": self.submit_reason,
+            "instance_id": self.instance_id,
+            "subquery_id": self.subquery_id,
+        }
+        return ret
+
+
+class SubDagInfo(JsonSerializable):
     subdag_id: str = StringField("subdag_id")
     status: DagStatus = EnumField("status", DagStatus, FieldTypes.int8, default=None)
     progress: float = Float64Field("progress", default=None)
@@ -355,9 +422,52 @@ class SubDagInfo(Serializable):
         FieldTypes.reference,
         default_factory=dict,
     )
+    start_timestamp: Optional[float] = Float64Field("start_timestamp", default=None)
+    end_timestamp: Optional[float] = Float64Field("end_timestamp", default=None)
+    submit_instances: List[SubDagSubmitInstanceInfo] = ListField(
+        "submit_instances",
+        FieldTypes.reference,
+        default_factory=list,
+    )
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "SubDagInfo":
+        kw = serialized.copy()
+        kw["status"] = DagStatus(kw["status"])
+        if kw.get("tileable_to_result_infos"):
+            kw["tileable_to_result_infos"] = {
+                k: ResultInfo.from_json(s)
+                for k, s in kw["tileable_to_result_infos"].items()
+            }
+        if kw.get("error_info"):
+            kw["error_info"] = ErrorInfo.from_json(kw["error_info"])
+        if kw.get("submit_instances"):
+            kw["submit_instances"] = [
+                SubDagSubmitInstanceInfo.from_json(s) for s in kw["submit_instances"]
+            ]
+        return SubDagInfo(**kw)
+
+    def to_json(self) -> dict:
+        ret = {
+            "subdag_id": self.subdag_id,
+            "status": self.status.value,
+            "progress": self.progress,
+            "start_timestamp": self.start_timestamp,
+            "end_timestamp": self.end_timestamp,
+        }
+        if self.error_info:
+            ret["error_info"] = self.error_info.to_json()
+        if self.tileable_to_result_infos:
+            ret["tileable_to_result_infos"] = {
+                k: v.to_json() for k, v in self.tileable_to_result_infos.items()
+            }
+        if self.submit_instances:
+            ret["submit_instances"] = [i.to_json() for i in self.submit_instances]
+        return ret
 
 
 class ExecuteSubDagRequest(Serializable):
+    subdag_id: str = StringField("subdag_id")
     dag: TileableGraph = ReferenceField(
         "dag",
         on_serialize=SerializableGraph.from_graph,
@@ -371,7 +481,7 @@ class DecrefRequest(Serializable):
     keys: List[str] = ListField("keys", FieldTypes.string, default=None)
 
 
-class DataFrameTableMeta(Serializable):
+class DataFrameTableMeta(JsonSerializable):
     __slots__ = "_pd_column_names", "_pd_index_level_names"
 
     table_name: Optional[str] = StringField("table_name", default=None)
@@ -402,7 +512,7 @@ class DataFrameTableMeta(Serializable):
             self._pd_index_level_names = self.pd_index_dtypes.index.tolist()
             return self._pd_index_level_names
 
-    def __eq__(self, other: "Serializable") -> bool:
+    def __eq__(self, other: "DataFrameTableMeta") -> bool:
         if not isinstance(other, type(self)):
             return False
         for k in self._FIELDS:
@@ -413,3 +523,29 @@ class DataFrameTableMeta(Serializable):
             if not is_same:
                 return False
         return True
+
+    def to_json(self) -> dict:
+        b64_pk = lambda x: base64.b64encode(pickle.dumps(x))
+        ret = {
+            "table_name": self.table_name,
+            "type": self.type.value,
+            "table_column_names": self.table_column_names,
+            "table_index_column_names": self.table_index_column_names,
+            "pd_column_dtypes": b64_pk(self.pd_column_dtypes),
+            "pd_column_level_names": b64_pk(self.pd_column_level_names),
+            "pd_index_dtypes": b64_pk(self.pd_index_dtypes),
+        }
+        return ret
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "DataFrameTableMeta":
+        b64_upk = lambda x: pickle.loads(base64.b64decode(x))
+        serialized.update(
+            {
+                "type": OutputType(serialized["type"]),
+                "pd_column_dtypes": b64_upk(serialized["pd_column_dtypes"]),
+                "pd_column_level_names": b64_upk(serialized["pd_column_level_names"]),
+                "pd_index_dtypes": b64_upk(serialized["pd_index_dtypes"]),
+            }
+        )
+        return DataFrameTableMeta(**serialized)

maxframe/serialization/core.pxd

@@ -18,6 +18,9 @@ from libc.stdint cimport int32_t, uint64_t
 cdef class Serializer:
     cdef int _serializer_id
 
+    cpdef bint is_public_data_exist(self, dict context, object key)
+    cpdef put_public_data(self, dict context, object key, object value)
+    cpdef get_public_data(self, dict context, object key)
     cpdef serial(self, object obj, dict context)
     cpdef deserial(self, list serialized, dict context, list subs)
     cpdef on_deserial_error(

maxframe/serialization/core.pyi

@@ -29,6 +29,9 @@ class PickleContainer:
 
 class Serializer:
     serializer_id: int
+    def is_public_data_exist(self, context: Dict, key: Any) -> bool: ...
+    def put_public_data(self, context: Dict, key: Any, value: Any) -> None: ...
+    def get_public_data(self, context: Dict, key: Any) -> Any: ...
     def serial(self, obj: Any, context: Dict): ...
     def deserial(self, serialized: List, context: Dict, subs: List[Any]): ...
     def on_deserial_error(

maxframe/serialization/core.pyx

@@ -130,11 +130,30 @@ cdef Serializer get_deserializer(int32_t deserializer_id):
 
 cdef class Serializer:
     serializer_id = None
+    _public_data_context_key = 0x7fffffff - 1
 
     def __cinit__(self):
         # make the value can be referenced with C code
         self._serializer_id = self.serializer_id
 
+    cpdef bint is_public_data_exist(self, dict context, object key):
+        cdef dict public_dict = context.get(self._public_data_context_key, None)
+        if public_dict is None:
+            return False
+        return key in public_dict
+
+    cpdef put_public_data(self, dict context, object key, object value):
+        cdef dict public_dict = context.get(self._public_data_context_key, None)
+        if public_dict is None:
+            public_dict = context[self._public_data_context_key] = {}
+        public_dict[key] = value
+
+    cpdef get_public_data(self, dict context, object key):
+        cdef dict public_dict = context.get(self._public_data_context_key, None)
+        if public_dict is None:
+            return None
+        return public_dict.get(key)
+
     cpdef serial(self, object obj, dict context):
         """
         Returns intermediate serialization result of certain object.
@@ -993,17 +1012,20 @@ def serialize(obj, dict context = None):
     cdef list subs
     cdef bint final
     cdef _IdContextHolder id_context_holder = _IdContextHolder()
+    cdef tuple result
 
     context = context if context is not None else dict()
     serialized, subs, final = _serial_single(obj, context, id_context_holder)
     if final or not subs:
         # marked as a leaf node, return directly
-        return [{}, serialized], subs
-
-    serial_stack.append(_SerialStackItem(serialized, subs))
-    return _serialize_with_stack(
-        serial_stack, None, context, id_context_holder, result_bufs_list
-    )
+        result = [{}, serialized], subs
+    else:
+        serial_stack.append(_SerialStackItem(serialized, subs))
+        result = _serialize_with_stack(
+            serial_stack, None, context, id_context_holder, result_bufs_list
+        )
+    result[0][0]["_PUB"] = context.get(Serializer._public_data_context_key)
+    return result
 
 
 async def serialize_with_spawn(
@@ -1036,31 +1058,38 @@ async def serialize_with_spawn(
     cdef list subs
     cdef bint final
     cdef _IdContextHolder id_context_holder = _IdContextHolder()
+    cdef tuple result
 
     context = context if context is not None else dict()
     serialized, subs, final = _serial_single(obj, context, id_context_holder)
     if final or not subs:
         # marked as a leaf node, return directly
-        return [{}, serialized], subs
-
-    serial_stack.append(_SerialStackItem(serialized, subs))
+        result = [{}, serialized], subs
+    else:
+        serial_stack.append(_SerialStackItem(serialized, subs))
 
-    try:
-        result = _serialize_with_stack(
-            serial_stack, None, context, id_context_holder, result_bufs_list, spawn_threshold
-        )
-    except _SerializeObjectOverflow as ex:
-        result = await asyncio.get_running_loop().run_in_executor(
-            executor,
-            _serialize_with_stack,
-            serial_stack,
-            ex.cur_serialized,
-            context,
-            id_context_holder,
-            result_bufs_list,
-            0,
-            ex.num_total_serialized,
-        )
+        try:
+            result = _serialize_with_stack(
+                serial_stack,
+                None,
+                context,
+                id_context_holder,
+                result_bufs_list,
+                spawn_threshold,
+            )
+        except _SerializeObjectOverflow as ex:
+            result = await asyncio.get_running_loop().run_in_executor(
+                executor,
+                _serialize_with_stack,
+                serial_stack,
+                ex.cur_serialized,
+                context,
+                id_context_holder,
+                result_bufs_list,
+                0,
+                ex.num_total_serialized,
+            )
+    result[0][0]["_PUB"] = context.get(Serializer._public_data_context_key)
     return result
 
 

maxframe/serialization/exception.py

@@ -35,7 +35,7 @@ class RemoteException(MaxFrameError):
     def from_exception(cls, exc: Exception):
         try:
             buffers = pickle_buffers(exc)
-        except (TypeError, pickle.PicklingError):
+        except:
             logger.exception("Cannot pickle exception %s", exc)
             buffers = []
 

maxframe/serialization/pandas.py

@@ -176,11 +176,16 @@ class PdTimestampSerializer(Serializer):
 
 class PdTimedeltaSerializer(Serializer):
     def serial(self, obj: pd.Timedelta, context: Dict):
-        return [int(obj.seconds), obj.microseconds, obj.nanoseconds], [], True
+        return [int(obj.seconds), obj.microseconds, obj.nanoseconds, obj.days], [], True
 
     def deserial(self, serialized: List, context: Dict, subs: List):
+        days = 0 if len(serialized) < 4 else serialized[3]
+        seconds, microseconds, nanoseconds = serialized[:3]
         return pd.Timedelta(
-            seconds=serialized[0], microseconds=serialized[1], nanoseconds=serialized[2]
+            days=days,
+            seconds=seconds,
+            microseconds=microseconds,
+            nanoseconds=nanoseconds,
         )