Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""manage PyTables query interface via Expressions"""
3from __future__ import annotations
5import ast
6from decimal import (
7 Decimal,
8 InvalidOperation,
9)
10from functools import partial
11from typing import (
12 TYPE_CHECKING,
13 Any,
14 ClassVar,
15 Self,
16 cast,
17)
19import numpy as np
21from pandas._libs import lib
22from pandas._libs.tslibs import (
23 Timedelta,
24 Timestamp,
25 iNaT,
26)
27from pandas.errors import UndefinedVariableError
29from pandas.core.dtypes.common import is_list_like
31import pandas.core.common as com
32from pandas.core.computation import (
33 expr,
34 ops,
35 scope as _scope,
36)
37from pandas.core.computation.common import ensure_decoded
38from pandas.core.computation.expr import BaseExprVisitor
39from pandas.core.computation.ops import (
40 ARITH_OPS_SYMS,
41 is_term,
42)
43from pandas.core.construction import extract_array
44from pandas.core.indexes.base import Index
46from pandas.io.formats.printing import (
47 pprint_thing,
48 pprint_thing_encoded,
49)
51if TYPE_CHECKING:
52 from pandas._typing import (
53 TimeUnit,
54 npt,
55 )
58class PyTablesScope(_scope.Scope):
59 __slots__ = ("queryables",)
61 queryables: dict[str, Any]
63 def __init__(
64 self,
65 level: int,
66 global_dict=None,
67 local_dict=None,
68 queryables: dict[str, Any] | None = None,
69 ) -> None:
70 super().__init__(level + 1, global_dict=global_dict, local_dict=local_dict)
71 self.queryables = queryables or {}
74class Term(ops.Term):
75 env: PyTablesScope
77 def __new__(cls, name, env, side=None, encoding=None):
78 if isinstance(name, str):
79 klass = cls
80 else:
81 klass = Constant
82 return object.__new__(klass)
84 def __init__(self, name, env: PyTablesScope, side=None, encoding=None) -> None:
85 super().__init__(name, env, side=side, encoding=encoding)
87 def _resolve_name(self):
88 # must be a queryables
89 if self.side == "left":
90 # Note: The behavior of __new__ ensures that self.name is a str here
91 if self.name not in self.env.queryables:
92 raise NameError(f"name {self.name!r} is not defined")
93 return self.name
95 # resolve the rhs (and allow it to be None)
96 try:
97 return self.env.resolve(self.name, is_local=False)
98 except UndefinedVariableError:
99 return self.name
101 # read-only property overwriting read/write property
102 @property # type: ignore[misc]
103 def value(self):
104 return self._value
107class Constant(Term):
108 def __init__(self, name, env: PyTablesScope, side=None, encoding=None) -> None:
109 assert isinstance(env, PyTablesScope), type(env)
110 super().__init__(name, env, side=side, encoding=encoding)
112 def _resolve_name(self):
113 return self._name
116class BinOp(ops.BinOp):
117 _max_selectors = 31
119 op: str
120 queryables: dict[str, Any]
121 condition: str | None
123 def __init__(self, op: str, lhs, rhs, queryables: dict[str, Any], encoding) -> None:
124 super().__init__(op, lhs, rhs)
125 self.queryables = queryables
126 self.encoding = encoding
127 self.condition = None
129 def _disallow_scalar_only_bool_ops(self) -> None:
130 pass
132 def prune(self, klass):
133 if self.op in ARITH_OPS_SYMS:
134 # GH#41100: arithmetic in a where-clause is not supported. PyTables
135 # support is in maintenance mode, so rather than grow the query
136 # grammar we raise with a pointer to a working alternative.
137 raise NotImplementedError(
138 "arithmetic operations are not supported inside an HDFStore "
139 "'where' filter; instead store a precomputed column as a "
140 "data_column and query that, or read the data and apply the "
141 "filter in pandas (e.g. df[df['A'] % 3 == 0])."
142 )
144 def pr(left, right):
145 """create and return a new specialized BinOp from myself"""
146 if left is None:
147 return right
148 elif right is None:
149 return left
151 k = klass
152 if isinstance(left, ConditionBinOp):
153 if isinstance(right, ConditionBinOp):
154 k = JointConditionBinOp
155 elif isinstance(left, k):
156 return left
157 elif isinstance(right, k):
158 return right
160 elif isinstance(left, FilterBinOp):
161 if isinstance(right, FilterBinOp):
162 k = JointFilterBinOp
163 elif isinstance(left, k):
164 return left
165 elif isinstance(right, k):
166 return right
168 return k(
169 self.op, left, right, queryables=self.queryables, encoding=self.encoding
170 ).evaluate()
172 left, right = self.lhs, self.rhs
174 if is_term(left) and is_term(right):
175 res = pr(left, right)
176 elif not is_term(left) and is_term(right):
177 res = pr(left.prune(klass), right)
178 elif is_term(left) and not is_term(right):
179 res = pr(left, right.prune(klass))
180 elif not (is_term(left) or is_term(right)):
181 res = pr(left.prune(klass), right.prune(klass))
183 return res
185 def conform(self, rhs):
186 """inplace conform rhs"""
187 rhs = rhs.value
188 if rhs is None:
189 # Guarded here rather than in convert_value so both query paths are
190 # covered; bool(None) is False, Timestamp(None) is NaT. GH#64348
191 raise TypeError(f"Cannot compare [{self.lhs.value}] to None")
192 if not is_list_like(rhs):
193 rhs = [rhs]
194 if isinstance(rhs, np.ndarray):
195 rhs = rhs.ravel()
196 return rhs
198 @property
199 def is_valid(self) -> bool:
200 """return True if this is a valid field"""
201 return self.lhs.value in self.queryables
203 @property
204 def is_in_table(self) -> bool:
205 """
206 return True if this is a valid column name for generation (e.g. an
207 actual column in the table)
208 """
209 return self.queryables.get(self.lhs.value) is not None
211 @property
212 def kind(self):
213 """the kind of my field"""
214 return getattr(self.queryables.get(self.lhs.value), "kind", None)
216 @property
217 def meta(self):
218 """the meta of my field"""
219 return getattr(self.queryables.get(self.lhs.value), "meta", None)
221 @property
222 def metadata(self):
223 """the metadata of my field"""
224 return getattr(self.queryables.get(self.lhs.value), "metadata", None)
226 @property
227 def ordered(self):
228 """whether my field is an ordered Categorical (None if unrecorded)"""
229 return getattr(self.queryables.get(self.lhs.value), "ordered", None)
231 def generate(self, v) -> str:
232 """create and return the op string for this TermValue"""
233 val = v.tostring(self.encoding)
234 lhs = self.lhs.value
235 if v.truncated:
236 # The requested value fell strictly between `val` and the next
237 # storable one, so no row can equal it. Compare against `val` with
238 # an adjusted operator rather than against the value we truncated to.
239 if self.op == "==":
240 # matches nothing
241 return f"(({lhs} > {val}) & ({lhs} <= {val}))"
242 if self.op == "!=":
243 # matches everything
244 return f"(({lhs} <= {val}) | ({lhs} > {val}))"
245 op = "<=" if self.op in ("<", "<=") else ">"
246 return f"({lhs} {op} {val})"
247 kind = ensure_decoded(self.kind) or ""
248 if self.op in ("<", "<=") and kind.startswith(("datetime", "timedelta")):
249 # datetime64/timedelta64 columns are stored as int64 with NaT as
250 # iNaT, which sorts below every real value, so an unguarded "less
251 # than" would match the NaT rows. ">"/">=" need no such guard.
252 return f"(({lhs} != {iNaT}) & ({lhs} {self.op} {val}))"
253 return f"({lhs} {self.op} {val})"
255 def convert_value(self, conv_val) -> TermValue:
256 """
257 convert the expression that is in the term to something that is
258 accepted by pytables
259 """
261 def stringify(value):
262 if self.encoding is not None:
263 return pprint_thing_encoded(value, encoding=self.encoding)
264 return pprint_thing(value)
266 kind = ensure_decoded(self.kind)
267 meta = ensure_decoded(self.meta)
268 if kind == "datetime" or (kind and kind.startswith("datetime64")):
269 if isinstance(conv_val, (int, float)):
270 conv_val = stringify(conv_val)
271 conv_val = ensure_decoded(conv_val)
272 unit: TimeUnit = "ns"
273 if "[" in kind:
274 unit = cast("TimeUnit", kind.split("[")[-1][:-1])
275 conv_val = Timestamp(conv_val).as_unit(unit)
276 if conv_val.tz is not None:
277 conv_val = conv_val.tz_convert("UTC")
278 return TermValue(conv_val, conv_val._value, kind)
279 elif kind.startswith("timedelta"):
280 unit = "ns"
281 if "[" in kind:
282 unit = cast("TimeUnit", kind.split("[")[-1][:-1])
283 if isinstance(conv_val, str):
284 conv_val = Timedelta(conv_val)
285 elif lib.is_integer(conv_val) or lib.is_float(conv_val):
286 conv_val = Timedelta(conv_val, unit="s")
287 else:
288 conv_val = Timedelta(conv_val)
289 conv_val = conv_val.as_unit(unit)._value
290 return TermValue(int(conv_val), conv_val, kind)
292 elif meta == "category":
293 # `ordered` may be a np.bool_, and is None for files written before
294 # the flag was recorded, where we cannot tell.
295 ordered = self.ordered
296 is_unordered = ordered is not None and not ordered
297 if is_unordered and self.op in ["<", "<=", ">", ">="]:
298 # GH#68040 the stored codes are orderable, but the categories
299 # they stand for are not; match the in-memory comparison.
300 raise TypeError(
301 "Unordered Categoricals can only compare equality or not"
302 )
303 metadata = extract_array(self.metadata, extract_numpy=True)
304 result: npt.NDArray[np.intp] | np.intp | int
305 if conv_val not in metadata:
306 # GH#22977 value is not a category; use -2 as a code that
307 # matches no row, unlike -1 which is the code for NaN values.
308 result = -2
309 else:
310 # Find the index of the first match of conv_val in metadata
311 result = np.flatnonzero(metadata == conv_val)[0]
312 return TermValue(result, result, "integer")
313 elif kind == "integer":
314 try:
315 v_dec = Decimal(conv_val)
316 except InvalidOperation:
317 # GH 54186
318 # convert v to float to raise float's ValueError
319 float(conv_val)
320 else:
321 # Round toward -inf rather than to nearest, and flag the loss, so
322 # that generate() can adjust the operator instead of silently
323 # querying for a neighboring integer.
324 floored = v_dec.to_integral_exact(rounding="ROUND_FLOOR")
325 truncated = floored != v_dec
326 conv_val = int(floored)
327 return TermValue(conv_val, conv_val, kind, truncated=truncated)
328 return TermValue(conv_val, conv_val, kind)
329 elif kind == "float":
330 conv_val = float(conv_val)
331 return TermValue(conv_val, conv_val, kind)
332 elif kind == "bool":
333 if isinstance(conv_val, str):
334 conv_val = conv_val.strip().lower() not in [
335 "false",
336 "f",
337 "no",
338 "n",
339 "none",
340 "0",
341 "[]",
342 "{}",
343 "",
344 ]
345 else:
346 conv_val = bool(conv_val)
347 return TermValue(conv_val, conv_val, kind)
348 elif isinstance(conv_val, str):
349 # string quoting
350 return TermValue(conv_val, stringify(conv_val), "string")
351 else:
352 raise TypeError(
353 f"Cannot compare {conv_val} of type {type(conv_val)} to {kind} column"
354 )
356 def convert_values(self) -> None:
357 pass
360class FilterBinOp(BinOp):
361 filter: tuple[Any, Any, Index] | None = None
363 def __repr__(self) -> str:
364 if self.filter is None:
365 return "Filter: Not Initialized"
366 return pprint_thing(f"[Filter : [{self.filter[0]}] -> [{self.filter[1]}]")
368 def invert(self) -> Self:
369 """invert the filter"""
370 if self.filter is not None:
371 self.filter = (
372 self.filter[0],
373 self.generate_filter_op(invert=True),
374 self.filter[2],
375 )
376 return self
378 def format(self):
379 """return the actual filter format"""
380 return [self.filter]
382 # error: Signature of "evaluate" incompatible with supertype "BinOp"
383 def evaluate(self) -> Self | None: # type: ignore[override]
384 if not self.is_valid:
385 raise ValueError(f"query term is not valid [{self}]")
387 rhs = self.conform(self.rhs)
388 values = list(rhs)
390 if self.op not in ["==", "!="]:
391 if not self.is_in_table:
392 raise TypeError(
393 f"passing a filterable condition to a non-table indexer [{self}]"
394 )
395 return None
397 if self.is_in_table and len(values) <= self._max_selectors:
398 return None
399 filter_op = self.generate_filter_op()
400 self.filter = (self.lhs.value, filter_op, Index(values))
401 return self
403 def generate_filter_op(self, invert: bool = False):
404 if (self.op == "!=" and not invert) or (self.op == "==" and invert):
405 return lambda axis, vals: ~axis.isin(vals)
406 else:
407 return lambda axis, vals: axis.isin(vals)
410class JointFilterBinOp(FilterBinOp):
411 def format(self):
412 raise NotImplementedError("unable to collapse Joint Filters")
414 # error: Signature of "evaluate" incompatible with supertype "BinOp"
415 def evaluate(self) -> Self: # type: ignore[override]
416 return self
419class ConditionBinOp(BinOp):
420 def __repr__(self) -> str:
421 return pprint_thing(f"[Condition : [{self.condition}]]")
423 def invert(self):
424 """invert the condition"""
425 # if self.condition is not None:
426 # self.condition = "~(%s)" % self.condition
427 # return self
428 raise NotImplementedError(
429 "cannot use an invert condition when passing to numexpr"
430 )
432 def format(self):
433 """return the actual ne format"""
434 return self.condition
436 # error: Signature of "evaluate" incompatible with supertype "BinOp"
437 def evaluate(self) -> Self | None: # type: ignore[override]
438 if not self.is_valid:
439 raise ValueError(f"query term is not valid [{self}]")
441 # convert values if we are in the table
442 if not self.is_in_table:
443 return None
445 rhs = self.conform(self.rhs)
446 values = [self.convert_value(v) for v in rhs]
448 # equality conditions
449 if self.op in ["==", "!="]:
450 # too many values to create the expression?
451 if len(values) <= self._max_selectors:
452 vs = [self.generate(v) for v in values]
453 # De Morgan: "col != [a, b]" means "col != a AND col != b".
454 # OR-joining "!=" comparisons is true for every row.
455 joiner = " & " if self.op == "!=" else " | "
456 self.condition = f"({joiner.join(vs)})"
458 # use a filter after reading
459 else:
460 return None
461 else:
462 self.condition = self.generate(values[0])
464 return self
467class JointConditionBinOp(ConditionBinOp):
468 # error: Signature of "evaluate" incompatible with supertype "BinOp"
469 def evaluate(self) -> Self: # type: ignore[override]
470 self.condition = f"({self.lhs.condition} {self.op} {self.rhs.condition})"
471 return self
474class UnaryOp(ops.UnaryOp):
475 def prune(self, klass):
476 if self.op != "~":
477 raise NotImplementedError("UnaryOp only support invert type ops")
479 operand = self.operand
480 operand = operand.prune(klass)
482 if operand is not None and (
483 (issubclass(klass, ConditionBinOp) and operand.condition is not None)
484 or (
485 not issubclass(klass, ConditionBinOp)
486 and issubclass(klass, FilterBinOp)
487 and operand.filter is not None
488 )
489 ):
490 return operand.invert()
491 return None
494class PyTablesExprVisitor(BaseExprVisitor):
495 const_type: ClassVar[type[ops.Term]] = Constant
496 term_type: ClassVar[type[Term]] = Term
498 def __init__(self, env, engine, parser, **kwargs) -> None:
499 super().__init__(env, engine, parser)
500 for bin_op in self.binary_ops:
501 bin_node = self.binary_op_nodes_map[bin_op]
502 setattr(
503 self,
504 f"visit_{bin_node}",
505 lambda node, bin_op=bin_op: partial(BinOp, bin_op, **kwargs),
506 )
508 def visit_UnaryOp(self, node, **kwargs) -> ops.Term | UnaryOp | None:
509 if isinstance(node.op, (ast.Not, ast.Invert)):
510 return UnaryOp("~", self.visit(node.operand))
511 elif isinstance(node.op, ast.USub):
512 return self.const_type(-self.visit(node.operand).value, self.env)
513 elif isinstance(node.op, ast.UAdd):
514 raise NotImplementedError("Unary addition not supported")
515 # TODO: return None might never be reached
516 return None
518 def visit_Index(self, node, **kwargs):
519 return self.visit(node.value).value
521 def visit_Assign(self, node, **kwargs):
522 cmpr = ast.Compare(
523 ops=[ast.Eq()], left=node.targets[0], comparators=[node.value]
524 )
525 return self.visit(cmpr)
527 def visit_Subscript(self, node, **kwargs) -> ops.Term:
528 # only allow simple subscripts
530 value = self.visit(node.value)
531 slobj = self.visit(node.slice)
532 try:
533 value = value.value
534 except AttributeError:
535 pass
537 if isinstance(slobj, Term):
538 # In py39 np.ndarray lookups with Term containing int raise
539 slobj = slobj.value
541 try:
542 return self.const_type(value[slobj], self.env)
543 except TypeError as err:
544 raise ValueError(f"cannot subscript {value!r} with {slobj!r}") from err
546 def visit_Attribute(self, node, **kwargs):
547 attr = node.attr
548 value = node.value
550 ctx = type(node.ctx)
551 if ctx == ast.Load:
552 # resolve the value
553 resolved = self.visit(value)
555 # try to get the value to see if we are another expression
556 try:
557 resolved = resolved.value
558 except AttributeError:
559 pass
561 try:
562 return self.term_type(getattr(resolved, attr), self.env)
563 except AttributeError:
564 # something like datetime.datetime where scope is overridden
565 if isinstance(value, ast.Name) and value.id == attr:
566 return resolved
568 raise ValueError(f"Invalid Attribute context {ctx.__name__}")
570 def translate_In(self, op):
571 return ast.Eq() if isinstance(op, ast.In) else op
573 def _rewrite_membership_op(self, node, left, right):
574 return self.visit(node.op), node.op, left, right
577def _validate_where(w):
578 """
579 Validate that the where statement is of the right type.
581 The type may either be String, Expr, or list-like of Exprs.
583 Parameters
584 ----------
585 w : String term expression, Expr, or list-like of Exprs.
587 Returns
588 -------
589 where : The original where clause if the check was successful.
591 Raises
592 ------
593 TypeError : An invalid data type was passed in for w (e.g. dict).
594 """
595 if not (isinstance(w, (PyTablesExpr, str)) or is_list_like(w)):
596 raise TypeError(
597 "where must be passed as a string, PyTablesExpr, "
598 "or list-like of PyTablesExpr"
599 )
601 return w
604class PyTablesExpr(expr.Expr):
605 """
606 Hold a pytables-like expression, comprised of possibly multiple 'terms'.
608 Parameters
609 ----------
610 where : string term expression, PyTablesExpr, or list-like of PyTablesExprs
611 queryables : a "kinds" map (dict of column name -> kind), or None if column
612 is non-indexable
613 encoding : an encoding that will encode the query terms
615 Returns
616 -------
617 a PyTablesExpr object
619 Examples
620 --------
621 'index>=date'
622 "columns=['A', 'D']"
623 'columns=A'
624 'columns==A'
625 "~(columns=['A','B'])"
626 'index>df.index[3] & string="bar"'
627 '(index>df.index[3] & index<=df.index[6]) | string="bar"'
628 "ts>=Timestamp('2012-02-01')"
629 "major_axis>=20130101"
630 """
632 _visitor: PyTablesExprVisitor | None
633 env: PyTablesScope
634 expr: str
636 def __init__(
637 self,
638 where,
639 queryables: dict[str, Any] | None = None,
640 encoding=None,
641 scope_level: int = 0,
642 ) -> None:
643 where = _validate_where(where)
645 self.encoding = encoding
646 self.condition = None
647 self.filter = None
648 self.terms = None
649 self._visitor = None
651 # capture the environment if needed
652 local_dict: _scope.DeepChainMap[Any, Any] | None = None
654 if isinstance(where, PyTablesExpr):
655 local_dict = where.env.scope
656 _where = where.expr
658 elif is_list_like(where):
659 where = list(where)
660 for idx, w in enumerate(where):
661 if isinstance(w, PyTablesExpr):
662 local_dict = w.env.scope
663 else:
664 where[idx] = _validate_where(w)
665 _where = " & ".join([f"({w})" for w in com.flatten(where)])
666 else:
667 # _validate_where ensures we otherwise have a string
668 _where = where
670 self.expr = _where
671 self.env = PyTablesScope(scope_level + 1, local_dict=local_dict)
673 if queryables is not None and isinstance(self.expr, str):
674 self.env.queryables.update(queryables)
675 self._visitor = PyTablesExprVisitor(
676 self.env,
677 queryables=queryables,
678 parser="pytables",
679 engine="pytables",
680 encoding=encoding,
681 )
682 self.terms = self.parse()
684 def __repr__(self) -> str:
685 if self.terms is not None:
686 return pprint_thing(self.terms)
687 return pprint_thing(self.expr)
689 def evaluate(self):
690 """create and return the numexpr condition and filter"""
691 try:
692 self.condition = self.terms.prune(ConditionBinOp)
693 except AttributeError as err:
694 raise ValueError(
695 f"cannot process expression [{self.expr}], [{self}] "
696 "is not a valid condition"
697 ) from err
698 try:
699 self.filter = self.terms.prune(FilterBinOp)
700 except AttributeError as err:
701 raise ValueError(
702 f"cannot process expression [{self.expr}], [{self}] "
703 "is not a valid filter"
704 ) from err
706 return self.condition, self.filter
709class TermValue:
710 """hold a term value that we use to construct a condition/filter"""
712 def __init__(self, value, converted, kind: str, truncated: bool = False) -> None:
713 assert isinstance(kind, str), kind
714 self.value = value
715 self.converted = converted
716 self.kind = kind
717 # whether `converted` lost information relative to the value the user
718 # asked for; see BinOp.generate
719 self.truncated = truncated
721 def tostring(self, encoding) -> str:
722 """quote the string if not encoded else encode and return"""
723 if self.kind == "string":
724 if encoding is not None:
725 return str(self.converted)
726 return f'"{self.converted}"'
727 elif self.kind == "float":
728 return repr(self.converted)
729 return str(self.converted)
732def maybe_expression(s) -> bool:
733 """loose checking if s is a pytables-acceptable expression"""
734 if not isinstance(s, str):
735 return False
736 operations = PyTablesExprVisitor.binary_ops + PyTablesExprVisitor.unary_ops + ("=",)
738 # make sure we have an op at least
739 return any(op in s for op in operations)