ngio 0.5.0a2__py3-none-any.whl → 0.5.0b1__py3-none-any.whl

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)

ngio/tables/v1/_roi_table.py

@@ -4,6 +4,7 @@ This class follows the roi_table specification at:
 https://fractal-analytics-platform.github.io/fractal-tasks-core/tables/
 """
 
+import warnings
 from collections.abc import Iterable
 from typing import Literal
 from uuid import uuid4
@@ -26,7 +27,6 @@ from ngio.utils import (
     NgioTableValidationError,
     NgioValueError,
     ZarrGroupHandler,
-    ngio_warn,
 )
 
 REQUIRED_COLUMNS = [
@@ -77,7 +77,9 @@ OPTIONAL_COLUMNS = ORIGIN_COLUMNS + TRANSLATION_COLUMNS + PLATE_COLUMNS + INDEX_
 def _check_optional_columns(col_name: str) -> None:
     """Check if the column name is in the optional columns."""
     if col_name not in OPTIONAL_COLUMNS + TIME_COLUMNS:
-        ngio_warn(f"Column {col_name} is not in the optional columns.")
+        warnings.warn(
+            f"Column {col_name} is not in the optional columns.", stacklevel=2
+        )
 
 
 def _dataframe_to_rois(
@@ -120,17 +122,17 @@ def _dataframe_to_rois(
         else:
             label = getattr(row, "label", None)
 
-        roi = Roi(
+        slices = {
+            "x": (row.x_micrometer, row.len_x_micrometer),
+            "y": (row.y_micrometer, row.len_y_micrometer),
+            "z": (z_micrometer, z_length_micrometer),
+        }
+        if t_second is not None or t_length_second is not None:
+            slices["t"] = (t_second, t_length_second)
+        roi = Roi.from_values(
             name=str(row.Index),
-            x=row.x_micrometer,  # type: ignore (type can not be known here)
-            y=row.y_micrometer,  # type: ignore (type can not be known here)
-            z=z_micrometer,
-            t=t_second,
-            x_length=row.len_x_micrometer,  # type: ignore (type can not be known here)
-            y_length=row.len_y_micrometer,  # type: ignore (type can not be known here)
-            z_length=z_length_micrometer,
-            t_length=t_length_second,
-            unit="micrometer",
+            slices=slices,
+            space="world",
             label=label,
             **extras,
         )
@@ -143,24 +145,39 @@ def _rois_to_dataframe(rois: dict[str, Roi], index_key: str | None) -> pd.DataFr
     data = []
     for roi in rois.values():
         # This normalization is necessary for backward compatibility
-        z_micrometer = roi.z if roi.z is not None else 0.0
-        len_z_micrometer = roi.z_length if roi.z_length is not None else 1.0
+        if roi.space != "world":
+            raise NotImplementedError(
+                "Only ROIs in world coordinates can be serialized."
+            )
 
+        z_slice = roi.get("z")
+        if z_slice is None:
+            z_micrometer = 0.0
+            len_z_micrometer = 1.0
+        else:
+            z_micrometer = z_slice.start if z_slice.start is not None else 0.0
+            len_z_micrometer = z_slice.length if z_slice.length is not None else 1.0
+
+        x_slice = roi.get("x")
+        if x_slice is None:
+            raise NgioValueError("ROI is missing 'x' slice.")
+        y_slice = roi.get("y")
+        if y_slice is None:
+            raise NgioValueError("ROI is missing 'y' slice.")
         row = {
             index_key: roi.get_name(),
-            "x_micrometer": roi.x,
-            "y_micrometer": roi.y,
+            "x_micrometer": x_slice.start if x_slice.start is not None else 0.0,
+            "y_micrometer": y_slice.start if y_slice.start is not None else 0.0,
             "z_micrometer": z_micrometer,
-            "len_x_micrometer": roi.x_length,
-            "len_y_micrometer": roi.y_length,
+            "len_x_micrometer": x_slice.length if x_slice.length is not None else 1.0,
+            "len_y_micrometer": y_slice.length if y_slice.length is not None else 1.0,
             "len_z_micrometer": len_z_micrometer,
         }
 
-        if roi.t is not None:
-            row["t_second"] = roi.t
-
-        if roi.t_length is not None:
-            row["len_t_second"] = roi.t_length
+        t_slice = roi.get("t")
+        if t_slice is not None:
+            row["t_second"] = t_slice.start if t_slice.start is not None else 0.0
+            row["len_t_second"] = t_slice.length if t_slice.length is not None else 1.0
 
         if roi.label is not None and index_key != "label":
             row["label"] = roi.label
@@ -183,7 +200,7 @@ class RoiDictWrapper:
         self._rois_by_name = {}
         self._rois_by_label = {}
         for roi in rois:
-            name = roi.get_name()
+            name = roi.name
             if name in self._rois_by_name:
                 name = f"{name}_{uuid4().hex[:8]}"
             self._rois_by_name[name] = roi

ngio/utils/__init__.py

@@ -1,13 +1,12 @@
 """Various utilities for the ngio package."""
 
-import os
-
 from ngio.utils._datasets import (
     download_ome_zarr_dataset,
     list_ome_zarr_datasets,
     print_datasets_infos,
 )
 from ngio.utils._errors import (
+    NgioError,
     NgioFileExistsError,
     NgioFileNotFoundError,
     NgioTableValidationError,
@@ -15,22 +14,19 @@ from ngio.utils._errors import (
     NgioValueError,
 )
 from ngio.utils._fractal_fsspec_store import fractal_fsspec_store
-from ngio.utils._logger import ngio_logger, ngio_warn, set_logger_level
 from ngio.utils._zarr_utils import (
     AccessModeLiteral,
     NgioCache,
     StoreOrGroup,
     ZarrGroupHandler,
+    copy_group,
     open_group_wrapper,
 )
 
-set_logger_level(os.getenv("NGIO_LOGGER_LEVEL", "WARNING"))
-
 __all__ = [
-    # Zarr
     "AccessModeLiteral",
     "NgioCache",
-    # Errors
+    "NgioError",
     "NgioFileExistsError",
     "NgioFileNotFoundError",
     "NgioTableValidationError",
@@ -38,14 +34,10 @@ __all__ = [
     "NgioValueError",
     "StoreOrGroup",
     "ZarrGroupHandler",
-    # Other
+    "copy_group",
     "download_ome_zarr_dataset",
     "fractal_fsspec_store",
     "list_ome_zarr_datasets",
-    # Logger
-    "ngio_logger",
-    "ngio_warn",
     "open_group_wrapper",
     "print_datasets_infos",
-    "set_logger_level",
 ]