ngio 0.2.7__py3-none-any.whl → 0.3.0a0__py3-none-any.whl

ngio/tables/backends/_table_backends.py

@@ -10,17 +10,19 @@ from polars import LazyFrame
 from ngio.tables.backends._anndata_v1 import AnnDataBackend
 from ngio.tables.backends._csv_v1 import CsvTableBackend
 from ngio.tables.backends._json_v1 import JsonTableBackend
+from ngio.tables.backends._parquet_v1 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"] | None = None,
-    ):
-        """Backend constructor.
+    ) -> 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
@@ -42,12 +44,21 @@ class TableBackendProtocol(Protocol):
         """
         ...
 
+    @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
-        `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`.
@@ -59,7 +70,7 @@ class TableBackendProtocol(Protocol):
         """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`.
@@ -71,7 +82,7 @@ class TableBackendProtocol(Protocol):
         """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`.
@@ -90,6 +101,16 @@ class TableBackendProtocol(Protocol):
         """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."""
         ...
@@ -104,7 +125,7 @@ class TableBackendProtocol(Protocol):
 
     def write(
         self,
-        table: DataFrame | AnnData | PolarsDataFrame | LazyFrame,
+        table_data: DataFrame | AnnData | PolarsDataFrame | LazyFrame,
         metadata: dict[str, str] | None = None,
         mode: Literal["pandas", "anndata", "polars"] | None = None,
     ) -> None:
@@ -144,23 +165,20 @@ class ImplementedTableBackends:
 
     def get_backend(
         self,
-        backend_name: str | None,
+        *,
         group_handler: ZarrGroupHandler,
+        backend_name: str = "anndata_v1",
         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 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(
         self,
@@ -180,3 +198,4 @@ class ImplementedTableBackends:
 ImplementedTableBackends().add_backend(AnnDataBackend)
 ImplementedTableBackends().add_backend(JsonTableBackend)
 ImplementedTableBackends().add_backend(CsvTableBackend)
+ImplementedTableBackends().add_backend(ParquetTableBackend)

ngio/tables/backends/_utils.py

@@ -9,7 +9,6 @@ These functions are used to validate and normalize the tables
 to ensure that conversion between formats is consistent.
 """
 
-# %%
 from copy import deepcopy
 from typing import Literal
 
@@ -24,6 +23,8 @@ from polars import LazyFrame
 
 from ngio.utils import NgioTableValidationError, NgioValueError
 
+TabularData = AnnData | DataFrame | PolarsDataFrame | LazyFrame
+
 # -----------------
 # Validation utils
 # -----------------
@@ -460,3 +461,148 @@ def convert_polars_to_anndata(
         pandas_df,
         index_key=index_key,
     )
+
+
+# -----------------
+# Conversion functions
+# -----------------
+
+
+def normalize_table(
+    table_data: TabularData,
+    index_key: str | None = None,
+    index_type: Literal["int", "str"] | None = None,
+) -> TabularData:
+    """Normalize a table to a specific format.
+
+    Args:
+        table_data (TabularData): The table to normalize.
+        index_key (str | None): The column name to use as the index of the DataFrame.
+            Default is None.
+        index_type (str | None): The type of the index column in the DataFrame.
+            Either 'str' or 'int'. Default is None.
+
+    Returns:
+        DataFrame | AnnData | PolarsDataFrame | LazyFrame: Normalized table.
+    """
+    if isinstance(table_data, DataFrame):
+        return normalize_pandas_df(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+            reset_index=False,
+        )
+    if isinstance(table_data, AnnData):
+        return normalize_anndata(table_data, index_key=index_key)
+    if isinstance(table_data, PolarsDataFrame) or isinstance(table_data, LazyFrame):
+        return normalize_polars_lf(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+        )
+    raise NgioValueError(f"Unsupported table type: {type(table_data)}")
+
+
+def convert_to_anndata(
+    table_data: TabularData,
+    index_key: str | None = None,
+) -> AnnData:
+    """Convert a table to an AnnData object.
+
+    Args:
+        table_data (TabularData): The table to convert.
+        index_key (str | None): The column name to use as the index of the DataFrame.
+            Default is None.
+
+    Returns:
+        AnnData: Converted AnnData object.
+    """
+    if isinstance(table_data, AnnData):
+        return normalize_anndata(table_data, index_key=index_key)
+    if isinstance(table_data, DataFrame):
+        return convert_pandas_to_anndata(table_data, index_key=index_key)
+    if isinstance(table_data, PolarsDataFrame) or isinstance(table_data, LazyFrame):
+        return convert_polars_to_anndata(table_data, index_key=index_key)
+    raise NgioValueError(f"Unsupported table type: {type(table_data)}")
+
+
+def convert_to_pandas(
+    table_data: TabularData,
+    index_key: str | None = None,
+    index_type: Literal["int", "str"] | None = None,
+    reset_index: bool = False,
+) -> DataFrame:
+    """Convert a table to a pandas DataFrame.
+
+    Args:
+        table_data (TabularData): The table to convert.
+        index_key (str | None): The column name to use as the index of the DataFrame.
+            Default is None.
+        index_type (str | None): The type of the index column in the DataFrame.
+            Either 'str' or 'int'. Default is None.
+        reset_index (bool): If True the index will be reset (i.e., the index will be
+            converted to a column). If False, the index will be kept as is.
+
+    Returns:
+        DataFrame: Converted pandas DataFrame.
+    """
+    if isinstance(table_data, DataFrame):
+        return normalize_pandas_df(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+            reset_index=reset_index,
+        )
+    if isinstance(table_data, AnnData):
+        return convert_anndata_to_pandas(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+            reset_index=reset_index,
+        )
+    if isinstance(table_data, PolarsDataFrame) or isinstance(table_data, LazyFrame):
+        return convert_polars_to_pandas(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+            reset_index=reset_index,
+        )
+    raise NgioValueError(f"Unsupported table type: {type(table_data)}")
+
+
+def convert_to_polars(
+    table_data: TabularData,
+    index_key: str | None = None,
+    index_type: Literal["int", "str"] | None = None,
+) -> LazyFrame:
+    """Convert a table to a polars LazyFrame.
+
+    Args:
+        table_data (TabularData): The table to convert.
+        index_key (str | None): The column name to use as the index of the DataFrame.
+            Default is None.
+        index_type (str | None): The type of the index column in the DataFrame.
+            Either 'str' or 'int'. Default is None.
+
+    Returns:
+        LazyFrame: Converted polars LazyFrame.
+    """
+    if isinstance(table_data, PolarsDataFrame) or isinstance(table_data, LazyFrame):
+        return normalize_polars_lf(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+        )
+    if isinstance(table_data, DataFrame):
+        return convert_pandas_to_polars(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+        )
+    if isinstance(table_data, AnnData):
+        return convert_anndata_to_polars(
+            table_data,
+            index_key=index_key,
+            index_type=index_type,
+        )
+    raise NgioValueError(f"Unsupported table type: {type(table_data)}")

ngio/tables/tables_container.py

@@ -1,10 +1,24 @@
 """Module for handling the /tables group in an OME-NGFF file."""
 
-from typing import Literal, Protocol
+from typing import Literal, Protocol, TypeVar
 
-from ngio.tables.v1 import FeatureTableV1, MaskingRoiTableV1, RoiTableV1
-from ngio.tables.v1._generic_table import GenericTable
-from ngio.tables.v1._roi_table import _GenericRoiTableV1
+import anndata as ad
+import pandas as pd
+import polars as pl
+
+from ngio.tables.backends import (
+    BackendMeta,
+    TableBackendProtocol,
+    TabularData,
+)
+from ngio.tables.v1 import (
+    ConditionTableV1,
+    FeatureTableV1,
+    GenericTable,
+    MaskingRoiTableV1,
+    RoiTableV1,
+)
+from ngio.tables.v1._roi_table import GenericRoiTableV1
 from ngio.utils import (
     AccessModeLiteral,
     NgioValidationError,
@@ -13,22 +27,23 @@ from ngio.utils import (
     ZarrGroupHandler,
 )
 
-GenericRoiTable = _GenericRoiTableV1
+GenericRoiTable = GenericRoiTableV1
 RoiTable = RoiTableV1
 MaskingRoiTable = MaskingRoiTableV1
 FeatureTable = FeatureTableV1
+ConditionTable = ConditionTableV1
 
 
 class Table(Protocol):
     """Placeholder class for a table."""
 
     @staticmethod
-    def type() -> str | None:
+    def table_type() -> str:
         """Return the type of the table."""
         ...
 
     @staticmethod
-    def version() -> str | None:
+    def version() -> str:
         """Return the version of the table."""
         ...
 
@@ -37,19 +52,72 @@ class Table(Protocol):
         """The name of the backend."""
         ...
 
+    @property
+    def meta(self) -> BackendMeta:
+        """Return the metadata for the table."""
+        ...
+
+    @property
+    def dataframe(self) -> pd.DataFrame:
+        """Return the table as a DataFrame."""
+        ...
+
+    @property
+    def lazy_frame(self) -> pl.LazyFrame:
+        """Return the table as a LazyFrame."""
+        ...
+
+    @property
+    def anndata(self) -> ad.AnnData:
+        """Return the table as an AnnData object."""
+        ...
+
+    def set_table_data(
+        self,
+        table_data: TabularData | None = None,
+        refresh: bool = False,
+    ) -> None:
+        """Make sure that the table data is set (exist in memory).
+
+        If an object is passed, it will be used as the table.
+        If None is passed, the table will be loaded from the backend.
+
+        If refresh is True, the table will be reloaded from the backend.
+            If table is not None, this will be ignored.
+        """
+        ...
+
+    def set_backend(
+        self,
+        handler: ZarrGroupHandler | None = None,
+        backend: str | TableBackendProtocol = "anndata_v1",
+    ) -> None:
+        """Set the backend store and path for the table.
+
+        Either a handler or a backend must be provided.
+
+        If the hanlder in none it will be inferred from the backend.
+        If the backend is none, it will be inferred from the group attrs
+        """
+        ...
+
     @classmethod
-    def _from_handler(
-        cls, handler: ZarrGroupHandler, backend_name: str | None = None
+    def from_handler(
+        cls,
+        handler: ZarrGroupHandler,
+        backend: str | TableBackendProtocol | None = None,
     ) -> "Table":
         """Create a new table from a Zarr group handler."""
         ...
 
-    def _set_backend(
-        self,
-        handler: ZarrGroupHandler,
-        backend_name: str | None = None,
-    ) -> None:
-        """Set the backend store and path for the table."""
+    @classmethod
+    def from_table_data(cls, table_data: TabularData, meta: BackendMeta) -> "Table":
+        """Create a new table from a DataFrame."""
+        ...
+
+    @property
+    def table_data(self) -> TabularData:
+        """Return the table."""
         ...
 
     def consolidate(self) -> None:
@@ -58,13 +126,32 @@ class Table(Protocol):
 
 
 TypedTable = Literal[
-    "roi_table", "masking_roi_table", "feature_table", "generic_roi_table"
+    "roi_table",
+    "masking_roi_table",
+    "feature_table",
+    "generic_roi_table",
+    "condition_table",
 ]
 
+TableType = TypeVar("TableType", bound=Table)
+
+
+class TableMeta(BackendMeta):
+    """Base class for table metadata."""
+
+    fractal_table_version: str = "1"
+    type: str = "generic_table"
 
-def _unique_table_name(type_name, version) -> str:
-    """Return the unique name for a table."""
-    return f"{type_name}_v{version}"
+    def unique_name(self) -> str:
+        """Return the unique name for the table."""
+        return f"{self.type}_v{self.fractal_table_version}"
+
+
+def _get_meta(handler: ZarrGroupHandler) -> TableMeta:
+    """Get the metadata from the handler."""
+    attrs = handler.load_attrs()
+    meta = TableMeta(**attrs)
+    return meta
 
 
 class ImplementedTables:
@@ -86,77 +173,38 @@ class ImplementedTables:
 
     def get_table(
         self,
-        type: str,
-        version: str,
+        meta: TableMeta,
         handler: ZarrGroupHandler,
-        backend_name: str | None = None,
+        backend: str | TableBackendProtocol | None = None,
         strict: bool = True,
     ) -> Table:
         """Try to get a handler for the given store based on the metadata version."""
-        _errors = {}
-        for name, table_cls in self._implemented_tables.items():
-            if name != _unique_table_name(type, version):
-                continue
-            try:
-                table = table_cls._from_handler(
-                    handler=handler, backend_name=backend_name
-                )
-                return table
-            except Exception as e:
-                if strict:
-                    raise NgioValidationError(
-                        f"Could not load table {name} from handler. Error: {e}"
-                    ) from e
-                else:
-                    _errors[name] = e
-        # If no table was found, we can try to load the table from a generic table
-        try:
-            table = GenericTable._from_handler(
-                handler=handler, backend_name=backend_name
-            )
-            return table
-        except Exception as e:
-            _errors["generic"] = e
+        if strict:
+            default = None
+        else:
+            default = GenericTable
 
-        if len(_errors) == 0:
-            raise NgioValidationError(
-                f"Could not find a table implementation for {type} v{version}. "
-                f"Available tables: {self.available_implementations()}"
+        table_cls = self._implemented_tables.get(meta.unique_name(), default)
+        if table_cls is None:
+            raise NgioValueError(
+                f"Table handler for {meta.unique_name()} not implemented."
             )
-
-        raise NgioValidationError(
-            f"Could not load table from any known version. Errors: {_errors}"
-        )
+        table = table_cls.from_handler(handler=handler, backend=backend)
+        return table
 
     def add_implementation(self, handler: type[Table], overwrite: bool = False):
         """Register a new table handler."""
-        table_type = handler.type()
-        version = handler.version()
-        if table_type is None:
-            raise NgioValueError("Table handler must have a type.")
-
-        if version is None:
-            raise NgioValueError("Table handler must have a version.")
+        meta = TableMeta(
+            type=handler.table_type(),
+            fractal_table_version=handler.version(),
+        )
 
-        table_unique_name = _unique_table_name(table_type, version)
-        if table_unique_name in self._implemented_tables and not overwrite:
+        if meta.unique_name() in self._implemented_tables and not overwrite:
             raise NgioValueError(
-                f"Table handler for {table_unique_name} already exists. "
+                f"Table handler for {meta.unique_name()} already implemented. "
                 "Use overwrite=True to replace it."
             )
-        self._implemented_tables[table_unique_name] = handler
-
-
-def _get_table_type(handler: ZarrGroupHandler) -> str:
-    """Get the type of the table from the handler."""
-    attrs = handler.load_attrs()
-    return attrs.get("type", "None")
-
-
-def _get_table_version(handler: ZarrGroupHandler) -> str:
-    """Get the version of the table from the handler."""
-    attrs = handler.load_attrs()
-    return attrs.get("fractal_table_version", "None")
+        self._implemented_tables[meta.unique_name()] = handler
 
 
 class TablesContainer:
@@ -208,34 +256,52 @@ class TablesContainer:
         filtered_tables = []
         for table_name in tables:
             tb_handler = self._get_table_group_handler(table_name)
-            table_type = _get_table_type(tb_handler)
+            table_type = _get_meta(tb_handler).type
             if table_type == filter_types:
                 filtered_tables.append(table_name)
         return filtered_tables
 
     def get(
-        self, name: str, backend_name: str | None = None, strict: bool = True
+        self,
+        name: str,
+        backend: str | TableBackendProtocol | None = None,
+        strict: bool = True,
     ) -> Table:
         """Get a label from the group."""
         if name not in self.list():
-            raise KeyError(f"Table '{name}' not found in the group.")
+            raise NgioValueError(f"Table '{name}' not found in the group.")
 
         table_handler = self._get_table_group_handler(name)
-        table_type = _get_table_type(table_handler)
-        table_version = _get_table_version(table_handler)
+
+        meta = _get_meta(table_handler)
         return ImplementedTables().get_table(
-            type=table_type,
-            version=table_version,
+            meta=meta,
             handler=table_handler,
-            backend_name=backend_name,
+            backend=backend,
             strict=strict,
         )
 
+    def get_as(
+        self,
+        name: str,
+        table_cls: type[TableType],
+        backend: str | TableBackendProtocol | None = None,
+    ) -> TableType:
+        """Get a table from the group as a specific type."""
+        if name not in self.list():
+            raise NgioValueError(f"Table '{name}' not found in the group.")
+
+        table_handler = self._get_table_group_handler(name)
+        return table_cls.from_handler(
+            handler=table_handler,
+            backend=backend,
+        )  # type: ignore[return-value]
+
     def add(
         self,
         name: str,
         table: Table,
-        backend: str | None = None,
+        backend: str | TableBackendProtocol = "anndata_v1",
         overwrite: bool = False,
     ) -> None:
         """Add a table to the group."""
@@ -253,9 +319,10 @@ class TablesContainer:
         if backend is None:
             backend = table.backend_name
 
-        table._set_backend(
+        table.set_table_data()
+        table.set_backend(
             handler=table_handler,
-            backend_name=backend,
+            backend=backend,
         )
         table.consolidate()
         if name not in existing_tables:
@@ -266,6 +333,7 @@ class TablesContainer:
 ImplementedTables().add_implementation(RoiTableV1)
 ImplementedTables().add_implementation(MaskingRoiTableV1)
 ImplementedTables().add_implementation(FeatureTableV1)
+ImplementedTables().add_implementation(ConditionTableV1)
 
 ###################################################################################
 #
@@ -289,6 +357,7 @@ def open_tables_container(
 
 def open_table(
     store: StoreOrGroup,
+    backend: str | TableBackendProtocol | None = None,
     cache: bool = False,
     mode: AccessModeLiteral = "a",
     parallel_safe: bool = False,
@@ -297,15 +366,34 @@ def open_table(
     handler = ZarrGroupHandler(
         store=store, cache=cache, mode=mode, parallel_safe=parallel_safe
     )
+    meta = _get_meta(handler)
     return ImplementedTables().get_table(
-        _get_table_type(handler), _get_table_version(handler), handler
+        meta=meta, handler=handler, backend=backend, strict=False
     )
 
 
+def open_table_as(
+    store: StoreOrGroup,
+    table_cls: type[TableType],
+    backend: str | TableBackendProtocol = "anndata_v1",
+    cache: bool = False,
+    mode: AccessModeLiteral = "a",
+    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
+    )
+    return table_cls.from_handler(
+        handler=handler,
+        backend=backend,
+    )  # type: ignore[return-value]
+
+
 def write_table(
     store: StoreOrGroup,
     table: Table,
-    backend: str | None = None,
+    backend: str | TableBackendProtocol = "anndata_v1",
     cache: bool = False,
     mode: AccessModeLiteral = "a",
     parallel_safe: bool = False,
@@ -314,8 +402,8 @@ def write_table(
     handler = ZarrGroupHandler(
         store=store, cache=cache, mode=mode, parallel_safe=parallel_safe
     )
-    table._set_backend(
+    table.set_backend(
         handler=handler,
-        backend_name=backend,
+        backend=backend,
     )
     table.consolidate()