Skip to content

Commit e37bea5

Browse files
TST: run frame/series arithmetic tests with+without numexpr (#41178)
1 parent 59b2db1 commit e37bea5

File tree

2 files changed

+42
-7
lines changed

2 files changed

+42
-7
lines changed

pandas/tests/frame/test_arithmetic.py

+31-7
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
)
1818
import pandas._testing as tm
1919
import pandas.core.common as com
20+
from pandas.core.computation import expressions as expr
2021
from pandas.core.computation.expressions import (
2122
_MIN_ELEMENTS,
2223
NUMEXPR_INSTALLED,
@@ -27,6 +28,16 @@
2728
)
2829

2930

31+
@pytest.fixture(
32+
autouse=True, scope="module", params=[0, 1000000], ids=["numexpr", "python"]
33+
)
34+
def switch_numexpr_min_elements(request):
35+
_MIN_ELEMENTS = expr._MIN_ELEMENTS
36+
expr._MIN_ELEMENTS = request.param
37+
yield request.param
38+
expr._MIN_ELEMENTS = _MIN_ELEMENTS
39+
40+
3041
class DummyElement:
3142
def __init__(self, value, dtype):
3243
self.value = value
@@ -514,7 +525,12 @@ def f(x, y):
514525

515526
@pytest.mark.parametrize("op", ["__add__", "__sub__", "__mul__"])
516527
def test_arith_flex_frame_mixed(
517-
self, op, int_frame, mixed_int_frame, mixed_float_frame
528+
self,
529+
op,
530+
int_frame,
531+
mixed_int_frame,
532+
mixed_float_frame,
533+
switch_numexpr_min_elements,
518534
):
519535
f = getattr(operator, op)
520536

@@ -528,6 +544,12 @@ def test_arith_flex_frame_mixed(
528544
dtype = {"B": "uint64", "C": None}
529545
elif op in ["__add__", "__mul__"]:
530546
dtype = {"C": None}
547+
if expr.USE_NUMEXPR and switch_numexpr_min_elements == 0:
548+
# when using numexpr, the casting rules are slightly different:
549+
# in the `2 + mixed_int_frame` operation, int32 column becomes
550+
# and int64 column (not preserving dtype in operation with Python
551+
# scalar), and then the int32/int64 combo results in int64 result
552+
dtype["A"] = (2 + mixed_int_frame)["A"].dtype
531553
tm.assert_frame_equal(result, expected)
532554
_check_mixed_int(result, dtype=dtype)
533555

@@ -892,7 +914,7 @@ def test_frame_with_frame_reindex(self):
892914
],
893915
ids=lambda x: x.__name__,
894916
)
895-
def test_binop_other(self, op, value, dtype):
917+
def test_binop_other(self, op, value, dtype, switch_numexpr_min_elements):
896918

897919
skip = {
898920
(operator.truediv, "bool"),
@@ -941,11 +963,13 @@ def test_binop_other(self, op, value, dtype):
941963
elif (op, dtype) in skip:
942964

943965
if op in [operator.add, operator.mul]:
944-
# TODO we should assert this or not depending on whether
945-
# numexpr is used or not
946-
# with tm.assert_produces_warning(UserWarning):
947-
# # "evaluating in Python space because ..."
948-
op(s, e.value)
966+
if expr.USE_NUMEXPR and switch_numexpr_min_elements == 0:
967+
# "evaluating in Python space because ..."
968+
warn = UserWarning
969+
else:
970+
warn = None
971+
with tm.assert_produces_warning(warn):
972+
op(s, e.value)
949973

950974
else:
951975
msg = "operator '.*' not implemented for .* dtypes"

pandas/tests/series/test_arithmetic.py

+11
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,17 @@
2828
nanops,
2929
ops,
3030
)
31+
from pandas.core.computation import expressions as expr
32+
33+
34+
@pytest.fixture(
35+
autouse=True, scope="module", params=[0, 1000000], ids=["numexpr", "python"]
36+
)
37+
def switch_numexpr_min_elements(request):
38+
_MIN_ELEMENTS = expr._MIN_ELEMENTS
39+
expr._MIN_ELEMENTS = request.param
40+
yield request.param
41+
expr._MIN_ELEMENTS = _MIN_ELEMENTS
3142

3243

3344
def _permute(obj):

0 commit comments

Comments
 (0)