Support indexing by pyarrow.Array

Author: johnkerl

apis/python/src/tiledbsoma/types.py

@@ -26,7 +26,7 @@
     pa.SparseCSCMatrix,
 ]
 
-SOMADenseCoordinates = Union[int, slice]
+SOMADenseCoordinates = Union[int, slice, pa.Array]
 SOMADenseNdCoordinates = Tuple[SOMADenseCoordinates, ...]
 
 SOMASparseCoordinates = Union[int, slice, Tuple[int, ...], List[int], pa.IntegerArray]

apis/python/src/tiledbsoma/util.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 import scipy.sparse as sp
 
 T = TypeVar("T", np.ndarray, pd.Series, pd.DataFrame, sp.spmatrix)
@@ -93,7 +94,7 @@ def uri_joinpath(base: str, path: str) -> str:
     return urllib.parse.urlunparse(parts)
 
 
-def ids_to_list(ids: Union[slice, List[int]]) -> List[int]:
+def ids_to_list(ids: Union[slice, List[int], pa.Array]) -> List[int]:
     """
     For the interface between ``SOMADataFrame::read`` et al. (Python) and ``SOMAReader`` (C++): the
     ``ids`` argument to the former can be slice or list; the argument to
@@ -112,4 +113,6 @@ def ids_to_list(ids: Union[slice, List[int]]) -> List[int]:
                 step = -1
         stop = ids.stop + step
         return list(range(ids.start, stop, step))
+    if isinstance(pa.Array):
+        return [id.as_py() for id in ids]
     raise Exception(f"expected ids as slice or list; got {type(ids)}")

apis/python/tests/test_util.py

@@ -1,3 +1,4 @@
+import pyarrow as pa
 import pytest
 
 from tiledbsoma.util import ids_to_list, uri_joinpath
@@ -35,8 +36,11 @@ def test_uri_joinpath_tiledb():
 
 
 def test_ids_to_list():
+    assert ids_to_list([]) == []
     assert ids_to_list([]) == []
     assert ids_to_list([1, 2, 3]) == [1, 2, 3]
+    assert ids_to_list(pa.array([])) == []
+    assert ids_to_list(pa.array([1, 2, 3])) == [1, 2, 3]
     assert ids_to_list(slice(None)) == []
     assert ids_to_list(slice(1, 5)) == [1, 2, 3, 4, 5]
     assert ids_to_list(slice(5, 1)) == [5, 4, 3, 2, 1]