✨ Introduce the using argument in ln.save()

lamindb/curators/core.py

@@ -664,17 +664,21 @@ def __init__(
         schema: Schema,
         slot: str | None = None,
         require_saved_schema: bool = True,
+        instance: str | None = None,
     ) -> None:
         super().__init__(
             dataset=dataset, schema=schema, require_saved_schema=require_saved_schema
         )
+        self._instance = instance
 
         categoricals = []
         features = []
         feature_ids: set[int] = set()
 
         if schema.flexible:
-            features += Feature.filter(name__in=self._dataset.keys()).to_list()
+            features += Feature.filter(
+                name__in=self._dataset.keys(), _using_key=instance
+            ).to_list()
             feature_ids = {feature.id for feature in features}
 
         if schema.n_members and schema.n_members > 0:
@@ -847,6 +851,7 @@ def __init__(
             categoricals=categoricals,
             index=schema.index,
             slot=slot,
+            instance=instance,
             maximal_set=schema.maximal_set,
             schema=schema,
         )
@@ -978,6 +983,7 @@ def __init__(
         slot: str | None = None,
         features: dict[str, Any] | None = None,
         require_saved_schema: bool = True,
+        instance: str | None = None,
     ) -> None:
         # loads or opens dataset, dataset may be an artifact
         super().__init__(
@@ -986,12 +992,14 @@ def __init__(
             features=features,
             require_saved_schema=require_saved_schema,
         )
+        self._instance = instance
         # uses open dataset at self._dataset
         self._atomic_curator = ComponentCurator(
             dataset=self._dataset,
             schema=schema,
             slot=slot,
             require_saved_schema=require_saved_schema,
+            instance=instance,
         )
         # Handle (nested) attrs
         if slot is None and schema.slots:
@@ -1013,6 +1021,7 @@ def __init__(
                             slot_schema,
                             slot=slot_name,
                             require_saved_schema=require_saved_schema,
+                            instance=instance,
                         )
                 elif slot_name != "__external__":
                     raise ValueError(
@@ -1066,6 +1075,7 @@ def __init__(
         schema: Schema,
         slot: str | None = None,
         require_saved_schema: bool = False,
+        instance: str | None = None,
     ) -> None:
         if not isinstance(dataset, dict) and not isinstance(dataset, Artifact):
             raise InvalidArgument("The dataset must be a dict or dict-like artifact.")
@@ -1076,7 +1086,11 @@ def __init__(
             d = dataset
         df = convert_dict_to_dataframe_for_validation(d, schema)  # type: ignore
         super().__init__(
-            df, schema, slot=slot, require_saved_schema=require_saved_schema
+            df,
+            schema,
+            slot=slot,
+            require_saved_schema=require_saved_schema,
+            instance=instance,
         )
 
 
@@ -1545,6 +1559,7 @@ def __init__(
         type_uid: str | None = None,
         maximal_set: bool = True,  # whether unvalidated categoricals cause validation failure.
         schema: Schema = None,
+        instance: str | None = None,  # target instance for cross-instance curation
     ) -> None:
         self._values_getter = values_getter
         self._values_setter = values_setter
@@ -1561,14 +1576,17 @@ def __init__(
         self.records = None
         self._maximal_set = maximal_set
         self._type_record = None
+        self._instance = instance
         self._registry = self._field.field.model
         self._field_name = self._field.field.name
         self._filter_kwargs = {}
         self._schema = schema
         if filter_str and filter_str != "unsaved":
             self._filter_kwargs.update(
                 resolve_relation_filters(
-                    parse_filter_string(filter_str), self._registry
+                    parse_filter_string(filter_str),
+                    self._registry,
+                    using=self._instance,
                 )  # type: ignore
             )
         if self._registry.__base__.__name__ == "BioRecord":
@@ -1590,6 +1608,7 @@ def __init__(
             self._type_record = get_record_type_from_uid(
                 self._registry,
                 self._type_uid,
+                using=self._instance,
             )
 
         if hasattr(self._registry, "_name_field"):
@@ -1748,7 +1767,9 @@ def _add_validated(self) -> tuple[list, list]:
                 if filter_str:
                     parsed_filters = parse_filter_string(filter_str)
                     filter_kwargs.update(
-                        resolve_relation_filters(parsed_filters, registry)
+                        resolve_relation_filters(
+                            parsed_filters, registry, using=self._instance
+                        )
                     )
                 if registry.__base__.__name__ == "BioRecord":
                     organism_record = get_organism_record_from_field(
@@ -1774,7 +1795,7 @@ def _add_validated(self) -> tuple[list, list]:
                     # When we have a Schema with typed members,
                     # scope the query to the types present in the schema's members (plus untyped features)
                     # to avoid ambiguous matches across different feature types.
-                    qs = registry.filter()
+                    qs = registry.filter(_using_key=self._instance)
                     if self._schema and self._schema.n_members:
                         type_ids = {
                             m.type_id
@@ -1783,7 +1804,8 @@ def _add_validated(self) -> tuple[list, list]:
                         }
                         if type_ids:
                             qs = registry.filter(
-                                Q(type_id__in=type_ids) | Q(type_id__isnull=True)
+                                Q(type_id__in=type_ids) | Q(type_id__isnull=True),
+                                _using_key=self._instance,
                             )
                     self._subtype_query_set = qs
                 else:
@@ -1837,6 +1859,7 @@ def _add_validated(self) -> tuple[list, list]:
                     remaining_values,
                     field=field,
                     mute=True,
+                    using=self._instance,
                     **filter_kwargs,  # type: ignore
                 )
                 existing_and_public_values = [
@@ -1939,12 +1962,20 @@ def _validate(
             filter_str = result.get("filter_str", "")
             if filter_str:
                 parsed_filters = parse_filter_string(filter_str)
-                filter_kwargs.update(resolve_relation_filters(parsed_filters, registry))
+                filter_kwargs.update(
+                    resolve_relation_filters(
+                        parsed_filters, registry, using=self._instance
+                    )
+                )
             registry_or_queryset = registry
             if self._subtype_query_set is not None and registry == self._registry:
                 registry_or_queryset = self._subtype_query_set
             # first inspect against the registry
-            inspect_result = registry_or_queryset.filter(**filter_kwargs).inspect(
+            if registry_or_queryset is registry:
+                queryset = registry.filter(_using_key=self._instance, **filter_kwargs)
+            else:
+                queryset = registry_or_queryset.filter(**filter_kwargs)
+            inspect_result = queryset.inspect(
                 non_validated,
                 field=field,
                 mute=True,
@@ -2074,10 +2105,12 @@ def __init__(
         slot: str | None = None,
         maximal_set: bool = False,
         schema: Schema | None = None,
+        instance: str | None = None,
     ) -> None:
         self._non_validated = None
         self._index = index
         self._schema = schema
+        self._instance = instance
         self._artifact: Artifact = None  # pass the dataset as an artifact
         self._dataset: Any = df  # pass the dataset as an AnyPathStr or data object
         if isinstance(self._dataset, Artifact):
@@ -2115,6 +2148,7 @@ def __init__(
             if schema.id is None
             else f"schemas__id={schema.id}",
             schema=schema,
+            instance=instance,
         )
         for feature in self._categoricals:
             result = parse_dtype(feature._dtype_str)[0]
@@ -2135,6 +2169,7 @@ def __init__(
                     cat_manager=self,
                     filter_str=result["filter_str"],
                     type_uid=result.get("type_uid"),
+                    instance=instance,
                 )
         if index is not None and index._dtype_str.startswith("cat"):
             result = parse_dtype(index._dtype_str)[0]
@@ -2150,6 +2185,7 @@ def __init__(
                 cat_manager=self,
                 filter_str=result["filter_str"],
                 type_uid=result.get("type_uid"),
+                instance=instance,
             )
 
     @property

lamindb/models/__init__.py

@@ -239,3 +239,7 @@
 )
 
 FeatureValue = JsonValue  # backward compatibility
+
+from .sqlrecord import _ensure_lamindb_router
+
+_ensure_lamindb_router()

lamindb/models/_feature_manager.py

@@ -1700,7 +1700,10 @@ def add_values(
             )
             schema = Schema(feature_objects)
         ExperimentalDictCurator(
-            dictionary, schema, require_saved_schema=False
+            dictionary,
+            schema,
+            require_saved_schema=False,
+            instance=self._host._state.db,
         ).validate()
         if host_is_record and schema.index is not None:
             from .record import strip_index_for_record_persistence
@@ -1730,6 +1733,7 @@ def _add_values(
         from ..base.dtypes import is_iterable_of_sqlrecord
         from .can_curate import CanCurate
 
+        host_db = self._host._state.db
         host_is_record = self._host.__class__.__name__ == "Record"
         if host_is_record:
             feature_json_values: list[SQLRecord] = []
@@ -1746,12 +1750,12 @@ def _add_values(
             )
             self._raise_not_validated_values(record_not_validated_values)
             if feature_json_values:
-                save(feature_json_values)
+                save(feature_json_values, using=host_db)
             for links in links_by_model.values():
                 try:
-                    save(links, ignore_conflicts=False)
+                    save(links, ignore_conflicts=False, using=host_db)
                 except Exception:
-                    save(links, ignore_conflicts=True)
+                    save(links, ignore_conflicts=True, using=host_db)
             from .record import get_type_schema_index, persist_record_name
 
             if (
@@ -1918,7 +1922,9 @@ def set_values(
         if host_is_artifact:
             schema = self._get_external_schema()
         if schema is not None:
-            ExperimentalDictCurator(dictionary, schema).validate()
+            ExperimentalDictCurator(
+                dictionary, schema, instance=self._host._state.db
+            ).validate()
             member_ids = set(schema.members.values_list("id", flat=True))
             features_not_in_schema = [
                 feature.name
@@ -2221,7 +2227,9 @@ def _add_from(self, data: Artifact | Collection, transfer_logs: dict = None):
                 self._host.features._add_schema(schema_self, slot)
 
 
-def bulk_set_features_in_records(records: Iterable[Record]) -> None:
+def bulk_set_features_in_records(
+    records: Iterable[Record], using: str | None = None
+) -> None:
     import numpy as np
 
     """Bulk-set lazy feature dictionaries for records.
@@ -2241,6 +2249,8 @@ def bulk_set_features_in_records(records: Iterable[Record]) -> None:
     if len(records_with_features) == 0:
         return None
 
+    instance = using if using not in (None, "default") else None
+
     batch_schema: Schema | None = None
     batch_schema_index: Feature | None = None
     prepared_records: list[
@@ -2308,6 +2318,7 @@ def bulk_set_features_in_records(records: Iterable[Record]) -> None:
     )
 
     data: dict[str, pd.Series] = {}
+    multivalued_columns: list[str] = []
     for column in ordered_columns:
         # None from entirely-null columns is not a valid sentinel for extension
         # dtypes (StringDtype, BooleanDtype, Int64Dtype) — convert to pd.NA
@@ -2316,22 +2327,30 @@ def bulk_set_features_in_records(records: Iterable[Record]) -> None:
             for v in (row.get(column, pd.NA) for row in prepared_rows)
         ]
         target_dtype = feature_dtype_by_name.get(column)
-        if target_dtype is not None:
+        is_multivalued = any(
+            isinstance(v, (list, tuple, set, np.ndarray)) for v in values
+        )
+        if target_dtype is not None and not is_multivalued:
             data[column] = pd.Series(values, dtype=target_dtype)
         else:
-            data[column] = pd.Series(values)
+            data[column] = pd.Series(values, dtype="object")
+        if is_multivalued:
+            multivalued_columns.append(column)
     if data:
         dataframe = pd.DataFrame(data)
     else:
         dataframe = pd.DataFrame(index=range(len(prepared_rows)))
 
     dataframe = move_schema_index_column_to_dataframe_index(dataframe, batch_schema)
+    # set after the index move, which does not preserve `.attrs`
+    for column in multivalued_columns:
+        dataframe.attrs[column] = "list_of_categories"
     # Single-pass dataframe curation:
     # validate schema and resolve categoricals once for the entire batch.
     #
     # The resolved label records are then reused below when creating per-record
     # link rows, avoiding repeated registry calls for each row.
-    curator = DataFrameCurator(dataframe, batch_schema)
+    curator = DataFrameCurator(dataframe, batch_schema, instance=instance)
     curator.validate()
 
     members_by_name: dict[str, list[Feature]] = defaultdict(list)
@@ -2412,18 +2431,18 @@ def bulk_set_features_in_records(records: Iterable[Record]) -> None:
         )
     FeatureManager._raise_not_validated_values(not_validated_values)
     if feature_json_values:
-        save(feature_json_values)
+        save(feature_json_values, using=instance)
     for links in links_by_model.values():
         try:
-            save(links, ignore_conflicts=False)
+            save(links, ignore_conflicts=False, using=instance)
         except Exception:
-            save(links, ignore_conflicts=True)
+            save(links, ignore_conflicts=True, using=instance)
     from .save import bulk_update
 
     if batch_schema_index is not None:
         # only `name` was modified (via strip_index_for_record_persistence)
         # updating all fields generates a massive CASE WHEN SQL for large batches
-        bulk_update(records_with_features, update_fields=["name"])
+        bulk_update(records_with_features, update_fields=["name"], using=instance)
     for record in records_with_features:
         del record._features
     return None

lamindb/models/_from_values.py

@@ -57,6 +57,7 @@ def _from_values(
     standardize: bool = True,
     from_source: bool = True,
     mute: bool = False,
+    using: str | None = None,
     **filter_kwargs,
 ) -> SQLRecordList:
     """Get or create records from iterables."""
@@ -84,6 +85,7 @@ def _from_values(
         field=field,
         organism=organism_record,
         mute=mute,
+        using=using,
         **filter_kwargs,
     )
 
@@ -135,6 +137,7 @@ def get_existing_records(
     organism: SQLRecord | None = None,
     standardize: bool = True,
     mute: bool = False,
+    using: str | None = None,
     **filter_kwargs,
 ) -> tuple[list, Index, str]:
     """Get existing records from the database."""
@@ -144,7 +147,7 @@ def get_existing_records(
 
     # NOTE: existing records matching is agnostic to the source
     registry = field.field.model  # type: ignore
-    queryset = registry.filter(**filter_kwargs)
+    queryset = registry.filter(_using_key=using, **filter_kwargs)
 
     if standardize:
         # log synonyms mapped terms