lsst-daf-butler 29.2025.2100__tar.gz → 29.2025.2300__tar.gz

{lsst_daf_butler-29.2025.2100/python/lsst_daf_butler.egg-info → lsst_daf_butler-29.2025.2300}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lsst-daf-butler
-Version: 29.2025.2100
+Version: 29.2025.2300
 Summary: An abstraction layer for reading and writing astronomical data to datastores.
 Author-email: Rubin Observatory Data Management <dm-admin@lists.lsst.org>
 License: BSD 3-Clause License

{lsst_daf_butler-29.2025.2100 → lsst_daf_butler-29.2025.2300}/python/lsst/daf/butler/cli/cmd/commands.py

@@ -424,6 +424,21 @@ def prune_datasets(**kwargs: Any) -> None:
         case_sensitive=False,
     ),
 )
+@click.option(
+    "-t",
+    "--show-dataset-types",
+    is_flag=True,
+    help="Also show the dataset types registered within each collection.",
+)
+@click.option(
+    "--exclude-dataset-types",
+    type=click.STRING,
+    multiple=True,
+    default=["*_config,*_log,*_metadata,packages"],
+    callback=split_commas,
+    show_default=True,
+    help="Dataset types (comma-separated) to exclude. Only valid with --show-dataset-types.",
+)
 @options_file_option()
 def query_collections(*args: Any, **kwargs: Any) -> None:
     """Get the collections whose names match an expression."""
@@ -454,7 +469,7 @@ def query_dataset_types(*args: Any, **kwargs: Any) -> None:
     """Get the dataset types in a repository."""
     table = script.queryDatasetTypes(*args, **kwargs)
     if table:
-        table.pprint_all()
+        table.pprint_all(align="<")
     else:
         print("No results. Try --help for more information.")
 

{lsst_daf_butler-29.2025.2100 → lsst_daf_butler-29.2025.2300}/python/lsst/daf/butler/cli/utils.py

@@ -1057,7 +1057,7 @@ class MWCommand(click.Command):
         return ret
 
     @epilog.setter
-    def epilog(self, val: str) -> None:
+    def epilog(self, val: str | None) -> None:
         self._epilog = val
 
 

{lsst_daf_butler-29.2025.2100 → lsst_daf_butler-29.2025.2300}/python/lsst/daf/butler/direct_butler/_direct_butler.py

@@ -44,7 +44,7 @@ import os
 import uuid
 import warnings
 from collections import Counter, defaultdict
-from collections.abc import Iterable, Iterator, MutableMapping, Sequence
+from collections.abc import Collection, Iterable, Iterator, MutableMapping, Sequence
 from types import EllipsisType
 from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, TextIO, cast
 
@@ -1491,36 +1491,15 @@ class DirectButler(Butler):  # numpydoc ignore=PR02
             # Get all the datasets from these runs.
             refs = self._query_all_datasets(names, find_first=False, limit=None)
 
-        # Chunk the deletions using a size that is reasonably efficient whilst
-        # also giving reasonable feedback to the user. Chunking also minimizes
-        # what needs to rollback if there is a failure and should allow
-        # incremental re-running of the command. The only issue will be if the
-        # Ctrl-C comes during emptyTrash since an admin command would need to
-        # run to finish the emptying of that chunk.
-        progress = Progress("lsst.daf.butler.Butler.ingest", level=_LOG.INFO)
-        chunk_size = 50_000
-        n_chunks = math.ceil(len(refs) / chunk_size)
-        chunk_num = 0
-        for chunked_refs in progress.wrap(
-            chunk_iterable(refs, chunk_size=chunk_size), desc="Deleting datasets from RUNs", total=n_chunks
+        # Call pruneDatasets since we are deliberately removing
+        # datasets in chunks from the RUN collections rather than
+        # attempting to remove everything at once.
+        with time_this(
+            _LOG,
+            msg="Removing %d dataset%s from %s",
+            args=(len(refs), "s" if len(refs) != 1 else "", ", ".join(names)),
         ):
-            chunk_num += 1
-            _LOG.verbose(
-                "Deleting %d dataset%s from requested RUN collections in chunk %d/%d",
-                len(chunked_refs),
-                "s" if len(chunked_refs) != 1 else "",
-                chunk_num,
-                n_chunks,
-            )
-            # Call pruneDatasets since we are deliberately removing
-            # datasets in chunks from the RUN collections rather than
-            # attempting to remove everything at once.
-            with time_this(
-                _LOG,
-                msg="Removing %d datasets for chunk %d/%d",
-                args=(len(chunked_refs), chunk_num, n_chunks),
-            ):
-                self.pruneDatasets(chunked_refs, unstore=True, purge=True, disassociate=True)
+            self.pruneDatasets(refs, unstore=True, purge=True, disassociate=True)
 
         # Now can remove the actual RUN collection and unlink from chains.
         with self._registry.transaction():
@@ -1577,6 +1556,49 @@ class DirectButler(Butler):  # numpydoc ignore=PR02
         for ref in refs:
             if ref.datasetType.component():
                 raise ValueError(f"Can not prune a component of a dataset (ref={ref})")
+
+        # Chunk the deletions using a size that is reasonably efficient whilst
+        # also giving reasonable feedback to the user. Chunking also minimizes
+        # what needs to rollback if there is a failure and should allow
+        # incremental re-running of the pruning (assuming the query is
+        # repeated). The only issue will be if the Ctrl-C comes during
+        # emptyTrash since an admin command would need to run to finish the
+        # emptying of that chunk.
+        progress = Progress("lsst.daf.butler.Butler.pruneDatasets", level=_LOG.INFO)
+        chunk_size = 50_000
+        n_chunks = math.ceil(len(refs) / chunk_size)
+        if n_chunks > 1:
+            _LOG.verbose("Pruning a total of %d datasets", len(refs))
+        chunk_num = 0
+        for chunked_refs in progress.wrap(
+            chunk_iterable(refs, chunk_size=chunk_size), desc="Deleting datasets", total=n_chunks
+        ):
+            chunk_num += 1
+            _LOG.verbose(
+                "Pruning %d dataset%s in chunk %d/%d",
+                len(chunked_refs),
+                "s" if len(chunked_refs) != 1 else "",
+                chunk_num,
+                n_chunks,
+            )
+            with time_this(
+                _LOG,
+                msg="Removing %d datasets for chunk %d/%d",
+                args=(len(chunked_refs), chunk_num, n_chunks),
+            ):
+                self._prune_datasets(
+                    chunked_refs, tags=tags, unstore=unstore, purge=purge, disassociate=disassociate
+                )
+
+    def _prune_datasets(
+        self,
+        refs: Collection[DatasetRef],
+        *,
+        disassociate: bool = True,
+        unstore: bool = False,
+        tags: Iterable[str] = (),
+        purge: bool = False,
+    ) -> None:
         # We don't need an unreliable Datastore transaction for this, because
         # we've been extra careful to ensure that Datastore.trash only involves
         # mutating the Registry (it can _look_ at Datastore-specific things,

{lsst_daf_butler-29.2025.2100 → lsst_daf_butler-29.2025.2300}/python/lsst/daf/butler/direct_query_driver/_driver.py

@@ -571,12 +571,15 @@ class DirectQueryDriver(QueryDriver):
                 union_dataset_dimensions=tree.any_dataset.dimensions,
                 projection_columns=projection_columns,
                 final_columns=final_columns,
+                find_first_dataset=find_first_dataset,
             )
         else:
+            assert find_first_dataset != qt.AnyDatasetType.ANY_DATASET
             builder = SingleSelectQueryBuilder(
                 tree_analysis=query_tree_analysis,
                 projection_columns=projection_columns,
                 final_columns=final_columns,
+                find_first_dataset=find_first_dataset,
             )
         # Finish setting up the projection part of the builder.
         builder.analyze_projection()
@@ -590,7 +593,7 @@ class DirectQueryDriver(QueryDriver):
         builder.joins_analysis.columns.update(builder.projection_columns)
         # Set up the find-first part of the builder.
         if find_first_dataset is not None:
-            builder.analyze_find_first(find_first_dataset)
+            builder.analyze_find_first()
         # At this point, analysis is complete, and we can proceed to making
         # the select_builder(s) reflect that analysis.
         if not analyze_only:
@@ -1060,6 +1063,7 @@ class DirectQueryDriver(QueryDriver):
         # it depends on the kinds of collection(s) we're searching and whether
         # it's a find-first query.
         for dataset_type, fields_for_dataset in projection_columns.dataset_fields.items():
+            is_find_first = dataset_type == find_first_dataset
             dataset_search: ResolvedDatasetSearch[Any]
             if dataset_type is qt.ANY_DATASET:
                 assert union_datasets is not None
@@ -1077,7 +1081,7 @@ class DirectQueryDriver(QueryDriver):
                         derived_fields.append((dataset_type, "collection_key"))
                 elif dataset_field == "timespan" and dataset_search.is_calibration_search:
                     # The timespan is also a unique key...
-                    if dataset_type == find_first_dataset:
+                    if is_find_first:
                         # ...unless we're doing a find-first search on this
                         # dataset, in which case we need to use ANY_VALUE on
                         # the timespan and check that _VALIDITY_MATCH_COUNT
@@ -1096,6 +1100,20 @@ class DirectQueryDriver(QueryDriver):
                         ].apply_any_aggregate(self.db.apply_any_aggregate)
                     else:
                         unique_keys.extend(select_builder.joins.timespans[dataset_type].flatten())
+                elif (
+                    dataset_field == "dataset_id"
+                    and len(dataset_search.collection_records) > 1
+                    and not is_find_first
+                ):
+                    # If we have more than one collection in the search, we can
+                    # find multiple dataset IDs with the same data ID, in
+                    # different collections.
+                    # In a non-find-first search, we have to make dataset ID a
+                    # unique key to prevent de-duplication for rows with the
+                    # same data ID but different dataset IDs.
+                    # We don't do this for a find-first search because the
+                    # window function will take care of it.
+                    unique_keys.append(select_builder.joins.fields[dataset_type]["dataset_id"])
                 else:
                     # Other dataset fields derive their uniqueness from key
                     # fields.

{lsst_daf_butler-29.2025.2100 → lsst_daf_butler-29.2025.2300}/python/lsst/daf/butler/direct_query_driver/_query_builder.py

@@ -78,6 +78,10 @@ class QueryBuilder(ABC):
         or DISTINCT may be performed.
     final_columns : `.queries.tree.ColumnSet`
         Columns to include in the final query.
+    find_first_dataset : `str` or ``...`` or None
+        Name of the dataset type that needs a find-first search.  ``...``
+        is used to indicate the dataset types in a union dataset query.
+        `None` means find-first is not used.
     """
 
     def __init__(
@@ -86,13 +90,14 @@ class QueryBuilder(ABC):
         *,
         projection_columns: qt.ColumnSet,
         final_columns: qt.ColumnSet,
+        find_first_dataset: str | qt.AnyDatasetType | None,
     ):
         self.joins_analysis = tree_analysis.joins
         self.postprocessing = tree_analysis.postprocessing
         self.projection_columns = projection_columns
         self.final_columns = final_columns
         self.needs_dimension_distinct = False
-        self.find_first_dataset = None
+        self.find_first_dataset = find_first_dataset
 
     joins_analysis: QueryJoinsAnalysis
     """Description of the "joins" stage of query construction."""
@@ -161,7 +166,7 @@ class QueryBuilder(ABC):
             self.needs_dimension_distinct = True
 
     @abstractmethod
-    def analyze_find_first(self, find_first_dataset: str | qt.AnyDatasetType) -> None:
+    def analyze_find_first(self) -> None:
         """Analyze the "find first" stage of query construction, in  which a
         Common Table Expression with PARTITION ON may be used to find the first
         dataset for each data ID and dataset type in an ordered collection
@@ -169,12 +174,6 @@ class QueryBuilder(ABC):
 
         This modifies the builder in place, and should be called immediately
         after `analyze_projection`.
-
-        Parameters
-        ----------
-        find_first_dataset : `str` or ``...``
-            Name of the dataset type that needs a find-first search.  ``...``
-            is used to indicate the dataset types in a union dataset query.
         """
         raise NotImplementedError()
 
@@ -284,6 +283,47 @@ class QueryBuilder(ABC):
         """
         raise NotImplementedError()
 
+    def _needs_collection_key_field(
+        self, dataset_search: ResolvedDatasetSearch, fields_for_dataset: set[str]
+    ) -> bool:
+        """Return `True` if the ``collection_key`` dataset field is needed to
+        provide uniqueness for rows.
+        """
+        # For a dataset search, we sometimes want just one row for each dataset
+        # and sometimes we need multiple rows, one for each collection that
+        # the dataset was found in.
+        #
+        # We need multiple rows if any of the following are true:
+        # - This is a find-first dataset search.  The rows will be ranked using
+        # a window function to determine the first collection containing a
+        # matching dataset, so we need a row for each collection to feed into
+        # the window.
+        # - The user requested dataset fields that differ depending on which
+        # collection the dataset was found in, so we need a row for each
+        # collection to get all the possible values for the dataset fields.
+        #
+        # To ensure that we keep the necessary rows after DISTINCT or GROUP BY
+        # is applied, we add a "collection_key" field that is unique for each
+        # collection.
+
+        # If there is only one collection, there will only be one row per
+        # dataset, so we don't need to disambiguate.
+        if len(dataset_search.collection_records) > 1:
+            if (
+                # We need a row for each collection, which will later
+                # be filtered down using the window function.
+                self.find_first_dataset is not None
+                # We might have multiple calibration collections containing the
+                # same dataset with the same timespan.
+                or "timespan" in fields_for_dataset
+                # The user specifically asked for a row for each collection we
+                # found the dataset in.
+                or "collection" in fields_for_dataset
+            ):
+                return True
+
+        return False
+
 
 class SingleSelectQueryBuilder(QueryBuilder):
     """An implementation of `QueryBuilder` for queries that are structured as
@@ -304,6 +344,9 @@ class SingleSelectQueryBuilder(QueryBuilder):
         or DISTINCT may be performed.
     final_columns : `.queries.tree.ColumnSet`
         Columns to include in the final query.
+    find_first_dataset : `str` or None
+        Name of the dataset type that needs a find-first search.
+        `None` means find-first is not used.
     """
 
     def __init__(
@@ -312,11 +355,13 @@ class SingleSelectQueryBuilder(QueryBuilder):
         *,
         projection_columns: qt.ColumnSet,
         final_columns: qt.ColumnSet,
+        find_first_dataset: str | None,
     ) -> None:
         super().__init__(
             tree_analysis=tree_analysis,
             projection_columns=projection_columns,
             final_columns=final_columns,
+            find_first_dataset=find_first_dataset,
         )
         assert not tree_analysis.union_datasets, "UnionQueryPlan should be used instead."
         self._select_builder = tree_analysis.initial_select_builder
@@ -359,13 +404,16 @@ class SingleSelectQueryBuilder(QueryBuilder):
         # or GROUP BY columns.
         for dataset_type, fields_for_dataset in self.projection_columns.dataset_fields.items():
             assert dataset_type is not qt.ANY_DATASET, "Union dataset in non-dataset-union query."
-            if len(self.joins_analysis.datasets[dataset_type].collection_records) > 1:
+            if self._needs_collection_key_field(
+                self.joins_analysis.datasets[dataset_type], fields_for_dataset
+            ):
                 fields_for_dataset.add("collection_key")
 
-    def analyze_find_first(self, find_first_dataset: str | qt.AnyDatasetType) -> None:
+    def analyze_find_first(self) -> None:
         # Docstring inherited.
-        assert find_first_dataset is not qt.ANY_DATASET, "No dataset union in this query"
-        self.find_first = QueryFindFirstAnalysis(self.joins_analysis.datasets[find_first_dataset])
+        assert self.find_first_dataset is not qt.ANY_DATASET, "No dataset union in this query"
+        assert self.find_first_dataset is not None
+        self.find_first = QueryFindFirstAnalysis(self.joins_analysis.datasets[self.find_first_dataset])
         # If we're doing a find-first search and there's a calibration
         # collection in play, we need to make sure the rows coming out of
         # the base query have only one timespan for each data ID +
@@ -502,6 +550,10 @@ class UnionQueryBuilder(QueryBuilder):
         Columns to include in the final query.
     union_dataset_dimensions : `DimensionGroup`
         Dimensions of the dataset types that comprise the union.
+    find_first_dataset : `str` or ``...`` or None
+        Name of the dataset type that needs a find-first search.  ``...``
+        is used to indicate the dataset types in a union dataset query.
+        `None` means find-first is not used.
 
     Notes
     -----
@@ -522,11 +574,13 @@ class UnionQueryBuilder(QueryBuilder):
         projection_columns: qt.ColumnSet,
         final_columns: qt.ColumnSet,
         union_dataset_dimensions: DimensionGroup,
+        find_first_dataset: str | qt.AnyDatasetType | None,
     ):
         super().__init__(
             tree_analysis=tree_analysis,
             projection_columns=projection_columns,
             final_columns=final_columns,
+            find_first_dataset=find_first_dataset,
         )
         self._initial_select_builder: SqlSelectBuilder | None = tree_analysis.initial_select_builder
         self.union_dataset_dimensions = union_dataset_dimensions
@@ -580,15 +634,18 @@ class UnionQueryBuilder(QueryBuilder):
                     # If there is more than one collection for one union term,
                     # we need to add collection_key to all of them to keep the
                     # SELECT columns uniform.
-                    if len(union_term.datasets.collection_records) > 1:
+                    if self._needs_collection_key_field(union_term.datasets, fields_for_dataset):
                         fields_for_dataset.add("collection_key")
                         break
-            elif len(self.joins_analysis.datasets[dataset_type].collection_records) > 1:
+            elif self._needs_collection_key_field(
+                self.joins_analysis.datasets[dataset_type], fields_for_dataset
+            ):
                 fields_for_dataset.add("collection_key")
 
-    def analyze_find_first(self, find_first_dataset: str | qt.AnyDatasetType) -> None:
+    def analyze_find_first(self) -> None:
         # Docstring inherited.
-        if find_first_dataset is qt.ANY_DATASET:
+        assert self.find_first_dataset is not None
+        if self.find_first_dataset is qt.ANY_DATASET:
             for union_term in self.union_terms:
                 union_term.find_first = QueryFindFirstAnalysis(union_term.datasets)
                 # If we're doing a find-first search and there's a calibration
@@ -609,7 +666,7 @@ class UnionQueryBuilder(QueryBuilder):
             # like it'd be useful, so it's better not to have to maintain that
             # logic branch.
             raise NotImplementedError(
-                f"Additional dataset search {find_first_dataset!r} can only be joined into a "
+                f"Additional dataset search {self.find_first_dataset!r} can only be joined into a "
                 "union dataset query as a constraint in data IDs, not as a find-first result."
             )
 

{lsst_daf_butler-29.2025.2100 → lsst_daf_butler-29.2025.2300}/python/lsst/daf/butler/registry/interfaces/_database.py

@@ -1696,9 +1696,8 @@ class Database(ABC):
             for k, v in content.items():
                 if k == name:
                     continue
-                column = table.columns[k]
                 # The set only has one element
-                clauses.append(column == v.pop())
+                clauses.append(table.columns[k] == v.pop())
 
             # The IN operator will not work for "infinite" numbers of
             # rows so must batch it up into distinct calls.