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 import get_option
19from pandas.compat.numpy import np_version_gt2_3
20from pandas.util._exceptions import find_stack_level
22from pandas.core import roperator
23from pandas.core.computation.check import NUMEXPR_INSTALLED
24from pandas.util.version import Version
26if NUMEXPR_INSTALLED:
27 import numexpr as ne
29 ne_gt_211 = Version(ne.__version__) >= Version("2.11.0")
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 if np_version_gt2_3 and not ne_gt_211:
55 # incompatibility of numexpr 2.10 with newer pandas resulting in wrong data
56 # https://github.com/pandas-dev/pandas/issues/63320
57 USE_NUMEXPR = False
58 else:
59 USE_NUMEXPR = v
61 # choose what we are going to do
62 global _evaluate, _where
64 _evaluate = _evaluate_numexpr if USE_NUMEXPR else _evaluate_standard
65 _where = _where_numexpr if USE_NUMEXPR else _where_standard
68def set_numexpr_threads(n=None) -> None:
69 # if we are using numexpr, set the threads to n
70 # otherwise reset
71 if NUMEXPR_INSTALLED and USE_NUMEXPR:
72 if n is None:
73 n = ne.detect_number_of_cores()
74 ne.set_num_threads(n)
77def _evaluate_standard(op, op_str, left_op, right_op):
78 """
79 Standard evaluation.
80 """
81 if _TEST_MODE:
82 _store_test_result(False)
83 return op(left_op, right_op)
86def _can_use_numexpr(op, op_str, left_op, right_op, dtype_check) -> bool:
87 """return left_op boolean if we WILL be using numexpr"""
88 if op_str is not None:
89 # required min elements (otherwise we are adding overhead)
90 if left_op.size > _MIN_ELEMENTS:
91 # check for dtype compatibility
92 dtypes: set[str] = set()
93 for o in [left_op, right_op]:
94 # ndarray and Series Case
95 if hasattr(o, "dtype"):
96 dtypes |= {o.dtype.name}
98 # allowed are a superset
99 if not len(dtypes) or _ALLOWED_DTYPES[dtype_check] >= dtypes:
100 return True
102 return False
105def _evaluate_numexpr(op, op_str, left_op, right_op):
106 result = None
108 if _can_use_numexpr(op, op_str, left_op, right_op, "evaluate"):
109 is_reversed = op.__name__.strip("_").startswith("r")
110 if is_reversed:
111 # we were originally called by a reversed op method
112 left_op, right_op = right_op, left_op
114 left_value = left_op
115 right_value = right_op
117 try:
118 result = ne.evaluate(
119 f"left_value {op_str} right_value",
120 local_dict={"left_value": left_value, "right_value": right_value},
121 casting="safe",
122 )
123 except TypeError:
124 # numexpr raises eg for array ** array with integers
125 # (https://github.com/pydata/numexpr/issues/379)
126 pass
127 except NotImplementedError:
128 if _bool_arith_fallback(op_str, left_op, right_op):
129 pass
130 else:
131 raise
133 if is_reversed:
134 # reverse order to original for fallback
135 left_op, right_op = right_op, left_op
137 if _TEST_MODE:
138 _store_test_result(result is not None)
140 if result is None:
141 result = _evaluate_standard(op, op_str, left_op, right_op)
143 return result
146_op_str_mapping = {
147 operator.add: "+",
148 roperator.radd: "+",
149 operator.mul: "*",
150 roperator.rmul: "*",
151 operator.sub: "-",
152 roperator.rsub: "-",
153 operator.truediv: "/",
154 roperator.rtruediv: "/",
155 # floordiv not supported by numexpr 2.x
156 operator.floordiv: None,
157 roperator.rfloordiv: None,
158 # we require Python semantics for mod of negative for backwards compatibility
159 # see https://github.com/pydata/numexpr/issues/365
160 # so sticking with unaccelerated for now GH#36552
161 operator.mod: None,
162 roperator.rmod: None,
163 operator.pow: "**",
164 roperator.rpow: "**",
165 operator.eq: "==",
166 operator.ne: "!=",
167 operator.le: "<=",
168 operator.lt: "<",
169 operator.ge: ">=",
170 operator.gt: ">",
171 operator.and_: "&",
172 roperator.rand_: "&",
173 operator.or_: "|",
174 roperator.ror_: "|",
175 operator.xor: "^",
176 roperator.rxor: "^",
177 divmod: None,
178 roperator.rdivmod: None,
179}
182def _where_standard(cond, left_op, right_op):
183 # Caller is responsible for extracting ndarray if necessary
184 return np.where(cond, left_op, right_op)
187def _where_numexpr(cond, left_op, right_op):
188 # Caller is responsible for extracting ndarray if necessary
189 result = None
191 if _can_use_numexpr(None, "where", left_op, right_op, "where"):
192 result = ne.evaluate(
193 "where(cond_value, a_value, b_value)",
194 local_dict={"cond_value": cond, "a_value": left_op, "b_value": right_op},
195 casting="safe",
196 )
198 if result is None:
199 result = _where_standard(cond, left_op, right_op)
201 return result
204# turn myself on
205set_use_numexpr(get_option("compute.use_numexpr"))
208def _has_bool_dtype(x):
209 try:
210 return x.dtype == bool
211 except AttributeError:
212 return isinstance(x, (bool, np.bool_))
215_BOOL_OP_UNSUPPORTED = {"+": "|", "*": "&", "-": "^"}
218def _bool_arith_fallback(op_str, left_op, right_op) -> bool:
219 """
220 Check if we should fallback to the python `_evaluate_standard` in case
221 of an unsupported operation by numexpr, which is the case for some
222 boolean ops.
223 """
224 if _has_bool_dtype(left_op) and _has_bool_dtype(right_op):
225 if op_str in _BOOL_OP_UNSUPPORTED:
226 warnings.warn(
227 f"evaluating in Python space because the {op_str!r} "
228 "operator is not supported by numexpr for the bool dtype, "
229 f"use {_BOOL_OP_UNSUPPORTED[op_str]!r} instead.",
230 stacklevel=find_stack_level(),
231 )
232 return True
233 return False
236def evaluate(op, left_op, right_op, use_numexpr: bool = True):
237 """
238 Evaluate and return the expression of the op on left_op and right_op.
240 Parameters
241 ----------
242 op : the actual operand
243 left_op : left operand
244 right_op : right operand
245 use_numexpr : bool, default True
246 Whether to try to use numexpr.
247 """
248 op_str = _op_str_mapping[op]
249 if op_str is not None:
250 if use_numexpr:
251 # error: "None" not callable
252 return _evaluate(op, op_str, left_op, right_op) # type: ignore[misc]
253 return _evaluate_standard(op, op_str, left_op, right_op)
256def where(cond, left_op, right_op, use_numexpr: bool = True):
257 """
258 Evaluate the where condition cond on left_op and right_op.
260 Parameters
261 ----------
262 cond : np.ndarray[bool]
263 left_op : return if cond is True
264 right_op : return if cond is False
265 use_numexpr : bool, default True
266 Whether to try to use numexpr.
267 """
268 assert _where is not None
269 if use_numexpr:
270 return _where(cond, left_op, right_op)
271 else:
272 return _where_standard(cond, left_op, right_op)
275def set_test_mode(v: bool = True) -> None:
276 """
277 Keeps track of whether numexpr was used.
279 Stores an additional ``True`` for every successful use of evaluate with
280 numexpr since the last ``get_test_result``.
281 """
282 global _TEST_MODE, _TEST_RESULT
283 _TEST_MODE = v
284 _TEST_RESULT = []
287def _store_test_result(used_numexpr: bool) -> None:
288 if used_numexpr:
289 _TEST_RESULT.append(used_numexpr)
292def get_test_result() -> list[bool]:
293 """
294 Get test result and reset test_results.
295 """
296 global _TEST_RESULT
297 res = _TEST_RESULT
298 _TEST_RESULT = []
299 return res