Support nD-dim sparse arrays (#130)

* Replace scipy.sparse with pydata/sparse

* Enable tests >2D sparse arrays

* Merge base abstract TileDBTensorGenerator with TileDBNumpyGenerator

Author: gsakkis

Diff for: setup.cfg

@@ -34,7 +34,7 @@ zip_safe = False
 packages = find_namespace:
 python_requires = >=3.7
 test_suite = tests
-install_requires = tiledb == 0.12.3
+install_requires = sparse; tiledb == 0.12.3
 setup_requires = setuptools_scm <= 6.0.0
 
 [options.extras_require]

Diff for: tests/readers/utils.py

@@ -22,13 +22,6 @@ def parametrize_for_dataset(
     buffer_bytes=(1024, None),
     shuffle_buffer_size=(0, 16),
 ):
-    def is_valid_combination(t):
-        x_sparse_, y_sparse_, x_shape_, y_shape_, *_ = t
-        # sparse not supported with multi-dimensional arrays
-        if x_sparse_ and len(x_shape_) > 1 or y_sparse_ and len(y_shape_) > 1:
-            return False
-        return True
-
     argnames = [
         "x_sparse",
         "y_sparse",
@@ -40,19 +33,16 @@ def is_valid_combination(t):
         "batch_size",
         "shuffle_buffer_size",
     ]
-    argvalues = filter(
-        is_valid_combination,
-        it.product(
-            x_sparse,
-            y_sparse,
-            x_shape,
-            y_shape,
-            num_attrs,
-            pass_attrs,
-            buffer_bytes,
-            batch_size,
-            shuffle_buffer_size,
-        ),
+    argvalues = it.product(
+        x_sparse,
+        y_sparse,
+        x_shape,
+        y_shape,
+        num_attrs,
+        pass_attrs,
+        buffer_bytes,
+        batch_size,
+        shuffle_buffer_size,
     )
     return pytest.mark.parametrize(argnames, argvalues)
 

Diff for: tiledb/ml/readers/_tensor_gen.py

@@ -2,7 +2,7 @@
 from typing import Generic, Iterator, Sequence, Type, TypeVar, Union
 
 import numpy as np
-import scipy.sparse as sp
+import sparse
 
 import tiledb
 
@@ -11,7 +11,7 @@
 Tensor = TypeVar("Tensor")
 
 
-class TileDBTensorGenerator(ABC, Generic[Tensor]):
+class TileDBNumpyGenerator:
     """Base class for generating tensors read from a TileDB array."""
 
     def __init__(self, array: tiledb.Array, attrs: Sequence[str]) -> None:
@@ -21,83 +21,50 @@ def __init__(self, array: tiledb.Array, attrs: Sequence[str]) -> None:
         """
         self._query = array.query(attrs=attrs)
 
-    @abstractmethod
     def read_buffer(self, array_slice: slice) -> None:
         """
         Read an array slice and save it as the current buffer.
 
         :param array_slice: Requested array slice.
         """
+        self._buf_arrays = tuple(self._query[array_slice].values())
 
-    @abstractmethod
-    def iter_tensors(self, buffer_slice: slice) -> Iterator[Tensor]:
+    def iter_tensors(self, buffer_slice: slice) -> Iterator[np.ndarray]:
         """
         Return an iterator of tensors for the given slice, one tensor per attribute
 
         Must be called after `read_buffer`.
 
         :param buffer_slice: Slice of the current buffer to convert to tensors.
         """
+        return (buf_array[buffer_slice] for buf_array in self._buf_arrays)
 
 
-class TileDBNumpyGenerator(TileDBTensorGenerator[np.ndarray]):
-    def read_buffer(self, array_slice: slice) -> None:
-        self._buf_arrays = tuple(self._query[array_slice].values())
-
-    def iter_tensors(self, buffer_slice: slice) -> Iterator[np.ndarray]:
-        for buf_array in self._buf_arrays:
-            yield buf_array[buffer_slice]
-
-
-class TileDBSparseTensorGenerator(TileDBTensorGenerator[Tensor]):
+class TileDBSparseTensorGenerator(TileDBNumpyGenerator, ABC, Generic[Tensor]):
     def __init__(self, array: tiledb.Array, attrs: Sequence[str]) -> None:
-        schema = array.schema
-        if schema.ndim != 2:
-            raise NotImplementedError("Only 2D sparse tensors are currently supported")
-        self._row_dim = schema.domain.dim(0).name
-        self._col_dim = schema.domain.dim(1).name
-        self._row_shape = schema.shape[1:]
-        self._attr_dtypes = tuple(schema.attr(attr).dtype for attr in attrs)
+        self._dims = tuple(array.domain.dim(i).name for i in range(array.ndim))
+        self._row_shape = array.shape[1:]
         super().__init__(array, attrs)
 
     def read_buffer(self, array_slice: slice) -> None:
         buffer = self._query[array_slice]
-        # COO to CSR transformation for batching and row slicing
-        row = buffer.pop(self._row_dim)
-        col = buffer.pop(self._col_dim)
-        # Normalize indices: We want the coords indices to be in the [0, array_slice size]
-        # range. If we do not normalize the sparse tensor is being created but with a
-        # dimension [0, max(coord_index)], which is overkill
+        coords = [buffer.pop(dim) for dim in self._dims]
+        # normalize the first coordinate dimension to start at start_offset
         start_offset = array_slice.start
-        stop_offset = array_slice.stop
-        shape = (stop_offset - start_offset, *self._row_shape)
-        self._buf_csrs = tuple(
-            sp.csr_matrix((data, (row - start_offset, col)), shape=shape)
-            for data in buffer.values()
+        if start_offset:
+            coords[0] -= start_offset
+        shape = (array_slice.stop - start_offset, *self._row_shape)
+        self._buf_arrays = tuple(
+            sparse.COO(coords, data, shape) for data in buffer.values()
         )
 
     def iter_tensors(self, buffer_slice: slice) -> Iterator[Tensor]:
-        for buf_csr, dtype in zip(self._buf_csrs, self._attr_dtypes):
-            batch_csr = buf_csr[buffer_slice]
-            batch_coo = batch_csr.tocoo()
-            data = batch_coo.data
-            coords = np.stack((batch_coo.row, batch_coo.col), axis=-1)
-            dense_shape = (batch_csr.shape[0], *self._row_shape)
-            yield self._tensor_from_coo(data, coords, dense_shape, dtype)
+        return map(self._tensor_from_coo, super().iter_tensors(buffer_slice))
 
     @staticmethod
     @abstractmethod
-    def _tensor_from_coo(
-        data: np.ndarray,
-        coords: np.ndarray,
-        dense_shape: Sequence[int],
-        dtype: np.dtype,
-    ) -> Tensor:
-        """Convert a scipy.sparse.coo_matrix to a Tensor"""
-
-
-DT = TypeVar("DT")
-ST = TypeVar("ST")
+    def _tensor_from_coo(coo: sparse.COO) -> Tensor:
+        """Convert a sparse.COO to a Tensor"""
 
 
 def tensor_generator(
@@ -107,11 +74,10 @@ def tensor_generator(
     y_buffer_size: int,
     x_attrs: Sequence[str],
     y_attrs: Sequence[str],
+    sparse_generator_cls: Type[TileDBSparseTensorGenerator[Tensor]],
     start_offset: int = 0,
     stop_offset: int = 0,
-    dense_generator_cls: Type[TileDBTensorGenerator[DT]] = TileDBNumpyGenerator,
-    sparse_generator_cls: Type[TileDBTensorGenerator[ST]] = TileDBSparseTensorGenerator,
-) -> Iterator[Sequence[Union[DT, ST]]]:
+) -> Iterator[Sequence[Union[np.ndarray, Tensor]]]:
     """
     Generator for batches of tensors.
 
@@ -126,18 +92,17 @@ def tensor_generator(
     :param y_attrs: Attribute names of y_array.
     :param start_offset: Start row offset; defaults to 0.
     :param stop_offset: Stop row offset; defaults to number of rows.
-    :param dense_generator_cls: Dense tensor generator type.
     :param sparse_generator_cls: Sparse tensor generator type.
     """
-    x_gen: Union[TileDBTensorGenerator[DT], TileDBTensorGenerator[ST]] = (
+    x_gen: Union[TileDBNumpyGenerator, TileDBSparseTensorGenerator[Tensor]] = (
         sparse_generator_cls(x_array, x_attrs)
         if x_array.schema.sparse
-        else dense_generator_cls(x_array, x_attrs)
+        else TileDBNumpyGenerator(x_array, x_attrs)
     )
-    y_gen: Union[TileDBTensorGenerator[DT], TileDBTensorGenerator[ST]] = (
+    y_gen: Union[TileDBNumpyGenerator, TileDBSparseTensorGenerator[Tensor]] = (
         sparse_generator_cls(y_array, y_attrs)
         if y_array.schema.sparse
-        else dense_generator_cls(y_array, y_attrs)
+        else TileDBNumpyGenerator(y_array, y_attrs)
     )
     if not stop_offset:
         stop_offset = x_array.shape[0]

Diff for: tiledb/ml/readers/pytorch.py

@@ -6,6 +6,7 @@
 from typing import Any, Callable, Iterable, Iterator, Optional, Sequence, TypeVar
 
 import numpy as np
+import sparse
 import torch
 
 import tiledb
@@ -168,12 +169,5 @@ def iter_shuffled(iterable: Iterable[T], buffer_size: int) -> Iterator[T]:
 
 class PyTorchSparseTensorGenerator(TileDBSparseTensorGenerator[torch.Tensor]):
     @staticmethod
-    def _tensor_from_coo(
-        data: np.ndarray,
-        coords: np.ndarray,
-        dense_shape: Sequence[int],
-        dtype: np.dtype,
-    ) -> torch.Tensor:
-        return torch.sparse_coo_tensor(
-            torch.tensor(coords).t(), data, dense_shape, requires_grad=False
-        )
+    def _tensor_from_coo(coo: sparse.COO) -> torch.Tensor:
+        return torch.sparse_coo_tensor(coo.coords, coo.data, coo.shape)

Diff for: tiledb/ml/readers/tensorflow.py

@@ -3,7 +3,7 @@
 from functools import partial
 from typing import Iterator, Optional, Sequence, Union
 
-import numpy as np
+import sparse
 import tensorflow as tf
 
 import tiledb
@@ -79,14 +79,5 @@ def _iter_tensor_specs(
 
 class TensorflowSparseTensorGenerator(TileDBSparseTensorGenerator[tf.SparseTensor]):
     @staticmethod
-    def _tensor_from_coo(
-        data: np.ndarray,
-        coords: np.ndarray,
-        dense_shape: Sequence[int],
-        dtype: np.dtype,
-    ) -> tf.SparseTensor:
-        return tf.SparseTensor(
-            indices=tf.constant(coords, dtype=tf.int64),
-            values=tf.constant(data, dtype=dtype),
-            dense_shape=dense_shape,
-        )
+    def _tensor_from_coo(coo: sparse.COO) -> tf.SparseTensor:
+        return tf.SparseTensor(coo.coords.T, coo.data, coo.shape)