Coverage for python/lsst/daf/butler/registry/queries/_query_backend.py: 37%
83 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-06 01:32 -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
29__all__ = ("QueryBackend",)
31from abc import abstractmethod
32from collections.abc import Iterable, Mapping, Sequence, Set
33from contextlib import AbstractContextManager
34from typing import TYPE_CHECKING, Any, Generic, TypeVar
36from lsst.daf.relation import (
37 BinaryOperationRelation,
38 ColumnTag,
39 LeafRelation,
40 MarkerRelation,
41 Relation,
42 UnaryOperationRelation,
43)
45from ..._collection_type import CollectionType
46from ..._column_tags import DatasetColumnTag, DimensionKeyColumnTag
47from ..._dataset_type import DatasetType
48from ..._exceptions import MissingDatasetTypeError
49from ..._exceptions_legacy import DatasetTypeError
50from ...dimensions import DimensionGroup, DimensionRecordSet, DimensionUniverse
51from ..wildcards import CollectionWildcard
52from ._query_context import QueryContext
54if TYPE_CHECKING:
55 from ..interfaces import CollectionRecord
58_C = TypeVar("_C", bound=QueryContext)
61class QueryBackend(Generic[_C]):
62 """An interface for constructing and evaluating the
63 `~lsst.daf.relation.Relation` objects that comprise registry queries.
65 This ABC is expected to have a concrete subclass for each concrete registry
66 type, and most subclasses will be paired with a `QueryContext` subclass.
67 See `QueryContext` for the division of responsibilities between these two
68 interfaces.
69 """
71 @property
72 @abstractmethod
73 def universe(self) -> DimensionUniverse:
74 """Definition of all dimensions and dimension elements for this
75 registry (`DimensionUniverse`).
76 """
77 raise NotImplementedError()
79 @abstractmethod
80 def caching_context(self) -> AbstractContextManager[None]:
81 """Enable caching of collection records and summaries for the duration
82 of the returned context manager.
83 """
84 raise NotImplementedError()
86 def context(self) -> _C:
87 """Return a context manager that can be used to execute queries with
88 this backend.
90 Returns
91 -------
92 context : `QueryContext`
93 Context manager that manages state and connections needed to
94 execute queries.
95 """
96 raise NotImplementedError()
98 @abstractmethod
99 def get_collection_name(self, key: Any) -> str:
100 """Return the collection name associated with a collection primary key
101 value.
103 Parameters
104 ----------
105 key : `~typing.Any`
106 Collection primary key value.
108 Returns
109 -------
110 name : `str`
111 Collection name.
112 """
113 raise NotImplementedError()
115 @abstractmethod
116 def resolve_collection_wildcard(
117 self,
118 expression: Any,
119 *,
120 collection_types: Set[CollectionType] = CollectionType.all(),
121 flatten_chains: bool = True,
122 include_chains: bool | None = None,
123 ) -> list[CollectionRecord]:
124 """Return the collection records that match a wildcard expression.
126 Parameters
127 ----------
128 expression : `~typing.Any`
129 Names and/or patterns for collections; will be passed to
130 `CollectionWildcard.from_expression`.
131 collection_types : `collections.abc.Set` [ `CollectionType` ], optional
132 If provided, only yield collections of these types.
133 flatten_chains : `bool`, optional
134 If `True` (default) recursively yield the child collections of
135 `~CollectionType.CHAINED` collections.
136 include_chains : `bool`, optional
137 If `False`, return records for `~CollectionType.CHAINED`
138 collections themselves. The default is the opposite of
139 ``flattenChains``: either return records for CHAINED collections or
140 their children, but not both.
142 Returns
143 -------
144 records : `list` [ `CollectionRecord` ]
145 Matching collection records.
146 """
147 raise NotImplementedError()
149 @abstractmethod
150 def resolve_dataset_type_wildcard(
151 self,
152 expression: Any,
153 missing: list[str] | None = None,
154 explicit_only: bool = False,
155 ) -> list[DatasetType]:
156 """Return the dataset types that match a wildcard expression.
158 Parameters
159 ----------
160 expression : `~typing.Any`
161 Names and/or patterns for dataset types; will be passed to
162 `DatasetTypeWildcard.from_expression`.
163 missing : `list` of `str`, optional
164 String dataset type names that were explicitly given (i.e. not
165 regular expression patterns) but not found will be appended to this
166 list, if it is provided.
167 explicit_only : `bool`, optional
168 If `True`, require explicit `DatasetType` instances or `str` names,
169 with `re.Pattern` instances deprecated and ``...`` prohibited.
171 Returns
172 -------
173 dataset_types : `list` [ `DatasetType` ]
174 A list of resolved dataset types.
175 """
176 raise NotImplementedError()
178 def resolve_single_dataset_type_wildcard(
179 self,
180 expression: Any,
181 explicit_only: bool = False,
182 ) -> DatasetType:
183 """Return a single dataset type that matches a wildcard expression.
185 Parameters
186 ----------
187 expression : `~typing.Any`
188 Names and/or patterns for the dataset type; will be passed to
189 `DatasetTypeWildcard.from_expression`.
190 explicit_only : `bool`, optional
191 If `True`, require explicit `DatasetType` instances or `str` names,
192 with `re.Pattern` instances deprecated and ``...`` prohibited.
194 Returns
195 -------
196 single : `DatasetType`
197 The matched dataset type.
198 """
199 missing: list[str] = []
200 matching = self.resolve_dataset_type_wildcard(
201 expression, missing=missing, explicit_only=explicit_only
202 )
203 if not matching:
204 if missing:
205 raise MissingDatasetTypeError(
206 "\n".join(
207 f"Dataset type {t!r} is not registered, so no instances of it can exist."
208 for t in missing
209 )
210 )
211 else:
212 raise MissingDatasetTypeError(
213 f"No registered dataset types matched expression {expression!r}, "
214 "so no datasets will be found."
215 )
216 if len(matching) > 1:
217 raise DatasetTypeError(
218 f"Expression {expression!r} matched multiple parent dataset types: "
219 f"{[t.name for t in matching]}, but only one is allowed."
220 )
221 (single_parent,) = matching
222 if missing:
223 raise DatasetTypeError(
224 f"Expression {expression!r} appears to involve multiple dataset types, even though only "
225 f"one ({single_parent.name}) is registered, and only one is allowed here."
226 )
227 return single_parent
229 @abstractmethod
230 def filter_dataset_collections(
231 self,
232 dataset_types: Iterable[DatasetType],
233 collections: Sequence[CollectionRecord],
234 *,
235 governor_constraints: Mapping[str, Set[str]],
236 rejections: list[str] | None = None,
237 ) -> dict[DatasetType, list[CollectionRecord]]:
238 """Filter a sequence of collections to those for which a dataset query
239 might succeed.
241 Parameters
242 ----------
243 dataset_types : `~collections.abc.Iterable` [ `DatasetType` ]
244 Dataset types that are being queried. Must include only parent
245 or standalone dataset types, not components.
246 collections : `~collections.abc.Sequence` [ `CollectionRecord` ]
247 Sequence of collections that will be searched.
248 governor_constraints : `~collections.abc.Mapping` [ `str`, \
249 `~collections.abc.Set` [ `str` ] ], optional
250 Constraints imposed by other aspects of the query on governor
251 dimensions; collections inconsistent with these constraints will be
252 skipped.
253 rejections : `list` [ `str` ], optional
254 If not `None`, a `list` that diagnostic messages will be appended
255 to, for any collection that matches ``collections`` that is not
256 returned. At least one message is guaranteed whenever the result
257 is empty.
259 Returns
260 -------
261 dataset_collections : `dict` [ `DatasetType`, \
262 `list` [ `CollectionRecord` ] ]
263 The collections to search for each dataset. The dictionary's keys
264 are always exactly ``dataset_types`` (in the same order), and each
265 nested `list` of collections is ordered consistently with the
266 given ``collections``.
268 Notes
269 -----
270 This method accepts multiple dataset types and multiple collections at
271 once to enable implementations to batch up the fetching of summary
272 information needed to relate them.
273 """
274 raise NotImplementedError()
276 def resolve_dataset_collections(
277 self,
278 dataset_type: DatasetType,
279 collections: CollectionWildcard,
280 *,
281 governor_constraints: Mapping[str, Set[str]],
282 rejections: list[str] | None = None,
283 collection_types: Set[CollectionType] = CollectionType.all(),
284 allow_calibration_collections: bool = False,
285 ) -> list[CollectionRecord]:
286 """Resolve the sequence of collections to query for a dataset type.
288 Parameters
289 ----------
290 dataset_type : `DatasetType`
291 Dataset type to be queried in the returned collections.
292 collections : `CollectionWildcard`
293 Expression for the collections to be queried.
294 governor_constraints : `~collections.abc.Mapping` [ `str`, \
295 `~collections.abc.Set` ], optional
296 Constraints imposed by other aspects of the query on governor
297 dimensions; collections inconsistent with these constraints will be
298 skipped.
299 rejections : `list` [ `str` ], optional
300 If not `None`, a `list` that diagnostic messages will be appended
301 to, for any collection that matches ``collections`` that is not
302 returned. At least one message is guaranteed whenever the result
303 is empty.
304 collection_types : `~collections.abc.Set` [ `CollectionType` ], \
305 optional
306 Collection types to consider when resolving the collection
307 expression.
308 allow_calibration_collections : `bool`, optional
309 If `False`, skip (with a ``rejections`` message) any calibration
310 collections that match ``collections`` are not given explicitly by
311 name, and raise `NotImplementedError` for any calibration
312 collection that is given explicitly. This is a temporary option
313 that will be removed when the query system can handle temporal
314 joins involving calibration collections.
316 Returns
317 -------
318 records : `list` [ `CollectionRecord` ]
319 A new list of `CollectionRecord` instances, for collections that
320 both match ``collections`` and may have datasets of the given type.
322 Notes
323 -----
324 This is a higher-level driver for `resolve_collection_wildcard` and
325 `filter_dataset_collections` that is mostly concerned with handling
326 queries against `~Collection.Type.CALIBRATION` collections that aren't
327 fully supported yet. Once that support improves, this method may be
328 removed.
329 """
330 if collections == CollectionWildcard() and collection_types == CollectionType.all():
331 collection_types = {CollectionType.RUN}
332 explicit_collections = frozenset(collections.strings)
333 matching_collection_records = self.resolve_collection_wildcard(
334 collections, collection_types=collection_types
335 )
336 ((_, filtered_collection_records),) = self.filter_dataset_collections(
337 [dataset_type],
338 matching_collection_records,
339 governor_constraints=governor_constraints,
340 rejections=rejections,
341 ).items()
342 if not allow_calibration_collections:
343 supported_collection_records: list[CollectionRecord] = []
344 for record in filtered_collection_records:
345 if record.type is CollectionType.CALIBRATION:
346 # If collection name was provided explicitly then raise,
347 # since this is a kind of query we don't support yet;
348 # otherwise collection is a part of a chained one or regex
349 # match, and we skip it to not break queries of other
350 # included collections.
351 if record.name in explicit_collections:
352 raise NotImplementedError(
353 f"Query for dataset type {dataset_type.name!r} in CALIBRATION-type "
354 f"collection {record.name!r} is not yet supported."
355 )
356 else:
357 if rejections is not None:
358 rejections.append(
359 f"Not searching for dataset {dataset_type.name!r} in CALIBRATION "
360 f"collection {record.name!r} because calibration queries aren't fully "
361 "implemented; this is not an error only because the query structure "
362 "implies that searching this collection may be incidental."
363 )
364 supported_collection_records.append(record)
365 else:
366 supported_collection_records.append(record)
367 else:
368 supported_collection_records = filtered_collection_records
369 if not supported_collection_records and rejections is not None and not rejections:
370 rejections.append(f"No collections to search matching expression {collections!r}.")
371 return supported_collection_records
373 def make_doomed_dataset_relation(
374 self,
375 dataset_type: DatasetType,
376 columns: Set[str],
377 messages: Iterable[str],
378 context: _C,
379 ) -> Relation:
380 """Construct a relation that represents a doomed query for datasets.
382 Parameters
383 ----------
384 dataset_type : `DatasetType`
385 Dataset type being queried.
386 columns : `~collections.abc.Set` [ `str` ]
387 Dataset columns to include (dimension key columns are always
388 included). See `make_dataset_query_relation` for allowed values.
389 messages : `~collections.abc.Iterable` [ `str` ]
390 Diagnostic messages that explain why the query is doomed to yield
391 no rows.
392 context : `QueryContext`
393 Context that manages per-query state.
395 Returns
396 -------
397 relation : `lsst.daf.relation.Relation`
398 Relation with the requested columns and no rows.
399 """
400 column_tags: set[ColumnTag] = set(DimensionKeyColumnTag.generate(dataset_type.dimensions.required))
401 column_tags.update(DatasetColumnTag.generate(dataset_type.name, columns))
402 return context.preferred_engine.make_doomed_relation(columns=column_tags, messages=list(messages))
404 @abstractmethod
405 def resolve_governor_constraints(
406 self, dimensions: DimensionGroup, constraints: Mapping[str, Set[str]]
407 ) -> Mapping[str, Set[str]]:
408 """Resolve governor dimension constraints provided by user input to
409 a query against the content in the `Registry`.
411 Parameters
412 ----------
413 dimensions : `DimensionGroup`
414 Dimensions that bound the governor dimensions to consider (via
415 ``dimensions.governors``, more specifically).
416 constraints : `~collections.abc.Mapping` [ `str`, \
417 `~collections.abc.Set` [ `str` ] ]
418 Constraints from user input to the query (e.g. from data IDs and
419 string expression predicates).
421 Returns
422 -------
423 resolved : `~collections.abc.Mapping` [ `str`, \
424 `~collections.abc.Set` [ `str` ] ]
425 A shallow copy of ``constraints`` with keys equal to
426 ``dimensions.governors.names`` and value sets constrained by the
427 Registry content if they were not already in ``constraints``.
429 Raises
430 ------
431 DataIdValueError
432 Raised if ``constraints`` includes governor dimension values that
433 are not present in the `Registry`.
434 """
435 raise NotImplementedError()
437 @abstractmethod
438 def get_dimension_record_cache(self, element_name: str) -> DimensionRecordSet | None:
439 """Return a local cache of all `DimensionRecord` objects for a
440 dimension element, fetching it if necessary.
442 Parameters
443 ----------
444 element_name : `str`
445 Name of the dimension element.
447 Returns
448 -------
449 cache : `~collections.abc.Mapping` [ `DataCoordinate`, \
450 `DimensionRecord` ] or `None`
451 Mapping from data ID to dimension record, or `None` if this
452 element's records are never cached.
453 """
454 raise NotImplementedError()
456 def extract_dimension_relationships(self, relation: Relation) -> set[frozenset[str]]:
457 """Extract the dimension key relationships encoded in a relation tree.
459 Parameters
460 ----------
461 relation : `Relation`
462 Relation tree to process.
464 Returns
465 -------
466 relationships : `set` [ `frozenset` [ `str` ] ]
467 Set of sets of dimension names, where each inner set represents a
468 relationship between dimensions.
470 Notes
471 -----
472 Dimension relationships include both many-to-one implied dependencies
473 and many-to-many joins backed by "always-join" dimension elements, and
474 it's important to join in the dimension table that defines a
475 relationship in any query involving dimensions that are a superset of
476 that relationship. For example, let's consider a relation tree that
477 joins dataset existence-check relations for two dataset types, with
478 dimensions ``{instrument, exposure, detector}`` and ``{instrument,
479 physical_filter}``. The joined relation appears to have all dimension
480 keys in its expanded graph present except ``band``, and the system
481 could easily correct this by joining that dimension in directly. But
482 it's also missing the ``{instrument, exposure, physical_filter}``
483 relationship we'd get from the ``exposure`` dimension's own relation
484 (``exposure`` implies ``physical_filter``) and the similar
485 ``{instrument, physical_filter, band}`` relationship from the
486 ``physical_filter`` dimension relation; we need the relationship logic
487 to recognize that those dimensions need to be joined in as well in
488 order for the full relation to have rows that represent valid data IDs.
490 The implementation of this method relies on the assumption that
491 `LeafRelation` objects always have rows that are consistent with all
492 defined relationships (i.e. are valid data IDs). This is true for not
493 just dimension relations themselves, but anything created from queries
494 based on them, including datasets and query results. It is possible to
495 construct `LeafRelation` objects that don't satisfy this criteria (e.g.
496 when accepting in user-provided data IDs), and in this case
497 higher-level guards or warnings must be provided.``
498 """
499 return {
500 frozenset(
501 tag.dimension
502 for tag in DimensionKeyColumnTag.filter_from(leaf_relation.columns & relation.columns)
503 )
504 for leaf_relation in self._extract_leaf_relations(relation).values()
505 }
507 def _extract_leaf_relations(self, relation: Relation) -> dict[str, LeafRelation]:
508 """Recursively extract leaf relations from a relation tree.
510 Parameters
511 ----------
512 relation : `Relation`
513 Tree to process.
515 Returns
516 -------
517 leaves : `dict` [ `str`, `LeafRelation` ]
518 Leaf relations, keyed and deduplicated by name.
519 """
520 match relation:
521 case LeafRelation() as leaf:
522 return {leaf.name: leaf}
523 case UnaryOperationRelation(target=target):
524 return self._extract_leaf_relations(target)
525 case BinaryOperationRelation(lhs=lhs, rhs=rhs):
526 return self._extract_leaf_relations(lhs) | self._extract_leaf_relations(rhs)
527 case MarkerRelation(target=target):
528 return self._extract_leaf_relations(target)
529 raise AssertionError("Match should be exhaustive and all branches should return.")