Coverage for python/lsst/daf/butler/queries/_general_query_results.py: 26%
103 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:23 -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__ = ("GeneralQueryResults", "GeneralResultTuple")
32import itertools
33from collections.abc import Iterator
34from typing import Any, NamedTuple, final
36from .._dataset_ref import DatasetRef
37from .._dataset_type import DatasetType
38from ..dimensions import DataCoordinate, DimensionElement, DimensionGroup, DimensionRecord, DimensionRecordSet
39from ._base import QueryResultsBase
40from .driver import QueryDriver
41from .result_specs import GeneralResultSpec
42from .tree import QueryTree, ResultColumn
45class GeneralResultTuple(NamedTuple):
46 """Helper class for general result that represents the result row as a
47 data coordinate and optionally a set of dataset refs extracted from a row.
48 """
50 data_id: DataCoordinate
51 """Data coordinate for current row."""
53 refs: list[DatasetRef]
54 """Dataset refs extracted from the current row, the order matches the order
55 of arguments in ``iter_tuples`` call."""
57 raw_row: dict[str, Any]
58 """Original result row, the keys are the names of the dimensions,
59 dimension fields (separated from dimension by dot) or dataset type fields
60 (separated from dataset type name by dot).
61 """
64@final
65class GeneralQueryResults(QueryResultsBase):
66 """A query for `DatasetRef` results with a single dataset type.
68 Parameters
69 ----------
70 driver : `QueryDriver`
71 Implementation object that knows how to actually execute queries.
72 tree : `QueryTree`
73 Description of the query as a tree of joins and column expressions. The
74 instance returned directly by the `Butler._query` entry point should be
75 constructed via `make_unit_query_tree`.
76 spec : `GeneralResultSpec`
77 Specification of the query result rows, including output columns,
78 ordering, and slicing.
80 Notes
81 -----
82 This class should never be constructed directly by users; use `Query`
83 methods instead.
84 """
86 def __init__(self, driver: QueryDriver, tree: QueryTree, spec: GeneralResultSpec):
87 spec.validate_tree(tree)
88 super().__init__(driver, tree)
89 self._spec = spec
91 def __iter__(self) -> Iterator[dict[str, Any]]:
92 """Iterate over result rows.
94 Yields
95 ------
96 row_dict : `dict` [`str`, `Any`]
97 Result row as dictionary, the keys are the names of the dimensions,
98 dimension fields (separated from dimension by dot) or dataset type
99 fields (separated from dataset type name by dot).
100 """
101 for page in self._driver.execute(self._spec, self._tree):
102 columns = tuple(str(column) for column in page.spec.get_result_columns())
103 for row in page.rows:
104 result = dict(zip(columns, row, strict=True))
105 if page.dimension_records:
106 records = self._get_cached_dimension_records(result, page.dimension_records)
107 self._add_dimension_records(result, records)
108 yield result
110 def iter_tuples(self, *dataset_types: DatasetType) -> Iterator[GeneralResultTuple]:
111 """Iterate over result rows and return data coordinate, and dataset
112 refs constructed from each row, and an original row.
114 This object has to include "dataset_id" and "run" columns for each type
115 in ``dataset_types``.
117 Parameters
118 ----------
119 *dataset_types : `DatasetType`
120 Zero or more types of the datasets to return.
122 Yields
123 ------
124 row_tuple : `GeneralResultTuple`
125 Structure containing data coordinate, refs, and a copy of the row.
126 """
127 all_dimensions = self._spec.dimensions
128 dataset_keys: list[tuple[DatasetType, DimensionGroup, str, str]] = []
129 for dataset_type in dataset_types:
130 dimensions = dataset_type.dimensions
131 id_key = f"{dataset_type.name}.dataset_id"
132 run_key = f"{dataset_type.name}.run"
133 dataset_keys.append((dataset_type, dimensions, id_key, run_key))
134 for page in self._driver.execute(self._spec, self._tree):
135 columns = tuple(str(column) for column in page.spec.get_result_columns())
136 for page_row in page.rows:
137 row = dict(zip(columns, page_row, strict=True))
138 if page.dimension_records:
139 cached_records = self._get_cached_dimension_records(row, page.dimension_records)
140 self._add_dimension_records(row, cached_records)
141 else:
142 cached_records = {}
143 data_coordinate = self._make_data_id(row, all_dimensions, cached_records)
144 refs = []
145 for dataset_type, dimensions, id_key, run_key in dataset_keys:
146 data_id = data_coordinate.subset(dimensions)
147 refs.append(DatasetRef(dataset_type, data_id, row[run_key], id=row[id_key]))
148 yield GeneralResultTuple(data_id=data_coordinate, refs=refs, raw_row=row)
150 @property
151 def dimensions(self) -> DimensionGroup:
152 # Docstring inherited
153 return self._spec.dimensions
155 @property
156 def has_dimension_records(self) -> bool:
157 """Whether all data IDs in this iterable contain dimension records."""
158 return self._spec.include_dimension_records
160 def with_dimension_records(self) -> GeneralQueryResults:
161 """Return a results object for which `has_dimension_records` is
162 `True`.
163 """
164 if self.has_dimension_records:
165 return self
166 return self._copy(tree=self._tree, include_dimension_records=True)
168 def count(self, *, exact: bool = True, discard: bool = False) -> int:
169 # Docstring inherited.
170 return self._driver.count(self._tree, self._spec, exact=exact, discard=discard)
172 def _copy(self, tree: QueryTree, **kwargs: Any) -> GeneralQueryResults:
173 # Docstring inherited.
174 return GeneralQueryResults(self._driver, tree, self._spec.model_copy(update=kwargs))
176 def _get_datasets(self) -> frozenset[str]:
177 # Docstring inherited.
178 return frozenset(self._spec.dataset_fields)
180 def _make_data_id(
181 self,
182 row: dict[str, Any],
183 dimensions: DimensionGroup,
184 cached_row_records: dict[DimensionElement, DimensionRecord],
185 ) -> DataCoordinate:
186 values = tuple(row[key] for key in itertools.chain(dimensions.required, dimensions.implied))
187 data_coordinate = DataCoordinate.from_full_values(dimensions, values)
188 if self.has_dimension_records:
189 records = {}
190 for name in dimensions.elements:
191 element = dimensions.universe[name]
192 record = cached_row_records.get(element)
193 if record is None:
194 record = self._make_dimension_record(row, dimensions.universe[name])
195 records[name] = record
196 data_coordinate = data_coordinate.expanded(records)
197 return data_coordinate
199 def _make_dimension_record(self, row: dict[str, Any], element: DimensionElement) -> DimensionRecord:
200 column_map = list(
201 zip(
202 element.schema.dimensions.names,
203 element.dimensions.names,
204 )
205 )
206 for field in element.schema.remainder.names:
207 column_map.append((field, str(ResultColumn(element.name, field))))
208 d = {k: row[v] for k, v in column_map}
209 record_cls = element.RecordClass
210 return record_cls(**d)
212 def _get_cached_dimension_records(
213 self, row: dict[str, Any], dimension_records: dict[DimensionElement, DimensionRecordSet]
214 ) -> dict[DimensionElement, DimensionRecord]:
215 """Find cached dimension records matching this row."""
216 records = {}
217 for element, element_records in dimension_records.items():
218 required_values = tuple(row[key] for key in element.required.names)
219 records[element] = element_records.find_with_required_values(required_values)
220 return records
222 def _add_dimension_records(
223 self, row: dict[str, Any], records: dict[DimensionElement, DimensionRecord]
224 ) -> None:
225 """Extend row with the fields from cached dimension records."""
226 for element, record in records.items():
227 for name, value in record.toDict().items():
228 if name not in element.schema.required.names:
229 row[f"{element.name}.{name}"] = value