lamindb 0.77.0__py3-none-any.whl → 0.77.2__py3-none-any.whl

lamindb/_query_set.py

@@ -1,19 +1,22 @@
 from __future__ import annotations
 
+import re
 from collections import UserList
 from collections.abc import Iterable
 from collections.abc import Iterable as IterableType
-from typing import TYPE_CHECKING, Any, NamedTuple
+from typing import TYPE_CHECKING, Any, Generic, NamedTuple, TypeVar
 
 import pandas as pd
 from django.db import models
-from django.db.models import F
-from lamin_utils import colors, logger
+from django.db.models import F, ForeignKey, ManyToManyField
+from django.db.models.fields.related import ForeignObjectRel
+from lamin_utils import logger
 from lamindb_setup.core._docs import doc_args
 from lnschema_core.models import (
     Artifact,
     CanCurate,
     Collection,
+    Feature,
     IsVersioned,
     Record,
     Registry,
@@ -24,6 +27,8 @@ from lnschema_core.models import (
 
 from .core.exceptions import DoesNotExist
 
+T = TypeVar("T")
+
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
@@ -34,6 +39,9 @@ class MultipleResultsFound(Exception):
     pass
 
 
+pd.set_option("display.max_columns", 200)
+
+
 # def format_and_convert_to_local_time(series: pd.Series):
 #     tzinfo = datetime.now().astimezone().tzinfo
 #     timedelta = tzinfo.utcoffset(datetime.now())  # type: ignore
@@ -155,21 +163,295 @@ def get(
         return registry.objects.using(qs.db).get(**expressions)
 
 
-class RecordsList(UserList):
+class RecordList(UserList, Generic[T]):
     """Is ordered, can't be queried, but has `.df()`."""
 
-    def __init__(self, records: Iterable[Record]):
-        super().__init__(record for record in records)
+    def __init__(self, records: Iterable[T]):
+        if isinstance(records, list):
+            self.data = records  # Direct assignment if already a list, no copy
+        else:
+            super().__init__(records)  # Let UserList handle the conversion
 
     def df(self) -> pd.DataFrame:
         keys = get_keys_from_df(self.data, self.data[0].__class__)
         values = [record.__dict__ for record in self.data]
         return pd.DataFrame(values, columns=keys)
 
-    def one(self) -> Record:
+    def one(self) -> T:
         """Exactly one result. Throws error if there are more or none."""
         return one_helper(self)
 
+    def save(self) -> RecordList[T]:
+        """Save all records to the database."""
+        from lamindb._save import save
+
+        save(self)
+        return self
+
+
+def get_basic_field_names(
+    qs: QuerySet, include: list[str], features: bool | list[str] = False
+) -> list[str]:
+    exclude_field_names = ["updated_at"]
+    field_names = [
+        field.name
+        for field in qs.model._meta.fields
+        if (
+            not isinstance(field, models.ForeignKey)
+            and field.name not in exclude_field_names
+        )
+    ]
+    field_names += [
+        f"{field.name}_id"
+        for field in qs.model._meta.fields
+        if isinstance(field, models.ForeignKey)
+    ]
+    for field_name in [
+        "version",
+        "is_latest",
+        "run_id",
+        "created_at",
+        "created_by_id",
+        "updated_at",
+    ]:
+        if field_name in field_names:
+            field_names.remove(field_name)
+            field_names.append(field_name)
+    if field_names[0] != "uid" and "uid" in field_names:
+        field_names.remove("uid")
+        field_names.insert(0, "uid")
+    if include or features:
+        subset_field_names = field_names[:4]
+        intersection = set(field_names) & set(include)
+        subset_field_names += list(intersection)
+        field_names = subset_field_names
+    return field_names
+
+
+def get_feature_annotate_kwargs(show_features: bool | list[str]) -> dict[str, Any]:
+    features = Feature.filter()
+    if isinstance(show_features, list):
+        features.filter(name__in=show_features)
+    # Get the categorical features
+    cat_feature_types = {
+        feature.dtype.replace("cat[", "").replace("]", "")
+        for feature in features
+        if feature.dtype.startswith("cat[")
+    }
+    # Get relationships of labels and features
+    link_models_on_models = {
+        getattr(
+            Artifact, obj.related_name
+        ).through.__get_name_with_schema__(): obj.related_model.__get_name_with_schema__()
+        for obj in Artifact._meta.related_objects
+        if obj.related_model.__get_name_with_schema__() in cat_feature_types
+    }
+    link_models_on_models["ArtifactULabel"] = "ULabel"
+    link_attributes_on_models = {
+        obj.related_name: link_models_on_models[
+            obj.related_model.__get_name_with_schema__()
+        ]
+        for obj in Artifact._meta.related_objects
+        if obj.related_model.__get_name_with_schema__() in link_models_on_models
+    }
+    # Prepare Django's annotate for features
+    annotate_kwargs = {}
+    for link_attr, feature_type in link_attributes_on_models.items():
+        annotate_kwargs[f"{link_attr}__feature__name"] = F(
+            f"{link_attr}__feature__name"
+        )
+        field_name = (
+            feature_type.split(".")[1] if "." in feature_type else feature_type
+        ).lower()
+        annotate_kwargs[f"{link_attr}__{field_name}__name"] = F(
+            f"{link_attr}__{field_name}__name"
+        )
+
+    annotate_kwargs["_feature_values__feature__name"] = F(
+        "_feature_values__feature__name"
+    )
+    annotate_kwargs["_feature_values__value"] = F("_feature_values__value")
+    return annotate_kwargs
+
+
+# https://claude.ai/share/16280046-6ae5-4f6a-99ac-dec01813dc3c
+def analyze_lookup_cardinality(
+    model_class: Record, lookup_paths: list[str] | None
+) -> dict[str, str]:
+    """Analyze lookup cardinality.
+
+    Analyzes Django model lookups to determine if they will result in
+    one-to-one or one-to-many relationships when used in annotations.
+
+    Args:
+        model_class: The Django model class to analyze
+        include: List of lookup paths (e.g. ["created_by__name", "ulabels__name"])
+
+    Returns:
+        Dictionary mapping lookup paths to either 'one' or 'many'
+    """
+    result = {}  # type: ignore
+    if lookup_paths is None:
+        return result
+    for lookup_path in lookup_paths:
+        parts = lookup_path.split("__")
+        current_model = model_class
+        is_many = False
+
+        # Walk through each part of the lookup path
+        for part in parts[:-1]:  # Exclude the last part as it's an attribute
+            field = None
+
+            # Handle reverse relations
+            for f in current_model._meta.get_fields():
+                if isinstance(f, ForeignObjectRel) and f.get_accessor_name() == part:
+                    field = f
+                    is_many = not f.one_to_one
+                    if hasattr(f, "field"):
+                        current_model = f.field.model
+                    break
+
+            # Handle forward relations
+            if field is None:
+                field = current_model._meta.get_field(part)
+                if isinstance(field, ManyToManyField):
+                    is_many = True
+                    current_model = field.remote_field.model
+                elif isinstance(field, ForeignKey):
+                    current_model = field.remote_field.model
+
+        result[lookup_path] = "many" if is_many else "one"
+
+    return result
+
+
+# https://lamin.ai/laminlabs/lamindata/transform/BblTiuKxsb2g0003
+# https://claude.ai/chat/6ea2498c-944d-4e7a-af08-29e5ddf637d2
+def reshape_annotate_result(
+    field_names: list[str],
+    df: pd.DataFrame,
+    extra_columns: dict[str, str] | None = None,
+    features: bool | list[str] = False,
+) -> pd.DataFrame:
+    """Reshapes experimental data with optional feature handling.
+
+    Parameters:
+    field_names: List of basic fields to include in result
+    df: Input dataframe with experimental data
+    extra_columns: Dict specifying additional columns to process with types ('one' or 'many')
+                  e.g., {'ulabels__name': 'many', 'created_by__name': 'one'}
+    features: If False, skip feature processing. If True, process all features.
+             If list of strings, only process specified features.
+
+    Returns:
+    DataFrame with reshaped data
+    """
+    extra_columns = extra_columns or {}
+
+    # Initialize result with basic fields
+    result = df[field_names].drop_duplicates(subset=["id"])
+
+    # Process features if requested
+    if features:
+        # Handle _feature_values if columns exist
+        feature_cols = ["_feature_values__feature__name", "_feature_values__value"]
+        if all(col in df.columns for col in feature_cols):
+            feature_values = process_feature_values(df, features)
+            if not feature_values.empty:
+                for col in feature_values.columns:
+                    if col in result.columns:
+                        continue
+                    result.insert(4, col, feature_values[col])
+
+        # Handle links features if they exist
+        links_features = [
+            col
+            for col in df.columns
+            if "feature__name" in col and col.startswith("links_")
+        ]
+
+        if links_features:
+            result = process_links_features(df, result, links_features, features)
+
+    # Process extra columns
+    if extra_columns:
+        result = process_extra_columns(df, result, extra_columns)
+
+    return result
+
+
+def process_feature_values(
+    df: pd.DataFrame, features: bool | list[str]
+) -> pd.DataFrame:
+    """Process _feature_values columns."""
+    feature_values = df.groupby(["id", "_feature_values__feature__name"])[
+        "_feature_values__value"
+    ].agg(set)
+
+    # Filter features if specific ones requested
+    if isinstance(features, list):
+        feature_values = feature_values[
+            feature_values.index.get_level_values(
+                "_feature_values__feature__name"
+            ).isin(features)
+        ]
+
+    return feature_values.unstack().reset_index()
+
+
+def process_links_features(
+    df: pd.DataFrame,
+    result: pd.DataFrame,
+    feature_cols: list[str],
+    features: bool | list[str],
+) -> pd.DataFrame:
+    """Process links_XXX feature columns."""
+    # this loops over different entities that might be linked under a feature
+    for feature_col in feature_cols:
+        prefix = re.match(r"links_(.+?)__feature__name", feature_col).group(1)
+
+        value_cols = [
+            col
+            for col in df.columns
+            if col.startswith(f"links_{prefix}__")
+            and col.endswith("__name")
+            and "feature__name" not in col
+        ]
+
+        if not value_cols:
+            continue
+
+        value_col = value_cols[0]
+        feature_names = df[feature_col].unique()
+        feature_names = feature_names[~pd.isna(feature_names)]
+
+        # Filter features if specific ones requested
+        if isinstance(features, list):
+            feature_names = [f for f in feature_names if f in features]
+
+        for feature_name in feature_names:
+            mask = df[feature_col] == feature_name
+            feature_values = df[mask].groupby("id")[value_col].agg(set)
+            result.insert(4, feature_name, result["id"].map(feature_values))
+
+    return result
+
+
+def process_extra_columns(
+    df: pd.DataFrame, result: pd.DataFrame, extra_columns: dict[str, str]
+) -> pd.DataFrame:
+    """Process additional columns based on their specified types."""
+    for col, col_type in extra_columns.items():
+        if col not in df.columns:
+            continue
+        if col in result.columns:
+            continue
+
+        values = df.groupby("id")[col].agg(set if col_type == "many" else "first")
+        result.insert(4, col, result["id"].map(values))
+
+    return result
+
 
 class QuerySet(models.QuerySet):
     """Sets of records returned by queries.
@@ -180,108 +462,45 @@ class QuerySet(models.QuerySet):
 
     Examples:
 
-        >>> ln.ULabel(name="my label").save()
-        >>> queryset = ln.ULabel.filter(name="my label")
+        >>> ULabel(name="my label").save()
+        >>> queryset = ULabel.filter(name="my label")
         >>> queryset
     """
 
     @doc_args(Record.df.__doc__)
     def df(
-        self, include: str | list[str] | None = None, join: str = "inner"
+        self,
+        include: str | list[str] | None = None,
+        features: bool | list[str] = False,
     ) -> pd.DataFrame:
         """{}"""  # noqa: D415
-        # re-order the columns
-        exclude_field_names = ["updated_at"]
-        field_names = [
-            field.name
-            for field in self.model._meta.fields
-            if (
-                not isinstance(field, models.ForeignKey)
-                and field.name not in exclude_field_names
-            )
-        ]
-        field_names += [
-            f"{field.name}_id"
-            for field in self.model._meta.fields
-            if isinstance(field, models.ForeignKey)
-        ]
-        for field_name in ["run_id", "created_at", "created_by_id", "updated_at"]:
-            if field_name in field_names:
-                field_names.remove(field_name)
-                field_names.append(field_name)
-        if field_names[0] != "uid" and "uid" in field_names:
-            field_names.remove("uid")
-            field_names.insert(0, "uid")
-        # create the dataframe
-        df = pd.DataFrame(self.values(), columns=field_names)
-        # if len(df) > 0 and "updated_at" in df:
-        #     df.updated_at = format_and_convert_to_local_time(df.updated_at)
-        # if len(df) > 0 and "started_at" in df:
-        #     df.started_at = format_and_convert_to_local_time(df.started_at)
-        pk_name = self.model._meta.pk.name
-        pk_column_name = pk_name if pk_name in df.columns else f"{pk_name}_id"
-        if pk_column_name in df.columns:
-            df = df.set_index(pk_column_name)
+        if include is None:
+            include = []
+        elif isinstance(include, str):
+            include = [include]
+        field_names = get_basic_field_names(self, include, features)
+        annotate_kwargs = {}
+        if features:
+            annotate_kwargs.update(get_feature_annotate_kwargs(features))
+        if include:
+            include = include.copy()[::-1]
+            include_kwargs = {s: F(s) for s in include if s not in field_names}
+            annotate_kwargs.update(include_kwargs)
+        if annotate_kwargs:
+            queryset = self.annotate(**annotate_kwargs)
+        else:
+            queryset = self
+        df = pd.DataFrame(queryset.values(*field_names, *list(annotate_kwargs.keys())))
         if len(df) == 0:
-            logger.warning(colors.yellow("No records found"))
+            df = pd.DataFrame({}, columns=field_names)
             return df
-        if include is not None:
-            if isinstance(include, str):
-                include = [include]
-            # fix ordering
-            include = include[::-1]
-            for expression in include:
-                split = expression.split("__")
-                field_name = split[0]
-                if len(split) > 1:
-                    lookup_str = "__".join(split[1:])
-                else:
-                    lookup_str = "id"
-                Record = self.model
-                field = getattr(Record, field_name)
-                if isinstance(field.field, models.ManyToManyField):
-                    related_ORM = (
-                        field.field.model
-                        if field.field.model != Record
-                        else field.field.related_model
-                    )
-                    if Record == related_ORM:
-                        left_side_link_model = f"from_{Record.__name__.lower()}"
-                        values_expression = (
-                            f"to_{Record.__name__.lower()}__{lookup_str}"
-                        )
-                    else:
-                        left_side_link_model = f"{Record.__name__.lower()}"
-                        values_expression = (
-                            f"{related_ORM.__name__.lower()}__{lookup_str}"
-                        )
-                    link_df = pd.DataFrame(
-                        field.through.objects.using(self.db).values(
-                            left_side_link_model, values_expression
-                        )
-                    )
-                    if link_df.shape[0] == 0:
-                        logger.warning(
-                            f"{colors.yellow(expression)} is not shown because no values are found"
-                        )
-                        continue
-                    link_groupby = link_df.groupby(left_side_link_model)[
-                        values_expression
-                    ].apply(list)
-                    df = pd.concat((link_groupby, df), axis=1, join=join)
-                    df.rename(columns={values_expression: expression}, inplace=True)
-                else:
-                    # the F() based implementation could also work for many-to-many,
-                    # would need to test what is faster
-                    df_anno = pd.DataFrame(
-                        self.annotate(expression=F(expression)).values(
-                            pk_column_name, "expression"
-                        )
-                    )
-                    df_anno = df_anno.set_index(pk_column_name)
-                    df_anno.rename(columns={"expression": expression}, inplace=True)
-                    df = pd.concat((df_anno, df), axis=1, join=join)
-        return df
+        extra_cols = analyze_lookup_cardinality(self.model, include)  # type: ignore
+        df_reshaped = reshape_annotate_result(field_names, df, extra_cols, features)
+        pk_name = self.model._meta.pk.name
+        pk_column_name = pk_name if pk_name in df.columns else f"{pk_name}_id"
+        if pk_column_name in df_reshaped.columns:
+            df_reshaped = df_reshaped.set_index(pk_column_name)
+        return df_reshaped
 
     def delete(self, *args, **kwargs):
         """Delete all records in the query set."""
@@ -335,8 +554,8 @@ class QuerySet(models.QuerySet):
         """At most one result. Returns it if there is one, otherwise returns ``None``.
 
         Examples:
-            >>> ln.ULabel.filter(name="benchmark").one_or_none()
-            >>> ln.ULabel.filter(name="non existing label").one_or_none()
+            >>> ULabel.filter(name="benchmark").one_or_none()
+            >>> ULabel.filter(name="non existing label").one_or_none()
         """
         if len(self) == 0:
             return None

lamindb/_record.py

@@ -264,14 +264,14 @@ def get(
 def df(
     cls,
     include: str | list[str] | None = None,
-    join: str = "inner",
+    features: bool | list[str] = False,
     limit: int = 100,
 ) -> pd.DataFrame:
     """{}"""  # noqa: D415
     query_set = cls.filter()
     if hasattr(cls, "updated_at"):
         query_set = query_set.order_by("-updated_at")
-    return query_set[:limit].df(include=include, join=join)
+    return query_set[:limit].df(include=include, features=features)
 
 
 def _search(
@@ -345,7 +345,9 @@ def _search(
         ranks.append(sub_rank)
         # startswith and avoid matching string with " " on the right
         # mostly for truncated
-        startswith_expr = regex_lookup(field_expr, rf"(?:^|\|){string}[^ ]*(\||$)")
+        startswith_expr = regex_lookup(
+            field_expr, rf"(?:^|.*\|){string}[^ ]*(?:\|.*|$)"
+        )
         startswith_rank = Cast(startswith_expr, output_field=IntegerField()) * 8
         ranks.append(startswith_rank)
         # match as sub-phrase from the left, mostly for truncated

lamindb/_save.py

@@ -112,6 +112,7 @@ def bulk_create(records: Iterable[Record], ignore_conflicts: bool | None = False
         records_by_orm[record.__class__].append(record)
     for registry, records in records_by_orm.items():
         registry.objects.bulk_create(records, ignore_conflicts=ignore_conflicts)
+        # records[:] = created  # In-place list update; does not seem to be necessary
 
 
 def bulk_update(records: Iterable[Record], ignore_conflicts: bool | None = False):

lamindb/_view.py

@@ -3,22 +3,107 @@ from __future__ import annotations
 import builtins
 import importlib
 import inspect
+from typing import TYPE_CHECKING
 
+from IPython.display import HTML, display
 from lamin_utils import colors, logger
 from lamindb_setup import settings
 from lamindb_setup._init_instance import get_schema_module_name
-from lnschema_core import Record
+from lnschema_core import Feature, Record
+
+from lamindb.core import FeatureValue, ParamValue
+
+from ._feature import convert_pandas_dtype_to_lamin_dtype
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 is_run_from_ipython = getattr(builtins, "__IPYTHON__", False)
 
 
+def display_df_with_descriptions(
+    df: pd.DataFrame, descriptions: dict[str, str] | None = None
+):
+    if descriptions is None:
+        display(df)
+        return None
+
+    # Start building HTML table
+    html = '<table class="dataframe">'
+
+    # Create header with title and description rows
+    html += "<thead>"
+
+    # Column names row
+    html += "<tr>"
+    html += '<th class="header-title index-header"></th>'  # Index header
+    for col in df.columns:
+        html += f'<th class="header-title">{col}</th>'
+    html += "</tr>"
+
+    # Descriptions row
+    html += "<tr>"
+    html += f'<th class="header-desc index-header">{df.index.name or ""}</th>'  # Index column
+    for col in df.columns:
+        desc = descriptions.get(col, "")
+        html += f'<th class="header-desc">{desc}</th>'
+    html += "</tr>"
+
+    html += "</thead>"
+
+    # Add body rows
+    html += "<tbody>"
+    for idx, row in df.iterrows():
+        html += "<tr>"
+        html += f'<th class="row-index">{idx}</th>'  # Index value
+        for col in df.columns:
+            html += f"<td>{row[col]}</td>"
+        html += "</tr>"
+    html += "</tbody>"
+    html += "</table>"
+
+    # Add CSS styles
+    styled_html = f"""
+    <style>
+        .dataframe {{
+            border-collapse: collapse;
+            margin: 10px 0;
+        }}
+        .dataframe th, .dataframe td {{
+            border: 1px solid #ddd;
+            padding: 8px;
+            text-align: left;
+        }}
+        .header-title {{
+            font-weight: bold;
+        }}
+        .header-desc {{
+            color: #666;
+            font-weight: normal;
+        }}
+        .row-index {{
+            font-weight: bold;
+        }}
+        .index-header {{
+            font-weight: bold;
+        }}
+    </style>
+    {html}
+    """
+    return display(HTML(styled_html))
+
+
 def view(
-    n: int = 7, schema: str | None = None, registries: list[str] | None = None
+    df: pd.DataFrame | None = None,
+    limit: int = 7,
+    schema: str | None = None,
+    registries: list[str] | None = None,
 ) -> None:
-    """View latest metadata state.
+    """View metadata.
 
     Args:
-        n: Display the last `n` rows of a registry.
+        df: A DataFrame to display.
+        limit: Display the latest `n` records
         schema: Schema module to view. Default's to
             `None` and displays all schema modules.
         registries: List of Record names. Defaults to
@@ -27,6 +112,16 @@ def view(
     Examples:
         >>> ln.view()
     """
+    if df is not None:
+        descriptions = {
+            col_name: convert_pandas_dtype_to_lamin_dtype(dtype)
+            for col_name, dtype in df.dtypes.to_dict().items()
+        }
+        feature_dtypes = dict(Feature.objects.values_list("name", "dtype"))
+        descriptions.update(feature_dtypes)
+        display_df_with_descriptions(df, descriptions)
+        return None
+
     if is_run_from_ipython:
         from IPython.display import display as show
     else:
@@ -39,6 +134,9 @@ def view(
 
     for schema_name in schema_names:
         schema_module = importlib.import_module(get_schema_module_name(schema_name))
+        # the below is necessary because a schema module might not have been
+        # explicitly accessed
+        importlib.reload(schema_module)
 
         all_registries = {
             registry
@@ -47,6 +145,8 @@ def view(
             and issubclass(registry, Record)
             and registry is not Record
         }
+        if schema_name == "core":
+            all_registries.update({FeatureValue, ParamValue})
         if registries is not None:
             filtered_registries = {
                 registry
@@ -62,11 +162,7 @@ def view(
             logger.print(section)
             logger.print("*" * len(section_no_color))
         for registry in sorted(filtered_registries, key=lambda x: x.__name__):
-            if hasattr(registry, "updated_at"):
-                df = registry.filter().order_by("-updated_at")[:n].df()
-            else:
-                # need to adjust in the future
-                df = registry.df().iloc[-n:]
+            df = registry.df(limit=limit)
             if df.shape[0] > 0:
                 logger.print(colors.blue(colors.bold(registry.__name__)))
                 show(df)