673 lines
23 KiB
Python
673 lines
23 KiB
Python
|
from datetime import datetime
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from pandas.core.dtypes.common import is_scalar
|
||
|
|
|
||
|
|
import pandas as pd
|
||
|
|
from pandas import DataFrame, DatetimeIndex, Series, Timestamp, date_range, isna
|
||
|
|
import pandas._testing as tm
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture(params=["default", "float_string", "mixed_float", "mixed_int"])
|
||
|
|
def where_frame(request, float_string_frame, mixed_float_frame, mixed_int_frame):
|
||
|
|
if request.param == "default":
|
||
|
|
return DataFrame(np.random.randn(5, 3), columns=["A", "B", "C"])
|
||
|
|
if request.param == "float_string":
|
||
|
|
return float_string_frame
|
||
|
|
if request.param == "mixed_float":
|
||
|
|
return mixed_float_frame
|
||
|
|
if request.param == "mixed_int":
|
||
|
|
return mixed_int_frame
|
||
|
|
|
||
|
|
|
||
|
|
def _safe_add(df):
|
||
|
|
# only add to the numeric items
|
||
|
|
def is_ok(s):
|
||
|
|
return (
|
||
|
|
issubclass(s.dtype.type, (np.integer, np.floating)) and s.dtype != "uint8"
|
||
|
|
)
|
||
|
|
|
||
|
|
return DataFrame(dict((c, s + 1) if is_ok(s) else (c, s) for c, s in df.items()))
|
||
|
|
|
||
|
|
|
||
|
|
class TestDataFrameIndexingWhere:
|
||
|
|
def test_where_get(self, where_frame, float_string_frame):
|
||
|
|
def _check_get(df, cond, check_dtypes=True):
|
||
|
|
other1 = _safe_add(df)
|
||
|
|
rs = df.where(cond, other1)
|
||
|
|
rs2 = df.where(cond.values, other1)
|
||
|
|
for k, v in rs.items():
|
||
|
|
exp = Series(np.where(cond[k], df[k], other1[k]), index=v.index)
|
||
|
|
tm.assert_series_equal(v, exp, check_names=False)
|
||
|
|
tm.assert_frame_equal(rs, rs2)
|
||
|
|
|
||
|
|
# dtypes
|
||
|
|
if check_dtypes:
|
||
|
|
assert (rs.dtypes == df.dtypes).all()
|
||
|
|
|
||
|
|
# check getting
|
||
|
|
df = where_frame
|
||
|
|
if df is float_string_frame:
|
||
|
|
msg = "'>' not supported between instances of 'str' and 'int'"
|
||
|
|
with pytest.raises(TypeError, match=msg):
|
||
|
|
df > 0
|
||
|
|
return
|
||
|
|
cond = df > 0
|
||
|
|
_check_get(df, cond)
|
||
|
|
|
||
|
|
def test_where_upcasting(self):
|
||
|
|
# upcasting case (GH # 2794)
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
c: Series([1] * 3, dtype=c)
|
||
|
|
for c in ["float32", "float64", "int32", "int64"]
|
||
|
|
}
|
||
|
|
)
|
||
|
|
df.iloc[1, :] = 0
|
||
|
|
result = df.dtypes
|
||
|
|
expected = Series(
|
||
|
|
[
|
||
|
|
np.dtype("float32"),
|
||
|
|
np.dtype("float64"),
|
||
|
|
np.dtype("int32"),
|
||
|
|
np.dtype("int64"),
|
||
|
|
],
|
||
|
|
index=["float32", "float64", "int32", "int64"],
|
||
|
|
)
|
||
|
|
|
||
|
|
# when we don't preserve boolean casts
|
||
|
|
#
|
||
|
|
# expected = Series({ 'float32' : 1, 'float64' : 3 })
|
||
|
|
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
def test_where_alignment(self, where_frame, float_string_frame):
|
||
|
|
# aligning
|
||
|
|
def _check_align(df, cond, other, check_dtypes=True):
|
||
|
|
rs = df.where(cond, other)
|
||
|
|
for i, k in enumerate(rs.columns):
|
||
|
|
result = rs[k]
|
||
|
|
d = df[k].values
|
||
|
|
c = cond[k].reindex(df[k].index).fillna(False).values
|
||
|
|
|
||
|
|
if is_scalar(other):
|
||
|
|
o = other
|
||
|
|
else:
|
||
|
|
if isinstance(other, np.ndarray):
|
||
|
|
o = Series(other[:, i], index=result.index).values
|
||
|
|
else:
|
||
|
|
o = other[k].values
|
||
|
|
|
||
|
|
new_values = d if c.all() else np.where(c, d, o)
|
||
|
|
expected = Series(new_values, index=result.index, name=k)
|
||
|
|
|
||
|
|
# since we can't always have the correct numpy dtype
|
||
|
|
# as numpy doesn't know how to downcast, don't check
|
||
|
|
tm.assert_series_equal(result, expected, check_dtype=False)
|
||
|
|
|
||
|
|
# dtypes
|
||
|
|
# can't check dtype when other is an ndarray
|
||
|
|
|
||
|
|
if check_dtypes and not isinstance(other, np.ndarray):
|
||
|
|
assert (rs.dtypes == df.dtypes).all()
|
||
|
|
|
||
|
|
df = where_frame
|
||
|
|
if df is float_string_frame:
|
||
|
|
msg = "'>' not supported between instances of 'str' and 'int'"
|
||
|
|
with pytest.raises(TypeError, match=msg):
|
||
|
|
df > 0
|
||
|
|
return
|
||
|
|
|
||
|
|
# other is a frame
|
||
|
|
cond = (df > 0)[1:]
|
||
|
|
_check_align(df, cond, _safe_add(df))
|
||
|
|
|
||
|
|
# check other is ndarray
|
||
|
|
cond = df > 0
|
||
|
|
_check_align(df, cond, (_safe_add(df).values))
|
||
|
|
|
||
|
|
# integers are upcast, so don't check the dtypes
|
||
|
|
cond = df > 0
|
||
|
|
check_dtypes = all(not issubclass(s.type, np.integer) for s in df.dtypes)
|
||
|
|
_check_align(df, cond, np.nan, check_dtypes=check_dtypes)
|
||
|
|
|
||
|
|
def test_where_invalid(self):
|
||
|
|
# invalid conditions
|
||
|
|
df = DataFrame(np.random.randn(5, 3), columns=["A", "B", "C"])
|
||
|
|
cond = df > 0
|
||
|
|
|
||
|
|
err1 = (df + 1).values[0:2, :]
|
||
|
|
msg = "other must be the same shape as self when an ndarray"
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.where(cond, err1)
|
||
|
|
|
||
|
|
err2 = cond.iloc[:2, :].values
|
||
|
|
other1 = _safe_add(df)
|
||
|
|
msg = "Array conditional must be same shape as self"
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.where(err2, other1)
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.mask(True)
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.mask(0)
|
||
|
|
|
||
|
|
def test_where_set(self, where_frame, float_string_frame):
|
||
|
|
# where inplace
|
||
|
|
|
||
|
|
def _check_set(df, cond, check_dtypes=True):
|
||
|
|
dfi = df.copy()
|
||
|
|
econd = cond.reindex_like(df).fillna(True)
|
||
|
|
expected = dfi.mask(~econd)
|
||
|
|
|
||
|
|
return_value = dfi.where(cond, np.nan, inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(dfi, expected)
|
||
|
|
|
||
|
|
# dtypes (and confirm upcasts)x
|
||
|
|
if check_dtypes:
|
||
|
|
for k, v in df.dtypes.items():
|
||
|
|
if issubclass(v.type, np.integer) and not cond[k].all():
|
||
|
|
v = np.dtype("float64")
|
||
|
|
assert dfi[k].dtype == v
|
||
|
|
|
||
|
|
df = where_frame
|
||
|
|
if df is float_string_frame:
|
||
|
|
msg = "'>' not supported between instances of 'str' and 'int'"
|
||
|
|
with pytest.raises(TypeError, match=msg):
|
||
|
|
df > 0
|
||
|
|
return
|
||
|
|
|
||
|
|
cond = df > 0
|
||
|
|
_check_set(df, cond)
|
||
|
|
|
||
|
|
cond = df >= 0
|
||
|
|
_check_set(df, cond)
|
||
|
|
|
||
|
|
# aligning
|
||
|
|
cond = (df >= 0)[1:]
|
||
|
|
_check_set(df, cond)
|
||
|
|
|
||
|
|
def test_where_series_slicing(self):
|
||
|
|
# GH 10218
|
||
|
|
# test DataFrame.where with Series slicing
|
||
|
|
df = DataFrame({"a": range(3), "b": range(4, 7)})
|
||
|
|
result = df.where(df["a"] == 1)
|
||
|
|
expected = df[df["a"] == 1].reindex(df.index)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("klass", [list, tuple, np.array])
|
||
|
|
def test_where_array_like(self, klass):
|
||
|
|
# see gh-15414
|
||
|
|
df = DataFrame({"a": [1, 2, 3]})
|
||
|
|
cond = [[False], [True], [True]]
|
||
|
|
expected = DataFrame({"a": [np.nan, 2, 3]})
|
||
|
|
|
||
|
|
result = df.where(klass(cond))
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
df["b"] = 2
|
||
|
|
expected["b"] = [2, np.nan, 2]
|
||
|
|
cond = [[False, True], [True, False], [True, True]]
|
||
|
|
|
||
|
|
result = df.where(klass(cond))
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"cond",
|
||
|
|
[
|
||
|
|
[[1], [0], [1]],
|
||
|
|
Series([[2], [5], [7]]),
|
||
|
|
DataFrame({"a": [2, 5, 7]}),
|
||
|
|
[["True"], ["False"], ["True"]],
|
||
|
|
[[Timestamp("2017-01-01")], [pd.NaT], [Timestamp("2017-01-02")]],
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_where_invalid_input_single(self, cond):
|
||
|
|
# see gh-15414: only boolean arrays accepted
|
||
|
|
df = DataFrame({"a": [1, 2, 3]})
|
||
|
|
msg = "Boolean array expected for the condition"
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.where(cond)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"cond",
|
||
|
|
[
|
||
|
|
[[0, 1], [1, 0], [1, 1]],
|
||
|
|
Series([[0, 2], [5, 0], [4, 7]]),
|
||
|
|
[["False", "True"], ["True", "False"], ["True", "True"]],
|
||
|
|
DataFrame({"a": [2, 5, 7], "b": [4, 8, 9]}),
|
||
|
|
[
|
||
|
|
[pd.NaT, Timestamp("2017-01-01")],
|
||
|
|
[Timestamp("2017-01-02"), pd.NaT],
|
||
|
|
[Timestamp("2017-01-03"), Timestamp("2017-01-03")],
|
||
|
|
],
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_where_invalid_input_multiple(self, cond):
|
||
|
|
# see gh-15414: only boolean arrays accepted
|
||
|
|
df = DataFrame({"a": [1, 2, 3], "b": [2, 2, 2]})
|
||
|
|
msg = "Boolean array expected for the condition"
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.where(cond)
|
||
|
|
|
||
|
|
def test_where_dataframe_col_match(self):
|
||
|
|
df = DataFrame([[1, 2, 3], [4, 5, 6]])
|
||
|
|
cond = DataFrame([[True, False, True], [False, False, True]])
|
||
|
|
|
||
|
|
result = df.where(cond)
|
||
|
|
expected = DataFrame([[1.0, np.nan, 3], [np.nan, np.nan, 6]])
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# this *does* align, though has no matching columns
|
||
|
|
cond.columns = ["a", "b", "c"]
|
||
|
|
result = df.where(cond)
|
||
|
|
expected = DataFrame(np.nan, index=df.index, columns=df.columns)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_where_ndframe_align(self):
|
||
|
|
msg = "Array conditional must be same shape as self"
|
||
|
|
df = DataFrame([[1, 2, 3], [4, 5, 6]])
|
||
|
|
|
||
|
|
cond = [True]
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.where(cond)
|
||
|
|
|
||
|
|
expected = DataFrame([[1, 2, 3], [np.nan, np.nan, np.nan]])
|
||
|
|
|
||
|
|
out = df.where(Series(cond))
|
||
|
|
tm.assert_frame_equal(out, expected)
|
||
|
|
|
||
|
|
cond = np.array([False, True, False, True])
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.where(cond)
|
||
|
|
|
||
|
|
expected = DataFrame([[np.nan, np.nan, np.nan], [4, 5, 6]])
|
||
|
|
|
||
|
|
out = df.where(Series(cond))
|
||
|
|
tm.assert_frame_equal(out, expected)
|
||
|
|
|
||
|
|
def test_where_bug(self):
|
||
|
|
# see gh-2793
|
||
|
|
df = DataFrame(
|
||
|
|
{"a": [1.0, 2.0, 3.0, 4.0], "b": [4.0, 3.0, 2.0, 1.0]}, dtype="float64"
|
||
|
|
)
|
||
|
|
expected = DataFrame(
|
||
|
|
{"a": [np.nan, np.nan, 3.0, 4.0], "b": [4.0, 3.0, np.nan, np.nan]},
|
||
|
|
dtype="float64",
|
||
|
|
)
|
||
|
|
result = df.where(df > 2, np.nan)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(result > 2, np.nan, inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_where_bug_mixed(self, sint_dtype):
|
||
|
|
# see gh-2793
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"a": np.array([1, 2, 3, 4], dtype=sint_dtype),
|
||
|
|
"b": np.array([4.0, 3.0, 2.0, 1.0], dtype="float64"),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
expected = DataFrame(
|
||
|
|
{"a": [np.nan, np.nan, 3.0, 4.0], "b": [4.0, 3.0, np.nan, np.nan]},
|
||
|
|
dtype="float64",
|
||
|
|
)
|
||
|
|
|
||
|
|
result = df.where(df > 2, np.nan)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(result > 2, np.nan, inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_where_bug_transposition(self):
|
||
|
|
# see gh-7506
|
||
|
|
a = DataFrame({0: [1, 2], 1: [3, 4], 2: [5, 6]})
|
||
|
|
b = DataFrame({0: [np.nan, 8], 1: [9, np.nan], 2: [np.nan, np.nan]})
|
||
|
|
do_not_replace = b.isna() | (a > b)
|
||
|
|
|
||
|
|
expected = a.copy()
|
||
|
|
expected[~do_not_replace] = b
|
||
|
|
|
||
|
|
result = a.where(do_not_replace, b)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
a = DataFrame({0: [4, 6], 1: [1, 0]})
|
||
|
|
b = DataFrame({0: [np.nan, 3], 1: [3, np.nan]})
|
||
|
|
do_not_replace = b.isna() | (a > b)
|
||
|
|
|
||
|
|
expected = a.copy()
|
||
|
|
expected[~do_not_replace] = b
|
||
|
|
|
||
|
|
result = a.where(do_not_replace, b)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_where_datetime(self):
|
||
|
|
|
||
|
|
# GH 3311
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"A": date_range("20130102", periods=5),
|
||
|
|
"B": date_range("20130104", periods=5),
|
||
|
|
"C": np.random.randn(5),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
stamp = datetime(2013, 1, 3)
|
||
|
|
msg = "'>' not supported between instances of 'float' and 'datetime.datetime'"
|
||
|
|
with pytest.raises(TypeError, match=msg):
|
||
|
|
df > stamp
|
||
|
|
|
||
|
|
result = df[df.iloc[:, :-1] > stamp]
|
||
|
|
|
||
|
|
expected = df.copy()
|
||
|
|
expected.loc[[0, 1], "A"] = np.nan
|
||
|
|
expected.loc[:, "C"] = np.nan
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_where_none(self):
|
||
|
|
# GH 4667
|
||
|
|
# setting with None changes dtype
|
||
|
|
df = DataFrame({"series": Series(range(10))}).astype(float)
|
||
|
|
df[df > 7] = None
|
||
|
|
expected = DataFrame(
|
||
|
|
{"series": Series([0, 1, 2, 3, 4, 5, 6, 7, np.nan, np.nan])}
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(df, expected)
|
||
|
|
|
||
|
|
# GH 7656
|
||
|
|
df = DataFrame(
|
||
|
|
[
|
||
|
|
{"A": 1, "B": np.nan, "C": "Test"},
|
||
|
|
{"A": np.nan, "B": "Test", "C": np.nan},
|
||
|
|
]
|
||
|
|
)
|
||
|
|
msg = "boolean setting on mixed-type"
|
||
|
|
|
||
|
|
with pytest.raises(TypeError, match=msg):
|
||
|
|
df.where(~isna(df), None, inplace=True)
|
||
|
|
|
||
|
|
def test_where_empty_df_and_empty_cond_having_non_bool_dtypes(self):
|
||
|
|
# see gh-21947
|
||
|
|
df = DataFrame(columns=["a"])
|
||
|
|
cond = df
|
||
|
|
assert (cond.dtypes == object).all()
|
||
|
|
|
||
|
|
result = df.where(cond)
|
||
|
|
tm.assert_frame_equal(result, df)
|
||
|
|
|
||
|
|
def test_where_align(self):
|
||
|
|
def create():
|
||
|
|
df = DataFrame(np.random.randn(10, 3))
|
||
|
|
df.iloc[3:5, 0] = np.nan
|
||
|
|
df.iloc[4:6, 1] = np.nan
|
||
|
|
df.iloc[5:8, 2] = np.nan
|
||
|
|
return df
|
||
|
|
|
||
|
|
# series
|
||
|
|
df = create()
|
||
|
|
expected = df.fillna(df.mean())
|
||
|
|
result = df.where(pd.notna(df), df.mean(), axis="columns")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
return_value = df.where(pd.notna(df), df.mean(), inplace=True, axis="columns")
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(df, expected)
|
||
|
|
|
||
|
|
df = create().fillna(0)
|
||
|
|
expected = df.apply(lambda x, y: x.where(x > 0, y), y=df[0])
|
||
|
|
result = df.where(df > 0, df[0], axis="index")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
result = df.where(df > 0, df[0], axis="rows")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# frame
|
||
|
|
df = create()
|
||
|
|
expected = df.fillna(1)
|
||
|
|
result = df.where(
|
||
|
|
pd.notna(df), DataFrame(1, index=df.index, columns=df.columns)
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_where_complex(self):
|
||
|
|
# GH 6345
|
||
|
|
expected = DataFrame([[1 + 1j, 2], [np.nan, 4 + 1j]], columns=["a", "b"])
|
||
|
|
df = DataFrame([[1 + 1j, 2], [5 + 1j, 4 + 1j]], columns=["a", "b"])
|
||
|
|
df[df.abs() >= 5] = np.nan
|
||
|
|
tm.assert_frame_equal(df, expected)
|
||
|
|
|
||
|
|
def test_where_axis(self):
|
||
|
|
# GH 9736
|
||
|
|
df = DataFrame(np.random.randn(2, 2))
|
||
|
|
mask = DataFrame([[False, False], [False, False]])
|
||
|
|
s = Series([0, 1])
|
||
|
|
|
||
|
|
expected = DataFrame([[0, 0], [1, 1]], dtype="float64")
|
||
|
|
result = df.where(mask, s, axis="index")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, s, axis="index", inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
expected = DataFrame([[0, 1], [0, 1]], dtype="float64")
|
||
|
|
result = df.where(mask, s, axis="columns")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, s, axis="columns", inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# Upcast needed
|
||
|
|
df = DataFrame([[1, 2], [3, 4]], dtype="int64")
|
||
|
|
mask = DataFrame([[False, False], [False, False]])
|
||
|
|
s = Series([0, np.nan])
|
||
|
|
|
||
|
|
expected = DataFrame([[0, 0], [np.nan, np.nan]], dtype="float64")
|
||
|
|
result = df.where(mask, s, axis="index")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, s, axis="index", inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
expected = DataFrame([[0, np.nan], [0, np.nan]])
|
||
|
|
result = df.where(mask, s, axis="columns")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
expected = DataFrame(
|
||
|
|
{
|
||
|
|
0: np.array([0, 0], dtype="int64"),
|
||
|
|
1: np.array([np.nan, np.nan], dtype="float64"),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, s, axis="columns", inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# Multiple dtypes (=> multiple Blocks)
|
||
|
|
df = pd.concat(
|
||
|
|
[
|
||
|
|
DataFrame(np.random.randn(10, 2)),
|
||
|
|
DataFrame(np.random.randint(0, 10, size=(10, 2)), dtype="int64"),
|
||
|
|
],
|
||
|
|
ignore_index=True,
|
||
|
|
axis=1,
|
||
|
|
)
|
||
|
|
mask = DataFrame(False, columns=df.columns, index=df.index)
|
||
|
|
s1 = Series(1, index=df.columns)
|
||
|
|
s2 = Series(2, index=df.index)
|
||
|
|
|
||
|
|
result = df.where(mask, s1, axis="columns")
|
||
|
|
expected = DataFrame(1.0, columns=df.columns, index=df.index)
|
||
|
|
expected[2] = expected[2].astype("int64")
|
||
|
|
expected[3] = expected[3].astype("int64")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, s1, axis="columns", inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result = df.where(mask, s2, axis="index")
|
||
|
|
expected = DataFrame(2.0, columns=df.columns, index=df.index)
|
||
|
|
expected[2] = expected[2].astype("int64")
|
||
|
|
expected[3] = expected[3].astype("int64")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, s2, axis="index", inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# DataFrame vs DataFrame
|
||
|
|
d1 = df.copy().drop(1, axis=0)
|
||
|
|
expected = df.copy()
|
||
|
|
expected.loc[1, :] = np.nan
|
||
|
|
|
||
|
|
result = df.where(mask, d1)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
result = df.where(mask, d1, axis="index")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, d1, inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, d1, inplace=True, axis="index")
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
d2 = df.copy().drop(1, axis=1)
|
||
|
|
expected = df.copy()
|
||
|
|
expected.loc[:, 1] = np.nan
|
||
|
|
|
||
|
|
result = df.where(mask, d2)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
result = df.where(mask, d2, axis="columns")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, d2, inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
result = df.copy()
|
||
|
|
return_value = result.where(mask, d2, inplace=True, axis="columns")
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_where_callable(self):
|
||
|
|
# GH 12533
|
||
|
|
df = DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
||
|
|
result = df.where(lambda x: x > 4, lambda x: x + 1)
|
||
|
|
exp = DataFrame([[2, 3, 4], [5, 5, 6], [7, 8, 9]])
|
||
|
|
tm.assert_frame_equal(result, exp)
|
||
|
|
tm.assert_frame_equal(result, df.where(df > 4, df + 1))
|
||
|
|
|
||
|
|
# return ndarray and scalar
|
||
|
|
result = df.where(lambda x: (x % 2 == 0).values, lambda x: 99)
|
||
|
|
exp = DataFrame([[99, 2, 99], [4, 99, 6], [99, 8, 99]])
|
||
|
|
tm.assert_frame_equal(result, exp)
|
||
|
|
tm.assert_frame_equal(result, df.where(df % 2 == 0, 99))
|
||
|
|
|
||
|
|
# chain
|
||
|
|
result = (df + 2).where(lambda x: x > 8, lambda x: x + 10)
|
||
|
|
exp = DataFrame([[13, 14, 15], [16, 17, 18], [9, 10, 11]])
|
||
|
|
tm.assert_frame_equal(result, exp)
|
||
|
|
tm.assert_frame_equal(result, (df + 2).where((df + 2) > 8, (df + 2) + 10))
|
||
|
|
|
||
|
|
def test_where_tz_values(self, tz_naive_fixture):
|
||
|
|
df1 = DataFrame(
|
||
|
|
DatetimeIndex(["20150101", "20150102", "20150103"], tz=tz_naive_fixture),
|
||
|
|
columns=["date"],
|
||
|
|
)
|
||
|
|
df2 = DataFrame(
|
||
|
|
DatetimeIndex(["20150103", "20150104", "20150105"], tz=tz_naive_fixture),
|
||
|
|
columns=["date"],
|
||
|
|
)
|
||
|
|
mask = DataFrame([True, True, False], columns=["date"])
|
||
|
|
exp = DataFrame(
|
||
|
|
DatetimeIndex(["20150101", "20150102", "20150105"], tz=tz_naive_fixture),
|
||
|
|
columns=["date"],
|
||
|
|
)
|
||
|
|
result = df1.where(mask, df2)
|
||
|
|
tm.assert_frame_equal(exp, result)
|
||
|
|
|
||
|
|
def test_df_where_change_dtype(self):
|
||
|
|
# GH#16979
|
||
|
|
df = DataFrame(np.arange(2 * 3).reshape(2, 3), columns=list("ABC"))
|
||
|
|
mask = np.array([[True, False, False], [False, False, True]])
|
||
|
|
|
||
|
|
result = df.where(mask)
|
||
|
|
expected = DataFrame(
|
||
|
|
[[0, np.nan, np.nan], [np.nan, np.nan, 5]], columns=list("ABC")
|
||
|
|
)
|
||
|
|
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("kwargs", [{}, {"other": None}])
|
||
|
|
def test_df_where_with_category(self, kwargs):
|
||
|
|
# GH#16979
|
||
|
|
df = DataFrame(np.arange(2 * 3).reshape(2, 3), columns=list("ABC"))
|
||
|
|
mask = np.array([[True, False, False], [False, False, True]])
|
||
|
|
|
||
|
|
# change type to category
|
||
|
|
df.A = df.A.astype("category")
|
||
|
|
df.B = df.B.astype("category")
|
||
|
|
df.C = df.C.astype("category")
|
||
|
|
|
||
|
|
result = df.where(mask, **kwargs)
|
||
|
|
A = pd.Categorical([0, np.nan], categories=[0, 3])
|
||
|
|
B = pd.Categorical([np.nan, np.nan], categories=[1, 4])
|
||
|
|
C = pd.Categorical([np.nan, 5], categories=[2, 5])
|
||
|
|
expected = DataFrame({"A": A, "B": B, "C": C})
|
||
|
|
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# Check Series.where while we're here
|
||
|
|
result = df.A.where(mask[:, 0], **kwargs)
|
||
|
|
expected = Series(A, name="A")
|
||
|
|
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
def test_where_categorical_filtering(self):
|
||
|
|
# GH#22609 Verify filtering operations on DataFrames with categorical Series
|
||
|
|
df = DataFrame(data=[[0, 0], [1, 1]], columns=["a", "b"])
|
||
|
|
df["b"] = df["b"].astype("category")
|
||
|
|
|
||
|
|
result = df.where(df["a"] > 0)
|
||
|
|
expected = df.copy()
|
||
|
|
expected.loc[0, :] = np.nan
|
||
|
|
|
||
|
|
tm.assert_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_where_copies_with_noop(frame_or_series):
|
||
|
|
# GH-39595
|
||
|
|
result = frame_or_series([1, 2, 3, 4])
|
||
|
|
expected = result.copy()
|
||
|
|
col = result[0] if frame_or_series is DataFrame else result
|
||
|
|
|
||
|
|
where_res = result.where(col < 5)
|
||
|
|
where_res *= 2
|
||
|
|
|
||
|
|
tm.assert_equal(result, expected)
|
||
|
|
|
||
|
|
where_res = result.where(col > 5, [1, 2, 3, 4])
|
||
|
|
where_res *= 2
|
||
|
|
|
||
|
|
tm.assert_equal(result, expected)
|