first attempt to support awkward arrays

anndata/_core/aligned_mapping.py

@@ -4,15 +4,17 @@
 from typing import Iterator, Mapping, Sequence  # ABCs
 from typing import Tuple, List, Dict  # Generic base types
 
+
 import numpy as np
 import pandas as pd
 from scipy.sparse import spmatrix
 
-from ..utils import deprecated, ensure_df_homogeneous
+from ..utils import deprecated, ensure_df_homogeneous, dim_len
 from . import raw, anndata
 from .views import as_view
 from .access import ElementRef
 from .index import _subset
+from anndata.compat import AwkArray
 
 
 OneDIdx = Union[Sequence[int], Sequence[bool], slice]
@@ -47,7 +49,9 @@ def _ipython_key_completions_(self) -> List[str]:
     def _validate_value(self, val: V, key: str) -> V:
         """Raises an error if value is invalid"""
         for i, axis in enumerate(self.axes):
-            if self.parent.shape[axis] != val.shape[i]:
+            if (not isinstance(val, AwkArray)) and (  # don't validate awk arrays
+                self.parent.shape[axis] != dim_len(val, i)
+            ):
                 right_shape = tuple(self.parent.shape[a] for a in self.axes)
                 raise ValueError(
                     f"Value passed for key {key!r} is of incorrect shape. "
@@ -84,7 +88,10 @@ def parent(self) -> Union["anndata.AnnData", "raw.Raw"]:
     def copy(self):
         d = self._actual_class(self.parent, self._axis)
         for k, v in self.items():
-            d[k] = v.copy()
+            if isinstance(v, AwkArray):
+                d[k] = v
+            else:
+                d[k] = v.copy()
         return d
 
     def _view(self, parent: "anndata.AnnData", subset_idx: I):

anndata/_core/anndata.py

@@ -44,7 +44,7 @@
 )
 from .sparse_dataset import SparseDataset
 from .. import utils
-from ..utils import convert_to_dict, ensure_df_homogeneous
+from ..utils import convert_to_dict, ensure_df_homogeneous, dim_len
 from ..logging import anndata_logger as logger
 from ..compat import (
     ZarrArray,
@@ -1846,7 +1846,7 @@ def _check_dimensions(self, key=None):
         if "obsm" in key:
             obsm = self._obsm
             if (
-                not all([o.shape[0] == self._n_obs for o in obsm.values()])
+                not all([dim_len(o, 0) == self._n_obs for o in obsm.values()])
                 and len(obsm.dim_names) != self._n_obs
             ):
                 raise ValueError(
@@ -1856,7 +1856,7 @@ def _check_dimensions(self, key=None):
         if "varm" in key:
             varm = self._varm
             if (
-                not all([v.shape[0] == self._n_vars for v in varm.values()])
+                not all([dim_len(v, 0) == self._n_vars for v in varm.values()])
                 and len(varm.dim_names) != self._n_vars
             ):
                 raise ValueError(

anndata/_core/index.py

@@ -7,7 +7,7 @@
 import numpy as np
 import pandas as pd
 from scipy.sparse import spmatrix, issparse
-
+from anndata.compat import AwkArray
 
 Index1D = Union[slice, int, str, np.int64, np.ndarray]
 Index = Union[Index1D, Tuple[Index1D, Index1D], spmatrix]
@@ -140,6 +140,13 @@ def _subset_df(df: pd.DataFrame, subset_idx: Index):
     return df.iloc[subset_idx]
 
 
+@_subset.register(AwkArray)
+def _subset_awkarray(a: AwkArray, subset_idx: Index):
+    if all(isinstance(x, cabc.Iterable) for x in subset_idx):
+        subset_idx = np.ix_(*subset_idx)
+    return a[subset_idx]
+
+
 # Registration for SparseDataset occurs in sparse_dataset.py
 @_subset.register(h5py.Dataset)
 def _subset_dataset(d, subset_idx):

anndata/_core/merge.py

@@ -16,8 +16,8 @@
 from scipy.sparse.base import spmatrix
 
 from .anndata import AnnData
-from ..compat import Literal
-from ..utils import asarray
+from ..compat import Literal, AwkArray
+from ..utils import asarray, dim_len
 
 T = TypeVar("T")
 
@@ -122,6 +122,16 @@ def equal_sparse(a, b) -> bool:
         return False
 
 
+@equal.register(AwkArray)
+def equal_awkward(a, b) -> bool:
+    import awkward as ak
+
+    if dim_len(a, 0) == dim_len(b, 0):
+        return ak.all(a == b)
+    else:
+        return False
+
+
 def as_sparse(x):
     if not isinstance(x, sparse.spmatrix):
         return sparse.csr_matrix(x)
@@ -282,12 +292,14 @@ def apply(self, el, *, axis, fill_value=None):
 
         Missing values are to be replaced with `fill_value`.
         """
-        if self.no_change and (el.shape[axis] == len(self.old_idx)):
+        if self.no_change and (dim_len(el, axis) == len(self.old_idx)):
             return el
         if isinstance(el, pd.DataFrame):
             return self._apply_to_df(el, axis=axis, fill_value=fill_value)
         elif isinstance(el, sparse.spmatrix):
             return self._apply_to_sparse(el, axis=axis, fill_value=fill_value)
+        elif isinstance(el, AwkArray):
+            return self._apply_to_awkward(el, axis=axis, fill_value=fill_value)
         else:
             return self._apply_to_array(el, axis=axis, fill_value=fill_value)
 
@@ -365,6 +377,9 @@ def _apply_to_sparse(self, el: spmatrix, *, axis, fill_value=None) -> spmatrix:
 
         return out
 
+    def _apply_to_awkward(self, el: AwkArray, *, axis, fill_value=None):
+        return el
+
 
 def merge_indices(
     inds: Iterable[pd.Index], join: Literal["inner", "outer"]
@@ -438,6 +453,10 @@ def concat_arrays(arrays, reindexers, axis=0, index=None, fill_value=None):
             ],
             format="csr",
         )
+    elif any(isinstance(a, AwkArray) for a in arrays):
+        import awkward as ak
+
+        return ak.concatenate([f(a) for f, a in zip(reindexers, arrays)])
     else:
         return np.concatenate(
             [
@@ -450,7 +469,6 @@ def concat_arrays(arrays, reindexers, axis=0, index=None, fill_value=None):
 
 def inner_concat_aligned_mapping(mappings, reindexers=None, index=None, axis=0):
     result = {}
-
     for k in intersect_keys(mappings):
         els = [m[k] for m in mappings]
         if reindexers is None:
@@ -474,6 +492,8 @@ def gen_inner_reindexers(els, new_index, axis: Literal[0, 1] = 0):
             lambda x, y: x.intersection(y), (df_indices(el) for el in els)
         )
         reindexers = [Reindexer(df_indices(el), common_ind) for el in els]
+    elif all(isinstance(el, AwkArray) for el in els if not_missing(el)):
+        reindexers = [gen_reindexer(pd.RangeIndex(0), pd.RangeIndex(0)) for _ in els]
     else:
         min_ind = min(el.shape[alt_axis] for el in els)
         reindexers = [
@@ -491,6 +511,8 @@ def gen_outer_reindexers(els, shapes, new_index: pd.Index, *, axis=0):
             else (lambda x: pd.DataFrame(index=range(shape)))
             for el, shape in zip(els, shapes)
         ]
+    elif all(isinstance(el, AwkArray) for el in els if not_missing(el)):
+        reindexers = [gen_reindexer(pd.RangeIndex(0), pd.RangeIndex(0)) for _ in els]
     else:
         # if fill_value is None:
         # fill_value = default_fill_value(els)

anndata/_core/views.py

@@ -10,7 +10,7 @@
 
 from .access import ElementRef
 from ..logging import anndata_logger as logger
-from ..compat import ZappyArray
+from ..compat import ZappyArray, AwkArray
 
 
 class _SetItemMixin:
@@ -112,6 +112,14 @@ def drop(self, *args, inplace: bool = False, **kw):
             df.drop(*args, inplace=True, **kw)
 
 
+class AwkwardArrayView(_ViewMixin, AwkArray):
+    def copy(self, order: str = "C") -> np.ndarray:
+        import awkward as ak
+
+        # we want a copy of an akward array
+        return ak.copy(self)
+
+
 @singledispatch
 def as_view(obj, view_args):
     raise NotImplementedError(f"No view type has been registered for {type(obj)}")
@@ -149,6 +157,11 @@ def as_view_zappy(z, view_args):
     return z
 
 
+@as_view.register(AwkArray)
+def as_view_awkarray(array, view_args):
+    return AwkwardArrayView(array, view_args=view_args)
+
+
 def _resolve_idxs(old, new, adata):
     t = tuple(_resolve_idx(old[i], new[i], adata.shape[i]) for i in (0, 1))
     return t

anndata/compat/__init__.py

@@ -26,6 +26,16 @@ def __repr__():
             return "mock zarr.core.Array"
 
 
+try:
+    from awkward import Array as AwkArray
+except ImportError:
+
+    class AwkArray:
+        @staticmethod
+        def __repr__():
+            return "mock awkward.highlevel.Array"
+
+
 try:
     from zappy.base import ZappyArray
 except ImportError:

anndata/tests/helpers.py

@@ -17,6 +17,8 @@
 from anndata._core.aligned_mapping import AlignedMapping
 from anndata.utils import asarray
 
+from anndata.compat import AwkArray
+
 
 def gen_vstr_recarray(m, n, dtype=None):
     size = m * n
@@ -46,6 +48,17 @@ def gen_typed_df(n, index=None):
     )
 
 
+def gen_awkward(n, index=None):
+    rng = np.random.default_rng(42)
+    arr = AwkArray(
+        [
+            rng.standard_normal(size=(rng.integers(1, 10), rng.integers(1, 10)))
+            for _ in range(n)
+        ]
+    )
+    return arr
+
+
 def gen_typed_df_t2_size(m, n, index=None, columns=None) -> pd.DataFrame:
     s = 0
     df = pd.DataFrame()
@@ -68,8 +81,18 @@ def gen_adata(
     X_dtype=np.float32,
     # obs_dtypes,
     # var_dtypes,
-    obsm_types: "Collection[Type]" = (sparse.csr_matrix, np.ndarray, pd.DataFrame),
-    varm_types: "Collection[Type]" = (sparse.csr_matrix, np.ndarray, pd.DataFrame),
+    obsm_types: "Collection[Type]" = (
+        sparse.csr_matrix,
+        np.ndarray,
+        pd.DataFrame,
+        AwkArray,
+    ),
+    varm_types: "Collection[Type]" = (
+        sparse.csr_matrix,
+        np.ndarray,
+        pd.DataFrame,
+        AwkArray,
+    ),
     layers_types: "Collection[Type]" = (sparse.csr_matrix, np.ndarray, pd.DataFrame),
 ) -> AnnData:
     """\
@@ -112,12 +135,14 @@ def gen_adata(
         array=np.random.random((M, 50)),
         sparse=sparse.random(M, 100, format="csr"),
         df=gen_typed_df(M, obs_names),
+        awk=gen_awkward(M),
     )
     obsm = {k: v for k, v in obsm.items() if type(v) in obsm_types}
     varm = dict(
         array=np.random.random((N, 50)),
         sparse=sparse.random(N, 100, format="csr"),
         df=gen_typed_df(N, var_names),
+        awk=gen_awkward(N),
     )
     varm = {k: v for k, v in varm.items() if type(v) in varm_types}
     layers = dict(
@@ -328,6 +353,13 @@ def are_equal_dataframe(a, b, exact=False, elem_name=None):
     )
 
 
+@assert_equal.register(AwkArray)
+def assert_equal_awkarray(a, b, exact=False, elem_name=None):
+    import awkward as ak
+
+    assert ak.all(a == b)
+
+
 @assert_equal.register(Mapping)
 def assert_equal_mapping(a, b, exact=False, elem_name=None):
     assert set(a.keys()) == set(b.keys()), format_msg(elem_name)

anndata/tests/test_concatenate.py

@@ -4,6 +4,7 @@
 from functools import partial, singledispatch
 import warnings
 
+
 import numpy as np
 from numpy import ma
 import pandas as pd
@@ -18,7 +19,8 @@
 from anndata._core import merge
 from anndata.tests import helpers
 from anndata.tests.helpers import assert_equal, gen_adata
-from anndata.utils import asarray
+from anndata.utils import asarray, dim_len
+from anndata.compat import AwkArray
 
 
 @singledispatch
@@ -430,19 +432,21 @@ def get_obs_els(adata):
 
     adata1 = gen_adata((10, 10))
     adata1.obsm = {
-        k: v for k, v in adata1.obsm.items() if not isinstance(v, pd.DataFrame)
+        k: v
+        for k, v in adata1.obsm.items()
+        if not isinstance(v, (pd.DataFrame, ak.Array))
     }
     adata2 = gen_adata((10, 5))
     adata2.obsm = {
         k: v[:, : v.shape[1] // 2]
         for k, v in adata2.obsm.items()
-        if not isinstance(v, pd.DataFrame)
+        if not isinstance(v, (pd.DataFrame, ak.Array))
     }
     adata3 = gen_adata((7, 3))
     adata3.obsm = {
         k: v[:, : v.shape[1] // 3]
         for k, v in adata3.obsm.items()
-        if not isinstance(v, pd.DataFrame)
+        if not isinstance(v, (pd.DataFrame, ak.Array))
     }
     joined = adata1.concatenate([adata2, adata3], join="outer", fill_value=fill_val)