Avoid intermediate sparse.COO instances if possible (#172)

* SparseTensorSchema: yield lightweight SparseData instances instead of sparse.COO

Convert to sparse.COO only for >2D sparse arrays, otherwise convert to scipy.sparse.{csr,coo}_matrix

Author: gsakkis

tiledb/ml/readers/_tensor_schema.py

@@ -5,7 +5,7 @@
 from collections import Counter
 from dataclasses import dataclass
 from math import ceil
-from operator import itemgetter, methodcaller
+from operator import itemgetter
 from typing import (
     Any,
     Callable,
@@ -14,6 +14,7 @@
     Iterable,
     Sequence,
     Tuple,
+    Type,
     TypeVar,
     Union,
     cast,
@@ -234,9 +235,22 @@ def max_partition_weight(self) -> int:
         return max(1, int(rows_per_slice * num_slices))
 
 
-class SparseTensorSchema(TensorSchema[sparse.COO]):
+@dataclass(frozen=True)
+class SparseData:
+    coords: np.ndarray
+    data: np.ndarray
+    shape: Sequence[int]
+
+    def to_sparse_array(self) -> Union[scipy.sparse.csr_matrix, sparse.COO]:
+        if len(self.shape) == 2:
+            return scipy.sparse.csr_matrix((self.data, self.coords), self.shape)
+        else:
+            return sparse.COO(self.coords, self.data, self.shape)
+
+
+class SparseTensorSchema(TensorSchema[SparseData]):
     """
-    TensorSchema for reading sparse TileDB arrays as sparse.COO instances.
+    TensorSchema for reading sparse TileDB arrays as SparseData instances.
     """
 
     def __init__(self, **kwargs: Any):
@@ -263,7 +277,7 @@ def key_range(self) -> InclusiveRange[Any, int]:
 
     def iter_tensors(
         self, key_ranges: Iterable[InclusiveRange[Any, int]]
-    ) -> Union[Iterable[sparse.COO], Iterable[Sequence[sparse.COO]]]:
+    ) -> Union[Iterable[SparseData], Iterable[Sequence[SparseData]]]:
         shape = list(self.shape)
         query = self.query
         get_data = itemgetter(*self._fields)
@@ -285,12 +299,13 @@ def iter_tensors(
                 field_arrays.pop(dim) - dim_start
                 for dim, dim_start in zip(non_key_dims, non_key_dim_starts)
             )
+            coords = np.array(coords)
 
-            # yield either a single tensor or a sequence of tensors, one for each field
+            # yield either a single SparseData or one SparseData per field
             if single_field:
-                yield sparse.COO(coords, data, shape)
+                yield SparseData(coords, data, shape)
             else:
-                yield tuple(sparse.COO(coords, d, shape) for d in data)
+                yield tuple(SparseData(coords, d, shape) for d in data)
 
     @property
     def max_partition_weight(self) -> int:
@@ -322,18 +337,10 @@ def max_partition_weight(self) -> int:
         return max(1, memory_budget // ceil(max(bytes_per_cell)))
 
 
-def SparseCSRTensorSchema(**kwargs: Any) -> TensorSchema[scipy.sparse.csr_matrix]:
-    """
-    Return a TensorSchema for reading sparse 2D TileDB arrays as scipy.sparse.csr_matrix
-    instances.
-    """
-    return MappedTensorSchema(SparseTensorSchema(**kwargs), methodcaller("tocsr"))
-
-
-TensorSchemaFactories = {
+TensorSchemaFactories: Dict[TensorKind, Type[TensorSchema[Any]]] = {
     TensorKind.DENSE: DenseTensorSchema,
     TensorKind.SPARSE_COO: SparseTensorSchema,
-    TensorKind.SPARSE_CSR: SparseCSRTensorSchema,
+    TensorKind.SPARSE_CSR: SparseTensorSchema,
 }
 
 

tiledb/ml/readers/pytorch.py

@@ -3,12 +3,14 @@
 import itertools
 import random
 from dataclasses import dataclass
+from operator import methodcaller
 from typing import (
     Any,
     Callable,
     Dict,
     Iterable,
     Iterator,
+    Mapping,
     Sequence,
     Tuple,
     TypeVar,
@@ -61,7 +63,7 @@ def PyTorchTileDBDataLoader(
     the following arguments: 'shuffle', 'sampler', 'batch_sampler', 'worker_init_fn' and 'collate_fn'.
     """
     schemas = tuple(
-        array_params.to_tensor_schema() for array_params in all_array_params
+        array_params.to_tensor_schema(_transforms) for array_params in all_array_params
     )
     key_range = schemas[0].key_range
     if not all(key_range.equal_values(schema.key_range) for schema in schemas[1:]):
@@ -140,13 +142,20 @@ def _ndarray_collate(arrays: Sequence[np.ndarray]) -> torch.Tensor:
     return torch.from_numpy(np.stack(arrays))
 
 
-def _sparse_coo_collate(arrays: Sequence[sparse.COO]) -> torch.Tensor:
+def _coo_collate(arrays: Sequence[sparse.COO]) -> torch.Tensor:
     """Collate multiple sparse.COO arrays to a torch.Tensor with sparse_coo layout."""
     stacked = sparse.stack(arrays)
     return torch.sparse_coo_tensor(stacked.coords, stacked.data, stacked.shape)
 
 
-def _sparse_csr_collate(arrays: Sequence[scipy.sparse.csr_matrix]) -> torch.Tensor:
+def _csr_to_coo_collate(arrays: Sequence[scipy.sparse.csr_matrix]) -> torch.Tensor:
+    """Collate multiple Scipy CSR matrices to a torch.Tensor with sparse_coo layout."""
+    stacked = scipy.sparse.vstack(arrays).tocoo()
+    coords = np.stack((stacked.row, stacked.col))
+    return torch.sparse_coo_tensor(coords, stacked.data, stacked.shape)
+
+
+def _csr_collate(arrays: Sequence[scipy.sparse.csr_matrix]) -> torch.Tensor:
     """Collate multiple Scipy CSR matrices to a torch.Tensor with sparse_csr layout."""
     stacked = scipy.sparse.vstack(arrays)
     return torch.sparse_csr_tensor(
@@ -157,24 +166,37 @@ def _sparse_csr_collate(arrays: Sequence[scipy.sparse.csr_matrix]) -> torch.Tens
     )
 
 
-_collators = {
-    TensorKind.DENSE: _ndarray_collate,
-    TensorKind.SPARSE_COO: _sparse_coo_collate,
-    TensorKind.SPARSE_CSR: _sparse_csr_collate,
-}
-
-
 def _get_tensor_collator(
     schema: TensorSchema[Tensor],
 ) -> Union[_SingleCollator, _CompositeCollator]:
-    collator = _collators[schema.kind]
+    if schema.kind is TensorKind.DENSE:
+        collator = _ndarray_collate
+    elif schema.kind is TensorKind.SPARSE_COO:
+        if len(schema.shape) != 2:
+            collator = _coo_collate
+        else:
+            collator = _csr_to_coo_collate
+    elif schema.kind is TensorKind.SPARSE_CSR:
+        if len(schema.shape) != 2:
+            raise ValueError("SPARSE_CSR is supported only for 2D tensors")
+        collator = _csr_collate
+    else:
+        assert False, schema.kind
+
     num_fields = schema.num_fields
     if num_fields == 1:
         return collator
     else:
         return _CompositeCollator(*itertools.repeat(collator, num_fields))
 
 
+_transforms: Mapping[TensorKind, Union[Callable[[Any], Any], bool]] = {
+    TensorKind.DENSE: True,
+    TensorKind.SPARSE_COO: methodcaller("to_sparse_array"),
+    TensorKind.SPARSE_CSR: methodcaller("to_sparse_array"),
+}
+
+
 _T = TypeVar("_T")
 
 

tiledb/ml/readers/tensorflow.py

@@ -5,7 +5,7 @@
 import numpy as np
 import tensorflow as tf
 
-from ._tensor_schema import TensorKind, TensorSchema
+from ._tensor_schema import SparseData, TensorKind, TensorSchema
 from .types import ArrayParams
 
 Tensor = Union[np.ndarray, tf.SparseTensor]
@@ -70,10 +70,18 @@ def _get_tensor_specs(
     return specs if len(specs) > 1 else specs[0]
 
 
+def _to_sparse_tensor(sd: SparseData) -> tf.SparseTensor:
+    sa = sd.to_sparse_array()
+    coords = getattr(sa, "coords", None)
+    if coords is None:
+        # sa is a scipy.sparse.csr_matrix
+        coo = sa.tocoo()
+        coords = np.array((coo.row, coo.col))
+    return tf.SparseTensor(coords.T, sa.data, sa.shape)
+
+
 _transforms: Mapping[TensorKind, Union[Callable[[Any], Any], bool]] = {
     TensorKind.DENSE: True,
-    TensorKind.SPARSE_COO: (
-        lambda coo: tf.SparseTensor(coo.coords.T, coo.data, coo.shape)
-    ),
+    TensorKind.SPARSE_COO: _to_sparse_tensor,
     TensorKind.SPARSE_CSR: False,
 }