lamindb 0.76.9__py3-none-any.whl → 0.76.11__py3-none-any.whl

lamindb/core/_django.py

@@ -0,0 +1,209 @@
+from django.contrib.postgres.aggregates import ArrayAgg
+from django.db import connection
+from django.db.models import F, OuterRef, Q, Subquery
+from django.db.models.fields.related import ForeignKey, ManyToManyField
+from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel
+from django.db.models.functions import JSONObject
+from lnschema_core.models import Artifact, FeatureSet, Record
+
+from .schema import dict_related_model_to_related_name
+
+
+def get_related_model(model, field_name):
+    try:
+        field = model._meta.get_field(field_name)
+
+        if isinstance(field, (ForeignKey, ManyToManyField)):
+            # Forward ForeignKey or ManyToManyField
+            return field.remote_field.model
+        elif isinstance(field, (ManyToOneRel, ManyToManyRel)):
+            # Reverse ForeignKey or ManyToManyField
+            return field.related_model
+        else:
+            return f"Unexpected field type: {type(field)}"
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+def get_artifact_with_related(
+    artifact: Record,
+    include_fk: bool = False,
+    include_m2m: bool = False,
+    include_feature_link: bool = False,
+    include_featureset: bool = False,
+) -> dict:
+    """Fetch an artifact with its related data."""
+    from lamindb._can_validate import get_name_field
+
+    model = artifact.__class__
+    foreign_key_fields = [f.name for f in model._meta.fields if f.is_relation]
+
+    m2m_relations = (
+        []
+        if not include_m2m
+        else [
+            v
+            for v in dict_related_model_to_related_name(model).values()
+            if not v.startswith("_")
+        ]
+    )
+    link_tables = (
+        []
+        if not include_feature_link
+        else list(dict_related_model_to_related_name(model, links=True).values())
+    )
+
+    # Clear previous queries
+    connection.queries_log.clear()
+
+    annotations = {}
+
+    if include_fk:
+        for fk in foreign_key_fields:
+            name_field = get_name_field(get_related_model(model, fk))
+            annotations[f"fkfield_{fk}"] = JSONObject(
+                id=F(f"{fk}__id"), name=F(f"{fk}__{name_field}")
+            )
+
+    for name in m2m_relations:
+        related_model = get_related_model(model, name)
+        name_field = get_name_field(related_model)
+        annotations[f"m2mfield_{name}"] = ArrayAgg(
+            JSONObject(id=F(f"{name}__id"), name=F(f"{name}__{name_field}")),
+            filter=Q(**{f"{name}__isnull": False}),
+            distinct=True,
+        )
+
+    for link in link_tables:
+        link_model = getattr(model, link).rel.related_model
+        if not hasattr(link_model, "feature"):
+            continue
+        label_field = link.removeprefix("links_").replace("_", "")
+        annotations[f"linkfield_{link}"] = Subquery(
+            link_model.objects.filter(artifact=OuterRef("pk"))
+            .annotate(
+                data=JSONObject(
+                    id=F("id"),
+                    feature=F("feature"),
+                    **{label_field: F(label_field)},
+                )
+            )
+            .values("artifact")
+            .annotate(json_agg=ArrayAgg("data"))
+            .values("json_agg")
+        )
+
+    if include_featureset:
+        annotations["featuresets"] = Subquery(
+            model.feature_sets.through.objects.filter(artifact=OuterRef("pk"))
+            .annotate(
+                data=JSONObject(
+                    id=F("id"),
+                    slot=F("slot"),
+                    featureset=F("featureset"),
+                )
+            )
+            .values("artifact")
+            .annotate(json_agg=ArrayAgg("data"))
+            .values("json_agg")
+        )
+
+    artifact_meta = (
+        model.objects.using(artifact._state.db)
+        .filter(uid=artifact.uid)
+        .annotate(**annotations)
+        .values(*["id", "uid"], *annotations.keys())
+        .first()
+    )
+
+    if not artifact_meta:
+        return None
+
+    related_data: dict = {"m2m": {}, "fk": {}, "link": {}, "featuresets": {}}
+    for k, v in artifact_meta.items():
+        if k.startswith("m2mfield_"):
+            related_data["m2m"][k[9:]] = v
+        elif k.startswith("fkfield_"):
+            related_data["fk"][k[8:]] = v
+        elif k.startswith("linkfield_"):
+            related_data["link"][k[10:]] = v
+        elif k == "featuresets":
+            if v:
+                related_data["featuresets"] = get_featureset_m2m_relations(
+                    artifact, {i["featureset"]: i["slot"] for i in v}
+                )
+
+    related_data["m2m"] = {
+        k: {item["id"]: item["name"] for item in v}
+        for k, v in related_data["m2m"].items()
+        if v
+    }
+
+    return {
+        **{name: artifact_meta[name] for name in ["id", "uid"]},
+        "related_data": related_data,
+    }
+
+
+def get_featureset_m2m_relations(
+    artifact: Artifact, slot_featureset: dict, limit: int = 20
+):
+    """Fetch all many-to-many relationships for given feature sets."""
+    from lamindb._can_validate import get_name_field
+
+    m2m_relations = [
+        v
+        for v in dict_related_model_to_related_name(FeatureSet).values()
+        if not v.startswith("_") and v != "artifacts"
+    ]
+
+    annotations = {}
+    related_names = {}
+    for name in m2m_relations:
+        related_model = get_related_model(FeatureSet, name)
+        name_field = get_name_field(related_model)
+
+        # Get the correct field names for the through table
+        through_model = getattr(FeatureSet, name).through
+        related_field = (
+            through_model.__name__.replace("FeatureSet", "").lower().replace("_", "")
+        )
+
+        # Subquery to get limited related records
+        limited_related = Subquery(
+            through_model.objects.filter(featureset=OuterRef("pk")).values(
+                related_field
+            )[:limit]
+        )
+
+        annotations[f"m2mfield_{name}"] = ArrayAgg(
+            JSONObject(id=F(f"{name}__id"), name=F(f"{name}__{name_field}")),
+            filter=Q(
+                **{
+                    f"{name}__id__in": limited_related,
+                }
+            ),
+            distinct=True,
+        )
+        related_names[name] = related_model.__get_name_with_schema__()
+
+    featureset_m2m = (
+        FeatureSet.objects.using(artifact._state.db)
+        .filter(id__in=slot_featureset.keys())
+        .annotate(**annotations)
+        .values("id", *annotations.keys())
+    )
+
+    result = {}
+    for fs in featureset_m2m:
+        slot = slot_featureset.get(fs["id"])
+        result[fs["id"]] = (
+            slot,
+            {
+                related_names.get(k[9:]): [item["name"] for item in v]
+                for k, v in fs.items()
+                if k.startswith("m2mfield_") and v
+            },
+        )
+
+    return result

lamindb/core/_feature_manager.py

@@ -43,6 +43,7 @@ from lamindb._save import save
 from lamindb.core.exceptions import ValidationError
 from lamindb.core.storage import LocalPathClasses
 
+from ._django import get_artifact_with_related
 from ._label_manager import get_labels_as_dict
 from ._settings import settings
 from .schema import (
@@ -132,12 +133,81 @@ def custom_aggregate(field, using: str):
         return GroupConcat(field)
 
 
-def print_features(
+def _print_categoricals_postgres(
     self: Artifact | Collection,
+    related_data: dict | None = None,
     print_types: bool = False,
     to_dict: bool = False,
     print_params: bool = False,
-) -> str | dict[str, Any]:
+):
+    from lamindb._from_values import _print_values
+
+    if not related_data:
+        artifact_meta = get_artifact_with_related(
+            self, include_feature_link=True, include_m2m=True
+        )
+        related_data = artifact_meta.get("related_data", {})
+
+    m2m_data = related_data.get("m2m", {}) if related_data else {}
+    m2m_name = {}
+    for related_name, values in m2m_data.items():
+        link_model = getattr(self.__class__, related_name).through
+        related_model_name = link_model.__name__.replace(
+            self.__class__.__name__, ""
+        ).lower()
+        m2m_name[related_model_name] = values
+    links_data = related_data.get("link", {}) if related_data else {}
+    feature_dict = {
+        id: (name, dtype)
+        for id, name, dtype in Feature.objects.using(self._state.db).values_list(
+            "id", "name", "dtype"
+        )
+    }
+
+    msg = ""
+    dictionary = {}
+
+    # categorical feature values
+    if not print_params:
+        labels_msg = ""
+        labels_msgs = []
+        feature_values: dict = {}
+        for link_name, link_values in links_data.items():
+            related_name = link_name.removeprefix("links_").replace("_", "")
+            link_model = getattr(self.__class__, link_name).rel.related_model
+            if not link_values:
+                continue
+            for link_value in link_values:
+                feature_id = link_value.get("feature")
+                if feature_id is None:
+                    continue
+                feature_name = feature_dict.get(feature_id)[0]
+                if feature_name not in feature_values:
+                    feature_values[feature_name] = (feature_dict.get(feature_id)[1], [])
+                label_id = link_value.get(related_name)
+                feature_values[feature_name][1].append(
+                    m2m_name.get(related_name, {}).get(label_id)
+                )
+        for feature_name, (dtype, labels_list) in feature_values.items():
+            print_values = _print_values(labels_list, n=10)
+            type_str = f": {dtype}" if print_types else ""
+            if to_dict:
+                dictionary[feature_name] = (
+                    labels_list if len(labels_list) > 1 else labels_list[0]
+                )
+            labels_msgs.append(f"    '{feature_name}'{type_str} = {print_values}")
+        if len(labels_msgs) > 0:
+            labels_msg = "\n".join(sorted(labels_msgs)) + "\n"
+            msg += labels_msg
+    return msg, dictionary
+
+
+def _print_categoricals(
+    self: Artifact | Collection,
+    print_types: bool = False,
+    to_dict: bool = False,
+    print_params: bool = False,
+):
     from lamindb._from_values import _print_values
 
     msg = ""
@@ -166,6 +236,56 @@ def print_features(
         if len(labels_msgs) > 0:
             labels_msg = "\n".join(sorted(labels_msgs)) + "\n"
             msg += labels_msg
+    return msg, dictionary
+
+
+def _print_featuresets_postgres(
+    self: Artifact | Collection,
+    related_data: dict | None = None,
+    print_types: bool = False,
+):
+    from lamindb._from_values import _print_values
+
+    if not related_data:
+        artifact_meta = get_artifact_with_related(self, include_featureset=True)
+        related_data = artifact_meta.get("related_data", {})
+
+    fs_data = related_data.get("featuresets", {}) if related_data else {}
+    feature_set_msg = ""
+    for _, (slot, data) in fs_data.items():
+        for type_str, feature_names in data.items():
+            type_str = f": {type_str}" if print_types else ""
+            feature_set_msg += (
+                f"    '{slot}'{type_str} = {_print_values(feature_names)}\n"
+            )
+
+    return feature_set_msg
+
+
+def print_features(
+    self: Artifact | Collection,
+    related_data: dict | None = None,
+    print_types: bool = False,
+    to_dict: bool = False,
+    print_params: bool = False,
+) -> str | dict[str, Any]:
+    from lamindb._from_values import _print_values
+
+    if not self._state.adding and connections[self._state.db].vendor == "postgresql":
+        msg, dictionary = _print_categoricals_postgres(
+            self,
+            related_data=related_data,
+            print_types=print_types,
+            to_dict=to_dict,
+            print_params=print_params,
+        )
+    else:
+        msg, dictionary = _print_categoricals(
+            self,
+            print_types=print_types,
+            to_dict=to_dict,
+            print_params=print_params,
+        )
 
     # non-categorical feature values
     non_labels_msg = ""
@@ -203,15 +323,20 @@ def print_features(
     # feature sets
     if not print_params:
         feature_set_msg = ""
-        for slot, feature_set in get_feature_set_by_slot_(self).items():
-            features = feature_set.members
-            # features.first() is a lot slower than features[0] here
-            name_field = get_name_field(features[0])
-            feature_names = list(features.values_list(name_field, flat=True)[:20])
-            type_str = f": {feature_set.registry}" if print_types else ""
-            feature_set_msg += (
-                f"    '{slot}'{type_str} = {_print_values(feature_names)}\n"
+        if self.id is not None and connections[self._state.db].vendor == "postgresql":
+            feature_set_msg = _print_featuresets_postgres(
+                self, related_data=related_data
             )
+        else:
+            for slot, feature_set in get_feature_set_by_slot_(self).items():
+                features = feature_set.members
+                # features.first() is a lot slower than features[0] here
+                name_field = get_name_field(features[0])
+                feature_names = list(features.values_list(name_field, flat=True)[:20])
+                type_str = f": {feature_set.registry}" if print_types else ""
+                feature_set_msg += (
+                    f"    '{slot}'{type_str} = {_print_values(feature_names)}\n"
+                )
         if feature_set_msg:
             msg += f"  {colors.italic('Feature sets')}\n"
             msg += feature_set_msg
@@ -372,6 +497,7 @@ def filter_base(cls, **expression):
         )
     new_expression = {}
     features = model.filter(name__in=keys_normalized).all().distinct()
+    feature_param = "param" if model is Param else "feature"
     for key, value in expression.items():
         split_key = key.split("__")
         normalized_key = split_key[0]
@@ -380,9 +506,9 @@ def filter_base(cls, **expression):
             comparator = f"__{split_key[1]}"
         feature = features.get(name=normalized_key)
         if not feature.dtype.startswith("cat"):
-            expression = {"feature": feature, f"value{comparator}": value}
+            expression = {feature_param: feature, f"value{comparator}": value}
             feature_value = value_model.filter(**expression)
-            new_expression["_feature_values__in"] = feature_value
+            new_expression[f"_{feature_param}_values__in"] = feature_value
         else:
             if isinstance(value, str):
                 expression = {f"name{comparator}": value}
@@ -792,7 +918,7 @@ def _add_from(self, data: Artifact | Collection, transfer_logs: dict = None):
     """Transfer features from a artifact or collection."""
     # This only covers feature sets
     if transfer_logs is None:
-        transfer_logs = {"mapped": [], "transferred": []}
+        transfer_logs = {"mapped": [], "transferred": [], "run": None}
     using_key = settings._using_key
     for slot, feature_set in data.features._feature_set_by_slot.items():
         members = feature_set.members
@@ -865,3 +991,4 @@ FeatureManager.filter = filter
 FeatureManager.get = get
 ParamManager.add_values = add_values_params
 ParamManager.get_values = get_values
+ParamManager.filter = filter

lamindb/core/_label_manager.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING
 
 import numpy as np
-from lamin_utils import colors
-from lnschema_core.models import Feature
+from django.db import connections
+from lamin_utils import colors, logger
+from lnschema_core.models import CanValidate, Feature
 
 from lamindb._from_values import _print_values
 from lamindb._record import (
@@ -16,6 +17,7 @@ from lamindb._record import (
 )
 from lamindb._save import save
 
+from ._django import get_artifact_with_related, get_related_model
 from ._settings import settings
 from .schema import dict_related_model_to_related_name
 
@@ -57,22 +59,46 @@ def get_labels_as_dict(self: Artifact | Collection, links: bool = False):
     return labels
 
 
+def _print_labels_postgres(
+    self: Artifact | Collection, m2m_data: dict | None = None, print_types: bool = False
+) -> str:
+    labels_msg = ""
+    if not m2m_data:
+        artifact_meta = get_artifact_with_related(self, include_m2m=True)
+        m2m_data = artifact_meta.get("related_data", {}).get("m2m", {})
+    if m2m_data:
+        for related_name, labels in m2m_data.items():
+            if not labels or related_name == "feature_sets":
+                continue
+            related_model = get_related_model(self, related_name)
+            print_values = _print_values(labels.values(), n=10)
+            type_str = f": {related_model}" if print_types else ""
+            labels_msg += f"    .{related_name}{type_str} = {print_values}\n"
+    return labels_msg
+
+
 def print_labels(
-    self: Artifact | Collection, field: str = "name", print_types: bool = False
+    self: Artifact | Collection,
+    m2m_data: dict | None = None,
+    print_types: bool = False,
 ):
-    labels_msg = ""
-    for related_name, (related_model, labels) in get_labels_as_dict(self).items():
-        # there is a try except block here to deal with schema inconsistencies
-        # during transfer between instances
-        try:
-            labels_list = list(labels.values_list(field, flat=True))
-            if len(labels_list) > 0:
-                get_name_field(labels)
-                print_values = _print_values(labels_list, n=10)
-                type_str = f": {related_model}" if print_types else ""
-                labels_msg += f"    .{related_name}{type_str} = {print_values}\n"
-        except Exception:  # noqa: S112
-            continue
+    if not self._state.adding and connections[self._state.db].vendor == "postgresql":
+        labels_msg = _print_labels_postgres(self, m2m_data, print_types)
+    else:
+        labels_msg = ""
+        for related_name, (related_model, labels) in get_labels_as_dict(self).items():
+            # there is a try except block here to deal with schema inconsistencies
+            # during transfer between instances
+            try:
+                field = get_name_field(labels)
+                labels_list = list(labels.values_list(field, flat=True))
+                if len(labels_list) > 0:
+                    print_values = _print_values(labels_list, n=10)
+                    type_str = f": {related_model}" if print_types else ""
+                    labels_msg += f"    .{related_name}{type_str} = {print_values}\n"
+            except Exception:  # noqa: S112
+                continue
+
     msg = ""
     if labels_msg:
         msg += f"  {colors.italic('Labels')}\n"
@@ -106,13 +132,19 @@ def validate_labels(labels: QuerySet | list | dict):
             label_uids = np.array(
                 [getattr(label, field) for label in labels if label is not None]
             )
-        validated = registry.validate(label_uids, field=field, mute=True)
-        validated_uids = label_uids[validated]
-        validated_labels = registry.filter(**{f"{field}__in": validated_uids}).list()
-        new_labels = [labels[int(i)] for i in np.argwhere(~validated).flatten()]
+        if issubclass(registry, CanValidate):
+            validated = registry.validate(label_uids, field=field, mute=True)
+            validated_uids = label_uids[validated]
+            validated_labels = registry.filter(
+                **{f"{field}__in": validated_uids}
+            ).list()
+            new_labels = [labels[int(i)] for i in np.argwhere(~validated).flatten()]
+        else:
+            validated_labels = []
+            new_labels = list(labels)
         return validated_labels, new_labels
 
-    if isinstance(labels, Dict):
+    if isinstance(labels, dict):
         result = {}
         for registry, labels_registry in labels.items():
             result[registry] = validate_labels_registry(labels_registry)
@@ -185,7 +217,7 @@ class LabelManager:
         from django.db.utils import ProgrammingError
 
         if transfer_logs is None:
-            transfer_logs = {"mapped": [], "transferred": []}
+            transfer_logs = {"mapped": [], "transferred": [], "run": None}
         using_key = settings._using_key
         for related_name, (_, labels) in get_labels_as_dict(data).items():
             labels = labels.all()
@@ -250,4 +282,7 @@ class LabelManager:
                         )
             # ProgrammingError is raised when schemas don't match between source and target instances
             except ProgrammingError:
+                logger.warning(
+                    f"{related_name} labels cannot be transferred because schema module does not exist in target instance: {labels}"
+                )
                 continue

lamindb/core/_mapped_collection.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from collections import Counter
 from functools import reduce
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 import pandas as pd

lamindb/core/_settings.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Literal, Mapping
+from typing import TYPE_CHECKING, Literal
 
 import lamindb_setup as ln_setup
 from lamin_utils import logger
@@ -13,6 +13,7 @@ from .subsettings._creation_settings import CreationSettings, creation_settings
 from .subsettings._transform_settings import TransformSettings, transform_settings
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
     from pathlib import Path
 
     from lamindb_setup.core._settings_storage import StorageSettings

lamindb/core/exceptions.py

@@ -6,7 +6,7 @@
    InvalidArgument
    DoesNotExist
    ValidationError
-   NotebookNotSavedError
+   NotebookNotSaved
    NoTitleError
    MissingContextUID
    UpdateContext
@@ -26,7 +26,7 @@ class InvalidArgument(SystemExit):
 
 
 class TrackNotCalled(SystemExit):
-    """ln.context.track() wasn't called."""
+    """`ln.track()` wasn't called."""
 
     pass
 
@@ -51,8 +51,14 @@ class DoesNotExist(Exception):
     pass
 
 
+class InconsistentKey(Exception):
+    """Inconsistent transform or artifact `key`."""
+
+    pass
+
+
 # -------------------------------------------------------------------------------------
-# ln.context.track() AKA context
+# run context
 # -------------------------------------------------------------------------------------
 
 
@@ -66,12 +72,6 @@ class IntegrityError(Exception):
     pass
 
 
-class NotebookNotSavedError(Exception):
-    """Notebook wasn't saved."""
-
-    pass
-
-
 class NoTitleError(Exception):
     """Notebook has no title."""
 

lamindb/core/storage/_anndata_accessor.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import inspect
 from functools import cached_property
 from itertools import chain
-from typing import TYPE_CHECKING, Callable, Literal, Mapping, Union
+from typing import TYPE_CHECKING, Callable, Literal, Union
 
 import h5py
 import numpy as np
@@ -21,6 +21,7 @@ from lamindb_setup.core.upath import UPath, create_mapper, infer_filesystem
 from packaging import version
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
     from pathlib import Path
 
     from fsspec.core import OpenFile

lamindb/core/versioning.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Literal
 
 from lamin_utils import logger
-from lamin_utils._base62 import CHARSET_DEFAULT as BASE62_CHARS
+from lamin_utils._base62 import increment_base62
 from lamindb_setup.core.upath import LocalPathClasses, UPath
 from lnschema_core import ids
 
@@ -18,19 +18,7 @@ def message_update_key_in_version_family(
     registry: str,
     new_key: str,
 ) -> str:
-    return f'Or update key "{existing_key}" to "{new_key}":\n\nln.{registry}.filter(uid__startswith="{suid}").update(key="{new_key}")\n'
-
-
-def increment_base62(s: str) -> str:
-    # we don't need to throw an error for zzzz because uids are enforced to be unique
-    # on the db level and have an enforced maximum length
-    value = sum(BASE62_CHARS.index(c) * (62**i) for i, c in enumerate(reversed(s)))
-    value += 1
-    result = ""
-    while value:
-        value, remainder = divmod(value, 62)
-        result = BASE62_CHARS[remainder] + result
-    return result.zfill(len(s))
+    return f'Or update key "{existing_key}" to "{new_key}" for all previous versions:\n\nln.{registry}.filter(uid__startswith="{suid}").update(key="{new_key}")\n'
 
 
 def bump_version(