ngio 0.5.0b6__py3-none-any.whl

ngio/tables/backends/__init__.py

@@ -0,0 +1,57 @@
+"""Ngio Tables backend implementations."""
+
+from ngio.tables.backends._abstract_backend import AbstractTableBackend, BackendMeta
+from ngio.tables.backends._anndata import AnnDataBackend, AnnDataBackendV1
+from ngio.tables.backends._csv import CsvTableBackend
+from ngio.tables.backends._json import JsonTableBackend
+from ngio.tables.backends._parquet import ParquetTableBackend
+from ngio.tables.backends._table_backends import (
+    DefaultTableBackend,
+    ImplementedTableBackends,
+    TableBackend,
+    TableBackendProtocol,
+)
+from ngio.tables.backends._utils import (
+    TabularData,
+    convert_anndata_to_pandas,
+    convert_anndata_to_polars,
+    convert_pandas_to_anndata,
+    convert_pandas_to_polars,
+    convert_polars_to_anndata,
+    convert_polars_to_pandas,
+    convert_to_anndata,
+    convert_to_pandas,
+    convert_to_polars,
+    normalize_anndata,
+    normalize_pandas_df,
+    normalize_polars_lf,
+    normalize_table,
+)
+
+__all__ = [
+    "AbstractTableBackend",
+    "AnnDataBackend",
+    "AnnDataBackendV1",
+    "BackendMeta",
+    "CsvTableBackend",
+    "DefaultTableBackend",
+    "ImplementedTableBackends",
+    "JsonTableBackend",
+    "ParquetTableBackend",
+    "TableBackend",
+    "TableBackendProtocol",
+    "TabularData",
+    "convert_anndata_to_pandas",
+    "convert_anndata_to_polars",
+    "convert_pandas_to_anndata",
+    "convert_pandas_to_polars",
+    "convert_polars_to_anndata",
+    "convert_polars_to_pandas",
+    "convert_to_anndata",
+    "convert_to_pandas",
+    "convert_to_polars",
+    "normalize_anndata",
+    "normalize_pandas_df",
+    "normalize_polars_lf",
+    "normalize_table",
+]

ngio/tables/backends/_abstract_backend.py

@@ -0,0 +1,240 @@
+from abc import ABC, abstractmethod
+from typing import Literal
+
+from anndata import AnnData
+from pandas import DataFrame
+from polars import DataFrame as PolarsDataFrame
+from polars import LazyFrame
+from pydantic import BaseModel, ConfigDict
+
+from ngio.tables.backends._utils import (
+    TabularData,
+    convert_to_anndata,
+    convert_to_pandas,
+    convert_to_polars,
+)
+from ngio.utils import NgioValueError, ZarrGroupHandler
+
+
+class BackendMeta(BaseModel):
+    """Metadata for the backend."""
+
+    backend: str = "anndata_v1"
+    index_key: str | None = None
+    index_type: Literal["int", "str"] | None = None
+
+    model_config = ConfigDict(extra="allow")
+
+
+class AbstractTableBackend(ABC):
+    """Abstract class for table backends."""
+
+    def set_group_handler(
+        self,
+        group_handler: ZarrGroupHandler,
+        index_key: str | None = None,
+        index_type: Literal["int", "str"] | None = None,
+    ) -> None:
+        """Attach a group handler to the backend.
+
+        Index keys and index types are used to ensure that the
+        serialization and deserialization of the table
+        is consistent across different backends.
+
+        Making sure that this is consistent is
+        a duty of the backend implementations.
+        """
+        self._group_handler = group_handler
+        self._index_key = index_key
+        self._index_type = index_type
+
+    @staticmethod
+    @abstractmethod
+    def backend_name() -> str:
+        """Return the name of the backend.
+
+        As a convention we set name as:
+            {backend_name}_v{version}
+
+        Where the version is a integer.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def implements_anndata() -> bool:
+        """Check if the backend implements the anndata protocol.
+
+        If this is True, the backend should implement the
+        `write_from_anndata` method.
+
+        AnnData objects are more complex than DataFrames,
+        so if this is true the backend should implement the
+        full serialization of the AnnData object.
+
+        If this is False, these methods should raise a
+        `NotImplementedError`.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def implements_pandas() -> bool:
+        """Check if the backend implements the pandas protocol.
+
+        If this is True, the backend should implement the
+        `write_from_dataframe` methods.
+
+        If this is False, these methods should raise a
+        `NotImplementedError`.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def implements_polars() -> bool:
+        """Check if the backend implements the polars protocol.
+
+        If this is True, the backend should implement the
+        `write_from_polars` methods.
+
+        If this is False, these methods should raise a
+        `NotImplementedError`.
+        """
+        raise NotImplementedError
+
+    @property
+    def group_handler(self) -> ZarrGroupHandler:
+        """Get the group handler."""
+        return self._group_handler
+
+    @property
+    def index_key(self) -> str | None:
+        """Get the index key."""
+        return self._index_key
+
+    @property
+    def index_type(self) -> Literal["int", "str"] | None:
+        """Get the index type."""
+        if self._index_type is None:
+            return None
+
+        if self._index_type not in ["int", "str"]:
+            raise NgioValueError(
+                f"Invalid index type {self._index_type}. Must be 'int' or 'str'."
+            )
+        return self._index_type  # type: ignore[return-value]
+
+    @abstractmethod
+    def load(self) -> TabularData:
+        """Load the table from the store.
+
+        This is a generic load method.
+        Based on the explicit mode or the type of the table,
+        it will call the appropriate load method.
+        """
+        ...
+
+    def load_as_anndata(self) -> AnnData:
+        """Load the table as an AnnData object.
+
+        Since the AnnData object is more complex than a DataFrame,
+        selecting columns is not implemented, because it is not
+        straightforward to do so for an arbitrary AnnData object.
+        """
+        table = self.load()
+        return convert_to_anndata(
+            table,
+            index_key=self.index_key,
+        )
+
+    def load_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame.
+
+        If columns are provided, the table should be filtered
+        """
+        table = self.load()
+        return convert_to_pandas(
+            table,
+            index_key=self.index_key,
+            index_type=self.index_type,
+        )
+
+    def load_as_polars_lf(self) -> LazyFrame:
+        """Load the table as a polars LazyFrame.
+
+        If columns are provided, the table should be filtered
+        """
+        table = self.load()
+        return convert_to_polars(
+            table,
+            index_key=self.index_key,
+            index_type=self.index_type,
+        )
+
+    def write_from_pandas(self, table: DataFrame) -> None:
+        """Serialize the table from a pandas DataFrame."""
+        raise NotImplementedError(
+            f"Backend {self.backend_name()} does not support "
+            "serialization of DataFrame objects."
+        )
+
+    def write_from_anndata(self, table: AnnData) -> None:
+        """Serialize the table from an AnnData object."""
+        raise NotImplementedError(
+            f"Backend {self.backend_name()} does not support "
+            "serialization of AnnData objects."
+        )
+
+    def write_from_polars(self, table: PolarsDataFrame | LazyFrame) -> None:
+        """Serialize the table from a polars DataFrame or LazyFrame."""
+        raise NotImplementedError(
+            f"Backend {self.backend_name()} does not support "
+            "serialization of Polars objects."
+        )
+
+    def write_metadata(self, metadata: dict | None = None) -> None:
+        """Write the metadata to the store."""
+        if metadata is None:
+            metadata = {}
+
+        attrs = self._group_handler.reopen_group().attrs.asdict()
+        # This is required by anndata to identify the format
+        if "encoding-type" in attrs:
+            metadata["encoding-type"] = attrs["encoding-type"]
+        if "encoding-version" in attrs:
+            metadata["encoding-version"] = attrs["encoding-version"]
+
+        backend_metadata = BackendMeta(
+            backend=self.backend_name(),
+            index_key=self.index_key,
+            index_type=self.index_type,
+        ).model_dump(exclude_none=True)
+        metadata.update(backend_metadata)
+        self._group_handler.write_attrs(metadata)
+
+    def write(
+        self,
+        table_data: TabularData,
+        metadata: dict | None = None,
+    ) -> None:
+        """Serialize the table to the store, and write the metadata.
+
+        This is a generic write method.
+        Based on the explicit mode or the type of the table,
+        it will call the appropriate write method.
+        """
+        if isinstance(table_data, DataFrame):
+            self.write_from_pandas(table_data)
+        elif isinstance(table_data, AnnData):
+            self.write_from_anndata(table_data)
+        elif isinstance(table_data, PolarsDataFrame | LazyFrame):
+            self.write_from_polars(table_data)
+        else:
+            raise NgioValueError(
+                f"Unsupported table type {type(table_data)}. "
+                "Please specify the mode explicitly. "
+                "Supported serialization modes are: "
+                "'pandas', 'anndata', 'polars'."
+            )
+        self.write_metadata(metadata)

ngio/tables/backends/_anndata.py

@@ -0,0 +1,139 @@
+import zarr
+from anndata import AnnData
+from anndata._settings import settings
+from pandas import DataFrame
+from polars import DataFrame as PolarsDataFrame
+from polars import LazyFrame
+from zarr.storage import FsspecStore, LocalStore, MemoryStore
+
+from ngio.tables.backends._abstract_backend import AbstractTableBackend
+from ngio.tables.backends._anndata_utils import (
+    custom_anndata_read_zarr,
+)
+from ngio.tables.backends._utils import (
+    convert_pandas_to_anndata,
+    convert_polars_to_anndata,
+    normalize_anndata,
+)
+from ngio.utils import NgioValueError, copy_group
+
+
+class AnnDataBackend(AbstractTableBackend):
+    """A class to load and write tables from/to an AnnData object."""
+
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "anndata"
+
+    @staticmethod
+    def implements_anndata() -> bool:
+        """Check if the backend implements the anndata protocol."""
+        return True
+
+    @staticmethod
+    def implements_pandas() -> bool:
+        """Whether the handler implements the dataframe protocol."""
+        return True
+
+    @staticmethod
+    def implements_polars() -> bool:
+        """Whether the handler implements the polars protocol."""
+        return True
+
+    def load_as_anndata(self) -> AnnData:
+        """Load the table as an AnnData object."""
+        settings.zarr_write_format = self._group_handler.zarr_format
+        anndata = custom_anndata_read_zarr(self._group_handler._group)
+        anndata = normalize_anndata(anndata, index_key=self.index_key)
+        return anndata
+
+    def load(self) -> AnnData:
+        """Load the table as an AnnData object."""
+        return self.load_as_anndata()
+
+    def _write_to_local_store(
+        self, store: LocalStore, path: str, table: AnnData
+    ) -> None:
+        """Write the AnnData table to a LocalStore."""
+        store_path = f"{store.root}/{path}"
+        table.write_zarr(store_path)
+
+    def _write_to_fsspec_store(
+        self, store: FsspecStore, path: str, table: AnnData
+    ) -> None:
+        """Write the AnnData table to a FsspecStore."""
+        full_url = f"{store.path}/{path}"
+        fs = store.fs
+        mapper = fs.get_mapper(full_url)
+        table.write_zarr(mapper)
+
+    def _write_to_memory_store(
+        self, store: MemoryStore, path: str, table: AnnData
+    ) -> None:
+        """Write the AnnData table to a MemoryStore."""
+        store = MemoryStore()
+        table.write_zarr(store)
+        anndata_group = zarr.open_group(store, mode="r")
+        copy_group(
+            anndata_group,
+            self._group_handler._group,
+            suppress_warnings=True,
+        )
+
+    def write_from_anndata(self, table: AnnData) -> None:
+        """Serialize the table from an AnnData object."""
+        # Make sure to use the correct zarr format
+        settings.zarr_write_format = self._group_handler.zarr_format
+        store = self._group_handler.store
+        path = self._group_handler.group.path
+        if isinstance(store, LocalStore):
+            self._write_to_local_store(
+                store,
+                path,
+                table,
+            )
+        elif isinstance(store, FsspecStore):
+            self._write_to_fsspec_store(
+                store,
+                path,
+                table,
+            )
+        elif isinstance(store, MemoryStore):
+            self._write_to_memory_store(
+                store,
+                path,
+                table,
+            )
+        else:
+            raise NgioValueError(
+                f"Ngio does not support writing an AnnData table to a "
+                f"store of type {type(store)}. "
+                "Please make sure to use a compatible "
+                "store like a LocalStore, or FsspecStore."
+            )
+
+    def write_from_pandas(self, table: DataFrame) -> None:
+        """Serialize the table from a pandas DataFrame."""
+        anndata = convert_pandas_to_anndata(
+            table,
+            index_key=self.index_key,
+        )
+        self.write_from_anndata(anndata)
+
+    def write_from_polars(self, table: PolarsDataFrame | LazyFrame) -> None:
+        """Consolidate the metadata in the store."""
+        anndata = convert_polars_to_anndata(
+            table,
+            index_key=self.index_key,
+        )
+        self.write_from_anndata(anndata)
+
+
+class AnnDataBackendV1(AnnDataBackend):
+    """A wrapper for the AnnData backend that for backwards compatibility."""
+
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "anndata_v1"

ngio/tables/backends/_anndata_utils.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import zarr
+from anndata import AnnData
+from anndata._io.specs import read_elem
+from anndata._io.utils import _read_legacy_raw
+from anndata._io.zarr import read_dataframe
+from anndata.compat import _clean_uns
+from anndata.experimental import read_dispatched
+
+from ngio.utils import (
+    NgioValueError,
+    StoreOrGroup,
+    open_group_wrapper,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+
+def custom_anndata_read_zarr(
+    store: StoreOrGroup, elem_to_read: Sequence[str] | None = None
+) -> AnnData:
+    """Read from a hierarchical Zarr array store.
+
+    # Implementation originally from https://github.com/scverse/anndata/blob/main/src/anndata/_io/zarr.py
+    # Original implementation would not work with remote storages so we had to copy it
+    # here and slightly modified it to work with remote storages.
+
+    Args:
+        store (StoreOrGroup): A store or group to read the AnnData from.
+        elem_to_read (Sequence[str] | None): The elements to read from the store.
+    """
+    group = open_group_wrapper(store=store, mode="r")
+    if elem_to_read is None:
+        elem_to_read = [
+            "X",
+            "obs",
+            "var",
+            "uns",
+            "obsm",
+            "varm",
+            "obsp",
+            "varp",
+            "layers",
+        ]
+
+    # Read with handling for backwards compat
+    def callback(func: Callable, elem_name: str, elem: Any, iospec: Any) -> Any:
+        if iospec.encoding_type == "anndata" or elem_name.endswith("/"):
+            ad_kwargs = {}
+            # Some of these elem fail on https
+            # So we only include the ones that are strictly necessary
+            # for fractal tables
+            # This fails on some https
+            # base_elem += list(elem.keys())
+            for k in elem_to_read:
+                v = elem.get(k)
+                if v is not None and not k.startswith("raw."):
+                    ad_kwargs[k] = read_dispatched(v, callback)  # type: ignore
+            return AnnData(**ad_kwargs)
+
+        elif elem_name.startswith("/raw."):
+            return None
+        elif elem_name in {"/obs", "/var"}:
+            return read_dataframe(elem)
+        elif elem_name == "/raw":
+            # Backwards compat
+            return _read_legacy_raw(group, func(elem), read_dataframe, func)
+        return func(elem)
+
+    adata = read_dispatched(group, callback=callback)  # type: ignore
+
+    # Backwards compat (should figure out which version)
+    if "raw.X" in group:
+        raw = AnnData(**_read_legacy_raw(group, adata.raw, read_dataframe, read_elem))  # type: ignore
+        raw.obs_names = adata.obs_names  # type: ignore
+        adata.raw = raw  # type: ignore
+
+    # Backwards compat for <0.7
+    if isinstance(group["obs"], zarr.Array):
+        _clean_uns(adata)
+
+    if isinstance(adata, dict):
+        adata = AnnData(**adata)  # type: ignore
+    if not isinstance(adata, AnnData):
+        raise NgioValueError(f"Expected an AnnData object, but got {type(adata)}")
+    return adata

ngio/tables/backends/_csv.py

@@ -0,0 +1,19 @@
+from ngio.tables.backends._py_arrow_backends import PyArrowBackend
+
+
+class CsvTableBackend(PyArrowBackend):
+    """A class to load and write small tables in CSV format."""
+
+    def __init__(
+        self,
+    ):
+        """Initialize the CsvTableBackend."""
+        super().__init__(
+            table_name="table.csv",
+            table_format="csv",
+        )
+
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "csv"

ngio/tables/backends/_json.py

@@ -0,0 +1,92 @@
+import pandas as pd
+from pandas import DataFrame
+from polars import DataFrame as PolarsDataFrame
+from polars import LazyFrame
+
+from ngio.tables.backends._abstract_backend import AbstractTableBackend
+from ngio.tables.backends._utils import (
+    normalize_pandas_df,
+    normalize_polars_lf,
+)
+from ngio.utils import NgioError
+
+
+class JsonTableBackend(AbstractTableBackend):
+    """A class to load and write small tables in the zarr group .attrs (json)."""
+
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "json"
+
+    @staticmethod
+    def implements_anndata() -> bool:
+        """Whether the handler implements the anndata protocol."""
+        return False
+
+    @staticmethod
+    def implements_pandas() -> bool:
+        """Whether the handler implements the dataframe protocol."""
+        return True
+
+    @staticmethod
+    def implements_polars() -> bool:
+        """Whether the handler implements the polars protocol."""
+        return True
+
+    def _get_table_group(self):
+        """Get the table group, creating it if it doesn't exist."""
+        try:
+            table_group = self._group_handler.get_group(path="table", create_mode=True)
+        except NgioError as e:
+            raise NgioError(
+                "Could not get or create a 'table' group in the store "
+                f"{self._group_handler.store} path "
+                f"{self._group_handler.group.path}/table."
+            ) from e
+        return table_group
+
+    def load_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame."""
+        table_dict = self._get_table_group().attrs.asdict()
+        data_frame = pd.DataFrame.from_dict(table_dict)
+        data_frame = normalize_pandas_df(
+            data_frame,
+            index_key=self.index_key,
+            index_type=self.index_type,
+            reset_index=False,
+        )
+        return data_frame
+
+    def load(self) -> DataFrame:
+        return self.load_as_pandas_df()
+
+    def _write_from_dict(self, table: dict) -> None:
+        """Write the table from a dictionary to the store."""
+        table_group = self._get_table_group()
+        table_group.attrs.clear()
+        table_group.attrs.update(table)
+
+    def write_from_pandas(self, table: DataFrame) -> None:
+        """Write the table from a pandas DataFrame."""
+        table = normalize_pandas_df(
+            table,
+            index_key=self.index_key,
+            index_type=self.index_type,
+            reset_index=True,
+        )
+        table_dict = table.to_dict(orient="list")
+        self._write_from_dict(table=table_dict)
+
+    def write_from_polars(self, table: PolarsDataFrame | LazyFrame) -> None:
+        """Write the table from a polars DataFrame or LazyFrame."""
+        table = normalize_polars_lf(
+            table,
+            index_key=self.index_key,
+            index_type=self.index_type,
+        )
+        if isinstance(table, LazyFrame):
+            table = table.collect()
+
+        table_dict = table.to_dict(as_series=False)
+        self._write_from_dict(table=table_dict)

ngio/tables/backends/_parquet.py

@@ -0,0 +1,19 @@
+from ngio.tables.backends._py_arrow_backends import PyArrowBackend
+
+
+class ParquetTableBackend(PyArrowBackend):
+    """A class to load and write small tables in Parquet format."""
+
+    def __init__(
+        self,
+    ):
+        """Initialize the ParquetTableBackend."""
+        super().__init__(
+            table_name="table.parquet",
+            table_format="parquet",
+        )
+
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "parquet"