Coverage for python/lsst/daf/butler/registry/queries/expressions/_predicate.py: 12%
207 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-04-07 00:58 -0700
« prev ^ index » next coverage.py v6.5.0, created at 2023-04-07 00:58 -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 program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
21from __future__ import annotations
23__all__ = ("make_string_expression_predicate", "ExpressionTypeError")
25import builtins
26import datetime
27import types
28import warnings
29from collections.abc import Mapping, Set
30from typing import Any, Union, cast
32import astropy.time
33import astropy.utils.exceptions
34from lsst.daf.relation import (
35 ColumnContainer,
36 ColumnExpression,
37 ColumnExpressionSequence,
38 ColumnLiteral,
39 ColumnTag,
40 Predicate,
41 sql,
42)
44# We import the timespan module rather than types within it because match
45# syntax uses qualified names with periods to distinguish literals from
46# captures.
47from ....core import (
48 ColumnTypeInfo,
49 DataCoordinate,
50 DatasetColumnTag,
51 Dimension,
52 DimensionGraph,
53 DimensionKeyColumnTag,
54 DimensionRecordColumnTag,
55 DimensionUniverse,
56 timespan,
57)
58from ..._exceptions import UserExpressionError, UserExpressionSyntaxError
59from .categorize import ExpressionConstant, categorizeConstant, categorizeElementId
60from .check import CheckVisitor
61from .normalForm import NormalForm, NormalFormExpression
62from .parser import Node, ParserYacc, TreeVisitor # type: ignore
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: DimensionGraph,
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 : `DimensionGraph`
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 : `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 : `Mapping` [ `str` , `~collections.abc.Set` ]
124 Constraints on dimension values derived from the expression and data
125 ID.
126 """
127 governor_constraints: dict[str, Set[str]] = {}
128 if data_id is None:
129 data_id = DataCoordinate.makeEmpty(dimensions.universe)
130 if not string:
131 for dimension in data_id.graph.governors:
132 governor_constraints[dimension.name] = {cast(str, data_id[dimension])}
133 return None, governor_constraints
134 try:
135 parser = ParserYacc()
136 tree = parser.parse(string)
137 except Exception as exc:
138 raise UserExpressionSyntaxError(f"Failed to parse user expression {string!r}.") from exc
139 if bind is None:
140 bind = {}
141 if bind:
142 for identifier in bind:
143 if identifier in dimensions.universe.getStaticElements().names:
144 raise RuntimeError(f"Bind parameter key {identifier!r} conflicts with a dimension element.")
145 table, _, column = identifier.partition(".")
146 if column and table in dimensions.universe.getStaticElements().names:
147 raise RuntimeError(f"Bind parameter key {identifier!r} looks like a dimension column.")
148 if defaults is None:
149 defaults = DataCoordinate.makeEmpty(dimensions.universe)
150 # Convert the expression to disjunctive normal form (ORs of ANDs).
151 # That's potentially super expensive in the general case (where there's
152 # a ton of nesting of ANDs and ORs). That won't be the case for the
153 # expressions we expect, and we actually use disjunctive normal instead
154 # of conjunctive (i.e. ANDs of ORs) because I think the worst-case is
155 # a long list of OR'd-together data IDs, which is already in or very
156 # close to disjunctive normal form.
157 expr = NormalFormExpression.fromTree(tree, NormalForm.DISJUNCTIVE)
158 # Check the expression for consistency and completeness.
159 visitor = CheckVisitor(data_id, dimensions, bind, defaults, allow_orphans=allow_orphans)
160 try:
161 summary = expr.visit(visitor)
162 except UserExpressionError as err:
163 exprOriginal = str(tree)
164 exprNormal = str(expr.toTree())
165 if exprNormal == exprOriginal:
166 msg = f'Error in query expression "{exprOriginal}": {err}'
167 else:
168 msg = f'Error in query expression "{exprOriginal}" (normalized to "{exprNormal}"): {err}'
169 raise UserExpressionError(msg) from None
170 for dimension_name, values in summary.dimension_constraints.items():
171 if dimension_name in dimensions.universe.getGovernorDimensions().names:
172 governor_constraints[dimension_name] = cast(Set[str], values)
173 converter = PredicateConversionVisitor(bind, dataset_type_name, dimensions.universe, column_types)
174 predicate = tree.visit(converter)
175 return predicate, governor_constraints
178VisitorResult = Union[Predicate, ColumnExpression, ColumnContainer]
181class PredicateConversionVisitor(TreeVisitor[VisitorResult]):
182 def __init__(
183 self,
184 bind: Mapping[str, Any],
185 dataset_type_name: str | None,
186 universe: DimensionUniverse,
187 column_types: ColumnTypeInfo,
188 ):
189 self.bind = bind
190 self.dataset_type_name = dataset_type_name
191 self.universe = universe
192 self.column_types = column_types
194 OPERATOR_MAP = {
195 "=": "__eq__",
196 "!=": "__ne__",
197 "<": "__lt__",
198 ">": "__gt__",
199 "<=": "__le__",
200 ">=": "__ge__",
201 "+": "__add__",
202 "-": "__sub__",
203 "/": "__mul__",
204 }
206 def to_datetime(self, time: astropy.time.Time) -> datetime.datetime:
207 with warnings.catch_warnings():
208 warnings.simplefilter("ignore", category=astropy.utils.exceptions.AstropyWarning)
209 if erfa is not None:
210 warnings.simplefilter("ignore", category=erfa.ErfaWarning)
211 return time.to_datetime()
213 def visitBinaryOp(
214 self, operator: str, lhs: VisitorResult, rhs: VisitorResult, node: Node
215 ) -> VisitorResult:
216 # Docstring inherited.
217 b = builtins
218 match (operator, lhs, rhs):
219 case ["OR", Predicate() as lhs, Predicate() as rhs]:
220 return lhs.logical_or(rhs)
221 case ["AND", Predicate() as lhs, Predicate() as rhs]:
222 return lhs.logical_and(rhs)
223 # Allow all comparisons between expressions of the same type for
224 # sortable types.
225 case [
226 "=" | "!=" | "<" | ">" | "<=" | ">=",
227 ColumnExpression(
228 dtype=b.int | b.float | b.str | astropy.time.Time | datetime.datetime
229 ) as lhs,
230 ColumnExpression() as rhs,
231 ] if lhs.dtype is rhs.dtype:
232 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
233 # Allow comparisons between datetime expressions and
234 # astropy.time.Time literals/binds (only), by coercing the
235 # astropy.time.Time version to datetime.
236 case [
237 "=" | "!=" | "<" | ">" | "<=" | ">=",
238 ColumnLiteral(dtype=astropy.time.Time) as lhs,
239 ColumnExpression(dtype=datetime.datetime) as rhs,
240 ]:
241 lhs = ColumnLiteral(self.to_datetime(lhs.value), datetime.datetime)
242 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
243 case [
244 "=" | "!=" | "<" | ">" | "<=" | ">=",
245 ColumnExpression(dtype=datetime.datetime) as lhs,
246 ColumnLiteral(dtype=astropy.time.Time) as rhs,
247 ]:
248 rhs = ColumnLiteral(self.to_datetime(rhs.value), datetime.datetime)
249 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
250 # Allow comparisons between astropy.time.Time expressions and
251 # datetime literals/binds, by coercing the
252 # datetime literals to astropy.time.Time (in UTC scale).
253 case [
254 "=" | "!=" | "<" | ">" | "<=" | ">=",
255 ColumnLiteral(dtype=datetime.datetime) as lhs,
256 ColumnExpression(dtype=astropy.time.Time) as rhs,
257 ]:
258 lhs = ColumnLiteral(astropy.time.Time(lhs.value, scale="utc"), astropy.time.Time)
259 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
260 case [
261 "=" | "!=" | "<" | ">" | "<=" | ">=",
262 ColumnExpression(dtype=astropy.time.Time) as lhs,
263 ColumnLiteral(dtype=datetime.datetime) as rhs,
264 ]:
265 rhs = ColumnLiteral(astropy.time.Time(rhs.value, scale="utc"), astropy.time.Time)
266 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
267 # Allow equality comparisons with None/NULL. We don't have an 'IS'
268 # operator.
269 case ["=" | "!=", ColumnExpression(dtype=types.NoneType) as lhs, ColumnExpression() as rhs]:
270 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
271 case ["=" | "!=", ColumnExpression() as lhs, ColumnExpression(dtype=types.NoneType) as rhs]:
272 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
273 # Comparisions between Time and Timespan need have the Timespan on
274 # the lhs, since that (actually TimespanDatabaseRepresentation) is
275 # what actually has the methods.
276 case [
277 "<",
278 ColumnExpression(dtype=astropy.time.Time) as lhs,
279 ColumnExpression(dtype=timespan.Timespan) as rhs,
280 ]:
281 return rhs.predicate_method(self.OPERATOR_MAP[">"], lhs)
282 case [
283 ">",
284 ColumnExpression(dtype=astropy.time.Time) as lhs,
285 ColumnExpression(dtype=timespan.Timespan) as rhs,
286 ]:
287 return rhs.predicate_method(self.OPERATOR_MAP["<"], lhs)
288 # Enable other comparisons between times and Timespans (many of the
289 # combinations matched by this branch will have already been
290 # covered by a preceding branch).
291 case [
292 "<" | ">",
293 ColumnExpression(dtype=timespan.Timespan | astropy.time.Time) as lhs,
294 ColumnExpression(dtype=timespan.Timespan | astropy.time.Time) as rhs,
295 ]:
296 return lhs.predicate_method(self.OPERATOR_MAP[operator], rhs)
297 # Enable "overlaps" operations between timespans, and between times
298 # and timespans. The latter resolve to the `Timespan.contains` or
299 # `TimespanDatabaseRepresentation.contains` methods, but we use
300 # OVERLAPS in the string expression language to keep that simple.
301 case [
302 "OVERLAPS",
303 ColumnExpression(dtype=timespan.Timespan) as lhs,
304 ColumnExpression(dtype=timespan.Timespan) as rhs,
305 ]:
306 return lhs.predicate_method("overlaps", rhs)
307 case [
308 "OVERLAPS",
309 ColumnExpression(dtype=timespan.Timespan) as lhs,
310 ColumnExpression(dtype=astropy.time.Time) as rhs,
311 ]:
312 return lhs.predicate_method("overlaps", rhs)
313 case [
314 "OVERLAPS",
315 ColumnExpression(dtype=astropy.time.Time) as lhs,
316 ColumnExpression(dtype=timespan.Timespan) as rhs,
317 ]:
318 return rhs.predicate_method("overlaps", lhs)
319 # Enable arithmetic operators on numeric types, without any type
320 # coercion or broadening.
321 case [
322 "+" | "-" | "*",
323 ColumnExpression(dtype=b.int | b.float) as lhs,
324 ColumnExpression() as rhs,
325 ] if lhs.dtype is rhs.dtype:
326 return lhs.method(self.OPERATOR_MAP[operator], rhs, dtype=lhs.dtype)
327 case ["/", ColumnExpression(dtype=b.float) as lhs, ColumnExpression(dtype=b.float) as rhs]:
328 return lhs.method("__truediv__", rhs, dtype=b.float)
329 case ["/", ColumnExpression(dtype=b.int) as lhs, ColumnExpression(dtype=b.int) as rhs]:
330 # SQLAlchemy maps Python's '/' (__truediv__) operator directly
331 # to SQL's '/', despite those being defined differently for
332 # integers. Our expression language uses the SQL definition,
333 # and we only care about these expressions being evaluated in
334 # SQL right now, but we still want to guard against it being
335 # evaluated in Python and producing a surprising answer, so we
336 # mark it as being supported only by a SQL engine.
337 return lhs.method(
338 "__truediv__",
339 rhs,
340 dtype=b.int,
341 supporting_engine_types={sql.Engine},
342 )
343 case ["%", ColumnExpression(dtype=b.int) as lhs, ColumnExpression(dtype=b.int) as rhs]:
344 return lhs.method("__mod__", rhs, dtype=b.int)
345 assert (
346 lhs.dtype is not None and rhs.dtype is not None
347 ), "Expression converter should not yield untyped nodes."
348 raise ExpressionTypeError(
349 f"Invalid types {lhs.dtype.__name__}, {rhs.dtype.__name__} for binary operator {operator!r} "
350 f"in expression {node!s}."
351 )
353 def visitIdentifier(self, name: str, node: Node) -> VisitorResult:
354 # Docstring inherited.
355 if name in self.bind:
356 value = self.bind[name]
357 if isinstance(value, (list, tuple, Set)):
358 elements = []
359 all_dtypes = set()
360 for item in value:
361 dtype = type(item)
362 all_dtypes.add(dtype)
363 elements.append(ColumnExpression.literal(item, dtype=dtype))
364 if len(all_dtypes) > 1:
365 raise ExpressionTypeError(
366 f"Mismatched types in bind iterable: {value} has a mix of {all_dtypes}."
367 )
368 elif not elements:
369 # Empty container
370 return ColumnContainer.sequence([])
371 else:
372 (dtype,) = all_dtypes
373 return ColumnContainer.sequence(elements, dtype=dtype)
374 return ColumnExpression.literal(value, dtype=type(value))
375 tag: ColumnTag
376 match categorizeConstant(name):
377 case ExpressionConstant.INGEST_DATE:
378 assert self.dataset_type_name is not None
379 tag = DatasetColumnTag(self.dataset_type_name, "ingest_date")
380 return ColumnExpression.reference(tag, self.column_types.ingest_date_pytype)
381 case ExpressionConstant.NULL:
382 return ColumnExpression.literal(None, type(None))
383 case None:
384 pass
385 case _:
386 raise AssertionError("Check for enum values should be exhaustive.")
387 element, column = categorizeElementId(self.universe, name)
388 if column is not None:
389 tag = DimensionRecordColumnTag(element.name, column)
390 dtype = (
391 timespan.Timespan
392 if column == timespan.TimespanDatabaseRepresentation.NAME
393 else element.RecordClass.fields.standard[column].getPythonType()
394 )
395 return ColumnExpression.reference(tag, dtype)
396 else:
397 tag = DimensionKeyColumnTag(element.name)
398 assert isinstance(element, Dimension)
399 return ColumnExpression.reference(tag, element.primaryKey.getPythonType())
401 def visitIsIn(
402 self, lhs: VisitorResult, values: list[VisitorResult], not_in: bool, node: Node
403 ) -> VisitorResult:
404 # Docstring inherited.
405 clauses: list[Predicate] = []
406 items: list[ColumnExpression] = []
407 assert isinstance(lhs, ColumnExpression), "LHS of IN guaranteed to be scalar by parser."
408 for rhs_item in values:
409 match rhs_item:
410 case ColumnExpressionSequence(
411 items=rhs_items, dtype=rhs_dtype
412 ) if rhs_dtype is None or rhs_dtype == lhs.dtype:
413 items.extend(rhs_items)
414 case ColumnContainer(dtype=lhs.dtype):
415 clauses.append(rhs_item.contains(lhs))
416 case ColumnExpression(dtype=lhs.dtype):
417 items.append(rhs_item)
418 case _:
419 raise ExpressionTypeError(
420 f"Invalid type {rhs_item.dtype} for element in {lhs.dtype} IN expression '{node}'."
421 )
422 if items:
423 clauses.append(ColumnContainer.sequence(items, dtype=lhs.dtype).contains(lhs))
424 result = Predicate.logical_or(*clauses)
425 if not_in:
426 result = result.logical_not()
427 return result
429 def visitNumericLiteral(self, value: str, node: Node) -> VisitorResult:
430 # Docstring inherited.
431 try:
432 return ColumnExpression.literal(int(value), dtype=int)
433 except ValueError:
434 return ColumnExpression.literal(float(value), dtype=float)
436 def visitParens(self, expression: VisitorResult, node: Node) -> VisitorResult:
437 # Docstring inherited.
438 return expression
440 def visitPointNode(self, ra: VisitorResult, dec: VisitorResult, node: Node) -> VisitorResult:
441 # Docstring inherited.
443 # this is a placeholder for future extension, we enabled syntax but
444 # do not support actual use just yet.
445 raise NotImplementedError("POINT() function is not supported yet")
447 def visitRangeLiteral(self, start: int, stop: int, stride: int | None, node: Node) -> VisitorResult:
448 # Docstring inherited.
449 return ColumnContainer.range_literal(range(start, stop + 1, stride or 1))
451 def visitStringLiteral(self, value: str, node: Node) -> VisitorResult:
452 # Docstring inherited.
453 return ColumnExpression.literal(value, dtype=str)
455 def visitTimeLiteral(self, value: astropy.time.Time, node: Node) -> VisitorResult:
456 # Docstring inherited.
457 return ColumnExpression.literal(value, dtype=astropy.time.Time)
459 def visitTupleNode(self, items: tuple[VisitorResult, ...], node: Node) -> VisitorResult:
460 # Docstring inherited.
461 match items:
462 case [
463 ColumnLiteral(value=begin, dtype=astropy.time.Time | types.NoneType),
464 ColumnLiteral(value=end, dtype=astropy.time.Time | types.NoneType),
465 ]:
466 return ColumnExpression.literal(timespan.Timespan(begin, end), dtype=timespan.Timespan)
467 raise ExpressionTypeError(
468 f'Invalid type(s) ({items[0].dtype}, {items[1].dtype}) in timespan tuple "{node}" '
469 '(Note that date/time strings must be preceded by "T" to be recognized).'
470 )
472 def visitUnaryOp(self, operator: str, operand: VisitorResult, node: Node) -> VisitorResult:
473 # Docstring inherited.
474 match (operator, operand):
475 case ["NOT", Predicate() as operand]:
476 return operand.logical_not()
477 case ["+", ColumnExpression(dtype=builtins.int | builtins.float) as operand]:
478 return operand.method("__pos__")
479 case ["-", ColumnExpression(dtype=builtins.int | builtins.float) as operand]:
480 return operand.method("__neg__")
481 raise ExpressionTypeError(
482 f"Unary operator {operator!r} is not valid for operand of type {operand.dtype!s} in {node!s}."
483 )