| # 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 __future__ import absolute_import |
| |
| import collections |
| import six |
| import sys |
| |
| import numpy as np |
| |
| import pyarrow |
| from pyarrow.compat import builtin_pickle |
| from pyarrow.lib import SerializationContext, py_buffer |
| |
| try: |
| import cloudpickle |
| except ImportError: |
| cloudpickle = builtin_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): |
| 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'), obj.dtype.str |
| else: |
| return obj.tolist(), obj.dtype.str |
| |
| |
| def _deserialize_numpy_array_list(data): |
| if data[1] != '|O': |
| assert data[0].dtype == np.uint8 |
| return data[0].view(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'), obj.A.dtype.str |
| else: |
| return obj.A.tolist(), obj.A.dtype.str |
| |
| |
| def _deserialize_numpy_matrix(data): |
| if data[1] != '|O': |
| assert data[0].dtype == np.uint8 |
| return np.matrix(data[0].view(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 = pyarrow.BufferOutputStream() |
| writer = pyarrow.RecordBatchStreamWriter(output_stream, |
| schema=batch.schema) |
| writer.write_batch(batch) |
| writer.close() |
| return output_stream.getvalue() # This will also close the stream. |
| |
| |
| def _deserialize_pyarrow_recordbatch(buf): |
| reader = pyarrow.RecordBatchStreamReader(buf) |
| batch = reader.read_next_batch() |
| return batch |
| |
| |
| # ---------------------------------------------------------------------- |
| # pyarrow.Array-specific serialization matters |
| |
| def _serialize_pyarrow_array(array): |
| # TODO(suquark): implement more effcient array serialization. |
| batch = pyarrow.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 = pyarrow.BufferOutputStream() |
| writer = pyarrow.RecordBatchStreamWriter(output_stream, |
| schema=table.schema) |
| writer.write_table(table) |
| writer.close() |
| return output_stream.getvalue() # This will also close the stream. |
| |
| |
| def _deserialize_pyarrow_table(buf): |
| reader = pyarrow.RecordBatchStreamReader(buf) |
| table = reader.read_all() |
| return table |
| |
| |
| 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) |
| if six.PY2: |
| return builtin_pickle.loads(as_memoryview.tobytes()) |
| else: |
| 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 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 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 |
| |
| try: |
| import torch |
| |
| def _serialize_torch_tensor(obj): |
| if obj.is_sparse: |
| # TODO(pcm): Once ARROW-4453 is resolved, return sparse |
| # tensor representation here |
| return (obj._indices().detach().numpy(), |
| obj._values().detach().numpy(), list(obj.shape)) |
| else: |
| return obj.detach().numpy() |
| |
| def _deserialize_torch_tensor(data): |
| if isinstance(data, tuple): |
| return torch.sparse_coo_tensor(data[0], data[1], data[2]) |
| 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) |
| |
| |
| 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)) |
| |
| if (sys.version_info < (3, 0)): |
| serialization_context.register_type( |
| long, "long", # noqa: F821 |
| custom_serializer=lambda obj: str(obj), |
| custom_deserializer=lambda data: long(data)) # noqa: F821 |
| |
| 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( |
| pyarrow.Array, 'pyarrow.Array', |
| custom_serializer=_serialize_pyarrow_array, |
| custom_deserializer=_deserialize_pyarrow_array) |
| |
| serialization_context.register_type( |
| pyarrow.RecordBatch, 'pyarrow.RecordBatch', |
| custom_serializer=_serialize_pyarrow_recordbatch, |
| custom_deserializer=_deserialize_pyarrow_recordbatch) |
| |
| serialization_context.register_type( |
| pyarrow.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) |
| |
| |
| def default_serialization_context(): |
| context = SerializationContext() |
| register_default_serialization_handlers(context) |
| return context |
