ngio 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

ngio/tables/backends/_abstract_backend.py

@@ -1,11 +1,29 @@
 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
 
-from ngio.utils import ZarrGroupHandler
+from ngio.tables.backends._utils import (
+    convert_anndata_to_pandas,
+    convert_anndata_to_polars,
+    convert_pandas_to_anndata,
+    convert_pandas_to_polars,
+    convert_polars_to_anndata,
+    convert_polars_to_pandas,
+)
+from ngio.utils import NgioValueError, ZarrGroupHandler
+
+
+class BackendMeta(BaseModel):
+    """Metadata for the backend."""
+
+    backend: str | None = None
+    index_key: str | None = None
+    index_type: Literal["int", "str"] | None = None
 
 
 class AbstractTableBackend(ABC):
@@ -15,10 +33,12 @@ class AbstractTableBackend(ABC):
         self,
         group_handler: ZarrGroupHandler,
         index_key: str | None = None,
-        index_type: Literal["int", "str"] = "int",
+        index_type: Literal["int", "str"] | None = None,
     ):
         """Initialize the handler.
 
+        This is a base class for the table backends protocol.
+
         Args:
             group_handler (ZarrGroupHandler): An object to handle the Zarr group
                 containing the table data.
@@ -32,40 +52,205 @@ 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
+        `load_as_anndata` and `write_from_anndata` methods.
+
+        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
+        `load_as_dataframe` and `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
+        `load_as_polars` and `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]
+
+    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.
+        """
+        if self.implements_pandas():
+            return convert_pandas_to_anndata(
+                self.load_as_pandas_df(),
+                index_key=self.index_key,
+            )
+        elif self.implements_polars():
+            return convert_polars_to_anndata(
+                self.load_as_polars_lf(),
+                index_key=self.index_key,
+            )
+        else:
+            raise NgioValueError(
+                "Backend does not implement any of the protocols. "
+                "A backend should implement at least one of the "
+                "following protocols: anndata, pandas, polars."
+            )
+
+    def load_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame.
+
+        If columns are provided, the table should be filtered
+        """
+        if self.implements_anndata():
+            return convert_anndata_to_pandas(
+                self.load_as_anndata(),
+                index_key=self.index_key,
+                index_type=self.index_type,
+            )
+        elif self.implements_polars():
+            return convert_polars_to_pandas(
+                self.load_as_polars_lf(),
+                index_key=self.index_key,
+                index_type=self.index_type,
+            )
+        else:
+            raise NgioValueError(
+                "Backend does not implement any of the protocols. "
+                "A backend should implement at least one of the "
+                "following protocols: anndata, pandas, polars."
+            )
+
+    def load_as_polars_lf(self) -> LazyFrame:
+        """Load the table as a polars LazyFrame.
+
+        If columns are provided, the table should be filtered
+        """
+        if self.implements_anndata():
+            return convert_anndata_to_polars(
+                self.load_as_anndata(),
+                index_key=self.index_key,
+                index_type=self.index_type,
+            ).lazy()
+        elif self.implements_pandas():
+            return convert_pandas_to_polars(
+                self.load_as_pandas_df(),
+                index_key=self.index_key,
+                index_type=self.index_type,
+            ).lazy()
+        else:
+            raise NgioValueError(
+                "Backend does not implement any of the protocols. "
+                "A backend should implement at least one of the "
+                "following protocols: anndata, pandas, polars."
+            )
+
+    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 = {}
+
+        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: DataFrame | AnnData | PolarsDataFrame | LazyFrame,
+        metadata: dict | None = None,
+        mode: Literal["pandas", "anndata", "polars"] | 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 mode == "pandas" or isinstance(table, DataFrame):
+            self.write_from_pandas(table)  # type: ignore[arg-type]
+        elif mode == "anndata" or isinstance(table, AnnData):
+            self.write_from_anndata(table)  # type: ignore[arg-type]
+        elif mode == "polars" or isinstance(table, PolarsDataFrame | LazyFrame):
+            self.write_from_polars(table)
+        else:
+            raise NgioValueError(
+                f"Unsupported table type {type(table)}. "
+                "Please specify the mode explicitly. "
+                "Supported serialization modes are: "
+                "'pandas', 'anndata', 'polars'."
+            )
+        self.write_metadata(metadata)

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,9 +10,7 @@ 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,
@@ -26,7 +20,7 @@ if TYPE_CHECKING:
     from collections.abc import Callable, Collection
 
 
-def custom_read_zarr(
+def custom_anndata_read_zarr(
     store: StoreOrGroup, elem_to_read: Collection[str] | None = None
 ) -> AnnData:
     """Read from a hierarchical Zarr array store.
@@ -39,7 +33,7 @@ def custom_read_zarr(
         store (StoreOrGroup): A store or group to read the AnnData from.
         elem_to_read (Collection[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 not isinstance(group.store, zarr.DirectoryStore):
         elem_to_read = ["X", "obs", "var"]
@@ -96,103 +90,3 @@ def custom_read_zarr(
     if not isinstance(adata, AnnData):
         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/_anndata_v1.py

@@ -1,14 +1,16 @@
-from collections.abc import Collection
-from pathlib import Path
-
 from anndata import AnnData
 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._anndata_utils import (
-    anndata_to_dataframe,
-    custom_read_zarr,
-    dataframe_to_anndata,
+    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
 
@@ -18,59 +20,54 @@ class AnnDataBackend(AbstractTableBackend):
 
     @staticmethod
     def backend_name() -> str:
-        """The name of the backend."""
+        """Return the name of the backend."""
         return "anndata_v1"
 
     @staticmethod
     def implements_anndata() -> bool:
-        """Whether the handler implements the anndata protocol."""
+        """Check if the backend implements the anndata protocol."""
         return True
 
     @staticmethod
-    def implements_dataframe() -> bool:
+    def implements_pandas() -> bool:
         """Whether the handler implements the dataframe protocol."""
         return True
 
-    def load_columns(self) -> list[str]:
-        """List all labels in the group."""
-        return list(self.load_as_dataframe().columns)
+    @staticmethod
+    def implements_polars() -> bool:
+        """Whether the handler implements the polars protocol."""
+        return True
 
-    def load_as_anndata(self, columns: Collection[str] | None = None) -> AnnData:
-        """Load the metadata in the store."""
-        anndata = custom_read_zarr(self._group_handler._group)
-        if columns is not None:
-            raise NotImplementedError(
-                "Selecting columns is not implemented for AnnData."
-            )
+    def load_as_anndata(self) -> AnnData:
+        """Load the table as an AnnData object."""
+        anndata = custom_anndata_read_zarr(self._group_handler._group)
+        anndata = normalize_anndata(anndata, index_key=self.index_key)
         return anndata
 
-    def load_as_dataframe(self, columns: Collection[str] | None = None) -> DataFrame:
-        """List all labels in the group."""
-        dataframe = anndata_to_dataframe(
-            self.load_as_anndata(),
-            index_key=self._index_key,
-            index_type=self._index_type,
-        )
-        if columns is not None:
-            dataframe = dataframe[columns]
-        return dataframe
-
-    def write_from_dataframe(
-        self, table: DataFrame, metadata: dict | None = None
-    ) -> None:
-        """Consolidate the metadata in the store."""
-        anndata = dataframe_to_anndata(table, index_key=self._index_key)
-        self.write_from_anndata(anndata, metadata)
-
-    def write_from_anndata(self, table: AnnData, metadata: dict | None = None) -> None:
-        """Consolidate the metadata in the store."""
-        store = self._group_handler.store
-        if not isinstance(store, str | Path):
+    def write_from_anndata(self, table: AnnData) -> None:
+        """Serialize the table from an AnnData object."""
+        full_url = self._group_handler.full_url
+        if full_url is None:
             raise NgioValueError(
-                "To write an AnnData object the store must be a local path/str."
+                f"Ngio does not support writing file from a "
+                f"store of type {type(self._group_handler)}. "
+                "Please make sure to use a compatible "
+                "store like a zarr.DirectoryStore."
             )
+        table.write_zarr(full_url)  # type: ignore
 
-        store = Path(store) / self._group_handler.group.path
-        table.write_zarr(store)
-        if metadata is not None:
-            self._group_handler.write_attrs(metadata)
+    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)