17
17
)
18
18
import pandas ._testing as tm
19
19
import pandas .core .common as com
20
+ from pandas .core .computation import expressions as expr
20
21
from pandas .core .computation .expressions import (
21
22
_MIN_ELEMENTS ,
22
23
NUMEXPR_INSTALLED ,
27
28
)
28
29
29
30
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
+
30
41
class DummyElement :
31
42
def __init__ (self , value , dtype ):
32
43
self .value = value
@@ -514,7 +525,12 @@ def f(x, y):
514
525
515
526
@pytest .mark .parametrize ("op" , ["__add__" , "__sub__" , "__mul__" ])
516
527
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 ,
518
534
):
519
535
f = getattr (operator , op )
520
536
@@ -528,6 +544,12 @@ def test_arith_flex_frame_mixed(
528
544
dtype = {"B" : "uint64" , "C" : None }
529
545
elif op in ["__add__" , "__mul__" ]:
530
546
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
531
553
tm .assert_frame_equal (result , expected )
532
554
_check_mixed_int (result , dtype = dtype )
533
555
@@ -892,7 +914,7 @@ def test_frame_with_frame_reindex(self):
892
914
],
893
915
ids = lambda x : x .__name__ ,
894
916
)
895
- def test_binop_other (self , op , value , dtype ):
917
+ def test_binop_other (self , op , value , dtype , switch_numexpr_min_elements ):
896
918
897
919
skip = {
898
920
(operator .truediv , "bool" ),
@@ -941,11 +963,13 @@ def test_binop_other(self, op, value, dtype):
941
963
elif (op , dtype ) in skip :
942
964
943
965
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 )
949
973
950
974
else :
951
975
msg = "operator '.*' not implemented for .* dtypes"
0 commit comments