Coverage for python / lsst / pipe / base / quantum_graph_skeleton.py: 34%
206 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:14 -0800
1# This file is part of pipe_base.
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/>.
28"""An under-construction version of QuantumGraph and various helper
29classes.
30"""
32from __future__ import annotations
34__all__ = (
35 "DatasetKey",
36 "PrerequisiteDatasetKey",
37 "QuantumGraphSkeleton",
38 "QuantumKey",
39 "TaskInitKey",
40)
42import dataclasses
43from collections import defaultdict
44from collections.abc import Iterable, Iterator, MutableMapping, Set
45from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias
47import networkx
49from lsst.daf.butler import (
50 Butler,
51 DataCoordinate,
52 DataIdValue,
53 DatasetRef,
54 DimensionDataAttacher,
55 DimensionDataExtractor,
56 DimensionGroup,
57 DimensionRecordSet,
58)
59from lsst.utils.logging import getLogger
61if TYPE_CHECKING:
62 pass
64_LOG = getLogger(__name__)
67@dataclasses.dataclass(slots=True, eq=True, frozen=True)
68class QuantumKey:
69 """Identifier type for quantum keys in a `QuantumGraphSkeleton`."""
71 task_label: str
72 """Label of the task in the pipeline."""
74 data_id_values: tuple[DataIdValue, ...]
75 """Data ID values of the quantum.
77 Note that keys are fixed given `task_label`, so using only the values here
78 speeds up comparisons.
79 """
81 is_task: ClassVar[Literal[True]] = True
82 """Whether this node represents a quantum or task initialization rather
83 than a dataset (always `True`).
84 """
87@dataclasses.dataclass(slots=True, eq=True, frozen=True)
88class TaskInitKey:
89 """Identifier type for task init keys in a `QuantumGraphSkeleton`."""
91 task_label: str
92 """Label of the task in the pipeline."""
94 is_task: ClassVar[Literal[True]] = True
95 """Whether this node represents a quantum or task initialization rather
96 than a dataset (always `True`).
97 """
100@dataclasses.dataclass(slots=True, eq=True, frozen=True)
101class DatasetKey:
102 """Identifier type for dataset keys in a `QuantumGraphSkeleton`."""
104 parent_dataset_type_name: str
105 """Name of the dataset type (never a component)."""
107 data_id_values: tuple[DataIdValue, ...]
108 """Data ID values of the dataset.
110 Note that keys are fixed given `parent_dataset_type_name`, so using only
111 the values here speeds up comparisons.
112 """
114 is_task: ClassVar[Literal[False]] = False
115 """Whether this node represents a quantum or task initialization rather
116 than a dataset (always `False`).
117 """
119 is_prerequisite: ClassVar[Literal[False]] = False
122@dataclasses.dataclass(slots=True, eq=True, frozen=True)
123class PrerequisiteDatasetKey:
124 """Identifier type for prerequisite dataset keys in a
125 `QuantumGraphSkeleton`.
127 Unlike regular datasets, prerequisites are not actually required to come
128 from a find-first search of `input_collections`, so we don't want to
129 assume that the same data ID implies the same dataset. Happily we also
130 don't need to search for them by data ID in the graph, so we can use the
131 dataset ID (UUID) instead.
132 """
134 parent_dataset_type_name: str
135 """Name of the dataset type (never a component)."""
137 dataset_id_bytes: bytes
138 """Dataset ID (UUID) as raw bytes."""
140 is_task: ClassVar[Literal[False]] = False
141 """Whether this node represents a quantum or task initialization rather
142 than a dataset (always `False`).
143 """
145 is_prerequisite: ClassVar[Literal[True]] = True
148Key: TypeAlias = QuantumKey | TaskInitKey | DatasetKey | PrerequisiteDatasetKey
151class QuantumGraphSkeleton:
152 """An under-construction quantum graph.
154 QuantumGraphSkeleton is intended for use inside `QuantumGraphBuilder` and
155 its subclasses.
157 Parameters
158 ----------
159 task_labels : `~collections.abc.Iterable` [ `str` ]
160 The labels of all tasks whose quanta may be included in the graph, in
161 topological order.
163 Notes
164 -----
165 QuantumGraphSkeleton models a bipartite version of the quantum graph, in
166 which both quanta and datasets are represented as nodes and each type of
167 node only has edges to the other type.
169 Square-bracket (`getitem`) indexing returns a mutable mapping of a node's
170 flexible attributes.
172 The details of the `QuantumGraphSkeleton` API (e.g. which operations
173 operate on multiple nodes vs. a single node) are set by what's actually
174 needed by current quantum graph generation algorithms. New variants can be
175 added as needed, but adding all operations that *might* be useful for some
176 future algorithm seems premature.
177 """
179 def __init__(self, task_labels: Iterable[str]):
180 self._tasks: dict[str, tuple[TaskInitKey, set[QuantumKey]]] = {}
181 self._xgraph: networkx.DiGraph = networkx.DiGraph()
182 self._global_init_outputs: set[DatasetKey] = set()
183 self._dimension_data: dict[str, DimensionRecordSet] = {}
184 for task_label in task_labels:
185 task_init_key = TaskInitKey(task_label)
186 self._tasks[task_label] = (task_init_key, set())
187 self._xgraph.add_node(task_init_key)
189 def __contains__(self, key: Key) -> bool:
190 return key in self._xgraph.nodes
192 def __getitem__(self, key: Key) -> MutableMapping[str, Any]:
193 return self._xgraph.nodes[key]
195 def __iter__(self) -> Iterator[Key]:
196 return iter(self._xgraph.nodes)
198 @property
199 def n_nodes(self) -> int:
200 """The total number of nodes of all types."""
201 return len(self._xgraph.nodes)
203 @property
204 def n_edges(self) -> int:
205 """The total number of edges."""
206 return len(self._xgraph.edges)
208 @property
209 def has_any_quanta(self) -> bool:
210 """Test whether this graph has any quanta."""
211 for _ in self.iter_all_quanta():
212 return True
213 return False
215 def has_task(self, task_label: str) -> bool:
216 """Test whether the given task is in this skeleton.
218 Tasks are only added to the skeleton at initialization, but may be
219 removed by `remove_task` if they end up having no quanta.
221 Parameters
222 ----------
223 task_label : `str`
224 Task to check for.
226 Returns
227 -------
228 has : `bool`
229 `True` if the task is in this skeleton.
230 """
231 return task_label in self._tasks
233 def get_task_init_node(self, task_label: str) -> TaskInitKey:
234 """Return the graph node that represents a task's initialization.
236 Parameters
237 ----------
238 task_label : `str`
239 The task label to use.
241 Returns
242 -------
243 node : `TaskInitKey`
244 The graph node representing this task's initialization.
245 """
246 return self._tasks[task_label][0]
248 def get_quanta(self, task_label: str) -> Set[QuantumKey]:
249 """Return the quanta for the given task label.
251 Parameters
252 ----------
253 task_label : `str`
254 Label for the task.
256 Returns
257 -------
258 quanta : `~collections.abc.Set` [ `QuantumKey` ]
259 A set-like object with the identifiers of all quanta for the given
260 task. *The skeleton object's set of quanta must not be modified
261 while iterating over this container; make a copy if mutation during
262 iteration is necessary*.
263 """
264 return self._tasks[task_label][1]
266 @property
267 def global_init_outputs(self) -> Set[DatasetKey]:
268 """The set of dataset nodes that are not associated with any task."""
269 return self._global_init_outputs
271 def iter_all_quanta(self) -> Iterator[QuantumKey]:
272 """Iterate over all quanta from any task, in topological (but otherwise
273 unspecified) order.
274 """
275 for _, quanta in self._tasks.values():
276 yield from quanta
278 def iter_outputs_of(self, quantum_key: QuantumKey | TaskInitKey) -> Iterator[DatasetKey]:
279 """Iterate over the datasets produced by the given quantum.
281 Parameters
282 ----------
283 quantum_key : `QuantumKey` or `TaskInitKey`
284 Quantum to iterate over.
286 Returns
287 -------
288 datasets : `~collections.abc.Iterator` of `DatasetKey`
289 Datasets produced by the given quanta.
290 """
291 return self._xgraph.successors(quantum_key)
293 def iter_inputs_of(
294 self, quantum_key: QuantumKey | TaskInitKey
295 ) -> Iterator[DatasetKey | PrerequisiteDatasetKey]:
296 """Iterate over the datasets consumed by the given quantum.
298 Parameters
299 ----------
300 quantum_key : `QuantumKey` or `TaskInitKey`
301 Quantum to iterate over.
303 Returns
304 -------
305 datasets : `~collections.abc.Iterator` of `DatasetKey` \
306 or `PrequisiteDatasetKey`
307 Datasets consumed by the given quanta.
308 """
309 return self._xgraph.predecessors(quantum_key)
311 def update(self, other: QuantumGraphSkeleton) -> None:
312 """Copy all nodes from ``other`` to ``self``.
314 Parameters
315 ----------
316 other : `QuantumGraphSkeleton`
317 Source of nodes. The tasks in ``other`` must be a subset of the
318 tasks in ``self`` (this method is expected to be used to populate
319 a skeleton for a full from independent-subgraph skeletons).
320 """
321 for task_label, (_, quanta) in other._tasks.items():
322 self._tasks[task_label][1].update(quanta)
323 self._xgraph.update(other._xgraph)
324 for record_set in other._dimension_data.values():
325 self._dimension_data.setdefault(
326 record_set.element.name, DimensionRecordSet(record_set.element)
327 ).update(record_set)
329 def add_quantum_node(self, task_label: str, data_id: DataCoordinate, **attrs: Any) -> QuantumKey:
330 """Add a new node representing a quantum.
332 Parameters
333 ----------
334 task_label : `str`
335 Name of task.
336 data_id : `~lsst.daf.butler.DataCoordinate`
337 The data ID of the quantum.
338 **attrs : `~typing.Any`
339 Additional attributes.
340 """
341 key = QuantumKey(task_label, data_id.required_values)
342 self._xgraph.add_node(key, data_id=data_id, **attrs)
343 self._tasks[key.task_label][1].add(key)
344 return key
346 def add_dataset_node(
347 self,
348 parent_dataset_type_name: str,
349 data_id: DataCoordinate,
350 is_global_init_output: bool = False,
351 **attrs: Any,
352 ) -> DatasetKey:
353 """Add a new node representing a dataset.
355 Parameters
356 ----------
357 parent_dataset_type_name : `str`
358 Name of the parent dataset type.
359 data_id : `~lsst.daf.butler.DataCoordinate`
360 The dataset data ID.
361 is_global_init_output : `bool`, optional
362 Whether this dataset is a global init output.
363 **attrs : `~typing.Any`
364 Additional attributes for the node.
365 """
366 key = DatasetKey(parent_dataset_type_name, data_id.required_values)
367 self._xgraph.add_node(key, data_id=data_id, **attrs)
368 if is_global_init_output:
369 assert isinstance(key, DatasetKey), str(key)
370 self._global_init_outputs.add(key)
371 return key
373 def add_prerequisite_node(
374 self,
375 ref: DatasetRef,
376 **attrs: Any,
377 ) -> PrerequisiteDatasetKey:
378 """Add a new node representing a prerequisite input dataset.
380 Parameters
381 ----------
382 ref : `~lsst.daf.butler.DatasetRef`
383 The dataset ref of the prerequisite.
384 **attrs : `~typing.Any`
385 Additional attributes for the node.
387 Notes
388 -----
389 This automatically sets the 'existing_input' ref attribute (see
390 `set_existing_input_ref`), since prerequisites are always overall
391 inputs.
392 """
393 key = PrerequisiteDatasetKey(ref.datasetType.name, ref.id.bytes)
394 self._xgraph.add_node(key, data_id=ref.dataId, ref=ref, **attrs)
395 return key
397 def remove_quantum_node(self, key: QuantumKey, remove_outputs: bool) -> None:
398 """Remove a node representing a quantum.
400 Parameters
401 ----------
402 key : `QuantumKey`
403 Identifier for the node.
404 remove_outputs : `bool`
405 If `True`, also remove all dataset nodes produced by this quantum.
406 If `False`, any such dataset nodes will become overall inputs.
407 """
408 _, quanta = self._tasks[key.task_label]
409 quanta.remove(key)
410 if remove_outputs:
411 to_remove = list(self._xgraph.successors(key))
412 to_remove.append(key)
413 self._xgraph.remove_nodes_from(to_remove)
414 else:
415 self._xgraph.remove_node(key)
417 def remove_dataset_nodes(self, keys: Iterable[DatasetKey | PrerequisiteDatasetKey]) -> None:
418 """Remove nodes representing datasets.
420 Parameters
421 ----------
422 keys : `~collections.abc.Iterable` of `DatasetKey`\
423 or `PrerequisiteDatasetKey`
424 Nodes to remove.
425 """
426 self._xgraph.remove_nodes_from(keys)
428 def remove_task(self, task_label: str) -> None:
429 """Fully remove a task from the skeleton.
431 All init-output datasets and quanta for the task must already have been
432 removed.
434 Parameters
435 ----------
436 task_label : `str`
437 Name of task to remove.
438 """
439 task_init_key, quanta = self._tasks.pop(task_label)
440 assert not quanta, "Cannot remove task unless all quanta have already been removed."
441 assert not list(self._xgraph.successors(task_init_key))
442 self._xgraph.remove_node(task_init_key)
444 def add_input_edges(
445 self,
446 task_key: QuantumKey | TaskInitKey,
447 dataset_keys: Iterable[DatasetKey | PrerequisiteDatasetKey],
448 ) -> None:
449 """Add edges connecting datasets to a quantum that consumes them.
451 Parameters
452 ----------
453 task_key : `QuantumKey` or `TaskInitKey`
454 Quantum to connect.
455 dataset_keys : `~collections.abc.Iterable` of `DatasetKey`\
456 or `PrequisiteDatasetKey`
457 Datasets to join to the quantum.
459 Notes
460 -----
461 This must only be called if the task node has already been added.
462 Use `add_input_edge` if this cannot be assumed.
464 Dataset nodes that are not already present will be created.
465 """
466 assert task_key in self._xgraph, str(task_key)
467 self._xgraph.add_edges_from((dataset_key, task_key) for dataset_key in dataset_keys)
469 def remove_input_edges(
470 self,
471 task_key: QuantumKey | TaskInitKey,
472 dataset_keys: Iterable[DatasetKey | PrerequisiteDatasetKey],
473 ) -> None:
474 """Remove edges connecting datasets to a quantum that consumes them.
476 Parameters
477 ----------
478 task_key : `QuantumKey` or `TaskInitKey`
479 Quantum to disconnect.
480 dataset_keys : `~collections.abc.Iterable` of `DatasetKey`\
481 or `PrequisiteDatasetKey`
482 Datasets to remove from the quantum.
483 """
484 self._xgraph.remove_edges_from((dataset_key, task_key) for dataset_key in dataset_keys)
486 def add_input_edge(
487 self,
488 task_key: QuantumKey | TaskInitKey,
489 dataset_key: DatasetKey | PrerequisiteDatasetKey,
490 ignore_unrecognized_quanta: bool = False,
491 ) -> bool:
492 """Add an edge connecting a dataset to a quantum that consumes it.
494 Parameters
495 ----------
496 task_key : `QuantumKey` or `TaskInitKey`
497 Identifier for the quantum node.
498 dataset_key : `DatasetKey` or `PrerequisiteKey`
499 Identifier for the dataset node.
500 ignore_unrecognized_quanta : `bool`, optional
501 If `False`, do nothing if the quantum node is not already present.
502 If `True`, the quantum node is assumed to be present.
504 Returns
505 -------
506 added : `bool`
507 `True` if an edge was actually added, `False` if the quantum was
508 not recognized and the edge was not added as a result.
510 Notes
511 -----
512 Dataset nodes that are not already present will be created.
513 """
514 if ignore_unrecognized_quanta and task_key not in self._xgraph:
515 return False
516 self._xgraph.add_edge(dataset_key, task_key)
517 return True
519 def add_output_edge(self, task_key: QuantumKey | TaskInitKey, dataset_key: DatasetKey) -> None:
520 """Add an edge connecting a dataset to the quantum that produces it.
522 Parameters
523 ----------
524 task_key : `QuantumKey` or `TaskInitKey`
525 Identifier for the quantum node. Must identify a node already
526 present in the graph.
527 dataset_key : `DatasetKey`
528 Identifier for the dataset node. Must identify a node already
529 present in the graph.
530 """
531 assert task_key in self._xgraph, str(task_key)
532 assert dataset_key in self._xgraph, str(dataset_key)
533 self._xgraph.add_edge(task_key, dataset_key)
535 def remove_output_edge(self, dataset_key: DatasetKey) -> None:
536 """Remove the edge connecting a dataset to the quantum that produces
537 it.
539 Parameters
540 ----------
541 dataset_key : `DatasetKey`
542 Identifier for the dataset node. Must identify a node already
543 present in the graph.
544 """
545 (task_key,) = self._xgraph.predecessors(dataset_key)
546 assert dataset_key in self._xgraph, str(dataset_key)
547 self._xgraph.remove_edge(task_key, dataset_key)
549 def remove_orphan_datasets(self) -> None:
550 """Remove any dataset nodes that do not have any edges."""
551 for orphan in list(networkx.isolates(self._xgraph)):
552 if not orphan.is_task and orphan not in self._global_init_outputs:
553 self._xgraph.remove_node(orphan)
555 def extract_overall_inputs(self) -> dict[DatasetKey | PrerequisiteDatasetKey, DatasetRef]:
556 """Find overall input datasets.
558 Returns
559 -------
560 datasets : `dict` [ `DatasetKey` or `PrerequisiteDatasetKey`, \
561 `~lsst.daf.butler.DatasetRef` ]
562 Overall-input datasets, including prerequisites and init-inputs.
563 """
564 result = {}
565 for generation in networkx.algorithms.topological_generations(self._xgraph):
566 for dataset_key in generation:
567 if dataset_key.is_task:
568 continue
569 if (ref := self.get_dataset_ref(dataset_key)) is None:
570 raise AssertionError(
571 f"Logic bug in QG generation: dataset {dataset_key} was never resolved."
572 )
573 result[dataset_key] = ref
574 break
575 return result
577 def set_dataset_ref(
578 self, ref: DatasetRef, key: DatasetKey | PrerequisiteDatasetKey | None = None
579 ) -> None:
580 """Associate a dataset node with a `DatasetRef` instance.
582 Parameters
583 ----------
584 ref : `DatasetRef`
585 `DatasetRef` to associate with the node.
586 key : `DatasetKey` or `PrerequisiteDatasetKey`, optional
587 Identifier for the graph node. If not provided, a `DatasetKey`
588 is constructed from the dataset type name and data ID of ``ref``.
589 """
590 if key is None:
591 key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
592 self._xgraph.nodes[key]["ref"] = ref
594 def set_output_for_skip(self, ref: DatasetRef) -> None:
595 """Associate a dataset node with a `DatasetRef` that represents an
596 existing output in a collection where such outputs can cause a quantum
597 to be skipped.
599 Parameters
600 ----------
601 ref : `DatasetRef`
602 `DatasetRef` to associate with the node.
603 """
604 key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
605 self._xgraph.nodes[key]["output_for_skip"] = ref
607 def set_output_in_the_way(self, ref: DatasetRef) -> None:
608 """Associate a dataset node with a `DatasetRef` that represents an
609 existing output in the output RUN collectoin.
611 Parameters
612 ----------
613 ref : `DatasetRef`
614 `DatasetRef` to associate with the node.
615 """
616 key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
617 self._xgraph.nodes[key]["output_in_the_way"] = ref
619 def get_dataset_ref(self, key: DatasetKey | PrerequisiteDatasetKey) -> DatasetRef | None:
620 """Return the `DatasetRef` associated with the given node.
622 This does not return "output for skip" and "output in the way"
623 datasets.
625 Parameters
626 ----------
627 key : `DatasetKey` or `PrerequisiteDatasetKey`
628 Identifier for the graph node.
630 Returns
631 -------
632 ref : `DatasetRef` or `None`
633 Dataset reference associated with the node.
634 """
635 return self._xgraph.nodes[key].get("ref")
637 def get_output_for_skip(self, key: DatasetKey) -> DatasetRef | None:
638 """Return the `DatasetRef` associated with the given node in a
639 collection where it could lead to a quantum being skipped.
641 Parameters
642 ----------
643 key : `DatasetKey`
644 Identifier for the graph node.
646 Returns
647 -------
648 ref : `DatasetRef` or `None`
649 Dataset reference associated with the node.
650 """
651 return self._xgraph.nodes[key].get("output_for_skip")
653 def get_output_in_the_way(self, key: DatasetKey) -> DatasetRef | None:
654 """Return the `DatasetRef` associated with the given node in the
655 output RUN collection.
657 Parameters
658 ----------
659 key : `DatasetKey`
660 Identifier for the graph node.
662 Returns
663 -------
664 ref : `DatasetRef` or `None`
665 Dataset reference associated with the node.
666 """
667 return self._xgraph.nodes[key].get("output_in_the_way")
669 def discard_output_in_the_way(self, key: DatasetKey) -> None:
670 """Drop any `DatasetRef` associated with this node in the output RUN
671 collection.
673 Does nothing if there is no such `DatasetRef`.
675 Parameters
676 ----------
677 key : `DatasetKey`
678 Identifier for the graph node.
679 """
680 self._xgraph.nodes[key].pop("output_in_the_way", None)
682 def set_data_id(self, key: Key, data_id: DataCoordinate) -> None:
683 """Set the data ID associated with a node.
685 This updates the data ID in any `DatasetRef` objects associated with
686 the node via `set_ref`, `set_output_for_skip`, or
687 `set_output_in_the_way` as well, assuming it is an expanded version
688 of the original data ID.
690 Parameters
691 ----------
692 key : `Key`
693 Identifier for the graph node.
694 data_id : `DataCoordinate`
695 Data ID for the node.
696 """
697 state: MutableMapping[str, Any] = self._xgraph.nodes[key]
698 state["data_id"] = data_id
699 ref: DatasetRef | None
700 if (ref := state.get("ref")) is not None:
701 state["ref"] = ref.expanded(data_id)
702 output_for_skip: DatasetRef | None
703 if (output_for_skip := state.get("output_for_skip")) is not None:
704 state["output_for_skip"] = output_for_skip.expanded(data_id)
705 output_in_the_way: DatasetRef | None
706 if (output_in_the_way := state.get("output_in_the_way")) is not None:
707 state["output_in_the_way"] = output_in_the_way.expanded(data_id)
709 def get_data_id(self, key: Key) -> DataCoordinate:
710 """Return the full data ID for a quantum or dataset, if available.
712 Parameters
713 ----------
714 key : `Key`
715 Identifier for the graph node.
717 Returns
718 -------
719 data_id : `DataCoordinate`
720 Expanded data ID for the node, if one is available.
722 Raises
723 ------
724 KeyError
725 Raised if this node does not have an expanded data ID.
726 """
727 return self._xgraph.nodes[key]["data_id"]
729 def attach_dimension_records(
730 self,
731 butler: Butler,
732 dimensions: DimensionGroup,
733 dimension_records: Iterable[DimensionRecordSet] = (),
734 ) -> None:
735 """Attach dimension records to the data IDs in the skeleton.
737 This both attaches records to data IDs in the skeleton and aggregates
738 any existing records on data IDS, so `get_dimension_data` returns all
739 dimension records used in the skeleton. It can be called multiple
740 times.
742 Parameters
743 ----------
744 butler : `lsst.daf.butler.Butler`
745 Butler to use to query for missing dimension records.
746 dimensions : `lsst.daf.butler.DimensionGroup`
747 Superset of all of the dimensions of all data IDs.
748 dimension_records : `~collections.abc.Iterable` [ \
749 `lsst.daf.butler.DimensionRecordSet` ], optional
750 Iterable of sets of dimension records to attach.
751 """
752 for record_set in dimension_records:
753 self._dimension_data.setdefault(
754 record_set.element.name, DimensionRecordSet(record_set.element)
755 ).update(record_set)
756 # Group all nodes by data ID (and dimensions of data ID).
757 data_ids_to_expand: defaultdict[DimensionGroup, defaultdict[DataCoordinate, list[Key]]] = defaultdict(
758 lambda: defaultdict(list)
759 )
760 extractor = DimensionDataExtractor.from_dimension_group(dimensions)
761 data_id: DataCoordinate | None
762 for node_key in self:
763 if data_id := self[node_key].get("data_id"):
764 if data_id.hasRecords():
765 extractor.update([data_id])
766 else:
767 data_ids_to_expand[data_id.dimensions][data_id].append(node_key)
768 # Add records we extracted from data IDs that were already expanded, in
769 # case other nodes want them.
770 for record_set in extractor.records.values():
771 self._dimension_data.setdefault(
772 record_set.element.name, DimensionRecordSet(record_set.element)
773 ).update(record_set)
774 attacher = DimensionDataAttacher(records=self._dimension_data.values(), dimensions=dimensions)
775 for dimensions, data_ids in data_ids_to_expand.items():
776 with butler.query() as query:
777 # Butler query will be used as-needed to get dimension records
778 # (from prerequisites) we didn't fetch in advance. These are
779 # cached in the attacher so we don't look them up multiple
780 # times.
781 expanded_data_ids = attacher.attach(dimensions, data_ids.keys(), query=query)
782 for expanded_data_id, node_keys in zip(expanded_data_ids, data_ids.values()):
783 for node_key in node_keys:
784 self.set_data_id(node_key, expanded_data_id)
785 # Hold on to any records that we had to query for or extracted.
786 self._dimension_data = attacher.records
788 def get_dimension_data(self) -> list[DimensionRecordSet]:
789 """Return the dimension records attached to data IDs."""
790 return list(self._dimension_data.values())