Coverage for python/lsst/daf/butler/registry/queries/expressions/_predicate.py: 15%
224 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:23 -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__ = ("ExpressionTypeError", "make_string_expression_predicate")
31import builtins
32import datetime
33import types
34import warnings
35from collections.abc import Mapping, Set
36from typing import Any, cast
37from uuid import UUID
39import astropy.time
40import astropy.utils.exceptions
42from lsst.daf.relation import (
43 ColumnContainer,
44 ColumnExpression,
45 ColumnExpressionSequence,
46 ColumnLiteral,
47 ColumnTag,
48 Predicate,
49 sql,
50)
52# We import the some modules rather than types within them because match syntax
53# uses qualified names with periods to distinguish literals from captures.
54from .... import _timespan, timespan_database_representation
55from ...._column_tags import DatasetColumnTag, DimensionKeyColumnTag, DimensionRecordColumnTag
56from ...._column_type_info import ColumnTypeInfo
57from ....dimensions import DataCoordinate, Dimension, DimensionGroup, DimensionUniverse
58from ..._exceptions import UserExpressionError, UserExpressionSyntaxError
59from .categorize import ExpressionConstant, categorizeConstant, categorizeElementId
60from .check import CheckVisitor
61from .normalForm import NormalForm, NormalFormExpression
62from .parser import Node, TreeVisitor, parse_expression
64# As of astropy 4.2, the erfa interface is shipped independently and
65# ErfaWarning is no longer an AstropyWarning
66try:
67 import erfa
68except ImportError:
69 erfa = None
72class ExpressionTypeError(TypeError):
73 """Exception raised when the types in a query expression are not
74 compatible with the operators or other syntax.
75 """
78def make_string_expression_predicate(
79 string: str,
80 dimensions: DimensionGroup,
81 *,
82 column_types: ColumnTypeInfo,
83 bind: Mapping[str, Any] | None = None,
84 data_id: DataCoordinate | None = None,
85 defaults: DataCoordinate | None = None,
86 dataset_type_name: str | None = None,
87 allow_orphans: bool = False,
88) -> tuple[Predicate | None, Mapping[str, Set[str]]]:
89 """Create a predicate by parsing and analyzing a string expression.
91 Parameters
92 ----------
93 string : `str`
94 String to parse.
95 dimensions : `DimensionGroup`
96 The dimensions the query would include in the absence of this WHERE
97 expression.
98 column_types : `ColumnTypeInfo`
99 Information about column types.
100 bind : `~collections.abc.Mapping` [ `str`, `Any` ], optional
101 Literal values referenced in the expression.
102 data_id : `DataCoordinate`, optional
103 A fully-expanded data ID identifying dimensions known in advance.
104 If not provided, will be set to an empty data ID.
105 ``dataId.hasRecords()`` must return `True`.
106 defaults : `DataCoordinate`, optional
107 A data ID containing default for governor dimensions. Ignored
108 unless ``check=True``.
109 dataset_type_name : `str` or `None`, optional
110 The name of the dataset type to assume for unqualified dataset
111 columns, or `None` if there are no such identifiers.
112 allow_orphans : `bool`, optional
113 If `True`, permit expressions to refer to dimensions without
114 providing a value for their governor dimensions (e.g. referring to
115 a visit without an instrument). Should be left to default to
116 `False` in essentially all new code.
118 Returns
119 -------
120 predicate : `lsst.daf.relation.colum_expressions.Predicate` or `None`
121 New predicate derived from the string expression, or `None` if the
122 string is empty.
123 governor_constraints : `~collections.abc.Mapping` [ `str` , \
124 `~collections.abc.Set` ]
125 Constraints on dimension values derived from the expression and data
126 ID.
127 """
128 governor_constraints: dict[str, Set[str]] = {}
129 if data_id is None:
130 data_id = DataCoordinate.make_empty(dimensions.universe)
131 if not string:
132 for dimension in data_id.dimensions.governors:
133 governor_constraints[dimension] = {cast(str, data_id[dimension])}
134 return None, governor_constraints
135 try:
136 tree = parse_expression(string)
137 except Exception as exc:
138 raise UserExpressionSyntaxError(f"Failed to parse user expression {string!r}.") from exc
139 assert tree is not None, "Should only happen if the string is empty, and that's handled above."
140 if bind is None:
141 bind = {}
142 if bind:
143 for identifier in bind:
144 if identifier in dimensions.universe.elements.names:
145 raise RuntimeError(f"Bind parameter key {identifier!r} conflicts with a dimension element.")
146 table, _, column = identifier.partition(".")
147 if column and table in dimensions.universe.elements.names:
148 raise RuntimeError(f"Bind parameter key {identifier!r} looks like a dimension column.")
149 if defaults is None:
150 defaults = DataCoordinate.make_empty(dimensions.universe)
151 # Convert the expression to disjunctive normal form (ORs of ANDs).
152 # That's potentially super expensive in the general case (where there's
153 # a ton of nesting of ANDs and ORs). That won't be the case for the
154 # expressions we expect, and we actually use disjunctive normal instead
155 # of conjunctive (i.e. ANDs of ORs) because I think the worst-case is
156 # a long list of OR'd-together data IDs, which is already in or very
157 # close to disjunctive normal form.
158 expr = NormalFormExpression.fromTree(tree, NormalForm.DISJUNCTIVE)
159 # Check the expression for consistency and completeness.
160 visitor = CheckVisitor(data_id, dimensions, bind, defaults, allow_orphans=allow_orphans)
161 try:
162 summary = expr.visit(visitor)
163 except UserExpressionError as err:
164 exprOriginal = str(tree)
165 exprNormal = str(expr.toTree())
166 if exprNormal == exprOriginal:
167 msg = f'Error in query expression "{exprOriginal}": {err}'
168 else:
169 msg = f'Error in query expression "{exprOriginal}" (normalized to "{exprNormal}"): {err}'
170 raise UserExpressionError(msg) from None
171 for dimension_name, values in summary.dimension_constraints.items():
172 if dimension_name in dimensions.universe.governor_dimensions.names:
173 governor_constraints[dimension_name] = cast(Set[str], values)
174 converter = PredicateConversionVisitor(bind, dataset_type_name, dimensions.universe, column_types)
175 predicate = tree.visit(converter)
176 return predicate, governor_constraints
179VisitorResult = Predicate | ColumnExpression | ColumnContainer
182class PredicateConversionVisitor(TreeVisitor[VisitorResult]):
183 def __init__(
184 self,
185 bind: Mapping[str, Any],
186 dataset_type_name: str | None,
187 universe: DimensionUniverse,
188 column_types: ColumnTypeInfo,
189 ):
190 self.bind = bind
191 self.dataset_type_name = dataset_type_name
192 self.universe = universe
193 self.column_types = column_types
195 OPERATOR_MAP = {
196 "=": "__eq__",
197 "!=": "__ne__",
198 "<": "__lt__",
199 ">": "__gt__",
200 "<=": "__le__",
201 ">=": "__ge__",
202 "+": "__add__",
203 "-": "__sub__",
204 "*": "__mul__",
205 }
207 def to_datetime(self, time: astropy.time.Time) -> datetime.datetime:
208 with warnings.catch_warnings():
209 warnings.simplefilter("ignore", category=astropy.utils.exceptions.AstropyWarning)
210 if erfa is not None:
211 warnings.simplefilter("ignore", category=erfa.ErfaWarning)
212 return time.to_datetime()
214 def visitBinaryOp(
215 self, operator: str, lhs: VisitorResult, rhs: VisitorResult, node: Node
216 ) -> VisitorResult:
217 # Docstring inherited.
218 b = builtins
219 match (operator, lhs, rhs):
220 case ["OR", Predicate() as lhs, Predicate() as rhs]:
221 return lhs.logical_or(rhs)
222 case ["AND", Predicate() as lhs, Predicate() as rhs]:
223 return lhs.logical_and(rhs)
224 # Allow all comparisons between expressions of the same type for
225 # sortable types.
226 case [
227 "=" | "!=" | "<" | ">" | "<=" | ">=",
228 ColumnExpression(
229 dtype=b.int | b.float | b.str | astropy.time.Time | datetime.datetime
230 ) as lhs,
231 ColumnExpression() as rhs,
232 ] if lhs.dtype is rhs.dtype:
233 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
234 # Allow comparisons between datetime expressions and
235 # astropy.time.Time literals/binds (only), by coercing the
236 # astropy.time.Time version to datetime.
237 case [
238 "=" | "!=" | "<" | ">" | "<=" | ">=",
239 ColumnLiteral(dtype=astropy.time.Time) as lhs,
240 ColumnExpression(dtype=datetime.datetime) as rhs,
241 ]:
242 lhs = ColumnLiteral(self.to_datetime(lhs.value), datetime.datetime)
243 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
244 case [
245 "=" | "!=" | "<" | ">" | "<=" | ">=",
246 ColumnExpression(dtype=datetime.datetime) as lhs,
247 ColumnLiteral(dtype=astropy.time.Time) as rhs,
248 ]:
249 rhs = ColumnLiteral(self.to_datetime(rhs.value), datetime.datetime)
250 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
251 # Allow comparisons between astropy.time.Time expressions and
252 # datetime literals/binds, by coercing the
253 # datetime literals to astropy.time.Time (in UTC scale).
254 case [
255 "=" | "!=" | "<" | ">" | "<=" | ">=",
256 ColumnLiteral(dtype=datetime.datetime) as lhs,
257 ColumnExpression(dtype=astropy.time.Time) as rhs,
258 ]:
259 lhs = ColumnLiteral(astropy.time.Time(lhs.value, scale="utc"), astropy.time.Time)
260 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
261 case [
262 "=" | "!=" | "<" | ">" | "<=" | ">=",
263 ColumnExpression(dtype=astropy.time.Time) as lhs,
264 ColumnLiteral(dtype=datetime.datetime) as rhs,
265 ]:
266 rhs = ColumnLiteral(astropy.time.Time(rhs.value, scale="utc"), astropy.time.Time)
267 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
268 # Allow equality comparisons with None/NULL. We don't have an 'IS'
269 # operator.
270 case ["=" | "!=", ColumnExpression(dtype=types.NoneType) as lhs, ColumnExpression() as rhs]:
271 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
272 case ["=" | "!=", ColumnExpression() as lhs, ColumnExpression(dtype=types.NoneType) as rhs]:
273 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
274 # Comparisions between Time and Timespan need have the Timespan on
275 # the lhs, since that (actually TimespanDatabaseRepresentation) is
276 # what actually has the methods.
277 case [
278 "<",
279 ColumnExpression(dtype=astropy.time.Time) as lhs,
280 ColumnExpression(dtype=_timespan.Timespan) as rhs,
281 ]:
282 return rhs.predicate_method(self.OPERATOR_MAP[">"], lhs)
283 case [
284 ">",
285 ColumnExpression(dtype=astropy.time.Time) as lhs,
286 ColumnExpression(dtype=_timespan.Timespan) as rhs,
287 ]:
288 return rhs.predicate_method(self.OPERATOR_MAP["<"], lhs)
289 # Enable other comparisons between times and Timespans (many of the
290 # combinations matched by this branch will have already been
291 # covered by a preceding branch).
292 case [
293 "<" | ">",
294 ColumnExpression(dtype=_timespan.Timespan | astropy.time.Time) as lhs,
295 ColumnExpression(dtype=_timespan.Timespan | astropy.time.Time) as rhs,
296 ]:
297 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
298 # Enable "overlaps" operations between timespans, and between times
299 # and timespans. The latter resolve to the `Timespan.contains` or
300 # `TimespanDatabaseRepresentation.contains` methods, but we use
301 # OVERLAPS in the string expression language to keep that simple.
302 case [
303 "OVERLAPS",
304 ColumnExpression(dtype=_timespan.Timespan) as lhs,
305 ColumnExpression(dtype=_timespan.Timespan) as rhs,
306 ]:
307 return lhs.predicate_method("overlaps", rhs)
308 case [
309 "OVERLAPS",
310 ColumnExpression(dtype=_timespan.Timespan) as lhs,
311 ColumnExpression(dtype=astropy.time.Time) as rhs,
312 ]:
313 return lhs.predicate_method("overlaps", rhs)
314 case [
315 "OVERLAPS",
316 ColumnExpression(dtype=astropy.time.Time) as lhs,
317 ColumnExpression(dtype=_timespan.Timespan) as rhs,
318 ]:
319 return rhs.predicate_method("overlaps", lhs)
320 # Enable arithmetic operators on numeric types, without any type
321 # coercion or broadening.
322 case [
323 "+" | "-" | "*",
324 ColumnExpression(dtype=b.int | b.float) as lhs,
325 ColumnExpression() as rhs,
326 ] if lhs.dtype is rhs.dtype:
327 return lhs.method(self.OPERATOR_MAP[operator], rhs, dtype=lhs.dtype)
328 case ["/", ColumnExpression(dtype=b.float) as lhs, ColumnExpression(dtype=b.float) as rhs]:
329 return lhs.method("__truediv__", rhs, dtype=b.float)
330 case ["/", ColumnExpression(dtype=b.int) as lhs, ColumnExpression(dtype=b.int) as rhs]:
331 # SQLAlchemy maps Python's '/' (__truediv__) operator directly
332 # to SQL's '/', despite those being defined differently for
333 # integers. Our expression language uses the SQL definition,
334 # and we only care about these expressions being evaluated in
335 # SQL right now, but we still want to guard against it being
336 # evaluated in Python and producing a surprising answer, so we
337 # mark it as being supported only by a SQL engine.
338 return lhs.method(
339 "__truediv__",
340 rhs,
341 dtype=b.int,
342 supporting_engine_types={sql.Engine},
343 )
344 case ["%", ColumnExpression(dtype=b.int) as lhs, ColumnExpression(dtype=b.int) as rhs]:
345 return lhs.method("__mod__", rhs, dtype=b.int)
346 assert lhs.dtype is not None and rhs.dtype is not None, (
347 "Expression converter should not yield untyped nodes."
348 )
349 raise ExpressionTypeError(
350 f"Invalid types {lhs.dtype.__name__}, {rhs.dtype.__name__} for binary operator {operator!r} "
351 f"in expression {node!s}."
352 )
354 def visitIdentifier(self, name: str, node: Node) -> VisitorResult:
355 # Docstring inherited.
356 if name in self.bind:
357 return self.visitBind(name, node)
358 tag: ColumnTag
359 match categorizeConstant(name):
360 case ExpressionConstant.INGEST_DATE:
361 assert self.dataset_type_name is not None
362 tag = DatasetColumnTag(self.dataset_type_name, "ingest_date")
363 return ColumnExpression.reference(tag, self.column_types.ingest_date_pytype)
364 case ExpressionConstant.NULL:
365 return ColumnExpression.literal(None, type(None))
366 case None:
367 pass
368 case _:
369 raise AssertionError("Check for enum values should be exhaustive.")
370 element, column = categorizeElementId(self.universe, name)
371 if column is not None:
372 tag = DimensionRecordColumnTag(element.name, column)
373 dtype = (
374 _timespan.Timespan
375 if column == timespan_database_representation.TimespanDatabaseRepresentation.NAME
376 else element.RecordClass.fields.standard[column].getPythonType()
377 )
378 if dtype is bool:
379 # ColumnExpression is for non-boolean columns only. Booleans
380 # are represented as Predicate.
381 return Predicate.reference(tag)
382 else:
383 return ColumnExpression.reference(tag, dtype)
384 else:
385 tag = DimensionKeyColumnTag(element.name)
386 assert isinstance(element, Dimension)
387 return ColumnExpression.reference(tag, element.primaryKey.getPythonType())
389 def visitBind(self, name: str, node: Node) -> VisitorResult:
390 # Docstring inherited.
391 if name not in self.bind:
392 raise UserExpressionError(f"Name {name!r} is not in the bind map.")
394 value = self.bind[name]
395 if isinstance(value, list | tuple | Set):
396 elements = []
397 all_dtypes = set()
398 for item in value:
399 dtype = type(item)
400 all_dtypes.add(dtype)
401 elements.append(ColumnExpression.literal(item, dtype=dtype))
402 if len(all_dtypes) > 1:
403 raise ExpressionTypeError(
404 f"Mismatched types in bind iterable: {value} has a mix of {all_dtypes}."
405 )
406 elif not elements:
407 # Empty container
408 return ColumnContainer.sequence([])
409 else:
410 (dtype,) = all_dtypes
411 return ColumnContainer.sequence(elements, dtype=dtype)
412 return ColumnExpression.literal(value, dtype=type(value))
414 def visitIsIn(
415 self, lhs: VisitorResult, values: list[VisitorResult], not_in: bool, node: Node
416 ) -> VisitorResult:
417 # Docstring inherited.
418 clauses: list[Predicate] = []
419 items: list[ColumnExpression] = []
420 assert isinstance(lhs, ColumnExpression), "LHS of IN guaranteed to be scalar by parser."
421 for rhs_item in values:
422 match rhs_item:
423 case ColumnExpressionSequence(items=rhs_items, dtype=rhs_dtype) if (
424 rhs_dtype is None or rhs_dtype == lhs.dtype
425 ):
426 items.extend(rhs_items)
427 case ColumnContainer(dtype=lhs.dtype):
428 clauses.append(rhs_item.contains(lhs))
429 case ColumnExpression(dtype=lhs.dtype):
430 items.append(rhs_item)
431 case _:
432 raise ExpressionTypeError(
433 f"Invalid type {rhs_item.dtype} for element in {lhs.dtype} IN expression '{node}'."
434 )
435 if items:
436 clauses.append(ColumnContainer.sequence(items, dtype=lhs.dtype).contains(lhs))
437 result = Predicate.logical_or(*clauses)
438 if not_in:
439 result = result.logical_not()
440 return result
442 def visitNumericLiteral(self, value: str, node: Node) -> VisitorResult:
443 # Docstring inherited.
444 try:
445 return ColumnExpression.literal(int(value), dtype=int)
446 except ValueError:
447 return ColumnExpression.literal(float(value), dtype=float)
449 def visitParens(self, expression: VisitorResult, node: Node) -> VisitorResult:
450 # Docstring inherited.
451 return expression
453 def visitPointNode(self, ra: VisitorResult, dec: VisitorResult, node: Node) -> VisitorResult:
454 # Docstring inherited.
456 # this is a placeholder for future extension, we enabled syntax but
457 # do not support actual use just yet.
458 raise NotImplementedError("POINT() function is not supported yet")
460 def visitCircleNode(
461 self, ra: VisitorResult, dec: VisitorResult, radius: VisitorResult, node: Node
462 ) -> VisitorResult:
463 # Docstring inherited.
465 # this is a placeholder for future extension, we enabled syntax but
466 # do not support actual use just yet.
467 raise NotImplementedError("CIRCLE() function is not supported yet")
469 def visitBoxNode(
470 self, ra: VisitorResult, dec: VisitorResult, width: VisitorResult, height: VisitorResult, node: Node
471 ) -> VisitorResult:
472 # Docstring inherited.
474 # this is a placeholder for future extension, we enabled syntax but
475 # do not support actual use just yet.
476 raise NotImplementedError("BOX() function is not supported yet")
478 def visitPolygonNode(
479 self, vertices: list[tuple[VisitorResult, VisitorResult]], node: Node
480 ) -> VisitorResult:
481 # Docstring inherited.
483 # this is a placeholder for future extension, we enabled syntax but
484 # do not support actual use just yet.
485 raise NotImplementedError("POLYGON() function is not supported yet")
487 def visitRegionNode(self, pos: VisitorResult, node: Node) -> VisitorResult:
488 # Docstring inherited.
490 # this is a placeholder for future extension, we enabled syntax but
491 # do not support actual use just yet.
492 raise NotImplementedError("REGION() function is not supported yet")
494 def visitRangeLiteral(self, start: int, stop: int, stride: int | None, node: Node) -> VisitorResult:
495 # Docstring inherited.
496 return ColumnContainer.range_literal(range(start, stop + 1, stride or 1))
498 def visitStringLiteral(self, value: str, node: Node) -> VisitorResult:
499 # Docstring inherited.
500 return ColumnExpression.literal(value, dtype=str)
502 def visitTimeLiteral(self, value: astropy.time.Time, node: Node) -> VisitorResult:
503 # Docstring inherited.
504 return ColumnExpression.literal(value, dtype=astropy.time.Time)
506 def visitUuidLiteral(self, value: UUID, node: Node) -> VisitorResult:
507 # Docstring inherited.
508 return ColumnExpression.literal(value, dtype=UUID)
510 def visitTupleNode(self, items: tuple[VisitorResult, ...], node: Node) -> VisitorResult:
511 # Docstring inherited.
512 match items:
513 case [
514 ColumnLiteral(value=begin, dtype=astropy.time.Time | types.NoneType),
515 ColumnLiteral(value=end, dtype=astropy.time.Time | types.NoneType),
516 ]:
517 return ColumnExpression.literal(_timespan.Timespan(begin, end), dtype=_timespan.Timespan)
518 raise ExpressionTypeError(
519 f'Invalid type(s) ({items[0].dtype}, {items[1].dtype}) in timespan tuple "{node}" '
520 '(Note that date/time strings must be preceded by "T" to be recognized).'
521 )
523 def visitUnaryOp(self, operator: str, operand: VisitorResult, node: Node) -> VisitorResult:
524 # Docstring inherited.
525 match (operator, operand):
526 case ["NOT", Predicate() as operand]:
527 return operand.logical_not()
528 case ["+", ColumnExpression(dtype=builtins.int | builtins.float) as operand]:
529 return operand.method("__pos__")
530 case ["-", ColumnExpression(dtype=builtins.int | builtins.float) as operand]:
531 return operand.method("__neg__")
532 raise ExpressionTypeError(
533 f"Unary operator {operator!r} is not valid for operand of type {operand.dtype!s} in {node!s}."
534 )
536 def visitGlobNode(self, expression: VisitorResult, pattern: VisitorResult, node: Node) -> VisitorResult:
537 # Docstring inherited.
538 raise NotImplementedError("GLOB() function is not supported yet")