Coverage for python/lsst/daf/butler/registry/queries/_readers.py: 39%
109 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__ = (
31 "DataCoordinateReader",
32 "DatasetRefReader",
33 "DimensionRecordReader",
34)
36from abc import ABC, abstractmethod
37from collections.abc import Callable, Mapping, Set
38from typing import TYPE_CHECKING, Any
40from lsst.utils.classes import cached_getter
42from ..._column_tags import DatasetColumnTag, DimensionKeyColumnTag
43from ..._dataset_ref import DatasetRef
44from ..._dataset_type import DatasetType
45from ...dimensions import (
46 DataCoordinate,
47 DimensionElement,
48 DimensionGroup,
49 DimensionRecord,
50 DimensionRecordSet,
51)
53if TYPE_CHECKING:
54 from lsst.daf.relation import ColumnTag
57class DataCoordinateReader(ABC):
58 """Base class and factory for reader objects that extract `DataCoordinate`
59 instances from query result rows.
60 """
62 @staticmethod
63 def make(
64 dimensions: DimensionGroup,
65 full: bool = True,
66 records: bool = False,
67 record_caches: Mapping[str, DimensionRecordSet] | None = None,
68 ) -> DataCoordinateReader:
69 """Construct a concrete reader for a set of dimensions.
71 Parameters
72 ----------
73 dimensions : `DimensionGroup`
74 Dimensions of the `DataCoordinate` instances the new reader will
75 read.
76 full : `bool`, optional
77 Whether to expect and extract implied dimensions as well as
78 required dimensions.
79 records : `bool`, optional
80 Whether to attach dimension records.
81 record_caches : `~collections.abc.Mapping`, optional
82 Mapping of cached dimension records. Ignored unless
83 ``records=True``.
85 Returns
86 -------
87 reader : `DataCoordinateReader`
88 Concrete reader instance.
89 """
90 if full:
91 full_reader = _FullDataCoordinateReader(dimensions)
92 if records:
93 if record_caches is None:
94 record_caches = {}
95 else:
96 record_caches = {
97 e: cache for e, cache in record_caches.items() if e in dimensions.elements
98 }
99 record_readers = {}
100 for element_name in dimensions.elements:
101 element = dimensions.universe[element_name]
102 if element_name not in record_caches:
103 record_readers[element] = DimensionRecordReader(element)
104 return _ExpandedDataCoordinateReader(full_reader, record_caches, record_readers)
105 return full_reader
106 else:
107 assert not records, "Cannot add records unless full=True."
108 return _BasicDataCoordinateReader(dimensions)
110 __slots__ = ()
112 @abstractmethod
113 def read(self, row: Mapping[ColumnTag, Any]) -> DataCoordinate:
114 """Read a `DataCoordinate` from a query result row.
116 Parameters
117 ----------
118 row : `~collections.abc.Mapping`
119 Mapping with `ColumnTag` keys representing a query result row.
121 Returns
122 -------
123 data_coordinate : `DataCoordinate`
124 New data ID.
125 """
126 raise NotImplementedError()
128 @property
129 @abstractmethod
130 def columns_required(self) -> Set[ColumnTag]:
131 raise NotImplementedError()
134class _BasicDataCoordinateReader(DataCoordinateReader):
135 """Private subclass of `DataCoordinateReader` for the ``full=False`` case.
137 Parameters
138 ----------
139 dimensions : `DimensionGroup`
140 Dimensions of the `DataCoordinate` instances read.
141 """
143 def __init__(self, dimensions: DimensionGroup):
144 self._dimensions = dimensions
145 self._tags = tuple(DimensionKeyColumnTag(name) for name in self._dimensions.required)
147 __slots__ = ("_dimensions", "_tags")
149 def read(self, row: Mapping[ColumnTag, Any]) -> DataCoordinate:
150 # Docstring inherited.
151 return DataCoordinate.from_required_values(
152 self._dimensions,
153 tuple(row[tag] for tag in self._tags),
154 )
156 @property
157 def columns_required(self) -> Set[ColumnTag]:
158 return frozenset(self._tags)
161class _FullDataCoordinateReader(DataCoordinateReader):
162 """Private subclass of `DataCoordinateReader` for the ``full=True`` case.
164 Parameters
165 ----------
166 dimensions : `DimensionGroup`
167 Dimensions of the `DataCoordinate` instances read.
168 """
170 def __init__(self, dimensions: DimensionGroup):
171 self._dimensions = dimensions
172 self._tags = tuple(DimensionKeyColumnTag(name) for name in self._dimensions.data_coordinate_keys)
174 __slots__ = ("_dimensions", "_tags")
176 def read(self, row: Mapping[ColumnTag, Any]) -> DataCoordinate:
177 # Docstring inherited.
178 return DataCoordinate.from_full_values(
179 self._dimensions,
180 tuple(row[tag] for tag in self._tags),
181 )
183 @property
184 def columns_required(self) -> Set[ColumnTag]:
185 return frozenset(self._tags)
188class _ExpandedDataCoordinateReader(DataCoordinateReader):
189 """Private subclass of `DataCoordinateReader` for the ``full=True`` case.
191 Parameters
192 ----------
193 full_reader : `_FullDataCoordinateReader`
194 Reader for full data IDs that don't have records.
195 record_caches : `~collections.abc.Mapping`
196 Mapping of cached dimension records.
197 record_readers : `~collections.abc.Mapping`
198 Mapping from `DimensionElement` to `DimensionRecordReaders`. Should
199 include all elements in the data coordinate's dimensions that are not
200 in ``record_cache``.
201 """
203 def __init__(
204 self,
205 full_reader: _FullDataCoordinateReader,
206 record_caches: Mapping[str, DimensionRecordSet],
207 record_readers: Mapping[DimensionElement, DimensionRecordReader],
208 ):
209 self._full_reader = full_reader
210 self._record_readers = record_readers
211 self._record_caches = record_caches
213 __slots__ = ("_full_reader", "_record_readers", "_record_caches", "_cached_columns_required")
215 def read(self, row: Mapping[ColumnTag, Any]) -> DataCoordinate:
216 # Docstring inherited.
217 full = self._full_reader.read(row)
218 records = {}
219 for element_name, cache in self._record_caches.items():
220 records[element_name] = cache.find(full.subset(cache.element.minimal_group))
221 for element, reader in self._record_readers.items():
222 records[element.name] = reader.read(row)
223 return full.expanded(records)
225 @property
226 @cached_getter
227 def columns_required(self) -> Set[ColumnTag]:
228 result = set(self._full_reader.columns_required)
229 for reader in self._record_readers.values():
230 result.update(reader.columns_required)
231 return result
234class DatasetRefReader:
235 """Reader class that extracts `DatasetRef` objects from query result rows.
237 Parameters
238 ----------
239 dataset_type : `DatasetType`
240 Dataset type for extracted references.
241 full : `bool`, optional
242 Whether to expect and extract implied dimensions as well as required
243 dimensions.
244 translate_collection : `~collections.abc.Callable`, optional
245 Callable that returns `str` collection names given collection primary
246 key values. Optional only for registries that use names as primary
247 keys, or if ``run`` is always passed to `read`.
248 records : `bool`, optional
249 Whether to attach dimension records to data IDs.
250 record_caches : `~collections.abc.Mapping`, optional
251 Nested mapping (outer keys are dimension element names, inner keys
252 are data IDs for that element) of cached dimension records.
253 Ignored unless ``records=True``.
254 """
256 def __init__(
257 self,
258 dataset_type: DatasetType,
259 *,
260 full: bool = True,
261 translate_collection: Callable[[Any], str] | None = None,
262 records: bool = False,
263 record_caches: Mapping[str, DimensionRecordSet] | None = None,
264 ):
265 self._data_coordinate_reader = DataCoordinateReader.make(
266 dataset_type.dimensions, full=full, records=records, record_caches=record_caches
267 )
268 self._dataset_type = dataset_type
269 self._translate_collection = translate_collection
270 self._id_tag = DatasetColumnTag(dataset_type.name, "dataset_id")
271 self._run_tag = DatasetColumnTag(dataset_type.name, "run")
273 __slots__ = (
274 "_data_coordinate_reader",
275 "_dataset_type",
276 "_translate_collection",
277 "_id_tag",
278 "_run_tag",
279 "_cached_columns_required",
280 )
282 def read(
283 self,
284 row: Mapping[ColumnTag, Any],
285 *,
286 run: str | None = None,
287 data_id: DataCoordinate | None = None,
288 ) -> DatasetRef:
289 """Read a `DatasetRef` from a query result row.
291 Parameters
292 ----------
293 row : `~collections.abc.Mapping`
294 Mapping with `ColumnTag` keys representing a query result row.
295 run : `str`, optional
296 Name of the `~CollectionType.RUN` collection; when provided the run
297 key does not need to be present in the result row, and
298 ``translate_collection`` does not need to be provided at
299 construction.
300 data_id : `DataCoordinate`, optional
301 Data ID; when provided the dimensions do not need to be present in
302 the result row.
303 """
304 if data_id is None:
305 data_id = self._data_coordinate_reader.read(row)
306 if run is None:
307 run_key = row[self._run_tag]
308 if self._translate_collection is not None:
309 run = self._translate_collection(run_key)
310 else:
311 run = run_key
312 return DatasetRef(
313 self._dataset_type,
314 data_id,
315 run=run,
316 id=row[self._id_tag],
317 )
319 @property
320 @cached_getter
321 def columns_required(self) -> Set[ColumnTag]:
322 result = set(self._data_coordinate_reader.columns_required)
323 result.add(self._id_tag)
324 result.add(self._run_tag)
325 return result
328class DimensionRecordReader:
329 """Read dimension records.
331 Parameters
332 ----------
333 element : `DimensionElement`
334 The element to read.
335 """
337 def __init__(self, element: DimensionElement):
338 self._cls = element.RecordClass
339 self._tags = element.RecordClass.fields.columns
341 __slots__ = ("_cls", "_tags")
343 def read(self, row: Mapping[ColumnTag, Any]) -> DimensionRecord:
344 return self._cls(**{name: row[tag] for tag, name in self._tags.items()})
346 @property
347 def columns_required(self) -> Set[ColumnTag]:
348 return self._tags.keys()