test_where.py

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(False)
            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 econd[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(
            dict(
                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 = pd.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", [dict(), dict(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)

    @pytest.mark.parametrize("inplace", [True, False])
    def test_where_nullable_boolean_mask(self, inplace):
        # https://github.com/pandas-dev/pandas/issues/35429
        df = DataFrame([1, 2, 3])
        mask = Series([True, False, None], dtype="boolean")
        expected = DataFrame([1, 999, 999])

        if inplace:
            result = df.copy()
            result.where(mask, 999, inplace=True)
        else:
            result = df.where(mask, 999, inplace=False)

        tm.assert_frame_equal(result, expected)