forked from dask/distributed
-
Notifications
You must be signed in to change notification settings - Fork 0
/
sparse.py
31 lines (21 loc) · 879 Bytes
/
sparse.py
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
from .serialize import dask_serialize, dask_deserialize, serialize, deserialize
import sparse
@dask_serialize.register(sparse.COO)
def serialize_sparse(x):
coords_header, coords_frames = serialize(x.coords)
data_header, data_frames = serialize(x.data)
header = {
"coords-header": coords_header,
"data-header": data_header,
"shape": x.shape,
"nframes": [len(coords_frames), len(data_frames)],
}
return header, coords_frames + data_frames
@dask_deserialize.register(sparse.COO)
def deserialize_sparse(header, frames):
coords_frames = frames[: header["nframes"][0]]
data_frames = frames[header["nframes"][0] :]
coords = deserialize(header["coords-header"], coords_frames)
data = deserialize(header["data-header"], data_frames)
shape = header["shape"]
return sparse.COO(coords, data, shape=shape)