make custom serialization handlers accessible outside of the tests

python/pyarrow/__init__.py

@@ -116,6 +116,8 @@ from pyarrow.ipc import (Message, MessageReader,
 
 localfs = LocalFileSystem.get_instance()
 
+from pyarrow.serialization import _default_serialization_context
+
 import pyarrow.types as types
 
 # Entry point for starting the plasma store

python/pyarrow/ipc.py

@@ -187,50 +187,3 @@ def deserialize_pandas(buf, nthreads=None):
     reader = pa.RecordBatchStreamReader(buffer_reader)
     table = reader.read_all()
     return table.to_pandas(nthreads=nthreads)
-
-
-# ----------------------------------------------------------------------
-# Set up default serialization context
-
-def _serialize_pandas_series(s):
-    import pandas as pd
-    # TODO: serializing Series without extra copy
-    serialized = serialize_pandas(pd.DataFrame({s.name: s}))
-    return {
-        'type': 'Series',
-        'data': serialized.to_pybytes()
-    }
-
-
-def _serialize_pandas_dataframe(df):
-    return {
-        'type': 'DataFrame',
-        'data': serialize_pandas(df).to_pybytes()
-    }
-
-
-def _deserialize_callback_pandas(data):
-    deserialized = deserialize_pandas(data['data'])
-    type_ = data['type']
-    if type_ == 'Series':
-        return deserialized[deserialized.columns[0]]
-    elif type_ == 'DataFrame':
-        return deserialized
-    else:
-        raise ValueError(type_)
-
-
-try:
-    import pandas as pd
-    lib._default_serialization_context.register_type(
-        pd.Series, 'pandas.Series',
-        custom_serializer=_serialize_pandas_series,
-        custom_deserializer=_deserialize_callback_pandas)
-
-    lib._default_serialization_context.register_type(
-        pd.DataFrame, 'pandas.DataFrame',
-        custom_serializer=_serialize_pandas_dataframe,
-        custom_deserializer=_deserialize_callback_pandas)
-except ImportError:
-    # no pandas
-    pass

python/pyarrow/serialization.pxi

@@ -43,6 +43,9 @@ class DeserializationCallbackError(ArrowSerializationError):
         self.type_id = type_id
 
 
+_default_serialization_context = SerializationContext()
+
+
 cdef class SerializationContext:
     cdef:
         object type_to_type_id
@@ -138,9 +141,6 @@ cdef class SerializationContext:
         return obj
 
 
-_default_serialization_context = SerializationContext()
-
-
 cdef class SerializedPyObject:
     """
     Arrow-serialized representation of Python object

python/pyarrow/serialization.py

@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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 collections import OrderedDict, defaultdict
+import sys
+
+from pyarrow import serialize_pandas, deserialize_pandas
+from pyarrow.lib import _default_serialization_context
+
+# ----------------------------------------------------------------------
+# Set up serialization for primitive datatypes
+
+# TODO(pcm): This is currently a workaround until arrow supports
+# arbitrary precision integers. This is only called on long integers,
+# see the associated case in the append method in python_to_arrow.cc
+_default_serialization_context.register_type(
+    int, "int",
+    custom_serializer=lambda obj: str(obj),
+    custom_deserializer=lambda data: int(data))
+
+if (sys.version_info < (3, 0)):
+    _default_serialization_context.register_type(
+        long, "long",  # noqa: F821
+        custom_serializer=lambda obj: str(obj),
+        custom_deserializer=lambda data: long(data))  # noqa: F821
+
+
+def _serialize_ordered_dict(obj):
+    return list(obj.keys()), list(obj.values())
+
+
+def _deserialize_ordered_dict(data):
+    return OrderedDict(zip(data[0], data[1]))
+
+
+_default_serialization_context.register_type(
+    OrderedDict, "OrderedDict",
+    custom_serializer=_serialize_ordered_dict,
+    custom_deserializer=_deserialize_ordered_dict)
+
+
+def _serialize_default_dict(obj):
+    return list(obj.keys()), list(obj.values()), obj.default_factory
+
+
+def _deserialize_default_dict(data):
+    return defaultdict(data[2], zip(data[0], data[1]))
+
+
+_default_serialization_context.register_type(
+     defaultdict, "defaultdict",
+     custom_serializer=_serialize_default_dict,
+     custom_deserializer=_deserialize_default_dict)
+
+
+_default_serialization_context.register_type(
+     type(lambda: 0), "function",
+     pickle=True)
+
+# ----------------------------------------------------------------------
+# Set up serialization for numpy with dtype object (primitive types are
+# handled efficiently with Arrow's Tensor facilities, see python_to_arrow.cc)
+
+try:
+    import numpy as np
+
+    def _serialize_numpy_array(obj):
+        return obj.tolist(), obj.dtype.str
+
+    def _deserialize_numpy_array(data):
+        return np.array(data[0], dtype=np.dtype(data[1]))
+
+    _default_serialization_context.register_type(
+        np.ndarray, 'np.array',
+        custom_serializer=_serialize_numpy_array,
+        custom_deserializer=_deserialize_numpy_array)
+
+except ImportError:
+    # no numpy
+    pass
+
+# ----------------------------------------------------------------------
+# Set up serialization for pandas Series and DataFrame
+
+try:
+    import pandas as pd
+
+    def _serialize_pandas_series(obj):
+        # TODO: serializing Series without extra copy
+        return serialize_pandas(pd.DataFrame({obj.name: obj})).to_pybytes()
+
+    def _deserialize_pandas_series(data):
+        deserialized = deserialize_pandas(data)
+        return deserialized[deserialized.columns[0]]
+
+    def _serialize_pandas_dataframe(obj):
+        return serialize_pandas(obj).to_pybytes()
+
+    def _deserialize_pandas_dataframe(data):
+        return deserialize_pandas(data)
+
+    _default_serialization_context.register_type(
+        pd.Series, 'pd.Series',
+        custom_serializer=_serialize_pandas_series,
+        custom_deserializer=_deserialize_pandas_series)
+
+    _default_serialization_context.register_type(
+        pd.DataFrame, 'pd.DataFrame',
+        custom_serializer=_serialize_pandas_dataframe,
+        custom_deserializer=_deserialize_pandas_dataframe)
+except ImportError:
+    # no pandas
+    pass

python/pyarrow/tests/test_serialization.py

@@ -180,67 +180,18 @@ CUSTOM_OBJECTS = [Exception("Test object."), CustomError(), Point(11, y=22),
 
 def make_serialization_context():
 
-    def array_custom_serializer(obj):
-        return obj.tolist(), obj.dtype.str
+    context = pa._default_serialization_context
 
-    def array_custom_deserializer(serialized_obj):
-        return np.array(serialized_obj[0], dtype=np.dtype(serialized_obj[1]))
-
-    context = pa.SerializationContext()
-
-    # This is for numpy arrays of "object" only; primitive types are handled
-    # efficiently with Arrow's Tensor facilities (see python_to_arrow.cc)
-    context.register_type(np.ndarray, 20 * b"\x00",
-                          custom_serializer=array_custom_serializer,
-                          custom_deserializer=array_custom_deserializer)
-
-    context.register_type(Foo, 20 * b"\x01")
-    context.register_type(Bar, 20 * b"\x02")
-    context.register_type(Baz, 20 * b"\x03")
-    context.register_type(Qux, 20 * b"\x04")
-    context.register_type(SubQux, 20 * b"\x05")
-    context.register_type(SubQuxPickle, 20 * b"\x05", pickle=True)
-    context.register_type(Exception, 20 * b"\x06")
-    context.register_type(CustomError, 20 * b"\x07")
-    context.register_type(Point, 20 * b"\x08")
-    context.register_type(NamedTupleExample, 20 * b"\x09")
-
-    # TODO(pcm): This is currently a workaround until arrow supports
-    # arbitrary precision integers. This is only called on long integers,
-    # see the associated case in the append method in python_to_arrow.cc
-    context.register_type(int, 20 * b"\x10", pickle=False,
-                          custom_serializer=lambda obj: str(obj),
-                          custom_deserializer=(
-                              lambda serialized_obj: int(serialized_obj)))
-
-    if (sys.version_info < (3, 0)):
-        deserializer = (
-            lambda serialized_obj: long(serialized_obj))  # noqa: E501,F821
-        context.register_type(long, 20 * b"\x11", pickle=False,  # noqa: E501,F821
-                              custom_serializer=lambda obj: str(obj),
-                              custom_deserializer=deserializer)
-
-    def ordered_dict_custom_serializer(obj):
-        return list(obj.keys()), list(obj.values())
-
-    def ordered_dict_custom_deserializer(obj):
-        return OrderedDict(zip(obj[0], obj[1]))
-
-    context.register_type(OrderedDict, 20 * b"\x12", pickle=False,
-                          custom_serializer=ordered_dict_custom_serializer,
-                          custom_deserializer=ordered_dict_custom_deserializer)
-
-    def default_dict_custom_serializer(obj):
-        return list(obj.keys()), list(obj.values()), obj.default_factory
-
-    def default_dict_custom_deserializer(obj):
-        return defaultdict(obj[2], zip(obj[0], obj[1]))
-
-    context.register_type(defaultdict, 20 * b"\x13", pickle=False,
-                          custom_serializer=default_dict_custom_serializer,
-                          custom_deserializer=default_dict_custom_deserializer)
-
-    context.register_type(type(lambda: 0), 20 * b"\x14", pickle=True)
+    context.register_type(Foo, "Foo")
+    context.register_type(Bar, "Bar")
+    context.register_type(Baz, "Baz")
+    context.register_type(Qux, "Quz")
+    context.register_type(SubQux, "SubQux")
+    context.register_type(SubQuxPickle, "SubQuxPickle", pickle=True)
+    context.register_type(Exception, "Exception")
+    context.register_type(CustomError, "CustomError")
+    context.register_type(Point, "Point")
+    context.register_type(NamedTupleExample, "NamedTupleExample")
 
     return context