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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
|
import numpy as np
from suitesparse_graphblas import check_status, ffi, lib
from suitesparse_graphblas.utils import claim_buffer
def free_desc(desc):
"""Free a descriptor."""
check_status(desc, lib.GrB_Descriptor_free(desc))
def get_serialize_desc(compression=lib.GxB_COMPRESSION_DEFAULT, level=None, nthreads=None):
"""Create a descriptor for serializing or deserializing.
This returns None (for NULL descriptor) or a pointer to a GrB_Descriptor.
"""
if nthreads is None and (compression is None or compression == lib.GxB_COMPRESSION_DEFAULT):
return None
desc = ffi.new("GrB_Descriptor*")
check_status(desc, lib.GrB_Descriptor_new(desc))
desc = ffi.gc(desc, free_desc)
if nthreads is not None:
check_status(
desc,
lib.GxB_Desc_set_INT32(desc[0], lib.GxB_NTHREADS, ffi.cast("int32_t", nthreads)),
)
if compression is not None:
if level is not None and compression in {
lib.GxB_COMPRESSION_LZ4HC,
lib.GxB_COMPRESSION_ZSTD,
}:
compression += level
check_status(
desc,
lib.GxB_Desc_set_INT32(desc[0], lib.GxB_COMPRESSION, ffi.cast("int32_t", compression)),
)
return desc
def serialize_matrix(A, compression=lib.GxB_COMPRESSION_DEFAULT, level=None, *, nthreads=None):
"""Serialize a Matrix into an array of bytes.
Parameters
----------
compression : int, optional
One of None, GxB_COMPRESSION_NONE, GxB_COMPRESSION_DEFAULT,
GxB_COMPRESSION_LZ4, GxB_COMPRESSION_LZ4HC, or GxB_COMPRESSION_ZSTD
level : int, optional
For GxB_COMPRESSION_LZ4HC, should be between 1 and 9, where 9 is most compressed.
For GxB_COMPRESSION_ZSTD, should be between 1 and 19, where 19 is most compressed.
nthreads : int, optional
The maximum number of OpenMP threads to use.
"""
desc = get_serialize_desc(compression, level, nthreads)
data_ptr = ffi.new("void**")
size_ptr = ffi.new("GrB_Index*")
check_status(
A, lib.GxB_Matrix_serialize(data_ptr, size_ptr, A[0], ffi.NULL if desc is None else desc[0])
)
return claim_buffer(ffi, data_ptr[0], size_ptr[0], np.dtype(np.uint8))
def serialize_vector(v, compression=lib.GxB_COMPRESSION_DEFAULT, level=None, *, nthreads=None):
"""Serialize a Vector into an array of bytes.
Parameters
----------
compression : int, optional
One of None, GxB_COMPRESSION_NONE, GxB_COMPRESSION_DEFAULT,
GxB_COMPRESSION_LZ4, GxB_COMPRESSION_LZ4HC, or GxB_COMPRESSION_ZSTD
level : int, optional
For GxB_COMPRESSION_LZ4HC, should be between 1 and 9, where 9 is most compressed.
For GxB_COMPRESSION_ZSTD, should be between 1 and 19, where 19 is most compressed.
nthreads : int, optional
The maximum number of OpenMP threads to use.
"""
desc = get_serialize_desc(compression, level, nthreads)
data_ptr = ffi.new("void**")
size_ptr = ffi.new("GrB_Index*")
check_status(
v, lib.GxB_Vector_serialize(data_ptr, size_ptr, v[0], ffi.NULL if desc is None else desc[0])
)
return claim_buffer(ffi, data_ptr[0], size_ptr[0], np.dtype(np.uint8))
def deserialize_matrix(data, *, free=True, nthreads=None):
"""Deserialize a Matrix from bytes.
The `free` argument is called when the object is garbage
collected, the default is `matrix.free()`. If `free` is None then
there is no automatic garbage collection and it is up to the user
to free the matrix.
"""
data = np.frombuffer(data, np.uint8)
desc = get_serialize_desc(None, nthreads)
A = ffi.new("GrB_Matrix*")
check_status(
A,
lib.GxB_Matrix_deserialize(
A,
ffi.NULL, # dtype; we don't check for now
ffi.from_buffer("void*", data),
data.nbytes,
ffi.NULL if desc is None else desc[0],
),
)
if free:
if callable(free):
return ffi.gc(A, free)
return ffi.gc(A, matrix.free)
return A
def deserialize_vector(data, *, free=True, nthreads=None):
"""Deserialize a Vector from bytes.
The `free` argument is called when the object is garbage
collected, the default is `vector.free()`. If `free` is None then
there is no automatic garbage collection and it is up to the user
to free the vector.
"""
data = np.frombuffer(data, np.uint8)
desc = get_serialize_desc(None, nthreads)
v = ffi.new("GrB_Vector*")
check_status(
v,
lib.GxB_Vector_deserialize(
v,
ffi.NULL, # dtype; we don't check for now
ffi.from_buffer("void*", data),
data.nbytes,
ffi.NULL if desc is None else desc[0],
),
)
if free:
if callable(free):
return ffi.gc(v, free)
return ffi.gc(v, vector.free)
return v
from suitesparse_graphblas import matrix, vector # noqa isort:skip
|
