maxframe 2.0.0b2__cp310-cp310-macosx_10_9_universal2.whl → 2.2.0__cp310-cp310-macosx_10_9_universal2.whl

maxframe/lib/filesystem/tests/test_oss.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import time
 from io import BytesIO
 
@@ -21,7 +22,7 @@ import pytest
 from .. import oss
 from .._oss_lib import glob as og
 from .._oss_lib.common import OSSFileEntry
-from ..oss import build_oss_path
+from ..oss import HostEnforceType, _rewrite_internal_endpoint, build_oss_path
 
 
 class OSSObjInfo:
@@ -54,19 +55,25 @@ class MockObject:
 
 
 class SideEffectBucket:
+    cached_ctx = {}
+
     def __init__(self, *_, **__):
-        self.obj_dict = {
-            "file.csv": "id1,id2,id3\n1,2,3\n",
-            "dir/": "",
-            "dir/file1.csv": "2",
-            "dir/file2.csv": "3",
-            "dir/subdir/": "",
-            "dir/subdir/file3.csv": "s4",
-            "dir/subdir/file4.csv": "s5",
-            "dir2/": "",
-            "dir2/file6.csv": "6",
-            "dir2/file7.csv": "7",
-        }
+        cur_test = os.environ["PYTEST_CURRENT_TEST"]
+        if cur_test in self.cached_ctx:
+            self.obj_dict = self.cached_ctx[cur_test]
+        else:
+            self.obj_dict = self.cached_ctx[cur_test] = {
+                "file.csv": "id1,id2,id3\n1,2,3\n",
+                "dir/": "",
+                "dir/file1.csv": "2",
+                "dir/file2.csv": "3",
+                "dir/subdir/": "",
+                "dir/subdir/file3.csv": "s4",
+                "dir/subdir/file4.csv": "s5",
+                "dir2/": "",
+                "dir2/file6.csv": "6",
+                "dir2/file7.csv": "7",
+            }
 
     def get_object_meta(self, key):
         return ObjectMeta(key, self.obj_dict)
@@ -77,6 +84,17 @@ class SideEffectBucket:
     def get_object(self, key, byte_range):
         return MockObject(self.obj_dict, key, byte_range)
 
+    def copy_object(self, bucket, src_key, dst_key):
+        self.obj_dict[dst_key] = self.obj_dict[src_key]
+
+    def delete_object(self, key):
+        from oss2.exceptions import NoSuchKey
+
+        try:
+            del self.obj_dict[key]
+        except KeyError:
+            raise NoSuchKey(404, {}, key, {})
+
 
 class SideEffectObjIter:
     def __init__(self, *args, **kwargs):
@@ -96,28 +114,29 @@ def test_oss_filesystem(fake_obj_iter, fake_oss_bucket):
     access_key_secret = "your_access_key_secret"
     end_point = "your_endpoint"
 
-    file_path = f"oss://bucket/file.csv"
-    dir_path = f"oss://bucket/dir/"
-    dir_path_content_magic = f"oss://bucket/dir*/"
+    file_path = f"oss://your_endpoint/bucket/file.csv"
+    new_file_path = f"oss://your_endpoint/bucket/file1.csv"
+    dir_path = f"oss://your_endpoint/bucket/dir/"
+    dir_path_content_magic = f"oss://your_endpoint/bucket/dir*/"
     other_scheme_path = f"scheme://netloc/path"
-    not_exist_file_path = f"oss://bucket/not_exist.csv"
+    not_exist_file_path = f"oss://your_endpoint/bucket/not_exist.csv"
 
     fake_file_path = build_oss_path(
-        file_path, access_key_id, access_key_secret, end_point
+        file_path, end_point, access_key_id, access_key_secret
+    )
+    fake_new_file_path = build_oss_path(
+        new_file_path, end_point, access_key_id, access_key_secret
     )
     fake_dir_path = build_oss_path(
-        dir_path, access_key_id, access_key_secret, end_point
+        dir_path, end_point, access_key_id, access_key_secret
     )
     fake_dir_path_contains_magic = build_oss_path(
-        dir_path_content_magic, access_key_id, access_key_secret, end_point
-    )
-    fake_other_scheme_path = build_oss_path(
-        other_scheme_path, access_key_id, access_key_secret, end_point
+        dir_path_content_magic, end_point, access_key_id, access_key_secret
     )
     fake_not_exist_file_path = build_oss_path(
-        not_exist_file_path, access_key_id, access_key_secret, end_point
+        not_exist_file_path, end_point, access_key_id, access_key_secret
     )
-    fs = oss.OSSFileSystem.get_instance()
+    fs = oss.OSSFileSystem()
 
     # Test OSSFileSystem.
     assert len(fs.ls(fake_dir_path)) == 4
@@ -131,21 +150,15 @@ def test_oss_filesystem(fake_obj_iter, fake_oss_bucket):
     assert fs.stat(fake_dir_path)["type"] == "directory"
     assert fs.glob(fake_dir_path) == [fake_dir_path]
 
-    with pytest.raises(ValueError) as e:
-        fs.exists(fake_other_scheme_path)
-    msg1 = e.value.args[0]
-    assert (
-        msg1 == f"Except scheme oss, but got scheme: "
-        f"scheme in path: {fake_other_scheme_path}"
-    )
+    msg1 = f"Except scheme oss, but got scheme: scheme in path: {other_scheme_path}"
+    with pytest.raises(ValueError, match=msg1):
+        fs.exists(other_scheme_path)
 
-    with pytest.raises(RuntimeError) as e:
+    with pytest.raises(ValueError, match="No credentials provided"):
         fs.exists(file_path)
-    msg2 = e.value.args[0]
-    assert msg2 == "Please use build_oss_path to add OSS info"
 
     with pytest.raises(OSError):
-        print(fs.ls(fake_file_path))
+        fs.ls(fake_file_path)
 
     assert len(fs.glob(fake_file_path)) == 1
     assert len(fs.glob(fake_dir_path + "*", recursive=True)) == 4
@@ -180,3 +193,28 @@ def test_oss_filesystem(fake_obj_iter, fake_oss_bucket):
 
     fe = OSSFileEntry(fake_file_path)
     assert fe.path == fake_file_path
+
+    fs.rename(fake_file_path, fake_new_file_path)
+    assert not fs.exists(fake_file_path)
+    assert fs.exists(fake_new_file_path)
+
+    with pytest.raises(FileNotFoundError):
+        fs.delete(fake_not_exist_file_path)
+
+
+def test_host_rewrite():
+    assert "cn-shanghai.oss.service.com" == _rewrite_internal_endpoint(
+        "cn-shanghai.oss.service.com", HostEnforceType.force_external
+    )
+    assert "cn-shanghai.oss.service.com" == _rewrite_internal_endpoint(
+        "cn-shanghai-internal.oss.service.com", HostEnforceType.force_external
+    )
+    assert "cn-shanghai-internal.oss.service.com" == _rewrite_internal_endpoint(
+        "cn-shanghai.oss.service.com", HostEnforceType.force_internal
+    )
+    assert "cn-shanghai-internal.oss.service.com" == _rewrite_internal_endpoint(
+        "cn-shanghai-internal.oss.service.com", HostEnforceType.force_internal
+    )
+    assert "1.2.3.4" == _rewrite_internal_endpoint(
+        "1.2.3.4", HostEnforceType.force_internal
+    )

maxframe/lib/wrapped_pickle.py

@@ -75,6 +75,16 @@ class Unpickler(pickle_mod.Unpickler):
             raise ValueError("Unpickle is forbidden here")
         return super().load()
 
+    def find_class(self, module, name):
+        try:
+            return super().find_class(module, name)
+        except ImportError:
+            # workaround for pickle incompatibility since numpy>=2.0
+            if not module.startswith("numpy._core"):
+                raise
+            module = module.replace("numpy._core", "numpy.core")
+            return super().find_class(module, name)
+
 
 @functools.wraps(pickle_mod.load)
 def load(file, **kwargs):

maxframe/opcodes.py

@@ -271,6 +271,9 @@ SEM = 352
 STR_CONCAT = 353
 MAD = 354
 MEDIAN = 355
+RANK = 356
+IDXMAX = 357
+IDXMIN = 358
 
 # tensor operator
 RESHAPE = 401
@@ -389,6 +392,12 @@ ALIGN = 741
 CASE_WHEN = 742
 PIVOT = 743
 PIVOT_TABLE = 744
+TO_NUMERIC = 745
+DATAFRAME_FILTER = 746
+REORDER_LEVELS = 747
+DATAFRAME_COMPARE = 748
+DROPLEVEL = 749
+DATAFRAME_UPDATE = 750
 
 FUSE = 801
 
@@ -417,6 +426,7 @@ DATAFRAME_ILOC_GETITEM = 2021
 DATAFRAME_ILOC_SETITEM = 2022
 DATAFRAME_LOC_GETITEM = 2023
 DATAFRAME_LOC_SETITEM = 2024
+GET_LEVEL_VALUES = 2025
 
 # merge
 DATAFRAME_MERGE = 2010
@@ -426,7 +436,7 @@ DATAFRAME_SHUFFLE_MERGE_ALIGN = 2011
 DATAFRAME_BLOOM_FILTER = 2014
 
 # append
-APPEND = 2015
+APPEND = 2015  # deprecated since v2.2
 
 # reset index
 RESET_INDEX = 2028
@@ -439,12 +449,10 @@ GROUPBY_AGG = 2033
 GROUPBY_CONCAT = 2034
 GROUPBY_HEAD = 2035
 GROUPBY_SAMPLE_ILOC = 2036
-GROUPBY_SORT_REGULAR_SAMPLE = 2037
-GROUPBY_SORT_PIVOT = 2038
 GROUPBY_SORT_SHUFFLE = 2039
 
 # parallel sorting by regular sampling
-PSRS_SORT_REGULAR_SMAPLE = 2040
+PSRS_SORT_REGULAR_SAMPLE = 2040
 PSRS_CONCAT_PIVOT = 2041
 PSRS_SHUFFLE = 2042
 PSRS_ALIGN = 2043
@@ -475,6 +483,8 @@ TO_ODPS_TABLE = 20112
 READ_ODPS_VOLUME = 20113
 TO_ODPS_VOLUME = 20114
 READ_ODPS_QUERY = 20115
+READ_ODPS_MODEL = 20116  # reserved
+TO_ODPS_MODEL = 20117
 
 TO_CSV_STAT = 2102
 
@@ -506,9 +516,10 @@ FAISS_TRAIN_SAMPLED_INDEX = 2235
 FAISS_QUERY = 2236
 PROXIMA_SIMPLE_BUILDER = 2238
 PROXIMA_SIMPLE_SEARCHER = 2239
-KNEIGHBORS_GRAPH = 2237
+K_NEIGHBORS_GRAPH = 2237
 
 # cluster
+# k-means related
 KMEANS_PLUS_PLUS_INIT = 2250
 KMEANS_SCALABLE_PLUS_PLUS_INIT = 2251
 KMEANS_ELKAN_INIT_BOUNDS = 2252
@@ -517,7 +528,12 @@ KMEANS_ELKAN_POSTPROCESS = 2254
 KMEANS_LLOYD_UPDATE = 2255
 KMEANS_LLOYD_POSTPROCESS = 2256
 KMEANS_INERTIA = 2257
-KMEANS_RELOCASTE_EMPTY_CLUSTERS = 2258
+KMEANS_RELOCATE_EMPTY_CLUSTERS = 2258
+KMEANS_FIT = 2259
+KMEANS_PREDICT = 2260
+
+# linear models
+LOGISTIC_REGRESSION_FIT = 2270
 
 # XGBoost
 XGBOOST_TRAIN = 3001
@@ -580,6 +596,10 @@ RUN_SCRIPT = 5002
 # learn metrics
 MULTILABEL_CONFUSION_MATRIX = 5201
 PRECISION_RECALL_F_SCORE_SUPPORT = 5202
+AUC = 5203
+ROC_CURVE = 5204
+ROC_AUC_SCORE = 5205
+LOG_LOSS = 5206
 
 CHOLESKY_FUSE = 999988
 
@@ -588,23 +608,21 @@ DATAFRAME_RESHUFFLE = 10001
 FLATMAP = 10002
 FLATJSON = 10003
 APPLY_CHUNK = 10004
+EXTRACT_KV = 10013
+COLLECT_KV = 10014
+
+SERIES_DICT_METHOD = 10011
+SERIES_LIST_METHOD = 10012
+SERIES_STRUCT_METHOD = 10015
 
+# placeholders for compatibility, DO NOT REMOVE
 SERIES_DICT_GETITEM = 10005
 SERIES_DICT_SETITEM = 10006
 SERIES_DICT_LENGTH = 10007
 SERIES_DICT_REMOVE = 10008
 SERIES_DICT_CONTAINS = 10009
-SERIES_DICT_FLATTEN = 10010
-
 SERIES_LIST_GETITEM = 10020
-SERIES_LIST_SETITEM = 10021
-SERIES_LIST_CONTAINS = 10022
 SERIES_LIST_LENGTH = 10023
-SERIES_LIST_INSERT = 10024
-SERIES_LIST_EXTEND = 10025
-SERIES_LIST_POP = 10026
-SERIES_LIST_SORT = 10027
-SERIES_LIST_FLATTEN = 10028
 
 # MaxFrame internal operators
 DATAFRAME_PROJECTION_SAME_INDEX_MERGE = 100001

maxframe/protocol.py

@@ -39,6 +39,7 @@ from .serialization.serializables import (
     SeriesField,
     StringField,
 )
+from .utils import combine_error_message_and_traceback
 
 pickling_support.install()
 
@@ -244,6 +245,9 @@ class ErrorInfo(JsonSerializable):
         "raw_error_source", ErrorSource, FieldTypes.int8, default=None
     )
     raw_error_data: Optional[Exception] = AnyField("raw_error_data", default=None)
+    displayed_error_message: Optional[str] = StringField(
+        "displayed_error_message", default=None
+    )
 
     @classmethod
     def from_exception(cls, exc: Exception):
@@ -282,6 +286,7 @@ class ErrorInfo(JsonSerializable):
             "error_messages": self.error_messages,
             "error_tracebacks": self.error_tracebacks,
             "raw_error_source": self.raw_error_source.value,
+            "displayed_error_message": self.displayed_error_message,
         }
         err_data_bufs = None
         if isinstance(self.raw_error_data, (PickleContainer, RemoteException)):
@@ -299,6 +304,13 @@ class ErrorInfo(JsonSerializable):
             ]
         return ret
 
+    def get_displayed_error_message(self) -> str:
+        if self.displayed_error_message is not None:
+            return self.displayed_error_message
+        return combine_error_message_and_traceback(
+            self.error_messages, self.error_tracebacks
+        )
+
 
 class DagInfo(JsonSerializable):
     session_id: str = StringField("session_id", default=None)

maxframe/serialization/__init__.py

@@ -12,7 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from . import arrow, exception, maxframe_objects, numpy, pandas, scipy, serializables
+from . import (
+    arrow,
+    blob,
+    exception,
+    maxframe_objects,
+    numpy,
+    pandas,
+    scipy,
+    serializables,
+)
 from .core import (
     PickleContainer,
     PickleHookOptions,
@@ -27,4 +36,4 @@ from .core import (
 )
 from .exception import RemoteException
 
-del arrow, numpy, scipy, pandas, maxframe_objects, exception
+del arrow, blob, exception, maxframe_objects, numpy, pandas, scipy

maxframe/serialization/arrow.py

@@ -14,7 +14,7 @@
 
 from typing import Any, Dict, List, Union
 
-from ..utils import arrow_type_from_str
+from ..utils import arrow_type_from_str, extract_class_name
 from .core import Serializer, buffered
 
 try:
@@ -30,30 +30,54 @@ except ImportError:  # pragma: no cover
 
 _TYPE_CHAR_ARROW_ARRAY = "A"
 _TYPE_CHAR_ARROW_CHUNKED_ARRAY = "C"
+_TYPE_CHAR_ARROW_REDUCED = "R"
+
+
+class ArrowDataTypeSerializer(Serializer):
+    def serial(self, obj: pa.DataType, context):
+        return [str(obj)], [], True
+
+    def deserial(self, serialized, context, subs):
+        return arrow_type_from_str(serialized[0])
 
 
 class ArrowArraySerializer(Serializer):
     @buffered
     def serial(self, obj: PA_ARRAY_TYPES, context: Dict):
-        data_type = str(obj.type)
-        if isinstance(obj, pa.Array):
-            array_type = _TYPE_CHAR_ARROW_ARRAY
-            buffers = obj.buffers()
-            sizes = len(obj)
-        elif isinstance(obj, pa.ChunkedArray):
-            array_type = _TYPE_CHAR_ARROW_CHUNKED_ARRAY
-            buffers = [c.buffers() for c in obj.chunks]
-            sizes = [len(c) for c in obj.chunks]
-        else:  # pragma: no cover
+        if not isinstance(obj, (pa.Array, pa.ChunkedArray)):
             raise NotImplementedError(f"Array type {type(obj)} not supported")
-        return [array_type, data_type, sizes], buffers, True
+
+        if obj.type.num_fields == 0:
+            # use legacy serialization in case arrow changes deserializer method
+            data_type = str(obj.type)
+            if isinstance(obj, pa.Array):
+                array_type = _TYPE_CHAR_ARROW_ARRAY
+                buffers = obj.buffers()
+                sizes = len(obj)
+            else:  # ChunkedArray
+                array_type = _TYPE_CHAR_ARROW_CHUNKED_ARRAY
+                buffers = [c.buffers() for c in obj.chunks]
+                sizes = [len(c) for c in obj.chunks]
+            return [array_type, data_type, sizes], buffers, False
+
+        meth, extracted = obj.__reduce__()
+        meth_name = extract_class_name(meth)
+        return [_TYPE_CHAR_ARROW_REDUCED, meth_name, None], list(extracted), False
 
     def deserial(self, serialized: List, context: Dict, subs: List):
         array_type, data_type_str, sizes = serialized[:3]
-        data_type = arrow_type_from_str(data_type_str)
+        if array_type == _TYPE_CHAR_ARROW_REDUCED:
+            if data_type_str == "pyarrow.lib#chunked_array":
+                return pa.chunked_array(*subs)
+            elif data_type_str == "pyarrow.lib#_restore_array":
+                return pa.lib._restore_array(*subs)
+            else:
+                raise NotImplementedError(f"Unknown array type: {array_type}")
         if array_type == _TYPE_CHAR_ARROW_ARRAY:
+            data_type = arrow_type_from_str(data_type_str)
             return pa.Array.from_buffers(data_type, sizes, subs)
         elif array_type == _TYPE_CHAR_ARROW_CHUNKED_ARRAY:
+            data_type = arrow_type_from_str(data_type_str)
             chunks = [
                 pa.Array.from_buffers(data_type, size, bufs)
                 for size, bufs in zip(sizes, subs)
@@ -89,6 +113,7 @@ class ArrowBatchSerializer(Serializer):
 
 
 if pa is not None:  # pragma: no branch
+    ArrowDataTypeSerializer.register(pa.DataType)
     ArrowArraySerializer.register(pa.Array)
     ArrowArraySerializer.register(pa.ChunkedArray)
     ArrowBatchSerializer.register(pa.Table)

maxframe/serialization/blob.py

@@ -0,0 +1,32 @@
+# Copyright 1999-2025 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 typing import Dict
+
+from ..lib.dtypes_extension.blob import AbstractExternalBlob, SolidBlob
+from .core import Serializer
+
+
+class ExternalBlobSerializer(Serializer):
+    def serial(self, obj: AbstractExternalBlob, context: Dict):
+        _, vals = obj.__reduce__()
+        return [type(obj).__name__], list(vals), False
+
+    def deserial(self, serialized, context, subs):
+        cls_name = serialized[0]
+        cls = AbstractExternalBlob.get_cls_by_name(cls_name)
+        return cls(*subs)
+
+
+ExternalBlobSerializer.register(SolidBlob)

maxframe/serialization/core.pyx

@@ -52,6 +52,11 @@ try:
 except (ImportError, AttributeError):
     pass
 
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
+
 try:
     import pytz
     from pytz import BaseTzInfo as PyTZ_BaseTzInfo
@@ -95,6 +100,8 @@ cdef:
     int SLICE_SERIALIZER = 13
     int REGEX_SERIALIZER = 14
     int NO_DEFAULT_SERIALIZER = 15
+    int ARROW_BUFFER_SERIALIZER = 16
+    int RANGE_SERIALIZER = 17
     int PLACEHOLDER_SERIALIZER = 4096
 
 
@@ -874,12 +881,28 @@ cdef class SliceSerializer(Serializer):
     serializer_id = SLICE_SERIALIZER
 
     cpdef serial(self, object obj: slice, dict context):
-        return [obj.start, obj.stop, obj.step], [], True
+        cdef list elems = [obj.start, obj.stop, obj.step]
+        for x in elems:
+            if x is not None and not isinstance(x, int):
+                return [], elems, False
+        return elems, [], True
 
     cpdef deserial(self, list serialized, dict context, list subs):
+        if len(serialized) == 0:
+            return slice(subs[0], subs[1], subs[2])
         return slice(*serialized[:3])
 
 
+cdef class RangeSerializer(Serializer):
+    serializer_id = RANGE_SERIALIZER
+
+    cpdef serial(self, object obj: range, dict context):
+        return [obj.start, obj.stop, obj.step], [], True
+
+    cpdef deserial(self, list serialized, dict context, list subs):
+        return range(*serialized[:3])
+
+
 cdef class RegexSerializer(Serializer):
     serializer_id = REGEX_SERIALIZER
 
@@ -906,6 +929,18 @@ cdef class NoDefaultSerializer(Serializer):
         return no_default
 
 
+cdef class ArrowBufferSerializer(Serializer):
+    serializer_id = ARROW_BUFFER_SERIALIZER
+
+    cpdef serial(self, object obj, dict context):
+        return [], [obj], True
+
+    cpdef deserial(self, list obj, dict context, list subs):
+        if not isinstance(subs[0], pa.Buffer):
+            return pa.py_buffer(subs[0])
+        return subs[0]
+
+
 cdef class Placeholder:
     """
     Placeholder object to reduce duplicated serialization
@@ -959,8 +994,11 @@ DtypeSerializer.register(np.dtype)
 DtypeSerializer.register(ExtensionDtype)
 ComplexSerializer.register(complex)
 SliceSerializer.register(slice)
+RangeSerializer.register(range)
 RegexSerializer.register(re.Pattern)
 NoDefaultSerializer.register(NoDefault)
+if pa is not None:
+    ArrowBufferSerializer.register(pa.Buffer)
 PlaceholderSerializer.register(Placeholder)
 
 

maxframe/serialization/exception.py

@@ -18,6 +18,7 @@ from typing import Dict, List
 
 from ..errors import MaxFrameError
 from ..lib import wrapped_pickle as pickle
+from ..utils import combine_error_message_and_traceback
 from .core import Serializer, buffered, pickle_buffers, unpickle_buffers
 
 logger = logging.getLogger(__name__)
@@ -53,10 +54,7 @@ class RemoteException(MaxFrameError):
         return unpickle_buffers(self.buffers) if self.buffers else self
 
     def __str__(self):
-        tbs = []
-        for msg, tb in zip(self.messages, self.tracebacks):
-            tbs.append("".join([msg + "\n"] + tb))
-        return "\nCaused by:\n".join(tbs)
+        return combine_error_message_and_traceback(self.messages, self.tracebacks)
 
 
 class ExceptionSerializer(Serializer):

maxframe/serialization/numpy.py

@@ -95,5 +95,16 @@ class NDArraySerializer(Serializer):
         return val
 
 
+class RandomStateSerializer(Serializer):
+    def serial(self, obj: np.random.RandomState, context: Dict):
+        return [], [obj.get_state()], False
+
+    def deserial(self, serialized, context: Dict, subs: List):
+        rs = np.random.RandomState()
+        rs.set_state(subs[0])
+        return rs
+
+
 NDArraySerializer.register(np.generic)
 NDArraySerializer.register(np.ndarray)
+RandomStateSerializer.register(np.random.RandomState)