Coverage for python/lsst/daf/butler/script/queryDatasets.py: 16%
131 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-06 01:33 -0800
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
27from __future__ import annotations
29import logging
30import warnings
31from collections import defaultdict
32from collections.abc import Iterable, Iterator
33from contextlib import contextmanager
34from itertools import chain, groupby
35from typing import TYPE_CHECKING
37import numpy as np
38from astropy.table import Table as AstropyTable
40from .._butler import Butler
41from .._exceptions import MissingDatasetTypeError
42from .._query_all_datasets import QueryAllDatasetsParameters
43from ..cli.utils import sortAstropyTable
44from ..utils import has_globs
46if TYPE_CHECKING:
47 from lsst.daf.butler import DatasetRef
48 from lsst.resources import ResourcePath
51_LOG = logging.getLogger(__name__)
54class _Table:
55 """Aggregates rows for a single dataset type, and creates an astropy table
56 with the aggregated data. Eliminates duplicate rows.
57 """
59 datasetRefs: dict[DatasetRef, str | None]
61 def __init__(self) -> None:
62 self.datasetRefs = {}
64 def add(self, datasetRef: DatasetRef, uri: ResourcePath | None = None) -> None:
65 """Add a row of information to the table.
67 ``uri`` is optional but must be the consistent; provided or not, for
68 every call to a ``_Table`` instance.
70 Parameters
71 ----------
72 datasetRef : `DatasetRef`
73 A dataset ref that will be added as a row in the table.
74 uri : `lsst.resources.ResourcePath`, optional
75 The URI to show as a file location in the table, by default `None`.
76 """
77 uri_str = str(uri) if uri else None
78 # Use a dict to retain ordering.
79 self.datasetRefs[datasetRef] = uri_str
81 def getAstropyTable(self, datasetTypeName: str, sort: bool = True) -> AstropyTable:
82 """Get the table as an astropy table.
84 Parameters
85 ----------
86 datasetTypeName : `str`
87 The dataset type name to show in the ``type`` column of the table.
88 sort : `bool`, optional
89 If `True` the table will be sorted.
91 Returns
92 -------
93 table : `astropy.table._Table`
94 The table with the provided column names and rows.
95 """
96 # Should never happen; adding a dataset should be the action that
97 # causes a _Table to be created.
98 if not self.datasetRefs:
99 raise RuntimeError(f"No DatasetRefs were provided for dataset type {datasetTypeName}")
101 ref = next(iter(self.datasetRefs))
102 dimensions = [ref.dataId.universe.dimensions[k] for k in ref.dataId.dimensions.data_coordinate_keys]
103 columnNames = ["type", "run", "id", *[str(item) for item in dimensions]]
105 # Need to hint the column types for numbers since the per-row
106 # constructor of Table does not work this out on its own and sorting
107 # will not work properly without.
108 typeMap = {float: np.float64, int: np.int64}
109 columnTypes = [
110 None,
111 None,
112 str,
113 *[typeMap.get(type(value)) for value in ref.dataId.full_values],
114 ]
115 if self.datasetRefs[ref]:
116 columnNames.append("URI")
117 columnTypes.append(None)
119 rows = []
120 for ref, uri in self.datasetRefs.items():
121 row = [
122 datasetTypeName,
123 ref.run,
124 str(ref.id),
125 *ref.dataId.full_values,
126 ]
127 if uri:
128 row.append(uri)
129 rows.append(row)
131 dataset_table = AstropyTable(np.array(rows), names=columnNames, dtype=columnTypes)
132 if sort:
133 return sortAstropyTable(dataset_table, dimensions, ["type", "run"])
134 else:
135 return dataset_table
138class QueryDatasets:
139 """Get dataset refs from a repository.
141 Parameters
142 ----------
143 glob : iterable [`str`]
144 A list of glob-style search string that fully or partially identify
145 the dataset type names to search for.
146 collections : iterable [`str`]
147 A list of glob-style search string that fully or partially identify
148 the collections to search for.
149 where : `str`
150 A string expression similar to a SQL WHERE clause. May involve any
151 column of a dimension table or (as a shortcut for the primary key
152 column of a dimension table) dimension name.
153 find_first : `bool`
154 For each result data ID, only yield one DatasetRef of each DatasetType,
155 from the first collection in which a dataset of that dataset type
156 appears (according to the order of `collections` passed in). If used,
157 `collections` must specify at least one expression and must not contain
158 wildcards.
159 show_uri : `bool`
160 If True, include the dataset URI in the output.
161 limit : `int`, optional
162 Limit the number of results to be returned. A value of 0 means
163 unlimited. A negative value is used to specify a cap where a warning
164 is issued if that cap is hit.
165 order_by : `tuple` of `str`
166 Dimensions to use for sorting results. If no ordering is given the
167 results of ``limit`` are undefined and default sorting of the resulting
168 datasets will be applied. It is an error if the requested ordering
169 is inconsistent with the dimensions of the dataset type being queried.
170 repo : `str` or `None`
171 URI to the location of the repo or URI to a config file describing the
172 repo and its location. One of `repo` and `butler` must be `None` and
173 the other must not be `None`.
174 butler : `lsst.daf.butler.Butler` or `None`
175 The butler to use to query. One of `repo` and `butler` must be `None`
176 and the other must not be `None`.
177 with_dimension_records : `bool`, optional
178 If `True` (default is `False`) then returned data IDs will have
179 dimension records.
180 """
182 def __init__(
183 self,
184 glob: Iterable[str],
185 collections: Iterable[str],
186 where: str,
187 find_first: bool,
188 show_uri: bool,
189 limit: int = 0,
190 order_by: tuple[str, ...] = (),
191 repo: str | None = None,
192 butler: Butler | None = None,
193 with_dimension_records: bool = False,
194 ):
195 if (repo and butler) or (not repo and not butler):
196 raise RuntimeError("One of repo and butler must be provided and the other must be None.")
197 collections = list(collections)
198 if not collections:
199 warnings.warn(
200 "No --collections specified. The --collections argument will become mandatory after v28.",
201 FutureWarning,
202 stacklevel=2,
203 )
204 glob = list(glob)
205 if not glob:
206 warnings.warn(
207 "No dataset types specified. Explicitly specifying dataset types will become mandatory"
208 " after v28. Specify '*' to match the current behavior of querying all dataset types.",
209 FutureWarning,
210 stacklevel=2,
211 )
212 glob = ["*"]
214 searches_multiple_dataset_types = len(glob) > 1 or has_globs(glob)
215 if order_by and searches_multiple_dataset_types:
216 raise NotImplementedError("--order-by is only supported for queries with a single dataset type.")
218 # show_uri requires a datastore.
219 without_datastore = not show_uri
220 self.butler = butler or Butler.from_config(repo, without_datastore=without_datastore)
221 self.showUri = show_uri
222 self._dataset_type_glob = glob
223 self._collections_wildcard = collections
224 self._where = where
225 self._find_first = find_first
226 self._limit = limit
227 self._order_by = order_by
228 self._searches_multiple_dataset_types = searches_multiple_dataset_types
229 self._with_dimension_records = with_dimension_records
231 def getTables(self) -> Iterator[AstropyTable]:
232 """Get the datasets as a list of astropy tables.
234 Yields
235 ------
236 datasetTables : `collections.abc.Iterator` [``astropy.table._Table``]
237 Astropy tables, one for each dataset type.
238 """
239 # Sort if we haven't been told to enforce an order.
240 sort_table = not bool(self._order_by)
242 if not self.showUri:
243 for refs in self.getDatasets():
244 table = _Table()
245 for ref in refs:
246 table.add(ref)
247 if refs:
248 yield table.getAstropyTable(refs[0].datasetType.name, sort=sort_table)
249 else:
250 for refs in self.getDatasets():
251 if not refs:
252 continue
253 # For URIs of disassembled composites we create a table per
254 # component.
255 tables: dict[str, _Table] = defaultdict(_Table)
256 dataset_type_name = refs[0].datasetType.name
257 ref_uris = self.butler.get_many_uris(refs, predict=True)
258 for ref, uris in ref_uris.items():
259 if uris.primaryURI:
260 tables[dataset_type_name].add(ref, uris.primaryURI)
261 for name, uri in uris.componentURIs.items():
262 tables[ref.datasetType.componentTypeName(name)].add(ref, uri)
263 for name in sorted(tables):
264 yield tables[name].getAstropyTable(name, sort=sort_table)
265 return
267 def getDatasets(self) -> Iterator[list[DatasetRef]]:
268 """Get the datasets as a list of lists.
270 Yields
271 ------
272 refs : `collections.abc.Iterator` [ `list [ `DatasetRef` ] ]
273 Dataset references matching the given query criteria grouped
274 by dataset type.
275 """
276 query_collections = self._collections_wildcard or ["*"]
278 warn_limit = False
279 limit_reached = False
280 if self._limit < 0:
281 # Negative limit means we should warn if the limit is exceeded.
282 warn_limit = True
283 limit = abs(self._limit) + 1 # +1 to tell us we hit the limit.
284 elif self._limit == 0:
285 # 0 means 'unlimited' in the CLI.
286 limit = None
287 else:
288 limit = self._limit
290 # We want to allow --order-by, but the query backend only supports
291 # it when there is a single dataset type.
292 # So we have to select different backends for single vs multiple
293 # dataset type queries.
294 if self._searches_multiple_dataset_types:
295 query_func = self._query_multiple_dataset_types
296 else:
297 query_func = self._query_single_dataset_type
299 try:
300 with query_func(query_collections, limit) as ref_iter:
301 datasets_found = 0
302 for dataset_type, chunk in groupby(ref_iter, lambda ref: ref.datasetType.name):
303 refs = list(chunk)
304 datasets_found += len(refs)
305 if warn_limit and limit is not None and datasets_found >= limit:
306 # We asked for one too many so must remove that from
307 # the list.
308 refs = refs[0:-1]
309 limit_reached = True
311 yield refs
313 _LOG.debug("Got %d results for dataset type %s", len(refs), dataset_type)
314 except MissingDatasetTypeError as e:
315 _LOG.info(str(e))
316 return
318 if limit is not None and limit_reached:
319 _LOG.warning(
320 "Requested limit of %d hit for number of datasets returned. "
321 "Use --limit to increase this limit.",
322 limit - 1,
323 )
325 @contextmanager
326 def _query_multiple_dataset_types(
327 self, collections: list[str], limit: int | None
328 ) -> Iterator[Iterator[DatasetRef]]:
329 args = QueryAllDatasetsParameters(
330 collections=collections,
331 find_first=self._find_first,
332 name=self._dataset_type_glob,
333 where=self._where,
334 limit=limit,
335 data_id={},
336 bind={},
337 with_dimension_records=self._with_dimension_records,
338 )
339 with self.butler._query_all_datasets_by_page(args) as iterator:
340 yield chain.from_iterable(iterator)
342 @contextmanager
343 def _query_single_dataset_type(
344 self, collections: list[str], limit: int | None
345 ) -> Iterator[Iterator[DatasetRef]]:
346 assert len(self._dataset_type_glob) == 1
347 dataset_type = self._dataset_type_glob[0]
349 refs = self.butler.query_datasets(
350 dataset_type,
351 collections=collections,
352 find_first=self._find_first,
353 where=self._where,
354 limit=limit,
355 order_by=self._order_by,
356 with_dimension_records=self._with_dimension_records,
357 )
359 yield iter(refs)