first attempt to support awkward arrays

.gitignore

@@ -27,4 +27,5 @@ test.h5ad
 
 # IDEs
 /.idea/
+/.vscode/
 

anndata/__init__.py

@@ -18,7 +18,12 @@
         read_mtx,
         read_zarr,
     )
-    from ._warnings import OldFormatWarning, WriteWarning, ImplicitModificationWarning
+    from ._warnings import (
+        OldFormatWarning,
+        WriteWarning,
+        ImplicitModificationWarning,
+        ExperimentalFeatureWarning,
+    )
 
     # backwards compat / shortcut for default format
     from ._io import read_h5ad as read

anndata/_core/aligned_mapping.py

@@ -1,18 +1,22 @@
 from abc import ABC, abstractmethod
 from collections import abc as cabc
+from copy import copy
 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
+import warnings
 
 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
+from anndata._warnings import ExperimentalFeatureWarning
 
 
 OneDIdx = Union[Sequence[int], Sequence[bool], slice]
@@ -46,15 +50,37 @@ def _ipython_key_completions_(self) -> List[str]:
 
     def _validate_value(self, val: V, key: str) -> V:
         """Raises an error if value is invalid"""
+        if isinstance(val, AwkArray):
+            warnings.warn(
+                "Support for Awkward Arrays is currently experimental. "
+                "Behavior may change in the future. Please report any issues you may encounter!",
+                ExperimentalFeatureWarning,
+                # stacklevel=3,
+            )
+            # Prevent from showing up every time an awkward array is used
+            # You'd think `once` works, but it doesn't at the repl and in notebooks
+            warnings.filterwarnings(
+                "ignore",
+                category=ExperimentalFeatureWarning,
+                message="Support for Awkward Arrays is currently experimental.*",
+            )
         for i, axis in enumerate(self.axes):
-            if self.parent.shape[axis] != val.shape[i]:
+            if 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. "
-                    f"Values of {self.attrname} must match dimensions "
-                    f"{self.axes} of parent. Value had shape {val.shape} while "
-                    f"it should have had {right_shape}."
-                )
+                actual_shape = tuple(dim_len(val, a) for a, _ in enumerate(self.axes))
+                if actual_shape[i] is None and isinstance(val, AwkArray):
+                    raise ValueError(
+                        f"The AwkwardArray is of variable length in dimension {i}.",
+                        f"Try ak.to_regular(array, {i}) before including the array in AnnData",
+                    )
+                else:
+                    raise ValueError(
+                        f"Value passed for key {key!r} is of incorrect shape. "
+                        f"Values of {self.attrname} must match dimensions "
+                        f"{self.axes} of parent. Value had shape {actual_shape} while "
+                        f"it should have had {right_shape}."
+                    )
+
         if not self._allow_df and isinstance(val, pd.DataFrame):
             name = self.attrname.title().rstrip("s")
             val = ensure_df_homogeneous(val, f"{name} {key!r}")
@@ -84,7 +110,11 @@ 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):
+                # Shallow copy since awkward array buffers are immutable
+                d[k] = copy(v)
+            else:
+                d[k] = v.copy()
         return d
 
     def _view(self, parent: "anndata.AnnData", subset_idx: I):

anndata/_core/anndata.py

@@ -45,7 +45,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,
@@ -55,6 +55,7 @@
     _move_adj_mtx,
     _overloaded_uns,
     OverloadedDict,
+    AwkArray,
 )
 
 
@@ -1861,7 +1862,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(
@@ -1871,7 +1872,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 ..compat import DaskArray, Index, Index1D
+from ..compat import AwkArray, DaskArray, Index, Index1D
 
 
 def _normalize_indices(
@@ -145,6 +145,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

@@ -27,8 +27,8 @@
 from scipy.sparse import spmatrix
 
 from .anndata import AnnData
-from ..utils import asarray
-from ..compat import DaskArray
+from ..compat import AwkArray, DaskArray
+from ..utils import asarray, dim_len
 from .index import _subset, make_slice
 
 T = TypeVar("T")
@@ -154,6 +154,16 @@ def equal_sparse(a, b) -> bool:
         return False
 
 
+@equal.register(AwkArray)
+def equal_awkward(a, b) -> bool:
+    from ..compat 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)
@@ -366,12 +376,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)
         elif isinstance(el, DaskArray):
             return self._apply_to_dask_array(el, axis=axis, fill_value=fill_value)
         else:
@@ -468,6 +480,23 @@ 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):
+        import awkward as ak
+
+        if self.no_change:
+            return el
+        elif axis == 1:  # Indexing by field
+            if self.new_idx.isin(self.old_idx).all():  # inner join
+                return el[self.new_idx]
+            else:  # outer join
+                for field in self.new_idx.difference(self.old_idx):
+                    el = ak.with_field(el, None, field)
+                return el
+        else:
+            raise NotImplementedError(
+                "Reindexing along axis 0 is not yet implemented for Awkward Arrays"
+            )
+
 
 def merge_indices(
     inds: Iterable[pd.Index], join: Literal["inner", "outer"]
@@ -534,6 +563,17 @@ def concat_arrays(arrays, reindexers, axis=0, index=None, fill_value=None):
         )
         df.index = index
         return df
+    elif any(isinstance(a, AwkArray) for a in arrays):
+        from ..compat import awkward as ak
+
+        if not all(
+            isinstance(a, AwkArray) or a is MissingVal or 0 in a.shape for a in arrays
+        ):
+            raise NotImplementedError(
+                "Cannot concatenate an AwkwardArray with other array types."
+            )
+
+        return ak.concatenate([f(a) for f, a in zip(reindexers, arrays)], axis=axis)
     elif any(isinstance(a, sparse.spmatrix) for a in arrays):
         sparse_stack = (sparse.vstack, sparse.hstack)[axis]
         return sparse_stack(
@@ -579,6 +619,15 @@ 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 any(isinstance(el, AwkArray) for el in els if not_missing(el)):
+        if not all(isinstance(el, AwkArray) for el in els if not_missing(el)):
+            raise NotImplementedError(
+                "Cannot concatenate an AwkwardArray with other array types."
+            )
+        common_keys = intersect_keys(el.fields for el in els)
+        reindexers = [
+            Reindexer(pd.Index(el.fields), pd.Index(list(common_keys))) for el in els
+        ]
     else:
         min_ind = min(el.shape[alt_axis] for el in els)
         reindexers = [
@@ -596,10 +645,27 @@ 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)
         ]
-    else:
-        # if fill_value is None:
-        # fill_value = default_fill_value(els)
+    elif any(isinstance(el, AwkArray) for el in els if not_missing(el)):
+        import awkward as ak
 
+        if not all(isinstance(el, AwkArray) for el in els if not_missing(el)):
+            raise NotImplementedError(
+                "Cannot concatenate an AwkwardArray with other array types."
+            )
+        # all_keys = union_keys(el.fields for el in els if not_missing(el))
+        reindexers = []
+        for el in els:
+            if not_missing(el):
+                reindexers.append(lambda x: x)
+            else:
+                reindexers.append(
+                    lambda x: ak.pad_none(
+                        ak.Array([]),
+                        len(x),
+                        0,
+                    )
+                )
+    else:
         max_col = max(el.shape[1] for el in els if not_missing(el))
         orig_cols = [el.shape[1] if not_missing(el) else 0 for el in els]
         reindexers = [

anndata/_core/views.py

@@ -1,5 +1,6 @@
 from contextlib import contextmanager
 from copy import deepcopy
+from enum import Enum
 from functools import reduce, singledispatch, wraps
 from typing import Any, KeysView, Optional, Sequence, Tuple
 import warnings
@@ -12,7 +13,7 @@
 import anndata
 from anndata._warnings import ImplicitModificationWarning
 from .access import ElementRef
-from ..compat import ZappyArray, DaskArray
+from ..compat import ZappyArray, AwkArray, DaskArray
 
 
 class _SetItemMixin:
@@ -198,6 +199,76 @@ def as_view_zappy(z, view_args):
     return z
 
 
+try:
+    from ..compat import awkward as ak
+    import weakref
+
+    # Registry to store weak references from AwkwardArrayViews to their parent AnnData container
+    _registry = weakref.WeakValueDictionary()
+    _PARAM_NAME = "_view_args"
+
+    class AwkwardArrayView(_ViewMixin, AwkArray):
+        @property
+        def _view_args(self):
+            """Override _view_args to retrieve the values from awkward arrays parameters.
+
+            Awkward arrays cannot be subclassed like other python objects. Instead subclasses need
+            to be attached as "behavior". These "behaviors" cannot take any additional parameters (as we do
+            for other data types to store `_view_args`). Therefore, we need to store `_view_args` using awkward's
+            parameter mechanism. These parameters need to be json-serializable, which is why we can't store
+            ElementRef directly, but need to replace the reference to the parent AnnDataView container with a weak
+            reference.
+            """
+            parent_key, attrname, keys = self.layout.parameter(_PARAM_NAME)
+            if parent_key is None or attrname is None or keys is None:
+                raise KeyError(
+                    "AwkwardArrayView does not hold reference to original AnnData object."
+                )
+            else:
+                try:
+                    parent = _registry[parent_key]
+                except KeyError:
+                    raise KeyError(
+                        "AwkwardArrayView has invalid reference to original AnnData object."
+                    )
+                else:
+                    return ElementRef(parent, attrname, keys)
+
+        def __copy__(self) -> AwkArray:
+            """
+            Turn the AwkwardArrayView into an actual AwkwardArray with no special behavior.
+
+            Need to override __copy__ instead of `.copy()` as awkward arrays don't implement `.copy()`
+            and are copied using python's standard copy mechanism in `aligned_mapping.py`.
+            """
+            array = self
+            # makes a shallow copy and removes the reference to the original AnnData object
+            array = ak.with_parameter(self, _PARAM_NAME, None)
+            array = ak.with_parameter(array, "__array__", None)
+            return array
+
+    @as_view.register(AwkArray)
+    def as_view_awkarray(array, view_args):
+        parent, attrname, keys = view_args
+        parent_key = f"target-{id(parent)}"
+        _registry[parent_key] = parent
+        # TODO: See https://github.com/scverse/anndata/pull/647#discussion_r963494798_ for more details and
+        # possible strategies to stack behaviors.
+        if type(array).__name__ != "Array":
+            raise NotImplementedError(
+                "Cannot create a view of an awkward array with __array__ parameter. "
+                "Please open an issue in the AnnData repo and describe your use-case."
+            )
+        array = ak.with_parameter(array, _PARAM_NAME, (parent_key, attrname, keys))
+        array = ak.with_parameter(array, "__array__", "AwkwardArrayView")
+        return array
+
+    ak.behavior["AwkwardArrayView"] = AwkwardArrayView
+
+except ImportError:
+    pass
+
+
 def _resolve_idxs(old, new, adata):
     t = tuple(_resolve_idx(old[i], new[i], adata.shape[i]) for i in (0, 1))
     return t