Add GxB serialize and deserialize for matrices and vectors. (#44)

These are thin wrappers as per typical for `python-suitesparse-graphblas`

Author: eriknw

setup.cfg

@@ -3,6 +3,7 @@ test=pytest
 
 [flake8]
 max-line-length = 100
+inline-quotes = "
 exclude =
     versioneer.py,
 ignore =

suitesparse_graphblas/__init__.py

@@ -117,11 +117,15 @@ def libget(name):
     lib.GrB_DOMAIN_MISMATCH: ex.DomainMismatch,
     lib.GrB_DIMENSION_MISMATCH: ex.DimensionMismatch,
     lib.GrB_OUTPUT_NOT_EMPTY: ex.OutputNotEmpty,
+    lib.GrB_EMPTY_OBJECT: ex.EmptyObject,
     # Execution Errors
     lib.GrB_OUT_OF_MEMORY: ex.OutOfMemory,
     lib.GrB_INSUFFICIENT_SPACE: ex.InsufficientSpace,
     lib.GrB_INDEX_OUT_OF_BOUNDS: ex.IndexOutOfBound,
     lib.GrB_PANIC: ex.Panic,
+    lib.GrB_NOT_IMPLEMENTED: ex.NotImplementedException,
+    # GxB Errors
+    lib.GxB_EXHAUSTED: StopIteration,
 }
 GrB_SUCCESS = lib.GrB_SUCCESS
 GrB_NO_VALUE = lib.GrB_NO_VALUE

suitesparse_graphblas/exceptions.py

@@ -38,6 +38,10 @@ class OutputNotEmpty(GraphBLASException):
     pass
 
 
+class EmptyObject(GraphBLASException):
+    pass
+
+
 class OutOfMemory(GraphBLASException):
     pass
 
@@ -52,3 +56,7 @@ class IndexOutOfBound(GraphBLASException):
 
 class Panic(GraphBLASException):
     pass
+
+
+class NotImplementedException(GraphBLASException):
+    pass

suitesparse_graphblas/io/serialize.py

@@ -0,0 +1,133 @@
+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(desc[0], lib.GxB_NTHREADS, ffi.cast("int", nthreads)))
+    if compression is not None:
+        if level is not None and compression == lib.GxB_COMPRESSION_LZ4HC:
+            compression += level
+        check_status(
+            desc, lib.GxB_Desc_set(desc[0], lib.GxB_COMPRESSION, ffi.cast("int", 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, or GxB_COMPRESSION_LZ4HC
+    level : int, optional
+        Used by GxB_COMPRESSION_LZ4HC.  Should be between 1 and 9, where 9 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, or GxB_COMPRESSION_LZ4HC
+    level : int, optional
+        Used by GxB_COMPRESSION_LZ4HC.  Should be between 1 and 9, where 9 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

suitesparse_graphblas/matrix.py

@@ -1,5 +1,8 @@
 from suitesparse_graphblas import check_status, ffi, lib
 
+from .io.serialize import deserialize_matrix as deserialize  # noqa
+from .io.serialize import serialize_matrix as serialize  # noqa
+
 
 def free(A):
     """Free a matrix."""

suitesparse_graphblas/tests/test_io.py

@@ -18,6 +18,7 @@
     signed_integer_types,
     supports_complex,
     unsigned_integer_types,
+    vector,
 )
 
 if platform.system() == "Windows":
@@ -87,6 +88,53 @@ def _test_elements(T):
     )
 
 
+def test_serialize_matrix():
+    T = lib.GrB_INT64
+    A = matrix.new(T, 2, 2)
+    for args in zip(*_test_elements(T)):
+        f = _element_setters[T]
+        check_status(A, f(A[0], *args))
+    data = matrix.serialize(A)
+    B = matrix.deserialize(data)
+
+    # Test equal
+    C = matrix.new(lib.GrB_BOOL, 2, 2)
+    check_status(
+        C,
+        lib.GrB_Matrix_eWiseAdd_BinaryOp(C[0], NULL, NULL, _eq_ops[T], A[0], B[0], NULL),
+    )
+    assert matrix.nvals(A) == matrix.nvals(B) == matrix.nvals(C)
+    is_eq = ffi.new("bool*")
+    check_status(
+        C,
+        lib.GrB_Matrix_reduce_BOOL(is_eq, NULL, lib.GrB_LAND_MONOID_BOOL, C[0], NULL),
+    )
+    assert is_eq[0]
+
+
+def test_serialize_vector():
+    T = lib.GrB_INT64
+    v = vector.new(T, 3)
+    check_status(v, lib.GrB_Vector_setElement_INT64(v[0], 2, 0))
+    check_status(v, lib.GrB_Vector_setElement_INT64(v[0], 10, 1))
+    data = vector.serialize(v, lib.GxB_COMPRESSION_LZ4HC, level=7)
+    w = vector.deserialize(data)
+
+    # Test equal
+    x = vector.new(lib.GrB_BOOL, 3)
+    check_status(
+        x,
+        lib.GrB_Vector_eWiseAdd_BinaryOp(x[0], NULL, NULL, _eq_ops[T], v[0], w[0], NULL),
+    )
+    assert vector.nvals(v) == vector.nvals(w) == vector.nvals(x)
+    is_eq = ffi.new("bool*")
+    check_status(
+        x,
+        lib.GrB_Vector_reduce_BOOL(is_eq, NULL, lib.GrB_LAND_MONOID_BOOL, x[0], NULL),
+    )
+    assert is_eq[0]
+
+
 def test_matrix_binfile_read_write(tmp_path):
     for opener in (Path.open, gzip.open, bz2.open, lzma.open):
         for format in (lib.GxB_BY_ROW, lib.GxB_BY_COL):

suitesparse_graphblas/vector.py

@@ -1,5 +1,8 @@
 from suitesparse_graphblas import check_status, ffi, lib
 
+from .io.serialize import deserialize_vector as deserialize  # noqa
+from .io.serialize import serialize_vector as serialize  # noqa
+
 
 def free(v):
     """Free a vector."""