lamindb 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

lamindb/models/_label_manager.py

@@ -24,7 +24,7 @@ from ._describe import (
     TYPE_WIDTH,
     VALUES_WIDTH,
     describe_header,
-    print_rich_tree,
+    format_rich_tree,
 )
 from ._django import get_artifact_with_related, get_related_model
 from ._relations import dict_related_model_to_related_name
@@ -182,8 +182,14 @@ class LabelManager:
         self._host = host
 
     def __repr__(self) -> str:
+        return self.describe(return_str=True)
+
+    def describe(self, return_str=True) -> str:
+        """Describe the labels."""
         tree = describe_labels(self._host)
-        return print_rich_tree(tree, fallback="no linked labels")
+        return format_rich_tree(
+            tree, fallback="no linked labels", return_str=return_str
+        )
 
     def add(
         self,

lamindb/models/artifact.py

@@ -5,7 +5,7 @@ import os
 import shutil
 from collections import defaultdict
 from pathlib import Path, PurePath, PurePosixPath
-from typing import TYPE_CHECKING, Any, Union, overload
+from typing import TYPE_CHECKING, Any, Literal, Union, overload
 
 import fsspec
 import lamindb_setup as ln_setup
@@ -48,6 +48,11 @@ from ..core.storage import (
     write_to_disk,
 )
 from ..core.storage._anndata_accessor import _anndata_n_observations
+from ..core.storage._backed_access import (
+    _track_writes_factory,
+    backed_access,
+)
+from ..core.storage._polars_lazy_df import POLARS_SUFFIXES
 from ..core.storage._pyarrow_dataset import PYARROW_SUFFIXES
 from ..core.storage._tiledbsoma import _soma_n_observations
 from ..core.storage.paths import (
@@ -105,9 +110,10 @@ except ImportError:
 
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Iterable, Iterator
 
     from mudata import MuData  # noqa: TC004
+    from polars import LazyFrame as PolarsLazyFrame
     from pyarrow.dataset import Dataset as PyArrowDataset
     from spatialdata import SpatialData  # noqa: TC004
     from tiledbsoma import Collection as SOMACollection
@@ -311,10 +317,9 @@ def get_stat_or_artifact(
         result = Artifact.objects.using(instance).filter(hash=hash).all()
         artifact_with_same_hash_exists = len(result) > 0
     else:
-        storage_id = settings.storage.id
         result = (
             Artifact.objects.using(instance)
-            .filter(Q(hash=hash) | Q(key=key, storage_id=storage_id))
+            .filter(Q(hash=hash) | Q(key=key, storage=settings.storage.record))
             .order_by("-created_at")
             .all()
         )
@@ -759,15 +764,15 @@ def _describe_sqlite(self, print_types: bool = False):  # for artifact & collect
         return tree
 
 
-def describe_artifact_collection(self):  # for artifact & collection
-    from ._describe import print_rich_tree
+def describe_artifact_collection(self, return_str: bool = False) -> str | None:
+    from ._describe import format_rich_tree
 
     if not self._state.adding and connections[self._state.db].vendor == "postgresql":
         tree = _describe_postgres(self)
     else:
         tree = _describe_sqlite(self)
 
-    print_rich_tree(tree)
+    return format_rich_tree(tree, return_str=return_str)
 
 
 def validate_feature(feature: Feature, records: list[Record]) -> None:
@@ -962,7 +967,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
 
         Create an artifact **from a local file or folder**::
 
-            artifact = ln.Artifact("./my_file.parquet", key="example_datasets/my_file.parquet").save()
+            artifact = ln.Artifact("./my_file.parquet", key="examples/my_file.parquet").save()
             artifact = ln.Artifact("./my_folder", key="project1/my_folder").save()
 
         Calling `.save()` copies or uploads the file to the default storage location of your lamindb instance.
@@ -977,29 +982,12 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
 
         You can make a **new version** of an artifact by passing an existing `key`::
 
-            artifact_v2 = ln.Artifact("./my_file.parquet", key="example_datasets/my_file.parquet").save()
+            artifact_v2 = ln.Artifact("./my_file.parquet", key="examples/my_file.parquet").save()
             artifact_v2.versions.df()  # see all versions
 
-        .. dropdown:: Why does the API look this way?
-
-            It's inspired by APIs building on AWS S3.
-
-            Both boto3 and quilt select a bucket (a storage location in LaminDB) and define a target path through a `key` argument.
+        You can write artifacts to other storage locations by switching the current default storage location (:attr:`~lamindb.core.Settings.storage`)::
 
-            In `boto3 <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3/bucket/upload_file.html>`__::
-
-                # signature: S3.Bucket.upload_file(filepath, key)
-                import boto3
-                s3 = boto3.resource('s3')
-                bucket = s3.Bucket('mybucket')
-                bucket.upload_file('/tmp/hello.txt', 'hello.txt')
-
-            In `quilt3 <https://docs.quiltdata.com/api-reference/bucket>`__::
-
-                # signature: quilt3.Bucket.put_file(key, filepath)
-                import quilt3
-                bucket = quilt3.Bucket('mybucket')
-                bucket.put_file('hello.txt', '/tmp/hello.txt')
+            ln.settings.storage = "s3://some-bucket"
 
         Sometimes you want to **avoid mapping the artifact into a path hierarchy**, and you only pass `description`::
 
@@ -1034,6 +1022,27 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             In concurrent workloads where the same artifact is created repeatedly at the exact same time, `.save()`
             detects the duplication and will return the existing artifact.
 
+        .. dropdown:: Why does the constructor look the way it looks?
+
+            It's inspired by APIs building on AWS S3.
+
+            Both boto3 and quilt select a bucket (a storage location in LaminDB) and define a target path through a `key` argument.
+
+            In `boto3 <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3/bucket/upload_file.html>`__::
+
+                # signature: S3.Bucket.upload_file(filepath, key)
+                import boto3
+                s3 = boto3.resource('s3')
+                bucket = s3.Bucket('mybucket')
+                bucket.upload_file('/tmp/hello.txt', 'hello.txt')
+
+            In `quilt3 <https://docs.quiltdata.com/api-reference/bucket>`__::
+
+                # signature: quilt3.Bucket.put_file(key, filepath)
+                import quilt3
+                bucket = quilt3.Bucket('mybucket')
+                bucket.put_file('hello.txt', '/tmp/hello.txt')
+
     See Also:
         :class:`~lamindb.Storage`
             Storage locations for artifacts.
@@ -1089,7 +1098,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
         ln.Artifact.filter(scientist="Barbara McClintock")
 
     Features may or may not be part of the artifact content in storage. For
-    instance, the :class:`~lamindb.Curator` flow validates the columns of a
+    instance, the :class:`~lamindb.curators.DataFrameCurator` flow validates the columns of a
     `DataFrame`-like artifact and annotates it with features corresponding to
     these columns. `artifact.features.add_values`, by contrast, does not
     validate the content of the artifact.
@@ -1525,7 +1534,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             ::
 
                 artifact = ln.Artifact.get("tCUkRcaEjTjhtozp0000")
-                artifact = ln.Arfifact.get(key="my_datasets/my_file.parquet")
+                artifact = ln.Arfifact.get(key="examples/my_file.parquet")
         """
         from .query_set import QuerySet
 
@@ -1550,7 +1559,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
 
             Query by fields::
 
-                ln.Arfifact.filter(key="my_datasets/my_file.parquet")
+                ln.Arfifact.filter(key="examples/my_file.parquet")
 
             Query by features::
 
@@ -1610,7 +1619,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
         schema: Schema | None = None,
         **kwargs,
     ) -> Artifact:
-        """Create from `DataFrame`, validate & link features.
+        """Create from `DataFrame`, optionally validate & annotate.
 
         Args:
             df: A `DataFrame` object.
@@ -1619,7 +1628,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             description: A description.
             revises: An old version of the artifact.
             run: The run that creates the artifact.
-            schema: A schema to validate & annotate.
+            schema: A schema that defines how to validate & annotate.
 
         See Also:
             :meth:`~lamindb.Collection`
@@ -1627,19 +1636,30 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             :class:`~lamindb.Feature`
                 Track features.
 
-        Example::
+        Example:
 
-            import lamindb as ln
+            No validation and annotation::
+
+                import lamindb as ln
+
+                df = ln.core.datasets.mini_immuno.get_dataset1()
+                artifact = ln.Artifact.from_df(df, key="examples/dataset1.parquet").save()
+
+            With validation and annotation.
+
+            .. literalinclude:: scripts/curate_dataframe_flexible.py
+               :language: python
+
+            Under-the-hood, this used the following schema.
+
+            .. literalinclude:: scripts/define_valid_features.py
+               :language: python
+
+            Valid features & labels were defined as:
+
+            .. literalinclude:: scripts/define_mini_immuno_features_labels.py
+               :language: python
 
-            df = ln.core.datasets.df_iris_in_meter_batch1()
-            df.head()
-            #>   sepal_length sepal_width petal_length petal_width iris_organism_code
-            #> 0        0.051       0.035        0.014       0.002                 0
-            #> 1        0.049       0.030        0.014       0.002                 0
-            #> 2        0.047       0.032        0.013       0.002                 0
-            #> 3        0.046       0.031        0.015       0.002                 0
-            #> 4        0.050       0.036        0.014       0.002                 0
-            artifact = ln.Artifact.from_df(df, key="iris/result_batch1.parquet").save()
         """
         artifact = Artifact(  # type: ignore
             data=df,
@@ -1673,7 +1693,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
         schema: Schema | None = None,
         **kwargs,
     ) -> Artifact:
-        """Create from ``AnnData``, validate & link features.
+        """Create from `AnnData`, optionally validate & annotate.
 
         Args:
             adata: An `AnnData` object or a path of AnnData-like.
@@ -1682,7 +1702,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             description: A description.
             revises: An old version of the artifact.
             run: The run that creates the artifact.
-            schema: A schema to validate & annotate.
+            schema: A schema that defines how to validate & annotate.
 
         See Also:
 
@@ -1691,12 +1711,31 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             :class:`~lamindb.Feature`
                 Track features.
 
-        Example::
+        Example:
 
-            import lamindb as ln
+            No validation and annotation::
+
+                import lamindb as ln
+
+                adata = ln.core.datasets.anndata_with_obs()
+                artifact = ln.Artifact.from_anndata(adata, key="mini_anndata_with_obs.h5ad").save()
+
+            With validation and annotation.
+
+            .. literalinclude:: scripts/curate_anndata_flexible.py
+               :language: python
+
+            Under-the-hood, this used the following schema.
+
+            .. literalinclude:: scripts/define_schema_anndata_ensembl_gene_ids_and_valid_features_in_obs.py
+               :language: python
+
+            This schema tranposes the `var` DataFrame during curation, so that one validates and annotates the `var.T` schema, i.e., `[ENSG00000153563, ENSG00000010610, ENSG00000170458]`.
+            If one doesn't transpose, one would annotate with the schema of `var`, i.e., `[gene_symbol, gene_type]`.
+
+            .. image:: https://lamin-site-assets.s3.amazonaws.com/.lamindb/gLyfToATM7WUzkWW0001.png
+               :width: 800px
 
-            adata = ln.core.datasets.anndata_with_obs()
-            artifact = ln.Artifact.from_anndata(adata, key="mini_anndata_with_obs.h5ad").save()
         """
         if not data_is_anndata(adata):
             raise ValueError(
@@ -1745,7 +1784,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
         schema: Schema | None = None,
         **kwargs,
     ) -> Artifact:
-        """Create from ``MuData``, validate & link features.
+        """Create from `MuData`, optionally validate & annotate.
 
         Args:
             mdata: A `MuData` object.
@@ -1754,7 +1793,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             description: A description.
             revises: An old version of the artifact.
             run: The run that creates the artifact.
-            schema: A schema to validate & annotate.
+            schema: A schema that defines how to validate & annotate.
 
         See Also:
             :meth:`~lamindb.Collection`
@@ -1804,16 +1843,16 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
         schema: Schema | None = None,
         **kwargs,
     ) -> Artifact:
-        """Create from ``SpatialData``, validate & link features.
+        """Create from `SpatialData`, optionally validate & annotate.
 
         Args:
-            mdata: A `SpatialData` object.
+            sdata: A `SpatialData` object.
             key: A relative path within default storage,
                 e.g., `"myfolder/myfile.zarr"`.
             description: A description.
             revises: An old version of the artifact.
             run: The run that creates the artifact.
-             schema: A schema to validate & annotate.
+            schema: A schema that defines how to validate & annotate.
 
         See Also:
             :meth:`~lamindb.Collection`
@@ -1821,11 +1860,21 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             :class:`~lamindb.Feature`
                 Track features.
 
-        Example::
+        Example:
 
-            import lamindb as ln
+            No validation and annotation::
+
+                import lamindb as ln
+
+                artifact = ln.Artifact.from_spatialdata(sdata, key="my_dataset.zarr").save()
 
-            artifact = ln.Artifact.from_spatialdata(sdata, key="my_dataset.zarr").save()
+            With validation and annotation.
+
+            .. literalinclude:: scripts/define_schema_spatialdata.py
+                :language: python
+
+            .. literalinclude:: scripts/curate_spatialdata.py
+                :language: python
         """
         if not data_is_spatialdata(sdata):
             raise ValueError(
@@ -2117,29 +2166,39 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
         self._old_suffix = self.suffix
 
     def open(
-        self, mode: str = "r", is_run_input: bool | None = None, **kwargs
-    ) -> Union[
-        AnnDataAccessor,
-        BackedAccessor,
-        SOMACollection,
-        SOMAExperiment,
-        SOMAMeasurement,
-        PyArrowDataset,
-    ]:
-        """Return a cloud-backed data object.
+        self,
+        mode: str = "r",
+        engine: Literal["pyarrow", "polars"] = "pyarrow",
+        is_run_input: bool | None = None,
+        **kwargs,
+    ) -> (
+        AnnDataAccessor
+        | BackedAccessor
+        | SOMACollection
+        | SOMAExperiment
+        | SOMAMeasurement
+        | PyArrowDataset
+        | Iterator[PolarsLazyFrame]
+    ):
+        """Open a dataset for streaming.
 
         Works for `AnnData` (`.h5ad` and `.zarr`), generic `hdf5` and `zarr`,
-        `tiledbsoma` objects (`.tiledbsoma`), `pyarrow` compatible formats.
+        `tiledbsoma` objects (`.tiledbsoma`), `pyarrow` or `polars` compatible formats
+        (`.parquet`, `.csv`, `.ipc` etc. files or directories with such files).
 
         Args:
             mode: can only be `"w"` (write mode) for `tiledbsoma` stores,
                 otherwise should be always `"r"` (read-only mode).
+            engine: Which module to use for lazy loading of a dataframe
+                from `pyarrow` or `polars` compatible formats.
+                This has no effect if the artifact is not a dataframe, i.e.
+                if it is an `AnnData,` `hdf5`, `zarr` or `tiledbsoma` object.
             is_run_input: Whether to track this artifact as run input.
             **kwargs: Keyword arguments for the accessor, i.e. `h5py` or `zarr` connection,
-                `pyarrow.dataset.dataset`.
+                `pyarrow.dataset.dataset`, `polars.scan_*` function.
 
         Notes:
-            For more info, see tutorial: :doc:`/arrays`.
+            For more info, see guide: :doc:`/arrays`.
 
         Example::
 
@@ -2152,6 +2211,10 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             #> AnnDataAccessor object with n_obs × n_vars = 70 × 765
             #>     constructed for the AnnData object pbmc68k.h5ad
             #>     ...
+            artifact = ln.Artifact.get(key="lndb-storage/df.parquet")
+            artifact.open()
+            #> pyarrow._dataset.FileSystemDataset
+
         """
         if self._overwrite_versions and not self.is_latest:
             raise ValueError(INCONSISTENT_STATE_MSG)
@@ -2159,6 +2222,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
         h5_suffixes = [".h5", ".hdf5", ".h5ad"]
         h5_suffixes += [s + ".gz" for s in h5_suffixes]
         # ignore empty suffix for now
+        df_suffixes = tuple(set(PYARROW_SUFFIXES).union(POLARS_SUFFIXES))
         suffixes = (
             (
                 "",
@@ -2167,7 +2231,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
                 ".tiledbsoma",
             )
             + tuple(h5_suffixes)
-            + PYARROW_SUFFIXES
+            + df_suffixes
             + tuple(
                 s + ".gz" for s in PYARROW_SUFFIXES
             )  # this doesn't work for externally gzipped files, REMOVE LATER
@@ -2175,10 +2239,10 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
         if self.suffix not in suffixes:
             raise ValueError(
                 "Artifact should have a zarr, h5, tiledbsoma object"
-                " or a compatible `pyarrow.dataset.dataset` directory"
+                " or a compatible `pyarrow.dataset.dataset` or `polars.scan_*` directory"
                 " as the underlying data, please use one of the following suffixes"
                 f" for the object name: {', '.join(suffixes[1:])}."
-                f" Or no suffix for a folder with {', '.join(PYARROW_SUFFIXES)} files"
+                f" Or no suffix for a folder with {', '.join(df_suffixes)} files"
                 " (no mixing allowed)."
             )
         if self.suffix != ".tiledbsoma" and self.key != "soma" and mode != "r":
@@ -2187,10 +2251,6 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             )
 
         from lamindb import settings
-        from lamindb.core.storage._backed_access import (
-            _track_writes_factory,
-            backed_access,
-        )
 
         using_key = settings._using_key
         filepath, cache_key = filepath_cache_key_from_artifact(
@@ -2211,14 +2271,22 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             ) and not filepath.synchronize(localpath, just_check=True)
         if open_cache:
             try:
-                access = backed_access(localpath, mode, using_key, **kwargs)
+                access = backed_access(
+                    localpath, mode, engine, using_key=using_key, **kwargs
+                )
             except Exception as e:
-                if isinstance(filepath, LocalPathClasses):
+                # also ignore ValueError here because
+                # such errors most probably just imply an incorrect argument
+                if isinstance(filepath, LocalPathClasses) or isinstance(
+                    e, (ImportError, ValueError)
+                ):
                     raise e
                 logger.warning(
                     f"The cache might be corrupted: {e}. Trying to open directly."
                 )
-                access = backed_access(filepath, mode, using_key, **kwargs)
+                access = backed_access(
+                    filepath, mode, engine, using_key=using_key, **kwargs
+                )
                 # happens only if backed_access has been successful
                 # delete the corrupted cache
                 if localpath.is_dir():
@@ -2226,7 +2294,9 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
                 else:
                     localpath.unlink(missing_ok=True)
         else:
-            access = backed_access(filepath, mode, using_key, **kwargs)
+            access = backed_access(
+                filepath, mode, engine, using_key=using_key, **kwargs
+            )
             if is_tiledbsoma_w:
 
                 def finalize():
@@ -2399,6 +2469,9 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             artifact = ln.Artifact.get(key="some.tiledbsoma". is_latest=True)
             artiact.delete() # delete all versions, the data will be deleted or prompted for deletion.
         """
+        # we're *not* running the line below because the case `storage is None` triggers user feedback in one case
+        # storage = True if storage is None else storage
+
         # this first check means an invalid delete fails fast rather than cascading through
         # database and storage permission errors
         if os.getenv("LAMINDB_MULTI_INSTANCE") is None:
@@ -2449,8 +2522,10 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
             # only delete in storage if DB delete is successful
             # DB delete might error because of a foreign key constraint violated etc.
             if self._overwrite_versions and self.is_latest:
-                # includes self
-                for version in self.versions.all():
+                logger.important(
+                    "deleting all versions of this artifact because they all share the same store"
+                )
+                for version in self.versions.all():  # includes self
                     _delete_skip_storage(version)
             else:
                 self._delete_skip_storage()
@@ -2460,7 +2535,7 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
                 delete_in_storage = False
                 if storage:
                     logger.warning(
-                        "Storage argument is ignored; can't delete storage on an previous version"
+                        "storage argument is ignored; can't delete store of a previous version if overwrite_versions is True"
                     )
             elif self.key is None or self._key_is_virtual:
                 # do not ask for confirmation also if storage is None
@@ -2577,14 +2652,13 @@ class Artifact(Record, IsVersioned, TracksRun, TracksUpdates):
         self._branch_code = 1
         self.save()
 
-    def describe(self) -> None:
-        """Describe relations of record.
+    def describe(self, return_str: bool = False) -> None:
+        """Describe record including linked records.
 
-        Example::
-
-            artifact.describe()
+        Args:
+            return_str: Return a string instead of printing.
         """
-        return describe_artifact_collection(self)
+        return describe_artifact_collection(self, return_str=return_str)
 
     def _populate_subsequent_runs(self, run: Run) -> None:
         _populate_subsequent_runs_(self, run)
@@ -2624,9 +2698,11 @@ def _save_skip_storage(artifact, **kwargs) -> None:
 
 class ArtifactFeatureValue(BasicRecord, LinkORM, TracksRun):
     id: int = models.BigAutoField(primary_key=True)
-    artifact: Artifact = ForeignKey(Artifact, CASCADE, related_name="+")
+    artifact: Artifact = ForeignKey(
+        Artifact, CASCADE, related_name="links_featurevalue"
+    )
     # we follow the lower() case convention rather than snake case for link models
-    featurevalue = ForeignKey(FeatureValue, PROTECT, related_name="+")
+    featurevalue = ForeignKey(FeatureValue, PROTECT, related_name="links_artifact")
 
     class Meta:
         unique_together = ("artifact", "featurevalue")
@@ -2634,9 +2710,11 @@ class ArtifactFeatureValue(BasicRecord, LinkORM, TracksRun):
 
 class ArtifactParamValue(BasicRecord, LinkORM, TracksRun):
     id: int = models.BigAutoField(primary_key=True)
-    artifact: Artifact = ForeignKey(Artifact, CASCADE, related_name="+")
+    artifact: Artifact = ForeignKey(Artifact, CASCADE, related_name="links_paramvalue")
     # we follow the lower() case convention rather than snake case for link models
-    paramvalue: ParamValue = ForeignKey(ParamValue, PROTECT, related_name="+")
+    paramvalue: ParamValue = ForeignKey(
+        ParamValue, PROTECT, related_name="links_artifact"
+    )
 
     class Meta:
         unique_together = ("artifact", "paramvalue")