Skip to content

Commit

Permalink
Attempt to avoid unbounded refcat lookups in QG generation.
Browse files Browse the repository at this point in the history
  • Loading branch information
TallJimbo committed Jul 28, 2023
1 parent 548ec45 commit fbbd0df
Showing 1 changed file with 49 additions and 0 deletions.
49 changes: 49 additions & 0 deletions python/lsst/pipe/base/graphBuilder.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,10 +50,12 @@
NamedValueSet,
Quantum,
Registry,
SkyPixDimension,
)
from lsst.daf.butler.registry import MissingCollectionError, MissingDatasetTypeError
from lsst.daf.butler.registry.queries import DataCoordinateQueryResults
from lsst.daf.butler.registry.wildcards import CollectionWildcard
from lsst.sphgeom import PixelizationABC, RangeSet

# -----------------------------
# Imports for other modules --
Expand Down Expand Up @@ -415,6 +417,27 @@ def __repr__(self) -> str:
inputs to this quantum.
"""

def computeSpatialExtent(self, pixelization: PixelizationABC) -> RangeSet:
"""Return the spatial extent of this quantum's inputs and outputs in
a skypix system.
Parameters
----------
pixelization : `lsst.sphgeom.PixelizationABC`
Pixelization system.
Returns
-------
extent : `lsst.sphgeom.RangeSet`
Ranges of sky pixels that touch this quantum's inputs and outputs.
"""
result = RangeSet()

Check warning on line 434 in python/lsst/pipe/base/graphBuilder.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/pipe/base/graphBuilder.py#L434

Added line #L434 was not covered by tests
for dataset_type, datasets in itertools.chain(self.inputs.items(), self.outputs.items()):
if dataset_type.dimensions.spatial:
for data_id in datasets.keys():
result |= pixelization.envelope(data_id.region)
return result

Check warning on line 439 in python/lsst/pipe/base/graphBuilder.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/pipe/base/graphBuilder.py#L438-L439

Added lines #L438 - L439 were not covered by tests

def makeQuantum(self, datastore_records: Mapping[str, DatastoreRecordData] | None = None) -> Quantum:
"""Transform the scaffolding object into a true `Quantum` instance.
Expand Down Expand Up @@ -1318,6 +1341,30 @@ def resolveDatasetRefs(
# which just means there are no datasets here.
prereq_refs = []
else:
where = ""
bind: dict[str, Any] = {}

Check warning on line 1345 in python/lsst/pipe/base/graphBuilder.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/pipe/base/graphBuilder.py#L1344-L1345

Added lines #L1344 - L1345 were not covered by tests
if not quantum.dataId.graph.spatial:
# This has skypix dimensions (probably a reference
# catalog), but the quantum's data is not spatial
# (it's probably a full-survey sequence point).
# Try to limit the spatial extent to the union of
# the spatial extent of the inputs and outputs.
for dimension in datasetType.dimensions:
if isinstance(dimension, SkyPixDimension):
extent = quantum.computeSpatialExtent(dimension.pixelization)
pixels: list[int] = []

Check warning on line 1355 in python/lsst/pipe/base/graphBuilder.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/pipe/base/graphBuilder.py#L1354-L1355

Added lines #L1354 - L1355 were not covered by tests
for begin, end in extent:
pixels.extend(range(begin, end))

Check warning on line 1357 in python/lsst/pipe/base/graphBuilder.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/pipe/base/graphBuilder.py#L1357

Added line #L1357 was not covered by tests
if not pixels:
_LOG.warning(

Check warning on line 1359 in python/lsst/pipe/base/graphBuilder.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/pipe/base/graphBuilder.py#L1359

Added line #L1359 was not covered by tests
"Prerequisite input %r to task %r may be unbounded.",
datasetType.name,
quantum.task.taskDef.label,
)
else:
bind["quantum_extent"] = pixels
where = f"{dimension.name} IN (quantum_extent)"
break

Check warning on line 1367 in python/lsst/pipe/base/graphBuilder.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/pipe/base/graphBuilder.py#L1365-L1367

Added lines #L1365 - L1367 were not covered by tests
# Most general case.
prereq_refs = [
prereq_ref if component is None else prereq_ref.makeComponentRef(component)
Expand All @@ -1326,6 +1373,8 @@ def resolveDatasetRefs(
collections=collections,
dataId=quantum.dataId,
findFirst=True,
where=where,
bind=bind,
).expanded()
]

Expand Down

0 comments on commit fbbd0df

Please sign in to comment.