[python] Fix context-type for tiledbsoma_build_index (#2224)

* Fix context-type for `tiledbsoma_build_index` * more * more testing
single-cell-data · Mar 6, 2024 · 92593fd · 92593fd
1 parent 409a487
commit 92593fd
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 9 deletions.
diff --git a/apis/python/src/tiledbsoma/_experiment.py b/apis/python/src/tiledbsoma/_experiment.py
@@ -95,6 +95,7 @@ def axis_query(  # type: ignore
             obs_query=obs_query or query.AxisQuery(),
             var_query=var_query or query.AxisQuery(),
             index_factory=functools.partial(
-                tiledbsoma_build_index, context=self.context.native_context
+                tiledbsoma_build_index,
+                context=self.context,
             ),
         )
diff --git a/apis/python/src/tiledbsoma/_index_util.py b/apis/python/src/tiledbsoma/_index_util.py
@@ -36,6 +36,7 @@ def tiledbsoma_build_index(
     Lifecycle:
         Experimental.
     """
-    reindexer = clib.IntIndexer(context)
+    native_context = None if context is None else context.native_context
+    reindexer = clib.IntIndexer(native_context)
     reindexer.map_locations(keys)
     return reindexer  # type: ignore[no-any-return]
diff --git a/apis/python/src/tiledbsoma/_read_iters.py b/apis/python/src/tiledbsoma/_read_iters.py
@@ -136,7 +136,8 @@ def __init__(
         assert context is not None
         self.minor_axes_indexer = {
             d: tiledbsoma_build_index(
-                self.joinids[d].to_numpy(), context=context.native_context
+                self.joinids[d].to_numpy(),
+                context=context,
             )
             for d in (self.axes_to_reindex - set((self.major_axis,)))
         }
@@ -256,7 +257,8 @@ def _reindexed_table_reader(
                     if d == self.major_axis:
                         assert self.context is not None
                         col = tiledbsoma_build_index(
-                            coords[self.major_axis], context=self.context.native_context
+                            coords[self.major_axis],
+                            context=self.context,
                         ).get_indexer(
                             col.to_numpy(),
                         )

diff --git a/apis/python/tests/test_indexer.py b/apis/python/tests/test_indexer.py
@@ -23,13 +23,17 @@ def test_duplicate_key_indexer_error(
 ):
     context = _validate_soma_tiledb_context(SOMATileDBContext())
     with pytest.raises(RuntimeError, match="There are duplicate keys."):
-        tiledbsoma_build_index(keys, context=context.native_context)
+        tiledbsoma_build_index(keys, context=context)
 
     pd_index = pd.Index(keys)
     with pytest.raises(pd.errors.InvalidIndexError):
         pd_index.get_indexer(lookups)
 
 
+@pytest.mark.parametrize(
+    "contextual",
+    [True, False],
+)
 @pytest.mark.parametrize(
     "keys, lookups",
     [
@@ -88,13 +92,16 @@ def test_duplicate_key_indexer_error(
         ),
     ],
 )
-def test_indexer(keys: np.array, lookups: np.array):
-    context = _validate_soma_tiledb_context(SOMATileDBContext())
+def test_indexer(contextual: bool, keys: np.array, lookups: np.array):
+    if contextual:
+        context = _validate_soma_tiledb_context(SOMATileDBContext())
+    else:
+        context = None
     all_results = []
     num_threads = 10
 
     def target():
-        indexer = tiledbsoma_build_index(keys, context=context.native_context)
+        indexer = tiledbsoma_build_index(keys, context=context)
         results = indexer.get_indexer(lookups)
         all_results.append(results)
 

diff --git a/apis/python/tests/test_reindexer_api.py b/apis/python/tests/test_reindexer_api.py
@@ -17,6 +17,6 @@ def test_reindexer_api_context():
     keys = np.arange(3, 10, 2)
     ids = np.arange(3, 10, 2)
     expected = np.array([0, 1, 2, 3])
-    indexer = tiledbsoma_build_index(keys, context=context.native_context)
+    indexer = tiledbsoma_build_index(keys, context=context)
     result = indexer.get_indexer(ids)
     assert np.equal(result.all(), expected.all())