ngio 0.2.9__py3-none-any.whl → 0.3.0a1__py3-none-any.whl

ngio/tables/backends/_table_backends.py

@@ -7,20 +7,22 @@ from pandas import DataFrame
 from polars import DataFrame as PolarsDataFrame
 from polars import LazyFrame
 
-from ngio.tables.backends._anndata_v1 import AnnDataBackend
-from ngio.tables.backends._csv_v1 import CsvTableBackend
-from ngio.tables.backends._json_v1 import JsonTableBackend
+from ngio.tables.backends._anndata import AnnDataBackend
+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._utils import TabularData
 from ngio.utils import NgioValueError, ZarrGroupHandler
 
 
 class TableBackendProtocol(Protocol):
-    def __init__(
+    def set_group_handler(
         self,
         group_handler: ZarrGroupHandler,
         index_key: str | None = None,
         index_type: Literal["int", "str"] | None = None,
-    ):
-        """Backend constructor.
+    ) -> 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
@@ -42,12 +44,21 @@ class TableBackendProtocol(Protocol):
         """
         ...
 
+    @property
+    def group_handler(self) -> ZarrGroupHandler:
+        """Return the group handler."""
+        ...
+
     @staticmethod
     def implements_anndata() -> bool:
         """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.
+        `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`.
@@ -59,7 +70,7 @@ class TableBackendProtocol(Protocol):
         """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.
+        `write_from_dataframe` methods.
 
         If this is False, these methods should raise a
         `NotImplementedError`.
@@ -71,7 +82,7 @@ class TableBackendProtocol(Protocol):
         """Check if the backend implements the polars protocol.
 
         If this is True, the backend should implement the
-        `load_as_polars` and `write_from_polars` methods.
+        `write_from_polars` methods.
 
         If this is False, these methods should raise a
         `NotImplementedError`.
@@ -90,6 +101,16 @@ class TableBackendProtocol(Protocol):
         """Load the table as a polars LazyFrame."""
         ...
 
+    def load(self) -> TabularData:
+        """The default load method.
+
+        This method will be default way to load the table
+        from the backend. This method should wrap one of the
+        `load_as_anndata`, `load_as_dataframe` or `load_as_polars`
+        methods depending on the backend implementation.
+        """
+        ...
+
     def write_from_pandas(self, table: DataFrame) -> None:
         """Serialize the table from a pandas DataFrame."""
         ...
@@ -104,7 +125,7 @@ class TableBackendProtocol(Protocol):
 
     def write(
         self,
-        table: DataFrame | AnnData | PolarsDataFrame | LazyFrame,
+        table_data: DataFrame | AnnData | PolarsDataFrame | LazyFrame,
         metadata: dict[str, str] | None = None,
         mode: Literal["pandas", "anndata", "polars"] | None = None,
     ) -> None:
@@ -144,23 +165,20 @@ class ImplementedTableBackends:
 
     def get_backend(
         self,
-        backend_name: str | None,
+        *,
         group_handler: ZarrGroupHandler,
+        backend_name: str = "anndata",
         index_key: str | None = None,
         index_type: Literal["int", "str"] | None = None,
     ) -> TableBackendProtocol:
         """Try to get a handler for the given store based on the metadata version."""
-        if backend_name is None:
-            # Default to anndata since it is currently
-            # the only backend in use.
-            backend_name = "anndata_v1"
-
         if backend_name not in self._implemented_backends:
             raise NgioValueError(f"Table backend {backend_name} not implemented.")
-        handler = self._implemented_backends[backend_name](
+        backend = self._implemented_backends[backend_name]()
+        backend.set_group_handler(
             group_handler=group_handler, index_key=index_key, index_type=index_type
         )
-        return handler
+        return backend
 
     def add_backend(
         self,
@@ -180,3 +198,6 @@ class ImplementedTableBackends:
 ImplementedTableBackends().add_backend(AnnDataBackend)
 ImplementedTableBackends().add_backend(JsonTableBackend)
 ImplementedTableBackends().add_backend(CsvTableBackend)
+ImplementedTableBackends().add_backend(ParquetTableBackend)
+
+TableBackend = Literal["anndata", "json", "csv", "parquet"] | str | TableBackendProtocol

ngio/tables/backends/_utils.py

@@ -9,7 +9,6 @@ These functions are used to validate and normalize the tables
 to ensure that conversion between formats is consistent.
 """
 
-# %%
 from copy import deepcopy
 from typing import Literal
 
@@ -24,6 +23,8 @@ from polars import LazyFrame
 
 from ngio.utils import NgioTableValidationError, NgioValueError
 
+TabularData = AnnData | DataFrame | PolarsDataFrame | LazyFrame
+
 # -----------------
 # Validation utils
 # -----------------
@@ -460,3 +461,148 @@ def convert_polars_to_anndata(
         pandas_df,
         index_key=index_key,
     )
+
+
+# -----------------
+# Conversion functions
+# -----------------
+
+
+def normalize_table(
+    table_data: TabularData,
+    index_key: str | None = None,
+    index_type: Literal["int", "str"] | None = None,
+) -> TabularData:
+    """Normalize a table to a specific format.
+
+    Args:
+        table_data (TabularData): The table to normalize.
+        index_key (str | None): The column name to use as the index of the DataFrame.
+            Default is None.
+        index_type (str | None): The type of the index column in the DataFrame.
+            Either 'str' or 'int'. Default is None.
+
+    Returns:
+        DataFrame | AnnData | PolarsDataFrame | LazyFrame: Normalized table.
+    """
+    if isinstance(table_data, DataFrame):
+        return normalize_pandas_df(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+            reset_index=False,
+        )
+    if isinstance(table_data, AnnData):
+        return normalize_anndata(table_data, index_key=index_key)
+    if isinstance(table_data, PolarsDataFrame) or isinstance(table_data, LazyFrame):
+        return normalize_polars_lf(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+        )
+    raise NgioValueError(f"Unsupported table type: {type(table_data)}")
+
+
+def convert_to_anndata(
+    table_data: TabularData,
+    index_key: str | None = None,
+) -> AnnData:
+    """Convert a table to an AnnData object.
+
+    Args:
+        table_data (TabularData): The table to convert.
+        index_key (str | None): The column name to use as the index of the DataFrame.
+            Default is None.
+
+    Returns:
+        AnnData: Converted AnnData object.
+    """
+    if isinstance(table_data, AnnData):
+        return normalize_anndata(table_data, index_key=index_key)
+    if isinstance(table_data, DataFrame):
+        return convert_pandas_to_anndata(table_data, index_key=index_key)
+    if isinstance(table_data, PolarsDataFrame) or isinstance(table_data, LazyFrame):
+        return convert_polars_to_anndata(table_data, index_key=index_key)
+    raise NgioValueError(f"Unsupported table type: {type(table_data)}")
+
+
+def convert_to_pandas(
+    table_data: TabularData,
+    index_key: str | None = None,
+    index_type: Literal["int", "str"] | None = None,
+    reset_index: bool = False,
+) -> DataFrame:
+    """Convert a table to a pandas DataFrame.
+
+    Args:
+        table_data (TabularData): The table to convert.
+        index_key (str | None): The column name to use as the index of the DataFrame.
+            Default is None.
+        index_type (str | None): The type of the index column in the DataFrame.
+            Either 'str' or 'int'. Default is None.
+        reset_index (bool): If True the index will be reset (i.e., the index will be
+            converted to a column). If False, the index will be kept as is.
+
+    Returns:
+        DataFrame: Converted pandas DataFrame.
+    """
+    if isinstance(table_data, DataFrame):
+        return normalize_pandas_df(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+            reset_index=reset_index,
+        )
+    if isinstance(table_data, AnnData):
+        return convert_anndata_to_pandas(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+            reset_index=reset_index,
+        )
+    if isinstance(table_data, PolarsDataFrame) or isinstance(table_data, LazyFrame):
+        return convert_polars_to_pandas(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+            reset_index=reset_index,
+        )
+    raise NgioValueError(f"Unsupported table type: {type(table_data)}")
+
+
+def convert_to_polars(
+    table_data: TabularData,
+    index_key: str | None = None,
+    index_type: Literal["int", "str"] | None = None,
+) -> LazyFrame:
+    """Convert a table to a polars LazyFrame.
+
+    Args:
+        table_data (TabularData): The table to convert.
+        index_key (str | None): The column name to use as the index of the DataFrame.
+            Default is None.
+        index_type (str | None): The type of the index column in the DataFrame.
+            Either 'str' or 'int'. Default is None.
+
+    Returns:
+        LazyFrame: Converted polars LazyFrame.
+    """
+    if isinstance(table_data, PolarsDataFrame) or isinstance(table_data, LazyFrame):
+        return normalize_polars_lf(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+        )
+    if isinstance(table_data, DataFrame):
+        return convert_pandas_to_polars(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+        )
+    if isinstance(table_data, AnnData):
+        return convert_anndata_to_polars(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+        )
+    raise NgioValueError(f"Unsupported table type: {type(table_data)}")