lamindb 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

lamindb/core/loaders.py

@@ -20,26 +20,30 @@ from __future__ import annotations
 import builtins
 import re
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-import anndata as ad
 import pandas as pd
+from anndata import read_h5ad
 from lamin_utils import logger
 from lamindb_setup.core.upath import (
     create_path,
     infer_filesystem,
 )
 
-from ._settings import settings
+from ..core._settings import settings
 
 if TYPE_CHECKING:
+    from anndata import AnnData
     from lamindb_setup.core.types import UPathStr
+    from mudata import MuData
+
+    from lamindb.core.types import ScverseDataStructures
 
 try:
-    from .storage._zarr import load_anndata_zarr
+    from ..core.storage._zarr import load_zarr
 except ImportError:
 
-    def load_anndata_zarr(storepath):  # type: ignore
+    def load_zarr(storepath):  # type: ignore
         raise ImportError("Please install zarr: pip install zarr<=2.18.4")
 
 
@@ -47,7 +51,7 @@ is_run_from_ipython = getattr(builtins, "__IPYTHON__", False)
 
 
 # tested in lamin-usecases
-def load_fcs(*args, **kwargs) -> ad.AnnData:
+def load_fcs(*args, **kwargs) -> AnnData:
     """Load an `.fcs` file to `AnnData`."""
     try:
         import readfcs
@@ -62,16 +66,16 @@ def load_tsv(path: UPathStr, **kwargs) -> pd.DataFrame:
     return pd.read_csv(path_sanitized, sep="\t", **kwargs)
 
 
-def load_h5ad(filepath, **kwargs) -> ad.AnnData:
+def load_h5ad(filepath, **kwargs) -> AnnData:
     """Load an `.h5ad` file to `AnnData`."""
     fs, filepath = infer_filesystem(filepath)
-
-    with fs.open(filepath, mode="rb") as file:
-        adata = ad.read_h5ad(file, backed=False, **kwargs)
+    compression = kwargs.pop("compression", "infer")
+    with fs.open(filepath, mode="rb", compression=compression) as file:
+        adata = read_h5ad(file, backed=False, **kwargs)
         return adata
 
 
-def load_h5mu(filepath: UPathStr, **kwargs):
+def load_h5mu(filepath: UPathStr, **kwargs) -> MuData:
     """Load an `.h5mu` file to `MuData`."""
     import mudata as md
 
@@ -100,7 +104,7 @@ def load_html(path: UPathStr) -> None | UPathStr:
         return path
 
 
-def load_json(path: UPathStr) -> dict:
+def load_json(path: UPathStr) -> dict[str, Any] | list[Any]:
     """Load `.json` to `dict`."""
     import json
 
@@ -109,7 +113,7 @@ def load_json(path: UPathStr) -> dict:
     return data
 
 
-def load_yaml(path: UPathStr) -> dict:
+def load_yaml(path: UPathStr) -> dict[str, Any] | list[Any]:
     """Load `.yaml` to `dict`."""
     import yaml  # type: ignore
 
@@ -148,11 +152,15 @@ def load_rds(path: UPathStr) -> UPathStr:
 
 FILE_LOADERS = {
     ".csv": pd.read_csv,
+    ".csv.gz": pd.read_csv,
     ".tsv": load_tsv,
+    ".tsv.gz": load_tsv,
     ".h5ad": load_h5ad,
+    ".h5ad.gz": load_h5ad,
     ".parquet": pd.read_parquet,
+    ".parquet.gz": pd.read_parquet,  # this doesn't work for externally gzipped files, REMOVE LATER
     ".fcs": load_fcs,
-    ".zarr": load_anndata_zarr,
+    ".zarr": load_zarr,
     ".html": load_html,
     ".json": load_json,
     ".yaml": load_yaml,
@@ -168,17 +176,32 @@ SUPPORTED_SUFFIXES = [sfx for sfx in FILE_LOADERS.keys() if sfx != ".rds"]
 """Suffixes with defined artifact loaders."""
 
 
-def load_to_memory(filepath: UPathStr, **kwargs):
+def load_to_memory(
+    filepath: UPathStr, **kwargs
+) -> (
+    pd.DataFrame | ScverseDataStructures | dict[str, Any] | list[Any] | UPathStr | None
+):
     """Load a file into memory.
 
     Returns the filepath if no in-memory form is found.
+    May return None in interactive sessions for images.
     """
     filepath = create_path(filepath)
 
-    filepath = settings._storage_settings.cloud_to_local(filepath, print_progress=True)
+    # infer the correct suffix when .gz is present
+    suffixes = filepath.suffixes
+    suffix = (
+        "".join(suffixes[-2:])
+        if len(suffixes) > 1 and ".gz" in suffixes
+        else filepath.suffix
+    )
 
-    loader = FILE_LOADERS.get(filepath.suffix)
+    loader = FILE_LOADERS.get(suffix, None)
     if loader is None:
-        return filepath
-    else:
-        return loader(filepath, **kwargs)
+        raise NotImplementedError(
+            f"There is no loader for {suffix} files. Use .cache() to get the path."
+        )
+
+    filepath = settings._storage_settings.cloud_to_local(filepath, print_progress=True)
+
+    return loader(filepath, **kwargs)

lamindb/core/storage/_anndata_accessor.py

@@ -16,6 +16,7 @@ from anndata._io.h5ad import read_dataframe_legacy as read_dataframe_legacy_h5
 from anndata._io.specs.registry import get_spec, read_elem, read_elem_partial
 from anndata.compat import _read_attr
 from fsspec.implementations.local import LocalFileSystem
+from fsspec.utils import infer_compression
 from lamin_utils import logger
 from lamindb_setup.core.upath import create_mapper, infer_filesystem
 from packaging import version
@@ -152,9 +153,13 @@ registry = AccessRegistry()
 
 
 @registry.register_open("h5py")
-def open(filepath: UPathStr, mode: str = "r"):
+def open(filepath: UPathStr, mode: str = "r", compression: str | None = "infer"):
     fs, file_path_str = infer_filesystem(filepath)
-    if isinstance(fs, LocalFileSystem):
+    # we don't open compressed files directly because we need fsspec to uncompress on .open
+    compression = (
+        infer_compression(file_path_str) if compression == "infer" else compression
+    )
+    if isinstance(fs, LocalFileSystem) and compression is None:
         assert mode in {"r", "r+", "a", "w", "w-"}, f"Unknown mode {mode}!"  #  noqa: S101
         return None, h5py.File(file_path_str, mode=mode)
     if mode == "r":
@@ -165,7 +170,7 @@ def open(filepath: UPathStr, mode: str = "r"):
         conn_mode = "ab"
     else:
         raise ValueError(f"Unknown mode {mode}! Should be 'r', 'w' or 'a'.")
-    conn = fs.open(file_path_str, mode=conn_mode)
+    conn = fs.open(file_path_str, mode=conn_mode, compression=compression)
     try:
         storage = h5py.File(conn, mode=mode)
     except Exception as e:

lamindb/core/storage/_backed_access.py

@@ -5,8 +5,6 @@ from typing import TYPE_CHECKING, Any, Callable
 
 from anndata._io.specs.registry import get_spec
 
-from lamindb.models import Artifact
-
 from ._anndata_accessor import AnnDataAccessor, StorageType, registry
 from ._pyarrow_dataset import _is_pyarrow_dataset, _open_pyarrow_dataset
 from ._tiledbsoma import _open_tiledbsoma
@@ -19,6 +17,8 @@ if TYPE_CHECKING:
     from tiledbsoma import Experiment as SOMAExperiment
     from upath import UPath
 
+    from lamindb.models.artifact import Artifact
+
 
 # this dynamically creates a subclass of a context manager class
 # and reassigns it to an instance of the superclass
@@ -70,9 +70,12 @@ def backed_access(
     artifact_or_filepath: Artifact | UPath,
     mode: str = "r",
     using_key: str | None = None,
+    **kwargs,
 ) -> (
     AnnDataAccessor | BackedAccessor | SOMACollection | SOMAExperiment | PyArrowDataset
 ):
+    from lamindb.models import Artifact
+
     if isinstance(artifact_or_filepath, Artifact):
         objectpath, _ = filepath_from_artifact(
             artifact_or_filepath, using_key=using_key
@@ -80,18 +83,22 @@ def backed_access(
     else:
         objectpath = artifact_or_filepath
     name = objectpath.name
-    suffix = objectpath.suffix
+    # ignore .gz, only check the real suffix
+    suffixes = objectpath.suffixes
+    suffix = (
+        suffixes[-2] if len(suffixes) > 1 and ".gz" in suffixes else objectpath.suffix
+    )
 
     if name == "soma" or suffix == ".tiledbsoma":
         if mode not in {"r", "w"}:
             raise ValueError("`mode` should be either 'r' or 'w' for tiledbsoma.")
-        return _open_tiledbsoma(objectpath, mode=mode)  # type: ignore
+        return _open_tiledbsoma(objectpath, mode=mode, **kwargs)  # type: ignore
     elif suffix in {".h5", ".hdf5", ".h5ad"}:
-        conn, storage = registry.open("h5py", objectpath, mode=mode)
+        conn, storage = registry.open("h5py", objectpath, mode=mode, **kwargs)
     elif suffix == ".zarr":
-        conn, storage = registry.open("zarr", objectpath, mode=mode)
+        conn, storage = registry.open("zarr", objectpath, mode=mode, **kwargs)
     elif _is_pyarrow_dataset(objectpath):
-        return _open_pyarrow_dataset(objectpath)
+        return _open_pyarrow_dataset(objectpath, **kwargs)
     else:
         raise ValueError(
             "The object should have .h5, .hdf5, .h5ad, .zarr, .tiledbsoma suffix "

lamindb/core/storage/_pyarrow_dataset.py

@@ -18,15 +18,30 @@ def _is_pyarrow_dataset(paths: UPath | list[UPath]) -> bool:
     # we don't check here that the filesystem is the same
     # but this is a requirement for pyarrow.dataset.dataset
     if isinstance(paths, list):
-        suffixes = {path.suffix for path in paths}
-    elif paths.is_file():
-        suffixes = {paths.suffix}
+        path_list = paths
+    elif paths.is_dir():
+        path_list = [path for path in paths.rglob("*") if path.suffix != ""]
     else:
-        suffixes = {path.suffix for path in paths.rglob("*") if path.suffix != ""}
-    return len(suffixes) == 1 and suffixes.pop() in PYARROW_SUFFIXES
-
-
-def _open_pyarrow_dataset(paths: UPath | list[UPath]) -> PyArrowDataset:
+        path_list = [paths]
+    suffix = None
+    for path in path_list:
+        path_suffixes = path.suffixes
+        # this doesn't work for externally gzipped files, REMOVE LATER
+        path_suffix = (
+            path_suffixes[-2]
+            if len(path_suffixes) > 1 and ".gz" in path_suffixes
+            else path.suffix
+        )
+        if path_suffix not in PYARROW_SUFFIXES:
+            return False
+        elif suffix is None:
+            suffix = path_suffix
+        elif path_suffix != suffix:
+            return False
+    return True
+
+
+def _open_pyarrow_dataset(paths: UPath | list[UPath], **kwargs) -> PyArrowDataset:
     if isinstance(paths, list):
         path0 = paths[0]
         if isinstance(path0, LocalPathClasses):
@@ -38,4 +53,4 @@ def _open_pyarrow_dataset(paths: UPath | list[UPath]) -> PyArrowDataset:
     else:
         paths_str, filesystem = paths.path, paths.fs
 
-    return pyarrow.dataset.dataset(paths_str, filesystem=filesystem)
+    return pyarrow.dataset.dataset(paths_str, filesystem=filesystem, **kwargs)

lamindb/core/storage/_tiledbsoma.py

@@ -12,8 +12,6 @@ from lamindb_setup.core._settings_storage import get_storage_region
 from lamindb_setup.core.upath import LocalPathClasses, create_path
 from packaging import version
 
-from lamindb.models import Artifact, Run
-
 if TYPE_CHECKING:
     from lamindb_setup.core.types import UPathStr
     from tiledbsoma import Collection as SOMACollection
@@ -21,12 +19,15 @@ if TYPE_CHECKING:
     from tiledbsoma import Measurement as SOMAMeasurement
     from upath import UPath
 
+    from lamindb.models.artifact import Artifact
+    from lamindb.models.run import Run
+
 
 def _load_h5ad_zarr(objpath: UPath):
-    from lamindb.core.loaders import load_anndata_zarr, load_h5ad
+    from lamindb.core.loaders import load_h5ad, load_zarr
 
     if objpath.is_dir():
-        adata = load_anndata_zarr(objpath)
+        adata = load_zarr(objpath, expected_type="anndata")
     else:
         # read only local in backed for now
         # in principle possible to read remote in backed also
@@ -134,9 +135,10 @@ def save_tiledbsoma_experiment(
     except ImportError as e:
         raise ImportError("Please install tiledbsoma: pip install tiledbsoma") from e
 
-    from lamindb.core._data import get_run
     from lamindb.core.storage.paths import auto_storage_key_from_artifact_uid
-    from lamindb.core.versioning import create_uid
+    from lamindb.models import Artifact
+    from lamindb.models._is_versioned import create_uid
+    from lamindb.models.artifact import get_run
 
     run = get_run(run)
 

lamindb/core/storage/_zarr.py

@@ -1,55 +1,134 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 import scipy.sparse as sparse
 import zarr
 from anndata import __version__ as anndata_version
 from anndata._io.specs import write_elem
-from anndata._io.specs.registry import get_spec
 from fsspec.implementations.local import LocalFileSystem
-from lamindb_setup.core.upath import create_mapper, infer_filesystem
+from lamin_utils import logger
+from lamindb_setup.core.upath import S3FSMap, create_mapper, infer_filesystem
 from packaging import version
 
+from lamindb.core._compat import with_package
+
 from ._anndata_sizes import _size_elem, _size_raw, size_adata
 
 if version.parse(anndata_version) < version.parse("0.11.0"):
-    from anndata._io import read_zarr
+    from anndata._io import read_zarr as read_anndata_zarr
 else:
-    from anndata.io import read_zarr
+    from anndata.io import read_zarr as read_anndata_zarr
 
 
 if TYPE_CHECKING:
     from anndata import AnnData
+    from fsspec import FSMap
     from lamindb_setup.core.types import UPathStr
 
+    from lamindb.core.types import ScverseDataStructures
 
-def zarr_is_adata(storepath: UPathStr) -> bool:
-    fs, storepath_str = infer_filesystem(storepath)
-    if isinstance(fs, LocalFileSystem):
-        # this is faster than through an fsspec mapper for local
-        open_obj = storepath_str
-    else:
-        open_obj = create_mapper(fs, storepath_str, check=True)
-    storage = zarr.open(open_obj, mode="r")
-    return get_spec(storage).encoding_type == "anndata"
 
-
-def load_anndata_zarr(storepath: UPathStr) -> AnnData:
+def create_zarr_open_obj(
+    storepath: UPathStr, *, check: bool = True
+) -> str | S3FSMap | FSMap:
+    """Creates the correct object that can be used to open a zarr file depending on local or remote location."""
     fs, storepath_str = infer_filesystem(storepath)
+
     if isinstance(fs, LocalFileSystem):
-        # this is faster than through an fsspec mapper for local
         open_obj = storepath_str
     else:
-        open_obj = create_mapper(fs, storepath_str, check=True)
-    adata = read_zarr(open_obj)
-    return adata
+        open_obj = create_mapper(fs, storepath_str, check=check)
+
+    return open_obj
+
+
+def _identify_zarr_type_from_storage(
+    storage: zarr.Group,
+) -> Literal["anndata", "mudata", "spatialdata", "unknown"]:
+    """Internal helper to identify zarr type from an open storage object."""
+    try:
+        if storage.attrs.get("encoding-type", "") == "anndata":
+            return "anndata"
+        elif storage.attrs.get("encoding-type", "") == "MuData":
+            return "mudata"
+        elif "spatialdata_attrs" in storage.attrs:
+            return "spatialdata"
+    except Exception as error:
+        logger.warning(f"an exception occurred {error}")
+    return "unknown"
+
+
+def identify_zarr_type(
+    storepath: UPathStr, *, check: bool = True
+) -> Literal["anndata", "mudata", "spatialdata", "unknown"]:
+    """Identify whether a zarr store is AnnData, SpatialData, or unknown type."""
+    # we can add these cheap suffix-based-checks later
+    # also need to check whether the .spatialdata.zarr suffix
+    # actually becomes a "standard"; currently we don't recognize it
+    # unlike ".anndata.zarr" in VALID_SUFFIXES
+    # suffixes = UPath(storepath).suffixes
+    # if ".spatialdata" in suffixes:
+    #     return "spatialdata"
+    # elif ".anndata" in suffixes:
+    #     return "anndata"
+
+    open_obj = create_zarr_open_obj(storepath, check=check)
+    try:
+        storage = zarr.open(open_obj, mode="r")
+        return _identify_zarr_type_from_storage(storage)
+    except Exception as error:
+        logger.warning(
+            f"an exception occured while trying to open the zarr store\n {error}"
+        )
+    return "unknown"
+
+
+def load_zarr(
+    storepath: UPathStr,
+    expected_type: Literal["anndata", "mudata", "spatialdata"] = None,
+) -> ScverseDataStructures:
+    """Loads a zarr store and returns the corresponding scverse data structure.
+
+    Args:
+        storepath: Path to the zarr store
+        expected_type: If provided, ensures the zarr store is of this type ("anndata", "mudata", "spatialdata")
+                       and raises ValueError if it's not
+    """
+    open_obj = create_zarr_open_obj(storepath, check=True)
+
+    # Open the storage once
+    try:
+        storage = zarr.open(open_obj, mode="r")
+    except Exception as error:
+        raise ValueError(f"Could not open zarr store: {error}") from None
+
+    actual_type = _identify_zarr_type_from_storage(storage)
+    if expected_type is not None and actual_type != expected_type:
+        raise ValueError(
+            f"Expected zarr store of type '{expected_type}', but found '{actual_type}'"
+        )
+
+    match actual_type:
+        case "anndata":
+            scverse_obj = read_anndata_zarr(open_obj)
+        case "mudata":
+            scverse_obj = with_package("mudata", lambda mod: mod.read_zarr(open_obj))
+        case "spatialdata":
+            scverse_obj = with_package(
+                "spatialdata", lambda mod: mod.read_zarr(open_obj)
+            )
+        case "unknown" | _:
+            raise ValueError(
+                "Unable to determine zarr store format and therefore cannot load Artifact."
+            )
+    return scverse_obj
 
 
 def write_adata_zarr(
     adata: AnnData, storepath: UPathStr, callback=None, chunks=None, **dataset_kwargs
-):
+) -> None:
     fs, storepath_str = infer_filesystem(storepath)
     store = create_mapper(fs, storepath_str, create=True)
 
@@ -65,7 +144,7 @@ def write_adata_zarr(
     adata_size = None
     cumulative_val = 0
 
-    def _cb(key_write: str | None = None):
+    def _report_progress(key_write: str | None = None):
         nonlocal adata_size
         nonlocal cumulative_val
 
@@ -91,9 +170,9 @@ def write_adata_zarr(
 
     def _write_elem_cb(f, k, elem, dataset_kwargs):
         write_elem(f, k, elem, dataset_kwargs=dataset_kwargs)
-        _cb(k)
+        _report_progress(k)
 
-    _cb(None)
+    _report_progress(None)
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", category=UserWarning, module="zarr")
 
@@ -114,4 +193,4 @@ def write_adata_zarr(
             )
         _write_elem_cb(f, "raw", adata.raw, dataset_kwargs=dataset_kwargs)
     # todo: fix size less than total at the end
-    _cb(None)
+    _report_progress(None)

lamindb/core/storage/objects.py

@@ -1,62 +1,97 @@
 from __future__ import annotations
 
 from pathlib import PurePosixPath
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeAlias
 
 from anndata import AnnData
 from pandas import DataFrame
 
+from lamindb.core._compat import (
+    with_package_obj,
+)
+from lamindb.core.types import ScverseDataStructures
+
 if TYPE_CHECKING:
     from lamindb_setup.core.types import UPathStr
 
-
-def _mudata_is_installed():
-    try:
-        import mudata  # noqa: F401c
-    except ImportError:
-        return False
-    return True
+SupportedDataTypes: TypeAlias = DataFrame | ScverseDataStructures
 
 
-def infer_suffix(dmem, adata_format: str | None = None):
+def infer_suffix(dmem: SupportedDataTypes, format: str | None = None):
     """Infer LaminDB storage file suffix from a data object."""
     if isinstance(dmem, AnnData):
-        if adata_format is not None:
-            if adata_format not in {"h5ad", "zarr", "anndata.zarr"}:
+        if format is not None:
+            if format not in {"h5ad", "zarr", "anndata.zarr"}:
                 raise ValueError(
                     "Error when specifying AnnData storage format, it should be"
-                    f" 'h5ad', 'zarr', not '{adata_format}'. Check 'format'"
+                    f" 'h5ad', 'zarr', not '{format}'. Check 'format'"
                     " or the suffix of 'key'."
                 )
-            return "." + adata_format
+            return "." + format
         return ".h5ad"
-    elif isinstance(dmem, DataFrame):
+
+    if isinstance(dmem, DataFrame):
         return ".parquet"
-    else:
-        if _mudata_is_installed():
-            from mudata import MuData
 
-            if isinstance(dmem, MuData):
-                return ".h5mu"
+    if with_package_obj(
+        dmem,
+        "MuData",
+        "mudata",
+        lambda obj: True,  # Just checking type, not calling any method
+    )[0]:
+        return ".h5mu"
+
+    has_spatialdata, spatialdata_suffix = with_package_obj(
+        dmem,
+        "SpatialData",
+        "spatialdata",
+        lambda obj: "."
+        + (
+            format
+            if format is not None and format in {"spatialdata.zarr", "zarr"}
+            else ".zarr"
+            if format is None
+            else (_ for _ in ()).throw(
+                ValueError(
+                    "Error when specifying SpatialData storage format, it should be"
+                    f" 'zarr', 'spatialdata.zarr', not '{format}'. Check 'format'"
+                    " or the suffix of 'key'."
+                )
+            )
+        ),
+    )
+    if has_spatialdata:
+        return spatialdata_suffix
+    else:
         raise NotImplementedError
 
 
-def write_to_disk(dmem, filepath: UPathStr):
+def write_to_disk(dmem: SupportedDataTypes, filepath: UPathStr) -> None:
+    """Writes the passed in memory data to disk to a specified path."""
     if isinstance(dmem, AnnData):
         suffix = PurePosixPath(filepath).suffix
         if suffix == ".h5ad":
             dmem.write_h5ad(filepath)
+            return
         elif suffix == ".zarr":
             dmem.write_zarr(filepath)
+            return
         else:
             raise NotImplementedError
-    elif isinstance(dmem, DataFrame):
+
+    if isinstance(dmem, DataFrame):
         dmem.to_parquet(filepath)
-    else:
-        if _mudata_is_installed():
-            from mudata import MuData
+        return
 
-            if isinstance(dmem, MuData):
-                dmem.write(filepath)
-                return
-        raise NotImplementedError
+    if with_package_obj(dmem, "MuData", "mudata", lambda obj: obj.write(filepath))[0]:
+        return
+
+    if with_package_obj(
+        dmem,
+        "SpatialData",
+        "spatialdata",
+        lambda obj: obj.write(filepath, overwrite=True),
+    )[0]:
+        return
+
+    raise NotImplementedError