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

lamindb/core/_feature_manager.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import warnings
 from collections import defaultdict
 from collections.abc import Iterable
+from datetime import date, datetime
 from itertools import compress
 from typing import TYPE_CHECKING, Any
 
@@ -31,9 +33,16 @@ from lnschema_core.models import (
     Run,
     ULabel,
 )
+from rich.table import Column, Table
+from rich.text import Text
 
-from lamindb._feature import FEATURE_TYPES, convert_numpy_dtype_to_lamin_feature_type
+from lamindb._feature import (
+    FEATURE_DTYPES,
+    convert_pandas_dtype_to_lamin_dtype,
+    suggest_categorical_for_str_iterable,
+)
 from lamindb._feature_set import DICT_KEYS_TYPE, FeatureSet
+from lamindb._from_values import _print_values
 from lamindb._record import (
     REGISTRY_UNIQUE_FIELD,
     get_name_field,
@@ -44,8 +53,15 @@ from lamindb._save import save
 from lamindb.core.exceptions import DoesNotExist, ValidationError
 from lamindb.core.storage import LocalPathClasses
 
+from ._describe import (
+    NAME_WIDTH,
+    TYPE_WIDTH,
+    VALUES_WIDTH,
+    describe_header,
+    print_rich_tree,
+)
 from ._django import get_artifact_with_related
-from ._label_manager import get_labels_as_dict
+from ._label_manager import _get_labels, describe_labels
 from ._settings import settings
 from .schema import (
     dict_related_model_to_related_name,
@@ -53,6 +69,7 @@ from .schema import (
 
 if TYPE_CHECKING:
     from lnschema_core.types import FieldAttr
+    from rich.tree import Tree
 
     from lamindb._query_set import QuerySet
 
@@ -75,7 +92,9 @@ def get_accessor_by_registry_(host: Artifact | Collection) -> dict:
     return dictionary
 
 
-def get_feature_set_by_slot_(host) -> dict:
+def get_feature_set_by_slot_(host: Artifact | Collection) -> dict:
+    if isinstance(host, Collection):
+        return {}
     # if the host is not yet saved
     if host._state.adding:
         if hasattr(host, "_feature_sets"):
@@ -134,14 +153,14 @@ def custom_aggregate(field, using: str):
         return GroupConcat(field)
 
 
-def _print_categoricals_postgres(
+def _get_categoricals_postgres(
     self: Artifact | Collection,
     related_data: dict | None = None,
-    print_types: bool = False,
-    to_dict: bool = False,
     print_params: bool = False,
-):
-    from lamindb._from_values import _print_values
+) -> dict[tuple[str, str], set[str]]:
+    """Get categorical features and their values using PostgreSQL-specific optimizations."""
+    if print_params:
+        return {}
 
     if not related_data:
         artifact_meta = get_artifact_with_related(
@@ -149,6 +168,7 @@ def _print_categoricals_postgres(
         )
         related_data = artifact_meta.get("related_data", {})
 
+    # Process m2m data
     m2m_data = related_data.get("m2m", {}) if related_data else {}
     m2m_name = {}
     for related_name, values in m2m_data.items():
@@ -157,6 +177,8 @@ def _print_categoricals_postgres(
             self.__class__.__name__, ""
         ).lower()
         m2m_name[related_model_name] = values
+
+    # Get feature information
     links_data = related_data.get("link", {}) if related_data else {}
     feature_dict = {
         id: (name, dtype)
@@ -165,188 +187,295 @@ def _print_categoricals_postgres(
         )
     }
 
-    msg = ""
-    dictionary = {}
+    # Build result dictionary
+    result = defaultdict(set)
+    for link_name, link_values in links_data.items():
+        related_name = link_name.removeprefix("links_").replace("_", "")
+        if not link_values:
+            continue
 
-    # 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:
+        for link_value in link_values:
+            feature_id = link_value.get("feature")
+            if feature_id is None:
                 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
 
+            feature_name, feature_dtype = feature_dict.get(feature_id)
+            label_id = link_value.get(related_name)
+            label_name = m2m_name.get(related_name, {}).get(label_id)
+            if label_name:
+                result[(feature_name, feature_dtype)].add(label_name)
+
+    return dict(result)
 
-def _print_categoricals(
+
+def _get_categoricals(
     self: Artifact | Collection,
-    print_types: bool = False,
-    to_dict: bool = False,
     print_params: bool = False,
-):
-    from lamindb._from_values import _print_values
+) -> dict[tuple[str, str], set[str]]:
+    """Get categorical features and their values using the default approach."""
+    if print_params:
+        return {}
 
-    msg = ""
-    dictionary = {}
-    # categorical feature values
-    if not print_params:
-        labels_msg = ""
-        labels_by_feature = defaultdict(list)
-        for _, (_, links) in get_labels_as_dict(
-            self, links=True, instance=self._state.db
-        ).items():
-            for link in links:
-                if hasattr(link, "feature_id") and link.feature_id is not None:
-                    link_attr = get_link_attr(link, self)
-                    labels_by_feature[link.feature_id].append(
-                        getattr(link, link_attr).name
-                    )
-        labels_msgs = []
-        for feature_id, labels_list in labels_by_feature.items():
-            feature = Feature.objects.using(self._state.db).get(id=feature_id)
-            print_values = _print_values(labels_list, n=10)
-            type_str = f": {feature.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
+    result = defaultdict(set)
+    for _, links in _get_labels(self, links=True, instance=self._state.db).items():
+        for link in links:
+            if hasattr(link, "feature_id") and link.feature_id is not None:
+                feature = Feature.objects.using(self._state.db).get(id=link.feature_id)
+                link_attr = get_link_attr(link, self)
+                label_name = getattr(link, link_attr).name
+                result[(feature.name, feature.dtype)].add(label_name)
 
+    return dict(result)
+
+
+def _get_non_categoricals(
+    self,
+    print_params: bool = False,
+) -> dict[tuple[str, str], set[Any]]:
+    """Get non-categorical features and their values."""
+    non_categoricals = {}
 
-def _print_featuresets_postgres(
+    if self.id is not None and isinstance(self, (Artifact, Run)):
+        attr_name = "param" if print_params else "feature"
+        _feature_values = (
+            getattr(self, f"_{attr_name}_values")
+            .values(f"{attr_name}__name", f"{attr_name}__dtype")
+            .annotate(values=custom_aggregate("value", self._state.db))
+            .order_by(f"{attr_name}__name")
+        )
+
+        for fv in _feature_values:
+            feature_name = fv[f"{attr_name}__name"]
+            feature_dtype = fv[f"{attr_name}__dtype"]
+            values = fv["values"]
+
+            # Convert single values to sets
+            if not isinstance(values, (list, dict, set)):
+                values = {values}
+            elif (
+                isinstance(values, list)
+                and feature_dtype != "dict"
+                and not feature_dtype.startswith("list")
+            ):
+                values = set(values)
+
+            # Handle special datetime types
+            if feature_dtype == "datetime":
+                values = {datetime.fromisoformat(value) for value in values}
+            if feature_dtype == "date":
+                values = {date.fromisoformat(value) for value in values}
+
+            non_categoricals[(feature_name, feature_dtype)] = values
+
+    return non_categoricals
+
+
+def _get_featuresets_postgres(
     self: Artifact | Collection,
     related_data: dict | None = None,
-    print_types: bool = False,
-):
-    from lamindb._from_values import _print_values
-
+) -> dict:
     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
+    return fs_data
+
+
+def _create_feature_table(name: str, registry_str: str, data: list) -> Table:
+    """Create a Rich table for a feature group."""
+    table = Table(
+        Column(name, style="", no_wrap=True, width=NAME_WIDTH),
+        Column(registry_str, style="dim", no_wrap=True, width=TYPE_WIDTH),
+        Column("", width=VALUES_WIDTH, no_wrap=True),
+        show_header=True,
+        box=None,
+        pad_edge=False,
+    )
+    for row in data:
+        table.add_row(*row)
+    return table
 
 
-def print_features(
+def describe_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
+    tree: Tree | None = None,
+    with_labels: bool = False,
+):
+    """Describe features of an artifact or collection."""
+    if print_types:
+        warnings.warn(
+            "`print_types` parameter is deprecated and will be removed in a future version. Types are now always printed.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+    # initialize tree
+    if tree is None:
+        tree = describe_header(self)
+
+    dictionary: dict[str, Any] = {}
+
+    if self._state.adding:
+        return dictionary if to_dict else tree
 
+    # feature sets
+    feature_set_data: dict[str, tuple[str, list[str]]] = {}
+    feature_data: dict[str, tuple[str, list[str]]] = {}
+    if not print_params and not to_dict:
+        if self.id is not None and connections[self._state.db].vendor == "postgresql":
+            fs_data = _get_featuresets_postgres(self, related_data=related_data)
+            for fs_id, (slot, data) in fs_data.items():
+                for registry_str, feature_names in data.items():
+                    feature_set = FeatureSet.get(id=fs_id)
+                    feature_set_data[slot] = (feature_set, feature_names)
+                    for feature_name in feature_names:
+                        feature_data[feature_name] = (slot, registry_str)
+        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])
+                feature_set_data[slot] = (feature_set, feature_names)
+                for feature_name in feature_names:
+                    feature_data[feature_name] = (slot, feature_set.registry)
+
+    internal_feature_names: set[str] = {}  # type: ignore
+    if isinstance(self, Artifact):
+        feature_sets = self.feature_sets.filter(registry="Feature").all()
+        internal_feature_names = set()  # type: ignore
+        if len(feature_sets) > 0:
+            for feature_set in feature_sets:
+                internal_feature_names = internal_feature_names.union(
+                    set(feature_set.members.values_list("name", flat=True))
+                )  # type: ignore
+
+    # categorical feature values
+    # Get the categorical data using the appropriate method
     if not self._state.adding and connections[self._state.db].vendor == "postgresql":
-        msg, dictionary = _print_categoricals_postgres(
+        categoricals = _get_categoricals_postgres(
             self,
             related_data=related_data,
-            print_types=print_types,
-            to_dict=to_dict,
             print_params=print_params,
         )
     else:
-        msg, dictionary = _print_categoricals(
+        categoricals = _get_categoricals(
             self,
-            print_types=print_types,
-            to_dict=to_dict,
             print_params=print_params,
         )
 
-    # non-categorical feature values
-    non_labels_msg = ""
-    if self.id is not None and self.__class__ == Artifact or self.__class__ == Run:
-        attr_name = "param" if print_params else "feature"
-        _feature_values = (
-            getattr(self, f"_{attr_name}_values")
-            .values(f"{attr_name}__name", f"{attr_name}__dtype")
-            .annotate(values=custom_aggregate("value", self._state.db))
-            .order_by(f"{attr_name}__name")
-        )
-        if len(_feature_values) > 0:
-            for fv in _feature_values:
-                feature_name = fv[f"{attr_name}__name"]
-                feature_dtype = fv[f"{attr_name}__dtype"]
-                values = fv["values"]
-                # TODO: understand why the below is necessary
-                if not isinstance(values, list):
-                    values = [values]
-                if to_dict:
-                    dictionary[feature_name] = values if len(values) > 1 else values[0]
-                type_str = f": {feature_dtype}" if print_types else ""
-                printed_values = (
-                    _print_values(values, n=10, quotes=False)
-                    if not feature_dtype.startswith("list")
-                    else values
-                )
-                non_labels_msg += f"    '{feature_name}'{type_str} = {printed_values}\n"
-            msg += non_labels_msg
+    # Get non-categorical features
+    non_categoricals = _get_non_categoricals(
+        self,
+        print_params=print_params,
+    )
 
-    if msg != "":
-        header = "Features" if not print_params else "Params"
-        msg = f"  {colors.italic(header)}\n" + msg
+    # Process all Features containing labels and sort into internal/external
+    internal_feature_labels = {}
+    external_data = []
+    for features, is_list_type in [(categoricals, False), (non_categoricals, True)]:
+        for (feature_name, feature_dtype), values in sorted(features.items()):
+            # Handle dictionary conversion
+            if to_dict:
+                dict_value = values if len(values) > 1 else next(iter(values))
+                dictionary[feature_name] = dict_value
+                continue
 
-    # feature sets
-    if not print_params:
-        feature_set_msg = ""
-        if self.id is not None and connections[self._state.db].vendor == "postgresql":
-            feature_set_msg = _print_featuresets_postgres(
-                self, related_data=related_data
+            # Format message
+            printed_values = (
+                _print_values(sorted(values), n=10, quotes=False)
+                if not is_list_type or not feature_dtype.startswith("list")
+                else sorted(values)
             )
-        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
+
+            # Sort into internal/external
+            feature_info = (
+                feature_name,
+                Text(feature_dtype, style="dim"),
+                printed_values,
+            )
+            if feature_name in internal_feature_names:
+                internal_feature_labels[feature_name] = feature_info
+            else:
+                external_data.append(feature_info)
+
     if to_dict:
         return dictionary
-    else:
-        return msg
+
+    # Dataset section
+    internal_features_slot: dict[
+        str, list
+    ] = {}  # internal features from the `Feature` registry that contain labels
+    for feature_name, feature_row in internal_feature_labels.items():
+        slot, _ = feature_data.get(feature_name)
+        internal_features_slot.setdefault(slot, []).append(feature_row)
+    dataset_tree_children = []
+
+    for slot, (feature_set, feature_names) in feature_set_data.items():
+        if slot in internal_features_slot:
+            feature_rows = internal_features_slot[slot]
+        else:
+            feature_rows = [
+                (feature_name, Text(str(feature_set.dtype), style="dim"), "")
+                for feature_name in feature_names
+                if feature_name
+            ]
+        dataset_tree_children.append(
+            _create_feature_table(
+                Text.assemble(
+                    (slot, "violet"),
+                    (" • ", "dim"),
+                    (str(feature_set.n), "pink1"),
+                ),
+                Text.assemble((f"[{feature_set.registry}]", "pink1")),
+                feature_rows,
+            )
+        )
+    ## internal features from the non-`Feature` registry
+    if dataset_tree_children:
+        dataset_tree = tree.add(
+            Text.assemble(
+                ("Dataset", "bold bright_magenta"),
+                ("/", "dim"),
+                (".feature_sets", "dim bold"),
+            )
+        )
+        for child in dataset_tree_children:
+            dataset_tree.add(child)
+
+    # Annotations section
+    ## external features
+    features_tree_children = []
+    if external_data:
+        features_tree_children.append(
+            _create_feature_table(
+                Text.assemble(
+                    ("Params" if print_params else "Features", "green_yellow")
+                ),
+                "",
+                external_data,
+            )
+        )
+    annotations_tree = None
+    if features_tree_children:
+        annotations_tree = tree.add(Text("Annotations", style="bold dark_orange"))
+        for child in features_tree_children:
+            annotations_tree.add(child)
+    if with_labels:
+        labels_tree = describe_labels(self, as_subtree=True)
+        if labels_tree:
+            if annotations_tree is None:
+                annotations_tree = tree.add(
+                    Text("Annotations", style="bold dark_orange")
+                )
+            annotations_tree.add(labels_tree)
+
+    return tree
 
 
 def parse_feature_sets_from_anndata(
@@ -371,7 +500,7 @@ def parse_feature_sets_from_anndata(
         type = (
             "float"
             if adata.X is None
-            else convert_numpy_dtype_to_lamin_feature_type(adata.X.dtype)
+            else convert_pandas_dtype_to_lamin_dtype(adata.X.dtype)
         )
     feature_sets = {}
     if var_field is not None:
@@ -409,51 +538,75 @@ def parse_feature_sets_from_anndata(
     return feature_sets
 
 
+def is_valid_datetime_str(date_string: str) -> bool | str:
+    try:
+        dt = datetime.fromisoformat(date_string)
+        return dt.isoformat()
+    except ValueError:
+        return False
+
+
 def infer_feature_type_convert_json(
-    value: Any, mute: bool = False, str_as_ulabel: bool = True
-) -> tuple[str, Any]:
+    key: str, value: Any, mute: bool = False, str_as_ulabel: bool = True
+) -> tuple[str, Any, str]:
+    message = ""
     if isinstance(value, bool):
-        return FEATURE_TYPES["bool"], value
+        return "bool", value, message
     elif isinstance(value, int):
-        return FEATURE_TYPES["int"], value
+        return "int", value, message
     elif isinstance(value, float):
-        return FEATURE_TYPES["float"], value
+        return "float", value, message
+    elif isinstance(value, date):
+        return "date", value.isoformat(), message
+    elif isinstance(value, datetime):
+        return "datetime", value.isoformat(), message
     elif isinstance(value, str):
-        if str_as_ulabel:
-            return FEATURE_TYPES["str"] + "[ULabel]", value
+        if datetime_str := is_valid_datetime_str(value):
+            dt_type = (
+                "date" if len(value) == 10 else "datetime"
+            )  # YYYY-MM-DD is exactly 10 characters
+            sanitized_value = datetime_str[:10] if dt_type == "date" else datetime_str  # type: ignore
+            return dt_type, sanitized_value, message  # type: ignore
         else:
-            return "str", value
+            return "cat ? str", value, message
     elif isinstance(value, Iterable) and not isinstance(value, (str, bytes)):
-        if isinstance(value, (pd.Series, np.ndarray)):
-            return convert_numpy_dtype_to_lamin_feature_type(
-                value.dtype, str_as_cat=str_as_ulabel
-            ), list(value)
+        if isinstance(value, (pd.Series, np.ndarray, pd.Categorical)):
+            dtype = convert_pandas_dtype_to_lamin_dtype(value.dtype)
+            if dtype == "str":
+                # ndarray doesn't know categorical, so there was no conscious choice
+                # offer both options
+                if isinstance(value, np.ndarray):
+                    dtype = "cat ? str"
+                else:
+                    # suggest to create a categorical if there are few unique values
+                    message = suggest_categorical_for_str_iterable(value, key)
+                    if message:
+                        message = f"  # {message}"
+            return dtype, list(value), message
         if isinstance(value, dict):
-            return "dict", value
+            return "dict", value, message
         if len(value) > 0:  # type: ignore
             first_element_type = type(next(iter(value)))
             if all(isinstance(elem, first_element_type) for elem in value):
                 if first_element_type is bool:
-                    return f"list[{FEATURE_TYPES['bool']}]", value
+                    return "list[bool]", value, message
                 elif first_element_type is int:
-                    return f"list[{FEATURE_TYPES['int']}]", value
+                    return "list[int]", value, message
                 elif first_element_type is float:
-                    return f"list[{FEATURE_TYPES['float']}]", value
+                    return "list[float]", value, message
                 elif first_element_type is str:
-                    if str_as_ulabel:
-                        return FEATURE_TYPES["str"] + "[ULabel]", value
-                    else:
-                        return "list[str]", value
+                    return ("list[cat ? str]", value, message)
                 elif first_element_type == Record:
                     return (
-                        f"cat[{first_element_type.__get_name_with_schema__()}]",
+                        f"list[cat[{first_element_type.__get_name_with_schema__()}]]",
                         value,
+                        message,
                     )
     elif isinstance(value, Record):
-        return (f"cat[{value.__class__.__get_name_with_schema__()}]", value)
+        return (f"cat[{value.__class__.__get_name_with_schema__()}]", value, message)
     if not mute:
         logger.warning(f"cannot infer feature type of: {value}, returning '?")
-    return ("?", value)
+    return "?", value, message
 
 
 def __init__(self, host: Artifact | Collection | Run):
@@ -463,12 +616,13 @@ def __init__(self, host: Artifact | Collection | Run):
 
 
 def __repr__(self) -> str:
-    return print_features(self._host, print_params=(self.__class__ == ParamManager))  # type: ignore
+    tree = describe_features(self._host, print_params=(self.__class__ == ParamManager))  # type: ignore
+    return print_rich_tree(tree, fallback="no linked features")
 
 
 def get_values(self) -> dict[str, Any]:
     """Get feature values as a dictionary."""
-    return print_features(
+    return describe_features(
         self._host, to_dict=True, print_params=(self.__class__ == ParamManager)
     )  # type: ignore
 
@@ -669,10 +823,14 @@ def _add_values(
     validated_keys = keys_array[validated]
     if validated.sum() != len(keys):
         not_validated_keys = keys_array[~validated]
+        not_validated_keys_dtype_message = [
+            (key, infer_feature_type_convert_json(key, features_values[key]))
+            for key in not_validated_keys
+        ]
         hint = "\n".join(
             [
-                f"  ln.{model_name}(name='{key}', dtype='{infer_feature_type_convert_json(features_values[key], str_as_ulabel=str_as_ulabel)[0]}').save()"
-                for key in not_validated_keys
+                f"  ln.{model_name}(name='{key}', dtype='{dtype}').save(){message}"
+                for key, (dtype, _, message) in not_validated_keys_dtype_message
             ]
         )
         msg = (
@@ -690,12 +848,13 @@ def _add_values(
     not_validated_values = []
     for key, value in features_values.items():
         feature = model.get(name=key)
-        inferred_type, converted_value = infer_feature_type_convert_json(
+        inferred_type, converted_value, _ = infer_feature_type_convert_json(
+            key,
             value,
             mute=True,
             str_as_ulabel=str_as_ulabel,
         )
-        if feature.dtype == "number":
+        if feature.dtype == "num":
             if inferred_type not in {"int", "float"}:
                 raise TypeError(
                     f"Value for feature '{key}' with type {feature.dtype} must be a number"
@@ -706,12 +865,13 @@ def _add_values(
                     raise TypeError(
                         f"Value for feature '{key}' with type '{feature.dtype}' must be a string or record."
                     )
-        elif not inferred_type == feature.dtype:
+        elif (feature.dtype == "str" and feature.dtype not in inferred_type) or (
+            feature.dtype != "str" and feature.dtype != inferred_type
+        ):
             raise ValidationError(
                 f"Expected dtype for '{key}' is '{feature.dtype}', got '{inferred_type}'"
             )
         if not feature.dtype.startswith("cat"):
-            # can remove the query once we have the unique constraint
             filter_kwargs = {model_name.lower(): feature, "value": converted_value}
             feature_value = value_model.filter(**filter_kwargs).one_or_none()
             if feature_value is None:
@@ -814,6 +974,59 @@ def add_values_params(
     _add_values(self, values, Param.name, str_as_ulabel=False)
 
 
+def remove_values(
+    self,
+    feature: str | Feature,
+    *,
+    value: Any | None = None,
+):
+    """Remove value annotations for a given feature.
+
+    Args:
+        feature: The feature for which to remove values.
+        value: An optional value to restrict removal to a single value.
+
+    """
+    if isinstance(feature, str):
+        feature = Feature.get(name=feature)
+    filter_kwargs = {"feature": feature}
+    if feature.dtype.startswith("cat["):
+        feature_registry = feature.dtype.replace("cat[", "").replace("]", "")
+        if value is not None:
+            assert isinstance(value, Record)  # noqa: S101
+            # the below uses our convention for field names in link models
+            link_name = (
+                feature_registry.split(".")[1]
+                if "." in feature_registry
+                else feature_registry
+            ).lower()
+            filter_kwargs[link_name] = value
+        if feature_registry == "ULabel":
+            link_attribute = "links_ulabel"
+        else:
+            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__() == feature_registry
+            }
+            link_attribute = {
+                obj.related_name
+                for obj in Artifact._meta.related_objects
+                if obj.related_model.__get_name_with_schema__() in link_models_on_models
+            }.pop()
+        getattr(self._host, link_attribute).filter(**filter_kwargs).all().delete()
+    else:
+        if value is not None:
+            filter_kwargs["value"] = value
+        feature_values = self._host._feature_values.filter(**filter_kwargs)
+        self._host._feature_values.remove(*feature_values)
+        # this might leave a dangling feature_value record
+        # but we don't want to pay the price of making another query just to remove this annotation
+        # we can clean the FeatureValue registry periodically if we want to
+
+
 def add_feature_set(self, feature_set: FeatureSet, slot: str) -> None:
     """Curate artifact with a feature set.
 
@@ -847,7 +1060,10 @@ def add_feature_set(self, feature_set: FeatureSet, slot: str) -> None:
 
 
 def _add_set_from_df(
-    self, field: FieldAttr = Feature.name, organism: str | None = None
+    self,
+    field: FieldAttr = Feature.name,
+    organism: str | None = None,
+    mute: bool = False,
 ):
     """Add feature set corresponding to column names of DataFrame."""
     if isinstance(self._host, Artifact):
@@ -855,21 +1071,14 @@ def _add_set_from_df(
     else:
         # Collection
         assert self._host.artifact._accessor == "DataFrame"  # noqa: S101
-
-    # parse and register features
-    registry = field.field.model
     df = self._host.load()
-    features = registry.from_values(df.columns, field=field, organism=organism)
-    if len(features) == 0:
-        logger.error(
-            "no validated features found in DataFrame! please register features first!"
-        )
-        return
-
-    # create and link feature sets
-    feature_set = FeatureSet(features=features)
-    feature_sets = {"columns": feature_set}
-    self._host._feature_sets = feature_sets
+    feature_set = FeatureSet.from_df(
+        df=df,
+        field=field,
+        mute=mute,
+        organism=organism,
+    )
+    self._host._feature_sets = {"columns": feature_set}
     self._host.save()
 
 
@@ -1056,6 +1265,7 @@ FeatureManager._add_from = _add_from
 FeatureManager.filter = filter
 FeatureManager.get = get
 FeatureManager.make_external = make_external
+FeatureManager.remove_values = remove_values
 ParamManager.add_values = add_values_params
 ParamManager.get_values = get_values
 ParamManager.filter = filter