Patch pandas to make pickle compatible between 1.2 and 1.3 (#3047)

Author: wjsi

benchmarks/asv_bench/benchmarks/graph_assigner.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import random
+
 import mars.tensor as mt
 import mars.dataframe as md
 from mars.core.graph import TileableGraph, TileableGraphBuilder, ChunkGraphBuilder
@@ -25,7 +27,11 @@ class ChunkGraphAssignerSuite:
     Benchmark that times performance of chunk graph assigner
     """
 
+    repeat = 10
+
     def setup(self):
+        random.seed()
+
         num_rows = 10000
         df1 = md.DataFrame(
             mt.random.rand(num_rows, 4, chunk_size=10), columns=list("abcd")

mars/_utils.pxd

@@ -17,7 +17,7 @@ cdef class TypeDispatcher:
     cdef dict _handlers
     cdef dict _lazy_handlers
     cdef dict _inherit_handlers
-    cpdef object __weakref__
+    cdef object __weakref__
 
     cpdef void register(self, object type_, object handler)
     cpdef void unregister(self, object type_)
@@ -29,3 +29,5 @@ cpdef str to_str(s, encoding=*)
 cpdef bytes to_binary(s, encoding=*)
 cpdef unicode to_text(s, encoding=*)
 cpdef register_tokenizer(cls, handler)
+cpdef void reset_id_random_seed() except *
+cpdef bytes new_random_id(int byte_len)

mars/_utils.pyx

@@ -1,3 +1,4 @@
+# distutils: language = c++
 # Copyright 1999-2021 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -22,12 +23,16 @@ import uuid
 from datetime import date, datetime, timedelta, tzinfo
 from enum import Enum
 from functools import lru_cache, partial
+from random import getrandbits
 from weakref import WeakSet
 
 import numpy as np
 import pandas as pd
 import cloudpickle
 cimport cython
+from libc.stdint cimport uint_fast64_t
+from libc.stdlib cimport malloc, free
+from .lib.cython.libcpp cimport mt19937_64
 try:
     from pandas.tseries.offsets import Tick as PDTick
 except ImportError:
@@ -420,5 +425,46 @@ cdef class Timer:
         self.duration = time.time() - self._start
 
 
+cdef mt19937_64 _rnd_gen
+cdef bint _rnd_is_seed_set = False
+
+
+cpdef void reset_id_random_seed() except *:
+    cdef bytes seed_bytes
+    global _rnd_is_seed_set
+
+    seed_bytes = getrandbits(64).to_bytes(8, "little")
+    _rnd_gen.seed((<uint_fast64_t *><char *>seed_bytes)[0])
+    _rnd_is_seed_set = True
+
+
+cpdef bytes new_random_id(int byte_len):
+    cdef uint_fast64_t *res_ptr
+    cdef uint_fast64_t res_data[4]
+    cdef int i, qw_num = byte_len >> 3
+    cdef bytes res
+
+    if not _rnd_is_seed_set:
+        reset_id_random_seed()
+
+    if (qw_num << 3) < byte_len:
+        qw_num += 1
+
+    if qw_num <= 4:
+        # use stack memory to accelerate
+        res_ptr = res_data
+    else:
+        res_ptr = <uint_fast64_t *>malloc(qw_num << 3)
+
+    try:
+        for i in range(qw_num):
+            res_ptr[i] = _rnd_gen()
+        return <bytes>((<char *>&(res_ptr[0]))[:byte_len])
+    finally:
+        # free memory if allocated by malloc
+        if res_ptr != res_data:
+            free(res_ptr)
+
+
 __all__ = ['to_str', 'to_binary', 'to_text', 'TypeDispatcher', 'tokenize', 'tokenize_int',
-           'register_tokenizer', 'ceildiv', 'Timer']
+           'register_tokenizer', 'ceildiv', 'Timer', 'reset_id_random_seed', 'new_random_id']

mars/oscar/backends/mars/pool.py

@@ -25,9 +25,9 @@
 from types import TracebackType
 from typing import List
 
-from ....utils import get_next_port, dataslots, ensure_coverage
+from ....utils import get_next_port, dataslots, ensure_coverage, reset_id_random_seed
 from ..config import ActorPoolConfig
-from ..message import CreateActorMessage, reset_random_seed as reset_message_seed
+from ..message import CreateActorMessage
 from ..pool import MainActorPoolBase, SubActorPoolBase, _register_message_handler
 
 
@@ -168,7 +168,7 @@ def _start_sub_pool(
 
         # make sure enough randomness for every sub pool
         random.seed(uuid.uuid1().bytes)
-        reset_message_seed()
+        reset_id_random_seed()
 
         conf = actor_config.get_pool_config(process_index)
         suspend_sigint = conf["suspend_sigint"]

mars/oscar/backends/message.pyi

@@ -210,5 +210,4 @@ class DeserializeMessageFailed(RuntimeError):
     def __init__(self, message_id: bytes): ...
     def __str__(self): ...
 
-def reset_random_seed(): ...
 def new_message_id() -> bytes: ...

mars/oscar/backends/message.pyx

@@ -1,4 +1,3 @@
-# distutils: language = c++
 # Copyright 1999-2022 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,15 +13,12 @@
 # limitations under the License.
 
 from enum import Enum
-from random import getrandbits
 from types import TracebackType
 from typing import Any, Type
 
-from libc.stdint cimport uint_fast64_t
-
-from ...lib.cython.libcpp cimport mt19937_64
 from ...lib.tblib import pickling_support
 from ...serialization.core cimport Serializer
+from ..._utils cimport new_random_id
 from ...utils import wrap_exception
 from ..core cimport ActorRef
 
@@ -32,9 +28,6 @@ pickling_support.install()
 cdef int _DEFAULT_PROTOCOL = 0
 DEFAULT_PROTOCOL = _DEFAULT_PROTOCOL
 
-cdef mt19937_64 _rnd_gen
-cdef bint _rnd_is_seed_set = False
-
 
 class MessageType(Enum):
     control = 0
@@ -551,22 +544,5 @@ cdef class MessageSerializer(Serializer):
 MessageSerializer.register(_MessageBase)
 
 
-cpdef reset_random_seed():
-    cdef bytes seed_bytes
-    global _rnd_is_seed_set
-
-    seed_bytes = getrandbits(64).to_bytes(8, "little")
-    _rnd_gen.seed((<uint_fast64_t *><char *>seed_bytes)[0])
-    _rnd_is_seed_set = True
-
-
 cpdef bytes new_message_id():
-    cdef uint_fast64_t res_array[4]
-    cdef int i
-
-    if not _rnd_is_seed_set:
-        reset_random_seed()
-
-    for i in range(4):
-        res_array[i] = _rnd_gen()
-    return <bytes>((<char *>&(res_array[0]))[:32])
+    return new_random_id(32)

mars/oscar/utils.pyx

@@ -12,15 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from random import getrandbits
 from typing import AsyncGenerator
 
-from .._utils cimport to_str
+from .._utils cimport to_str, new_random_id
 from .core cimport ActorRef, LocalActorRef
 
 
 cpdef bytes new_actor_id():
-    return getrandbits(256).to_bytes(32, "little")
+    return new_random_id(32)
 
 
 def create_actor_ref(*args, **kwargs):
@@ -34,6 +33,7 @@ def create_actor_ref(*args, **kwargs):
 
     cdef str address
     cdef object uid
+    cdef ActorRef existing_ref
 
     address = to_str(kwargs.pop('address', None))
     uid = kwargs.pop('uid', None)
@@ -49,8 +49,9 @@ def create_actor_ref(*args, **kwargs):
     elif len(args) == 1:
         tp0 = type(args[0])
         if tp0 is ActorRef or tp0 is LocalActorRef:
-            uid = args[0].uid
-            address = to_str(address or args[0].address)
+            existing_ref = <ActorRef>(args[0])
+            uid = existing_ref.uid
+            address = to_str(address or existing_ref.address)
         else:
             uid = args[0]
 

mars/serialization/core.pyx

@@ -40,6 +40,15 @@ if sys.version_info[:2] < (3, 8):  # pragma: no cover
 else:
     import pickle  # nosec  # pylint: disable=import_pickle
 
+# resolve pandas pickle compatibility between <1.2 and >=1.3
+try:
+    from pandas.core.internals import blocks as pd_blocks
+    if not hasattr(pd_blocks, "new_block") and hasattr(pd_blocks, "make_block"):
+        # register missing func that would cause errors
+        pd_blocks.new_block = pd_blocks.make_block
+except (ImportError, AttributeError):
+    pass
+
 BUFFER_PICKLE_PROTOCOL = max(pickle.DEFAULT_PROTOCOL, 5)
 cdef bint HAS_PICKLE_BUFFER = pickle.HIGHEST_PROTOCOL >= 5
 cdef bint _PANDAS_HAS_MGR = hasattr(pd.Series([0]), "_mgr")
@@ -195,7 +204,8 @@ def buffered(func):
 
 
 def pickle_buffers(obj):
-    buffers = [None]
+    cdef list buffers = [None]
+
     if HAS_PICKLE_BUFFER:
 
         def buffer_cb(x):

mars/tests/test_utils.py

@@ -610,3 +610,9 @@ def __call__(self, *args, **kwargs):
 
     func = Func()
     assert get_func_token_values(func) == [func]
+
+
+@pytest.mark.parametrize("id_length", [0, 5, 32, 63])
+def test_gen_random_id(id_length):
+    rnd_id = utils.new_random_id(id_length)
+    assert len(rnd_id) == id_length

mars/utils.py

@@ -68,6 +68,8 @@
     tokenize_int,
     register_tokenizer,
     ceildiv,
+    reset_id_random_seed,
+    new_random_id,
     Timer,
 )
 from .lib.version import parse as parse_version
@@ -84,6 +86,8 @@
 tokenize = tokenize
 register_tokenizer = register_tokenizer
 ceildiv = ceildiv
+reset_id_random_seed = reset_id_random_seed
+new_random_id = new_random_id
 _create_task = asyncio.create_task
 
 
@@ -163,7 +167,7 @@ def on_serialize_nsplits(value: Tuple[Tuple[int]]):
         return None
     new_nsplits = []
     for dim_splits in value:
-        new_nsplits.append(tuple(None if np.isnan(v) else v for v in dim_splits))
+        new_nsplits.append(tuple(None if pd.isna(v) else v for v in dim_splits))
     return tuple(new_nsplits)