ngio 0.2.8__py3-none-any.whl → 0.3.0a0__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_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"] | 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

@@ -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_v1.py

@@ -1,162 +1,35 @@
-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
+from ngio.tables.backends._non_zarr_backends_v1 import NonZarrBaseBackend
 
-    @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 write_lf_to_csv(path: str, table: pl.DataFrame) -> None:
+    """Write a polars DataFrame to a CSV file."""
+    table.write_csv(path)
 
-    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 write_df_to_csv(path: str, table: pd.DataFrame) -> None:
+    """Write a pandas DataFrame to a CSV file."""
+    table.to_csv(path, index=False)
 
-    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)
+class CsvTableBackend(NonZarrBaseBackend):
+    """A class to load and write small tables in CSV format."""
 
-    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,
+    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",
         )
 
-        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)
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "experimental_csv_v1"

ngio/tables/backends/_json_v1.py

@@ -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/_non_zarr_backends_v1.py

@@ -0,0 +1,196 @@
+import io
+from collections.abc import Callable
+from typing import Any
+
+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 NonZarrBaseBackend(AbstractTableBackend):
+    """A class to load and write small tables in CSV format."""
+
+    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:
+        """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
+
+    @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 {ext} table from a "
+                f"store of type {type(self._group_handler)}. "
+                "Please make sure to use a compatible "
+                "store like a zarr.DirectoryStore."
+            )
+        table_path = f"{url}/{self.table_name}"
+        dataframe = reader(table_path)
+        return dataframe
+
+    def _load_from_fs_store_df(self, reader):
+        """Load the table from an FS store."""
+        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 {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=self.df_reader)
+        elif isinstance(store, FSStore):
+            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 {ext} table 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(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=self.lf_reader)
+        elif isinstance(store, FSStore):
+            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 {ext} 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:
+            ext = self.table_name.split(".")[-1]
+            raise NgioValueError(
+                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 {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."
+            )
+        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()
+        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."""
+        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()
+        table_path = f"{full_url}/{self.table_name}"
+        self.lf_writer(table_path, table)

ngio/tables/backends/_parquet_v1.py

@@ -0,0 +1,47 @@
+import pandas as pd
+import polars as pl
+
+from ngio.tables.backends._non_zarr_backends_v1 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 "experimental_parquet_v1"