Coverage for python/lsst/daf/butler/queries/_query.py: 15%
208 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-10 01:29 -0700
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/>.
28from __future__ import annotations
30__all__ = ("Query",)
32from collections.abc import Iterable, Mapping, Set
33from types import EllipsisType
34from typing import Any, final
36import astropy.table
38from lsst.utils.iteration import ensure_iterable
40from .._dataset_type import DatasetType
41from .._exceptions import DimensionNameError, InvalidQueryError
42from .._storage_class import StorageClassFactory
43from ..dimensions import DataCoordinate, DataId, DataIdValue, DimensionGroup
44from ..registry import DatasetTypeError
45from ._base import QueryBase
46from ._data_coordinate_query_results import DataCoordinateQueryResults
47from ._dataset_query_results import DatasetRefQueryResults
48from ._dimension_record_query_results import DimensionRecordQueryResults
49from ._general_query_results import GeneralQueryResults
50from ._identifiers import IdentifierContext, interpret_identifier
51from .convert_args import convert_where_args
52from .driver import QueryDriver
53from .expression_factory import ExpressionFactory
54from .predicate_constraints_summary import PredicateConstraintsSummary
55from .result_specs import (
56 DataCoordinateResultSpec,
57 DatasetRefResultSpec,
58 DimensionRecordResultSpec,
59 GeneralResultSpec,
60)
61from .tree import (
62 ANY_DATASET,
63 DatasetFieldName,
64 DatasetFieldReference,
65 DatasetSearch,
66 DimensionFieldReference,
67 DimensionKeyReference,
68 Predicate,
69 QueryTree,
70 make_identity_query_tree,
71)
74@final
75class Query(QueryBase):
76 """A method-chaining builder for butler queries.
78 Parameters
79 ----------
80 driver : `~lsst.daf.butler.queries.driver.QueryDriver`
81 Implementation object that knows how to actually execute queries.
82 tree : `~lsst.daf.butler.queries.tree.QueryTree`, optional
83 Description of the query as a tree of joins and column expressions.
84 Defaults to the result of a call to
85 `~lsst.daf.butler.queries.tree.make_identity_query_tree`.
87 Notes
88 -----
89 `Query` objects should never be constructed directly by users; use
90 `Butler.query <lsst.daf.butler.Butler.query>` instead.
92 A `Query` object represents the first stage of query construction, in which
93 constraints and joins are defined (roughly corresponding to the WHERE and
94 FROM clauses in SQL). The various "results" objects represent the second
95 (and final) stage, where the columns returned are specified and any sorting
96 or integer slicing can be applied. Result objects are obtained from the
97 `data_ids`, `datasets`, and `dimension_records` methods.
99 `Query` and query-result objects are always immutable (except for caching
100 information fetched from the database or server), so modifier methods
101 always return a new object without modifying the current one.
102 """
104 def __init__(self, driver: QueryDriver, tree: QueryTree | None = None):
105 # __init__ defined here because there are multiple base classes and
106 # not all define __init__ (and hence inherit object.__init__, which
107 # just ignores its args). Even if we just delegate to super(), it
108 # seems less fragile to make it explicit here.
109 if tree is None:
110 tree = make_identity_query_tree(driver.universe)
111 super().__init__(driver, tree)
113 # If ``_allow_duplicate_overlaps`` is set to `True` then query will be
114 # allowed to generate non-distinct rows for spatial overlaps. This is
115 # not a part of public API for now, to be used by graph builder as
116 # optimization.
117 self._allow_duplicate_overlaps: bool = False
119 @property
120 def constraint_dataset_types(self) -> Set[str]:
121 """The names of all dataset types joined into the query.
123 The existence of datasets of these types constrains the data IDs of any
124 type of result. Fields for these dataset types are also usable in
125 'where' expressions.
126 """
127 # Note that this includes only dataset type names, not `DatasetType`
128 # instances; the `DatasetQueryResults` adapter returned by the
129 # `datasets` method does include `DatasetType` instances, since it is
130 # in a better position to track and respect any storage class override
131 # specified.
132 return self._tree.datasets.keys()
134 @property
135 def constraint_dimensions(self) -> DimensionGroup:
136 """Dimensions currently present in the query, either directly or
137 indirectly.
139 This includes dimensions that are present in any joined subquery (such
140 as a dataset search, materialization, or data ID upload) or `where`
141 argument, as well as any required or implied dependency of those
142 dimensions.
143 """
144 return self._tree.dimensions
146 @property
147 def expression_factory(self) -> ExpressionFactory:
148 """A factory for column expressions using overloaded operators.
149 (`~lsst.daf.butler.queries.expression_factory.ExpressionFactory`).
151 Notes
152 -----
153 Typically this attribute will be assigned to a single-character local
154 variable, and then its (dynamic) attributes can be used to obtain
155 references to columns that can be included in a query::
157 with butler.query() as query:
158 x = query.expression_factory
159 query = query.where(
160 x.instrument == "LSSTCam",
161 x.visit.day_obs > 20240701,
162 x.any(x.band == "u", x.band == "y"),
163 )
165 As shown above, the returned object also has an
166 `~lsst.daf.butler.queries.expression_factory.ExpressionFactory.any`
167 method to create combine expressions with logical OR (as well as
168 `~lsst.daf.butler.queries.expression_factory.ExpressionFactory.not_`
169 and
170 `~lsst.daf.butler.queries.expression_factory.ExpressionFactory.all`,
171 though the latter is rarely necessary since `where` already combines
172 its arguments with AND).
174 Proxies for fields associated with individual datasets but not
175 dimension records (``dataset_id``, ``ingest_date``, ``run``,
176 ``collection``, as well as ``timespan`` for
177 `~lsst.daf.butler.CollectionType.CALIBRATION` collection searches) can
178 be obtained with dict-like access instead::
180 with butler.query() as query:
181 query = query.order_by(x["raw"].ingest_date)
183 Expression proxy objects that correspond to scalar columns overload the
184 standard comparison operators (``==``, ``!=``, ``<``, ``>``, ``<=``,
185 ``>=``) and provide
186 `~lsst.daf.butler.queries.expression_factory.ScalarExpressionProxy.in_range`,
187 `~lsst.daf.butler.queries.expression_factory.ScalarExpressionProxy.in_iterable`, and
188 `~lsst.daf.butler.queries.expression_factory.ScalarExpressionProxy.in_query`
189 methods for membership tests. For ``order_by`` contexts, they also have a
190 `~lsst.daf.butler.queries.expression_factory.ScalarExpressionProxy.desc`
191 property to indicate that the sort order for that expression should be
192 reversed.
194 Proxy objects for
195 `region <lsst.daf.butler.queries.expression_factory.RegionProxy>` and
196 `timespan <lsst.daf.butler.queries.expression_factory.TimespanProxy>`
197 fields have an ``overlaps`` method, and timespans also have
198 `~lsst.daf.butler.queries.expression_factory.TimespanProxy.begin` and
199 `~lsst.daf.butler.queries.expression_factory.TimespanProxy.end`
200 properties to access scalar expression proxies for the bounds.
202 All proxy objects also have an
203 `~lsst.daf.butler.queries.expression_factory.ExpressionProxy.is_null`
204 property.
206 Literal values can be created by calling
207 `ExpressionFactory.literal <lsst.daf.butler.queries.expression_factory.ExpressionFactory.literal>`,
208 but can almost always be created implicitly via overloaded operators
209 instead.
210 """ # noqa: W505, long docstrings
211 return ExpressionFactory(self._driver.universe)
213 def data_ids(
214 self, dimensions: DimensionGroup | Iterable[str] | str | None = None
215 ) -> DataCoordinateQueryResults:
216 """Return a result object that is a `~lsst.daf.butler.DataCoordinate`
217 iterable.
219 Parameters
220 ----------
221 dimensions : `~lsst.daf.butler.DimensionGroup`, `str`, or \
222 `~collections.abc.Iterable` [`str`], optional
223 The dimensions of the data IDs to yield, as either
224 `~lsst.daf.butler.DimensionGroup` instances or `str` names. Will
225 be automatically expanded to a complete
226 `~lsst.daf.butler.DimensionGroup`. These dimensions do not need to
227 match the query's current dimensions. Default is
228 `constraint_dimensions`.
230 Returns
231 -------
232 data_ids : `~lsst.daf.butler.queries.DataCoordinateQueryResults`
233 Data IDs matching the given query parameters. These are guaranteed
234 to identify all dimensions (``DataCoordinate.hasFull`` returns
235 `True`), but will not contain `~lsst.daf.butler.DimensionRecord`
236 objects (``DataCoordinate.hasRecords`` returns `False`). Call
237 `~DataCoordinateQueryResults.with_dimension_records` on the
238 returned object to include dimension records as well.
239 """
240 tree = self._tree
241 if dimensions is None:
242 dimensions = self._tree.dimensions
243 else:
244 dimensions = self._driver.universe.conform(dimensions)
245 if not dimensions <= self._tree.dimensions:
246 tree = tree.join_dimensions(dimensions)
247 result_spec = DataCoordinateResultSpec(
248 dimensions=dimensions,
249 include_dimension_records=False,
250 allow_duplicate_overlaps=self._allow_duplicate_overlaps,
251 )
252 return DataCoordinateQueryResults(self._driver, tree, result_spec)
254 def datasets(
255 self,
256 dataset_type: str | DatasetType,
257 collections: str | Iterable[str] | None = None,
258 *,
259 find_first: bool = True,
260 ) -> DatasetRefQueryResults:
261 """Return a result object that is a `~lsst.daf.butler.DatasetRef`
262 iterable.
264 Parameters
265 ----------
266 dataset_type : `str` or `~lsst.daf.butler.DatasetType`
267 The dataset type to search for.
268 collections : `str` or `~collections.abc.Iterable` [ `str` ], optional
269 The collection or collections to search, in order. If not provided
270 or `None`, and the dataset has not already been joined into the
271 query, the default collection search path for this butler is used.
272 find_first : `bool`, optional
273 If `True` (default), for each result data ID, only yield one
274 `~lsst.daf.butler.DatasetRef` of each
275 `~lsst.daf.butler.DatasetType`, from the first collection in
276 which a dataset of that dataset type appears (according to the
277 order of ``collections`` passed in). If `True`, ``collections``
278 must not be ``...``.
280 Returns
281 -------
282 refs : `lsst.daf.butler.queries.DatasetRefQueryResults`
283 Dataset references matching the given query criteria. Nested data
284 IDs are guaranteed to include values for all implied dimensions
285 (i.e. ``DataCoordinate.hasFull`` will return `True`), but will not
286 include dimension records (``DataCoordinate.hasRecords`` will be
287 `False`) unless
288 `~.queries.DatasetRefQueryResults.with_dimension_records` is
289 called on the result object (which returns a new one).
291 Raises
292 ------
293 lsst.daf.butler.registry.DatasetTypeExpressionError
294 Raised when the ``dataset_type`` expression is invalid.
295 lsst.daf.butler.registry.NoDefaultCollectionError
296 Raised when ``collections`` is `None` and default butler
297 collections are not defined.
298 TypeError
299 Raised when the arguments are incompatible, such as when a
300 collection wildcard is passed when ``find_first`` is `True`
301 """
302 dataset_type_name, storage_class_name, query = self._join_dataset_search_impl(
303 dataset_type, collections
304 )
305 dataset_search = query._tree.datasets[dataset_type_name]
306 spec = DatasetRefResultSpec.model_construct(
307 dataset_type_name=dataset_type_name,
308 dimensions=dataset_search.dimensions,
309 storage_class_name=storage_class_name,
310 include_dimension_records=False,
311 find_first=find_first,
312 allow_duplicate_overlaps=self._allow_duplicate_overlaps,
313 )
314 return DatasetRefQueryResults(self._driver, tree=query._tree, spec=spec)
316 def dimension_records(self, element: str) -> DimensionRecordQueryResults:
317 """Return a result object that is a `~lsst.daf.butler.DimensionRecord`
318 iterable.
320 Parameters
321 ----------
322 element : `str`
323 The name of a dimension element to obtain records for.
325 Returns
326 -------
327 records : `lsst.daf.butler.queries.DimensionRecordQueryResults`
328 Data IDs matching the given query parameters.
329 """
330 if element not in self._driver.universe:
331 # Prefer an explicit exception over a KeyError below.
332 raise DimensionNameError(
333 f"No such dimension '{element}', available dimensions: " + str(self._driver.universe.elements)
334 )
335 tree = self._tree
336 if element not in tree.dimensions.elements:
337 tree = tree.join_dimensions(self._driver.universe[element].minimal_group)
338 result_spec = DimensionRecordResultSpec(
339 element=self._driver.universe[element], allow_duplicate_overlaps=self._allow_duplicate_overlaps
340 )
341 return DimensionRecordQueryResults(self._driver, tree, result_spec)
343 def general(
344 self,
345 dimensions: DimensionGroup | Iterable[str],
346 *names: str,
347 dimension_fields: Mapping[str, Set[str]] | None = None,
348 dataset_fields: Mapping[str, Set[DatasetFieldName] | EllipsisType] | None = None,
349 find_first: bool | None = None,
350 ) -> GeneralQueryResults:
351 """Execute query returning general result.
353 **This is an experimental interface and may change at any time.**
355 Parameters
356 ----------
357 dimensions : `~lsst.daf.butler.DimensionGroup` or \
358 `~collections.abc.Iterable` [ `str` ]
359 The dimensions that span all fields returned by this query.
360 *names : `str`
361 Names of dimensions fields (in "dimension.field" format), dataset
362 fields (in "dataset_type.field" format) to include in this query.
363 dimension_fields : `~collections.abc.Mapping` [`str`, \
364 `~collections.abc.Set` [`str`]], optional
365 Dimension record fields included in this query, keyed by dimension
366 element name.
367 dataset_fields : `~collections.abc.Mapping` [`str`, \
368 `~collections.abc.Set` | ``...`` ], optional
369 Dataset fields included in this query, the key in the mapping is
370 dataset type name. Ellipsis (``...``) can be used for value
371 to include all dataset fields needed to extract
372 `~lsst.daf.butler.DatasetRef` instances later.
373 find_first : `bool`, optional
374 Whether this query requires find-first resolution for a dataset.
375 This is ignored and can be omitted if the query has no dataset
376 fields. It must be explicitly set to `False` if there are multiple
377 dataset types with fields, or if any dataset's ``collections``
378 or ``timespan`` fields are included in the results.
380 Returns
381 -------
382 result : `~lsst.daf.butler.queries.GeneralQueryResults`
383 Query result that can be iterated over.
385 Notes
386 -----
387 The dimensions of the returned query are automatically expanded to
388 include those associated with all dimension and dataset fields; the
389 ``dimensions`` argument is just the minimal dimensions to return.
390 """
391 if dimension_fields is None:
392 dimension_fields = {}
393 if dataset_fields is None:
394 dataset_fields = {}
395 # Processing fields from mapping args, processing the special `...`
396 # wildcard and dropping keys with empty values.
397 dataset_fields_dict: dict[str, set[DatasetFieldName]] = {}
398 for dataset_type_name, fields_for_dataset_type in dataset_fields.items():
399 if fields_for_dataset_type is ...:
400 new_fields_for_dataset_type: set[DatasetFieldName] = {
401 "run",
402 "dataset_id",
403 } # all we need for DatasetRefs.
404 else:
405 new_fields_for_dataset_type = set(fields_for_dataset_type)
406 if new_fields_for_dataset_type:
407 dataset_fields_dict[dataset_type_name] = new_fields_for_dataset_type
408 dimension_fields_dict = {
409 element_name: new_fields_for_element
410 for element_name, fields_for_element in dimension_fields.items()
411 if (new_fields_for_element := set(fields_for_element))
412 }
413 # Parse all names passed as positional arguments, and start to
414 # accumulate additional dimension names we'll need in the results.
415 dimensions = self._driver.universe.conform(dimensions)
416 context = IdentifierContext(dimensions, set(self._tree.datasets))
417 extra_dimension_names: set[str] = set()
418 for name in names:
419 identifier = interpret_identifier(context, name)
420 match identifier:
421 case DimensionKeyReference(dimension=dimension):
422 # Could be because someone asked for the key field.
423 extra_dimension_names.add(dimension.name)
424 case DimensionFieldReference(element=element, field=field):
425 dimension_fields_dict.setdefault(element.name, set()).add(field)
426 case DatasetFieldReference(dataset_type=dataset_type, field=dataset_field):
427 if dataset_type is ANY_DATASET:
428 raise InvalidQueryError("Dataset wildcard fields are not supported by Query.general.")
429 dataset_fields_dict.setdefault(dataset_type, set()).add(dataset_field)
430 case _:
431 raise TypeError(f"Unexpected type of identifier ({name}): {identifier}")
432 # Add more dimension names from the field mappings (including those
433 # we just populated from args). Also check that the dataset fields
434 # are consistent with find_first.
435 for element_name, fields in dimension_fields_dict.items():
436 extra_dimension_names.update(self._driver.universe[element_name].minimal_group.names)
437 for dataset_type_name, fields_for_dataset_type in dataset_fields_dict.items():
438 if "collections" in fields_for_dataset_type and find_first is not False:
439 raise InvalidQueryError(
440 f"find_first=False must be passed explicitly if {dataset_type_name}.collections "
441 "is included in query results."
442 )
443 if "timespan" in fields_for_dataset_type and find_first is not False:
444 raise InvalidQueryError(
445 f"find_first=False must be passed explicitly if {dataset_type_name}.timespan "
446 "is included in query results."
447 )
448 try:
449 dataset_search = self._tree.datasets[dataset_type_name]
450 except KeyError:
451 raise InvalidQueryError(
452 f"A search for dataset type {dataset_type_name!r} must be explicitly joined into the "
453 "query before including its fields in query results."
454 ) from None
455 extra_dimension_names.update(dataset_search.dimensions.names)
456 if find_first is None:
457 if dataset_fields_dict:
458 raise InvalidQueryError(
459 "find_first must be passed if dataset fields are included in query results."
460 )
461 else:
462 find_first = False
463 if find_first and len(dataset_fields_dict) != 1:
464 raise InvalidQueryError(
465 "find_first=True is not valid unless exactly one dataset type's fields are requested."
466 )
467 # Combine extra dimensions with the original ones.
468 dimensions = self._driver.universe.conform(dimensions.names | extra_dimension_names)
469 # Merge missing dimensions into the tree.
470 tree = self._tree
471 if not dimensions <= tree.dimensions:
472 tree = tree.join_dimensions(dimensions)
473 result_spec = GeneralResultSpec(
474 dimensions=dimensions,
475 dimension_fields=dimension_fields_dict,
476 dataset_fields=dataset_fields_dict,
477 find_first=find_first,
478 allow_duplicate_overlaps=self._allow_duplicate_overlaps,
479 )
480 return GeneralQueryResults(self._driver, tree=tree, spec=result_spec)
482 def materialize(
483 self,
484 *,
485 dimensions: Iterable[str] | DimensionGroup | None = None,
486 datasets: Iterable[str] | None = None,
487 ) -> Query:
488 """Execute the query, save its results to a temporary location, and
489 return a new query that represents fetching or joining against those
490 saved results.
492 Parameters
493 ----------
494 dimensions : `~collections.abc.Iterable` [ `str` ] or \
495 `~lsst.daf.butler.DimensionGroup`, optional
496 Dimensions to include in the temporary results. Default is to
497 include all dimensions in the query.
498 datasets : `~collections.abc.Iterable` [ `str` ], optional
499 Names of dataset types that should be included in the new query;
500 default is to include `constraint_dataset_types`.
502 Returns
503 -------
504 query : `Query`
505 A new query object whose that represents the materialized rows.
507 Notes
508 -----
509 Only dimension key columns and (at the discretion of the
510 implementation) certain dataset columns are actually materialized,
511 since at this stage we do not know which dataset or dimension record
512 fields are actually needed in result rows, and these can be joined back
513 in on the materialized dimension keys. But all constraints on those
514 dimension keys (including dataset existence) are applied to the
515 materialized rows.
516 """
517 if datasets is None:
518 datasets = frozenset(self.constraint_dataset_types)
519 else:
520 datasets = frozenset(datasets)
521 if not (datasets <= self.constraint_dataset_types):
522 raise InvalidQueryError(
523 f"Dataset(s) {datasets - self.constraint_dataset_types} are present in the query."
524 )
525 if dimensions is None:
526 dimensions = self._tree.dimensions
527 else:
528 dimensions = self._driver.universe.conform(dimensions)
529 key = self._driver.materialize(
530 self._tree, dimensions, datasets, allow_duplicate_overlaps=self._allow_duplicate_overlaps
531 )
532 tree = make_identity_query_tree(self._driver.universe).join_materialization(
533 key, dimensions=dimensions
534 )
535 for dataset_type_name in datasets:
536 dataset_search = self._tree.datasets[dataset_type_name]
537 if not (dataset_search.dimensions <= tree.dimensions):
538 raise InvalidQueryError(
539 f"Materialization-backed query has dimensions {tree.dimensions}, which do not "
540 f"cover the dimensions {dataset_search.dimensions} of dataset {dataset_type_name!r}. "
541 "Expand the dimensions or drop this dataset type in the arguments to materialize to "
542 "avoid this error."
543 )
544 tree = tree.join_dataset(dataset_type_name, dataset_search)
545 return Query(self._driver, tree)
547 def join_dataset_search(
548 self,
549 dataset_type: str | DatasetType,
550 collections: Iterable[str] | None = None,
551 ) -> Query:
552 """Return a new query with a search for a dataset joined in.
554 Parameters
555 ----------
556 dataset_type : `str` or `~lsst.daf.butler.DatasetType`
557 Dataset type or name. May not refer to a dataset component.
558 collections : `~collections.abc.Iterable` [ `str` ], optional
559 Iterable of collections to search. Order is preserved, but will
560 not matter if the dataset search is only used as a constraint on
561 dimensions or if ``find_first=False`` when requesting results. If
562 not present or `None`, the default collection search path will be
563 used.
565 Returns
566 -------
567 query : `Query`
568 A new query object with dataset columns available and rows
569 restricted to those consistent with the found data IDs.
571 Raises
572 ------
573 DatasetTypeError
574 Raised if given dataset type is inconsistent with the registered
575 dataset type.
576 MissingDatasetTypeError
577 Raised if the dataset type has not been registered and only a
578 `str` dataset type name was given.
580 Notes
581 -----
582 This method may require communication with the server unless the
583 dataset type and collections have already been referenced by the same
584 query context.
585 """
586 _, _, query = self._join_dataset_search_impl(
587 dataset_type, collections, allow_storage_class_overrides=False
588 )
589 return query
591 def join_data_coordinates(self, iterable: Iterable[DataCoordinate]) -> Query:
592 """Return a new query that joins in an explicit iterable of data IDs.
594 Parameters
595 ----------
596 iterable : `~collections.abc.Iterable` \
597 [`~lsst.daf.butler.DataCoordinate`]
598 Iterable of `~lsst.daf.butler.DataCoordinate`. All items must have
599 the same dimensions. Must have at least one item.
601 Returns
602 -------
603 query : `Query`
604 A new query object with the data IDs joined in.
605 """
606 rows: set[tuple[DataIdValue, ...]] = set()
607 dimensions: DimensionGroup | None = None
608 for data_coordinate in iterable:
609 if dimensions is None:
610 dimensions = data_coordinate.dimensions
611 elif dimensions != data_coordinate.dimensions:
612 raise InvalidQueryError(
613 f"Inconsistent dimensions: {dimensions} != {data_coordinate.dimensions}."
614 )
615 rows.add(data_coordinate.required_values)
616 if dimensions is None:
617 raise InvalidQueryError("Cannot upload an empty data coordinate set.")
618 key = self._driver.upload_data_coordinates(dimensions, rows)
619 return Query(
620 tree=self._tree.join_data_coordinate_upload(dimensions=dimensions, key=key),
621 driver=self._driver,
622 )
624 def join_data_coordinate_table(self, table: astropy.table.Table) -> Query:
625 """Return a new query that joins in an explicit table of data IDs.
627 Parameters
628 ----------
629 table : `astropy.table.Table`
630 A table of data IDs to join. Columns must be dimension names, and
631 columns for dimensions whose values that are implied by others are
632 ignored. If there is no column for a required dimension is missing
633 but is fully constrained to a literal by a previous `where` call, a
634 constant-valued column will be added.
636 Returns
637 -------
638 query : `Query`
639 A new query object with the data IDs joined in.
640 """
641 if not len(table):
642 raise InvalidQueryError("Cannot upload an empty data coordinate set.")
643 column_names = set(table.colnames)
644 dimensions = self._driver.universe.conform(column_names)
645 # To avoid numpy scalar types that will upset SQLAlchemy, we turn the
646 # columns we care about into lists of regular Python scalars. We do
647 # this in dimensions.required order so we can zip the values of this
648 # dict later to make data ID 'required_values' tuples.
649 column_lists = {d: table[d].data.tolist() if d in column_names else None for d in dimensions.required}
650 if not column_names.issuperset(dimensions.required):
651 # If columns are missing, see if they're fixed by a previous
652 # `where` call or equivalent.
653 predicate_summary = PredicateConstraintsSummary(self._tree.predicate)
654 missing = dimensions.required - column_names
655 provided_by_predicate = predicate_summary.constraint_data_id.keys() & missing
656 missing -= provided_by_predicate
657 if missing:
658 raise InvalidQueryError(f"Data coordinate table is missing required dimension(s) {missing}.")
659 if provided_by_predicate:
660 for k in provided_by_predicate:
661 column_lists[k] = [predicate_summary.constraint_data_id[k]] * len(table)
662 key = self._driver.upload_data_coordinates(dimensions, zip(*column_lists.values(), strict=True))
663 return Query(
664 tree=self._tree.join_data_coordinate_upload(dimensions=dimensions, key=key),
665 driver=self._driver,
666 )
668 def join_dimensions(self, dimensions: Iterable[str] | DimensionGroup) -> Query:
669 """Return a new query that joins the logical tables for additional
670 dimensions.
672 Parameters
673 ----------
674 dimensions : `~collections.abc.Iterable` [ `str` ] or \
675 `~lsst.daf.butler.DimensionGroup`
676 Names of dimensions to join in.
678 Returns
679 -------
680 query : `Query`
681 A new query object with the dimensions joined in.
683 Notes
684 -----
685 Dimensions are automatically joined in whenever needed, so this method
686 should rarely need to be called directly.
687 """
688 dimensions = self._driver.universe.conform(dimensions)
689 return Query(tree=self._tree.join_dimensions(dimensions), driver=self._driver)
691 def where(
692 self,
693 *args: str | Predicate | DataId,
694 bind: Mapping[str, Any] | None = None,
695 **kwargs: Any,
696 ) -> Query:
697 """Return a query with a boolean-expression filter on its rows.
699 Parameters
700 ----------
701 *args
702 Constraints to apply, combined with logical AND. Arguments may be
703 `str` expressions to parse,
704 `~lsst.daf.butler.queries.tree.Predicate` objects (these are
705 typically constructed via `expression_factory`) or data IDs.
706 bind : `~collections.abc.Mapping`
707 Mapping from string identifier appearing in a string expression to
708 a literal value that should be substituted for it. This is
709 recommended instead of embedding literals directly into the
710 expression, especially for strings, timespans, or other types where
711 quoting or formatting is nontrivial.
712 **kwargs
713 Data ID key value pairs that extend and override any present in
714 ``*args``.
716 Returns
717 -------
718 query : `Query`
719 A new query object with the given row filters (as well as any
720 already present in ``self``). All row filters are combined with
721 logical AND.
723 Notes
724 -----
725 Expressions referring to dimensions or dimension elements are resolved
726 automatically. References to dataset fields (see `expression_factory`
727 for the distinction) cannot be resolved by default; they must either be
728 preceded by a call to `join_dataset_search` or must be passed to
729 `DatasetRefQueryResults.where <lsst.daf.butler.queries.DatasetRefQueryResults.where>`
730 instead.
732 Data ID values are not checked for consistency; they are extracted from
733 ``args`` and then ``kwargs`` and combined, with later values overriding
734 earlier ones.
735 """ # noqa: W505, long docstrings
736 return Query(
737 tree=self._tree.where(
738 convert_where_args(
739 self.constraint_dimensions,
740 self.constraint_dataset_types,
741 *args,
742 bind=bind,
743 **kwargs,
744 )
745 ),
746 driver=self._driver,
747 )
749 def _skip_governor_validation(self) -> Query:
750 tree = self._tree.model_copy(update={"validateGovernorConstraints": False})
751 return Query(tree=tree, driver=self._driver)
753 def _join_dataset_search_impl(
754 self,
755 dataset_type: str | DatasetType,
756 collections: Iterable[str] | None = None,
757 allow_storage_class_overrides: bool = True,
758 ) -> tuple[str, str, Query]:
759 """Implement `join_dataset_search`, and also return the dataset type
760 name and storage class, in addition to the modified Query.
761 """
762 # In this method we need the dimensions of the dataset type, but we
763 # might not need the storage class, since the dataset may only be used
764 # as an existence constraint. It depends on whether
765 # `join_dataset_search` or `datasets` is calling this method.
766 dimensions: DimensionGroup | None = None
767 storage_class_name: str | None = None
768 # Handle DatasetType vs. str arg.
769 if isinstance(dataset_type, DatasetType):
770 dataset_type_name = dataset_type.name
771 dimensions = dataset_type.dimensions
772 storage_class_name = dataset_type.storageClass_name
773 elif isinstance(dataset_type, str):
774 dataset_type_name = dataset_type
775 else:
776 raise TypeError(f"Invalid dataset type argument {dataset_type!r}.")
777 # See if this dataset has already been joined into the query.
778 if existing_search := self._tree.datasets.get(dataset_type_name):
779 if collections is None:
780 collections = existing_search.collections
781 else:
782 collections = tuple(ensure_iterable(collections))
783 if collections != existing_search.collections:
784 raise InvalidQueryError(
785 f"Dataset type {dataset_type_name!r} was already joined into this "
786 "query with a different collection search path (previously "
787 f"[{', '.join(existing_search.collections)}], now [{', '.join(collections)}])."
788 )
789 if dimensions is None:
790 dimensions = existing_search.dimensions
791 else:
792 if collections is None:
793 collections = self._driver.get_default_collections()
794 collections = tuple(ensure_iterable(collections))
795 # Look up the data repository definition of the dataset type to check
796 # for consistency, or get dimensions and storage class if we don't have
797 # them.
798 resolved_dataset_type = self._driver.get_dataset_type(dataset_type_name)
799 resolved_dimensions = resolved_dataset_type.dimensions
800 if dimensions is not None and dimensions != resolved_dimensions:
801 raise DatasetTypeError(
802 f"Given dimensions {dimensions} for dataset type {dataset_type_name!r} do not match the "
803 f"registered dimensions {resolved_dimensions}."
804 )
805 if storage_class_name is not None:
806 if storage_class_name != resolved_dataset_type.storageClass_name:
807 if not allow_storage_class_overrides:
808 raise InvalidQueryError(
809 f"Storage class {storage_class_name!r} for dataset type {dataset_type!r} differs "
810 f"from repository definition {resolved_dataset_type.storageClass_name!r}, but "
811 "join_dataset_search does not are about storage classes and cannot record this "
812 "override. Pass the override to `Query.datasets` instead."
813 )
814 if not (
815 StorageClassFactory()
816 .getStorageClass(storage_class_name)
817 .can_convert(resolved_dataset_type.storageClass)
818 ):
819 raise DatasetTypeError(
820 f"Given storage class {storage_class_name!r} for {dataset_type_name!r} is not "
821 f"compatible with repository storage class {resolved_dataset_type.storageClass_name}."
822 )
823 else:
824 storage_class_name = resolved_dataset_type.storageClass_name
825 dataset_search = DatasetSearch.model_construct(
826 collections=collections,
827 dimensions=resolved_dimensions,
828 )
829 return (
830 dataset_type_name,
831 storage_class_name,
832 Query(self._driver, self._tree.join_dataset(dataset_type_name, dataset_search)),
833 )