Get COO round-trip working.

Author: hameerabbasi

sparse/numba_backend/__init__.py

@@ -157,7 +157,7 @@
     where,
 )
 from ._dok import DOK
-from ._io import load_npz, save_npz
+from ._io import from_binsparse, load_npz, save_npz
 from ._umath import elemwise
 from ._utils import random
 
@@ -226,6 +226,7 @@
     "float64",
     "floor",
     "floor_divide",
+    "from_binsparse",
     "full",
     "full_like",
     "greater",

sparse/numba_backend/_common.py

@@ -35,7 +35,7 @@ def _check_device(func):
     def wrapped(*args, **kwargs):
         device = kwargs.get("device", None)
         if device not in {"cpu", None}:
-            raise ValueError("Device must be `'cpu'` or `None`.")
+            raise BufferError("Device must be `'cpu'` or `None`.")
         return func(*args, **kwargs)
 
     return wrapped

sparse/numba_backend/_io.py

@@ -1,7 +1,9 @@
 import numpy as np
 
+from ._common import _check_device
 from ._compressed import GCXS
 from ._coo.core import COO
+from ._sparse_array import SparseArray
 
 
 def save_npz(filename, matrix, compressed=True):
@@ -130,3 +132,127 @@ def load_npz(filename):
             )
         except KeyError as e:
             raise RuntimeError(f"The file {filename!s} does not contain a valid sparse matrix") from e
+
+
+@_check_device
+def from_binsparse(arr, /, *, device=None, copy: bool | None = None) -> SparseArray:
+    desc, arrs = arr.__binsparse__()
+
+    desc = desc["binsparse"]
+    version_tuple: tuple[int, ...] = tuple(int(v) for v in desc["version"].split("."))
+    if version_tuple != (0, 1):
+        raise RuntimeError("Unsupported `__binsparse__` protocol version.")
+
+    format = desc["format"]
+    format_err_str = f"Unsupported format: `{format!r}`."
+    invalid_dtype_str = "Invalid dtype: `{dtype!s}`, expected `{expected!s}`."
+
+    if isinstance(format, str):
+        match format:
+            case "COO" | "COOR":
+                desc["format"] = {
+                    "custom": {
+                        "transpose": [0, 1],
+                        "level": {
+                            "level_desc": "sparse",
+                            "rank": 2,
+                            "level": {
+                                "level_desc": "element",
+                            },
+                        },
+                    }
+                }
+            case "CSC" | "CSR":
+                desc["format"] = {
+                    "custom": {
+                        "transpose": [0, 1] if format == "CSR" else [0, 1],
+                        "level": {
+                            "level_desc": "dense",
+                            "level": {
+                                "level_desc": "sparse",
+                                "level": {
+                                    "level_desc": "element",
+                                },
+                            },
+                        },
+                    },
+                }
+            case _:
+                raise RuntimeError(format_err_str)
+
+    format = desc["format"]
+    if "transpose" not in format:
+        rank = 0
+        level = format
+        while "level" in level:
+            if "rank" not in level:
+                level["rank"] = 1
+            rank += level["rank"]
+
+        format["transpose"] = list(range(rank))
+
+    match desc:
+        case {
+            "format": {
+                "custom": {
+                    "transpose": transpose,
+                    "level": {
+                        "level_desc": "sparse",
+                        "rank": ndim,
+                        "level": {
+                            "level_desc": "element",
+                        },
+                    },
+                },
+            },
+            "shape": shape,
+            "number_of_stored_values": nnz,
+            "data_types": {
+                "pointers_to_1": _,
+                "indices_1": coords_dtype,
+                "values": value_dtype,
+            },
+            **_kwargs,
+        }:
+            if transpose != list(range(ndim)):
+                raise RuntimeError(format_err_str)
+
+            ptr_arr: np.ndarray = np.from_dlpack(arrs[0])
+            start, end = ptr_arr
+            if copy is False and not (start == 0 or end == nnz):
+                raise RuntimeError(format_err_str)
+
+            coord_arr: np.ndarray = np.from_dlpack(arrs[1])
+            value_arr: np.ndarray = np.from_dlpack(arrs[2])
+
+            if str(coord_arr.dtype) != coords_dtype:
+                raise BufferError(
+                    invalid_dtype_str.format(
+                        dtype=str(coord_arr.dtype),
+                        expected=coords_dtype,
+                    )
+                )
+
+            if value_dtype.startswith("complex[float") and value_dtype.endswith("]"):
+                complex_bits = 2 * int(value_arr[len("complex[float") : -len("]")])
+                value_dtype: str = f"complex{complex_bits}"
+
+            if str(value_arr.dtype) != value_dtype:
+                raise BufferError(
+                    invalid_dtype_str.format(
+                        dtype=str(coord_arr.dtype),
+                        expected=coords_dtype,
+                    )
+                )
+
+            return COO(
+                coord_arr[:, start:end],
+                value_arr,
+                shape=shape,
+                has_duplicates=False,
+                sorted=True,
+                prune=False,
+                idx_dtype=coord_arr.dtype,
+            )
+        case _:
+            raise RuntimeError(format_err_str)

sparse/numba_backend/tests/test_io.py

@@ -28,3 +28,11 @@ def test_load_wrong_format_exception(tmp_path):
     np.savez(filename, x)
     with pytest.raises(RuntimeError):
         load_npz(filename)
+
+
+@pytest.mark.parametrize("format", ["coo", "csr", "csc"])
+def test_round_trip_binsparse(format: str) -> None:
+    x = sparse.random((20, 30), density=0.25, format=format)
+    y = sparse.from_binsparse(x)
+
+    assert_eq(x, y)

sparse/numba_backend/tests/test_namespace.py

@@ -67,6 +67,7 @@ def test_namespace():
         "float64",
         "floor",
         "floor_divide",
+        "from_binsparse",
         "full",
         "full_like",
         "greater",