File: sparse.py

package info (click to toggle)
dask.distributed 2021.01.0%2Bds.1-2.1%2Bdeb11u1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 6,668 kB
  • sloc: python: 54,131; javascript: 1,549; makefile: 207; sh: 100
file content (31 lines) | stat: -rw-r--r-- 879 bytes parent folder | download
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)