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

ngio/tables/backends/_abstract_backend.py

@@ -5,15 +5,13 @@ from anndata import AnnData
 from pandas import DataFrame
 from polars import DataFrame as PolarsDataFrame
 from polars import LazyFrame
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 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,
+    TabularData,
+    convert_to_anndata,
+    convert_to_pandas,
+    convert_to_polars,
 )
 from ngio.utils import NgioValueError, ZarrGroupHandler
 
@@ -21,29 +19,30 @@ from ngio.utils import NgioValueError, ZarrGroupHandler
 class BackendMeta(BaseModel):
     """Metadata for the backend."""
 
-    backend: str | None = None
+    backend: str = "anndata"
     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"] | None = None,
-    ):
-        """Initialize the handler.
+    ) -> None:
+        """Attach a group handler to the backend.
 
-        This is a base class for the table backends protocol.
+        Index keys and index types are used to ensure that the
+        serialization and deserialization of the table
+        is consistent across different backends.
 
-        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.
+        Making sure that this is consistent is
+        a duty of the backend implementations.
         """
         self._group_handler = group_handler
         self._index_key = index_key
@@ -67,7 +66,11 @@ class AbstractTableBackend(ABC):
         """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`.
@@ -80,7 +83,7 @@ class AbstractTableBackend(ABC):
         """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`.
@@ -93,7 +96,7 @@ class AbstractTableBackend(ABC):
         """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`.
@@ -122,6 +125,16 @@ class AbstractTableBackend(ABC):
             )
         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.
 
@@ -129,70 +142,35 @@ class AbstractTableBackend(ABC):
         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."
-            )
+        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
         """
-        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."
-            )
+        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
         """
-        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."
-            )
+        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."""
@@ -230,7 +208,7 @@ class AbstractTableBackend(ABC):
 
     def write(
         self,
-        table: DataFrame | AnnData | PolarsDataFrame | LazyFrame,
+        table_data: TabularData,
         metadata: dict | None = None,
         mode: Literal["pandas", "anndata", "polars"] | None = None,
     ) -> None:
@@ -240,15 +218,15 @@ class AbstractTableBackend(ABC):
         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)
+        if mode == "pandas" or isinstance(table_data, DataFrame):
+            self.write_from_pandas(table_data)  # type: ignore[arg-type]
+        elif mode == "anndata" or isinstance(table_data, AnnData):
+            self.write_from_anndata(table_data)  # type: ignore[arg-type]
+        elif mode == "polars" or isinstance(table_data, PolarsDataFrame | LazyFrame):
+            self.write_from_polars(table_data)
         else:
             raise NgioValueError(
-                f"Unsupported table type {type(table)}. "
+                f"Unsupported table type {type(table_data)}. "
                 "Please specify the mode explicitly. "
                 "Supported serialization modes are: "
                 "'pandas', 'anndata', 'polars'."

ngio/tables/backends/{_anndata_v1.py → _anndata.py}

@@ -21,7 +21,7 @@ class AnnDataBackend(AbstractTableBackend):
     @staticmethod
     def backend_name() -> str:
         """Return the name of the backend."""
-        return "anndata_v1"
+        return "anndata"
 
     @staticmethod
     def implements_anndata() -> bool:
@@ -44,6 +44,10 @@ class AnnDataBackend(AbstractTableBackend):
         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_from_anndata(self, table: AnnData) -> None:
         """Serialize the table from an AnnData object."""
         full_url = self._group_handler.full_url

ngio/tables/backends/_csv.py

@@ -0,0 +1,35 @@
+import pandas as pd
+import polars as pl
+
+from ngio.tables.backends._non_zarr_backends import NonZarrBaseBackend
+
+
+def write_lf_to_csv(path: str, table: pl.DataFrame) -> None:
+    """Write a polars DataFrame to a CSV file."""
+    table.write_csv(path)
+
+
+def write_df_to_csv(path: str, table: pd.DataFrame) -> None:
+    """Write a pandas DataFrame to a CSV file."""
+    table.to_csv(path, index=False)
+
+
+class CsvTableBackend(NonZarrBaseBackend):
+    """A class to load and write small tables in CSV format."""
+
+    def __init__(
+        self,
+    ):
+        """Initialize the CsvTableBackend."""
+        super().__init__(
+            lf_reader=pl.scan_csv,
+            df_reader=pd.read_csv,
+            lf_writer=write_lf_to_csv,
+            df_writer=write_df_to_csv,
+            table_name="table.csv",
+        )
+
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "csv"

ngio/tables/backends/{_json_v1.py → _json.py}

@@ -17,7 +17,7 @@ class JsonTableBackend(AbstractTableBackend):
     @staticmethod
     def backend_name() -> str:
         """Return the name of the backend."""
-        return "experimental_json_v1"
+        return "json"
 
     @staticmethod
     def implements_anndata() -> bool:
@@ -61,6 +61,9 @@ class JsonTableBackend(AbstractTableBackend):
         )
         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()

ngio/tables/backends/{_csv_v1.py → _non_zarr_backends.py}

@@ -1,7 +1,7 @@
 import io
+from collections.abc import Callable
+from typing import Any
 
-import pandas as pd
-import polars as pl
 from pandas import DataFrame
 from polars import DataFrame as PolarsDataFrame
 from polars import LazyFrame
@@ -12,15 +12,22 @@ from ngio.tables.backends._utils import normalize_pandas_df, normalize_polars_lf
 from ngio.utils import NgioFileNotFoundError, NgioValueError
 
 
-class CsvTableBackend(AbstractTableBackend):
+class NonZarrBaseBackend(AbstractTableBackend):
     """A class to load and write small tables in CSV format."""
 
-    csv_name = "table.csv"
-
-    @staticmethod
-    def backend_name() -> str:
-        """Return the name of the backend."""
-        return "experimental_csv_v1"
+    def __init__(
+        self,
+        df_reader: Callable[[Any], DataFrame],
+        lf_reader: Callable[[Any], LazyFrame],
+        df_writer: Callable[[str, DataFrame], None],
+        lf_writer: Callable[[str, PolarsDataFrame], None],
+        table_name: str,
+    ):
+        self.df_reader = df_reader
+        self.lf_reader = lf_reader
+        self.df_writer = df_writer
+        self.lf_writer = lf_writer
+        self.table_name = table_name
 
     @staticmethod
     def implements_anndata() -> bool:
@@ -37,38 +44,58 @@ class CsvTableBackend(AbstractTableBackend):
         """Whether the handler implements the polars protocol."""
         return True
 
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        raise NotImplementedError(
+            "The backend_name method must be implemented in the subclass."
+        )
+
     def _load_from_directory_store(self, reader):
         """Load the table from a directory store."""
         url = self._group_handler.full_url
         if url is None:
+            ext = self.table_name.split(".")[-1]
             raise NgioValueError(
-                f"Ngio does not support reading a CSV file from a "
+                f"Ngio does not support reading a {ext} table from a "
                 f"store of type {type(self._group_handler)}. "
                 "Please make sure to use a compatible "
                 "store like a zarr.DirectoryStore."
             )
-        csv_path = f"{url}/{self.csv_name}"
-        dataframe = reader(csv_path)
+        table_path = f"{url}/{self.table_name}"
+        dataframe = reader(table_path)
         return dataframe
 
-    def _load_from_fs_store(self, reader):
+    def _load_from_fs_store_df(self, reader):
         """Load the table from an FS store."""
-        bytes_table = self._group_handler.store.get(self.csv_name)
+        path = self._group_handler.group.path
+        table_path = f"{path}/{self.table_name}"
+        bytes_table = self._group_handler.store.get(table_path)
         if bytes_table is None:
-            raise NgioFileNotFoundError(f"No table found at {self.csv_name}. ")
+            raise NgioFileNotFoundError(f"No table found at {table_path}. ")
         dataframe = reader(io.BytesIO(bytes_table))
         return dataframe
 
+    def _load_from_fs_store_lf(self, reader):
+        """Load the table from an FS store."""
+        full_url = self._group_handler.full_url
+        parquet_path = f"{full_url}/{self.table_name}"
+        store_fs = self._group_handler.store.fs  # type: ignore
+        with store_fs.open(parquet_path, "rb") as f:
+            dataframe = reader(f)
+        return dataframe
+
     def load_as_pandas_df(self) -> DataFrame:
         """Load the table as a pandas DataFrame."""
         store = self._group_handler.store
         if isinstance(store, DirectoryStore):
-            dataframe = self._load_from_directory_store(reader=pd.read_csv)
+            dataframe = self._load_from_directory_store(reader=self.df_reader)
         elif isinstance(store, FSStore):
-            dataframe = self._load_from_fs_store(reader=pd.read_csv)
+            dataframe = self._load_from_fs_store_df(reader=self.df_reader)
         else:
+            ext = self.table_name.split(".")[-1]
             raise NgioValueError(
-                f"Ngio does not support reading a CSV file from a "
+                f"Ngio does not support reading a {ext} table from a "
                 f"store of type {type(store)}. "
                 "Please make sure to use a compatible "
                 "store like a zarr.DirectoryStore or "
@@ -83,16 +110,21 @@ class CsvTableBackend(AbstractTableBackend):
         )
         return dataframe
 
+    def load(self) -> DataFrame:
+        """Load the table as a pandas DataFrame."""
+        return self.load_as_pandas_df()
+
     def load_as_polars_lf(self) -> LazyFrame:
         """Load the table as a polars LazyFrame."""
         store = self._group_handler.store
         if isinstance(store, DirectoryStore):
-            lazy_frame = self._load_from_directory_store(reader=pl.scan_csv)
+            lazy_frame = self._load_from_directory_store(reader=self.lf_reader)
         elif isinstance(store, FSStore):
-            lazy_frame = self._load_from_fs_store(reader=pl.scan_csv)
+            lazy_frame = self._load_from_fs_store_lf(reader=self.lf_reader)
         else:
+            ext = self.table_name.split(".")[-1]
             raise NgioValueError(
-                f"Ngio does not support reading a CSV file from a "
+                f"Ngio does not support reading a {ext} from a "
                 f"store of type {type(store)}. "
                 "Please make sure to use a compatible "
                 "store like a zarr.DirectoryStore or "
@@ -117,16 +149,18 @@ class CsvTableBackend(AbstractTableBackend):
         if isinstance(store, DirectoryStore):
             full_url = self._group_handler.full_url
         else:
+            ext = self.table_name.split(".")[-1]
             raise NgioValueError(
-                f"Ngio does not support writing a CSV file to a "
+                f"Ngio does not support writing a {ext} file to a "
                 f"store of type {type(store)}. "
                 "Please make sure to use a compatible "
                 "store like a zarr.DirectoryStore or "
                 "zarr.FSStore."
             )
         if full_url is None:
+            ext = self.table_name.split(".")[-1]
             raise NgioValueError(
-                f"Ngio does not support writing a CSV file to a "
+                f"Ngio does not support writing a {ext} file to a "
                 f"store of type {type(store)}. "
                 "Please make sure to use a compatible "
                 "store like a zarr.DirectoryStore or "
@@ -143,8 +177,8 @@ class CsvTableBackend(AbstractTableBackend):
             reset_index=True,
         )
         full_url = self._get_store_url()
-        csv_path = f"{full_url}/{self.csv_name}"
-        table.to_csv(csv_path, index=False)
+        table_path = f"{full_url}/{self.table_name}"
+        self.df_writer(table_path, table)
 
     def write_from_polars(self, table: PolarsDataFrame | LazyFrame) -> None:
         """Write the table from a polars DataFrame or LazyFrame."""
@@ -158,5 +192,5 @@ class CsvTableBackend(AbstractTableBackend):
             table = table.collect()
 
         full_url = self._get_store_url()
-        csv_path = f"{full_url}/{self.csv_name}"
-        table.write_csv(csv_path)
+        table_path = f"{full_url}/{self.table_name}"
+        self.lf_writer(table_path, table)

ngio/tables/backends/_parquet.py

@@ -0,0 +1,47 @@
+import pandas as pd
+import polars as pl
+
+from ngio.tables.backends._non_zarr_backends import NonZarrBaseBackend
+
+
+def write_lf_to_parquet(path: str, table: pl.DataFrame) -> None:
+    """Write a polars DataFrame to a Parquet file."""
+    # make categorical into string (for pandas compatibility)
+    schema = table.collect_schema()
+
+    categorical_columns = []
+    for name, dtype in zip(schema.names(), schema.dtypes(), strict=True):
+        if dtype == pl.Categorical:
+            categorical_columns.append(name)
+
+    for col in categorical_columns:
+        table = table.with_columns(pl.col(col).cast(pl.Utf8))
+
+    # write to parquet
+    table.write_parquet(path)
+
+
+def write_df_to_parquet(path: str, table: pd.DataFrame) -> None:
+    """Write a pandas DataFrame to a Parquet file."""
+    table.to_parquet(path, index=False)
+
+
+class ParquetTableBackend(NonZarrBaseBackend):
+    """A class to load and write small tables in Parquet format."""
+
+    def __init__(
+        self,
+    ):
+        """Initialize the ParquetTableBackend."""
+        super().__init__(
+            lf_reader=pl.scan_parquet,
+            df_reader=pd.read_parquet,
+            lf_writer=write_lf_to_parquet,
+            df_writer=write_df_to_parquet,
+            table_name="table.parquet",
+        )
+
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "parquet"

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