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"""
2Expressions
3-----------
5Offer fast expression evaluation through numexpr
7"""
9from __future__ import annotations
11import operator
12from typing import TYPE_CHECKING
13import warnings
15import numpy as np
17from pandas._config.config import _global_config as config
19from pandas.util._exceptions import find_stack_level
21from pandas.core import roperator
22from pandas.core.computation.check import (
23 NUMEXPR_BLOCKED_VERSION,
24 NUMEXPR_INSTALLED,
25 warn_numexpr_blocked,
26)
28if NUMEXPR_INSTALLED:
29 import numexpr as ne
31if TYPE_CHECKING:
32 from pandas._typing import FuncType
34_TEST_MODE: bool | None = None
35_TEST_RESULT: list[bool] = []
36USE_NUMEXPR = NUMEXPR_INSTALLED
37_evaluate: FuncType | None = None
38_where: FuncType | None = None
40# the set of dtypes that we will allow pass to numexpr
41_ALLOWED_DTYPES = {
42 "evaluate": {"int64", "int32", "float64", "float32", "bool"},
43 "where": {"int64", "float64", "bool"},
44}
46# the minimum prod shape that we will use numexpr
47_MIN_ELEMENTS = 1_000_000
50def set_use_numexpr(v: bool = True) -> None:
51 # set/unset to use numexpr
52 global USE_NUMEXPR
53 if NUMEXPR_INSTALLED:
54 USE_NUMEXPR = v
56 # choose what we are going to do
57 global _evaluate, _where
59 if USE_NUMEXPR:
60 _evaluate = _evaluate_numexpr
61 _where = _where_numexpr
62 elif v and NUMEXPR_BLOCKED_VERSION is not None:
63 _evaluate = _evaluate_blocked
64 _where = _where_blocked
65 else:
66 _evaluate = _evaluate_standard
67 _where = _where_standard
70def set_numexpr_threads(n=None) -> None:
71 # if we are using numexpr, set the threads to n
72 # otherwise reset
73 if NUMEXPR_INSTALLED and USE_NUMEXPR:
74 if n is None:
75 n = ne.detect_number_of_cores()
76 ne.set_num_threads(n)
79def _evaluate_standard(op, op_str, left_op, right_op):
80 """
81 Standard evaluation.
82 """
83 if _TEST_MODE:
84 _store_test_result(False)
85 return op(left_op, right_op)
88def _evaluate_blocked(op, op_str, left_op, right_op):
89 """
90 Standard evaluation, warning that an unusable numexpr was skipped.
91 """
92 if _can_use_numexpr(op, op_str, left_op, right_op, "evaluate"):
93 warn_numexpr_blocked()
94 return _evaluate_standard(op, op_str, left_op, right_op)
97def _can_use_numexpr(op, op_str, left_op, right_op, dtype_check) -> bool:
98 """return left_op boolean if we WILL be using numexpr"""
99 if op_str is not None:
100 # required min elements (otherwise we are adding overhead)
101 if left_op.size > _MIN_ELEMENTS:
102 # check for dtype compatibility
103 dtypes: set[str] = set()
104 for o in [left_op, right_op]:
105 # ndarray and Series Case
106 if hasattr(o, "dtype"):
107 dtypes |= {o.dtype.name}
109 # allowed are a superset
110 if not len(dtypes) or _ALLOWED_DTYPES[dtype_check] >= dtypes:
111 return True
113 return False
116def normalize_numexpr_result(result):
117 """
118 View a numexpr result with numpy's canonical dtype for its itemsize/kind.
120 numexpr may return e.g. ``longlong`` where numpy's int64 is ``long``; the
121 two compare equal but have different ``dtype.type``, see GH#17945.
122 """
123 if isinstance(result, np.ndarray):
124 canonical = np.dtype(result.dtype.str)
125 # only view when needed, so the result otherwise keeps owning its data
126 if canonical.type is not result.dtype.type:
127 result = result.view(canonical)
128 return result
131def _evaluate_numexpr(op, op_str, left_op, right_op):
132 result = None
134 if _can_use_numexpr(op, op_str, left_op, right_op, "evaluate"):
135 is_reversed = op.__name__.strip("_").startswith("r")
136 if is_reversed:
137 # we were originally called by a reversed op method
138 left_op, right_op = right_op, left_op
140 left_value = left_op
141 right_value = right_op
143 try:
144 result = ne.evaluate(
145 f"left_value {op_str} right_value",
146 local_dict={"left_value": left_value, "right_value": right_value},
147 casting="safe",
148 )
149 result = normalize_numexpr_result(result)
150 except TypeError:
151 # numexpr raises eg for array ** array with integers
152 # (https://github.com/pydata/numexpr/issues/379)
153 pass
154 except NotImplementedError:
155 if _bool_arith_fallback(op_str, left_op, right_op):
156 pass
157 else:
158 raise
160 if is_reversed:
161 # reverse order to original for fallback
162 left_op, right_op = right_op, left_op
164 if _TEST_MODE:
165 _store_test_result(result is not None)
167 if result is None:
168 result = _evaluate_standard(op, op_str, left_op, right_op)
170 return result
173_op_str_mapping = {
174 operator.add: "+",
175 roperator.radd: "+",
176 operator.mul: "*",
177 roperator.rmul: "*",
178 operator.sub: "-",
179 roperator.rsub: "-",
180 operator.truediv: "/",
181 roperator.rtruediv: "/",
182 # floordiv not supported by numexpr 2.x
183 operator.floordiv: None,
184 roperator.rfloordiv: None,
185 # we require Python semantics for mod of negative for backwards compatibility
186 # see https://github.com/pydata/numexpr/issues/365
187 # so sticking with unaccelerated for now GH#36552
188 operator.mod: None,
189 roperator.rmod: None,
190 operator.pow: "**",
191 roperator.rpow: "**",
192 operator.eq: "==",
193 operator.ne: "!=",
194 operator.le: "<=",
195 operator.lt: "<",
196 operator.ge: ">=",
197 operator.gt: ">",
198 operator.and_: "&",
199 roperator.rand_: "&",
200 operator.or_: "|",
201 roperator.ror_: "|",
202 operator.xor: "^",
203 roperator.rxor: "^",
204 divmod: None,
205 roperator.rdivmod: None,
206}
209def _where_standard(cond, left_op, right_op):
210 # Caller is responsible for extracting ndarray if necessary
211 return np.where(cond, left_op, right_op)
214def _where_blocked(cond, left_op, right_op):
215 if _can_use_numexpr(None, "where", left_op, right_op, "where"):
216 warn_numexpr_blocked()
217 return _where_standard(cond, left_op, right_op)
220def _where_numexpr(cond, left_op, right_op):
221 # Caller is responsible for extracting ndarray if necessary
222 result = None
224 if _can_use_numexpr(None, "where", left_op, right_op, "where"):
225 result = ne.evaluate(
226 "where(cond_value, a_value, b_value)",
227 local_dict={"cond_value": cond, "a_value": left_op, "b_value": right_op},
228 casting="safe",
229 )
230 result = normalize_numexpr_result(result)
232 if result is None:
233 result = _where_standard(cond, left_op, right_op)
235 return result
238# turn myself on
239set_use_numexpr(config["compute"]["use_numexpr"])
242def _has_bool_dtype(x):
243 try:
244 return x.dtype == bool
245 except AttributeError:
246 return isinstance(x, (bool, np.bool_))
249_BOOL_OP_UNSUPPORTED = {"+": "|", "*": "&", "-": "^"}
252def _bool_arith_fallback(op_str, left_op, right_op) -> bool:
253 """
254 Check if we should fallback to the python `_evaluate_standard` in case
255 of an unsupported operation by numexpr, which is the case for some
256 boolean ops.
257 """
258 if _has_bool_dtype(left_op) and _has_bool_dtype(right_op):
259 if op_str in _BOOL_OP_UNSUPPORTED:
260 warnings.warn(
261 f"evaluating in Python space because the {op_str!r} "
262 "operator is not supported by numexpr for the bool dtype, "
263 f"use {_BOOL_OP_UNSUPPORTED[op_str]!r} instead.",
264 stacklevel=find_stack_level(),
265 )
266 return True
267 return False
270def evaluate(op, left_op, right_op, use_numexpr: bool = True):
271 """
272 Evaluate and return the expression of the op on left_op and right_op.
274 Parameters
275 ----------
276 op : the actual operand
277 left_op : left operand
278 right_op : right operand
279 use_numexpr : bool, default True
280 Whether to try to use numexpr.
281 """
282 op_str = _op_str_mapping[op]
283 if op_str is not None:
284 if use_numexpr:
285 # error: "None" not callable
286 return _evaluate(op, op_str, left_op, right_op) # type: ignore[misc]
287 return _evaluate_standard(op, op_str, left_op, right_op)
290def where(cond, left_op, right_op, use_numexpr: bool = True):
291 """
292 Evaluate the where condition cond on left_op and right_op.
294 Parameters
295 ----------
296 cond : np.ndarray[bool]
297 left_op : return if cond is True
298 right_op : return if cond is False
299 use_numexpr : bool, default True
300 Whether to try to use numexpr.
301 """
302 assert _where is not None
303 if use_numexpr:
304 return _where(cond, left_op, right_op)
305 else:
306 return _where_standard(cond, left_op, right_op)
309def set_test_mode(v: bool = True) -> None:
310 """
311 Keeps track of whether numexpr was used.
313 Stores an additional ``True`` for every successful use of evaluate with
314 numexpr since the last ``get_test_result``.
315 """
316 global _TEST_MODE, _TEST_RESULT
317 _TEST_MODE = v
318 _TEST_RESULT = []
321def _store_test_result(used_numexpr: bool) -> None:
322 if used_numexpr:
323 _TEST_RESULT.append(used_numexpr)
326def get_test_result() -> list[bool]:
327 """
328 Get test result and reset test_results.
329 """
330 global _TEST_RESULT
331 res = _TEST_RESULT
332 _TEST_RESULT = []
333 return res