diff --git a/python-spec/src/somacore/query/_fast_csr.py b/python-spec/src/somacore/query/_fast_csr.py index 511547a..9927e27 100644 --- a/python-spec/src/somacore/query/_fast_csr.py +++ b/python-spec/src/somacore/query/_fast_csr.py @@ -82,8 +82,8 @@ class _CSRAccumulator: def __init__( self, - obs_joinids: npt.NDArray[np.int64], - var_joinids: npt.NDArray[np.int64], + obs_joinids: pa.Array, + var_joinids: pa.Array, pool: futures.Executor, index_factory: types.IndexFactory, ): diff --git a/python-spec/src/somacore/query/types.py b/python-spec/src/somacore/query/types.py index ca9d8de..6ffbba7 100644 --- a/python-spec/src/somacore/query/types.py +++ b/python-spec/src/somacore/query/types.py @@ -1,11 +1,14 @@ """Common types used across SOMA query modules.""" -from typing import Any, Callable +from typing import Any, Callable, Union import numpy as np import numpy.typing as npt +import pyarrow as pa from typing_extensions import Protocol +_Array = Union[npt.NDArray[np.int64], pa.Array] + class IndexLike(Protocol): """The basics of what we expect an Index to be. @@ -16,11 +19,11 @@ class IndexLike(Protocol): not as a full specification of the types and behavior of ``get_indexer``. """ - def get_indexer(self, target: npt.NDArray[np.int64]) -> Any: + def get_indexer(self, target: _Array) -> Any: """Something compatible with Pandas' Index.get_indexer method.""" -IndexFactory = Callable[[npt.NDArray[np.int64]], "IndexLike"] +IndexFactory = Callable[[_Array], "IndexLike"] """Function that builds an index over the given NDArray. This interface is implemented by the callable ``pandas.Index``.