Coverage for python / lsst / daf / butler / queries / predicate_constraints_summary.py: 24%
81 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-28 08:36 +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
30from typing import Any
32from .._exceptions import InvalidQueryError
33from ..dimensions import DataCoordinate, DataIdValue, DimensionGroup, DimensionUniverse
34from . import tree as qt
35from .visitors import ColumnExpressionVisitor, PredicateVisitFlags, SimplePredicateVisitor
38class PredicateConstraintsSummary:
39 """Summarizes information about the constraints on data ID values implied
40 by a Predicate.
42 Parameters
43 ----------
44 predicate : `Predicate`
45 Predicate to summarize.
46 """
48 predicate: qt.Predicate
49 """The predicate examined by this summary."""
51 constraint_data_id: dict[str, DataIdValue]
52 """Data ID values that will be identical in all result rows due to query
53 constraints.
54 """
56 messages: list[str]
57 """Diagnostic messages that report reasons the query may not return any
58 rows.
59 """
61 def __init__(self, predicate: qt.Predicate) -> None:
62 self.predicate = predicate
63 self.constraint_data_id = {}
64 self.messages = []
65 # Governor dimensions referenced directly in the predicate, but not
66 # necessarily constrained to the same value in all logic branches.
67 self._governors_referenced: set[str] = set()
69 self.predicate.visit(
70 _DataIdExtractionVisitor(self.constraint_data_id, self.messages, self._governors_referenced)
71 )
73 def apply_default_data_id(
74 self,
75 default_data_id: DataCoordinate,
76 query_dimensions: DimensionGroup,
77 validate_governor_constraints: bool,
78 ) -> None:
79 """Augment the predicate and summary by adding missing constraints for
80 governor dimensions using a default data ID.
82 Parameters
83 ----------
84 default_data_id : `DataCoordinate`
85 Data ID values that will be used to constrain the query if governor
86 dimensions have not already been constrained by the predicate.
87 query_dimensions : `DimensionGroup`
88 The set of dimensions returned in result rows from the query.
89 validate_governor_constraints : `bool`
90 If `True`, enforce the requirement that governor dimensions must be
91 constrained if any dimensions that depend on them have constraints.
92 """
93 # Find governor dimensions required by the predicate.
94 # If these are not constrained by the predicate or the default data ID,
95 # we will raise an exception.
96 where_governors: set[str] = set()
97 self.predicate.gather_governors(where_governors)
99 # Add in governor dimensions that are returned in result rows.
100 # We constrain these using a default data ID if one is available,
101 # but it's not an error to omit the constraint.
102 governors_used_by_query = where_governors | query_dimensions.governors
104 # For each governor dimension needed by the query, add a constraint
105 # from the default data ID if the existing predicate does not
106 # constrain it.
107 for governor in governors_used_by_query:
108 if governor not in self.constraint_data_id and governor not in self._governors_referenced:
109 if governor in default_data_id.dimensions:
110 data_id_value = default_data_id[governor]
111 self.constraint_data_id[governor] = data_id_value
112 self._governors_referenced.add(governor)
113 self.predicate = self.predicate.logical_and(
114 _create_data_id_predicate(governor, data_id_value, query_dimensions.universe)
115 )
116 elif governor in where_governors and validate_governor_constraints:
117 # Check that the predicate doesn't reference any dimensions
118 # without constraining their governor dimensions, since
119 # that's a particularly easy mistake to make and it's
120 # almost never intentional.
121 raise InvalidQueryError(
122 f"Query 'where' expression references a dimension dependent on {governor} without "
123 "constraining it directly."
124 )
127def _create_data_id_predicate(
128 dimension_name: str, value: DataIdValue, universe: DimensionUniverse
129) -> qt.Predicate:
130 """Create a Predicate that tests whether the given dimension primary key is
131 equal to the given literal value.
132 """
133 dimension = universe.dimensions[dimension_name]
134 return qt.Predicate.compare(
135 qt.DimensionKeyReference(dimension=dimension), "==", qt.make_column_literal(value)
136 )
139class _DataIdExtractionVisitor(
140 SimplePredicateVisitor,
141 ColumnExpressionVisitor[tuple[str, None] | tuple[None, Any] | tuple[None, None]],
142):
143 """A column-expression visitor that extracts quality constraints on
144 dimensions that are not OR'd with anything else.
146 Parameters
147 ----------
148 data_id : `dict`
149 Dictionary to populate in place.
150 messages : `list` [ `str` ]
151 List of diagnostic messages to populate in place.
152 governor_references : `set` [ `str` ]
153 Set of the names of governor dimension names that were referenced
154 directly. This includes dimensions that were constrained to different
155 values in different logic branches, and hence not included in
156 ``data_id``.
157 """
159 def __init__(self, data_id: dict[str, DataIdValue], messages: list[str], governor_references: set[str]):
160 self.data_id = data_id
161 self.messages = messages
162 self.governor_references = governor_references
164 def visit_comparison(
165 self,
166 a: qt.ColumnExpression,
167 operator: qt.ComparisonOperator,
168 b: qt.ColumnExpression,
169 flags: PredicateVisitFlags,
170 ) -> None:
171 k_a, v_a = a.visit(self)
172 k_b, v_b = b.visit(self)
173 if flags & PredicateVisitFlags.HAS_OR_SIBLINGS:
174 return None
175 if flags & PredicateVisitFlags.INVERTED:
176 if operator == "!=":
177 operator = "=="
178 else:
179 return None
180 if operator != "==":
181 return None
182 if k_a is not None and v_b is not None:
183 key = k_a
184 value = v_b
185 elif k_b is not None and v_a is not None:
186 key = k_b
187 value = v_a
188 else:
189 return None
190 if (old := self.data_id.setdefault(key, value)) != value:
191 self.messages.append(f"'where' expression requires both {key}={value!r} and {key}={old!r}.")
192 return None
194 def visit_binary_expression(self, expression: qt.BinaryExpression) -> tuple[None, None]:
195 expression.a.visit(self)
196 expression.b.visit(self)
197 return None, None
199 def visit_unary_expression(self, expression: qt.UnaryExpression) -> tuple[None, None]:
200 expression.operand.visit(self)
201 return None, None
203 def visit_literal(self, expression: qt.ColumnLiteral) -> tuple[None, Any]:
204 return None, expression.get_literal_value()
206 def visit_dimension_key_reference(self, expression: qt.DimensionKeyReference) -> tuple[str, None]:
207 if expression.dimension.governor is expression.dimension:
208 self.governor_references.add(expression.dimension.name)
209 return expression.dimension.name, None
211 def visit_dimension_field_reference(self, expression: qt.DimensionFieldReference) -> tuple[None, None]:
212 if (
213 expression.element.governor is expression.element
214 and expression.field in expression.element.alternate_keys.names
215 ):
216 self.governor_references.add(expression.element.name)
217 return None, None
219 def visit_dataset_field_reference(self, expression: qt.DatasetFieldReference) -> tuple[None, None]:
220 return None, None
222 def visit_reversed(self, expression: qt.Reversed) -> tuple[None, None]:
223 raise AssertionError("No Reversed expressions in predicates.")