ngio 0.2.0a2__py3-none-any.whl → 0.5.0b4__py3-none-any.whl

ngio/tables/backends/_json.py

@@ -0,0 +1,92 @@
+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 NgioError
+
+
+class JsonTableBackend(AbstractTableBackend):
+    """A class to load and write small tables in the zarr group .attrs (json)."""
+
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "json"
+
+    @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 _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", create_mode=True)
+        except NgioError as e:
+            raise NgioError(
+                "Could not get or create a 'table' group in the store "
+                f"{self._group_handler.store} path "
+                f"{self._group_handler.group.path}/table."
+            ) from e
+        return table_group
+
+    def load_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame."""
+        table_dict = self._get_table_group().attrs.asdict()
+        data_frame = pd.DataFrame.from_dict(table_dict)
+        data_frame = normalize_pandas_df(
+            data_frame,
+            index_key=self.index_key,
+            index_type=self.index_type,
+            reset_index=False,
+        )
+        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()
+        table_group.attrs.clear()
+        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/_parquet.py

@@ -0,0 +1,19 @@
+from ngio.tables.backends._py_arrow_backends import PyArrowBackend
+
+
+class ParquetTableBackend(PyArrowBackend):
+    """A class to load and write small tables in Parquet format."""
+
+    def __init__(
+        self,
+    ):
+        """Initialize the ParquetTableBackend."""
+        super().__init__(
+            table_name="table.parquet",
+            table_format="parquet",
+        )
+
+    @staticmethod
+    def backend_name() -> str:
+        """Return the name of the backend."""
+        return "parquet"

ngio/tables/backends/_py_arrow_backends.py

@@ -0,0 +1,222 @@
+from typing import Literal
+
+import polars as pl
+import pyarrow as pa
+import pyarrow.csv as pa_csv
+import pyarrow.dataset as pa_ds
+import pyarrow.fs as pa_fs
+import pyarrow.parquet as pa_parquet
+from pandas import DataFrame
+from polars import DataFrame as PolarsDataFrame
+from polars import LazyFrame
+from zarr.storage import FsspecStore, LocalStore, MemoryStore, ZipStore
+
+from ngio.tables.backends._abstract_backend import AbstractTableBackend
+from ngio.tables.backends._utils import normalize_pandas_df, normalize_polars_lf
+from ngio.utils import NgioValueError
+from ngio.utils._zarr_utils import _make_sync_fs
+
+
+class PyArrowBackend(AbstractTableBackend):
+    """A class to load and write small tables in CSV format."""
+
+    def __init__(
+        self,
+        table_name: str,
+        table_format: Literal["csv", "parquet"] = "parquet",
+    ):
+        self.table_name = table_name
+        self.table_format = table_format
+
+    @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 _raise_store_type_not_supported(self):
+        """Raise an error for unsupported store types."""
+        ext = self.table_name.split(".")[-1]
+        store = self._group_handler.store
+        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 LocalStore, or "
+            "FsspecStore, or MemoryStore, or ZipStore."
+        )
+
+    def _load_from_local_store(self, store: LocalStore, path: str) -> pa_ds.Dataset:
+        """Load the table from a directory store."""
+        root_path = store.root
+        table_path = f"{root_path}/{path}/{self.table_name}"
+        dataset = pa_ds.dataset(table_path, format=self.table_format)
+        return dataset
+
+    def _load_from_fsspec_store(self, store: FsspecStore, path: str) -> pa_ds.Dataset:
+        """Load the table from an FS store."""
+        table_path = f"{store.path}/{path}/{self.table_name}"
+        fs = _make_sync_fs(store.fs)
+        dataset = pa_ds.dataset(table_path, format=self.table_format, filesystem=fs)
+        return dataset
+
+    def _load_from_in_memory_store(
+        self, store: MemoryStore, path: str
+    ) -> pa_ds.Dataset:
+        """Load the table from an in-memory store."""
+        table_path = f"{path}/{self.table_name}"
+        table = store._store_dict.get(table_path, None)
+        if table is None:
+            raise NgioValueError(
+                f"Table {self.table_name} not found in the in-memory store at "
+                f"path {path}."
+            )
+        assert isinstance(table, pa.Table)
+        dataset = pa_ds.dataset(table)
+        return dataset
+
+    def _load_from_zip_store(self, store: ZipStore, path: str) -> pa_ds.Dataset:
+        """Load the table from a zip store."""
+        raise NotImplementedError("Zip store loading is not implemented yet.")
+
+    def _load_pyarrow_dataset(self) -> pa_ds.Dataset:
+        """Load the table as a pyarrow Dataset."""
+        store = self._group_handler.store
+        path = self._group_handler.group.path
+        if isinstance(store, LocalStore):
+            return self._load_from_local_store(store, path)
+        elif isinstance(store, FsspecStore):
+            return self._load_from_fsspec_store(store, path)
+        elif isinstance(store, MemoryStore):
+            return self._load_from_in_memory_store(store, path)
+        elif isinstance(store, ZipStore):
+            return self._load_from_zip_store(store, path)
+        self._raise_store_type_not_supported()
+
+    def load_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame."""
+        dataset = self._load_pyarrow_dataset()
+        dataframe = dataset.to_table().to_pandas()
+        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."""
+        dataset = self._load_pyarrow_dataset()
+        lazy_frame = pl.scan_pyarrow_dataset(dataset)
+        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 _write_to_stream(self, stream, table: pa.Table) -> None:
+        """Write the table to a stream."""
+        if self.table_format == "parquet":
+            pa_parquet.write_table(table, stream)
+        elif self.table_format == "csv":
+            pa_csv.write_csv(table, stream)
+        else:
+            raise NgioValueError(
+                f"Unsupported table format: {self.table_format}. "
+                "Supported formats are 'parquet' and 'csv'."
+            )
+
+    def _write_to_local_store(
+        self, store: LocalStore, path: str, table: pa.Table
+    ) -> None:
+        """Write the table to a directory store."""
+        root_path = store.root
+        table_path = f"{root_path}/{path}/{self.table_name}"
+        self._write_to_stream(table_path, table)
+
+    def _write_to_fsspec_store(
+        self, store: FsspecStore, path: str, table: pa.Table
+    ) -> None:
+        """Write the table to an FS store."""
+        table_path = f"{store.path}/{path}/{self.table_name}"
+        fs = _make_sync_fs(store.fs)
+        fs = pa_fs.PyFileSystem(pa_fs.FSSpecHandler(fs))
+        with fs.open_output_stream(table_path) as out_stream:
+            self._write_to_stream(out_stream, table)
+
+    def _write_to_in_memory_store(
+        self, store: MemoryStore, path: str, table: pa.Table
+    ) -> None:
+        """Write the table to an in-memory store."""
+        table_path = f"{path}/{self.table_name}"
+        store._store_dict[table_path] = table
+
+    def _write_to_zip_store(self, store: ZipStore, path: str, table: pa.Table) -> None:
+        """Write the table to a zip store."""
+        raise NotImplementedError("Writing to zip store is not implemented yet.")
+
+    def _write_pyarrow_dataset(self, dataset: pa.Table) -> None:
+        """Write the table from a pyarrow Dataset."""
+        store = self._group_handler.store
+        path = self._group_handler.group.path
+        if isinstance(store, LocalStore):
+            return self._write_to_local_store(store=store, path=path, table=dataset)
+        elif isinstance(store, FsspecStore):
+            return self._write_to_fsspec_store(store=store, path=path, table=dataset)
+        elif isinstance(store, MemoryStore):
+            return self._write_to_in_memory_store(store=store, path=path, table=dataset)
+        elif isinstance(store, ZipStore):
+            return self._write_to_zip_store(store=store, path=path, table=dataset)
+        self._raise_store_type_not_supported()
+
+    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 = pa.Table.from_pandas(table, preserve_index=False)
+        self._write_pyarrow_dataset(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()
+        table = table.to_arrow()
+        self._write_pyarrow_dataset(table)

ngio/tables/backends/_table_backends.py

@@ -1,48 +1,147 @@
 """Protocol for table backends handlers."""
 
-from collections.abc import Collection
 from typing import Literal, Protocol
 
 from anndata import AnnData
 from pandas import DataFrame
-
-from ngio.tables.backends._anndata_v1 import AnnDataBackend
-from ngio.tables.backends._json_v1 import JsonTableBackend
+from polars import DataFrame as PolarsDataFrame
+from polars import LazyFrame
+
+from ngio.tables.backends._anndata import AnnDataBackend, AnnDataBackendV1
+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"] = "int",
-    ): ...
+        index_type: Literal["int", "str"] | None = None,
+    ) -> 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
+        is consistent across different backends.
+
+        Making sure that this is consistent is
+        a duty of the backend implementations.
+        """
+        ...
 
     @staticmethod
-    def backend_name() -> str: ...
+    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.
+        """
+        ...
+
+    @property
+    def group_handler(self) -> ZarrGroupHandler:
+        """Return the group handler."""
+        ...
 
     @staticmethod
-    def implements_anndata() -> bool: ...
+    def implements_anndata() -> bool:
+        """Check if the backend implements the anndata protocol.
+
+        If this is True, the backend should implement the
+        `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`.
+        """
+        ...
+
+    @staticmethod
+    def implements_pandas() -> bool:
+        """Check if the backend implements the pandas protocol.
+
+        If this is True, the backend should implement the
+        `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
+        `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_columns(self) -> list[str]: ...
+    def load_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame."""
+        ...
 
-    def load_as_anndata(self, columns: Collection[str] | None = None) -> AnnData: ...
+    def load_as_polars_lf(self) -> LazyFrame:
+        """Load the table as a polars LazyFrame."""
+        ...
 
-    def load_as_dataframe(
-        self, columns: Collection[str] | None = None
-    ) -> DataFrame: ...
+    def load(self) -> TabularData:
+        """The default load method.
 
-    def write_from_dataframe(
-        self, table: DataFrame, metadata: dict | None = None
-    ) -> None: ...
+        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_anndata(
-        self, table: AnnData, metadata: dict | None = None
-    ) -> None: ...
+    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_data: DataFrame | AnnData | PolarsDataFrame | LazyFrame,
+        metadata: dict[str, str] | None = None,
+    ) -> None:
+        """This is a generic write method.
+
+        Will call the appropriate write method
+        depending on the type of the table.
+
+        Moreover it will also write the metadata
+        if provided, and the backend methadata
+
+        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
+
+        """
 
 
 class ImplementedTableBackends:
@@ -65,38 +164,63 @@ 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"] = "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:
-            # 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(
+    def _add_backend(
         self,
-        table_beckend: type[TableBackendProtocol],
+        table_backend: type[TableBackendProtocol],
+        name: str,
         overwrite: bool = False,
-    ):
+    ) -> None:
         """Register a new handler."""
-        backend_name = table_beckend.backend_name()
-        if backend_name in self._implemented_backends and not overwrite:
+        if name in self._implemented_backends and not overwrite:
             raise NgioValueError(
-                f"Table backend {backend_name} already implemented. "
+                f"Table backend {name} already implemented. "
                 "Use the `overwrite=True` parameter to overwrite it."
             )
-        self._implemented_backends[backend_name] = table_beckend
+        self._implemented_backends[name] = table_backend
+
+    def add_backend(
+        self,
+        table_backend: type[TableBackendProtocol],
+        overwrite: bool = False,
+        aliases: list[str] | None = None,
+    ) -> None:
+        """Register a new handler."""
+        self._add_backend(
+            table_backend=table_backend,
+            name=table_backend.backend_name(),
+            overwrite=overwrite,
+        )
+        if aliases is not None:
+            for alias in aliases:
+                self._add_backend(
+                    table_backend=table_backend, name=alias, overwrite=overwrite
+                )
 
 
 ImplementedTableBackends().add_backend(AnnDataBackend)
-ImplementedTableBackends().add_backend(JsonTableBackend)
+ImplementedTableBackends().add_backend(AnnDataBackendV1)
+ImplementedTableBackends().add_backend(
+    JsonTableBackend, aliases=["experimental_json_v1"]
+)
+ImplementedTableBackends().add_backend(CsvTableBackend, aliases=["experimental_csv_v1"])
+ImplementedTableBackends().add_backend(
+    ParquetTableBackend, aliases=["experimental_parquet_v1"]
+)
+
+TableBackend = Literal["anndata", "json", "csv", "parquet"] | str | TableBackendProtocol
+DefaultTableBackend = "anndata_v1"