# 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.
import collections
import warnings
import numpy as np
import pyarrow as pa
from pyarrow.lib import SerializationContext, py_buffer, builtin_pickle
try:
import cloudpickle
except ImportError:
cloudpickle = builtin_pickle
try:
# This function is available after numpy-0.16.0.
# See also: https://github.com/numpy/numpy/blob/master/numpy/lib/format.py
from numpy.lib.format import descr_to_dtype
except ImportError:
def descr_to_dtype(descr):
'''
descr may be stored as dtype.descr, which is a list of (name, format,
[shape]) tuples where format may be a str or a tuple. Offsets are not
explicitly saved, rather empty fields with name, format == '', '|Vn'
are added as padding. This function reverses the process, eliminating
the empty padding fields.
'''
if isinstance(descr, str):
# No padding removal needed
return np.dtype(descr)
elif isinstance(descr, tuple):
# subtype, will always have a shape descr[1]
dt = descr_to_dtype(descr[0])
return np.dtype((dt, descr[1]))
fields = []
offset = 0
for field in descr:
if len(field) == 2:
name, descr_str = field
dt = descr_to_dtype(descr_str)
else:
name, descr_str, shape = field
dt = np.dtype((descr_to_dtype(descr_str), shape))
# Ignore padding bytes, which will be void bytes with '' as name
# Once support for blank names is removed, only "if name == ''"
# needed)
is_pad = (name == '' and dt.type is np.void and dt.names is None)
if not is_pad:
fields.append((name, dt, offset))
offset += dt.itemsize
names, formats, offsets = zip(*fields)
# names may be (title, names) tuples
nametups = (n if isinstance(n, tuple) else (None, n) for n in names)
titles, names = zip(*nametups)
return np.dtype({'names': names, 'formats': formats, 'titles': titles,
'offsets': offsets, 'itemsize': offset})
def _deprecate_serialization(name):
msg = (
"'pyarrow.{}' is deprecated as of 2.0.0 and will be removed in a "
"future version. Use pickle or the pyarrow IPC functionality instead."
).format(name)
warnings.warn(msg, FutureWarning, stacklevel=3)
# ----------------------------------------------------------------------
# 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):
if obj.dtype.str != '|O':
# Make the array c_contiguous if necessary so that we can call change
# the view.
if not obj.flags.c_contiguous:
obj = np.ascontiguousarray(obj)
return obj.view('uint8'), np.lib.format.dtype_to_descr(obj.dtype)
else:
return obj.tolist(), np.lib.format.dtype_to_descr(obj.dtype)
def _deserialize_numpy_array_list(data):
if data[1] != '|O':
assert data[0].dtype == np.uint8
return data[0].view(descr_to_dtype(data[1]))
else:
return np.array(data[0], dtype=np.dtype(data[1]))
def _serialize_numpy_matrix(obj):
if obj.dtype.str != '|O':
# Make the array c_contiguous if necessary so that we can call change
# the view.
if not obj.flags.c_contiguous:
obj = np.ascontiguousarray(obj.A)
return obj.A.view('uint8'), np.lib.format.dtype_to_descr(obj.dtype)
else:
return obj.A.tolist(), np.lib.format.dtype_to_descr(obj.dtype)
def _deserialize_numpy_matrix(data):
if data[1] != '|O':
assert data[0].dtype == np.uint8
return np.matrix(data[0].view(descr_to_dtype(data[1])),
copy=False)
else:
return np.matrix(data[0], dtype=np.dtype(data[1]), copy=False)
# ----------------------------------------------------------------------
# pyarrow.RecordBatch-specific serialization matters
def _serialize_pyarrow_recordbatch(batch):
output_stream = pa.BufferOutputStream()
with pa.RecordBatchStreamWriter(output_stream, schema=batch.schema) as wr:
wr.write_batch(batch)
return output_stream.getvalue() # This will also close the stream.
def _deserialize_pyarrow_recordbatch(buf):
with pa.RecordBatchStreamReader(buf) as reader:
return reader.read_next_batch()
# ----------------------------------------------------------------------
# pyarrow.Array-specific serialization matters
def _serialize_pyarrow_array(array):
# TODO(suquark): implement more effcient array serialization.
batch = pa.RecordBatch.from_arrays([array], [''])
return _serialize_pyarrow_recordbatch(batch)
def _deserialize_pyarrow_array(buf):
# TODO(suquark): implement more effcient array deserialization.
batch = _deserialize_pyarrow_recordbatch(buf)
return batch.columns[0]
# ----------------------------------------------------------------------
# pyarrow.Table-specific serialization matters
def _serialize_pyarrow_table(table):
output_stream = pa.BufferOutputStream()
with pa.RecordBatchStreamWriter(output_stream, schema=table.schema) as wr:
wr.write_table(table)
return output_stream.getvalue() # This will also close the stream.
def _deserialize_pyarrow_table(buf):
with pa.RecordBatchStreamReader(buf) as reader:
return reader.read_all()
def _pickle_to_buffer(x):
pickled = builtin_pickle.dumps(x, protocol=builtin_pickle.HIGHEST_PROTOCOL)
return py_buffer(pickled)
def _load_pickle_from_buffer(data):
as_memoryview = memoryview(data)
return builtin_pickle.loads(as_memoryview)
# ----------------------------------------------------------------------
# pandas-specific serialization matters
def _register_custom_pandas_handlers(context):
# ARROW-1784, faster path for pandas-only visibility
try:
import pandas as pd
except ImportError:
return
import pyarrow.pandas_compat as pdcompat
sparse_type_error_msg = (
'{0} serialization is not supported.\n'
'Note that {0} is planned to be deprecated '
'in pandas future releases.\n'
'See https://github.com/pandas-dev/pandas/issues/19239 '
'for more information.'
)
def _serialize_pandas_dataframe(obj):
if (pdcompat._pandas_api.has_sparse and
isinstance(obj, pd.SparseDataFrame)):
raise NotImplementedError(
sparse_type_error_msg.format('SparseDataFrame')
)
return pdcompat.dataframe_to_serialized_dict(obj)
def _deserialize_pandas_dataframe(data):
return pdcompat.serialized_dict_to_dataframe(data)
def _serialize_pandas_series(obj):
if (pdcompat._pandas_api.has_sparse and
isinstance(obj, pd.SparseSeries)):
raise NotImplementedError(
sparse_type_error_msg.format('SparseSeries')
)
return _serialize_pandas_dataframe(pd.DataFrame({obj.name: obj}))
def _deserialize_pandas_series(data):
deserialized = _deserialize_pandas_dataframe(data)
return deserialized[deserialized.columns[0]]
context.register_type(
pd.Series, 'pd.Series',
custom_serializer=_serialize_pandas_series,
custom_deserializer=_deserialize_pandas_series)
context.register_type(
pd.Index, 'pd.Index',
custom_serializer=_pickle_to_buffer,
custom_deserializer=_load_pickle_from_buffer)
if hasattr(pd.core, 'arrays'):
if hasattr(pd.core.arrays, 'interval'):
context.register_type(
pd.core.arrays.interval.IntervalArray,
'pd.core.arrays.interval.IntervalArray',
custom_serializer=_pickle_to_buffer,
custom_deserializer=_load_pickle_from_buffer)
if hasattr(pd.core.arrays, 'period'):
context.register_type(
pd.core.arrays.period.PeriodArray,
'pd.core.arrays.period.PeriodArray',
custom_serializer=_pickle_to_buffer,
custom_deserializer=_load_pickle_from_buffer)
if hasattr(pd.core.arrays, 'datetimes'):
context.register_type(
pd.core.arrays.datetimes.DatetimeArray,
'pd.core.arrays.datetimes.DatetimeArray',
custom_serializer=_pickle_to_buffer,
custom_deserializer=_load_pickle_from_buffer)
context.register_type(
pd.DataFrame, 'pd.DataFrame',
custom_serializer=_serialize_pandas_dataframe,
custom_deserializer=_deserialize_pandas_dataframe)
def register_torch_serialization_handlers(serialization_context):
# ----------------------------------------------------------------------
# Set up serialization for pytorch tensors
_deprecate_serialization("register_torch_serialization_handlers")
try:
import torch
def _serialize_torch_tensor(obj):
if obj.is_sparse:
return pa.SparseCOOTensor.from_numpy(
obj._values().detach().numpy(),
obj._indices().detach().numpy().T,
shape=list(obj.shape))
else:
return obj.detach().numpy()
def _deserialize_torch_tensor(data):
if isinstance(data, pa.SparseCOOTensor):
return torch.sparse_coo_tensor(
indices=data.to_numpy()[1].T,
values=data.to_numpy()[0][:, 0],
size=data.shape)
else:
return torch.from_numpy(data)
for t in [torch.FloatTensor, torch.DoubleTensor, torch.HalfTensor,
torch.ByteTensor, torch.CharTensor, torch.ShortTensor,
torch.IntTensor, torch.LongTensor, torch.Tensor]:
serialization_context.register_type(
t, "torch." + t.__name__,
custom_serializer=_serialize_torch_tensor,
custom_deserializer=_deserialize_torch_tensor)
except ImportError:
# no torch
pass
def _register_collections_serialization_handlers(serialization_context):
def _serialize_deque(obj):
return list(obj)
def _deserialize_deque(data):
return collections.deque(data)
serialization_context.register_type(
collections.deque, "collections.deque",
custom_serializer=_serialize_deque,
custom_deserializer=_deserialize_deque)
def _serialize_ordered_dict(obj):
return list(obj.keys()), list(obj.values())
def _deserialize_ordered_dict(data):
return collections.OrderedDict(zip(data[0], data[1]))
serialization_context.register_type(
collections.OrderedDict, "collections.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 collections.defaultdict(data[2], zip(data[0], data[1]))
serialization_context.register_type(
collections.defaultdict, "collections.defaultdict",
custom_serializer=_serialize_default_dict,
custom_deserializer=_deserialize_default_dict)
def _serialize_counter(obj):
return list(obj.keys()), list(obj.values())
def _deserialize_counter(data):
return collections.Counter(dict(zip(data[0], data[1])))
serialization_context.register_type(
collections.Counter, "collections.Counter",
custom_serializer=_serialize_counter,
custom_deserializer=_deserialize_counter)
# ----------------------------------------------------------------------
# Set up serialization for scipy sparse matrices. Primitive types are handled
# efficiently with Arrow's SparseTensor facilities, see numpy_convert.cc)
def _register_scipy_handlers(serialization_context):
try:
from scipy.sparse import (csr_matrix, csc_matrix, coo_matrix,
isspmatrix_coo, isspmatrix_csr,
isspmatrix_csc, isspmatrix)
def _serialize_scipy_sparse(obj):
if isspmatrix_coo(obj):
return 'coo', pa.SparseCOOTensor.from_scipy(obj)
elif isspmatrix_csr(obj):
return 'csr', pa.SparseCSRMatrix.from_scipy(obj)
elif isspmatrix_csc(obj):
return 'csc', pa.SparseCSCMatrix.from_scipy(obj)
elif isspmatrix(obj):
return 'csr', pa.SparseCOOTensor.from_scipy(obj.to_coo())
else:
raise NotImplementedError(
"Serialization of {} is not supported.".format(obj[0]))
def _deserialize_scipy_sparse(data):
if data[0] == 'coo':
return data[1].to_scipy()
elif data[0] == 'csr':
return data[1].to_scipy()
elif data[0] == 'csc':
return data[1].to_scipy()
else:
return data[1].to_scipy()
serialization_context.register_type(
coo_matrix, 'scipy.sparse.coo.coo_matrix',
custom_serializer=_serialize_scipy_sparse,
custom_deserializer=_deserialize_scipy_sparse)
serialization_context.register_type(
csr_matrix, 'scipy.sparse.csr.csr_matrix',
custom_serializer=_serialize_scipy_sparse,
custom_deserializer=_deserialize_scipy_sparse)
serialization_context.register_type(
csc_matrix, 'scipy.sparse.csc.csc_matrix',
custom_serializer=_serialize_scipy_sparse,
custom_deserializer=_deserialize_scipy_sparse)
except ImportError:
# no scipy
pass
# ----------------------------------------------------------------------
# Set up serialization for pydata/sparse tensors.
def _register_pydata_sparse_handlers(serialization_context):
try:
import sparse
def _serialize_pydata_sparse(obj):
if isinstance(obj, sparse.COO):
return 'coo', pa.SparseCOOTensor.from_pydata_sparse(obj)
else:
raise NotImplementedError(
"Serialization of {} is not supported.".format(sparse.COO))
def _deserialize_pydata_sparse(data):
if data[0] == 'coo':
data_array, coords = data[1].to_numpy()
return sparse.COO(
data=data_array[:, 0],
coords=coords.T, shape=data[1].shape)
serialization_context.register_type(
sparse.COO, 'sparse.COO',
custom_serializer=_serialize_pydata_sparse,
custom_deserializer=_deserialize_pydata_sparse)
except ImportError:
# no pydata/sparse
pass
def _register_default_serialization_handlers(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
serialization_context.register_type(
int, "int",
custom_serializer=lambda obj: str(obj),
custom_deserializer=lambda data: int(data))
serialization_context.register_type(
type(lambda: 0), "function",
pickle=True)
serialization_context.register_type(type, "type", pickle=True)
serialization_context.register_type(
np.matrix, 'np.matrix',
custom_serializer=_serialize_numpy_matrix,
custom_deserializer=_deserialize_numpy_matrix)
serialization_context.register_type(
np.ndarray, 'np.array',
custom_serializer=_serialize_numpy_array_list,
custom_deserializer=_deserialize_numpy_array_list)
serialization_context.register_type(
pa.Array, 'pyarrow.Array',
custom_serializer=_serialize_pyarrow_array,
custom_deserializer=_deserialize_pyarrow_array)
serialization_context.register_type(
pa.RecordBatch, 'pyarrow.RecordBatch',
custom_serializer=_serialize_pyarrow_recordbatch,
custom_deserializer=_deserialize_pyarrow_recordbatch)
serialization_context.register_type(
pa.Table, 'pyarrow.Table',
custom_serializer=_serialize_pyarrow_table,
custom_deserializer=_deserialize_pyarrow_table)
_register_collections_serialization_handlers(serialization_context)
_register_custom_pandas_handlers(serialization_context)
_register_scipy_handlers(serialization_context)
_register_pydata_sparse_handlers(serialization_context)
def register_default_serialization_handlers(serialization_context):
_deprecate_serialization("register_default_serialization_handlers")
_register_default_serialization_handlers(serialization_context)
def default_serialization_context():
_deprecate_serialization("default_serialization_context")
context = SerializationContext()
_register_default_serialization_handlers(context)
return context
