Extending CSR Accumulator to reuse query context instead of creating …

…its own

python-spec/src/somacore/query/_fast_csr.py

@@ -8,7 +8,7 @@
 import numpy.typing as npt
 import pyarrow as pa
 from scipy import sparse
-
+from typing import Optional, Any
 from .. import data as scd
 from . import _eager_iter
 from . import types
@@ -19,6 +19,7 @@ def read_csr(
     obs_joinids: pa.Array,
     var_joinids: pa.Array,
     index_factory: types.IndexFactory,
+    context: Optional[Any],
 ) -> "AccumulatedCSR":
     if not isinstance(matrix, scd.SparseNDArray) or matrix.ndim != 2:
         raise TypeError("Can only read from a 2D SparseNDArray")
@@ -30,6 +31,7 @@ def read_csr(
             var_joinids=var_joinids,
             pool=pool,
             index_factory=index_factory,
+            context=context,
         )
         for tbl in _eager_iter.EagerIterator(
             matrix.read((obs_joinids, var_joinids)).tables(),
@@ -86,14 +88,15 @@ def __init__(
         var_joinids: npt.NDArray[np.int64],
         pool: futures.Executor,
         index_factory: types.IndexFactory,
+        context: Optional[Any],
     ):
         self.obs_joinids = obs_joinids
         self.var_joinids = var_joinids
         self.pool = pool
 
         self.shape: Tuple[int, int] = (len(self.obs_joinids), len(self.var_joinids))
-        self.obs_indexer = index_factory(self.obs_joinids)
-        self.var_indexer = index_factory(self.var_joinids)
+        self.obs_indexer = index_factory(self.obs_joinids, context)
+        self.var_indexer = index_factory(self.var_joinids, context)
         self.row_length: npt.NDArray[np.int64] = np.zeros(
             (self.shape[0],), dtype=_select_dtype(self.shape[1])
         )

python-spec/src/somacore/query/query.py

@@ -346,6 +346,7 @@ def _read(
                 self.obs_joinids(),
                 self.var_joinids(),
                 index_factory=self._index_factory,
+                context=self.experiment.context,
             ).to_scipy()
             for _xname in all_x_arrays
         }

python-spec/src/somacore/query/types.py

@@ -1,9 +1,9 @@
 """Common types used across SOMA query modules."""
 
-from typing import Any, Callable
-
 import numpy as np
 import numpy.typing as npt
+
+from typing import Optional, Any, Callable
 from typing_extensions import Protocol
 
 
@@ -26,7 +26,7 @@ def get_indexer(
         """Something compatible with Pandas' Index.get_indexer method."""
 
 
-IndexFactory = Callable[[npt.NDArray[np.int64]], "IndexLike"]
+IndexFactory = Callable[[npt.NDArray[np.int64], Optional[Any]], "IndexLike"]
 """Function that builds an index over the given NDArray.
 
 This interface is implemented by the callable ``pandas.Index``.