Coverage for python/lsst/daf/butler/queries/overlaps.py: 93%
145 statements
« 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/>.
28from __future__ import annotations
30__all__ = ("OverlapsVisitor",)
32import itertools
33from collections.abc import Hashable, Iterable, Mapping, Sequence, Set
34from typing import Generic, Literal, TypeVar, cast
36from lsst.sphgeom import Region
38from .._exceptions import InvalidQueryError
39from .._topology import TopologicalFamily, TopologicalRelationshipEndpoint, TopologicalSpace
40from ..dimensions import DimensionElement, DimensionGroup
41from . import tree
42from .visitors import PredicateVisitFlags, SimplePredicateVisitor
44_T = TypeVar("_T", bound=Hashable)
47class _NaiveDisjointSet(Generic[_T]):
48 """A very naive (but simple) implementation of a "disjoint set" data
49 structure for strings, with mostly O(N) performance.
51 This class should not be used in any context where the number of elements
52 in the data structure is large. It intentionally implements a subset of
53 the interface of `scipy.cluster.DisJointSet` so that non-naive
54 implementation could be swapped in if desired.
56 Parameters
57 ----------
58 superset : `~collections.abc.Iterable` [ `str` ]
59 Elements to initialize the disjoint set, with each in its own
60 single-element subset.
61 """
63 def __init__(self, superset: Iterable[_T]):
64 self._subsets = [{k} for k in superset]
65 self._subsets.sort(key=len, reverse=True)
67 def merge(self, a: _T, b: _T) -> bool: # numpydoc ignore=PR04
68 """Merge the subsets containing the given elements.
70 Parameters
71 ----------
72 a
73 Element whose subset should be merged.
74 b
75 Element whose subset should be merged.
77 Returns
78 -------
79 merged : `bool`
80 `True` if a merge occurred, `False` if the elements were already in
81 the same subset.
82 """
83 for i, subset in enumerate(self._subsets): 83 ↛ 87line 83 didn't jump to line 87 because the loop on line 83 didn't complete
84 if a in subset:
85 break
86 else:
87 raise KeyError(f"Merge argument {a!r} not in disjoin set {self._subsets}.")
88 for j, subset in enumerate(self._subsets): 88 ↛ 92line 88 didn't jump to line 92 because the loop on line 88 didn't complete
89 if b in subset:
90 break
91 else:
92 raise KeyError(f"Merge argument {b!r} not in disjoin set {self._subsets}.")
93 if i == j: 93 ↛ 94line 93 didn't jump to line 94 because the condition on line 93 was never true
94 return False
95 i, j = sorted((i, j))
96 self._subsets[i].update(self._subsets[j])
97 del self._subsets[j]
98 self._subsets.sort(key=len, reverse=True)
99 return True
101 def subsets(self) -> Sequence[Set[_T]]:
102 """Return the current subsets, ordered from largest to smallest."""
103 return self._subsets
105 @property
106 def n_subsets(self) -> int:
107 """The number of subsets."""
108 return len(self._subsets)
111class CalibrationTemporalEndpoint(TopologicalRelationshipEndpoint):
112 """An implementation of the "topological relationship endpoint" interface
113 for a calibration dataset search.
115 Parameters
116 ----------
117 dataset_type_name : `str` or ``ANY_DATASET``
118 Name of the dataset type.
120 Notes
121 -----
122 This lets validity range lookups participate in the logic that checks to
123 see if an explicit spatial/temporal join in the WHERE expression is present
124 and hence an automatic join is unnecessary. That logic is simple for
125 datasets, since each "family" is a single dataset type that only has one
126 endpoint (whereas different dimensions like tract and patch can belong to
127 the same family).
128 """
130 def __init__(self, dataset_type_name: str | tree.AnyDatasetType):
131 self.dataset_type_name: str | tree.AnyDatasetType = dataset_type_name
133 @property
134 def name(self) -> str:
135 return self.dataset_type_name if self.dataset_type_name is not tree.ANY_DATASET else "<calibrations>"
137 @property
138 def topology(self) -> Mapping[TopologicalSpace, TopologicalFamily]:
139 return {TopologicalSpace.TEMPORAL: CalibrationTemporalFamily(self.dataset_type_name)}
142class CalibrationTemporalFamily(TopologicalFamily):
143 """An implementation of the "topological relationship endpoint" interface
144 for a calibration dataset search.
146 See `CalibrationTemporalEndpoint` for rationale.
148 Parameters
149 ----------
150 dataset_type_name : `str` or ``ANY_DATASET``
151 Name of the dataset type.
152 """
154 def __init__(self, dataset_type_name: str | tree.AnyDatasetType):
155 super().__init__(
156 dataset_type_name if dataset_type_name is not tree.ANY_DATASET else "<calibrations>",
157 TopologicalSpace.TEMPORAL,
158 )
159 self.dataset_type_name: str | tree.AnyDatasetType = dataset_type_name
161 def choose(self, dimensions: DimensionGroup) -> CalibrationTemporalEndpoint:
162 return CalibrationTemporalEndpoint(self.dataset_type_name)
164 def make_column_reference(self, endpoint: TopologicalRelationshipEndpoint) -> tree.DatasetFieldReference:
165 return tree.DatasetFieldReference(dataset_type=self.dataset_type_name, field="timespan")
168class OverlapsVisitor(SimplePredicateVisitor):
169 """A helper class for dealing with spatial and temporal overlaps in a
170 query.
172 Parameters
173 ----------
174 dimensions : `DimensionGroup`
175 Dimensions of the query.
176 calibration_dataset_types : `~collections.abc.Set` [ `str` ]
177 The names of dataset types that have been joined into the query via
178 a search that includes at least one calibration collection.
180 Notes
181 -----
182 This class includes logic for extracting explicit spatial and temporal
183 joins from a WHERE-clause predicate and computing automatic joins given the
184 dimensions of the query. It is designed to be subclassed by query driver
185 implementations that want to rewrite the predicate at the same time.
186 """
188 def __init__(self, dimensions: DimensionGroup, calibration_dataset_types: Set[str | tree.AnyDatasetType]):
189 self.dimensions = dimensions
190 self._spatial_connections = _NaiveDisjointSet(self.dimensions.spatial)
191 temporal_families: list[TopologicalFamily] = [
192 CalibrationTemporalFamily(name) for name in calibration_dataset_types
193 ]
194 temporal_families.extend(self.dimensions.temporal)
195 self._temporal_connections = _NaiveDisjointSet(temporal_families)
197 def run(self, predicate: tree.Predicate, join_operands: Iterable[DimensionGroup]) -> tree.Predicate:
198 """Process the given predicate to extract spatial and temporal
199 overlaps.
201 Parameters
202 ----------
203 predicate : `tree.Predicate`
204 Predicate to process.
205 join_operands : `~collections.abc.Iterable` [ `DimensionGroup` ]
206 The dimensions of logical tables being joined into this query;
207 these can included embedded spatial and temporal joins that can
208 make it unnecessary to add new ones.
210 Returns
211 -------
212 predicate : `tree.Predicate`
213 A possibly-modified predicate that should replace the original.
214 """
215 result = predicate.visit(self)
216 if result is None:
217 result = predicate
218 for join_operand_dimensions in join_operands:
219 self._add_join_operand_connections(
220 join_operand_dimensions.spatial,
221 self._spatial_connections,
222 join_operand_dimensions,
223 )
224 self._add_join_operand_connections(
225 join_operand_dimensions.temporal,
226 self._temporal_connections,
227 join_operand_dimensions,
228 )
229 result = result.logical_and(self._add_automatic_joins("spatial", self._spatial_connections))
230 result = result.logical_and(self._add_automatic_joins("temporal", self._temporal_connections))
231 return result
233 def visit_comparison(
234 self,
235 a: tree.ColumnExpression,
236 operator: tree.ComparisonOperator,
237 b: tree.ColumnExpression,
238 flags: PredicateVisitFlags,
239 ) -> tree.Predicate | None:
240 # Docstring inherited.
241 if operator == "overlaps":
242 if tree.is_one_timespan_and_one_datetime(a, b) or tree.is_one_timespan_and_one_ingest_date(a, b):
243 # Can be transformed directly without special handling here.
244 return None
245 elif a.column_type == "region":
246 return self.visit_spatial_overlap(a, b, flags)
247 elif b.column_type == "timespan":
248 return self.visit_temporal_overlap(a, b, flags)
249 else:
250 raise AssertionError(f"Unexpected column type {a.column_type} for overlap.")
251 return None
253 def _add_join_operand_connections(
254 self,
255 families: Iterable[TopologicalFamily],
256 connections: _NaiveDisjointSet[TopologicalFamily],
257 operand_dimensions: DimensionGroup,
258 ) -> None:
259 """Add overlap connections implied by a table or subquery.
261 Parameters
262 ----------
263 families : `~collections.abc.Iterable` [ `TpologicalFamily` ]
264 Iterable of spatial or temporal families in this operand's
265 dimensions.
266 connections : `_NaiveDisjointSet`
267 Relationships between spatial or temporal families to update.
268 operand_dimensions : `DimensionGroup`
269 Dimensions of of the table or subquery.
271 Notes
272 -----
273 We assume each join operand to a `tree.Select` has its own
274 complete set of spatial and temporal joins that went into generating
275 its rows. That will naturally be true for relations originating from
276 the butler database, like dataset searches and materializations, and if
277 it isn't true for a data ID upload, that would represent an intentional
278 association between non-overlapping things that we'd want to respect by
279 *not* adding a more restrictive automatic join.
280 """
281 for a_family, b_family in itertools.combinations(families, 2):
282 a_element = a_family.choose(self.dimensions)
283 b_element = b_family.choose(self.dimensions)
284 if (
285 a_element.name in operand_dimensions.elements
286 and b_element.name in operand_dimensions.elements
287 ):
288 connections.merge(a_family, b_family)
290 def _add_automatic_joins(
291 self,
292 kind: Literal["spatial", "temporal"],
293 connections: _NaiveDisjointSet[TopologicalFamily],
294 ) -> tree.Predicate:
295 if connections.n_subsets <= 1:
296 # All of the joins we need are already present.
297 return tree.Predicate.from_bool(True)
298 if connections.n_subsets > 2:
299 raise InvalidQueryError(
300 f"Too many disconnected sets of {kind} families for an automatic "
301 f"join: {connections.subsets()}. Add explicit {kind} joins to avoid this error."
302 )
303 a_subset, b_subset = connections.subsets()
304 if len(a_subset) > 1 or len(b_subset) > 1:
305 raise InvalidQueryError(
306 f"A {kind} join is needed between {a_subset} and {b_subset}, but which join to "
307 "add is ambiguous. Add an explicit spatial or temporal join to avoid this error."
308 )
309 # We have a pair of families that are not explicitly or implicitly
310 # connected to any other families; add an automatic join between their
311 # most fine-grained members.
312 (a_family,) = a_subset
313 (b_family,) = b_subset
314 a = a_family.make_column_reference(a_family.choose(self.dimensions))
315 b = b_family.make_column_reference(b_family.choose(self.dimensions))
316 join_predicate = self.visit_comparison(a, "overlaps", b, PredicateVisitFlags.HAS_AND_SIBLINGS)
317 if join_predicate is None:
318 join_predicate = tree.Predicate.compare(a, "overlaps", b)
319 return join_predicate
321 def visit_spatial_overlap(
322 self,
323 a: tree.ColumnExpression,
324 b: tree.ColumnExpression,
325 flags: PredicateVisitFlags,
326 ) -> tree.Predicate | None:
327 """Dispatch a spatial overlap comparison predicate to handlers.
329 This method should rarely (if ever) need to be overridden.
331 Parameters
332 ----------
333 a : `tree.ColumnExpression`
334 First operand.
335 b : `tree.ColumnExpression`
336 Second operand.
337 flags : `tree.PredicateLeafFlags`
338 Information about where this overlap comparison appears in the
339 larger predicate tree.
341 Returns
342 -------
343 replaced : `tree.Predicate` or `None`
344 The predicate to be inserted instead in the processed tree, or
345 `None` if no substitution is needed.
346 """
347 match a, b:
348 case tree.DimensionFieldReference(element=a_element), tree.DimensionFieldReference(
349 element=b_element
350 ):
351 return self.visit_spatial_join(a_element, b_element, flags)
352 case tree.DimensionFieldReference(element=element), region_expression:
353 pass
354 case region_expression, tree.DimensionFieldReference(element=element):
355 pass
356 case _:
357 raise InvalidQueryError(
358 "Spatial overlap comparison requires at least one dimension region column; "
359 f"got {a} and {b}."
360 )
361 if region := region_expression.get_literal_value(): 361 ↛ 363line 361 didn't jump to line 363 because the condition on line 361 was always true
362 return self.visit_spatial_constraint(element, region, flags)
363 raise InvalidQueryError(
364 f"Spatial overlap comparison requires its other operand to be a region literal; "
365 f"got {region_expression}."
366 )
368 def visit_temporal_overlap(
369 self,
370 a: tree.ColumnExpression,
371 b: tree.ColumnExpression,
372 flags: PredicateVisitFlags,
373 ) -> tree.Predicate | None:
374 """Dispatch a temporal overlap comparison predicate to handlers.
376 This method should rarely (if ever) need to be overridden.
378 Parameters
379 ----------
380 a : `tree.ColumnExpression`-
381 First operand.
382 b : `tree.ColumnExpression`
383 Second operand.
384 flags : `tree.PredicateLeafFlags`
385 Information about where this overlap comparison appears in the
386 larger predicate tree.
388 Returns
389 -------
390 replaced : `tree.Predicate` or `None`
391 The predicate to be inserted instead in the processed tree, or
392 `None` if no substitution is needed.
393 """
394 match a, b:
395 case (
396 tree.DimensionFieldReference(element=a_element),
397 tree.DimensionFieldReference(element=b_element),
398 ):
399 return self.visit_temporal_dimension_join(a_element, b_element, flags)
400 case (
401 tree.DatasetFieldReference(dataset_type=a_dataset),
402 tree.DimensionFieldReference(element=b_element),
403 ):
404 return self.visit_validity_range_dimension_join(a_dataset, b_element, flags)
405 case (
406 tree.DimensionFieldReference(element=a_element),
407 tree.DatasetFieldReference(dataset_type=b_dataset),
408 ):
409 return self.visit_validity_range_dimension_join(b_dataset, a_element, flags)
410 case (
411 tree.DatasetFieldReference(dataset_type=a_dataset),
412 tree.DatasetFieldReference(dataset_type=b_dataset),
413 ):
414 return self.visit_validity_range_join(a_dataset, b_dataset, flags)
415 case _:
416 # Other cases do not participate in automatic join logic and
417 # do not require the predicate to be rewritten.
418 return None
420 def visit_spatial_join(
421 self, a: DimensionElement, b: DimensionElement, flags: PredicateVisitFlags
422 ) -> tree.Predicate | None:
423 """Handle a spatial overlap comparison between two dimension elements.
425 The default implementation updates the set of known spatial connections
426 (for use by `compute_automatic_spatial_joins`) and returns `None`.
428 Parameters
429 ----------
430 a : `DimensionElement`
431 One element in the join.
432 b : `DimensionElement`
433 The other element in the join.
434 flags : `tree.PredicateLeafFlags`
435 Information about where this overlap comparison appears in the
436 larger predicate tree.
438 Returns
439 -------
440 replaced : `tree.Predicate` or `None`
441 The predicate to be inserted instead in the processed tree, or
442 `None` if no substitution is needed.
443 """
444 if a.spatial == b.spatial:
445 raise InvalidQueryError(f"Spatial join between {a} and {b} is not necessary.")
446 self._spatial_connections.merge(
447 cast(TopologicalFamily, a.spatial), cast(TopologicalFamily, b.spatial)
448 )
449 return None
451 def visit_spatial_constraint(
452 self,
453 element: DimensionElement,
454 region: Region,
455 flags: PredicateVisitFlags,
456 ) -> tree.Predicate | None:
457 """Handle a spatial overlap comparison between a dimension element and
458 a literal region.
460 The default implementation just returns `None`.
462 Parameters
463 ----------
464 element : `DimensionElement`
465 The dimension element in the comparison.
466 region : `lsst.sphgeom.Region`
467 The literal region in the comparison.
468 flags : `tree.PredicateLeafFlags`
469 Information about where this overlap comparison appears in the
470 larger predicate tree.
472 Returns
473 -------
474 replaced : `tree.Predicate` or `None`
475 The predicate to be inserted instead in the processed tree, or
476 `None` if no substitution is needed.
477 """
478 return None
480 def visit_temporal_dimension_join(
481 self, a: DimensionElement, b: DimensionElement, flags: PredicateVisitFlags
482 ) -> tree.Predicate | None:
483 """Handle a temporal overlap comparison between two dimension elements.
485 The default implementation updates the set of known temporal
486 connections (for use by `compute_automatic_temporal_joins`) and returns
487 `None`.
489 Parameters
490 ----------
491 a : `DimensionElement`
492 One element in the join.
493 b : `DimensionElement`
494 The other element in the join.
495 flags : `tree.PredicateLeafFlags`
496 Information about where this overlap comparison appears in the
497 larger predicate tree.
499 Returns
500 -------
501 replaced : `tree.Predicate` or `None`
502 The predicate to be inserted instead in the processed tree, or
503 `None` if no substitution is needed.
504 """
505 if a.temporal == b.temporal: 505 ↛ 507line 505 didn't jump to line 507 because the condition on line 505 was always true
506 raise InvalidQueryError(f"Temporal join between {a} and {b} is not necessary.")
507 self._temporal_connections.merge(
508 cast(TopologicalFamily, a.temporal), cast(TopologicalFamily, b.temporal)
509 )
510 return None
512 def visit_validity_range_dimension_join(
513 self, a: str | tree.AnyDatasetType, b: DimensionElement, flags: PredicateVisitFlags
514 ) -> tree.Predicate | None:
515 """Handle a temporal overlap comparison between two dimension elements.
517 The default implementation updates the set of known temporal
518 connections (for use by `compute_automatic_temporal_joins`) and returns
519 `None`.
521 Parameters
522 ----------
523 a : `str` or ``tree.AnyDatasetType``
524 Name of a calibration dataset type.
525 b : `DimensionElement`
526 The dimension element to join the dataset validity range to.
527 flags : `tree.PredicateLeafFlags`
528 Information about where this overlap comparison appears in the
529 larger predicate tree.
531 Returns
532 -------
533 replaced : `tree.Predicate` or `None`
534 The predicate to be inserted instead in the processed tree, or
535 `None` if no substitution is needed.
536 """
537 self._temporal_connections.merge(CalibrationTemporalFamily(a), cast(TopologicalFamily, b.temporal))
538 return None
540 def visit_validity_range_join(
541 self, a: str | tree.AnyDatasetType, b: str | tree.AnyDatasetType, flags: PredicateVisitFlags
542 ) -> tree.Predicate | None:
543 """Handle a temporal overlap comparison between two dimension elements.
545 The default implementation updates the set of known temporal
546 connections (for use by `compute_automatic_temporal_joins`) and returns
547 `None`.
549 Parameters
550 ----------
551 a : `str` or ``tree.AnyDatasetType``
552 Name of a calibration dataset type.
553 b : `str` or ``tree.AnyDatasetType``
554 Another claibration dataset type to join to.
555 flags : `tree.PredicateLeafFlags`
556 Information about where this overlap comparison appears in the
557 larger predicate tree.
559 Returns
560 -------
561 replaced : `tree.Predicate` or `None`
562 The predicate to be inserted instead in the processed tree, or
563 `None` if no substitution is needed.
564 """
565 self._temporal_connections.merge(CalibrationTemporalFamily(a), CalibrationTemporalFamily(b))
566 return None