ngio 0.2.0a2__py3-none-any.whl → 0.5.0b4__py3-none-any.whl

ngio/tables/backends/__init__.py

@@ -1,8 +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__ = ["ImplementedTableBackends", "TableBackendProtocol"]
+__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

@@ -1,29 +1,48 @@
 from abc import ABC, abstractmethod
-from collections.abc import Collection
 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.utils import ZarrGroupHandler
+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 __init__(
+    def set_group_handler(
         self,
         group_handler: ZarrGroupHandler,
         index_key: str | None = None,
-        index_type: Literal["int", "str"] = "int",
-    ):
-        """Initialize the handler.
-
-        Args:
-            group_handler (ZarrGroupHandler): An object to handle the Zarr group
-                containing the table data.
-            index_key (str): The column name to use as the index of the DataFrame.
-            index_type (str): The type of the index column in the DataFrame.
+        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
@@ -32,40 +51,190 @@ class AbstractTableBackend(ABC):
     @staticmethod
     @abstractmethod
     def backend_name() -> str:
-        """The name of the backend."""
+        """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:
-        """Whether the handler implements the anndata protocol."""
+        """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_dataframe() -> bool:
-        """Whether the handler implements the dataframe protocol."""
+    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 load_columns(self) -> list[str]:
-        """List all labels in the group."""
-        raise NotImplementedError
+    def implements_polars() -> bool:
+        """Check if the backend implements the polars protocol.
 
-    def load_as_anndata(self, columns: Collection[str] | None = None) -> AnnData:
-        """Load the metadata in the store."""
-        raise NotImplementedError
+        If this is True, the backend should implement the
+        `write_from_polars` methods.
 
-    def load_as_dataframe(self, columns: Collection[str] | None = None) -> DataFrame:
-        """List all labels in the group."""
+        If this is False, these methods should raise a
+        `NotImplementedError`.
+        """
         raise NotImplementedError
 
-    def write_from_dataframe(
-        self, table: DataFrame, metadata: dict | None = None
+    @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:
-        """Consolidate the metadata in the store."""
-        raise NotImplementedError
+        """Serialize the table to the store, and write the metadata.
 
-    def write_from_anndata(self, table: AnnData, metadata: dict | None = None) -> None:
-        """Consolidate the metadata in the store."""
-        raise NotImplementedError
+        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

@@ -1,11 +1,7 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any
 
-import anndata as ad
-import numpy as np
-import pandas as pd
-import pandas.api.types as ptypes
 import zarr
 from anndata import AnnData
 from anndata._io.specs import read_elem
@@ -14,19 +10,18 @@ from anndata._io.zarr import read_dataframe
 from anndata.compat import _clean_uns
 from anndata.experimental import read_dispatched
 
-from ngio.tables._validators import validate_index_dtype, validate_index_key
 from ngio.utils import (
-    NgioTableValidationError,
+    NgioValueError,
     StoreOrGroup,
     open_group_wrapper,
 )
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Collection
+    from collections.abc import Callable, Sequence
 
 
-def custom_read_zarr(
-    store: StoreOrGroup, elem_to_read: Collection[str] | None = None
+def custom_anndata_read_zarr(
+    store: StoreOrGroup, elem_to_read: Sequence[str] | None = None
 ) -> AnnData:
     """Read from a hierarchical Zarr array store.
 
@@ -36,10 +31,9 @@ def custom_read_zarr(
 
     Args:
         store (StoreOrGroup): A store or group to read the AnnData from.
-        elem_to_read (Collection[str] | None): The elements to read from the store.
+        elem_to_read (Sequence[str] | None): The elements to read from the store.
     """
-    group, _ = open_group_wrapper(store=store, mode="r")
-
+    group = open_group_wrapper(store=store, mode="r")
     if elem_to_read is None:
         elem_to_read = [
             "X",
@@ -89,106 +83,8 @@ def custom_read_zarr(
     if isinstance(group["obs"], zarr.Array):
         _clean_uns(adata)
 
+    if isinstance(adata, dict):
+        adata = AnnData(**adata)  # type: ignore
     if not isinstance(adata, AnnData):
-        raise ValueError(f"Expected an AnnData object, but got {type(adata)}")
+        raise NgioValueError(f"Expected an AnnData object, but got {type(adata)}")
     return adata
-
-
-def _check_for_mixed_types(series: pd.Series) -> None:
-    """Check if the column has mixed types."""
-    if series.apply(type).nunique() > 1:  # type: ignore
-        raise NgioTableValidationError(
-            f"Column {series.name} has mixed types: "
-            f"{series.apply(type).unique()}. "  # type: ignore
-            "Type of all elements must be the same."
-        )
-
-
-def _check_for_supported_types(series: pd.Series) -> Literal["str", "int", "numeric"]:
-    """Check if the column has supported types."""
-    if ptypes.is_string_dtype(series):
-        return "str"
-    if ptypes.is_integer_dtype(series):
-        return "int"
-    if ptypes.is_numeric_dtype(series):
-        return "numeric"
-    raise NgioTableValidationError(
-        f"Column {series.name} has unsupported type: {series.dtype}."
-        " Supported types are string and numerics."
-    )
-
-
-def dataframe_to_anndata(
-    dataframe: pd.DataFrame,
-    index_key: str | None = None,
-    overwrite: bool = False,
-) -> ad.AnnData:
-    """Convert a table DataFrame to an AnnData object.
-
-    Args:
-        dataframe (pd.DataFrame): A pandas DataFrame representing a fractal table.
-        index_key (str): The column name to use as the index of the DataFrame.
-            Default is None.
-        overwrite (bool): Whether to overwrite the index if a different index is found.
-            Default is False.
-    """
-    # Check if the index_key is present in the data frame + optional validations
-    dataframe = validate_index_key(dataframe, index_key, overwrite=overwrite)
-    dataframe = validate_index_dtype(dataframe, index_type="str")
-
-    str_columns, int_columns, num_columns = [], [], []
-    for c_name in dataframe.columns:
-        column_df = dataframe[c_name]
-        _check_for_mixed_types(column_df)  # Mixed types are not allowed in the table
-        c_type = _check_for_supported_types(
-            column_df
-        )  # Only string and numeric types are allowed
-
-        if c_type == "str":
-            str_columns.append(c_name)
-
-        elif c_type == "int":
-            int_columns.append(c_name)
-
-        elif c_type == "numeric":
-            num_columns.append(c_name)
-
-    # Converting all observations to string
-    obs_df = dataframe[str_columns + int_columns]
-    obs_df.index = dataframe.index
-
-    x_df = dataframe[num_columns]
-
-    if x_df.dtypes.nunique() > 1:
-        x_df = x_df.astype("float64")
-
-    if x_df.empty:
-        # If there are no numeric columns, create an empty array
-        # to avoid AnnData failing to create the object
-        x_df = np.zeros((len(obs_df), 0), dtype="float64")
-
-    return ad.AnnData(X=x_df, obs=obs_df)
-
-
-def anndata_to_dataframe(
-    anndata: ad.AnnData,
-    index_key: str | None = "label",
-    index_type: str = "int",
-    overwrite: bool = False,
-) -> pd.DataFrame:
-    """Convert a AnnData object representing a fractal table to a pandas DataFrame.
-
-    Args:
-        anndata (ad.AnnData): An AnnData object representing a fractal table.
-        index_key (str): The column name to use as the index of the DataFrame.
-            Default is 'label'.
-        index_type (str): The type of the index column in the DataFrame.
-            Either 'str' or 'int'. Default is 'int'.
-        overwrite (bool): Whether to overwrite the index if a different index is found.
-            Default is False.
-    """
-    dataframe = anndata.to_df()
-    dataframe[anndata.obs_keys()] = anndata.obs
-    dataframe = validate_index_key(dataframe, index_key, overwrite=overwrite)
-    dataframe = validate_index_dtype(dataframe, index_type)
-    return dataframe

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"