ngio 0.5.0b6__py3-none-any.whl

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

@@ -0,0 +1,226 @@
+"""Protocol for table backends handlers."""
+
+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 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 set_group_handler(
+        self,
+        group_handler: ZarrGroupHandler,
+        index_key: str | None = None,
+        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:
+        """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:
+        """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_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_as_pandas_df(self) -> DataFrame:
+        """Load the table as a pandas DataFrame."""
+        ...
+
+    def load_as_polars_lf(self) -> LazyFrame:
+        """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."""
+        ...
+
+    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:
+    """A class to manage the available table backends."""
+
+    _instance = None
+    _implemented_backends: dict[str, type[TableBackendProtocol]]
+
+    def __new__(cls):
+        """Create a new instance of the class if it does not exist."""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._implemented_backends = {}
+        return cls._instance
+
+    @property
+    def available_backends(self) -> list[str]:
+        """Return the available table backends."""
+        return list(self._implemented_backends.keys())
+
+    def get_backend(
+        self,
+        *,
+        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 not in self._implemented_backends:
+            raise NgioValueError(f"Table backend {backend_name} not implemented.")
+        backend = self._implemented_backends[backend_name]()
+        backend.set_group_handler(
+            group_handler=group_handler, index_key=index_key, index_type=index_type
+        )
+        return backend
+
+    def _add_backend(
+        self,
+        table_backend: type[TableBackendProtocol],
+        name: str,
+        overwrite: bool = False,
+    ) -> None:
+        """Register a new handler."""
+        if name in self._implemented_backends and not overwrite:
+            raise NgioValueError(
+                f"Table backend {name} already implemented. "
+                "Use the `overwrite=True` parameter to overwrite it."
+            )
+        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(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"