try:

import torch

if torch.is_tensor(obj1) and torch.is_tensor(obj2):

assert torch.equal(obj1, obj2)

return

except ImportError:

pass

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))

try:

# Workaround to make comparison of OrderedDicts work on Python 2.7

if obj1 == obj2:

return

except Exception:

pass

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])

else:

assert obj1 == obj2, ("Objects {} and {} are different."

0, 0.0, 0.9, 1 << 62, 1 << 100, 1 << 999,

[1 << 100, [1 << 100]], "a", string.printable, "\u262F",

"hello world", u"hello world", u"\xff\xfe\x9c\x001\x000\x00",

None, True, False, [], (), {}, {(1, 2): 1}, {(): 2},

[1, "hello", 3.0], u"\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": set([2, 3]), "world": set([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([100, 100]),

np.random.normal(size=[100, 100]), np.array(["hi", 3]),

np.array(["hi", 3], dtype=object),

[[[[[[[[[[[[]]]]]]]]]]]],

{"obj{}".format(i): np.random.normal(size=[100, 100]) for i in range(10)},

# {(): {(): {(): {(): {(): {(): {(): {(): {(): {(): {

# (): {(): {}}}}}}}}}}}}},

((((((((((),),),),),),),),),),

def __init__(self, value=0):

self.value = value

def __hash__(self):

return hash(self.value)

def __eq__(self, other):

def __init__(self):

self.foo = Foo()

self.bar = Bar()

def method(self, arg):

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")

f.seek(0)

pa.serialize_to(value, f, serialization_context)

f.seek(0)

result = pa.deserialize_from(f, None, serialization_context)

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])

for nthreads in [1, 4]:

for value in COMPLEX_OBJECTS:

buf = pa.serialize(value).to_buffer(nthreads=nthreads)

result = pa.deserialize(buf)

pytest.importorskip("cloudpickle")

with pa.memory_map(large_memory_map, mode="r+") as mmap:

obj = defaultdict(lambda: 0, [("hello", 1), ("world", 2)])

with pa.memory_map(large_memory_map, mode="r+") as mmap:

for t in ["bool", "int8", "uint8", "int16", "uint16", "int32",

"uint32", "float16", "float32", "float64"]:

obj = np.random.randint(0, 10, size=(100, 100)).astype(t)

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)

]

with pa.memory_map(large_memory_map, mode="r+") as mmap:

for d in data:

pytest.importorskip("torch")

import torch

with pa.memory_map(large_memory_map, mode="r+") as mmap:

# 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))

with pa.memory_map(large_memory_map, mode="r+") as mmap:

obj = np.zeros([10])

mmap.seek(0)

pa.serialize_to(obj, mmap, serialization_context)

mmap.seek(0)

result = pa.deserialize_from(mmap, None, serialization_context)

with pytest.raises(ValueError):

class DummyClass(object):

pass

def serialize_dummy_class(obj):

x = np.zeros(4)

return x

def deserialize_dummy_class(serialized_obj):

return serialized_obj

pa._default_serialization_context.register_type(

DummyClass, "DummyClass", pickle=False,

custom_serializer=serialize_dummy_class,

custom_deserializer=deserialize_dummy_class)

class BufferClass(object):

pass

def serialize_buffer_class(obj):

return pa.frombuffer(b"hello")

def deserialize_buffer_class(serialized_obj):

return serialized_obj

pa._default_serialization_context.register_type(

BufferClass, "BufferClass", pickle=False,

custom_serializer=serialize_buffer_class,

custom_deserializer=deserialize_buffer_class)

b = pa.serialize(BufferClass()).to_buffer()

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)

class TempClass(object):

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, 20*b"\x00")

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)

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

)

# 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.

serialization_context.register_type(

Serializable, "Serializable",

custom_serializer=serialize_serializable,

custom_deserializer=deserialize_serializable)

a = SerializableClass()

serialized = pa.serialize(a)

deserialized = serialized.deserialize()

assert type(deserialized).__name__ == SerializableClass.__name__