[python] Ingest performance improvements (#3865)

* first cut at ingest tuning

* mp optional param

* fix tests, add load from S3

* reshape fix

* lint

* add multiprocessing config warning

* performance tuning

* pass context

* lint

* add test

* interim checkpoint on enum cleanup

* add test for subset

* add new extend enum code (commented out)

* migrate to new enum API

* bump typeguard version to work around typeguard#504

* improve doctrings

* more cleanup

* loosen existance checks when from_anndata/from_h5ad are run in schema_only mode

* cleanup and additional error checking

* more docstring cleanup

* more docstring cleanup

* more docstrings

* pr fb

* pr fb

* more pr fb

* increase test code coverage

* additional code coverage

* lint

* PR fb

* PR fb

* more PR fb

* lint

* remove extraneous line break in docstring

Author: bkmartinjr

apis/python/setup.py

@@ -344,6 +344,7 @@ def run(self):
     install_requires=[
         "anndata>=0.10.1",
         "attrs>=22.2",
+        "more-itertools",
         "numpy",
         "pandas",
         "pyarrow",

apis/python/src/tiledbsoma/io/_registration/ambient_label_mappings.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from typing import Sequence
+
+import pandas as pd
+import pyarrow as pa
+
+from tiledbsoma import DataFrame
+
+
+def get_enumerations(
+    df: DataFrame, column_names: Sequence[str]
+) -> dict[str, pd.CategoricalDtype]:
+    """Look up enum info in schema, and return as a Pandas CategoricalDType. This
+    is a convenience wrapper around ``DataFrame.get_enumeration_values``, for use
+    in the registration module."""
+
+    # skip columns which are not of type dictionary
+    column_names = [
+        c for c in column_names if pa.types.is_dictionary(df.schema.field(c).type)
+    ]
+    return {
+        k: pd.CategoricalDtype(categories=v, ordered=df.schema.field(k).type.ordered)
+        for k, v in df.get_enumeration_values(column_names).items()
+    }
+
+
+def extend_enumerations(df: DataFrame, columns: dict[str, pd.CategoricalDtype]) -> None:
+    """
+    Extend enumerations as needed, starting with a CategoricalDType for each
+    cat/enum/dict column. A convenience wrapper around ``DataFrame.extend_enumeration_values``,
+    for use in the registration module.
+
+    DataFrame must be open for write.
+    """
+
+    current_enums = get_enumerations(df, list(columns.keys()))
+    columns_to_extend = {}
+    for column_name, cat_dtype in columns.items():
+
+        # first confirm this is a dictionary. If it has been decategorical-ized, i.e.,
+        # are an array of the value type, don't extend.
+        if column_name not in current_enums:
+            assert not pa.types.is_dictionary(df.schema.field(column_name).type)
+            continue
+
+        # determine if we have any new enum values in this column
+        existing_dtype = current_enums[column_name]
+        new_enum_values = pd.Index(cat_dtype.categories).difference(
+            existing_dtype.categories, sort=False
+        )
+        if len(new_enum_values) == 0:
+            continue
+
+        # if there are new values, extend the array schema enum
+        new_enum_values = pa.array(new_enum_values.to_numpy())
+        columns_to_extend[column_name] = new_enum_values
+
+    # and evolve the schema
+    df.extend_enumeration_values(columns_to_extend, deduplicate=False)

apis/python/src/tiledbsoma/io/_registration/enum.py

@@ -1,102 +1,92 @@
 # Copyright (c) TileDB, Inc. and The Chan Zuckerberg Initiative Foundation
 #
 # Licensed under the MIT License.
+from __future__ import annotations
 
-from typing import Dict, List, Tuple
+from typing import cast
 
 import anndata as ad
 import attrs
+import numpy as np
+import numpy.typing as npt
 import pandas as pd
 from typing_extensions import Self
 
-import tiledbsoma
-import tiledbsoma.logging
 
-
-@attrs.define(kw_only=True)
+@attrs.define(kw_only=True, frozen=True)
 class AxisIDMapping:
-    """
-    For a single to-be-appended AnnData/H5AD input in SOMA multi-file append-mode ingestion, this
+    """For a single to-be-appended AnnData/H5AD input in SOMA multi-file append-mode ingestion, this
     class tracks the mapping of input-data ``obs`` or ``var`` 0-up offsets to SOMA join ID values
     for the destination SOMA experiment.
 
-    See module-level comments for more information.
+    Private class
     """
 
-    # Tuple not List so this can't be modified by accident when passed into some function somewhere
-    data: Tuple[int, ...]
+    data: npt.NDArray[np.int64]
 
-    def is_identity(self) -> bool:
-        for i, data in enumerate(self.data):
-            if data != i:
-                return False
-        return True
+    def __attrs_post_init__(self) -> None:
+        self.data.setflags(write=False)
 
     def get_shape(self) -> int:
         if len(self.data) == 0:
             return 0
         else:
-            return 1 + max(self.data)
+            return int(self.data.max() + 1)
+
+    def is_identity(self) -> bool:
+        # fast rejection first
+        if self.get_shape() != len(self.data) or self.data[0] != 0:
+            return False
+
+        return np.array_equal(self.data, np.arange(0, len(self.data)))
 
     @classmethod
     def identity(cls, n: int) -> Self:
         """This maps 0-up input-file offsets to 0-up soma_joinid values. This is
         important for uns arrays which we never grow on ingest --- rather, we
         sub-nest the entire recursive ``uns`` data structure.
         """
-        return cls(data=tuple(range(n)))
+        return cls(data=np.arange(n, dtype=np.int64))
 
 
-@attrs.define(kw_only=True)
+@attrs.define(kw_only=True, frozen=True)
 class ExperimentIDMapping:
-    """
-    For a single to-be-appended AnnData/H5AD input in SOMA multi-file append-mode ingestion, this
-    class contains an ``ExperimentIDMapping`` for ``obs``, and one ``ExperimentIDMapping`` for
+    """For a single to-be-appended AnnData/H5AD input in SOMA multi-file append-mode ingestion, this
+    class contains an ``AxisIDMapping`` for ``obs``, and one ``AxisIDMapping`` for
     ``var`` in each measurement.
 
-    See module-level comments for more information.
+    Private class
     """
 
     obs_axis: AxisIDMapping
-    var_axes: Dict[str, AxisIDMapping]
+    var_axes: dict[str, AxisIDMapping]
 
     @classmethod
-    def from_isolated_anndata(
-        cls,
-        adata: ad.AnnData,
-        measurement_name: str,
-    ) -> Self:
-        """Factory method to compute offset-to-SOMA-join-ID mappings for a single input file in
-        isolation. This is used when a user is ingesting a single AnnData/H5AD to a single SOMA
-        experiment, not in append mode, allowing us to still have the bulk of the ingestor code to
-        be non-duplicated between non-append mode and append mode.
-        """
-        tiledbsoma.logging.logger.info(
-            "Registration: registering isolated AnnData object."
-        )
+    def from_anndata(cls, adata: ad.AnnData, *, measurement_name: str = "RNA") -> Self:
+        """Create a new ID mapping from an AnnData.
 
-        obs_mapping = AxisIDMapping(data=tuple(range(len(adata.obs))))
-        var_axes = {}
-        var_axes[measurement_name] = AxisIDMapping(data=tuple(range(len(adata.var))))
+        This is useful for creating a new Experiment from a single AnnData.
+        """
+        obs_axis = AxisIDMapping.identity(len(adata.obs))
+        var_axes = {measurement_name: AxisIDMapping.identity(len(adata.var))}
         if adata.raw is not None:
-            var_axes["raw"] = AxisIDMapping(data=tuple(range(len(adata.raw.var))))
-
-        return cls(obs_axis=obs_mapping, var_axes=var_axes)
+            var_axes["raw"] = AxisIDMapping.identity(len(adata.raw.var))
+        return cls(obs_axis=obs_axis, var_axes=var_axes)
 
 
-def get_dataframe_values(df: pd.DataFrame, field_name: str) -> List[str]:
+def get_dataframe_values(df: pd.DataFrame, field_name: str) -> pd.Series:  # type: ignore[type-arg]
     """Extracts the label values (e.g. cell barcode, gene symbol) from an AnnData/H5AD
     ``obs`` or ``var`` dataframe."""
     if field_name in df:
-        values = [str(e) for e in df[field_name]]
+        values = cast(pd.Series, df[field_name].astype(str))  # type: ignore[type-arg]
     elif df.index.name in (field_name, "index", None):
-        values = list(df.index)
+        values = cast(pd.Series, df.index.to_series().astype(str))  # type: ignore[type-arg]
     else:
-        raise ValueError(f"could not find field name {field_name} in dataframe")
+        raise ValueError(f"Could not find field name {field_name} in dataframe.")
 
     # Check the values are unique.
-    if len(values) != len(set(values)):
+    if not values.is_unique:
         raise ValueError(
-            f"non-unique registration values have been provided in field {field_name}"
+            f"Non-unique registration values have been provided in field {field_name}."
         )
     return values

apis/python/src/tiledbsoma/io/_registration/id_mappings.py

@@ -113,7 +113,8 @@ def _hack_patch_anndata() -> ContextManager[object]:
 
     @file_backing.AnnDataFileManager.filename.setter  # type: ignore[misc]
     def filename(
-        self: file_backing.AnnDataFileManager, filename: Union[Path, _FSPathWrapper]
+        self: file_backing.AnnDataFileManager,
+        filename: Union[Path, _FSPathWrapper, None],
     ) -> None:
         self._filename = filename