Repository: lsst/daf_butler
Pull request #1169: DM-49513: Reduce number of queries required when importing datasets
Author: login: dhirving
Created: 2025-03-18T21:29:54Z | Merged: 2025-03-20T18:11:39Z

Reduce the number of queries required to import datasets to the registry in `Butler.transfer_from()`.  Instead of doing separate queries for each dataset type, we now group dataset types that have the same dimensions.

This reduces the number of queries in Prompt Processing transfers back to the central repo from around 750 to 50, and reduces the number of temp tables used from 70 to 5.

## Checklist

- [x] ran Jenkins
- [x] added a release note for user-visible changes to `doc/changes`
- [ ] (if changing dimensions.yaml) make a copy of dimensions.yaml in `configs/old_dimensions`



Diff:
diff --git a/doc/changes/DM-49513.perf.md b/doc/changes/DM-49513.perf.md
new file mode 100644
index 0000000000..8cedc4ee89
--- /dev/null
+++ b/doc/changes/DM-49513.perf.md
@@ -0,0 +1 @@
+`Butler.transfer_from()` now uses fewer database queries when inserting datasets of multiple dataset types, when some of the dataset types have the same dimensions.
diff --git a/python/lsst/daf/butler/direct_butler/_direct_butler.py b/python/lsst/daf/butler/direct_butler/_direct_butler.py
index 32eef89ccb..0dadc578e5 100644
--- a/python/lsst/daf/butler/direct_butler/_direct_butler.py
+++ b/python/lsst/daf/butler/direct_butler/_direct_butler.py
@@ -72,7 +72,7 @@
 from .._timespan import Timespan
 from ..datastore import Datastore, NullDatastore
 from ..datastores.file_datastore.retrieve_artifacts import ZipIndex, retrieve_and_zip
-from ..dimensions import DataCoordinate, Dimension
+from ..dimensions import DataCoordinate, Dimension, DimensionGroup
 from ..direct_query_driver import DirectQueryDriver
 from ..progress import Progress
 from ..queries import Query
@@ -1977,12 +1977,12 @@ def transfer_from(
                 filtered_count,
             )
 
-        # Importing requires that we group the refs by dataset type and run
+        # Importing requires that we group the refs by dimension group and run
         # before doing the import.
         source_dataset_types = set()
         grouped_refs: defaultdict[_RefGroup, list[DatasetRef]] = defaultdict(list)
         for ref in source_refs:
-            grouped_refs[_RefGroup(ref.datasetType.name, ref.run)].append(ref)
+            grouped_refs[_RefGroup(ref.datasetType.dimensions, ref.run)].append(ref)
             source_dataset_types.add(ref.datasetType)
 
         # Check to see if the dataset type in the source butler has
@@ -2098,7 +2098,7 @@ def transfer_from(
             # where Postgres can deadlock due to the unique index on collection
             # name. (See DM-47543).
             groups = sorted(grouped_refs.items(), key=lambda item: item[0].run)
-            for (datasetType, run), refs_to_import in progress.iter_item_chunks(
+            for (dimension_group, run), refs_to_import in progress.iter_item_chunks(
                 groups, desc="Importing to registry by run and dataset type"
             ):
                 if run not in handled_collections:
@@ -2117,10 +2117,10 @@ def transfer_from(
 
                 n_refs = len(refs_to_import)
                 _LOG.verbose(
-                    "Importing %d ref%s of dataset type %s into run %s",
+                    "Importing %d ref%s with dimensions %s into run %s",
                     n_refs,
                     "" if n_refs == 1 else "s",
-                    datasetType,
+                    dimension_group.names,
                     run,
                 )
 
@@ -2131,11 +2131,6 @@ def transfer_from(
                 else:
                     imported_refs = refs_to_import
 
-                # Check all the UUIDs were imported, ignoring dataset type
-                # storage class differences.
-                imported_uuids = set(ref.id for ref in imported_refs)
-                uuids_to_import = set(ref.id for ref in refs_to_import)
-                assert set(imported_uuids) == set(uuids_to_import)
                 n_imported += len(imported_refs)
 
             assert len(source_refs) == n_imported
@@ -2385,5 +2380,5 @@ class _RefGroup(NamedTuple):
     `Butler.transfer_from`.
     """
 
-    dataset_type: str
+    dimensions: DimensionGroup
     run: str
diff --git a/python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py b/python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py
index 8736a9187c..db3d85adad 100644
--- a/python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py
+++ b/python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py
@@ -698,19 +698,30 @@ def insert(
             for dataId, row in zip(data_id_list, rows, strict=True)
         ]
 
-    def import_(
-        self,
-        dataset_type: DatasetType,
-        run: RunRecord,
-        data_ids: Mapping[DatasetId, DataCoordinate],
-    ) -> list[DatasetRef]:
+    def import_(self, run: RunRecord, refs: list[DatasetRef]) -> None:
         # Docstring inherited from DatasetRecordStorageManager.
-        if not data_ids:
+        if not refs:
             # Just in case an empty mapping is provided we want to avoid
             # adding dataset type to summary tables.
-            return []
-        if (storage := self._find_storage(dataset_type.name)) is None:
-            raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
+            return
+        assert all(ref.run == run.name for ref in refs), (
+            "Run names in refs must match the run we are inserting into"
+        )
+
+        dataset_types = {ref.datasetType for ref in refs}
+        dimensions = _ensure_dimension_groups_match(dataset_types)
+        dataset_type_storage: dict[str, _DatasetRecordStorage] = {}
+        for dt in dataset_types:
+            if (storage := self._find_storage(dt.name)) is None:
+                raise MissingDatasetTypeError(f"Dataset type {dt.name!r} has not been registered.")
+            dataset_type_storage[dt.name] = storage
+
+        dynamic_tables = self._cache.get_by_dimensions(dimensions)
+        assert dynamic_tables is not None, (
+            "Table cache should have been populated when looking up dataset types"
+        )
+        tags_table = self._get_tags_table(dynamic_tables)
+
         # Current timestamp, type depends on schema version.
         if self._use_astropy_ingest_date:
             # Astropy `now()` precision should be the same as `now()` which
@@ -718,25 +729,25 @@ def import_(
             timestamp = sqlalchemy.sql.literal(astropy.time.Time.now(), type_=ddl.AstropyTimeNsecTai)
         else:
             timestamp = sqlalchemy.sql.literal(datetime.datetime.now(datetime.UTC))
+
         # We'll insert all new rows into a temporary table
-        table_spec = makeTagTableSpec(
-            storage.dataset_type.dimensions, type(self._collections), constraints=False
-        )
+        table_spec = makeTagTableSpec(dimensions, type(self._collections), constraints=False)
         collection_fkey_name = self._collections.getCollectionForeignKeyName()
-        proto_ags_row = {
-            "dataset_type_id": storage.dataset_type_id,
-            collection_fkey_name: run.key,
-        }
         tmpRows = [
-            dict(proto_ags_row, dataset_id=dataset_id, **data_id.required)
-            for dataset_id, data_id in data_ids.items()
+            {
+                "dataset_type_id": dataset_type_storage[ref.datasetType.name].dataset_type_id,
+                collection_fkey_name: run.key,
+                "dataset_id": ref.id,
+                **ref.dataId.required,
+            }
+            for ref in refs
         ]
         with self._db.transaction(for_temp_tables=True), self._db.temporary_table(table_spec) as tmp_tags:
             # store all incoming data in a temporary table
             self._db.insert(tmp_tags, *tmpRows)
             # There are some checks that we want to make for consistency
             # of the new datasets with existing ones.
-            self._validate_import(storage, tmp_tags, run)
+            self._validate_import(dimensions, tags_table, tmp_tags, run)
             # Before we merge temporary table into dataset/tags we need to
             # drop datasets which are already there (and do not conflict).
             self._db.deleteWhere(
@@ -753,27 +764,23 @@ def import_(
                     timestamp.label("ingest_date"),
                 ),
             )
-            refs = [
-                DatasetRef(
-                    datasetType=dataset_type,
-                    id=dataset_id,
-                    dataId=dataId,
-                    run=run.name,
-                )
-                for dataset_id, dataId in data_ids.items()
-            ]
             # Update the summary tables for this collection in case this
             # is the first time this dataset type or these governor values
             # will be inserted there.
             summary = CollectionSummary()
             summary.add_datasets(refs)
-            self._summaries.update(run, [storage.dataset_type_id], summary)
+            self._summaries.update(
+                run, [storage.dataset_type_id for storage in dataset_type_storage.values()], summary
+            )
             # Copy from temp table into tags table.
-            self._db.insert(self._get_tags_table(storage.dynamic_tables), select=tmp_tags.select())
-        return refs
+            self._db.insert(tags_table, select=tmp_tags.select())
 
     def _validate_import(
-        self, storage: _DatasetRecordStorage, tmp_tags: sqlalchemy.schema.Table, run: RunRecord
+        self,
+        dimensions: DimensionGroup,
+        tags: sqlalchemy.schema.Table,
+        tmp_tags: sqlalchemy.schema.Table,
+        run: RunRecord,
     ) -> None:
         """Validate imported refs against existing datasets.
 
@@ -793,7 +800,6 @@ def _validate_import(
             Raise if new datasets conflict with existing ones.
         """
         dataset = self._static.dataset
-        tags = self._get_tags_table(storage.dynamic_tables)
         collection_fkey_name = self._collections.getCollectionForeignKeyName()
 
         # Check that existing datasets have the same dataset type and
@@ -820,23 +826,16 @@ def _validate_import(
             if (row := result.first()) is not None:
                 existing_run = self._collections[row.run].name
                 new_run = self._collections[row.new_run].name
-                if row.dataset_type_id == storage.dataset_type_id:
-                    if row.new_dataset_type_id == storage.dataset_type_id:
-                        raise ConflictingDefinitionError(
-                            f"Current run {existing_run!r} and new run {new_run!r} do not agree for "
-                            f"dataset {row.dataset_id}."
-                        )
-                    else:
-                        raise ConflictingDefinitionError(
-                            f"Dataset {row.dataset_id} was provided with type {storage.dataset_type.name!r} "
-                            f"in run {new_run!r}, but was already defined with type ID {row.dataset_type_id} "
-                            f"in run {run!r}."
-                        )
+                if row.dataset_type_id == row.new_dataset_type_id:
+                    raise ConflictingDefinitionError(
+                        f"Current run {existing_run!r} and new run {new_run!r} do not agree for "
+                        f"dataset {row.dataset_id}."
+                    )
                 else:
                     raise ConflictingDefinitionError(
                         f"Dataset {row.dataset_id} was provided with type ID {row.new_dataset_type_id} "
-                        f"in run {new_run!r}, but was already defined with type "
-                        f"{storage.dataset_type.name!r} in run {run!r}."
+                        f"in run {new_run!r}, but was already defined with type ID "
+                        f"{row.dataset_type_id} in run {run!r}."
                     )
 
         # Check that matching dataset in tags table has the same DataId.
@@ -845,20 +844,14 @@ def _validate_import(
                 tags.columns.dataset_id,
                 tags.columns.dataset_type_id.label("type_id"),
                 tmp_tags.columns.dataset_type_id.label("new_type_id"),
-                *[tags.columns[dim] for dim in storage.dataset_type.dimensions.required],
-                *[
-                    tmp_tags.columns[dim].label(f"new_{dim}")
-                    for dim in storage.dataset_type.dimensions.required
-                ],
+                *[tags.columns[dim] for dim in dimensions.required],
+                *[tmp_tags.columns[dim].label(f"new_{dim}") for dim in dimensions.required],
             )
             .select_from(tags.join(tmp_tags, tags.columns.dataset_id == tmp_tags.columns.dataset_id))
             .where(
                 sqlalchemy.sql.or_(
                     tags.columns.dataset_type_id != tmp_tags.columns.dataset_type_id,
-                    *[
-                        tags.columns[dim] != tmp_tags.columns[dim]
-                        for dim in storage.dataset_type.dimensions.required
-                    ],
+                    *[tags.columns[dim] != tmp_tags.columns[dim] for dim in dimensions.required],
                 )
             )
             .limit(1)
@@ -874,9 +867,10 @@ def _validate_import(
         # Check that matching run+dataId have the same dataset ID.
         query = (
             sqlalchemy.sql.select(
-                *[tags.columns[dim] for dim in storage.dataset_type.dimensions.required],
+                *[tags.columns[dim] for dim in dimensions.required],
                 tags.columns.dataset_id,
                 tmp_tags.columns.dataset_id.label("new_dataset_id"),
+                tmp_tags.columns.dataset_type_id.label("new_dataset_type_id"),
                 tags.columns[collection_fkey_name],
                 tmp_tags.columns[collection_fkey_name].label(f"new_{collection_fkey_name}"),
             )
@@ -886,10 +880,7 @@ def _validate_import(
                     sqlalchemy.sql.and_(
                         tags.columns.dataset_type_id == tmp_tags.columns.dataset_type_id,
                         tags.columns[collection_fkey_name] == tmp_tags.columns[collection_fkey_name],
-                        *[
-                            tags.columns[dim] == tmp_tags.columns[dim]
-                            for dim in storage.dataset_type.dimensions.required
-                        ],
+                        *[tags.columns[dim] == tmp_tags.columns[dim] for dim in dimensions.required],
                     ),
                 )
             )
@@ -899,11 +890,11 @@ def _validate_import(
         with self._db.query(query) as result:
             # only include the first one in the exception message
             if (row := result.first()) is not None:
-                data_id = {dim: getattr(row, dim) for dim in storage.dataset_type.dimensions.required}
+                data_id = {dim: getattr(row, dim) for dim in dimensions.required}
                 existing_collection = self._collections[getattr(row, collection_fkey_name)].name
                 new_collection = self._collections[getattr(row, f"new_{collection_fkey_name}")].name
                 raise ConflictingDefinitionError(
-                    f"Dataset with type {storage.dataset_type.name!r} and data ID {data_id} "
+                    f"Dataset with type ID {row.new_dataset_type_id} and data ID {data_id} "
                     f"has ID {row.dataset_id} in existing collection {existing_collection!r} "
                     f"but ID {row.new_dataset_id} in new collection {new_collection!r}."
                 )
@@ -1674,3 +1665,15 @@ def _create_case_expression_for_collections(
         return sqlalchemy.cast(sqlalchemy.null(), sqlalchemy.String)
 
     return sqlalchemy.case(mapping, value=id_column)
+
+
+def _ensure_dimension_groups_match(dataset_types: Iterable[DatasetType]) -> DimensionGroup:
+    dimensions = set(dt.dimensions for dt in dataset_types)
+    assert len(dimensions) > 0, "At least one dataset type is required"
+    if len(dimensions) != 1:
+        raise DatasetTypeError(
+            "Dataset types have more than one dimension group.\n"
+            f"Dataset types: {dataset_types}\n"
+            f"Dimension groups: {dimensions}"
+        )
+    return dimensions.pop()
diff --git a/python/lsst/daf/butler/registry/interfaces/_datasets.py b/python/lsst/daf/butler/registry/interfaces/_datasets.py
index 41f1ccef39..e86203000d 100644
--- a/python/lsst/daf/butler/registry/interfaces/_datasets.py
+++ b/python/lsst/daf/butler/registry/interfaces/_datasets.py
@@ -424,29 +424,17 @@ def insert(
         raise NotImplementedError()
 
     @abstractmethod
-    def import_(
-        self,
-        dataset_type: DatasetType,
-        run: RunRecord,
-        data_ids: Mapping[DatasetId, DataCoordinate],
-    ) -> list[DatasetRef]:
+    def import_(self, run: RunRecord, refs: list[DatasetRef]) -> None:
         """Insert one or more dataset entries into the database.
 
         Parameters
         ----------
-        dataset_type : `DatasetType`
-            Type of dataset to import.  Also used as the dataset type for
-            the returned refs.
         run : `RunRecord`
             The record object describing the `~CollectionType.RUN` collection
             these datasets will be associated with.
-        data_ids : `~collections.abc.Mapping`
-            Mapping from dataset ID to data ID.
-
-        Returns
-        -------
-        datasets : `list` [ `DatasetRef` ]
-            References to the inserted or existing datasets.
+        refs : `list` [ `DatasetRef` ]
+            List of datasets to be be inserted.  All of the ``DatasetRef``
+            ``run`` attributes must match the ``run`` parameter.
         """
         raise NotImplementedError()
 
diff --git a/python/lsst/daf/butler/registry/sql_registry.py b/python/lsst/daf/butler/registry/sql_registry.py
index d4c0d45e24..cf8699b3de 100644
--- a/python/lsst/daf/butler/registry/sql_registry.py
+++ b/python/lsst/daf/butler/registry/sql_registry.py
@@ -76,12 +76,10 @@
     CollectionSummary,
     CollectionTypeError,
     ConflictingDefinitionError,
-    DatasetTypeError,
     MissingDatasetTypeError,
     NoDefaultCollectionError,
     OrphanedRecordError,
     RegistryConfig,
-    RegistryConsistencyError,
     RegistryDefaults,
     queries,
 )
@@ -1103,13 +1101,12 @@ def _importDatasets(
         ----------
         datasets : `~collections.abc.Iterable` of `DatasetRef`
             Datasets to be inserted. All `DatasetRef` instances must have
-            identical ``datasetType`` and ``run`` attributes. ``run``
+            identical ``run`` attributes. ``run``
             attribute can be `None` and defaults to ``self.defaults.run``.
             Datasets can specify ``id`` attribute which will be used for
-            inserted datasets. All dataset IDs must have the same type
-            (`int` or `uuid.UUID`), if type of dataset IDs does not match
-            configured backend then IDs will be ignored and new IDs will be
-            generated by backend.
+            inserted datasets.
+            Datasets can be of multiple dataset types, but all the dataset
+            types must have the same set of dimensions.
         expand : `bool`, optional
             If `True` (default), expand data IDs as they are inserted.  This is
             necessary in general, but it may be disabled if the caller can
@@ -1128,8 +1125,7 @@ def _importDatasets(
         lsst.daf.butler.registry.NoDefaultCollectionError
             Raised if ``run`` is `None` and ``self.defaults.run`` is `None`.
         lsst.daf.butler.registry.DatasetTypeError
-            Raised if datasets correspond to more than one dataset type or
-            dataset type is not known to registry.
+            Raised if a dataset type is not known to registry.
         lsst.daf.butler.registry.ConflictingDefinitionError
             If a dataset with the same dataset type and data ID as one of those
             given already exists in ``run``.
@@ -1147,12 +1143,6 @@ def _importDatasets(
             # nothing to do
             return []
 
-        # find dataset type
-        datasetTypes = {dataset.datasetType for dataset in datasets}
-        if len(datasetTypes) != 1:
-            raise DatasetTypeError(f"Multiple dataset types in input datasets: {datasetTypes}")
-        datasetType = datasetTypes.pop()
-
         # find run name
         runs = {dataset.run for dataset in datasets}
         if len(runs) != 1:
@@ -1169,34 +1159,24 @@ def _importDatasets(
 
         progress = Progress("daf.butler.Registry.importDatasets", level=logging.DEBUG)
         if expand:
-            data_ids = {
-                dataset.id: self.expandDataId(dataset.dataId, dimensions=datasetType.dimensions)
-                for dataset in progress.wrap(datasets, f"Expanding {datasetType.name} data IDs")
-            }
-        else:
-            data_ids = {dataset.id: dataset.dataId for dataset in datasets}
+            datasets = [
+                dataset.expanded(self.expandDataId(dataset.dataId, dimensions=dataset.datasetType.dimensions))
+                for dataset in progress.wrap(datasets, "Expanding data IDs")
+            ]
 
         try:
-            refs = list(self._managers.datasets.import_(datasetType, runRecord, data_ids))
+            self._managers.datasets.import_(runRecord, datasets)
             if self._managers.obscore:
-                self._managers.obscore.add_datasets(refs)
+                self._managers.obscore.add_datasets(datasets)
         except sqlalchemy.exc.IntegrityError as err:
             raise ConflictingDefinitionError(
                 "A database constraint failure was triggered by inserting "
-                f"one or more datasets of type {datasetType} into "
-                f"collection '{run}'. "
+                f"one or more datasets into collection '{run}'. "
                 "This probably means a dataset with the same data ID "
                 "and dataset type already exists, but it may also mean a "
                 "dimension row is missing."
             ) from err
-        # Check that imported dataset IDs match the input
-        for imported_ref, input_ref in zip(refs, datasets, strict=True):
-            if imported_ref.id != input_ref.id:
-                raise RegistryConsistencyError(
-                    "Imported dataset ID differs from input dataset ID, "
-                    f"input ref: {input_ref}, imported ref: {imported_ref}"
-                )
-        return refs
+        return datasets
 
     def getDataset(self, id: DatasetId) -> DatasetRef | None:
         """Retrieve a Dataset entry.
diff --git a/tests/test_butler.py b/tests/test_butler.py
index de2e419313..cb34556ce7 100644
--- a/tests/test_butler.py
+++ b/tests/test_butler.py
@@ -41,6 +41,7 @@
 import string
 import tempfile
 import unittest
+import unittest.mock
 import uuid
 from collections.abc import Callable, Mapping
 from typing import TYPE_CHECKING, Any, cast
@@ -1823,8 +1824,7 @@ def testPruneDatasets(self) -> None:
         storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
         butler = self.create_empty_butler(writeable=True)
         # Load registry data with dimensions to hang datasets off of.
-        registryDataDir = os.path.normpath(os.path.join(TESTDIR, "data", "registry"))
-        butler.import_(filename=os.path.join(registryDataDir, "base.yaml"))
+        butler.import_(filename=_get_test_data_path("base.yaml"))
         # Add some RUN-type collections.
         run1 = "run1"
         butler.collections.register(run1)
@@ -2873,6 +2873,64 @@ def assertButlerTransfers(
             to_transfer = [ref for ref in source_refs if ref.run == "run2"]
             self.target_butler.transfer_from(self.source_butler, to_transfer)
 
+    def test_shared_dimension_group(self):
+        """Test internal logic that divides dataset types by dimension group
+        when doing registry updates.
+        """
+        self.create_butlers()
+        self.source_butler.import_(filename=_get_test_data_path("base.yaml"), without_datastore=True)
+        self.source_butler.import_(filename=_get_test_data_path("datasets.yaml"), without_datastore=True)
+
+        source_butler = self.source_butler
+        target_butler = self.target_butler
+
+        # Create a dataset type with the same dimensions as the 'bias' dataset
+        # type from base.yaml
+        dataset_type = DatasetType(
+            "test_type", ["instrument", "detector"], "int", universe=source_butler.dimensions
+        )
+        source_butler.registry.registerDatasetType(dataset_type)
+        # This has the same data ID as one of the bias datasets in
+        # datasets.yaml.
+        test_ref = source_butler.registry.insertDatasets(
+            "test_type", [{"instrument": "Cam1", "detector": 2}], run="imported_g"
+        )[0]
+
+        biases = source_butler.query_datasets("bias", ["imported_g", "imported_r"])
+        flats = source_butler.query_datasets("flat", ["imported_g", "imported_r"])
+        refs = [test_ref, *biases, *flats]
+
+        # Test setup will be even more convoluted if we want the datastore to
+        # actually transfer files.  For testing the dimension group behavior,
+        # we really only care about the registry.
+        with unittest.mock.patch.object(target_butler._datastore, "transfer_from") as mock:
+            mock.return_value = (set(refs), set())
+            target_butler.transfer_from(
+                source_butler,
+                refs,
+                transfer=None,
+                register_dataset_types=True,
+                skip_missing=False,
+                transfer_dimensions=True,
+            )
+
+        transferred_test_ref = target_butler.find_dataset(
+            "test_type", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
+        )
+        self.assertEqual(transferred_test_ref.id, test_ref.id)
+
+        transferred_bias = target_butler.find_dataset(
+            "bias", {"instrument": "Cam1", "detector": 2}, collections="imported_g"
+        )
+        self.assertEqual(transferred_bias.id, uuid.UUID("51352db4-a47a-447c-b12d-a50b206b17cd"))
+
+        transferred_flat = target_butler.find_dataset(
+            "flat",
+            {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-R1", "band": "r"},
+            collections="imported_r",
+        )
+        self.assertEqual(transferred_flat.id, uuid.UUID("c1296796-56c5-4acf-9b49-40d920c6f840"))
+
 
 class ChainedDatastoreTransfers(PosixDatastoreTransfers):
     """Test transfers using a chained datastore."""
@@ -3033,6 +3091,11 @@ def setup_module(module: types.ModuleType) -> None:
     clean_environment()
 
 
+def _get_test_data_path(filename: str) -> str:
+    registryDataDir = os.path.normpath(os.path.join(TESTDIR, "data", "registry"))
+    return os.path.join(registryDataDir, filename)
+
+
 if __name__ == "__main__":
     clean_environment()
     unittest.main()


Reviews:
user: login: dhirving
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T23:05:39Z
user: login: dhirving
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T23:08:14Z
user: login: andy-slac
body: Looks great!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T00:08:58Z
user: login: dhirving
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T00:15:01Z
user: login: timj
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T00:17:05Z

Review comments:
diff_hunk: @@ -2131,11 +2131,6 @@ def transfer_from(
                 else:
                     imported_refs = refs_to_import
 
-                # Check all the UUIDs were imported, ignoring dataset type
-                # storage class differences.
-                imported_uuids = set(ref.id for ref in imported_refs)
-                uuids_to_import = set(ref.id for ref in refs_to_import)
-                assert set(imported_uuids) == set(uuids_to_import)
path: python/lsst/daf/butler/direct_butler/_direct_butler.py
user: login: dhirving
body: This assertion wasn't meaningful even before this change -- a complete set of refs with the same IDs is/was roundtripped in all cases.
created_at: 2025-03-19T23:05:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -820,23 +826,16 @@ def _validate_import(
             if (row := result.first()) is not None:
                 existing_run = self._collections[row.run].name
                 new_run = self._collections[row.new_run].name
-                if row.dataset_type_id == storage.dataset_type_id:
-                    if row.new_dataset_type_id == storage.dataset_type_id:
-                        raise ConflictingDefinitionError(
-                            f"Current run {existing_run!r} and new run {new_run!r} do not agree for "
-                            f"dataset {row.dataset_id}."
-                        )
-                    else:
-                        raise ConflictingDefinitionError(
-                            f"Dataset {row.dataset_id} was provided with type {storage.dataset_type.name!r} "
-                            f"in run {new_run!r}, but was already defined with type ID {row.dataset_type_id} "
-                            f"in run {run!r}."
-                        )
path: python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py
user: login: dhirving
body: This branch was unreachable and uncovered by tests in the original version.
created_at: 2025-03-19T23:08:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,3 @@
+`Butler.transfer_from()` now uses fewer database queries when inserting
+datasets of multiple dataset types, when some of the dataset types have the
+same dimensions.
path: doc/changes/DM-49513.perf.md
user: login: andy-slac
body: I think it should be one sentence per line.
created_at: 2025-03-19T23:30:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1169,34 +1159,24 @@ def _importDatasets(
 
         progress = Progress("daf.butler.Registry.importDatasets", level=logging.DEBUG)
         if expand:
-            data_ids = {
-                dataset.id: self.expandDataId(dataset.dataId, dimensions=datasetType.dimensions)
-                for dataset in progress.wrap(datasets, f"Expanding {datasetType.name} data IDs")
-            }
-        else:
-            data_ids = {dataset.id: dataset.dataId for dataset in datasets}
+            datasets = [
+                dataset.expanded(self.expandDataId(dataset.dataId, dimensions=dataset.datasetType.dimensions))
path: python/lsst/daf/butler/registry/sql_registry.py
user: login: andy-slac
body: This could also generate query per ref (or even multiple queries)?
created_at: 2025-03-19T23:39:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -424,29 +424,16 @@ def insert(
         raise NotImplementedError()
 
     @abstractmethod
-    def import_(
-        self,
-        dataset_type: DatasetType,
-        run: RunRecord,
-        data_ids: Mapping[DatasetId, DataCoordinate],
-    ) -> list[DatasetRef]:
+    def import_(self, run: RunRecord, refs: list[DatasetRef]) -> None:
         """Insert one or more dataset entries into the database.
 
         Parameters
         ----------
-        dataset_type : `DatasetType`
-            Type of dataset to import.  Also used as the dataset type for
-            the returned refs.
         run : `RunRecord`
             The record object describing the `~CollectionType.RUN` collection
             these datasets will be associated with.
-        data_ids : `~collections.abc.Mapping`
-            Mapping from dataset ID to data ID.
-
-        Returns
-        -------
-        datasets : `list` [ `DatasetRef` ]
-            References to the inserted or existing datasets.
+        refs : `list` [ `DatasetRef ` ]
path: python/lsst/daf/butler/registry/interfaces/_datasets.py
user: login: andy-slac
body: ```suggestion
        refs : `list` [ `DatasetRef` ]
```
created_at: 2025-03-19T23:42:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -424,29 +424,16 @@ def insert(
         raise NotImplementedError()
 
     @abstractmethod
-    def import_(
-        self,
-        dataset_type: DatasetType,
-        run: RunRecord,
-        data_ids: Mapping[DatasetId, DataCoordinate],
-    ) -> list[DatasetRef]:
+    def import_(self, run: RunRecord, refs: list[DatasetRef]) -> None:
         """Insert one or more dataset entries into the database.
 
         Parameters
         ----------
-        dataset_type : `DatasetType`
-            Type of dataset to import.  Also used as the dataset type for
-            the returned refs.
         run : `RunRecord`
             The record object describing the `~CollectionType.RUN` collection
             these datasets will be associated with.
-        data_ids : `~collections.abc.Mapping`
-            Mapping from dataset ID to data ID.
-
-        Returns
-        -------
-        datasets : `list` [ `DatasetRef` ]
-            References to the inserted or existing datasets.
+        refs : `list` [ `DatasetRef ` ]
+            List of datasets to be be inserted.
path: python/lsst/daf/butler/registry/interfaces/_datasets.py
user: login: andy-slac
body: Maybe mention that datasets' `run` must match ``run`` parameter?
created_at: 2025-03-19T23:43:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -698,45 +698,56 @@ def insert(
             for dataId, row in zip(data_id_list, rows, strict=True)
         ]
 
-    def import_(
-        self,
-        dataset_type: DatasetType,
-        run: RunRecord,
-        data_ids: Mapping[DatasetId, DataCoordinate],
-    ) -> list[DatasetRef]:
+    def import_(self, run: RunRecord, refs: list[DatasetRef]) -> None:
         # Docstring inherited from DatasetRecordStorageManager.
-        if not data_ids:
+        if not refs:
             # Just in case an empty mapping is provided we want to avoid
             # adding dataset type to summary tables.
-            return []
-        if (storage := self._find_storage(dataset_type.name)) is None:
-            raise MissingDatasetTypeError(f"Dataset type {dataset_type.name!r} has not been registered.")
+            return
+        assert all(ref.run == run.name for ref in refs), (
+            "Run names in refs must match the run we are inserting into"
+        )
+
+        dataset_types = {ref.datasetType for ref in refs}
+        dimensions = _ensure_dimension_groups_match(dataset_types)
+        dataset_type_storage: dict[str, _DatasetRecordStorage] = {}
+        for dt in dataset_types:
+            if (storage := self._find_storage(dt.name)) is None:
+                raise MissingDatasetTypeError(f"Dataset type {dt.name!r} has not been registered.")
+            dataset_type_storage[dt.name] = storage
+
+        dynamic_tables = self._cache.get_by_dimensions(dimensions)
+        assert dynamic_tables is not None, (
+            "Table cache should have been populated when looking up dataset types"
+        )
+        tags_table = self._get_tags_table(dynamic_tables)
+
         # Current timestamp, type depends on schema version.
         if self._use_astropy_ingest_date:
             # Astropy `now()` precision should be the same as `now()` which
             # should mean microsecond.
             timestamp = sqlalchemy.sql.literal(astropy.time.Time.now(), type_=ddl.AstropyTimeNsecTai)
         else:
             timestamp = sqlalchemy.sql.literal(datetime.datetime.now(datetime.UTC))
+
         # We'll insert all new rows into a temporary table
-        table_spec = makeTagTableSpec(
-            storage.dataset_type.dimensions, type(self._collections), constraints=False
-        )
+        table_spec = makeTagTableSpec(dimensions, type(self._collections), constraints=False)
         collection_fkey_name = self._collections.getCollectionForeignKeyName()
-        proto_ags_row = {
-            "dataset_type_id": storage.dataset_type_id,
-            collection_fkey_name: run.key,
-        }
         tmpRows = [
-            dict(proto_ags_row, dataset_id=dataset_id, **data_id.required)
-            for dataset_id, data_id in data_ids.items()
+            {
+                "dataset_type_id": dataset_type_storage[ref.datasetType.name].dataset_type_id,
+                collection_fkey_name: run.key,
+                "dataset_id": ref.id,
+                **ref.dataId.required,
+            }
+            for ref in refs
         ]
         with self._db.transaction(for_temp_tables=True), self._db.temporary_table(table_spec) as tmp_tags:
path: python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py
user: login: andy-slac
body: Would be nice if we could get rid of temp table here in case there are not very many records to insert.
created_at: 2025-03-19T23:59:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1169,34 +1159,24 @@ def _importDatasets(
 
         progress = Progress("daf.butler.Registry.importDatasets", level=logging.DEBUG)
         if expand:
-            data_ids = {
-                dataset.id: self.expandDataId(dataset.dataId, dimensions=datasetType.dimensions)
-                for dataset in progress.wrap(datasets, f"Expanding {datasetType.name} data IDs")
-            }
-        else:
-            data_ids = {dataset.id: dataset.dataId for dataset in datasets}
+            datasets = [
+                dataset.expanded(self.expandDataId(dataset.dataId, dimensions=dataset.datasetType.dimensions))
path: python/lsst/daf/butler/registry/sql_registry.py
user: login: dhirving
body: Yeah, I noticed that when I was in there as well.

In the prompt processing case it's not an issue -- their `DatasetRef`s are already expanded.  But I think it's worth following up on sometime soon, there are several places in the code base that do `expandDataId` in a loop over a potentially large number of refs.
created_at: 2025-03-20T00:15:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1169,34 +1159,24 @@ def _importDatasets(
 
         progress = Progress("daf.butler.Registry.importDatasets", level=logging.DEBUG)
         if expand:
-            data_ids = {
-                dataset.id: self.expandDataId(dataset.dataId, dimensions=datasetType.dimensions)
-                for dataset in progress.wrap(datasets, f"Expanding {datasetType.name} data IDs")
-            }
-        else:
-            data_ids = {dataset.id: dataset.dataId for dataset in datasets}
+            datasets = [
+                dataset.expanded(self.expandDataId(dataset.dataId, dimensions=dataset.datasetType.dimensions))
path: python/lsst/daf/butler/registry/sql_registry.py
user: login: timj
body: CM team can give you some batch graphs that contain 200k datasets to be imported by `transfer-from-graph` if you want to see the effect this has on that use case.
created_at: 2025-03-20T00:17:05Z
author_association: MEMBER

PR comments:
user: login: codecov[bot]
created_at: 2025-03-18T21:56:34Z
body: ## [Codecov](https://app.codecov.io/gh/lsst/daf_butler/pull/1169?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lsst) Report
Attention: Patch coverage is `90.32258%` with `6 lines` in your changes missing coverage. Please review.
> Project coverage is 89.37%. Comparing base [(`c91e1e7`)](https://app.codecov.io/gh/lsst/daf_butler/commit/c91e1e7985caa98bb7e2664ef51caaae84865285?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lsst) to head [(`25ecb75`)](https://app.codecov.io/gh/lsst/daf_butler/commit/25ecb7588a1061dc7e50de52882520894a633603?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lsst).
> Report is 9 commits behind head on main.

:white_check_mark: All tests successful. No failed tests found.

| [Files with missing lines](https://app.codecov.io/gh/lsst/daf_butler/pull/1169?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lsst) | Patch % | Lines |
|---|---|---|
| [.../butler/registry/datasets/byDimensions/\_manager.py](https://app.codecov.io/gh/lsst/daf_butler/pull/1169?src=pr&el=tree&filepath=python%2Flsst%2Fdaf%2Fbutler%2Fregistry%2Fdatasets%2FbyDimensions%2F_manager.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lsst#diff-cHl0aG9uL2xzc3QvZGFmL2J1dGxlci9yZWdpc3RyeS9kYXRhc2V0cy9ieURpbWVuc2lvbnMvX21hbmFnZXIucHk=) | 77.77% | [3 Missing and 3 partials :warning: ](https://app.codecov.io/gh/lsst/daf_butler/pull/1169?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lsst) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1169      +/-   ##
==========================================
+ Coverage   89.36%   89.37%   +0.01%     
==========================================
  Files         367      367              
  Lines       49666    49688      +22     
  Branches     6018     6016       -2     
==========================================
+ Hits        44383    44411      +28     
+ Misses       3860     3857       -3     
+ Partials     1423     1420       -3     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/lsst/daf_butler/pull/1169?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lsst).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lsst).

author_association: NONE
user: login: timj
created_at: 2025-03-20T00:20:31Z
body: Do you want to fix `butler.ingest()` whilst you are in there? That also has a loop over dataset type doing importDatasets.
author_association: MEMBER
user: login: dhirving
created_at: 2025-03-20T17:00:39Z
body: > Do you want to fix `butler.ingest()` whilst you are in there? That also has a loop over dataset type doing importDatasets.

I was going to open a separate ticket for that -- I want to get this one down so PP can benefit from it because they are running weekly large scale tests now.  

There is a lot of duplicated logic between `_ingest_file_datasets` and `transfer_from` that should be pushed down into `SqlRegistry.import_datasets` because this whole dimension group thing is a weird implementation detail that doesn't need to be duplicated multiple places.

I really don't want to get involved in totally optimizing every facet of this situation right now -- my goal was just to get PP to a temporarily workable state so I have a minute to help them with their architecture change, and this whole PP thing is an interrupt that I am spending more time on than I really have right now.
author_association: CONTRIBUTOR