Coverage for python/lsst/daf/butler/_dataset_ref.py: 86%
331 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-26 08:49 +0000
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-26 08:49 +0000
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__ = [
30 "AmbiguousDatasetError",
31 "DatasetDatastoreRecords",
32 "DatasetId",
33 "DatasetIdFactory",
34 "DatasetIdGenEnum",
35 "DatasetRef",
36 "SerializedDatasetRef",
37 "SerializedDatasetRefContainerV1",
38 "SerializedDatasetRefContainers",
39]
41import enum
42import logging
43import sys
44import uuid
45from collections.abc import Callable, Iterable, Mapping
46from typing import (
47 TYPE_CHECKING,
48 Annotated,
49 Any,
50 ClassVar,
51 Literal,
52 Protocol,
53 Self,
54 TypeAlias,
55 cast,
56 runtime_checkable,
57)
59import pydantic
60from pydantic import StrictStr
62from lsst.utils.classes import immutable
64from ._config_support import LookupKey
65from ._dataset_type import DatasetType, SerializedDatasetType
66from ._exceptions import InconsistentUniverseError
67from ._named import NamedKeyDict
68from ._uuid import generate_uuidv7
69from .datastore.stored_file_info import StoredDatastoreItemInfo
70from .dimensions import (
71 DataCoordinate,
72 DimensionDataAttacher,
73 DimensionDataExtractor,
74 DimensionGroup,
75 DimensionUniverse,
76 SerializableDimensionData,
77 SerializedDataCoordinate,
78 SerializedDataId,
79)
80from .json import from_json_pydantic, to_json_pydantic
81from .persistence_context import PersistenceContextVars
83if TYPE_CHECKING:
84 from ._storage_class import StorageClass
85 from .registry import Registry
87# Per-dataset records grouped by opaque table name, usually there is just one
88# opaque table.
89DatasetDatastoreRecords: TypeAlias = Mapping[str, list[StoredDatastoreItemInfo]]
92_LOG = logging.getLogger(__name__)
95class AmbiguousDatasetError(Exception):
96 """Raised when a `DatasetRef` is not resolved but should be.
98 This happens when the `DatasetRef` has no ID or run but the requested
99 operation requires one of them.
100 """
103@runtime_checkable
104class _DatasetRefGroupedIterable(Protocol):
105 """A package-private interface for iterables of `DatasetRef` that know how
106 to efficiently group their contents by `DatasetType`.
108 """
110 def _iter_by_dataset_type(self) -> Iterable[tuple[DatasetType, Iterable[DatasetRef]]]:
111 """Iterate over `DatasetRef` instances, one `DatasetType` at a time.
113 Returns
114 -------
115 grouped : `~collections.abc.Iterable` [ `tuple` [ `DatasetType`, \
116 `~collections.abc.Iterable` [ `DatasetRef` ]
117 An iterable of tuples, in which the first element is a dataset type
118 and the second is an iterable of `DatasetRef` objects with exactly
119 that dataset type.
120 """
121 ...
124class DatasetIdGenEnum(enum.Enum):
125 """Enum used to specify dataset ID generation options."""
127 UNIQUE = 0
128 """Unique mode generates unique ID for each inserted dataset, e.g.
129 auto-generated by database or random UUID.
130 """
132 DATAID_TYPE = 1
133 """In this mode ID is computed deterministically from a combination of
134 dataset type and dataId.
135 """
137 DATAID_TYPE_RUN = 2
138 """In this mode ID is computed deterministically from a combination of
139 dataset type, dataId, and run collection name.
140 """
143class DatasetIdFactory:
144 """Factory for dataset IDs (UUIDs).
146 For now the logic is hard-coded and is controlled by the user-provided
147 value of `DatasetIdGenEnum`. In the future we may implement a configurable
148 logic that can guess `DatasetIdGenEnum` value from other parameters.
149 """
151 NS_UUID = uuid.UUID("840b31d9-05cd-5161-b2c8-00d32b280d0f")
152 """Namespace UUID used for UUID5 generation. Do not change. This was
153 produced by ``uuid.uuid5(uuid.NAMESPACE_DNS, "lsst.org")``.
154 """
156 def makeDatasetId(
157 self,
158 run: str,
159 datasetType: DatasetType,
160 dataId: DataCoordinate,
161 idGenerationMode: DatasetIdGenEnum,
162 ) -> uuid.UUID:
163 """Generate dataset ID for a dataset.
165 Parameters
166 ----------
167 run : `str`
168 Name of the RUN collection for the dataset.
169 datasetType : `DatasetType`
170 Dataset type.
171 dataId : `DataCoordinate`
172 Expanded data ID for the dataset.
173 idGenerationMode : `DatasetIdGenEnum`
174 ID generation option. `~DatasetIdGenEnum.UNIQUE` makes a random
175 UUID4-type ID. `~DatasetIdGenEnum.DATAID_TYPE` makes a
176 deterministic UUID5-type ID based on a dataset type name and
177 ``dataId``. `~DatasetIdGenEnum.DATAID_TYPE_RUN` makes a
178 deterministic UUID5-type ID based on a dataset type name, run
179 collection name, and ``dataId``.
181 Returns
182 -------
183 datasetId : `uuid.UUID`
184 Dataset identifier.
185 """
186 if idGenerationMode is DatasetIdGenEnum.UNIQUE:
187 # Earlier versions of this code used UUIDv4. However, totally
188 # random IDs create problems for Postgres insert performance,
189 # because it scatters index updates randomly around the disk.
190 # UUIDv7 has similar uniqueness properties to v4, but IDs generated
191 # at the same time are close together in the index.
192 return generate_uuidv7()
193 else:
194 # WARNING: If you modify this code make sure that the order of
195 # items in the `items` list below never changes.
196 items: list[tuple[str, str]] = []
197 if idGenerationMode is DatasetIdGenEnum.DATAID_TYPE:
198 items = [
199 ("dataset_type", datasetType.name),
200 ]
201 elif idGenerationMode is DatasetIdGenEnum.DATAID_TYPE_RUN:
202 items = [
203 ("dataset_type", datasetType.name),
204 ("run", run),
205 ]
206 else:
207 raise ValueError(f"Unexpected ID generation mode: {idGenerationMode}")
209 for name, value in sorted(dataId.required.items()):
210 items.append((name, str(value)))
211 data = ",".join(f"{key}={value}" for key, value in items)
212 return uuid.uuid5(self.NS_UUID, data)
215# This is constant, so don't recreate a set for each instance
216_serializedDatasetRefFieldsSet = {"id", "datasetType", "dataId", "run", "component"}
219class SerializedDatasetRef(pydantic.BaseModel):
220 """Simplified model of a `DatasetRef` suitable for serialization."""
222 id: uuid.UUID
223 datasetType: SerializedDatasetType | None = None
224 dataId: SerializedDataCoordinate | None = None
225 run: StrictStr | None = None
226 component: StrictStr | None = None
228 # Can not use "after" validator since in some cases the validator
229 # seems to trigger with the datasetType field not yet set.
230 @pydantic.model_validator(mode="before") # type: ignore[attr-defined]
231 @classmethod
232 def check_consistent_parameters(cls, data: dict[str, Any]) -> dict[str, Any]:
233 has_datasetType = data.get("datasetType") is not None
234 has_dataId = data.get("dataId") is not None
235 if has_datasetType is not has_dataId: 235 ↛ 236line 235 didn't jump to line 236 because the condition on line 235 was never true
236 raise ValueError("If specifying datasetType or dataId, must specify both.")
238 if data.get("component") is not None and has_datasetType: 238 ↛ 239line 238 didn't jump to line 239 because the condition on line 238 was never true
239 raise ValueError("datasetType can not be set if component is given.")
240 return data
242 @classmethod
243 def direct(
244 cls,
245 *,
246 id: str,
247 run: str,
248 datasetType: dict[str, Any] | None = None,
249 dataId: dict[str, Any] | None = None,
250 component: str | None = None,
251 ) -> SerializedDatasetRef:
252 """Construct a `SerializedDatasetRef` directly without validators.
254 Parameters
255 ----------
256 id : `str`
257 The UUID in string form.
258 run : `str`
259 The run for this dataset.
260 datasetType : `dict` [`str`, `typing.Any`]
261 A representation of the dataset type.
262 dataId : `dict` [`str`, `typing.Any`]
263 A representation of the data ID.
264 component : `str` or `None`
265 Any component associated with this ref.
267 Returns
268 -------
269 serialized : `SerializedDatasetRef`
270 A Pydantic model representing the given parameters.
272 Notes
273 -----
274 This differs from the pydantic "construct" method in that the arguments
275 are explicitly what the model requires, and it will recurse through
276 members, constructing them from their corresponding `direct` methods.
278 The ``id`` parameter is a string representation of dataset ID, it is
279 converted to UUID by this method.
281 This method should only be called when the inputs are trusted.
282 """
283 serialized_datasetType = (
284 SerializedDatasetType.direct(**datasetType) if datasetType is not None else None
285 )
286 serialized_dataId = SerializedDataCoordinate.direct(**dataId) if dataId is not None else None
288 node = cls.model_construct(
289 _fields_set=_serializedDatasetRefFieldsSet,
290 id=uuid.UUID(id),
291 datasetType=serialized_datasetType,
292 dataId=serialized_dataId,
293 run=sys.intern(run),
294 component=component,
295 )
297 return node
300DatasetId: TypeAlias = uuid.UUID
301"""A type-annotation alias for dataset ID providing typing flexibility.
302"""
305@immutable
306class DatasetRef:
307 """Reference to a Dataset in a `Registry`.
309 A `DatasetRef` may point to a Dataset that currently does not yet exist
310 (e.g., because it is a predicted input for provenance).
312 Parameters
313 ----------
314 datasetType : `DatasetType`
315 The `DatasetType` for this Dataset.
316 dataId : `DataCoordinate`
317 A mapping of dimensions that labels the Dataset within a Collection.
318 run : `str`
319 The name of the run this dataset was associated with when it was
320 created.
321 id : `DatasetId`, optional
322 The unique identifier assigned when the dataset is created. If ``id``
323 is not specified, a new unique ID will be created.
324 conform : `bool`, optional
325 If `True` (default), call `DataCoordinate.standardize` to ensure that
326 the data ID's dimensions are consistent with the dataset type's.
327 `DatasetRef` instances for which those dimensions are not equal should
328 not be created in new code, but are still supported for backwards
329 compatibility. New code should only pass `False` if it can guarantee
330 that the dimensions are already consistent.
331 id_generation_mode : `DatasetIdGenEnum`
332 ID generation option. `~DatasetIdGenEnum.UNIQUE` makes a random
333 UUID4-type ID. `~DatasetIdGenEnum.DATAID_TYPE` makes a
334 deterministic UUID5-type ID based on a dataset type name and
335 ``dataId``. `~DatasetIdGenEnum.DATAID_TYPE_RUN` makes a
336 deterministic UUID5-type ID based on a dataset type name, run
337 collection name, and ``dataId``.
338 datastore_records : `DatasetDatastoreRecords` or `None`
339 Datastore records to attach.
341 Notes
342 -----
343 See also :ref:`daf_butler_organizing_datasets`
344 """
346 _serializedType: ClassVar[type[pydantic.BaseModel]] = SerializedDatasetRef
347 __slots__ = (
348 "_id",
349 "datasetType",
350 "dataId",
351 "run",
352 "_datastore_records",
353 )
355 def __init__(
356 self,
357 datasetType: DatasetType,
358 dataId: DataCoordinate,
359 run: str,
360 *,
361 id: DatasetId | None = None,
362 conform: bool = True,
363 id_generation_mode: DatasetIdGenEnum = DatasetIdGenEnum.UNIQUE,
364 datastore_records: DatasetDatastoreRecords | None = None,
365 ):
366 self.datasetType = datasetType
367 if conform:
368 self.dataId = DataCoordinate.standardize(dataId, dimensions=datasetType.dimensions)
369 else:
370 self.dataId = dataId
371 self.run = run
372 if id is not None:
373 self._id = id.int
374 else:
375 self._id = (
376 DatasetIdFactory()
377 .makeDatasetId(self.run, self.datasetType, self.dataId, id_generation_mode)
378 .int
379 )
380 self._datastore_records = datastore_records
382 @property
383 def id(self) -> DatasetId:
384 """Primary key of the dataset (`DatasetId`).
386 Cannot be changed after a `DatasetRef` is constructed.
387 """
388 return uuid.UUID(int=self._id)
390 def __eq__(self, other: Any) -> bool:
391 try:
392 return (self.datasetType, self.dataId, self.id) == (other.datasetType, other.dataId, other.id)
393 except AttributeError:
394 return NotImplemented
396 def __hash__(self) -> int:
397 return hash((self.datasetType, self.dataId, self.id))
399 @property
400 def dimensions(self) -> DimensionGroup:
401 """Dimensions associated with the underlying `DatasetType`."""
402 return self.datasetType.dimensions
404 def __repr__(self) -> str:
405 # We delegate to __str__ (i.e use "!s") for the data ID) below because
406 # DataCoordinate's __repr__ - while adhering to the guidelines for
407 # __repr__ - is much harder to users to read, while its __str__ just
408 # produces a dict that can also be passed to DatasetRef's constructor.
409 return f"DatasetRef({self.datasetType!r}, {self.dataId!s}, run={self.run!r}, id={self.id})"
411 def __str__(self) -> str:
412 s = (
413 f"{self.datasetType.name}@{self.dataId!s} [sc={self.datasetType.storageClass_name}]"
414 f" (run={self.run} id={self.id})"
415 )
416 return s
418 def __lt__(self, other: Any) -> bool:
419 # Sort by run, DatasetType name and then by DataCoordinate
420 # The __str__ representation is probably close enough but we
421 # need to ensure that sorting a DatasetRef matches what you would
422 # get if you sorted DatasetType+DataCoordinate
423 if not isinstance(other, type(self)):
424 return NotImplemented
426 # Group by run if defined, takes precedence over DatasetType
427 self_run = "" if self.run is None else self.run
428 other_run = "" if other.run is None else other.run
430 # Compare tuples in the priority order
431 return (self_run, self.datasetType, self.dataId) < (other_run, other.datasetType, other.dataId)
433 def to_simple(self, minimal: bool = False) -> SerializedDatasetRef:
434 """Convert this class to a simple python type.
436 This makes it suitable for serialization.
438 Parameters
439 ----------
440 minimal : `bool`, optional
441 Use minimal serialization. Requires Registry to convert
442 back to a full type.
444 Returns
445 -------
446 simple : `dict` or `int`
447 The object converted to a dictionary.
448 """
449 if minimal:
450 # The only thing needed to uniquely define a DatasetRef is its id
451 # so that can be used directly if it is not a component DatasetRef.
452 # Store is in a dict to allow us to easily add the planned origin
453 # information later without having to support an int and dict in
454 # simple form.
455 simple: dict[str, Any] = {"id": self.id}
456 if self.isComponent():
457 # We can still be a little minimalist with a component
458 # but we will also need to record the datasetType component
459 simple["component"] = self.datasetType.component()
460 return SerializedDatasetRef(**simple)
462 return SerializedDatasetRef(
463 datasetType=self.datasetType.to_simple(minimal=minimal),
464 dataId=self.dataId.to_simple(),
465 run=self.run,
466 id=self.id,
467 )
469 @classmethod
470 def from_simple(
471 cls,
472 simple: SerializedDatasetRef,
473 universe: DimensionUniverse | None = None,
474 registry: Registry | None = None,
475 datasetType: DatasetType | None = None,
476 ) -> DatasetRef:
477 """Construct a new object from simplified form.
479 Generally this is data returned from the `to_simple` method.
481 Parameters
482 ----------
483 simple : `dict` of [`str`, `typing.Any`]
484 The value returned by `to_simple()`.
485 universe : `DimensionUniverse`
486 The special graph of all known dimensions.
487 Can be `None` if a registry is provided.
488 registry : `lsst.daf.butler.Registry`, optional
489 Registry to use to convert simple form of a DatasetRef to
490 a full `DatasetRef`. Can be `None` if a full description of
491 the type is provided along with a universe.
492 datasetType : DatasetType, optional
493 If datasetType is supplied, this will be used as the datasetType
494 object in the resulting DatasetRef instead of being read from
495 the `SerializedDatasetRef`. This is useful when many refs share
496 the same type as memory can be saved. Defaults to None.
498 Returns
499 -------
500 ref : `DatasetRef`
501 Newly-constructed object.
502 """
503 cache = PersistenceContextVars.datasetRefs.get()
504 key = simple.id.int
505 if cache is not None and (ref := cache.get(key, None)) is not None: 505 ↛ 506line 505 didn't jump to line 506 because the condition on line 505 was never true
506 if datasetType is not None:
507 if (component := datasetType.component()) is not None:
508 ref = ref.makeComponentRef(component)
509 ref = ref.overrideStorageClass(datasetType.storageClass_name)
510 return ref
511 if simple.datasetType is not None:
512 _, component = DatasetType.splitDatasetTypeName(simple.datasetType.name)
513 if component is not None:
514 ref = ref.makeComponentRef(component)
515 if simple.datasetType.storageClass is not None:
516 ref = ref.overrideStorageClass(simple.datasetType.storageClass)
517 return ref
518 # If dataset type is not given ignore the cache, because we can't
519 # reliably return the right storage class.
520 # Minimalist component will just specify component and id and
521 # require registry to reconstruct
522 if simple.datasetType is None and simple.dataId is None and simple.run is None:
523 if registry is None: 523 ↛ 524line 523 didn't jump to line 524 because the condition on line 523 was never true
524 raise ValueError("Registry is required to construct component DatasetRef from integer id")
525 if simple.id is None: 525 ↛ 526line 525 didn't jump to line 526 because the condition on line 525 was never true
526 raise ValueError("For minimal DatasetRef the ID must be defined.")
527 ref = registry.getDataset(simple.id)
528 if ref is None: 528 ↛ 529line 528 didn't jump to line 529 because the condition on line 528 was never true
529 raise RuntimeError(f"No matching dataset found in registry for id {simple.id}")
530 if simple.component:
531 ref = ref.makeComponentRef(simple.component)
532 else:
533 if universe is None:
534 if registry is None: 534 ↛ 535line 534 didn't jump to line 535 because the condition on line 534 was never true
535 raise ValueError("One of universe or registry must be provided.")
536 universe = registry.dimensions
537 if datasetType is None:
538 if simple.datasetType is None: 538 ↛ 539line 538 didn't jump to line 539 because the condition on line 538 was never true
539 raise ValueError("Cannot determine Dataset type of this serialized class")
540 datasetType = DatasetType.from_simple(
541 simple.datasetType, universe=universe, registry=registry
542 )
543 if simple.dataId is None: 543 ↛ 545line 543 didn't jump to line 545 because the condition on line 543 was never true
544 # mypy
545 raise ValueError("The DataId must be specified to construct a DatasetRef")
546 dataId = DataCoordinate.from_simple(simple.dataId, universe=universe)
547 # Check that simple ref is resolved.
548 if simple.run is None: 548 ↛ 549line 548 didn't jump to line 549 because the condition on line 548 was never true
549 dstr = ""
550 if simple.datasetType is None:
551 dstr = f" (datasetType={datasetType.name!r})"
552 raise ValueError(
553 "Run collection name is missing from serialized representation. "
554 f"Encountered with {simple!r}{dstr}."
555 )
556 ref = cls(
557 datasetType,
558 dataId,
559 id=simple.id,
560 run=simple.run,
561 )
562 if cache is not None: 562 ↛ 563line 562 didn't jump to line 563 because the condition on line 562 was never true
563 if ref.datasetType.component() is not None:
564 cache[key] = ref.makeCompositeRef()
565 else:
566 cache[key] = ref
567 return ref
569 to_json = to_json_pydantic
570 from_json: ClassVar[Callable[..., Self]] = cast(Callable[..., Self], classmethod(from_json_pydantic))
572 @classmethod
573 def _unpickle(
574 cls,
575 datasetType: DatasetType,
576 dataId: DataCoordinate,
577 id: DatasetId,
578 run: str,
579 datastore_records: DatasetDatastoreRecords | None,
580 ) -> DatasetRef:
581 """Create new `DatasetRef`.
583 A custom factory method for use by `__reduce__` as a workaround for
584 its lack of support for keyword arguments.
585 """
586 return cls(datasetType, dataId, id=id, run=run, datastore_records=datastore_records)
588 def __reduce__(self) -> tuple:
589 return (
590 self._unpickle,
591 (self.datasetType, self.dataId, self.id, self.run, self._datastore_records),
592 )
594 def __deepcopy__(self, memo: dict) -> DatasetRef:
595 # DatasetRef is recursively immutable; see note in @immutable
596 # decorator.
597 return self
599 def expanded(self, dataId: DataCoordinate) -> DatasetRef:
600 """Return a new `DatasetRef` with the given expanded data ID.
602 Parameters
603 ----------
604 dataId : `DataCoordinate`
605 Data ID for the new `DatasetRef`. Must compare equal to the
606 original data ID.
608 Returns
609 -------
610 ref : `DatasetRef`
611 A new `DatasetRef` with the given data ID.
612 """
613 assert dataId == self.dataId
614 return DatasetRef(
615 datasetType=self.datasetType,
616 dataId=dataId,
617 id=self.id,
618 run=self.run,
619 conform=False,
620 datastore_records=self._datastore_records,
621 )
623 def isComponent(self) -> bool:
624 """Indicate whether this `DatasetRef` refers to a component.
626 Returns
627 -------
628 isComponent : `bool`
629 `True` if this `DatasetRef` is a component, `False` otherwise.
630 """
631 return self.datasetType.isComponent()
633 def isComposite(self) -> bool:
634 """Boolean indicating whether this `DatasetRef` is a composite type.
636 Returns
637 -------
638 isComposite : `bool`
639 `True` if this `DatasetRef` is a composite type, `False`
640 otherwise.
641 """
642 return self.datasetType.isComposite()
644 def _lookupNames(self) -> tuple[LookupKey, ...]:
645 """Name keys to use when looking up this DatasetRef in a configuration.
647 The names are returned in order of priority.
649 Returns
650 -------
651 names : `tuple` of `LookupKey`
652 Tuple of the `DatasetType` name and the `StorageClass` name.
653 If ``instrument`` is defined in the dataId, each of those names
654 is added to the start of the tuple with a key derived from the
655 value of ``instrument``.
656 """
657 # Special case the instrument Dimension since we allow configs
658 # to include the instrument name in the hierarchy.
659 names: tuple[LookupKey, ...] = self.datasetType._lookupNames()
661 if "instrument" in self.dataId:
662 names = tuple(n.clone(dataId={"instrument": self.dataId["instrument"]}) for n in names) + names
664 return names
666 @staticmethod
667 def groupByType(refs: Iterable[DatasetRef]) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
668 """Group an iterable of `DatasetRef` by `DatasetType`.
670 Parameters
671 ----------
672 refs : `~collections.abc.Iterable` [ `DatasetRef` ]
673 `DatasetRef` instances to group.
675 Returns
676 -------
677 grouped : `NamedKeyDict` [ `DatasetType`, `list` [ `DatasetRef` ] ]
678 Grouped `DatasetRef` instances.
680 Notes
681 -----
682 When lazy item-iterables are acceptable instead of a full mapping,
683 `iter_by_type` can in some cases be far more efficient.
684 """
685 result: NamedKeyDict[DatasetType, list[DatasetRef]] = NamedKeyDict()
686 for ref in refs:
687 result.setdefault(ref.datasetType, []).append(ref)
688 return result
690 @staticmethod
691 def iter_by_type(
692 refs: Iterable[DatasetRef],
693 ) -> Iterable[tuple[DatasetType, Iterable[DatasetRef]]]:
694 """Group an iterable of `DatasetRef` by `DatasetType` with special
695 hooks for custom iterables that can do this efficiently.
697 Parameters
698 ----------
699 refs : `~collections.abc.Iterable` [ `DatasetRef` ]
700 `DatasetRef` instances to group. If this satisfies the
701 `_DatasetRefGroupedIterable` protocol, its
702 `~_DatasetRefGroupedIterable._iter_by_dataset_type` method will
703 be called.
705 Returns
706 -------
707 grouped : `~collections.abc.Iterable` [ `tuple` [ `DatasetType`, \
708 `~collections.abc.Iterable` [ `DatasetRef` ] ]]
709 Grouped `DatasetRef` instances.
710 """
711 if isinstance(refs, _DatasetRefGroupedIterable):
712 return refs._iter_by_dataset_type()
713 return DatasetRef.groupByType(refs).items()
715 def makeCompositeRef(self) -> DatasetRef:
716 """Create a `DatasetRef` of the composite from a component ref.
718 Requires that this `DatasetRef` is a component.
720 Returns
721 -------
722 ref : `DatasetRef`
723 A `DatasetRef` with a dataset type that corresponds to the
724 composite parent of this component, and the same ID and run
725 (which may be `None`, if they are `None` in ``self``).
726 """
727 # Assume that the data ID does not need to be standardized
728 # and should match whatever this ref already has.
729 return DatasetRef(
730 self.datasetType.makeCompositeDatasetType(),
731 self.dataId,
732 id=self.id,
733 run=self.run,
734 conform=False,
735 datastore_records=self._datastore_records,
736 )
738 def makeComponentRef(self, name: str) -> DatasetRef:
739 """Create a `DatasetRef` that corresponds to a component.
741 Parameters
742 ----------
743 name : `str`
744 Name of the component.
746 Returns
747 -------
748 ref : `DatasetRef`
749 A `DatasetRef` with a dataset type that corresponds to the given
750 component, and the same ID and run
751 (which may be `None`, if they are `None` in ``self``).
752 """
753 # Assume that the data ID does not need to be standardized
754 # and should match whatever this ref already has.
755 return DatasetRef(
756 self.datasetType.makeComponentDatasetType(name),
757 self.dataId,
758 id=self.id,
759 run=self.run,
760 conform=False,
761 datastore_records=self._datastore_records,
762 )
764 def overrideStorageClass(self, storageClass: str | StorageClass) -> DatasetRef:
765 """Create a new `DatasetRef` from this one, but with a modified
766 `DatasetType` that has a different `StorageClass`.
768 Parameters
769 ----------
770 storageClass : `str` or `StorageClass`
771 The new storage class.
773 Returns
774 -------
775 modified : `DatasetRef`
776 A new dataset reference that is the same as the current one but
777 with a different storage class in the `DatasetType`.
778 """
779 return self.replace(storage_class=storageClass)
781 def replace(
782 self,
783 *,
784 id: DatasetId | None = None,
785 run: str | None = None,
786 storage_class: str | StorageClass | None = None,
787 datastore_records: DatasetDatastoreRecords | None | Literal[False] = False,
788 ) -> DatasetRef:
789 """Create a new `DatasetRef` from this one, but with some modified
790 attributes.
792 Parameters
793 ----------
794 id : `DatasetId` or `None`
795 If not `None` then update dataset ID.
796 run : `str` or `None`
797 If not `None` then update run collection name. If ``dataset_id`` is
798 `None` then this will also cause new dataset ID to be generated.
799 storage_class : `str` or `StorageClass` or `None`
800 The new storage class. If not `None`, replaces existing storage
801 class.
802 datastore_records : `DatasetDatastoreRecords` or `None`
803 New datastore records. If `None` remove all records. By default
804 datastore records are preserved.
806 Returns
807 -------
808 modified : `DatasetRef`
809 A new dataset reference with updated attributes.
810 """
811 if datastore_records is False:
812 datastore_records = self._datastore_records
813 if storage_class is None:
814 datasetType = self.datasetType
815 else:
816 datasetType = self.datasetType.overrideStorageClass(storage_class)
817 if run is None:
818 run = self.run
819 # Do not regenerate dataset ID if run is the same.
820 if id is None:
821 id = self.id
822 return DatasetRef(
823 datasetType=datasetType,
824 dataId=self.dataId,
825 run=run,
826 id=id,
827 conform=False,
828 datastore_records=datastore_records,
829 )
831 def conform_to(self, universe: DimensionUniverse) -> DatasetRef:
832 """Rebuild this dataset ref in a different dimension universe.
834 Parameters
835 ----------
836 universe : `DimensionUniverse`
837 Target dimension universe.
839 Returns
840 -------
841 ref : `DatasetRef`
842 This ref if its dataset type's dimensions already belong to
843 ``universe``, else an otherwise-identical ref whose dataset type
844 and data ID belong to ``universe``. The data ID of a rebuilt ref
845 holds only the required dimension values; any implied values or
846 attached dimension records are dropped and must be recovered by
847 expanding the data ID against a data repository that uses the
848 target universe.
850 Raises
851 ------
852 InconsistentUniverseError
853 Raised if the dataset type cannot be rebuilt in the target
854 universe (see `DatasetType.conform_to`).
855 """
856 if self.datasetType.dimensions.universe is universe:
857 return self
858 datasetType = self.datasetType.conform_to(universe)
859 dataId = DataCoordinate.standardize(dict(self.dataId.required), dimensions=datasetType.dimensions)
860 return DatasetRef(
861 datasetType=datasetType,
862 dataId=dataId,
863 run=self.run,
864 id=self.id,
865 conform=False,
866 datastore_records=self._datastore_records,
867 )
869 def is_compatible_with(self, other: DatasetRef) -> bool:
870 """Determine if the given `DatasetRef` is compatible with this one.
872 Parameters
873 ----------
874 other : `DatasetRef`
875 Dataset ref to check.
877 Returns
878 -------
879 is_compatible : `bool`
880 Returns `True` if the other dataset ref is either the same as this
881 or the dataset type associated with the other is compatible with
882 this one and the dataId and dataset ID match.
884 Notes
885 -----
886 Compatibility requires that the dataId and dataset ID match and the
887 `DatasetType` is compatible. Compatibility is defined as the storage
888 class associated with the dataset type of the other ref can be
889 converted to this storage class.
891 Specifically this means that if you have done:
893 .. code-block:: py
895 new_ref = ref.overrideStorageClass(sc)
897 and this is successful, then the guarantee is that:
899 .. code-block:: py
901 assert ref.is_compatible_with(new_ref) is True
903 since we know that the python type associated with the new ref can
904 be converted to the original python type. The reverse is not guaranteed
905 and depends on whether bidirectional converters have been registered.
906 """
907 if self.id != other.id:
908 return False
909 if self.dataId != other.dataId:
910 return False
911 if self.run != other.run:
912 return False
913 return self.datasetType.is_compatible_with(other.datasetType)
915 datasetType: DatasetType
916 """The definition of this dataset (`DatasetType`).
918 Cannot be changed after a `DatasetRef` is constructed.
919 """
921 dataId: DataCoordinate
922 """A mapping of `Dimension` primary key values that labels the dataset
923 within a Collection (`DataCoordinate`).
925 Cannot be changed after a `DatasetRef` is constructed.
926 """
928 run: str
929 """The name of the run that produced the dataset.
931 Cannot be changed after a `DatasetRef` is constructed.
932 """
934 datastore_records: DatasetDatastoreRecords | None
935 """Optional datastore records (`DatasetDatastoreRecords`).
937 Cannot be changed after a `DatasetRef` is constructed.
938 """
941class MinimalistSerializableDatasetRef(pydantic.BaseModel):
942 """Minimal information needed to define a DatasetRef.
944 The ID is not included and is presumed to be the key to a mapping
945 to this information.
946 """
948 model_config = pydantic.ConfigDict(frozen=True)
950 dataset_type_name: str
951 """Name of the dataset type."""
953 run: str
954 """Name of the RUN collection."""
956 data_id: SerializedDataId
957 """Data coordinate of this dataset."""
959 def to_dataset_ref(
960 self,
961 id: DatasetId,
962 *,
963 dataset_type: DatasetType,
964 universe: DimensionUniverse,
965 attacher: DimensionDataAttacher | None = None,
966 ) -> DatasetRef:
967 """Convert serialized object to a `DatasetRef`.
969 Parameters
970 ----------
971 id : `DatasetId`
972 UUID identifying the dataset.
973 dataset_type : `DatasetType`
974 `DatasetType` record corresponding to the dataset type name in the
975 serialized object.
976 universe : `DimensionUniverse`
977 Dimension universe for the dataset.
978 attacher : `DimensionDataAttacher`, optional
979 If provided, will be used to add dimension records to the
980 deserialized `DatasetRef` instance.
982 Returns
983 -------
984 ref : `DatasetRef`
985 The deserialized object.
986 """
987 assert dataset_type.name == self.dataset_type_name, (
988 "Given DatasetType does not match the serialized dataset type name"
989 )
990 simple_data_id = SerializedDataCoordinate(dataId=self.data_id)
991 data_id = DataCoordinate.from_simple(simple=simple_data_id, universe=universe)
992 if attacher:
993 data_ids = attacher.attach(dataset_type.dimensions, [data_id])
994 data_id = data_ids[0]
995 return DatasetRef(
996 id=id,
997 run=self.run,
998 datasetType=dataset_type,
999 dataId=data_id,
1000 )
1002 @staticmethod
1003 def from_dataset_ref(ref: DatasetRef) -> MinimalistSerializableDatasetRef:
1004 """Serialize a ``DatasetRef` to a simplified format.
1006 Parameters
1007 ----------
1008 ref : `DatasetRef`
1009 `DatasetRef` object to serialize.
1010 """
1011 return MinimalistSerializableDatasetRef(
1012 dataset_type_name=ref.datasetType.name, run=ref.run, data_id=dict(ref.dataId.mapping)
1013 )
1016class SerializedDatasetRefContainer(pydantic.BaseModel):
1017 """Serializable model for a collection of DatasetRef.
1019 Dimension records are not included.
1020 """
1022 model_config = pydantic.ConfigDict(extra="allow", frozen=True)
1023 container_version: str
1026class SerializedDatasetRefContainerV1(SerializedDatasetRefContainer):
1027 """Serializable model for a collection of DatasetRef.
1029 Dimension records are not included.
1030 """
1032 container_version: Literal["V1"] = "V1"
1034 universe_version: int
1035 """Dimension universe version."""
1037 universe_namespace: str
1038 """Dimension universe namespace."""
1040 dataset_types: dict[str, SerializedDatasetType]
1041 """Dataset types indexed by their name."""
1043 compact_refs: dict[uuid.UUID, MinimalistSerializableDatasetRef]
1044 """Minimal dataset ref information indexed by UUID."""
1046 dimension_records: SerializableDimensionData | None = None
1047 """Dimension record information"""
1049 def __len__(self) -> int:
1050 """Return the number of datasets in the container."""
1051 return len(self.compact_refs)
1053 @classmethod
1054 def from_refs(cls, refs: Iterable[DatasetRef]) -> Self:
1055 """Construct a serializable form from a list of `DatasetRef`.
1057 Parameters
1058 ----------
1059 refs : `~collections.abc.Iterable` [ `DatasetRef` ]
1060 The datasets to include in the container.
1061 """
1062 # The serialized DatasetRef contains a lot of duplicated information.
1063 # We also want to drop dimension records and assume that the records
1064 # are already in the registry.
1065 universe: DimensionUniverse | None = None
1066 dataset_types: dict[str, SerializedDatasetType] = {}
1067 compact_refs: dict[uuid.UUID, MinimalistSerializableDatasetRef] = {}
1068 data_ids: list[DataCoordinate] = []
1069 dimensions: list[DimensionGroup] = []
1070 for ref in refs:
1071 if universe is None:
1072 universe = ref.datasetType.dimensions.universe
1073 if (name := ref.datasetType.name) not in dataset_types:
1074 dataset_types[name] = ref.datasetType.to_simple()
1075 compact_refs[ref.id] = MinimalistSerializableDatasetRef.from_dataset_ref(ref)
1076 if ref.dataId.hasRecords():
1077 dimensions.append(ref.datasetType.dimensions)
1078 data_ids.append(ref.dataId)
1080 # Extract dimension record metadata if present.
1081 dimension_records = None
1082 if data_ids and len(compact_refs) == len(data_ids):
1083 dimension_group = DimensionGroup.union(*dimensions, universe=universe)
1085 # Records were attached to all refs. Store them.
1086 extractor = DimensionDataExtractor.from_dimension_group(
1087 dimension_group,
1088 ignore_cached=False,
1089 include_skypix=False,
1090 )
1091 extractor.update(data_ids)
1092 dimension_records = SerializableDimensionData.from_record_sets(extractor.records.values())
1094 if universe: 1094 ↛ 1099line 1094 didn't jump to line 1099 because the condition on line 1094 was always true
1095 universe_version = universe.version
1096 universe_namespace = universe.namespace
1097 else:
1098 # No refs so no universe.
1099 universe_version = 0
1100 universe_namespace = "unknown"
1101 return cls(
1102 universe_version=universe_version,
1103 universe_namespace=universe_namespace,
1104 dataset_types=dataset_types,
1105 compact_refs=compact_refs,
1106 dimension_records=dimension_records,
1107 )
1109 def to_refs(self, universe: DimensionUniverse) -> list[DatasetRef]:
1110 """Construct the original `DatasetRef`.
1112 Parameters
1113 ----------
1114 universe : `DimensionUniverse`
1115 The universe to use when constructing the `DatasetRef`.
1117 Returns
1118 -------
1119 refs : `list` [ `DatasetRef` ]
1120 The `DatasetRef` that were serialized.
1121 """
1122 if not self.compact_refs: 1122 ↛ 1123line 1122 didn't jump to line 1123 because the condition on line 1122 was never true
1123 return []
1125 if universe.namespace != self.universe_namespace: 1125 ↛ 1126line 1125 didn't jump to line 1126 because the condition on line 1125 was never true
1126 raise InconsistentUniverseError(
1127 f"Can not convert to refs in universe {universe.namespace} that were created from "
1128 f"universe {self.universe_namespace}"
1129 )
1131 if universe.version != self.universe_version:
1132 _LOG.warning(
1133 "Universe mismatch when attempting to reconstruct DatasetRef from serialized form. "
1134 "Serialized with version %d but asked to use version %d. "
1135 "There could be failures due to different universe versions.",
1136 self.universe_version,
1137 universe.version,
1138 )
1140 # Reconstruct the DatasetType objects.
1141 dataset_types: dict[str, DatasetType] = {}
1142 if universe.version == self.universe_version:
1143 dataset_types = {
1144 name: DatasetType.from_simple(dtype, universe=universe)
1145 for name, dtype in self.dataset_types.items()
1146 }
1147 else:
1148 # When versions are different the dimensions may either disappear
1149 # or new dimensions can be aedded to conforming set.
1150 for name, dtype in self.dataset_types.items():
1151 try:
1152 dataset_type = DatasetType.from_simple(dtype, universe=universe)
1153 except KeyError as exc:
1154 raise InconsistentUniverseError(
1155 f"Source dimensions {dtype.dimensions} are not compatible with "
1156 f"target universe dimensions {universe}."
1157 ) from exc
1158 if set(dataset_type.dimensions.required) != set(dtype.dimensions or []):
1159 raise InconsistentUniverseError(
1160 f"Source dimensions {dtype.dimensions} are different from a conforming "
1161 f"set of target universe dimensions {dataset_type.dimensions}."
1162 )
1163 dataset_types[name] = dataset_type
1165 # Dimension records can be attached if available.
1166 # We assume that all dimension information was stored.
1167 attacher = None
1168 if self.dimension_records:
1169 attacher = DimensionDataAttacher(
1170 deserializers=self.dimension_records.make_deserializers(universe)
1171 )
1173 refs: list[DatasetRef] = []
1174 for id_, minimal in self.compact_refs.items():
1175 ref = minimal.to_dataset_ref(
1176 id_,
1177 dataset_type=dataset_types[minimal.dataset_type_name],
1178 universe=universe,
1179 attacher=attacher,
1180 )
1181 refs.append(ref)
1182 return refs
1185SerializedDatasetRefContainers: TypeAlias = Annotated[
1186 SerializedDatasetRefContainerV1,
1187 pydantic.Field(discriminator="container_version"),
1188]