385 lines
13 KiB
Python
385 lines
13 KiB
Python
import operator
|
|
import re
|
|
|
|
import numpy as np
|
|
from numpy.random import randn
|
|
import pytest
|
|
|
|
import pandas._testing as tm
|
|
from pandas.core.api import DataFrame
|
|
from pandas.core.computation import expressions as expr
|
|
|
|
_frame = DataFrame(randn(10000, 4), columns=list("ABCD"), dtype="float64")
|
|
_frame2 = DataFrame(randn(100, 4), columns=list("ABCD"), dtype="float64")
|
|
_mixed = DataFrame(
|
|
{
|
|
"A": _frame["A"].copy(),
|
|
"B": _frame["B"].astype("float32"),
|
|
"C": _frame["C"].astype("int64"),
|
|
"D": _frame["D"].astype("int32"),
|
|
}
|
|
)
|
|
_mixed2 = DataFrame(
|
|
{
|
|
"A": _frame2["A"].copy(),
|
|
"B": _frame2["B"].astype("float32"),
|
|
"C": _frame2["C"].astype("int64"),
|
|
"D": _frame2["D"].astype("int32"),
|
|
}
|
|
)
|
|
_integer = DataFrame(
|
|
np.random.randint(1, 100, size=(10001, 4)), columns=list("ABCD"), dtype="int64"
|
|
)
|
|
_integer2 = DataFrame(
|
|
np.random.randint(1, 100, size=(101, 4)), columns=list("ABCD"), dtype="int64"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not expr._USE_NUMEXPR, reason="not using numexpr")
|
|
class TestExpressions:
|
|
def setup_method(self, method):
|
|
|
|
self.frame = _frame.copy()
|
|
self.frame2 = _frame2.copy()
|
|
self.mixed = _mixed.copy()
|
|
self.mixed2 = _mixed2.copy()
|
|
self._MIN_ELEMENTS = expr._MIN_ELEMENTS
|
|
|
|
def teardown_method(self, method):
|
|
expr._MIN_ELEMENTS = self._MIN_ELEMENTS
|
|
|
|
def run_arithmetic(self, df, other):
|
|
expr._MIN_ELEMENTS = 0
|
|
operations = ["add", "sub", "mul", "mod", "truediv", "floordiv"]
|
|
for test_flex in [True, False]:
|
|
for arith in operations:
|
|
# TODO: share with run_binary
|
|
if test_flex:
|
|
op = lambda x, y: getattr(x, arith)(y)
|
|
op.__name__ = arith
|
|
else:
|
|
op = getattr(operator, arith)
|
|
expr.set_use_numexpr(False)
|
|
expected = op(df, other)
|
|
expr.set_use_numexpr(True)
|
|
|
|
result = op(df, other)
|
|
if arith == "truediv":
|
|
if expected.ndim == 1:
|
|
assert expected.dtype.kind == "f"
|
|
else:
|
|
assert all(x.kind == "f" for x in expected.dtypes.values)
|
|
tm.assert_equal(expected, result)
|
|
|
|
def run_binary(self, df, other):
|
|
"""
|
|
tests solely that the result is the same whether or not numexpr is
|
|
enabled. Need to test whether the function does the correct thing
|
|
elsewhere.
|
|
"""
|
|
expr._MIN_ELEMENTS = 0
|
|
expr.set_test_mode(True)
|
|
operations = ["gt", "lt", "ge", "le", "eq", "ne"]
|
|
|
|
for test_flex in [True, False]:
|
|
for arith in operations:
|
|
if test_flex:
|
|
op = lambda x, y: getattr(x, arith)(y)
|
|
op.__name__ = arith
|
|
else:
|
|
op = getattr(operator, arith)
|
|
expr.set_use_numexpr(False)
|
|
expected = op(df, other)
|
|
expr.set_use_numexpr(True)
|
|
|
|
expr.get_test_result()
|
|
result = op(df, other)
|
|
used_numexpr = expr.get_test_result()
|
|
assert used_numexpr, "Did not use numexpr as expected."
|
|
tm.assert_equal(expected, result)
|
|
|
|
def run_frame(self, df, other, run_binary=True):
|
|
self.run_arithmetic(df, other)
|
|
if run_binary:
|
|
expr.set_use_numexpr(False)
|
|
binary_comp = other + 1
|
|
expr.set_use_numexpr(True)
|
|
self.run_binary(df, binary_comp)
|
|
|
|
for i in range(len(df.columns)):
|
|
self.run_arithmetic(df.iloc[:, i], other.iloc[:, i])
|
|
# FIXME: dont leave commented-out
|
|
# series doesn't uses vec_compare instead of numexpr...
|
|
# binary_comp = other.iloc[:, i] + 1
|
|
# self.run_binary(df.iloc[:, i], binary_comp)
|
|
|
|
@pytest.mark.parametrize(
|
|
"df",
|
|
[
|
|
_integer,
|
|
_integer2,
|
|
# randint to get a case with zeros
|
|
_integer * np.random.randint(0, 2, size=np.shape(_integer)),
|
|
_frame,
|
|
_frame2,
|
|
_mixed,
|
|
_mixed2,
|
|
],
|
|
)
|
|
def test_arithmetic(self, df):
|
|
# TODO: FIGURE OUT HOW TO GET RUN_BINARY TO WORK WITH MIXED=...
|
|
# can't do arithmetic because comparison methods try to do *entire*
|
|
# frame instead of by-column
|
|
kinds = {x.kind for x in df.dtypes.values}
|
|
should = len(kinds) == 1
|
|
|
|
self.run_frame(df, df, run_binary=should)
|
|
|
|
def test_invalid(self):
|
|
|
|
# no op
|
|
result = expr._can_use_numexpr(
|
|
operator.add, None, self.frame, self.frame, "evaluate"
|
|
)
|
|
assert not result
|
|
|
|
# mixed
|
|
result = expr._can_use_numexpr(
|
|
operator.add, "+", self.mixed, self.frame, "evaluate"
|
|
)
|
|
assert not result
|
|
|
|
# min elements
|
|
result = expr._can_use_numexpr(
|
|
operator.add, "+", self.frame2, self.frame2, "evaluate"
|
|
)
|
|
assert not result
|
|
|
|
# ok, we only check on first part of expression
|
|
result = expr._can_use_numexpr(
|
|
operator.add, "+", self.frame, self.frame2, "evaluate"
|
|
)
|
|
assert result
|
|
|
|
@pytest.mark.parametrize(
|
|
"opname,op_str",
|
|
[("add", "+"), ("sub", "-"), ("mul", "*"), ("truediv", "/"), ("pow", "**")],
|
|
)
|
|
@pytest.mark.parametrize("left,right", [(_frame, _frame2), (_mixed, _mixed2)])
|
|
def test_binary_ops(self, opname, op_str, left, right):
|
|
def testit():
|
|
|
|
if opname == "pow":
|
|
# TODO: get this working
|
|
return
|
|
|
|
op = getattr(operator, opname)
|
|
|
|
result = expr._can_use_numexpr(op, op_str, left, left, "evaluate")
|
|
assert result != left._is_mixed_type
|
|
|
|
result = expr.evaluate(op, op_str, left, left, use_numexpr=True)
|
|
expected = expr.evaluate(op, op_str, left, left, use_numexpr=False)
|
|
|
|
if isinstance(result, DataFrame):
|
|
tm.assert_frame_equal(result, expected)
|
|
else:
|
|
tm.assert_numpy_array_equal(result, expected.values)
|
|
|
|
result = expr._can_use_numexpr(op, op_str, right, right, "evaluate")
|
|
assert not result
|
|
|
|
expr.set_use_numexpr(False)
|
|
testit()
|
|
expr.set_use_numexpr(True)
|
|
expr.set_numexpr_threads(1)
|
|
testit()
|
|
expr.set_numexpr_threads()
|
|
testit()
|
|
|
|
@pytest.mark.parametrize(
|
|
"opname,op_str",
|
|
[
|
|
("gt", ">"),
|
|
("lt", "<"),
|
|
("ge", ">="),
|
|
("le", "<="),
|
|
("eq", "=="),
|
|
("ne", "!="),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("left,right", [(_frame, _frame2), (_mixed, _mixed2)])
|
|
def test_comparison_ops(self, opname, op_str, left, right):
|
|
def testit():
|
|
f12 = left + 1
|
|
f22 = right + 1
|
|
|
|
op = getattr(operator, opname)
|
|
|
|
result = expr._can_use_numexpr(op, op_str, left, f12, "evaluate")
|
|
assert result != left._is_mixed_type
|
|
|
|
result = expr.evaluate(op, op_str, left, f12, use_numexpr=True)
|
|
expected = expr.evaluate(op, op_str, left, f12, use_numexpr=False)
|
|
if isinstance(result, DataFrame):
|
|
tm.assert_frame_equal(result, expected)
|
|
else:
|
|
tm.assert_numpy_array_equal(result, expected.values)
|
|
|
|
result = expr._can_use_numexpr(op, op_str, right, f22, "evaluate")
|
|
assert not result
|
|
|
|
expr.set_use_numexpr(False)
|
|
testit()
|
|
expr.set_use_numexpr(True)
|
|
expr.set_numexpr_threads(1)
|
|
testit()
|
|
expr.set_numexpr_threads()
|
|
testit()
|
|
|
|
@pytest.mark.parametrize("cond", [True, False])
|
|
@pytest.mark.parametrize("df", [_frame, _frame2, _mixed, _mixed2])
|
|
def test_where(self, cond, df):
|
|
def testit():
|
|
c = np.empty(df.shape, dtype=np.bool_)
|
|
c.fill(cond)
|
|
result = expr.where(c, df.values, df.values + 1)
|
|
expected = np.where(c, df.values, df.values + 1)
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
expr.set_use_numexpr(False)
|
|
testit()
|
|
expr.set_use_numexpr(True)
|
|
expr.set_numexpr_threads(1)
|
|
testit()
|
|
expr.set_numexpr_threads()
|
|
testit()
|
|
|
|
@pytest.mark.parametrize(
|
|
"op_str,opname", [("/", "truediv"), ("//", "floordiv"), ("**", "pow")]
|
|
)
|
|
def test_bool_ops_raise_on_arithmetic(self, op_str, opname):
|
|
df = DataFrame({"a": np.random.rand(10) > 0.5, "b": np.random.rand(10) > 0.5})
|
|
|
|
msg = f"operator {repr(op_str)} not implemented for bool dtypes"
|
|
f = getattr(operator, opname)
|
|
err_msg = re.escape(msg)
|
|
|
|
with pytest.raises(NotImplementedError, match=err_msg):
|
|
f(df, df)
|
|
|
|
with pytest.raises(NotImplementedError, match=err_msg):
|
|
f(df.a, df.b)
|
|
|
|
with pytest.raises(NotImplementedError, match=err_msg):
|
|
f(df.a, True)
|
|
|
|
with pytest.raises(NotImplementedError, match=err_msg):
|
|
f(False, df.a)
|
|
|
|
with pytest.raises(NotImplementedError, match=err_msg):
|
|
f(False, df)
|
|
|
|
with pytest.raises(NotImplementedError, match=err_msg):
|
|
f(df, True)
|
|
|
|
@pytest.mark.parametrize(
|
|
"op_str,opname", [("+", "add"), ("*", "mul"), ("-", "sub")]
|
|
)
|
|
def test_bool_ops_warn_on_arithmetic(self, op_str, opname):
|
|
n = 10
|
|
df = DataFrame({"a": np.random.rand(n) > 0.5, "b": np.random.rand(n) > 0.5})
|
|
|
|
subs = {"+": "|", "*": "&", "-": "^"}
|
|
sub_funcs = {"|": "or_", "&": "and_", "^": "xor"}
|
|
|
|
f = getattr(operator, opname)
|
|
fe = getattr(operator, sub_funcs[subs[op_str]])
|
|
|
|
if op_str == "-":
|
|
# raises TypeError
|
|
return
|
|
|
|
with tm.use_numexpr(True, min_elements=5):
|
|
with tm.assert_produces_warning(check_stacklevel=False):
|
|
r = f(df, df)
|
|
e = fe(df, df)
|
|
tm.assert_frame_equal(r, e)
|
|
|
|
with tm.assert_produces_warning(check_stacklevel=False):
|
|
r = f(df.a, df.b)
|
|
e = fe(df.a, df.b)
|
|
tm.assert_series_equal(r, e)
|
|
|
|
with tm.assert_produces_warning(check_stacklevel=False):
|
|
r = f(df.a, True)
|
|
e = fe(df.a, True)
|
|
tm.assert_series_equal(r, e)
|
|
|
|
with tm.assert_produces_warning(check_stacklevel=False):
|
|
r = f(False, df.a)
|
|
e = fe(False, df.a)
|
|
tm.assert_series_equal(r, e)
|
|
|
|
with tm.assert_produces_warning(check_stacklevel=False):
|
|
r = f(False, df)
|
|
e = fe(False, df)
|
|
tm.assert_frame_equal(r, e)
|
|
|
|
with tm.assert_produces_warning(check_stacklevel=False):
|
|
r = f(df, True)
|
|
e = fe(df, True)
|
|
tm.assert_frame_equal(r, e)
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_input,expected",
|
|
[
|
|
(
|
|
DataFrame(
|
|
[[0, 1, 2, "aa"], [0, 1, 2, "aa"]], columns=["a", "b", "c", "dtype"]
|
|
),
|
|
DataFrame([[False, False], [False, False]], columns=["a", "dtype"]),
|
|
),
|
|
(
|
|
DataFrame(
|
|
[[0, 3, 2, "aa"], [0, 4, 2, "aa"], [0, 1, 1, "bb"]],
|
|
columns=["a", "b", "c", "dtype"],
|
|
),
|
|
DataFrame(
|
|
[[False, False], [False, False], [False, False]],
|
|
columns=["a", "dtype"],
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_bool_ops_column_name_dtype(self, test_input, expected):
|
|
# GH 22383 - .ne fails if columns containing column name 'dtype'
|
|
result = test_input.loc[:, ["a", "dtype"]].ne(test_input.loc[:, ["a", "dtype"]])
|
|
tm.assert_frame_equal(result, expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"arith", ("add", "sub", "mul", "mod", "truediv", "floordiv")
|
|
)
|
|
@pytest.mark.parametrize("axis", (0, 1))
|
|
def test_frame_series_axis(self, axis, arith):
|
|
# GH#26736 Dataframe.floordiv(Series, axis=1) fails
|
|
if axis == 1 and arith == "floordiv":
|
|
pytest.xfail("'floordiv' does not succeed with axis=1 #27636")
|
|
|
|
df = self.frame
|
|
if axis == 1:
|
|
other = self.frame.iloc[0, :]
|
|
else:
|
|
other = self.frame.iloc[:, 0]
|
|
|
|
expr._MIN_ELEMENTS = 0
|
|
|
|
op_func = getattr(df, arith)
|
|
|
|
expr.set_use_numexpr(False)
|
|
expected = op_func(other, axis=axis)
|
|
expr.set_use_numexpr(True)
|
|
|
|
result = op_func(other, axis=axis)
|
|
tm.assert_frame_equal(expected, result)
|