first attempt to support awkward arrays

anndata/_core/aligned_mapping.py

@@ -3,6 +3,17 @@
 from typing import Union, Optional, Type, ClassVar, TypeVar  # Special types
 from typing import Iterator, Mapping, Sequence  # ABCs
 from typing import Tuple, List, Dict  # Generic base types
+from functools import singledispatch
+
+import warnings
+
+try:
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        import awkward as ak
+except ImportError:
+    ak = None
+
 
 import numpy as np
 import pandas as pd
@@ -47,7 +58,7 @@ 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 self.parent.shape[axis] != dim_len(val, i):  # val.shape[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. "
@@ -349,3 +360,15 @@ def __init__(
 
 PairwiseArraysBase._view_class = PairwiseArraysView
 PairwiseArraysBase._actual_class = PairwiseArrays
+
+
+@singledispatch
+def dim_len(x, dim):
+    return x.shape[dim]
+
+
+@dim_len.register(ak.Array)
+def dim_len_array(x, dim):
+    if dim != 0:
+        raise IndexError()
+    return len(x)

anndata/_core/anndata.py

@@ -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([len(o) == self._n_obs for o in obsm.values()])
                 and len(obsm.dim_names) != self._n_obs
             ):
                 raise ValueError(

anndata/_core/index.py

@@ -8,6 +8,14 @@
 import pandas as pd
 from scipy.sparse import spmatrix, issparse
 
+import warnings
+
+try:
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        import awkward as ak
+except ImportError:
+    ak = None
 
 Index1D = Union[slice, int, str, np.int64, np.ndarray]
 Index = Union[Index1D, Tuple[Index1D, Index1D], spmatrix]
@@ -140,6 +148,13 @@ def _subset_df(df: pd.DataFrame, subset_idx: Index):
     return df.iloc[subset_idx]
 
 
+@_subset.register(ak.Array)
+def _subset_awkarray(a: ak.Array, 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/views.py

@@ -12,6 +12,15 @@
 from ..logging import anndata_logger as logger
 from ..compat import ZappyArray
 
+import warnings
+
+try:
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        import awkward as ak
+except ImportError:
+    ak = None
+
 
 class _SetItemMixin:
     """\
@@ -112,6 +121,12 @@ def drop(self, *args, inplace: bool = False, **kw):
             df.drop(*args, inplace=True, **kw)
 
 
+class AwkwardArrayView(_ViewMixin, ak.Array):
+    def copy(self, order: str = "C") -> np.ndarray:
+        # 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 +164,11 @@ def as_view_zappy(z, view_args):
     return z
 
 
+@as_view.register(ak.Array)
+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