Coverage for python/lsst/daf/butler/registry/queries/expressions/check.py: 27%
195 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-11 01:27 -0700
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
27from __future__ import annotations
29__all__ = (
30 "CheckVisitor",
31 "InspectionSummary",
32 "InspectionVisitor",
33)
35import dataclasses
36from collections.abc import Mapping, Sequence, Set
37from typing import TYPE_CHECKING, Any
38from uuid import UUID
40from ...._column_tags import DatasetColumnTag, DimensionKeyColumnTag, DimensionRecordColumnTag
41from ....dimensions import DataCoordinate, DataIdValue, Dimension, DimensionGroup, DimensionUniverse
42from ..._exceptions import UserExpressionError
43from .categorize import ExpressionConstant, categorizeConstant, categorizeElementId
44from .normalForm import NormalForm, NormalFormVisitor
45from .parser import Node, TreeVisitor
47if TYPE_CHECKING:
48 import astropy.time
50 from lsst.daf.relation import ColumnTag
53@dataclasses.dataclass
54class InspectionSummary:
55 """Base class for objects used by `CheckVisitor` and `InspectionVisitor`
56 to gather information about a parsed expression.
57 """
59 def update(self, other: InspectionSummary) -> None:
60 """Update ``self`` with all dimensions and columns from ``other``.
62 Parameters
63 ----------
64 other : `InspectionSummary`
65 The other summary object.
66 """
67 self.dimensions.update(other.dimensions)
68 for element, columns in other.columns.items():
69 self.columns.setdefault(element, set()).update(columns)
70 self.hasIngestDate = self.hasIngestDate or other.hasIngestDate
72 dimensions: set[str] = dataclasses.field(default_factory=set)
73 """Names of dimensions whose primary keys or dependencies were referenced
74 anywhere in this branch (`set` [ `str` ]).
75 """
77 columns: dict[str, set[str]] = dataclasses.field(default_factory=dict)
78 """Names of dimension element tables whose columns were referenced anywhere
79 in this branch (`dict` [ `str`, `set` [ `str` ] ]).
80 """
82 hasIngestDate: bool = False
83 """Whether this expression includes the special dataset ingest date
84 identifier (`bool`).
85 """
87 def make_column_tag_set(self, dataset_type_name: str | None) -> set[ColumnTag]:
88 """Transform the columns captured here into a set of `ColumnTag`
89 objects.
91 Parameters
92 ----------
93 dataset_type_name : `str` or `None`
94 Name of the dataset type to assume for unqualified dataset columns,
95 or `None` to reject any such identifiers.
97 Returns
98 -------
99 tag_set : `set` [ `ColumnTag` ]
100 Set of categorized column tags.
101 """
102 result: set[ColumnTag] = set()
103 if self.hasIngestDate:
104 if dataset_type_name is None:
105 raise UserExpressionError(
106 "Expression requires an ingest date, which requires exactly one dataset type."
107 )
108 result.add(DatasetColumnTag(dataset_type_name, "ingest_date"))
109 result.update(DimensionKeyColumnTag.generate(self.dimensions))
110 for dimension_element, columns in self.columns.items():
111 result.update(DimensionRecordColumnTag.generate(dimension_element, columns))
112 return result
115@dataclasses.dataclass
116class TreeSummary(InspectionSummary):
117 """Result object used by `InspectionVisitor` to gather information about
118 a parsed expression.
120 Notes
121 -----
122 TreeSummary adds attributes that allow dimension equivalence expressions
123 (e.g. "tract=4") to be recognized when they appear in simple contexts
124 (surrounded only by ANDs and ORs). When `InspectionVisitor` is used on its
125 own (i.e. when ``check=False`` in the query code), these don't do anything,
126 but they don't cost much, either. They are used by `CheckVisitor` when it
127 delegates to `InspectionVisitor` to see what governor dimension values are
128 set in a branch of the normal-form expression.
129 """
131 def merge(self, other: TreeSummary, isEq: bool = False) -> TreeSummary:
132 """Merge ``other`` into ``self``, making ``self`` a summary of both
133 expression tree branches.
135 Parameters
136 ----------
137 other : `TreeSummary`
138 The other summary object.
139 isEq : `bool`, optional
140 If `True` (`False` is default), these summaries are being combined
141 via the equality operator.
143 Returns
144 -------
145 self : `TreeSummary`
146 The merged summary (updated in-place).
147 """
148 self.update(other)
149 if isEq and self.isDataIdKeyOnly() and other.isDataIdValueOnly():
150 self.dataIdValue = other.dataIdValue
151 elif isEq and self.isDataIdValueOnly() and other.isDataIdKeyOnly():
152 self.dataIdKey = other.dataIdKey
153 else:
154 self.dataIdKey = None
155 self.dataIdValue = None
156 return self
158 def isDataIdKeyOnly(self) -> bool:
159 """Test whether this branch is _just_ a data ID key identifier."""
160 return self.dataIdKey is not None and self.dataIdValue is None
162 def isDataIdValueOnly(self) -> bool:
163 """Test whether this branch is _just_ a literal value that may be
164 used as the value in a data ID key-value pair.
165 """
166 return self.dataIdKey is None and self.dataIdValue is not None
168 dataIdKey: Dimension | None = None
169 """A `Dimension` that is (if `dataIdValue` is not `None`) or may be
170 (if `dataIdValue` is `None`) fully identified by a literal value in this
171 branch.
172 """
174 dataIdValue: str | None = None
175 """A literal value that constrains (if `dataIdKey` is not `None`) or may
176 constrain (if `dataIdKey` is `None`) a dimension in this branch.
178 This is always a `str` or `None`, but it may need to be coerced to `int`
179 to reflect the actual user intent.
180 """
183class InspectionVisitor(TreeVisitor[TreeSummary]):
184 """Implements TreeVisitor to identify dimension elements that need
185 to be included in a query, prior to actually constructing a SQLAlchemy
186 WHERE clause from it.
188 Parameters
189 ----------
190 universe : `DimensionUniverse`
191 All known dimensions.
192 bind : `~collections.abc.Mapping` [ `str`, `object` ]
193 Mapping containing literal values that should be injected into the
194 query expression, keyed by the identifiers they replace.
195 """
197 def __init__(self, universe: DimensionUniverse, bind: Mapping[str, Any]):
198 self.universe = universe
199 self.bind = bind
201 def visitNumericLiteral(self, value: str, node: Node) -> TreeSummary:
202 # Docstring inherited from TreeVisitor.visitNumericLiteral
203 return TreeSummary(dataIdValue=value)
205 def visitStringLiteral(self, value: str, node: Node) -> TreeSummary:
206 # Docstring inherited from TreeVisitor.visitStringLiteral
207 return TreeSummary(dataIdValue=value)
209 def visitTimeLiteral(self, value: astropy.time.Time, node: Node) -> TreeSummary:
210 # Docstring inherited from TreeVisitor.visitTimeLiteral
211 return TreeSummary()
213 def visitUuidLiteral(self, value: UUID, node: Node) -> TreeSummary:
214 # Docstring inherited from TreeVisitor.visitUuidLiteral
215 return TreeSummary()
217 def visitIdentifier(self, name: str, node: Node) -> TreeSummary:
218 # Docstring inherited from TreeVisitor.visitIdentifier
219 if name in self.bind:
220 return self.visitBind(name, node)
222 constant = categorizeConstant(name)
223 if constant is ExpressionConstant.INGEST_DATE:
224 return TreeSummary(hasIngestDate=True)
225 elif constant is ExpressionConstant.NULL:
226 return TreeSummary()
227 assert constant is None, "Enum variant conditionals should be exhaustive."
228 element, column = categorizeElementId(self.universe, name)
229 if column is None:
230 assert isinstance(element, Dimension)
231 return TreeSummary(
232 dimensions=set(element.minimal_group.names),
233 dataIdKey=element,
234 )
235 else:
236 return TreeSummary(dimensions=set(element.minimal_group.names), columns={element.name: {column}})
238 def visitBind(self, name: str, node: Node) -> TreeSummary:
239 if name not in self.bind:
240 raise UserExpressionError(f"Name {name!r} is not in the bind map.")
241 value = self.bind[name]
242 match value:
243 case list() | tuple() | Set():
244 # This can happen on rhs of IN operator, if there is only one
245 # element in the list then take it.
246 match len(value):
247 case 1:
248 (value,) = value
249 return TreeSummary(dataIdValue=value)
250 case _:
251 return TreeSummary()
252 case _:
253 return TreeSummary(dataIdValue=value)
255 def visitUnaryOp(self, operator: str, operand: TreeSummary, node: Node) -> TreeSummary:
256 # Docstring inherited from TreeVisitor.visitUnaryOp
257 return operand
259 def visitBinaryOp(self, operator: str, lhs: TreeSummary, rhs: TreeSummary, node: Node) -> TreeSummary:
260 # Docstring inherited from TreeVisitor.visitBinaryOp
261 return lhs.merge(rhs, isEq=(operator == "="))
263 def visitIsIn(self, lhs: TreeSummary, values: list[TreeSummary], not_in: bool, node: Node) -> TreeSummary:
264 # Docstring inherited from TreeVisitor.visitIsIn
265 for v in values:
266 lhs.merge(v)
267 return lhs
269 def visitParens(self, expression: TreeSummary, node: Node) -> TreeSummary:
270 # Docstring inherited from TreeVisitor.visitParens
271 return expression
273 def visitTupleNode(self, items: tuple[TreeSummary, ...], node: Node) -> TreeSummary:
274 # Docstring inherited from base class
275 result = TreeSummary()
276 for i in items:
277 result.merge(i)
278 return result
280 def visitRangeLiteral(self, start: int, stop: int, stride: int | None, node: Node) -> TreeSummary:
281 # Docstring inherited from TreeVisitor.visitRangeLiteral
282 return TreeSummary()
284 def visitPointNode(self, ra: TreeSummary, dec: TreeSummary, node: Node) -> TreeSummary:
285 # Docstring inherited from base class
286 return TreeSummary()
288 def visitCircleNode(
289 self, ra: TreeSummary, dec: TreeSummary, radius: TreeSummary, node: Node
290 ) -> TreeSummary:
291 # Docstring inherited from base class
292 return TreeSummary()
294 def visitPolygonNode(self, vertices: list[tuple[TreeSummary, TreeSummary]], node: Node) -> TreeSummary:
295 # Docstring inherited from base class
296 return TreeSummary()
298 def visitRegionNode(self, pos: TreeSummary, node: Node) -> TreeSummary:
299 # Docstring inherited from base class
300 return TreeSummary()
302 def visitBoxNode(
303 self, ra: TreeSummary, dec: TreeSummary, width: TreeSummary, height: TreeSummary, node: Node
304 ) -> TreeSummary:
305 # Docstring inherited from base class
306 return TreeSummary()
308 def visitGlobNode(self, expression: TreeSummary, pattern: TreeSummary, node: Node) -> TreeSummary:
309 # Docstring inherited from base class
310 # pattern is a literal, but expression should refer to a dimension.
311 return expression
314@dataclasses.dataclass
315class InnerSummary(InspectionSummary):
316 """Result object used by `CheckVisitor` to gather referenced dimensions
317 and tables from an inner group of AND'd together expression branches, and
318 check them for consistency and completeness.
319 """
321 dimension_values: dict[str, DataIdValue] = dataclasses.field(default_factory=dict)
322 """Mapping containing the values of all dimensions that are equated with
323 literal values in this expression branch.
324 """
326 defaultsNeeded: set[str] = dataclasses.field(default_factory=set)
327 """Governor dimensions whose values are needed by the query, not provided
328 in the query itself, and present in the default data ID.
330 These should be added to the query's data ID when finalizing the WHERE
331 clause.
332 """
335@dataclasses.dataclass
336class OuterSummary(InspectionSummary):
337 """Result object used by `CheckVisitor` to gather referenced dimensions,
338 tables, and governor dimension values from the entire expression.
339 """
341 dimension_constraints: dict[str, set[DataIdValue]] = dataclasses.field(default_factory=dict)
342 """Mapping containing all values that appear in this expression for
343 dimensions relevant to the query.
345 Dimensions that are absent from this dict are not constrained by this
346 expression.
347 """
349 defaultsNeeded: set[str] = dataclasses.field(default_factory=set)
350 """Governor dimensions whose values are needed by the query, not provided
351 in the query itself, and present in the default data ID.
353 These should be added to the query's data ID when finalizing the WHERE
354 clause.
355 """
358class CheckVisitor(NormalFormVisitor[TreeSummary, InnerSummary, OuterSummary]):
359 """An implementation of `NormalFormVisitor` that identifies the dimensions
360 and tables that need to be included in a query while performing some checks
361 for completeness and consistency.
363 Parameters
364 ----------
365 dataId : `DataCoordinate`
366 Dimension values that are fully known in advance.
367 dimensions : `DimensionGroup`
368 The dimensions the query would include in the absence of this
369 expression.
370 bind : `~collections.abc.Mapping` [ `str`, `object` ]
371 Mapping containing literal values that should be injected into the
372 query expression, keyed by the identifiers they replace.
373 defaults : `DataCoordinate`
374 A data ID containing default for governor dimensions.
375 allow_orphans : `bool`, optional
376 If `True`, permit expressions to refer to dimensions without providing
377 a value for their governor dimensions (e.g. referring to a visit
378 without an instrument). Should be left to default to `False` in
379 essentially all new code.
380 """
382 def __init__(
383 self,
384 dataId: DataCoordinate,
385 dimensions: DimensionGroup,
386 bind: Mapping[str, Any],
387 defaults: DataCoordinate,
388 allow_orphans: bool = False,
389 ):
390 self.dataId = dataId
391 self.dimensions = dimensions
392 self.defaults = defaults
393 self._branchVisitor = InspectionVisitor(dataId.universe, bind)
394 self._allow_orphans = allow_orphans
396 @property
397 def universe(self) -> DimensionUniverse:
398 """Object that defines all dimensions."""
399 return self.dimensions.universe
401 def visitBranch(self, node: Node) -> TreeSummary:
402 # Docstring inherited from NormalFormVisitor.
403 return node.visit(self._branchVisitor)
405 def visitInner(self, branches: Sequence[TreeSummary], form: NormalForm) -> InnerSummary:
406 # Docstring inherited from NormalFormVisitor.
407 # Disjunctive normal form means inner branches are AND'd together...
408 assert form is NormalForm.DISJUNCTIVE
409 # ...and that means each branch we iterate over together below
410 # constrains the others, and they all need to be consistent. Moreover,
411 # because outer branches are OR'd together, we also know that if
412 # something is missing from one of these branches (like a governor
413 # dimension value like the instrument or skymap needed to interpret a
414 # visit or tract number), it really is missing, because there's no way
415 # some other inner branch can constraint it.
416 #
417 # That is, except the data ID the visitor was passed at construction;
418 # that's AND'd to the entire expression later, and thus it affects all
419 # branches. To take care of that, we add any governor values it
420 # contains to the summary in advance.
421 summary = InnerSummary()
422 summary.dimension_values.update(self.dataId.mapping)
423 # Finally, we loop over those branches.
424 for branch in branches:
425 # Update the sets of dimensions and columns we've seen anywhere in
426 # the expression in any context.
427 summary.update(branch)
428 # Test whether this branch has a form like '<dimension>=<value>'
429 # (or equivalent; categorizeIdentifier is smart enough to see that
430 # e.g. 'detector.id=4' is equivalent to 'detector=4'). If so,
431 # remember that we've constrained it on this branch to later make
432 # sure it's consistent with any other constraints on any other
433 # branches its AND'd with.
434 if branch.dataIdKey is not None and branch.dataIdValue is not None:
435 new_value = branch.dataIdKey.primaryKey.getPythonType()(branch.dataIdValue)
436 value = summary.dimension_values.setdefault(branch.dataIdKey.name, new_value)
437 if value != new_value:
438 # Expression says something like "instrument='HSC' AND
439 # instrument='DECam'", or data ID has one and expression
440 # has the other.
441 if branch.dataIdKey.name in self.dataId:
442 raise UserExpressionError(
443 f"Conflict between expression containing {branch.dataIdKey.name}={new_value!r} "
444 f"and data ID with {branch.dataIdKey.name}={value!r}."
445 )
446 else:
447 raise UserExpressionError(
448 f"Conflicting literal values for {branch.dataIdKey.name} in expression: "
449 f"{value!r} != {branch.dataIdValue!r}."
450 )
451 # Now that we know which governor values we've constrained, see if any
452 # are missing, i.e. if the expression contains something like "visit=X"
453 # without saying what instrument that visit corresponds to. This rules
454 # out a lot of accidents, but it also rules out possibly-legitimate
455 # multi-instrument queries like "visit.seeing < 0.7". But it's not
456 # unreasonable to ask the user to be explicit about the instruments
457 # they want to consider to work around this restriction, and that's
458 # what we do. Note that if someone does write an expression like
459 #
460 # (instrument='HSC' OR instrument='DECam') AND visit.seeing < 0.7
461 #
462 # then in disjunctive normal form that will become
463 #
464 # (instrument='HSC' AND visit.seeing < 0.7)
465 # OR (instrument='DECam' AND visit.seeing < 0.7)
466 #
467 # i.e. each instrument will get its own outer branch and the logic here
468 # still works (that sort of thing is why we convert to normal form,
469 # after all).
470 governorsNeededInBranch: set[str] = set()
471 for dimension in summary.dimensions:
472 governorsNeededInBranch.update(self.universe.dimensions[dimension].minimal_group.governors)
473 if not governorsNeededInBranch.issubset(summary.dimension_values.keys()):
474 missing = governorsNeededInBranch - summary.dimension_values.keys()
475 if missing <= self.defaults.dimensions.required:
476 summary.defaultsNeeded.update(missing)
477 elif not self._allow_orphans:
478 still_missing = missing - self.defaults.dimensions.names
479 raise UserExpressionError(
480 f"No value(s) for governor dimensions {still_missing} in expression "
481 "that references dependent dimensions. 'Governor' dimensions must always be specified "
482 "completely in either the query expression (via simple 'name=<value>' terms, not 'IN' "
483 "terms) or in a data ID passed to the query method."
484 )
485 return summary
487 def visitOuter(self, branches: Sequence[InnerSummary], form: NormalForm) -> OuterSummary:
488 # Docstring inherited from NormalFormVisitor.
489 # Disjunctive normal form means outer branches are OR'd together.
490 assert form is NormalForm.DISJUNCTIVE
491 summary = OuterSummary()
492 if branches:
493 # Iterate over branches in first pass to gather all dimensions and
494 # columns referenced. This aggregation is for the full query, so
495 # we don't care whether things are joined by AND or OR (or + or -,
496 # etc). Also gather the set of dimensions directly constrained or
497 # pulled from defaults in _all_ branches. This is the set we will
498 # be able to bound overall; any dimensions not referenced by even
499 # one branch could be unbounded.
500 dimensions_in_all_branches = set(self.universe.dimensions.names)
501 for branch in branches:
502 summary.update(branch)
503 summary.defaultsNeeded.update(branch.defaultsNeeded)
504 dimensions_in_all_branches.intersection_update(branch.dimension_values)
505 # Go back through and set up the dimension bounds.
506 summary.dimension_constraints.update(
507 {dimension: set() for dimension in dimensions_in_all_branches}
508 )
509 for dim in dimensions_in_all_branches:
510 for branch in branches:
511 summary.dimension_constraints[dim].add(branch.dimension_values[dim])
512 # See if we've referenced any dimensions that weren't in the original
513 # query graph; if so, we update that to include them. This is what
514 # lets a user say "tract=X" on the command line (well, "skymap=Y AND
515 # tract=X" - logic in visitInner checks for that) when running a task
516 # like ISR that has nothing to do with skymaps.
517 if not summary.dimensions.issubset(self.dimensions.names):
518 self.dimensions = self.universe.conform(summary.dimensions | self.dimensions.names)
519 for dimension, values in summary.dimension_constraints.items():
520 if dimension in summary.defaultsNeeded:
521 # One branch contained an explicit value for this dimension
522 # while another needed to refer to the default data ID.
523 # Even if these refer to the same value, that inconsistency
524 # probably indicates user error.
525 raise UserExpressionError(
526 f"Governor dimension {dimension} is explicitly "
527 f"constrained to {values} in one or more branches of "
528 "this query where expression, but is left to default "
529 f"to {self.defaults[dimension]!r} in another branch. "
530 "Defaults and explicit constraints cannot be mixed."
531 )
532 # If any default data ID values were needed, update self.dataId with
533 # them, and then update the governor restriction with them.
534 if summary.defaultsNeeded:
535 defaultsNeededGraph = self.universe.conform(summary.defaultsNeeded)
536 self.dataId = self.dataId.union(self.defaults.subset(defaultsNeededGraph))
537 for dimension in summary.defaultsNeeded:
538 summary.dimension_constraints[dimension] = {self.defaults[dimension]}
540 return summary