ngio 0.4.8__py3-none-any.whl → 0.5.0__py3-none-any.whl

ngio/tables/_tables_container.py

@@ -229,10 +229,10 @@ class ImplementedTables:
 
 
 class TablesContainer:
-    """A class to handle the /labels group in an OME-NGFF file."""
+    """A class to handle the /tables group in an OME-NGFF file."""
 
     def __init__(self, group_handler: ZarrGroupHandler) -> None:
-        """Initialize the LabelGroupHandler."""
+        """Initialize the TablesContainer."""
         self._group_handler = group_handler
 
         # Validate the group
@@ -252,17 +252,24 @@ class TablesContainer:
             )
 
     def _get_tables_list(self) -> list[str]:
-        """Create the /tables group if it doesn't exist."""
+        """Return the list of table names from the group attributes."""
         attrs = self._group_handler.load_attrs()
         return attrs.get("tables", [])
 
     def _get_table_group_handler(self, name: str) -> ZarrGroupHandler:
         """Get the group handler for a table."""
-        handler = self._group_handler.derive_handler(path=name)
+        handler = self._group_handler.get_handler(path=name)
         return handler
 
     def list(self, filter_types: TypedTable | str | None = None) -> list[str]:
-        """List all labels in the group."""
+        """List all tables in the group.
+
+        Args:
+            filter_types: If provided, only return tables of this type.
+
+        Returns:
+            A list of table names.
+        """
         tables = self._get_tables_list()
         if filter_types is None:
             return tables
@@ -281,7 +288,16 @@ class TablesContainer:
         backend: TableBackend | None = None,
         strict: bool = True,
     ) -> Table:
-        """Get a label from the group."""
+        """Get a table from the group.
+
+        Args:
+            name: The name of the table.
+            backend: The backend to use for reading the table.
+            strict: If True, raise an error if the table type is not implemented.
+
+        Returns:
+            The table object.
+        """
         if name not in self.list():
             raise NgioValueError(f"Table '{name}' not found in the group.")
 
@@ -301,7 +317,16 @@ class TablesContainer:
         table_cls: type[TableType],
         backend: TableBackend | None = None,
     ) -> TableType:
-        """Get a table from the group as a specific type."""
+        """Get a table from the group as a specific type.
+
+        Args:
+            name: The name of the table.
+            table_cls: The table class to use for loading the table.
+            backend: The backend to use for reading the table.
+
+        Returns:
+            The table object of the specified type.
+        """
         if name not in self.list():
             raise NgioValueError(f"Table '{name}' not found in the group.")
 
@@ -311,6 +336,27 @@ class TablesContainer:
             backend=backend,
         )  # type: ignore[return-value]
 
+    def delete(self, name: str, missing_ok: bool = False) -> None:
+        """Delete a table from the group.
+
+        Args:
+            name (str): The name of the table to delete.
+            missing_ok (bool): If True, do not raise an error if
+                the table does not exist.
+        """
+        existing_tables = self._get_tables_list()
+        if name not in existing_tables:
+            if missing_ok:
+                return
+            raise NgioValueError(
+                f"Table '{name}' not found in the Tables group. "
+                f"Available tables: {existing_tables}"
+            )
+
+        self._group_handler.delete_group(name)
+        existing_tables.remove(name)
+        self._group_handler.write_attrs({"tables": existing_tables})
+
     def add(
         self,
         name: str,
@@ -318,7 +364,14 @@ class TablesContainer:
         backend: TableBackend = DefaultTableBackend,
         overwrite: bool = False,
     ) -> None:
-        """Add a table to the group."""
+        """Add a table to the group.
+
+        Args:
+            name: The name of the table.
+            table: The table object to add.
+            backend: The backend to use for writing the table.
+            overwrite: Whether to overwrite an existing table with the same name.
+        """
         existing_tables = self._get_tables_list()
         if name in existing_tables and not overwrite:
             raise NgioValueError(
@@ -326,9 +379,7 @@ class TablesContainer:
                 "Use overwrite=True to replace it."
             )
 
-        table_handler = self._group_handler.derive_handler(
-            path=name, overwrite=overwrite
-        )
+        table_handler = self._group_handler.get_handler(path=name, overwrite=overwrite)
 
         if backend is None:
             backend = table.backend_name
@@ -360,12 +411,9 @@ def open_tables_container(
     store: StoreOrGroup,
     cache: bool = False,
     mode: AccessModeLiteral = "r+",
-    parallel_safe: bool = False,
 ) -> TablesContainer:
     """Open a table handler from a Zarr store."""
-    handler = ZarrGroupHandler(
-        store=store, cache=cache, mode=mode, parallel_safe=parallel_safe
-    )
+    handler = ZarrGroupHandler(store=store, cache=cache, mode=mode)
     return TablesContainer(handler)
 
 
@@ -374,11 +422,12 @@ def open_table(
     backend: TableBackend | None = None,
     cache: bool = False,
     mode: AccessModeLiteral = "r+",
-    parallel_safe: bool = False,
 ) -> Table:
     """Open a table from a Zarr store."""
     handler = ZarrGroupHandler(
-        store=store, cache=cache, mode=mode, parallel_safe=parallel_safe
+        store=store,
+        cache=cache,
+        mode=mode,
     )
     meta = _get_meta(handler)
     return ImplementedTables().get_table(
@@ -392,11 +441,12 @@ def open_table_as(
     backend: TableBackend | None = None,
     cache: bool = False,
     mode: AccessModeLiteral = "r+",
-    parallel_safe: bool = False,
 ) -> TableType:
     """Open a table from a Zarr store as a specific type."""
     handler = ZarrGroupHandler(
-        store=store, cache=cache, mode=mode, parallel_safe=parallel_safe
+        store=store,
+        cache=cache,
+        mode=mode,
     )
     return table_cls.from_handler(
         handler=handler,
@@ -410,12 +460,20 @@ def write_table(
     backend: TableBackend = DefaultTableBackend,
     cache: bool = False,
     mode: AccessModeLiteral = "a",
-    parallel_safe: bool = False,
 ) -> None:
-    """Write a table to a Zarr store."""
-    handler = ZarrGroupHandler(
-        store=store, cache=cache, mode=mode, parallel_safe=parallel_safe
-    )
+    """Write a table to a Zarr store.
+
+    A table will be created at the given store location.
+
+    Args:
+        store (StoreOrGroup): The Zarr store or group to write the table to.
+        table (Table): The table to write.
+        backend (TableBackend): The backend to use for writing the table.
+        cache (bool): Whether to use caching for the Zarr group handler.
+        mode (AccessModeLiteral): The access mode to use for the Zarr group handler.
+
+    """
+    handler = ZarrGroupHandler(store=store, cache=cache, mode=mode)
     table.set_backend(
         handler=handler,
         backend=backend,

ngio/tables/backends/_abstract_backend.py

@@ -198,6 +198,13 @@ class AbstractTableBackend(ABC):
         if metadata is None:
             metadata = {}
 
+        attrs = self._group_handler.reopen_group().attrs.asdict()
+        # This is required by anndata to identify the format
+        if "encoding-type" in attrs:
+            metadata["encoding-type"] = attrs["encoding-type"]
+        if "encoding-version" in attrs:
+            metadata["encoding-version"] = attrs["encoding-version"]
+
         backend_metadata = BackendMeta(
             backend=self.backend_name(),
             index_key=self.index_key,

ngio/tables/backends/_anndata.py

@@ -1,7 +1,10 @@
+import zarr
 from anndata import AnnData
+from anndata._settings import settings
 from pandas import DataFrame
 from polars import DataFrame as PolarsDataFrame
 from polars import LazyFrame
+from zarr.storage import FsspecStore, LocalStore, MemoryStore
 
 from ngio.tables.backends._abstract_backend import AbstractTableBackend
 from ngio.tables.backends._anndata_utils import (
@@ -12,7 +15,7 @@ from ngio.tables.backends._utils import (
     convert_polars_to_anndata,
     normalize_anndata,
 )
-from ngio.utils import NgioValueError
+from ngio.utils import NgioValueError, copy_group
 
 
 class AnnDataBackend(AbstractTableBackend):
@@ -40,6 +43,7 @@ class AnnDataBackend(AbstractTableBackend):
 
     def load_as_anndata(self) -> AnnData:
         """Load the table as an AnnData object."""
+        settings.zarr_write_format = self._group_handler.zarr_format
         anndata = custom_anndata_read_zarr(self._group_handler._group)
         anndata = normalize_anndata(anndata, index_key=self.index_key)
         return anndata
@@ -48,17 +52,66 @@ class AnnDataBackend(AbstractTableBackend):
         """Load the table as an AnnData object."""
         return self.load_as_anndata()
 
+    def _write_to_local_store(
+        self, store: LocalStore, path: str, table: AnnData
+    ) -> None:
+        """Write the AnnData table to a LocalStore."""
+        store_path = f"{store.root}/{path}"
+        table.write_zarr(store_path)
+
+    def _write_to_fsspec_store(
+        self, store: FsspecStore, path: str, table: AnnData
+    ) -> None:
+        """Write the AnnData table to a FsspecStore."""
+        full_url = f"{store.path}/{path}"
+        fs = store.fs
+        mapper = fs.get_mapper(full_url)
+        table.write_zarr(mapper)
+
+    def _write_to_memory_store(
+        self, store: MemoryStore, path: str, table: AnnData
+    ) -> None:
+        """Write the AnnData table to a MemoryStore."""
+        store = MemoryStore()
+        table.write_zarr(store)
+        anndata_group = zarr.open_group(store, mode="r")
+        copy_group(
+            anndata_group,
+            self._group_handler._group,
+            suppress_warnings=True,
+        )
+
     def write_from_anndata(self, table: AnnData) -> None:
         """Serialize the table from an AnnData object."""
-        full_url = self._group_handler.full_url
-        if full_url is None:
+        # Make sure to use the correct zarr format
+        settings.zarr_write_format = self._group_handler.zarr_format
+        store = self._group_handler.store
+        path = self._group_handler.group.path
+        if isinstance(store, LocalStore):
+            self._write_to_local_store(
+                store,
+                path,
+                table,
+            )
+        elif isinstance(store, FsspecStore):
+            self._write_to_fsspec_store(
+                store,
+                path,
+                table,
+            )
+        elif isinstance(store, MemoryStore):
+            self._write_to_memory_store(
+                store,
+                path,
+                table,
+            )
+        else:
             raise NgioValueError(
-                f"Ngio does not support writing file from a "
-                f"store of type {type(self._group_handler)}. "
+                f"Ngio does not support writing an AnnData table to a "
+                f"store of type {type(store)}. "
                 "Please make sure to use a compatible "
-                "store like a zarr.DirectoryStore."
+                "store like a LocalStore, or FsspecStore."
             )
-        table.write_zarr(full_url)  # type: ignore (AnnData writer requires a str path)
 
     def write_from_pandas(self, table: DataFrame) -> None:
         """Serialize the table from a pandas DataFrame."""

ngio/tables/backends/_anndata_utils.py

@@ -34,10 +34,6 @@ def custom_anndata_read_zarr(
         elem_to_read (Sequence[str] | None): The elements to read from the store.
     """
     group = open_group_wrapper(store=store, mode="r")
-
-    if not isinstance(group.store, zarr.DirectoryStore):
-        elem_to_read = ["X", "obs", "var"]
-
     if elem_to_read is None:
         elem_to_read = [
             "X",
@@ -87,6 +83,8 @@ def custom_anndata_read_zarr(
     if isinstance(group["obs"], zarr.Array):
         _clean_uns(adata)
 
+    if isinstance(adata, dict):
+        adata = AnnData(**adata)  # type: ignore
     if not isinstance(adata, AnnData):
         raise NgioValueError(f"Expected an AnnData object, but got {type(adata)}")
     return adata

ngio/tables/backends/_csv.py

@@ -1,20 +1,7 @@
-import pandas as pd
-import polars as pl
+from ngio.tables.backends._py_arrow_backends import PyArrowBackend
 
-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):
+class CsvTableBackend(PyArrowBackend):
     """A class to load and write small tables in CSV format."""
 
     def __init__(
@@ -22,11 +9,8 @@ class CsvTableBackend(NonZarrBaseBackend):
     ):
         """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",
+            table_format="csv",
         )
 
     @staticmethod

ngio/tables/backends/_json.py

@@ -8,7 +8,7 @@ from ngio.tables.backends._utils import (
     normalize_pandas_df,
     normalize_polars_lf,
 )
-from ngio.utils import NgioFileNotFoundError
+from ngio.utils import NgioError
 
 
 class JsonTableBackend(AbstractTableBackend):
@@ -37,22 +37,19 @@ class JsonTableBackend(AbstractTableBackend):
     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")
+            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_group = self._get_table_group()
-        table_dict = dict(table_group.attrs)
-
-        data_frame = pd.DataFrame.from_dict(table_dict)
-        return data_frame
-
     def load_as_pandas_df(self) -> DataFrame:
         """Load the table as a pandas DataFrame."""
-        data_frame = self._load_as_pandas_df()
+        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,

ngio/tables/backends/_parquet.py

@@ -1,32 +1,7 @@
-import pandas as pd
-import polars as pl
+from ngio.tables.backends._py_arrow_backends import PyArrowBackend
 
-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):
+class ParquetTableBackend(PyArrowBackend):
     """A class to load and write small tables in Parquet format."""
 
     def __init__(
@@ -34,11 +9,8 @@ class ParquetTableBackend(NonZarrBaseBackend):
     ):
         """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",
+            table_format="parquet",
         )
 
     @staticmethod

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)