1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
|
from __future__ import annotations
import pytest
from distributed.protocol import deserialize, serialize
np = pytest.importorskip("numpy")
pd = pytest.importorskip("pandas")
@pytest.mark.parametrize("collection", [tuple, dict, list])
@pytest.mark.parametrize(
"y,y_serializer",
[
(np.arange(50), "dask"),
(pd.DataFrame({"C": ["a", "b", None], "D": [2.5, 3.5, 4.5]}), "pickle"),
(None, "pickle"),
],
)
def test_serialize_collection(collection, y, y_serializer):
x = np.arange(100)
if issubclass(collection, dict):
header, frames = serialize({"x": x, "y": y}, serializers=("dask", "pickle"))
else:
header, frames = serialize(collection((x, y)), serializers=("dask", "pickle"))
frames = tuple(frames) # verify that no mutation occurs
t = deserialize(header, frames, deserializers=("dask", "pickle", "error"))
assert isinstance(t, collection)
assert header["is-collection"] is True
sub_headers = header["sub-headers"]
if collection is not dict:
assert sub_headers[0]["serializer"] == "dask"
assert sub_headers[1]["serializer"] == y_serializer
if collection is dict:
assert (t["x"] == x).all()
assert str(t["y"]) == str(y)
else:
assert (t[0] == x).all()
assert str(t[1]) == str(y)
def test_large_collections_serialize_simply():
header, frames = serialize(tuple(range(1000)))
assert len(frames) == 1
def test_nested_types():
x = np.ones(5)
header, frames = serialize([[[x]]])
assert "dask" in str(header)
assert len(frames) == 1
assert x.data == np.frombuffer(frames[0]).data
|
