Add default serialization context and add serialization callbacks for pandas Series, DataFrame

Change-Id: I71dbacb1d069618a347bc0a8de0213e1abccf682

python/pyarrow/ipc.py

@@ -187,3 +187,50 @@ 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/pandas_compat.py

@@ -526,15 +526,3 @@ def _add_any_metadata(table, pandas_metadata):
         return pa.Table.from_arrays(columns)
     else:
         return table
-
-
-_default_serialization_context.register_type(
-    pd.Series, 'pandas.Series',
-    custom_serializer=_serialize_series,
-    custom_deserializer=_deserialize_series)
-
-
-_default_serialization_context.register_type(
-    pd.DataFrame, 'pandas.DataFrame',
-    custom_serializer=_serialize_dataframe,
-    custom_deserializer=_deserialize_dataframe)

python/pyarrow/serialization.pxi

@@ -209,7 +209,8 @@ def serialize(object value, SerializationContext context=None):
     value: object
         Python object for the sequence that is to be serialized.
     context : SerializationContext
-        Custom serialization and deserialization context
+        Custom serialization and deserialization context, uses a default
+        context with some standard type handlers if not specified
 
     Returns
     -------
@@ -236,7 +237,8 @@ def serialize_to(object value, sink, SerializationContext context=None):
     sink: NativeFile or file-like
         File the sequence will be written to.
     context : SerializationContext
-        Custom serialization and deserialization context
+        Custom serialization and deserialization context, uses a default
+        context with some standard type handlers if not specified
     """
     serialized = serialize(value, context)
     serialized.write_to(sink)

python/pyarrow/tests/test_ipc.py

@@ -22,7 +22,8 @@ import threading
 
 import numpy as np
 
-from pandas.util.testing import assert_frame_equal
+from pandas.util.testing import (assert_frame_equal,
+                                 assert_series_equal)
 import pandas as pd
 
 from pyarrow.compat import unittest
@@ -429,6 +430,20 @@ def test_serialize_pandas_no_preserve_index():
     assert_frame_equal(result, df)
 
 
+def test_serialize_with_pandas_objects():
+    df = pd.DataFrame({'a': [1, 2, 3]}, index=[1, 2, 3])
+
+    data = {
+        'a_series': df['a'],
+        'a_frame': df
+    }
+
+    serialized = pa.serialize(data).to_buffer()
+    deserialized = pa.deserialize(serialized)
+    assert_frame_equal(deserialized['a_frame'], df)
+    assert_series_equal(deserialized['a_series'], df['a'])
+
+
 def test_schema_batch_serialize_methods():
     nrows = 5
     df = pd.DataFrame({