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

ngio/tables/backends/_csv_v1.py

@@ -0,0 +1,162 @@
+import io
+
+import pandas as pd
+import polars as pl
+from pandas import DataFrame
+from polars import DataFrame as PolarsDataFrame
+from polars import LazyFrame
+from zarr.storage import DirectoryStore, FSStore
+
+from ngio.tables.backends._abstract_backend import AbstractTableBackend
+from ngio.tables.backends._utils import normalize_pandas_df, normalize_polars_lf
+from ngio.utils import NgioFileNotFoundError, NgioValueError
+
+
+class CsvTableBackend(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"
+
+    @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 _load_from_directory_store(self, reader):
+        """Load the table from a directory store."""
+        url = self._group_handler.full_url
+        if url is None:
+            raise NgioValueError(
+                f"Ngio does not support reading a CSV file 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)
+        return dataframe
+
+    def _load_from_fs_store(self, reader):
+        """Load the table from an FS store."""
+        bytes_table = self._group_handler.store.get(self.csv_name)
+        if bytes_table is None:
+            raise NgioFileNotFoundError(f"No table found at {self.csv_name}. ")
+        dataframe = reader(io.BytesIO(bytes_table))
+        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)
+        elif isinstance(store, FSStore):
+            dataframe = self._load_from_fs_store(reader=pd.read_csv)
+        else:
+            raise NgioValueError(
+                f"Ngio does not support reading a CSV file from a "
+                f"store of type {type(store)}. "
+                "Please make sure to use a compatible "
+                "store like a zarr.DirectoryStore or "
+                "zarr.FSStore."
+            )
+
+        dataframe = normalize_pandas_df(
+            dataframe,
+            index_key=self.index_key,
+            index_type=self.index_type,
+            reset_index=False,
+        )
+        return dataframe
+
+    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)
+        elif isinstance(store, FSStore):
+            lazy_frame = self._load_from_fs_store(reader=pl.scan_csv)
+        else:
+            raise NgioValueError(
+                f"Ngio does not support reading a CSV file from a "
+                f"store of type {type(store)}. "
+                "Please make sure to use a compatible "
+                "store like a zarr.DirectoryStore or "
+                "zarr.FSStore."
+            )
+        if not isinstance(lazy_frame, LazyFrame):
+            raise NgioValueError(
+                "Table is not a lazy frame. Please report this issue as an ngio bug."
+                f" {type(lazy_frame)}"
+            )
+
+        lazy_frame = normalize_polars_lf(
+            lazy_frame,
+            index_key=self.index_key,
+            index_type=self.index_type,
+        )
+        return lazy_frame
+
+    def _get_store_url(self) -> str:
+        """Get the store URL."""
+        store = self._group_handler.store
+        if isinstance(store, DirectoryStore):
+            full_url = self._group_handler.full_url
+        else:
+            raise NgioValueError(
+                f"Ngio does not support writing a CSV 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:
+            raise NgioValueError(
+                f"Ngio does not support writing a CSV file to a "
+                f"store of type {type(store)}. "
+                "Please make sure to use a compatible "
+                "store like a zarr.DirectoryStore or "
+                "zarr.FSStore."
+            )
+        return full_url
+
+    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,
+        )
+        full_url = self._get_store_url()
+        csv_path = f"{full_url}/{self.csv_name}"
+        table.to_csv(csv_path, index=False)
+
+    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()
+
+        full_url = self._get_store_url()
+        csv_path = f"{full_url}/{self.csv_name}"
+        table.write_csv(csv_path)

ngio/tables/backends/_json_v1.py

@@ -1,9 +1,13 @@
-from collections.abc import Collection
-
 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 NgioFileNotFoundError
 
 
@@ -12,7 +16,7 @@ class JsonTableBackend(AbstractTableBackend):
 
     @staticmethod
     def backend_name() -> str:
-        """The name of the backend."""
+        """Return the name of the backend."""
         return "experimental_json_v1"
 
     @staticmethod
@@ -21,36 +25,68 @@ class JsonTableBackend(AbstractTableBackend):
         return False
 
     @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 _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")
         except NgioFileNotFoundError:
             table_group = self._group_handler.group.create_group("table")
         return table_group
 
-    def load_as_dataframe(self, columns: Collection[str] | None = None) -> DataFrame:
-        """List all labels in the group."""
+    def _load_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame."""
         table_group = self._get_table_group()
         table_dict = dict(table_group.attrs)
+
         data_frame = pd.DataFrame.from_dict(table_dict)
-        if columns is not None:
-            data_frame = data_frame[columns]
         return data_frame
 
-    def write_from_dataframe(
-        self, table: DataFrame, metadata: dict | None = None
-    ) -> None:
-        """Consolidate the metadata in the store."""
+    def load_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame."""
+        data_frame = self._load_as_pandas_df()
+        data_frame = normalize_pandas_df(
+            data_frame,
+            index_key=self.index_key,
+            index_type=self.index_type,
+            reset_index=False,
+        )
+        return data_frame
+
+    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.to_dict())
-        if metadata is not None:
-            self._group_handler.write_attrs(metadata)
+        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/_table_backends.py

@@ -1,12 +1,14 @@
 """Protocol for table backends handlers."""
 
-from collections.abc import Collection
 from typing import Literal, Protocol
 
 from anndata import AnnData
 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.utils import NgioValueError, ZarrGroupHandler
 
@@ -16,33 +18,110 @@ class TableBackendProtocol(Protocol):
         self,
         group_handler: ZarrGroupHandler,
         index_key: str | None = None,
-        index_type: Literal["int", "str"] = "int",
-    ): ...
+        index_type: Literal["int", "str"] | None = None,
+    ):
+        """Backend constructor.
+
+        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.
+        """
+        ...
+
+    @staticmethod
+    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.
+        """
+        ...
 
     @staticmethod
-    def backend_name() -> str: ...
+    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.
+
+        If this is False, these methods should raise a
+        `NotImplementedError`.
+        """
+        ...
 
     @staticmethod
-    def implements_anndata() -> bool: ...
+    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`.
+        """
+        ...
 
     @staticmethod
-    def implements_dataframe() -> bool: ...
+    def implements_polars() -> bool:
+        """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.
+
+        If this is False, these methods should raise a
+        `NotImplementedError`.
+        """
+        ...
+
+    def load_as_anndata(self) -> AnnData:
+        """Load the table as an AnnData object."""
+        ...
+
+    def load_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame."""
+        ...
 
-    def load_columns(self) -> list[str]: ...
+    def load_as_polars_lf(self) -> LazyFrame:
+        """Load the table as a polars LazyFrame."""
+        ...
+
+    def write_from_pandas(self, table: DataFrame) -> None:
+        """Serialize the table from a pandas DataFrame."""
+        ...
+
+    def write_from_anndata(self, table: AnnData) -> None:
+        """Serialize the table from an AnnData object."""
+        ...
+
+    def write_from_polars(self, table: LazyFrame | PolarsDataFrame) -> None:
+        """Serialize the table from a polars DataFrame or LazyFrame."""
+        ...
+
+    def write(
+        self,
+        table: DataFrame | AnnData | PolarsDataFrame | LazyFrame,
+        metadata: dict[str, str] | None = None,
+        mode: Literal["pandas", "anndata", "polars"] | None = None,
+    ) -> None:
+        """This is a generic write method.
 
-    def load_as_anndata(self, columns: Collection[str] | None = None) -> AnnData: ...
+        Will call the appropriate write method
+        depending on the type of the table.
 
-    def load_as_dataframe(
-        self, columns: Collection[str] | None = None
-    ) -> DataFrame: ...
+        Moreover it will also write the metadata
+        if provided, and the backend methadata
 
-    def write_from_dataframe(
-        self, table: DataFrame, metadata: dict | None = None
-    ) -> None: ...
+        the backend should write in the zarr group attributes
+            - backend: the backend name (self.backend_name())
+            - index_key: the index key
+            - index_type: the index type
 
-    def write_from_anndata(
-        self, table: AnnData, metadata: dict | None = None
-    ) -> None: ...
+        """
 
 
 class ImplementedTableBackends:
@@ -68,7 +147,7 @@ class ImplementedTableBackends:
         backend_name: str | None,
         group_handler: ZarrGroupHandler,
         index_key: str | None = None,
-        index_type: Literal["int", "str"] = "int",
+        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:
@@ -100,3 +179,4 @@ class ImplementedTableBackends:
 
 ImplementedTableBackends().add_backend(AnnDataBackend)
 ImplementedTableBackends().add_backend(JsonTableBackend)
+ImplementedTableBackends().add_backend(CsvTableBackend)