/
usr
/
local
/
lib64
/
python3.6
/
site-packages
/
pyarrow
/
tests
/
/usr/local/lib64/python3.6/site-packages/pyarrow/tests
mkdir
upload
Name
Size
Mode
Actions
data/
-
0755
rm
parquet/
-
0755
rm
__pycache__/
-
0755
rm
arrow_7980.py
1094
0644
edit
dl
rm
bound_function_visit_strings.pyx
2047
0644
edit
dl
rm
conftest.py
7982
0644
edit
dl
rm
deserialize_buffer.py
961
0644
edit
dl
rm
pandas_examples.py
5115
0644
edit
dl
rm
pandas_threaded_import.py
1429
0644
edit
dl
rm
pyarrow_cython_example.pyx
1951
0644
edit
dl
rm
strategies.py
12525
0644
edit
dl
rm
test_adhoc_memory_leak.py
1410
0644
edit
dl
rm
test_array.py
100157
0644
edit
dl
rm
test_builder.py
2184
0644
edit
dl
rm
test_cffi.py
13090
0644
edit
dl
rm
test_compute.py
83551
0644
edit
dl
rm
test_convert_builtin.py
71768
0644
edit
dl
rm
test_csv.py
68943
0644
edit
dl
rm
test_cuda.py
27863
0644
edit
dl
rm
test_cuda_numba_interop.py
8730
0644
edit
dl
rm
test_cython.py
5948
0644
edit
dl
rm
test_dataset.py
143023
0644
edit
dl
rm
test_deprecations.py
891
0644
edit
dl
rm
test_extension_type.py
24437
0644
edit
dl
rm
test_feather.py
23209
0644
edit
dl
rm
test_filesystem.py
2106
0644
edit
dl
rm
test_flight.py
73715
0644
edit
dl
rm
test_fs.py
54133
0644
edit
dl
rm
test_gandiva.py
14385
0644
edit
dl
rm
test_hdfs.py
13452
0644
edit
dl
rm
test_io.py
53066
0644
edit
dl
rm
test_ipc.py
31267
0644
edit
dl
rm
test_json.py
11325
0644
edit
dl
rm
test_jvm.py
15471
0644
edit
dl
rm
test_memory.py
4928
0644
edit
dl
rm
test_misc.py
4927
0644
edit
dl
rm
test_orc.py
9221
0644
edit
dl
rm
test_pandas.py
157909
0644
edit
dl
rm
test_plasma.py
44011
0644
edit
dl
rm
test_plasma_tf_op.py
3523
0644
edit
dl
rm
test_scalars.py
21074
0644
edit
dl
rm
test_schema.py
21282
0644
edit
dl
rm
test_serialization.py
42026
0644
edit
dl
rm
test_serialization_deprecated.py
1775
0644
edit
dl
rm
test_sparse_tensor.py
17436
0644
edit
dl
rm
test_strategies.py
1739
0644
edit
dl
rm
test_table.py
53168
0644
edit
dl
rm
test_tensor.py
6270
0644
edit
dl
rm
test_types.py
31487
0644
edit
dl
rm
test_util.py
1791
0644
edit
dl
rm
util.py
9011
0644
edit
dl
rm
__init__.py
0
0644
edit
dl
rm
Edit:
/usr/local/lib64/python3.6/site-packages/pyarrow/tests/test_serialization.py
(42026B)
# 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 pytest import collections import datetime import os import pathlib import pickle import subprocess import string import sys import pyarrow as pa import numpy as np import pyarrow.tests.util as test_util try: import torch except ImportError: torch = None # Blacklist the module in case `import torch` is costly before # failing (ARROW-2071) sys.modules['torch'] = None try: from scipy.sparse import coo_matrix, csr_matrix, csc_matrix except ImportError: coo_matrix = None csr_matrix = None csc_matrix = None try: import sparse except ImportError: sparse = None # ignore all serialization deprecation warnings in this file, we test that the # warnings are actually raised in test_serialization_deprecated.py pytestmark = pytest.mark.filterwarnings("ignore:'pyarrow:FutureWarning") def assert_equal(obj1, obj2): if torch is not None and torch.is_tensor(obj1) and torch.is_tensor(obj2): if obj1.is_sparse: obj1 = obj1.to_dense() if obj2.is_sparse: obj2 = obj2.to_dense() assert torch.equal(obj1, obj2) return module_numpy = (type(obj1).__module__ == np.__name__ or type(obj2).__module__ == np.__name__) if module_numpy: empty_shape = ((hasattr(obj1, "shape") and obj1.shape == ()) or (hasattr(obj2, "shape") and obj2.shape == ())) if empty_shape: # This is a special case because currently np.testing.assert_equal # fails because we do not properly handle different numerical # types. assert obj1 == obj2, ("Objects {} and {} are " "different.".format(obj1, obj2)) else: np.testing.assert_equal(obj1, obj2) elif hasattr(obj1, "__dict__") and hasattr(obj2, "__dict__"): special_keys = ["_pytype_"] assert (set(list(obj1.__dict__.keys()) + special_keys) == set(list(obj2.__dict__.keys()) + special_keys)), ("Objects {} " "and {} are " "different." .format( obj1, obj2)) if obj1.__dict__ == {}: print("WARNING: Empty dict in ", obj1) for key in obj1.__dict__.keys(): if key not in special_keys: assert_equal(obj1.__dict__[key], obj2.__dict__[key]) elif type(obj1) is dict or type(obj2) is dict: assert_equal(obj1.keys(), obj2.keys()) for key in obj1.keys(): assert_equal(obj1[key], obj2[key]) elif type(obj1) is list or type(obj2) is list: assert len(obj1) == len(obj2), ("Objects {} and {} are lists with " "different lengths." .format(obj1, obj2)) for i in range(len(obj1)): assert_equal(obj1[i], obj2[i]) elif type(obj1) is tuple or type(obj2) is tuple: assert len(obj1) == len(obj2), ("Objects {} and {} are tuples with " "different lengths." .format(obj1, obj2)) for i in range(len(obj1)): assert_equal(obj1[i], obj2[i]) elif (pa.lib.is_named_tuple(type(obj1)) or pa.lib.is_named_tuple(type(obj2))): assert len(obj1) == len(obj2), ("Objects {} and {} are named tuples " "with different lengths." .format(obj1, obj2)) for i in range(len(obj1)): assert_equal(obj1[i], obj2[i]) elif isinstance(obj1, pa.Array) and isinstance(obj2, pa.Array): assert obj1.equals(obj2) elif isinstance(obj1, pa.Tensor) and isinstance(obj2, pa.Tensor): assert obj1.equals(obj2) elif isinstance(obj1, pa.Tensor) and isinstance(obj2, pa.Tensor): assert obj1.equals(obj2) elif isinstance(obj1, pa.SparseCOOTensor) and \ isinstance(obj2, pa.SparseCOOTensor): assert obj1.equals(obj2) elif isinstance(obj1, pa.SparseCSRMatrix) and \ isinstance(obj2, pa.SparseCSRMatrix): assert obj1.equals(obj2) elif isinstance(obj1, pa.SparseCSCMatrix) and \ isinstance(obj2, pa.SparseCSCMatrix): assert obj1.equals(obj2) elif isinstance(obj1, pa.SparseCSFTensor) and \ isinstance(obj2, pa.SparseCSFTensor): assert obj1.equals(obj2) elif isinstance(obj1, pa.RecordBatch) and isinstance(obj2, pa.RecordBatch): assert obj1.equals(obj2) elif isinstance(obj1, pa.Table) and isinstance(obj2, pa.Table): assert obj1.equals(obj2) else: assert type(obj1) == type(obj2) and obj1 == obj2, \ "Objects {} and {} are different.".format(obj1, obj2) PRIMITIVE_OBJECTS = [ 0, 0.0, 0.9, 1 << 62, 1 << 999, [1 << 100, [1 << 100]], "a", string.printable, "\u262F", "hello world", "hello world", "\xff\xfe\x9c\x001\x000\x00", None, True, False, [], (), {}, {(1, 2): 1}, {(): 2}, [1, "hello", 3.0], "\u262F", 42.0, (1.0, "hi"), [1, 2, 3, None], [(None,), 3, 1.0], ["h", "e", "l", "l", "o", None], (None, None), ("hello", None), (True, False), {True: "hello", False: "world"}, {"hello": "world", 1: 42, 2.5: 45}, {"hello": {2, 3}, "world": {42.0}, "this": None}, np.int8(3), np.int32(4), np.int64(5), np.uint8(3), np.uint32(4), np.uint64(5), np.float16(1.9), np.float32(1.9), np.float64(1.9), np.zeros([8, 20]), np.random.normal(size=[17, 10]), np.array(["hi", 3]), np.array(["hi", 3], dtype=object), np.random.normal(size=[15, 13]).T ] index_types = ('i1', 'i2', 'i4', 'i8', 'u1', 'u2', 'u4', 'u8') tensor_types = ('i1', 'i2', 'i4', 'i8', 'u1', 'u2', 'u4', 'u8', 'f2', 'f4', 'f8') PRIMITIVE_OBJECTS += [0, np.array([["hi", "hi"], [1.3, 1]])] COMPLEX_OBJECTS = [ [[[[[[[[[[[[]]]]]]]]]]]], {"obj{}".format(i): np.random.normal(size=[4, 4]) for i in range(5)}, # {(): {(): {(): {(): {(): {(): {(): {(): {(): {(): { # (): {(): {}}}}}}}}}}}}}, ((((((((((),),),),),),),),),), {"a": {"b": {"c": {"d": {}}}}}, ] class Foo: def __init__(self, value=0): self.value = value def __hash__(self): return hash(self.value) def __eq__(self, other): return other.value == self.value class Bar: def __init__(self): for i, val in enumerate(COMPLEX_OBJECTS): setattr(self, "field{}".format(i), val) class Baz: def __init__(self): self.foo = Foo() self.bar = Bar() def method(self, arg): pass class Qux: def __init__(self): self.objs = [Foo(1), Foo(42)] class SubQux(Qux): def __init__(self): Qux.__init__(self) class SubQuxPickle(Qux): def __init__(self): Qux.__init__(self) class CustomError(Exception): pass Point = collections.namedtuple("Point", ["x", "y"]) NamedTupleExample = collections.namedtuple( "Example", "field1, field2, field3, field4, field5") CUSTOM_OBJECTS = [Exception("Test object."), CustomError(), Point(11, y=22), Foo(), Bar(), Baz(), Qux(), SubQux(), SubQuxPickle(), NamedTupleExample(1, 1.0, "hi", np.zeros([3, 5]), [1, 2, 3]), collections.OrderedDict([("hello", 1), ("world", 2)]), collections.deque([1, 2, 3, "a", "b", "c", 3.5]), collections.Counter([1, 1, 1, 2, 2, 3, "a", "b"])] def make_serialization_context(): with pytest.warns(FutureWarning): context = pa.default_serialization_context() 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 global_serialization_context = make_serialization_context() def serialization_roundtrip(value, scratch_buffer, context=global_serialization_context): writer = pa.FixedSizeBufferWriter(scratch_buffer) pa.serialize_to(value, writer, context=context) reader = pa.BufferReader(scratch_buffer) result = pa.deserialize_from(reader, None, context=context) assert_equal(value, result) _check_component_roundtrip(value, context=context) def _check_component_roundtrip(value, context=global_serialization_context): # Test to/from components serialized = pa.serialize(value, context=context) components = serialized.to_components() from_comp = pa.SerializedPyObject.from_components(components) recons = from_comp.deserialize(context=context) assert_equal(value, recons) @pytest.fixture(scope='session') def large_buffer(size=32*1024*1024): yield pa.allocate_buffer(size) def large_memory_map(tmpdir_factory, size=100*1024*1024): path = (tmpdir_factory.mktemp('data') .join('pyarrow-serialization-tmp-file').strpath) # Create a large memory mapped file with open(path, 'wb') as f: f.write(np.random.randint(0, 256, size=size) .astype('u1') .tobytes() [:size]) return path def test_clone(): context = pa.SerializationContext() class Foo: pass def custom_serializer(obj): return 0 def custom_deserializer(serialized_obj): return (serialized_obj, 'a') context.register_type(Foo, 'Foo', custom_serializer=custom_serializer, custom_deserializer=custom_deserializer) new_context = context.clone() f = Foo() serialized = pa.serialize(f, context=context) deserialized = serialized.deserialize(context=context) assert deserialized == (0, 'a') serialized = pa.serialize(f, context=new_context) deserialized = serialized.deserialize(context=new_context) assert deserialized == (0, 'a') def test_primitive_serialization_notbroken(large_buffer): serialization_roundtrip({(1, 2): 2}, large_buffer) def test_primitive_serialization_broken(large_buffer): serialization_roundtrip({(): 2}, large_buffer) def test_primitive_serialization(large_buffer): for obj in PRIMITIVE_OBJECTS: serialization_roundtrip(obj, large_buffer) def test_integer_limits(large_buffer): # Check that Numpy scalars can be represented up to their limit values # (except np.uint64 which is limited to 2**63 - 1) for dt in [np.int8, np.int64, np.int32, np.int64, np.uint8, np.uint64, np.uint32, np.uint64]: scal = dt(np.iinfo(dt).min) serialization_roundtrip(scal, large_buffer) if dt is not np.uint64: scal = dt(np.iinfo(dt).max) serialization_roundtrip(scal, large_buffer) else: scal = dt(2**63 - 1) serialization_roundtrip(scal, large_buffer) for v in (2**63, 2**64 - 1): scal = dt(v) with pytest.raises(pa.ArrowInvalid): pa.serialize(scal) def test_serialize_to_buffer(): for nthreads in [1, 4]: for value in COMPLEX_OBJECTS: buf = pa.serialize(value).to_buffer(nthreads=nthreads) result = pa.deserialize(buf) assert_equal(value, result) def test_complex_serialization(large_buffer): for obj in COMPLEX_OBJECTS: serialization_roundtrip(obj, large_buffer) def test_custom_serialization(large_buffer): for obj in CUSTOM_OBJECTS: serialization_roundtrip(obj, large_buffer) def test_default_dict_serialization(large_buffer): pytest.importorskip("cloudpickle") obj = collections.defaultdict(lambda: 0, [("hello", 1), ("world", 2)]) serialization_roundtrip(obj, large_buffer) def test_numpy_serialization(large_buffer): for t in ["bool", "int8", "uint8", "int16", "uint16", "int32", "uint32", "float16", "float32", "float64", "<U1", "<U2", "<U3", "<U4", "|S1", "|S2", "|S3", "|S4", "|O", np.dtype([('a', 'int64'), ('b', 'float')]), np.dtype([('x', 'uint32'), ('y', '<U8')])]: obj = np.random.randint(0, 10, size=(100, 100)).astype(t) serialization_roundtrip(obj, large_buffer) obj = obj[1:99, 10:90] serialization_roundtrip(obj, large_buffer) def test_datetime_serialization(large_buffer): data = [ # Principia Mathematica published datetime.datetime(year=1687, month=7, day=5), # Some random date datetime.datetime(year=1911, month=6, day=3, hour=4, minute=55, second=44), # End of WWI datetime.datetime(year=1918, month=11, day=11), # Beginning of UNIX time datetime.datetime(year=1970, month=1, day=1), # The Berlin wall falls datetime.datetime(year=1989, month=11, day=9), # Another random date datetime.datetime(year=2011, month=6, day=3, hour=4, minute=0, second=3), # Another random date datetime.datetime(year=1970, month=1, day=3, hour=4, minute=0, second=0) ] for d in data: serialization_roundtrip(d, large_buffer) def test_torch_serialization(large_buffer): pytest.importorskip("torch") serialization_context = pa.default_serialization_context() pa.register_torch_serialization_handlers(serialization_context) # Dense tensors: # These are the only types that are supported for the # PyTorch to NumPy conversion for t in ["float32", "float64", "uint8", "int16", "int32", "int64"]: obj = torch.from_numpy(np.random.randn(1000).astype(t)) serialization_roundtrip(obj, large_buffer, context=serialization_context) tensor_requiring_grad = torch.randn(10, 10, requires_grad=True) serialization_roundtrip(tensor_requiring_grad, large_buffer, context=serialization_context) # Sparse tensors: # These are the only types that are supported for the # PyTorch to NumPy conversion for t in ["float32", "float64", "uint8", "int16", "int32", "int64"]: i = torch.LongTensor([[0, 2], [1, 0], [1, 2]]) v = torch.from_numpy(np.array([3, 4, 5]).astype(t)) obj = torch.sparse_coo_tensor(i.t(), v, torch.Size([2, 3])) serialization_roundtrip(obj, large_buffer, context=serialization_context) @pytest.mark.skipif(not torch or not torch.cuda.is_available(), reason="requires pytorch with CUDA") def test_torch_cuda(): # ARROW-2920: This used to segfault if torch is not imported # before pyarrow # Note that this test will only catch the issue if it is run # with a pyarrow that has been built in the manylinux1 environment torch.nn.Conv2d(64, 2, kernel_size=3, stride=1, padding=1, bias=False).cuda() def test_numpy_immutable(large_buffer): obj = np.zeros([10]) writer = pa.FixedSizeBufferWriter(large_buffer) pa.serialize_to(obj, writer, global_serialization_context) reader = pa.BufferReader(large_buffer) result = pa.deserialize_from(reader, None, global_serialization_context) with pytest.raises(ValueError): result[0] = 1.0 def test_numpy_base_object(tmpdir): # ARROW-2040: deserialized Numpy array should keep a reference to the # owner of its memory path = os.path.join(str(tmpdir), 'zzz.bin') data = np.arange(12, dtype=np.int32) with open(path, 'wb') as f: f.write(pa.serialize(data).to_buffer()) serialized = pa.read_serialized(pa.OSFile(path)) result = serialized.deserialize() assert_equal(result, data) serialized = None assert_equal(result, data) assert result.base is not None # see https://issues.apache.org/jira/browse/ARROW-1695 def test_serialization_callback_numpy(): class DummyClass: pass def serialize_dummy_class(obj): x = np.zeros(4) return x def deserialize_dummy_class(serialized_obj): return serialized_obj context = pa.default_serialization_context() context.register_type(DummyClass, "DummyClass", custom_serializer=serialize_dummy_class, custom_deserializer=deserialize_dummy_class) pa.serialize(DummyClass(), context=context) def test_numpy_subclass_serialization(): # Check that we can properly serialize subclasses of np.ndarray. class CustomNDArray(np.ndarray): def __new__(cls, input_array): array = np.asarray(input_array).view(cls) return array def serializer(obj): return {'numpy': obj.view(np.ndarray)} def deserializer(data): array = data['numpy'].view(CustomNDArray) return array context = pa.default_serialization_context() context.register_type(CustomNDArray, 'CustomNDArray', custom_serializer=serializer, custom_deserializer=deserializer) x = CustomNDArray(np.zeros(3)) serialized = pa.serialize(x, context=context).to_buffer() new_x = pa.deserialize(serialized, context=context) assert type(new_x) == CustomNDArray assert np.alltrue(new_x.view(np.ndarray) == np.zeros(3)) @pytest.mark.parametrize('tensor_type', tensor_types) @pytest.mark.parametrize('index_type', index_types) def test_sparse_coo_tensor_serialization(index_type, tensor_type): tensor_dtype = np.dtype(tensor_type) index_dtype = np.dtype(index_type) data = np.array([[1, 2, 3, 4, 5, 6]]).T.astype(tensor_dtype) coords = np.array([ [0, 0, 2, 3, 1, 3], [0, 2, 0, 4, 5, 5], ]).T.astype(index_dtype) shape = (4, 6) dim_names = ('x', 'y') sparse_tensor = pa.SparseCOOTensor.from_numpy(data, coords, shape, dim_names) context = pa.default_serialization_context() serialized = pa.serialize(sparse_tensor, context=context).to_buffer() result = pa.deserialize(serialized) assert_equal(result, sparse_tensor) assert isinstance(result, pa.SparseCOOTensor) data_result, coords_result = result.to_numpy() assert np.array_equal(data_result, data) assert np.array_equal(coords_result, coords) assert result.dim_names == dim_names @pytest.mark.parametrize('tensor_type', tensor_types) @pytest.mark.parametrize('index_type', index_types) def test_sparse_coo_tensor_components_serialization(large_buffer, index_type, tensor_type): tensor_dtype = np.dtype(tensor_type) index_dtype = np.dtype(index_type) data = np.array([[1, 2, 3, 4, 5, 6]]).T.astype(tensor_dtype) coords = np.array([ [0, 0, 2, 3, 1, 3], [0, 2, 0, 4, 5, 5], ]).T.astype(index_dtype) shape = (4, 6) dim_names = ('x', 'y') sparse_tensor = pa.SparseCOOTensor.from_numpy(data, coords, shape, dim_names) serialization_roundtrip(sparse_tensor, large_buffer) @pytest.mark.skipif(not coo_matrix, reason="requires scipy") def test_scipy_sparse_coo_tensor_serialization(): data = np.array([1, 2, 3, 4, 5, 6]) row = np.array([0, 0, 2, 3, 1, 3]) col = np.array([0, 2, 0, 4, 5, 5]) shape = (4, 6) sparse_array = coo_matrix((data, (row, col)), shape=shape) serialized = pa.serialize(sparse_array) result = serialized.deserialize() assert np.array_equal(sparse_array.toarray(), result.toarray()) @pytest.mark.skipif(not sparse, reason="requires pydata/sparse") def test_pydata_sparse_sparse_coo_tensor_serialization(): data = np.array([1, 2, 3, 4, 5, 6]) coords = np.array([ [0, 0, 2, 3, 1, 3], [0, 2, 0, 4, 5, 5], ]) shape = (4, 6) sparse_array = sparse.COO(data=data, coords=coords, shape=shape) serialized = pa.serialize(sparse_array) result = serialized.deserialize() assert np.array_equal(sparse_array.todense(), result.todense()) @pytest.mark.parametrize('tensor_type', tensor_types) @pytest.mark.parametrize('index_type', index_types) def test_sparse_csr_matrix_serialization(index_type, tensor_type): tensor_dtype = np.dtype(tensor_type) index_dtype = np.dtype(index_type) data = np.array([[8, 2, 5, 3, 4, 6]]).T.astype(tensor_dtype) indptr = np.array([0, 2, 3, 4, 6]).astype(index_dtype) indices = np.array([0, 2, 5, 0, 4, 5]).astype(index_dtype) shape = (4, 6) dim_names = ('x', 'y') sparse_tensor = pa.SparseCSRMatrix.from_numpy(data, indptr, indices, shape, dim_names) context = pa.default_serialization_context() serialized = pa.serialize(sparse_tensor, context=context).to_buffer() result = pa.deserialize(serialized) assert_equal(result, sparse_tensor) assert isinstance(result, pa.SparseCSRMatrix) data_result, indptr_result, indices_result = result.to_numpy() assert np.array_equal(data_result, data) assert np.array_equal(indptr_result, indptr) assert np.array_equal(indices_result, indices) assert result.dim_names == dim_names @pytest.mark.parametrize('tensor_type', tensor_types) @pytest.mark.parametrize('index_type', index_types) def test_sparse_csr_matrix_components_serialization(large_buffer, index_type, tensor_type): tensor_dtype = np.dtype(tensor_type) index_dtype = np.dtype(index_type) data = np.array([8, 2, 5, 3, 4, 6]).astype(tensor_dtype) indptr = np.array([0, 2, 3, 4, 6]).astype(index_dtype) indices = np.array([0, 2, 5, 0, 4, 5]).astype(index_dtype) shape = (4, 6) dim_names = ('x', 'y') sparse_tensor = pa.SparseCSRMatrix.from_numpy(data, indptr, indices, shape, dim_names) serialization_roundtrip(sparse_tensor, large_buffer) @pytest.mark.skipif(not csr_matrix, reason="requires scipy") def test_scipy_sparse_csr_matrix_serialization(): data = np.array([8, 2, 5, 3, 4, 6]) indptr = np.array([0, 2, 3, 4, 6]) indices = np.array([0, 2, 5, 0, 4, 5]) shape = (4, 6) sparse_array = csr_matrix((data, indices, indptr), shape=shape) serialized = pa.serialize(sparse_array) result = serialized.deserialize() assert np.array_equal(sparse_array.toarray(), result.toarray()) @pytest.mark.parametrize('tensor_type', tensor_types) @pytest.mark.parametrize('index_type', index_types) def test_sparse_csc_matrix_serialization(index_type, tensor_type): tensor_dtype = np.dtype(tensor_type) index_dtype = np.dtype(index_type) data = np.array([[8, 2, 5, 3, 4, 6]]).T.astype(tensor_dtype) indptr = np.array([0, 2, 3, 4, 6]).astype(index_dtype) indices = np.array([0, 2, 5, 0, 4, 5]).astype(index_dtype) shape = (6, 4) dim_names = ('x', 'y') sparse_tensor = pa.SparseCSCMatrix.from_numpy(data, indptr, indices, shape, dim_names) context = pa.default_serialization_context() serialized = pa.serialize(sparse_tensor, context=context).to_buffer() result = pa.deserialize(serialized) assert_equal(result, sparse_tensor) assert isinstance(result, pa.SparseCSCMatrix) data_result, indptr_result, indices_result = result.to_numpy() assert np.array_equal(data_result, data) assert np.array_equal(indptr_result, indptr) assert np.array_equal(indices_result, indices) assert result.dim_names == dim_names @pytest.mark.parametrize('tensor_type', tensor_types) @pytest.mark.parametrize('index_type', index_types) def test_sparse_csc_matrix_components_serialization(large_buffer, index_type, tensor_type): tensor_dtype = np.dtype(tensor_type) index_dtype = np.dtype(index_type) data = np.array([8, 2, 5, 3, 4, 6]).astype(tensor_dtype) indptr = np.array([0, 2, 3, 6]).astype(index_dtype) indices = np.array([0, 2, 2, 0, 1, 2]).astype(index_dtype) shape = (3, 3) dim_names = ('x', 'y') sparse_tensor = pa.SparseCSCMatrix.from_numpy(data, indptr, indices, shape, dim_names) serialization_roundtrip(sparse_tensor, large_buffer) @pytest.mark.skipif(not csc_matrix, reason="requires scipy") def test_scipy_sparse_csc_matrix_serialization(): data = np.array([8, 2, 5, 3, 4, 6]) indptr = np.array([0, 2, 3, 4, 6]) indices = np.array([0, 2, 5, 0, 4, 5]) shape = (6, 4) sparse_array = csc_matrix((data, indices, indptr), shape=shape) serialized = pa.serialize(sparse_array) result = serialized.deserialize() assert np.array_equal(sparse_array.toarray(), result.toarray()) @pytest.mark.parametrize('tensor_type', tensor_types) @pytest.mark.parametrize('index_type', index_types) def test_sparse_csf_tensor_serialization(index_type, tensor_type): tensor_dtype = np.dtype(tensor_type) index_dtype = np.dtype(index_type) data = np.array([[1, 2, 3, 4, 5, 6, 7, 8]]).T.astype(tensor_dtype) indptr = [ np.array([0, 2, 3]), np.array([0, 1, 3, 4]), np.array([0, 2, 4, 5, 8]), ] indices = [ np.array([0, 1]), np.array([0, 1, 1]), np.array([0, 0, 1, 1]), np.array([1, 2, 0, 2, 0, 0, 1, 2]), ] indptr = [x.astype(index_dtype) for x in indptr] indices = [x.astype(index_dtype) for x in indices] shape = (2, 3, 4, 5) axis_order = (0, 1, 2, 3) dim_names = ("a", "b", "c", "d") for ndim in [2, 3, 4]: sparse_tensor = pa.SparseCSFTensor.from_numpy(data, indptr[:ndim - 1], indices[:ndim], shape[:ndim], axis_order[:ndim], dim_names[:ndim]) context = pa.default_serialization_context() serialized = pa.serialize(sparse_tensor, context=context).to_buffer() result = pa.deserialize(serialized) assert_equal(result, sparse_tensor) assert isinstance(result, pa.SparseCSFTensor) @pytest.mark.parametrize('tensor_type', tensor_types) @pytest.mark.parametrize('index_type', index_types) def test_sparse_csf_tensor_components_serialization(large_buffer, index_type, tensor_type): tensor_dtype = np.dtype(tensor_type) index_dtype = np.dtype(index_type) data = np.array([[1, 2, 3, 4, 5, 6, 7, 8]]).T.astype(tensor_dtype) indptr = [ np.array([0, 2, 3]), np.array([0, 1, 3, 4]), np.array([0, 2, 4, 5, 8]), ] indices = [ np.array([0, 1]), np.array([0, 1, 1]), np.array([0, 0, 1, 1]), np.array([1, 2, 0, 2, 0, 0, 1, 2]), ] indptr = [x.astype(index_dtype) for x in indptr] indices = [x.astype(index_dtype) for x in indices] shape = (2, 3, 4, 5) axis_order = (0, 1, 2, 3) dim_names = ("a", "b", "c", "d") for ndim in [2, 3, 4]: sparse_tensor = pa.SparseCSFTensor.from_numpy(data, indptr[:ndim - 1], indices[:ndim], shape[:ndim], axis_order[:ndim], dim_names[:ndim]) serialization_roundtrip(sparse_tensor, large_buffer) @pytest.mark.filterwarnings( "ignore:the matrix subclass:PendingDeprecationWarning") def test_numpy_matrix_serialization(tmpdir): class CustomType: def __init__(self, val): self.val = val rec_type = np.dtype([('x', 'int64'), ('y', 'double'), ('z', '<U4')]) path = os.path.join(str(tmpdir), 'pyarrow_npmatrix_serialization_test.bin') array = np.random.randint(low=-1, high=1, size=(2, 2)) for data_type in [str, int, float, rec_type, CustomType]: matrix = np.matrix(array.astype(data_type)) with open(path, 'wb') as f: f.write(pa.serialize(matrix).to_buffer()) serialized = pa.read_serialized(pa.OSFile(path)) result = serialized.deserialize() assert_equal(result, matrix) assert_equal(result.dtype, matrix.dtype) serialized = None assert_equal(result, matrix) assert result.base is not None def test_pyarrow_objects_serialization(large_buffer): # NOTE: We have to put these objects inside, # or it will affect 'test_total_bytes_allocated'. pyarrow_objects = [ pa.array([1, 2, 3, 4]), pa.array(['1', 'never U+1F631', '', "233 * U+1F600"]), pa.array([1, None, 2, 3]), pa.Tensor.from_numpy(np.random.rand(2, 3, 4)), pa.RecordBatch.from_arrays( [pa.array([1, None, 2, 3]), pa.array(['1', 'never U+1F631', '', "233 * u1F600"])], ['a', 'b']), pa.Table.from_arrays([pa.array([1, None, 2, 3]), pa.array(['1', 'never U+1F631', '', "233 * u1F600"])], ['a', 'b']) ] for obj in pyarrow_objects: serialization_roundtrip(obj, large_buffer) def test_buffer_serialization(): class BufferClass: pass def serialize_buffer_class(obj): return pa.py_buffer(b"hello") def deserialize_buffer_class(serialized_obj): return serialized_obj context = pa.default_serialization_context() context.register_type( BufferClass, "BufferClass", custom_serializer=serialize_buffer_class, custom_deserializer=deserialize_buffer_class) b = pa.serialize(BufferClass(), context=context).to_buffer() assert pa.deserialize(b, context=context).to_pybytes() == b"hello" @pytest.mark.skip(reason="extensive memory requirements") def test_arrow_limits(self): def huge_memory_map(temp_dir): return large_memory_map(temp_dir, 100 * 1024 * 1024 * 1024) with pa.memory_map(huge_memory_map, mode="r+") as mmap: # Test that objects that are too large for Arrow throw a Python # exception. These tests give out of memory errors on Travis and need # to be run on a machine with lots of RAM. x = 2 ** 29 * [1.0] serialization_roundtrip(x, mmap) del x x = 2 ** 29 * ["s"] serialization_roundtrip(x, mmap) del x x = 2 ** 29 * [["1"], 2, 3, [{"s": 4}]] serialization_roundtrip(x, mmap) del x x = 2 ** 29 * [{"s": 1}] + 2 ** 29 * [1.0] serialization_roundtrip(x, mmap) del x x = np.zeros(2 ** 25) serialization_roundtrip(x, mmap) del x x = [np.zeros(2 ** 18) for _ in range(2 ** 7)] serialization_roundtrip(x, mmap) del x def test_serialization_callback_error(): class TempClass: pass # Pass a SerializationContext into serialize, but TempClass # is not registered serialization_context = pa.SerializationContext() val = TempClass() with pytest.raises(pa.SerializationCallbackError) as err: serialized_object = pa.serialize(val, serialization_context) assert err.value.example_object == val serialization_context.register_type(TempClass, "TempClass") serialized_object = pa.serialize(TempClass(), serialization_context) deserialization_context = pa.SerializationContext() # Pass a Serialization Context into deserialize, but TempClass # is not registered with pytest.raises(pa.DeserializationCallbackError) as err: serialized_object.deserialize(deserialization_context) assert err.value.type_id == "TempClass" class TempClass2: pass # Make sure that we receive an error when we use an inappropriate value for # the type_id argument. with pytest.raises(TypeError): serialization_context.register_type(TempClass2, 1) def test_fallback_to_subclasses(): class SubFoo(Foo): def __init__(self): Foo.__init__(self) # should be able to serialize/deserialize an instance # if a base class has been registered serialization_context = pa.SerializationContext() serialization_context.register_type(Foo, "Foo") subfoo = SubFoo() # should fallbact to Foo serializer serialized_object = pa.serialize(subfoo, serialization_context) reconstructed_object = serialized_object.deserialize( serialization_context ) assert type(reconstructed_object) == Foo class Serializable: pass def serialize_serializable(obj): return {"type": type(obj), "data": obj.__dict__} def deserialize_serializable(obj): val = obj["type"].__new__(obj["type"]) val.__dict__.update(obj["data"]) return val class SerializableClass(Serializable): def __init__(self): self.value = 3 def test_serialize_subclasses(): # This test shows how subclasses can be handled in an idiomatic way # by having only a serializer for the base class # This technique should however be used with care, since pickling # type(obj) with couldpickle will include the full class definition # in the serialized representation. # This means the class definition is part of every instance of the # object, which in general is not desirable; registering all subclasses # with register_type will result in faster and more memory # efficient serialization. context = pa.default_serialization_context() context.register_type( Serializable, "Serializable", custom_serializer=serialize_serializable, custom_deserializer=deserialize_serializable) a = SerializableClass() serialized = pa.serialize(a, context=context) deserialized = serialized.deserialize(context=context) assert type(deserialized).__name__ == SerializableClass.__name__ assert deserialized.value == 3 def test_serialize_to_components_invalid_cases(): buf = pa.py_buffer(b'hello') components = { 'num_tensors': 0, 'num_sparse_tensors': { 'coo': 0, 'csr': 0, 'csc': 0, 'csf': 0, 'ndim_csf': 0 }, 'num_ndarrays': 0, 'num_buffers': 1, 'data': [buf] } with pytest.raises(pa.ArrowInvalid): pa.deserialize_components(components) components = { 'num_tensors': 0, 'num_sparse_tensors': { 'coo': 0, 'csr': 0, 'csc': 0, 'csf': 0, 'ndim_csf': 0 }, 'num_ndarrays': 1, 'num_buffers': 0, 'data': [buf, buf] } with pytest.raises(pa.ArrowInvalid): pa.deserialize_components(components) def test_deserialize_components_in_different_process(): arr = pa.array([1, 2, 5, 6], type=pa.int8()) ser = pa.serialize(arr) data = pickle.dumps(ser.to_components(), protocol=-1) code = """if 1: import pickle import pyarrow as pa data = {!r} components = pickle.loads(data) arr = pa.deserialize_components(components) assert arr.to_pylist() == [1, 2, 5, 6], arr """.format(data) subprocess_env = test_util.get_modified_env_with_pythonpath() print("** sys.path =", sys.path) print("** setting PYTHONPATH to:", subprocess_env['PYTHONPATH']) subprocess.check_call([sys.executable, "-c", code], env=subprocess_env) def test_serialize_read_concatenated_records(): # ARROW-1996 -- see stream alignment work in ARROW-2840, ARROW-3212 f = pa.BufferOutputStream() pa.serialize_to(12, f) pa.serialize_to(23, f) buf = f.getvalue() f = pa.BufferReader(buf) pa.read_serialized(f).deserialize() pa.read_serialized(f).deserialize() def deserialize_regex(serialized, q): import pyarrow as pa q.put(pa.deserialize(serialized)) def test_deserialize_in_different_process(): from multiprocessing import Process, Queue import re regex = re.compile(r"\d+\.\d*") serialization_context = pa.SerializationContext() serialization_context.register_type(type(regex), "Regex", pickle=True) serialized = pa.serialize(regex, serialization_context) serialized_bytes = serialized.to_buffer().to_pybytes() q = Queue() p = Process(target=deserialize_regex, args=(serialized_bytes, q)) p.start() assert q.get().pattern == regex.pattern p.join() def test_deserialize_buffer_in_different_process(): import tempfile f = tempfile.NamedTemporaryFile(delete=False) b = pa.serialize(pa.py_buffer(b'hello')).to_buffer() f.write(b.to_pybytes()) f.close() test_util.invoke_script('deserialize_buffer.py', f.name) def test_set_pickle(): # Use a custom type to trigger pickling. class Foo: pass context = pa.SerializationContext() context.register_type(Foo, 'Foo', pickle=True) test_object = Foo() # Define a custom serializer and deserializer to use in place of pickle. def dumps1(obj): return b'custom' def loads1(serialized_obj): return serialized_obj + b' serialization 1' # Test that setting a custom pickler changes the behavior. context.set_pickle(dumps1, loads1) serialized = pa.serialize(test_object, context=context).to_buffer() deserialized = pa.deserialize(serialized.to_pybytes(), context=context) assert deserialized == b'custom serialization 1' # Define another custom serializer and deserializer. def dumps2(obj): return b'custom' def loads2(serialized_obj): return serialized_obj + b' serialization 2' # Test that setting another custom pickler changes the behavior again. context.set_pickle(dumps2, loads2) serialized = pa.serialize(test_object, context=context).to_buffer() deserialized = pa.deserialize(serialized.to_pybytes(), context=context) assert deserialized == b'custom serialization 2' def test_path_objects(tmpdir): # Test compatibility with PEP 519 path-like objects p = pathlib.Path(tmpdir) / 'zzz.bin' obj = 1234 pa.serialize_to(obj, p) res = pa.deserialize_from(p, None) assert res == obj def test_tensor_alignment(): # Deserialized numpy arrays should be 64-byte aligned. x = np.random.normal(size=(10, 20, 30)) y = pa.deserialize(pa.serialize(x).to_buffer()) assert y.ctypes.data % 64 == 0 xs = [np.random.normal(size=i) for i in range(100)] ys = pa.deserialize(pa.serialize(xs).to_buffer()) for y in ys: assert y.ctypes.data % 64 == 0 xs = [np.random.normal(size=i * (1,)) for i in range(20)] ys = pa.deserialize(pa.serialize(xs).to_buffer()) for y in ys: assert y.ctypes.data % 64 == 0 xs = [np.random.normal(size=i * (5,)) for i in range(1, 8)] xs = [xs[i][(i + 1) * (slice(1, 3),)] for i in range(len(xs))] ys = pa.deserialize(pa.serialize(xs).to_buffer()) for y in ys: assert y.ctypes.data % 64 == 0 def test_empty_tensor(): # ARROW-8122, serialize and deserialize empty tensors x = np.array([], dtype=np.float64) y = pa.deserialize(pa.serialize(x).to_buffer()) np.testing.assert_array_equal(x, y) x = np.array([[], [], []], dtype=np.float64) y = pa.deserialize(pa.serialize(x).to_buffer()) np.testing.assert_array_equal(x, y) x = np.array([[], [], []], dtype=np.float64).T y = pa.deserialize(pa.serialize(x).to_buffer()) np.testing.assert_array_equal(x, y) def test_serialization_determinism(): for obj in COMPLEX_OBJECTS: buf1 = pa.serialize(obj).to_buffer() buf2 = pa.serialize(obj).to_buffer() assert buf1.to_pybytes() == buf2.to_pybytes() def test_serialize_recursive_objects(): class ClassA: pass # Make a list that contains itself. lst = [] lst.append(lst) # Make an object that contains itself as a field. a1 = ClassA() a1.field = a1 # Make two objects that contain each other as fields. a2 = ClassA() a3 = ClassA() a2.field = a3 a3.field = a2 # Make a dictionary that contains itself. d1 = {} d1["key"] = d1 # Make a numpy array that contains itself. arr = np.array([None], dtype=object) arr[0] = arr # Create a list of recursive objects. recursive_objects = [lst, a1, a2, a3, d1, arr] # Check that exceptions are thrown when we serialize the recursive # objects. for obj in recursive_objects: with pytest.raises(Exception): pa.serialize(obj).deserialize()
Save
cmd:
run