Add SerializationContext.clone method. Add pandas_serialization_context member that uses pickle for NumPy arrays with unsupported tensor types

Change-Id: Ia70c26954ff9ab3af435281bcafbf298c8c0cf28

Author: wesm

python/pyarrow/__init__.py

@@ -125,6 +125,7 @@
 localfs = LocalFileSystem.get_instance()
 
 from pyarrow.serialization import (_default_serialization_context,
+                                   pandas_serialization_context,
                                    register_default_serialization_handlers)
 
 import pyarrow.types as types

python/pyarrow/serialization.pxi

@@ -57,6 +57,23 @@ cdef class SerializationContext:
         self.custom_serializers = dict()
         self.custom_deserializers = dict()
 
+    def clone(self):
+        """
+        Return copy of this SerializationContext
+
+        Returns
+        -------
+        clone : SerializationContext
+        """
+        result = SerializationContext()
+        result.type_to_type_id = self.type_to_type_id.copy()
+        result.whitelisted_types = self.whitelisted_types.copy()
+        result.types_to_pickle = self.types_to_pickle.copy()
+        result.custom_serializers = self.custom_serializers.copy()
+        result.custom_deserializers = self.custom_deserializers.copy()
+
+        return result
+
     def register_type(self, type_, type_id,
                       custom_serializer=None, custom_deserializer=None):
         """EXPERIMENTAL: Add type to the list of types we can serialize.

python/pyarrow/serialization.py

@@ -31,6 +31,28 @@
     cloudpickle = pickle
 
 
+# ----------------------------------------------------------------------
+# Set up serialization for numpy with dtype object (primitive types are
+# handled efficiently with Arrow's Tensor facilities, see
+# python_to_arrow.cc)
+
+def _serialize_numpy_array_list(obj):
+    return obj.tolist(), obj.dtype.str
+
+
+def _deserialize_numpy_array_list(data):
+    return np.array(data[0], dtype=np.dtype(data[1]))
+
+
+def _serialize_numpy_array_pickle(obj):
+    pickled = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
+    return frombuffer(pickled)
+
+
+def _deserialize_numpy_array_pickle(data):
+    return pickle.loads(memoryview(data))
+
+
 def register_default_serialization_handlers(serialization_context):
 
     # ----------------------------------------------------------------------
@@ -81,22 +103,10 @@ def _deserialize_default_dict(data):
                                         custom_serializer=cloudpickle.dumps,
                                         custom_deserializer=cloudpickle.loads)
 
-    # ----------------------------------------------------------------------
-    # Set up serialization for numpy with dtype object (primitive types are
-    # handled efficiently with Arrow's Tensor facilities, see
-    # python_to_arrow.cc)
-
-    def _serialize_numpy_array(obj):
-        pickled = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
-        return frombuffer(pickled)
-
-    def _deserialize_numpy_array(data):
-        return pickle.loads(memoryview(data))
-
     serialization_context.register_type(
         np.ndarray, 'np.array',
-        custom_serializer=_serialize_numpy_array,
-        custom_deserializer=_deserialize_numpy_array)
+        custom_serializer=_serialize_numpy_array_list,
+        custom_deserializer=_deserialize_numpy_array_list)
 
     # ----------------------------------------------------------------------
     # Set up serialization for pandas Series and DataFrame
@@ -155,3 +165,10 @@ def _deserialize_torch_tensor(data):
 
 
 register_default_serialization_handlers(_default_serialization_context)
+
+pandas_serialization_context = _default_serialization_context.clone()
+
+pandas_serialization_context.register_type(
+    np.ndarray, 'np.array',
+    custom_serializer=_serialize_numpy_array_pickle,
+    custom_deserializer=_deserialize_numpy_array_pickle)

python/pyarrow/tests/test_serialization.py

@@ -212,11 +212,11 @@ def make_serialization_context():
 serialization_context = make_serialization_context()
 
 
-def serialization_roundtrip(value, f):
+def serialization_roundtrip(value, f, ctx=serialization_context):
     f.seek(0)
-    pa.serialize_to(value, f, serialization_context)
+    pa.serialize_to(value, f, ctx)
     f.seek(0)
-    result = pa.deserialize_from(f, None, serialization_context)
+    result = pa.deserialize_from(f, None, ctx)
     assert_equal(value, result)
 
     _check_component_roundtrip(value)
@@ -249,6 +249,7 @@ def test_primitive_serialization(large_memory_map):
     with pa.memory_map(large_memory_map, mode="r+") as mmap:
         for obj in PRIMITIVE_OBJECTS:
             serialization_roundtrip(obj, mmap)
+            serialization_roundtrip(obj, mmap, pa.pandas_serialization_context)
 
 
 def test_serialize_to_buffer():