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

lamindb/models/artifact_set.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from typing import TYPE_CHECKING, Literal
+
+from lamin_utils import logger
+from lamindb_setup.core._docs import doc_args
+
+from ..core._mapped_collection import MappedCollection
+from ..core.storage._backed_access import _open_dataframe
+from .artifact import Artifact, _track_run_input
+from .collection import Collection, _load_concat_artifacts
+
+if TYPE_CHECKING:
+    from anndata import AnnData
+    from pandas import DataFrame
+    from polars import LazyFrame as PolarsLazyFrame
+    from pyarrow.dataset import Dataset as PyArrowDataset
+    from upath import UPath
+
+
+UNORDERED_WARNING = (
+    "this query set is unordered, consider using `.order_by()` first "
+    "to avoid opening the artifacts in an arbitrary order"
+)
+
+
+class ArtifactSet(Iterable):
+    """Abstract class representing sets of artifacts returned by queries.
+
+    This class automatically extends :class:`~lamindb.models.BasicQuerySet`
+    and :class:`~lamindb.models.QuerySet` when the base model is :class:`~lamindb.Artifact`.
+
+    Examples:
+
+        >>> artifacts = ln.Artifact.filter(otype="AnnData")
+        >>> artifacts # an instance of ArtifactQuerySet inheriting from ArtifactSet
+    """
+
+    @doc_args(Collection.load.__doc__)
+    def load(
+        self,
+        join: Literal["inner", "outer"] = "outer",
+        is_run_input: bool | None = None,
+        **kwargs,
+    ) -> DataFrame | AnnData:
+        """{}"""  # noqa: D415
+        if not self.ordered:  # type: ignore
+            logger.warning(UNORDERED_WARNING)
+
+        artifacts: list[Artifact] = list(self)
+        concat_object = _load_concat_artifacts(artifacts, join, **kwargs)
+        # track only if successful
+        _track_run_input(artifacts, is_run_input)
+        return concat_object
+
+    @doc_args(Collection.open.__doc__)
+    def open(
+        self,
+        engine: Literal["pyarrow", "polars"] = "pyarrow",
+        is_run_input: bool | None = None,
+        **kwargs,
+    ) -> PyArrowDataset | Iterator[PolarsLazyFrame]:
+        """{}"""  # noqa: D415
+        if not self.ordered:  # type: ignore
+            logger.warning(UNORDERED_WARNING)
+
+        artifacts: list[Artifact] = list(self)
+        paths: list[UPath] = [artifact.path for artifact in artifacts]
+
+        dataframe = _open_dataframe(paths, engine=engine, **kwargs)
+        # track only if successful
+        _track_run_input(artifacts, is_run_input)
+        return dataframe
+
+    @doc_args(Collection.mapped.__doc__)
+    def mapped(
+        self,
+        layers_keys: str | list[str] | None = None,
+        obs_keys: str | list[str] | None = None,
+        obsm_keys: str | list[str] | None = None,
+        obs_filter: dict[str, str | list[str]] | None = None,
+        join: Literal["inner", "outer"] | None = "inner",
+        encode_labels: bool | list[str] = True,
+        unknown_label: str | dict[str, str] | None = None,
+        cache_categories: bool = True,
+        parallel: bool = False,
+        dtype: str | None = None,
+        stream: bool = False,
+        is_run_input: bool | None = None,
+    ) -> MappedCollection:
+        """{}"""  # noqa: D415
+        if not self.ordered:  # type: ignore
+            logger.warning(UNORDERED_WARNING)
+
+        artifacts: list[Artifact] = []
+        paths: list[UPath] = []
+        for artifact in self:
+            if ".h5ad" not in artifact.suffix and ".zarr" not in artifact.suffix:
+                logger.warning(f"ignoring artifact with suffix {artifact.suffix}")
+                continue
+            elif not stream:
+                paths.append(artifact.cache())
+            else:
+                paths.append(artifact.path)
+            artifacts.append(artifact)
+        ds = MappedCollection(
+            paths,
+            layers_keys,
+            obs_keys,
+            obsm_keys,
+            obs_filter,
+            join,
+            encode_labels,
+            unknown_label,
+            cache_categories,
+            parallel,
+            dtype,
+        )
+        # track only if successful
+        _track_run_input(artifacts, is_run_input)
+        return ds

lamindb/models/collection.py

@@ -24,7 +24,7 @@ from lamindb.base.fields import (
 
 from ..base.ids import base62_20
 from ..core._mapped_collection import MappedCollection
-from ..core.storage._pyarrow_dataset import _is_pyarrow_dataset, _open_pyarrow_dataset
+from ..core.storage._backed_access import _open_dataframe
 from ..errors import FieldValidationError
 from ..models._is_versioned import process_revises
 from ._is_versioned import IsVersioned
@@ -48,8 +48,9 @@ from .record import (
 from .run import Run, TracksRun, TracksUpdates
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Iterable, Iterator
 
+    from polars import LazyFrame as PolarsLazyFrame
     from pyarrow.dataset import Dataset as PyArrowDataset
 
     from ..core.storage import UPath
@@ -94,6 +95,39 @@ if TYPE_CHECKING:
 #         return feature_sets_union
 
 
+def _load_concat_artifacts(
+    artifacts: list[Artifact], join: Literal["inner", "outer"] = "outer", **kwargs
+) -> pd.DataFrame | ad.AnnData:
+    suffixes = {artifact.suffix for artifact in artifacts}
+    # Why is that? - Sergei
+    if len(suffixes) != 1:
+        raise ValueError(
+            "Can only load collections where all artifacts have the same suffix"
+        )
+
+    # because we're tracking data flow on the collection-level, here, we don't
+    # want to track it on the artifact-level
+    first_object = artifacts[0].load(is_run_input=False)
+    is_dataframe = isinstance(first_object, pd.DataFrame)
+    is_anndata = isinstance(first_object, ad.AnnData)
+    if not is_dataframe and not is_anndata:
+        raise ValueError(f"Unable to concatenate {suffixes.pop()} objects.")
+
+    objects = [first_object]
+    artifact_uids = [artifacts[0].uid]
+    for artifact in artifacts[1:]:
+        objects.append(artifact.load(is_run_input=False))
+        artifact_uids.append(artifact.uid)
+
+    if is_dataframe:
+        concat_object = pd.concat(objects, join=join, **kwargs)
+    elif is_anndata:
+        label = kwargs.pop("label", "artifact_uid")
+        keys = kwargs.pop("keys", artifact_uids)
+        concat_object = ad.concat(objects, join=join, label=label, keys=keys, **kwargs)
+    return concat_object
+
+
 class Collection(Record, IsVersioned, TracksRun, TracksUpdates):
     """Collections of artifacts.
 
@@ -342,13 +376,25 @@ class Collection(Record, IsVersioned, TracksRun, TracksUpdates):
             run=run,
         )
 
-    def open(self, is_run_input: bool | None = None) -> PyArrowDataset:
-        """Return a cloud-backed pyarrow Dataset.
+    def open(
+        self,
+        engine: Literal["pyarrow", "polars"] = "pyarrow",
+        is_run_input: bool | None = None,
+        **kwargs,
+    ) -> PyArrowDataset | Iterator[PolarsLazyFrame]:
+        """Open a dataset for streaming.
 
-        Works for `pyarrow` compatible formats.
+        Works for `pyarrow` and `polars` compatible formats
+        (`.parquet`, `.csv`, `.ipc` etc. files or directories with such files).
+
+        Args:
+            engine: Which module to use for lazy loading of a dataframe
+                from `pyarrow` or `polars` compatible formats.
+            is_run_input: Whether to track this artifact as run input.
+            **kwargs: Keyword arguments for `pyarrow.dataset.dataset` or `polars.scan_*` functions.
 
         Notes:
-            For more info, see tutorial: :doc:`/arrays`.
+            For more info, see guide: :doc:`/arrays`.
         """
         if self._state.adding:
             artifacts = self._artifacts
@@ -356,31 +402,12 @@ class Collection(Record, IsVersioned, TracksRun, TracksUpdates):
         else:
             artifacts = self.ordered_artifacts.all()
         paths = [artifact.path for artifact in artifacts]
-        # this checks that the filesystem is the same for all paths
-        # this is a requirement of pyarrow.dataset.dataset
-        fs = paths[0].fs
-        for path in paths[1:]:
-            # this assumes that the filesystems are cached by fsspec
-            if path.fs is not fs:
-                raise ValueError(
-                    "The collection has artifacts with different filesystems, this is not supported."
-                )
-        if not _is_pyarrow_dataset(paths):
-            suffixes = {path.suffix for path in paths}
-            suffixes_str = ", ".join(suffixes)
-            err_msg = (
-                "This collection is not compatible with pyarrow.dataset.dataset(), "
-            )
-            err_msg += (
-                f"the artifacts have incompatible file types: {suffixes_str}"
-                if len(suffixes) > 1
-                else f"the file type {suffixes_str} is not supported by pyarrow."
-            )
-            raise ValueError(err_msg)
-        dataset = _open_pyarrow_dataset(paths)
+
+        dataframe = _open_dataframe(paths, engine=engine, **kwargs)
         # track only if successful
+        # is it really needed if tracking is done in self.ordered_artifacts.all()? - Sergei
         _track_run_input(self, is_run_input)
-        return dataset
+        return dataframe
 
     def mapped(
         self,
@@ -403,8 +430,8 @@ class Collection(Record, IsVersioned, TracksRun, TracksUpdates):
         <https://pytorch.org/docs/stable/data.html#map-style-datasets>`__ by
         virtually concatenating `AnnData` arrays.
 
-        If your `AnnData` collection is in the cloud, move them into a local
-        cache first via :meth:`~lamindb.Collection.cache`.
+        By default (`stream=False`) `AnnData` arrays are moved into a local
+        cache first.
 
         `__getitem__` of the `MappedCollection` object takes a single integer index
         and returns a dictionary with the observation data sample for this index from
@@ -416,7 +443,7 @@ class Collection(Record, IsVersioned, TracksRun, TracksUpdates):
 
             For a guide, see :doc:`docs:scrna-mappedcollection`.
 
-            This method currently only works for collections of `AnnData` artifacts.
+            This method currently only works for collections or query sets of `AnnData` artifacts.
 
         Args:
             layers_keys: Keys from the ``.layers`` slot. ``layers_keys=None`` or ``"X"`` in the list
@@ -445,6 +472,11 @@ class Collection(Record, IsVersioned, TracksRun, TracksUpdates):
             >>> ds = ln.Collection.get(description="my collection")
             >>> mapped = collection.mapped(obs_keys=["cell_type", "batch"])
             >>> dl = DataLoader(mapped, batch_size=128, shuffle=True)
+            >>> # also works for query sets of artifacts, '...' represents some filtering condition
+            >>> # additional filtering on artifacts of the collection
+            >>> mapped = collection.artifacts.all().filter(...).order_by("-created_at").mapped()
+            >>> # or directly from a query set of artifacts
+            >>> mapped = ln.Artifact.filter(..., otype="AnnData").order_by("-created_at").mapped()
         """
         path_list = []
         if self._state.adding:
@@ -474,6 +506,7 @@ class Collection(Record, IsVersioned, TracksRun, TracksUpdates):
             dtype,
         )
         # track only if successful
+        # is it really needed if tracking is done in self.ordered_artifacts.all()? - Sergei
         _track_run_input(self, is_run_input)
         return ds
 
@@ -490,6 +523,7 @@ class Collection(Record, IsVersioned, TracksRun, TracksUpdates):
         path_list = []
         for artifact in self.ordered_artifacts.all():
             path_list.append(artifact.cache())
+        # is it really needed if tracking is done in self.ordered_artifacts.all()? - Sergei
         _track_run_input(self, is_run_input)
         return path_list
 
@@ -498,29 +532,16 @@ class Collection(Record, IsVersioned, TracksRun, TracksUpdates):
         join: Literal["inner", "outer"] = "outer",
         is_run_input: bool | None = None,
         **kwargs,
-    ) -> Any:
-        """Stage and load to memory.
+    ) -> pd.DataFrame | ad.AnnData:
+        """Cache and load to memory.
 
-        Returns in-memory representation if possible such as a concatenated `DataFrame` or `AnnData` object.
+        Returns an in-memory concatenated `DataFrame` or `AnnData` object.
         """
         # cannot call _track_run_input here, see comment further down
-        all_artifacts = self.ordered_artifacts.all()
-        suffixes = [artifact.suffix for artifact in all_artifacts]
-        if len(set(suffixes)) != 1:
-            raise RuntimeError(
-                "Can only load collections where all artifacts have the same suffix"
-            )
-        # because we're tracking data flow on the collection-level, here, we don't
-        # want to track it on the artifact-level
-        objects = [artifact.load(is_run_input=False) for artifact in all_artifacts]
-        artifact_uids = [artifact.uid for artifact in all_artifacts]
-        if isinstance(objects[0], pd.DataFrame):
-            concat_object = pd.concat(objects, join=join)
-        elif isinstance(objects[0], ad.AnnData):
-            concat_object = ad.concat(
-                objects, join=join, label="artifact_uid", keys=artifact_uids
-            )
-        # only call it here because there might be errors during concat
+        artifacts = self.ordered_artifacts.all()
+        concat_object = _load_concat_artifacts(artifacts, join, **kwargs)
+        # only call it here because there might be errors during load or concat
+        # is it really needed if tracking is done in self.ordered_artifacts.all()? - Sergei
         _track_run_input(self, is_run_input)
         return concat_object
 

lamindb/models/core.py

@@ -24,7 +24,7 @@ if TYPE_CHECKING:
 
 
 class Storage(Record, TracksRun, TracksUpdates):
-    """Storage locations.
+    """Storage locations of artifacts such as S3 buckets or local directories.
 
     A storage location is either a directory/folder (local or in the cloud) or
     an entire S3/GCP bucket.

lamindb/models/feature.py

@@ -143,40 +143,50 @@ def parse_cat_dtype(
 
 
 def serialize_dtype(
-    dtype: Record | FieldAttr | list[Record], is_itype: bool = False
+    dtype: Registry | Record | FieldAttr | list[Record] | list[Registry] | str,
+    is_itype: bool = False,
 ) -> str:
     """Converts a data type object into its string representation."""
+    from .ulabel import ULabel
+
     if (
         not isinstance(dtype, list)
         and hasattr(dtype, "__name__")
         and dtype.__name__ in FEATURE_DTYPES
     ):
         dtype_str = dtype.__name__
+    elif dtype is dict:
+        dtype_str = "dict"
+    elif is_itype and isinstance(dtype, str):
+        if dtype not in "Feature":
+            parse_cat_dtype(
+                dtype_str=dtype, is_itype=True
+            )  # throws an error if invalid
+        dtype_str = dtype
     elif isinstance(dtype, (ExtensionDtype, np.dtype)):
         dtype_str = serialize_pandas_dtype(dtype)
     else:
-        error_message = (
-            "dtype has to be a record, a record field, or a list of records, not {}"
-        )
-        if isinstance(dtype, Registry):
-            dtype = [dtype]
-        elif isinstance(dtype, DeferredAttribute):
+        error_message = "dtype has to be a registry, a ulabel subtype, a registry field, or a list of registries or fields, not {}"
+        if isinstance(dtype, (Registry, DeferredAttribute, ULabel)):
             dtype = [dtype]
         elif not isinstance(dtype, list):
             raise ValueError(error_message.format(dtype))
         dtype_str = ""
-        for single_dtype in dtype:
-            if not isinstance(single_dtype, Registry) and not isinstance(
-                single_dtype, DeferredAttribute
-            ):
-                raise ValueError(error_message.format(single_dtype))
-            if isinstance(single_dtype, Registry):
-                dtype_str += single_dtype.__get_name_with_module__() + "|"
+        for one_dtype in dtype:
+            if not isinstance(one_dtype, (Registry, DeferredAttribute, ULabel)):
+                raise ValueError(error_message.format(one_dtype))
+            if isinstance(one_dtype, Registry):
+                dtype_str += one_dtype.__get_name_with_module__() + "|"
+            elif isinstance(one_dtype, ULabel):
+                assert one_dtype.is_type, (  # noqa: S101
+                    f"ulabel has to be a type if acting as dtype, {one_dtype} has `is_type` False"
+                )
+                dtype_str += f"ULabel[{one_dtype.name}]"
             else:
+                name = one_dtype.field.name
+                field_ext = f".{name}" if name != "name" else ""
                 dtype_str += (
-                    single_dtype.field.model.__get_name_with_module__()
-                    + f".{single_dtype.field.name}"
-                    + "|"
+                    one_dtype.field.model.__get_name_with_module__() + field_ext + "|"
                 )
         dtype_str = dtype_str.rstrip("|")
         if not is_itype:
@@ -571,6 +581,30 @@ class Feature(Record, CanCurate, TracksRun, TracksUpdates):
         self._aux = self._aux or {}
         self._aux.setdefault("af", {})["2"] = value
 
+    # we'll enable this later
+    # @property
+    # def observational_unit(self) -> Literal["Artifact", "Observation"]:
+    #     """Default observational unit on which the feature is measured.
+
+    #     Currently, we only make a distinction between artifact-level and observation-level features.
+
+    #     For example, a feature `"ml_split"` that stores `"test"` & `"train"` labels is typically defined on the artifact level.
+    #     When accessing `artifact.features.get_values(["ml_split"])`, you expect a single value, either `"test"` or `"train"`.
+
+    #     However, when accessing an artifact annotation with a feature that's defined on the observation-level, say `"cell_type"`, you expect a set of values. So,
+    #     `artifact.features.get_values(["cell_type_from_expert"])` should return a set: `{"T cell", "B cell"}`.
+
+    #     The value of `observational_unit` is currently auto-managed: if using `artifact.featueres.add_values()`,
+    #     it will be set to `Artifact`. In a curator, the value depends on whether it's an artifact- or observation-level slot
+    #     (e.g. `.uns` is artifact-level in `AnnData` whereas `.obs` is observation-level).
+
+    #     Note: This attribute might in the future be used to distinguish different types of observational units (e.g. single cells vs. physical samples vs. study subjects etc.).
+    #     """
+    #     if self._expect_many:
+    #         return "Observation"  # this here might be replaced with the specific observational unit
+    #     else:
+    #         return "Artifact"
+
 
 class FeatureValue(Record, TracksRun):
     """Non-categorical features values.

lamindb/models/has_parents.py

@@ -84,10 +84,44 @@ class HasParents:
         return view_parents(
             record=self,  # type: ignore
             field=field,
+            with_parents=True,
             with_children=with_children,
             distance=distance,
         )
 
+    def view_children(
+        self,
+        field: StrField | None = None,
+        distance: int = 5,
+    ):
+        """View children in an ontology.
+
+        Args:
+            field: Field to display on graph
+            distance: Maximum distance still shown.
+
+        Ontological hierarchies: :class:`~lamindb.ULabel` (project & sub-project), :class:`~bionty.CellType` (cell type & subtype).
+
+        Examples:
+            >>> import bionty as bt
+            >>> bt.Tissue.from_source(name="subsegmental bronchus").save()
+            >>> record = bt.Tissue.get(name="respiratory tube")
+            >>> record.view_parents()
+            >>> tissue.view_parents(with_children=True)
+        """
+        if field is None:
+            field = get_name_field(self)
+        if not isinstance(field, str):
+            field = field.field.name
+
+        return view_parents(
+            record=self,  # type: ignore
+            field=field,
+            with_parents=False,
+            with_children=True,
+            distance=distance,
+        )
+
     def query_parents(self) -> QuerySet:
         """Query parents in an ontology."""
         return _query_relatives([self], "parents", self.__class__)  # type: ignore
@@ -210,6 +244,7 @@ def view_lineage(
 def view_parents(
     record: Record,
     field: str,
+    with_parents: bool = True,
     with_children: bool = False,
     distance: int = 100,
     attr_name: Literal["parents", "predecessors"] = "parents",
@@ -223,11 +258,12 @@ def view_parents(
     import pandas as pd
 
     df_edges = None
-    df_edges_parents = _df_edges_from_parents(
-        record=record, field=field, distance=distance, attr_name=attr_name
-    )
-    if df_edges_parents is not None:
-        df_edges = df_edges_parents
+    df_edges_parents = None
+    df_edges_children = None
+    if with_parents:
+        df_edges_parents = _df_edges_from_parents(
+            record=record, field=field, distance=distance, attr_name=attr_name
+        )
     if with_children:
         df_edges_children = _df_edges_from_parents(
             record=record,
@@ -236,13 +272,32 @@ def view_parents(
             children=True,
             attr_name=attr_name,
         )
-        if df_edges_children is not None:
-            if df_edges is not None:
-                df_edges = pd.concat(
-                    [df_edges_parents, df_edges_children]
-                ).drop_duplicates()
-            else:
-                df_edges = df_edges_children
+        # Rename the columns to swap source and target
+        df_edges_children = df_edges_children.rename(
+            columns={
+                "source": "temp_target",
+                "source_label": "temp_target_label",
+                "source_record": "temp_target_record",
+                "target": "source",
+                "target_label": "source_label",
+                "target_record": "source_record",
+            }
+        )
+        df_edges_children = df_edges_children.rename(
+            columns={
+                "temp_target": "target",
+                "temp_target_label": "target_label",
+                "temp_target_record": "target_record",
+            }
+        )
+    if df_edges_parents is not None and df_edges_children is not None:
+        df_edges = pd.concat([df_edges_parents, df_edges_children]).drop_duplicates()
+    elif df_edges_parents is not None:
+        df_edges = df_edges_parents
+    elif df_edges_children is not None:
+        df_edges = df_edges_children
+    else:
+        return None
 
     record_label = _record_label(record, field)
 
@@ -520,14 +575,14 @@ def _get_all_child_runs(data: Artifact | Collection) -> list:
             run_inputs_outputs += [(r, outputs_run)]
 
             child_runs.update(
-                Run.filter(
+                Run.filter(  # type: ignore
                     **{f"input_{name}s__uid__in": [i.uid for i in outputs_run]}
                 ).list()
             )
             # for artifacts, also include collections in the lineage
             if name == "artifact":
                 child_runs.update(
-                    Run.filter(
+                    Run.filter(  # type: ignore
                         input_collections__uid__in=[i.uid for i in outputs_run]
                     ).list()
                 )

lamindb/models/project.py

@@ -36,7 +36,7 @@ if TYPE_CHECKING:
 
 
 class Person(Record, CanCurate, TracksRun, TracksUpdates, ValidateFields):
-    """People.
+    """People such as authors of a study or collaborators in a project.
 
     This registry is distinct from `User` and exists for project management.