(chore): refactor ZarrArray subset function

scverse · Jul 30, 2023 · 8ce994b · 8ce994b
1 parent 703812c
commit 8ce994b
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/anndata/_core/index.py b/anndata/_core/index.py
@@ -116,15 +116,18 @@ def unpack_index(index: Index) -> Tuple[Index1D, Index1D]:
 
 
 @singledispatch
-def _subset(a: Union[np.ndarray, pd.DataFrame], subset_idx: Index):
+def _subset(a: np.ndarray, subset_idx: Index):
     # Select as combination of indexes, not coordinates
     # Correcting for indexing behaviour of np.ndarray
     if all(isinstance(x, cabc.Iterable) for x in subset_idx):
         subset_idx = np.ix_(*subset_idx)
-    if isinstance(a, ZarrArray):
-        return a.oindex[subset_idx]
     return a[subset_idx]
 
+@_subset.register(ZarrArray)
+def _subset_zarr(a: ZarrArray, subset_idx: Index):
+    if all(isinstance(x, cabc.Iterable) for x in subset_idx):
+        subset_idx = np.ix_(*subset_idx)
+    return a.oindex[subset_idx]
 
 @_subset.register(DaskArray)
 def _subset_dask(a: DaskArray, subset_idx: Index):