from collections import deque from datetime import ( datetime, timezone, ) from enum import Enum import functools import operator import re import numpy as np import pytest from pandas._config import using_pyarrow_string_dtype import pandas.util._test_decorators as td import pandas as pd from pandas import ( DataFrame, Index, MultiIndex, Series, ) import pandas._testing as tm from pandas.core.computation import expressions as expr from pandas.tests.frame.common import ( _check_mixed_float, _check_mixed_int, ) @pytest.fixture def simple_frame(): """ Fixture for simple 3x3 DataFrame Columns are ['one', 'two', 'three'], index is ['a', 'b', 'c']. one two three a 1.0 2.0 3.0 b 4.0 5.0 6.0 c 7.0 8.0 9.0 """ arr = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) return DataFrame(arr, columns=["one", "two", "three"], index=["a", "b", "c"]) @pytest.fixture(autouse=True, params=[0, 100], ids=["numexpr", "python"]) def switch_numexpr_min_elements(request, monkeypatch): with monkeypatch.context() as m: m.setattr(expr, "_MIN_ELEMENTS", request.param) yield request.param class DummyElement: def __init__(self, value, dtype) -> None: self.value = value self.dtype = np.dtype(dtype) def __array__(self, dtype=None, copy=None): return np.array(self.value, dtype=self.dtype) def __str__(self) -> str: return f"DummyElement({self.value}, {self.dtype})" def __repr__(self) -> str: return str(self) def astype(self, dtype, copy=False): self.dtype = dtype return self def view(self, dtype): return type(self)(self.value.view(dtype), dtype) def any(self, axis=None): return bool(self.value) # ------------------------------------------------------------------- # Comparisons class TestFrameComparisons: # Specifically _not_ flex-comparisons def test_comparison_with_categorical_dtype(self): # GH#12564 df = DataFrame({"A": ["foo", "bar", "baz"]}) exp = DataFrame({"A": [True, False, False]}) res = df == "foo" tm.assert_frame_equal(res, exp) # casting to categorical shouldn't affect the result df["A"] = df["A"].astype("category") res = df == "foo" tm.assert_frame_equal(res, exp) def test_frame_in_list(self): # GH#12689 this should raise at the DataFrame level, not blocks df = DataFrame( np.random.default_rng(2).standard_normal((6, 4)), columns=list("ABCD") ) msg = "The truth value of a DataFrame is ambiguous" with pytest.raises(ValueError, match=msg): df in [None] @pytest.mark.parametrize( "arg, arg2", [ [ { "a": np.random.default_rng(2).integers(10, size=10), "b": pd.date_range("20010101", periods=10), }, { "a": np.random.default_rng(2).integers(10, size=10), "b": np.random.default_rng(2).integers(10, size=10), }, ], [ { "a": np.random.default_rng(2).integers(10, size=10), "b": np.random.default_rng(2).integers(10, size=10), }, { "a": np.random.default_rng(2).integers(10, size=10), "b": pd.date_range("20010101", periods=10), }, ], [ { "a": pd.date_range("20010101", periods=10), "b": pd.date_range("20010101", periods=10), }, { "a": np.random.default_rng(2).integers(10, size=10), "b": np.random.default_rng(2).integers(10, size=10), }, ], [ { "a": np.random.default_rng(2).integers(10, size=10), "b": pd.date_range("20010101", periods=10), }, { "a": pd.date_range("20010101", periods=10), "b": pd.date_range("20010101", periods=10), }, ], ], ) def test_comparison_invalid(self, arg, arg2): # GH4968 # invalid date/int comparisons x = DataFrame(arg) y = DataFrame(arg2) # we expect the result to match Series comparisons for # == and !=, inequalities should raise result = x == y expected = DataFrame( {col: x[col] == y[col] for col in x.columns}, index=x.index, columns=x.columns, ) tm.assert_frame_equal(result, expected) result = x != y expected = DataFrame( {col: x[col] != y[col] for col in x.columns}, index=x.index, columns=x.columns, ) tm.assert_frame_equal(result, expected) msgs = [ r"Invalid comparison between dtype=datetime64\[ns\] and ndarray", "invalid type promotion", ( # npdev 1.20.0 r"The DTypes and " r" do not have a common DType." ), ] msg = "|".join(msgs) with pytest.raises(TypeError, match=msg): x >= y with pytest.raises(TypeError, match=msg): x > y with pytest.raises(TypeError, match=msg): x < y with pytest.raises(TypeError, match=msg): x <= y @pytest.mark.parametrize( "left, right", [ ("gt", "lt"), ("lt", "gt"), ("ge", "le"), ("le", "ge"), ("eq", "eq"), ("ne", "ne"), ], ) def test_timestamp_compare(self, left, right): # make sure we can compare Timestamps on the right AND left hand side # GH#4982 df = DataFrame( { "dates1": pd.date_range("20010101", periods=10), "dates2": pd.date_range("20010102", periods=10), "intcol": np.random.default_rng(2).integers(1000000000, size=10), "floatcol": np.random.default_rng(2).standard_normal(10), "stringcol": [chr(100 + i) for i in range(10)], } ) df.loc[np.random.default_rng(2).random(len(df)) > 0.5, "dates2"] = pd.NaT left_f = getattr(operator, left) right_f = getattr(operator, right) # no nats if left in ["eq", "ne"]: expected = left_f(df, pd.Timestamp("20010109")) result = right_f(pd.Timestamp("20010109"), df) tm.assert_frame_equal(result, expected) else: msg = ( "'(<|>)=?' not supported between " "instances of 'numpy.ndarray' and 'Timestamp'" ) with pytest.raises(TypeError, match=msg): left_f(df, pd.Timestamp("20010109")) with pytest.raises(TypeError, match=msg): right_f(pd.Timestamp("20010109"), df) # nats if left in ["eq", "ne"]: expected = left_f(df, pd.Timestamp("nat")) result = right_f(pd.Timestamp("nat"), df) tm.assert_frame_equal(result, expected) else: msg = ( "'(<|>)=?' not supported between " "instances of 'numpy.ndarray' and 'NaTType'" ) with pytest.raises(TypeError, match=msg): left_f(df, pd.Timestamp("nat")) with pytest.raises(TypeError, match=msg): right_f(pd.Timestamp("nat"), df) @pytest.mark.xfail( using_pyarrow_string_dtype(), reason="can't compare string and int" ) def test_mixed_comparison(self): # GH#13128, GH#22163 != datetime64 vs non-dt64 should be False, # not raise TypeError # (this appears to be fixed before GH#22163, not sure when) df = DataFrame([["1989-08-01", 1], ["1989-08-01", 2]]) other = DataFrame([["a", "b"], ["c", "d"]]) result = df == other assert not result.any().any() result = df != other assert result.all().all() def test_df_boolean_comparison_error(self): # GH#4576, GH#22880 # comparing DataFrame against list/tuple with len(obj) matching # len(df.columns) is supported as of GH#22800 df = DataFrame(np.arange(6).reshape((3, 2))) expected = DataFrame([[False, False], [True, False], [False, False]]) result = df == (2, 2) tm.assert_frame_equal(result, expected) result = df == [2, 2] tm.assert_frame_equal(result, expected) def test_df_float_none_comparison(self): df = DataFrame( np.random.default_rng(2).standard_normal((8, 3)), index=range(8), columns=["A", "B", "C"], ) result = df.__eq__(None) assert not result.any().any() def test_df_string_comparison(self): df = DataFrame([{"a": 1, "b": "foo"}, {"a": 2, "b": "bar"}]) mask_a = df.a > 1 tm.assert_frame_equal(df[mask_a], df.loc[1:1, :]) tm.assert_frame_equal(df[-mask_a], df.loc[0:0, :]) mask_b = df.b == "foo" tm.assert_frame_equal(df[mask_b], df.loc[0:0, :]) tm.assert_frame_equal(df[-mask_b], df.loc[1:1, :]) class TestFrameFlexComparisons: # TODO: test_bool_flex_frame needs a better name @pytest.mark.parametrize("op", ["eq", "ne", "gt", "lt", "ge", "le"]) def test_bool_flex_frame(self, op): data = np.random.default_rng(2).standard_normal((5, 3)) other_data = np.random.default_rng(2).standard_normal((5, 3)) df = DataFrame(data) other = DataFrame(other_data) ndim_5 = np.ones(df.shape + (1, 3)) # DataFrame assert df.eq(df).values.all() assert not df.ne(df).values.any() f = getattr(df, op) o = getattr(operator, op) # No NAs tm.assert_frame_equal(f(other), o(df, other)) # Unaligned part_o = other.loc[3:, 1:].copy() rs = f(part_o) xp = o(df, part_o.reindex(index=df.index, columns=df.columns)) tm.assert_frame_equal(rs, xp) # ndarray tm.assert_frame_equal(f(other.values), o(df, other.values)) # scalar tm.assert_frame_equal(f(0), o(df, 0)) # NAs msg = "Unable to coerce to Series/DataFrame" tm.assert_frame_equal(f(np.nan), o(df, np.nan)) with pytest.raises(ValueError, match=msg): f(ndim_5) @pytest.mark.parametrize("box", [np.array, Series]) def test_bool_flex_series(self, box): # Series # list/tuple data = np.random.default_rng(2).standard_normal((5, 3)) df = DataFrame(data) idx_ser = box(np.random.default_rng(2).standard_normal(5)) col_ser = box(np.random.default_rng(2).standard_normal(3)) idx_eq = df.eq(idx_ser, axis=0) col_eq = df.eq(col_ser) idx_ne = df.ne(idx_ser, axis=0) col_ne = df.ne(col_ser) tm.assert_frame_equal(col_eq, df == Series(col_ser)) tm.assert_frame_equal(col_eq, -col_ne) tm.assert_frame_equal(idx_eq, -idx_ne) tm.assert_frame_equal(idx_eq, df.T.eq(idx_ser).T) tm.assert_frame_equal(col_eq, df.eq(list(col_ser))) tm.assert_frame_equal(idx_eq, df.eq(Series(idx_ser), axis=0)) tm.assert_frame_equal(idx_eq, df.eq(list(idx_ser), axis=0)) idx_gt = df.gt(idx_ser, axis=0) col_gt = df.gt(col_ser) idx_le = df.le(idx_ser, axis=0) col_le = df.le(col_ser) tm.assert_frame_equal(col_gt, df > Series(col_ser)) tm.assert_frame_equal(col_gt, -col_le) tm.assert_frame_equal(idx_gt, -idx_le) tm.assert_frame_equal(idx_gt, df.T.gt(idx_ser).T) idx_ge = df.ge(idx_ser, axis=0) col_ge = df.ge(col_ser) idx_lt = df.lt(idx_ser, axis=0) col_lt = df.lt(col_ser) tm.assert_frame_equal(col_ge, df >= Series(col_ser)) tm.assert_frame_equal(col_ge, -col_lt) tm.assert_frame_equal(idx_ge, -idx_lt) tm.assert_frame_equal(idx_ge, df.T.ge(idx_ser).T) idx_ser = Series(np.random.default_rng(2).standard_normal(5)) col_ser = Series(np.random.default_rng(2).standard_normal(3)) def test_bool_flex_frame_na(self): df = DataFrame(np.random.default_rng(2).standard_normal((5, 3))) # NA df.loc[0, 0] = np.nan rs = df.eq(df) assert not rs.loc[0, 0] rs = df.ne(df) assert rs.loc[0, 0] rs = df.gt(df) assert not rs.loc[0, 0] rs = df.lt(df) assert not rs.loc[0, 0] rs = df.ge(df) assert not rs.loc[0, 0] rs = df.le(df) assert not rs.loc[0, 0] def test_bool_flex_frame_complex_dtype(self): # complex arr = np.array([np.nan, 1, 6, np.nan]) arr2 = np.array([2j, np.nan, 7, None]) df = DataFrame({"a": arr}) df2 = DataFrame({"a": arr2}) msg = "|".join( [ "'>' not supported between instances of '.*' and 'complex'", r"unorderable types: .*complex\(\)", # PY35 ] ) with pytest.raises(TypeError, match=msg): # inequalities are not well-defined for complex numbers df.gt(df2) with pytest.raises(TypeError, match=msg): # regression test that we get the same behavior for Series df["a"].gt(df2["a"]) with pytest.raises(TypeError, match=msg): # Check that we match numpy behavior here df.values > df2.values rs = df.ne(df2) assert rs.values.all() arr3 = np.array([2j, np.nan, None]) df3 = DataFrame({"a": arr3}) with pytest.raises(TypeError, match=msg): # inequalities are not well-defined for complex numbers df3.gt(2j) with pytest.raises(TypeError, match=msg): # regression test that we get the same behavior for Series df3["a"].gt(2j) with pytest.raises(TypeError, match=msg): # Check that we match numpy behavior here df3.values > 2j def test_bool_flex_frame_object_dtype(self): # corner, dtype=object df1 = DataFrame({"col": ["foo", np.nan, "bar"]}, dtype=object) df2 = DataFrame({"col": ["foo", datetime.now(), "bar"]}, dtype=object) result = df1.ne(df2) exp = DataFrame({"col": [False, True, False]}) tm.assert_frame_equal(result, exp) def test_flex_comparison_nat(self): # GH 15697, GH 22163 df.eq(pd.NaT) should behave like df == pd.NaT, # and _definitely_ not be NaN df = DataFrame([pd.NaT]) result = df == pd.NaT # result.iloc[0, 0] is a np.bool_ object assert result.iloc[0, 0].item() is False result = df.eq(pd.NaT) assert result.iloc[0, 0].item() is False result = df != pd.NaT assert result.iloc[0, 0].item() is True result = df.ne(pd.NaT) assert result.iloc[0, 0].item() is True @pytest.mark.parametrize("opname", ["eq", "ne", "gt", "lt", "ge", "le"]) def test_df_flex_cmp_constant_return_types(self, opname): # GH 15077, non-empty DataFrame df = DataFrame({"x": [1, 2, 3], "y": [1.0, 2.0, 3.0]}) const = 2 result = getattr(df, opname)(const).dtypes.value_counts() tm.assert_series_equal( result, Series([2], index=[np.dtype(bool)], name="count") ) @pytest.mark.parametrize("opname", ["eq", "ne", "gt", "lt", "ge", "le"]) def test_df_flex_cmp_constant_return_types_empty(self, opname): # GH 15077 empty DataFrame df = DataFrame({"x": [1, 2, 3], "y": [1.0, 2.0, 3.0]}) const = 2 empty = df.iloc[:0] result = getattr(empty, opname)(const).dtypes.value_counts() tm.assert_series_equal( result, Series([2], index=[np.dtype(bool)], name="count") ) def test_df_flex_cmp_ea_dtype_with_ndarray_series(self): ii = pd.IntervalIndex.from_breaks([1, 2, 3]) df = DataFrame({"A": ii, "B": ii}) ser = Series([0, 0]) res = df.eq(ser, axis=0) expected = DataFrame({"A": [False, False], "B": [False, False]}) tm.assert_frame_equal(res, expected) ser2 = Series([1, 2], index=["A", "B"]) res2 = df.eq(ser2, axis=1) tm.assert_frame_equal(res2, expected) # ------------------------------------------------------------------- # Arithmetic class TestFrameFlexArithmetic: def test_floordiv_axis0(self): # make sure we df.floordiv(ser, axis=0) matches column-wise result arr = np.arange(3) ser = Series(arr) df = DataFrame({"A": ser, "B": ser}) result = df.floordiv(ser, axis=0) expected = DataFrame({col: df[col] // ser for col in df.columns}) tm.assert_frame_equal(result, expected) result2 = df.floordiv(ser.values, axis=0) tm.assert_frame_equal(result2, expected) def test_df_add_td64_columnwise(self): # GH 22534 Check that column-wise addition broadcasts correctly dti = pd.date_range("2016-01-01", periods=10) tdi = pd.timedelta_range("1", periods=10) tser = Series(tdi) df = DataFrame({0: dti, 1: tdi}) result = df.add(tser, axis=0) expected = DataFrame({0: dti + tdi, 1: tdi + tdi}) tm.assert_frame_equal(result, expected) def test_df_add_flex_filled_mixed_dtypes(self): # GH 19611 dti = pd.date_range("2016-01-01", periods=3) ser = Series(["1 Day", "NaT", "2 Days"], dtype="timedelta64[ns]") df = DataFrame({"A": dti, "B": ser}) other = DataFrame({"A": ser, "B": ser}) fill = pd.Timedelta(days=1).to_timedelta64() result = df.add(other, fill_value=fill) expected = DataFrame( { "A": Series( ["2016-01-02", "2016-01-03", "2016-01-05"], dtype="datetime64[ns]" ), "B": ser * 2, } ) tm.assert_frame_equal(result, expected) def test_arith_flex_frame( self, all_arithmetic_operators, float_frame, mixed_float_frame ): # one instance of parametrized fixture op = all_arithmetic_operators def f(x, y): # r-versions not in operator-stdlib; get op without "r" and invert if op.startswith("__r"): return getattr(operator, op.replace("__r", "__"))(y, x) return getattr(operator, op)(x, y) result = getattr(float_frame, op)(2 * float_frame) expected = f(float_frame, 2 * float_frame) tm.assert_frame_equal(result, expected) # vs mix float result = getattr(mixed_float_frame, op)(2 * mixed_float_frame) expected = f(mixed_float_frame, 2 * mixed_float_frame) tm.assert_frame_equal(result, expected) _check_mixed_float(result, dtype={"C": None}) @pytest.mark.parametrize("op", ["__add__", "__sub__", "__mul__"]) def test_arith_flex_frame_mixed( self, op, int_frame, mixed_int_frame, mixed_float_frame, switch_numexpr_min_elements, ): f = getattr(operator, op) # vs mix int result = getattr(mixed_int_frame, op)(2 + mixed_int_frame) expected = f(mixed_int_frame, 2 + mixed_int_frame) # no overflow in the uint dtype = None if op in ["__sub__"]: dtype = {"B": "uint64", "C": None} elif op in ["__add__", "__mul__"]: dtype = {"C": None} if expr.USE_NUMEXPR and switch_numexpr_min_elements == 0: # when using numexpr, the casting rules are slightly different: # in the `2 + mixed_int_frame` operation, int32 column becomes # and int64 column (not preserving dtype in operation with Python # scalar), and then the int32/int64 combo results in int64 result dtype["A"] = (2 + mixed_int_frame)["A"].dtype tm.assert_frame_equal(result, expected) _check_mixed_int(result, dtype=dtype) # vs mix float result = getattr(mixed_float_frame, op)(2 * mixed_float_frame) expected = f(mixed_float_frame, 2 * mixed_float_frame) tm.assert_frame_equal(result, expected) _check_mixed_float(result, dtype={"C": None}) # vs plain int result = getattr(int_frame, op)(2 * int_frame) expected = f(int_frame, 2 * int_frame) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("dim", range(3, 6)) def test_arith_flex_frame_raise(self, all_arithmetic_operators, float_frame, dim): # one instance of parametrized fixture op = all_arithmetic_operators # Check that arrays with dim >= 3 raise arr = np.ones((1,) * dim) msg = "Unable to coerce to Series/DataFrame" with pytest.raises(ValueError, match=msg): getattr(float_frame, op)(arr) def test_arith_flex_frame_corner(self, float_frame): const_add = float_frame.add(1) tm.assert_frame_equal(const_add, float_frame + 1) # corner cases result = float_frame.add(float_frame[:0]) expected = float_frame.sort_index() * np.nan tm.assert_frame_equal(result, expected) result = float_frame[:0].add(float_frame) expected = float_frame.sort_index() * np.nan tm.assert_frame_equal(result, expected) with pytest.raises(NotImplementedError, match="fill_value"): float_frame.add(float_frame.iloc[0], fill_value=3) with pytest.raises(NotImplementedError, match="fill_value"): float_frame.add(float_frame.iloc[0], axis="index", fill_value=3) @pytest.mark.parametrize("op", ["add", "sub", "mul", "mod"]) def test_arith_flex_series_ops(self, simple_frame, op): # after arithmetic refactor, add truediv here df = simple_frame row = df.xs("a") col = df["two"] f = getattr(df, op) op = getattr(operator, op) tm.assert_frame_equal(f(row), op(df, row)) tm.assert_frame_equal(f(col, axis=0), op(df.T, col).T) def test_arith_flex_series(self, simple_frame): df = simple_frame row = df.xs("a") col = df["two"] # special case for some reason tm.assert_frame_equal(df.add(row, axis=None), df + row) # cases which will be refactored after big arithmetic refactor tm.assert_frame_equal(df.div(row), df / row) tm.assert_frame_equal(df.div(col, axis=0), (df.T / col).T) @pytest.mark.parametrize("dtype", ["int64", "float64"]) def test_arith_flex_series_broadcasting(self, dtype): # broadcasting issue in GH 7325 df = DataFrame(np.arange(3 * 2).reshape((3, 2)), dtype=dtype) expected = DataFrame([[np.nan, np.inf], [1.0, 1.5], [1.0, 1.25]]) result = df.div(df[0], axis="index") tm.assert_frame_equal(result, expected) def test_arith_flex_zero_len_raises(self): # GH 19522 passing fill_value to frame flex arith methods should # raise even in the zero-length special cases ser_len0 = Series([], dtype=object) df_len0 = DataFrame(columns=["A", "B"]) df = DataFrame([[1, 2], [3, 4]], columns=["A", "B"]) with pytest.raises(NotImplementedError, match="fill_value"): df.add(ser_len0, fill_value="E") with pytest.raises(NotImplementedError, match="fill_value"): df_len0.sub(df["A"], axis=None, fill_value=3) def test_flex_add_scalar_fill_value(self): # GH#12723 dat = np.array([0, 1, np.nan, 3, 4, 5], dtype="float") df = DataFrame({"foo": dat}, index=range(6)) exp = df.fillna(0).add(2) res = df.add(2, fill_value=0) tm.assert_frame_equal(res, exp) def test_sub_alignment_with_duplicate_index(self): # GH#5185 dup aligning operations should work df1 = DataFrame([1, 2, 3, 4, 5], index=[1, 2, 1, 2, 3]) df2 = DataFrame([1, 2, 3], index=[1, 2, 3]) expected = DataFrame([0, 2, 0, 2, 2], index=[1, 1, 2, 2, 3]) result = df1.sub(df2) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("op", ["__add__", "__mul__", "__sub__", "__truediv__"]) def test_arithmetic_with_duplicate_columns(self, op): # operations df = DataFrame({"A": np.arange(10), "B": np.random.default_rng(2).random(10)}) expected = getattr(df, op)(df) expected.columns = ["A", "A"] df.columns = ["A", "A"] result = getattr(df, op)(df) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("level", [0, None]) def test_broadcast_multiindex(self, level): # GH34388 df1 = DataFrame({"A": [0, 1, 2], "B": [1, 2, 3]}) df1.columns = df1.columns.set_names("L1") df2 = DataFrame({("A", "C"): [0, 0, 0], ("A", "D"): [0, 0, 0]}) df2.columns = df2.columns.set_names(["L1", "L2"]) result = df1.add(df2, level=level) expected = DataFrame({("A", "C"): [0, 1, 2], ("A", "D"): [0, 1, 2]}) expected.columns = expected.columns.set_names(["L1", "L2"]) tm.assert_frame_equal(result, expected) def test_frame_multiindex_operations(self): # GH 43321 df = DataFrame( {2010: [1, 2, 3], 2020: [3, 4, 5]}, index=MultiIndex.from_product( [["a"], ["b"], [0, 1, 2]], names=["scen", "mod", "id"] ), ) series = Series( [0.4], index=MultiIndex.from_product([["b"], ["a"]], names=["mod", "scen"]), ) expected = DataFrame( {2010: [1.4, 2.4, 3.4], 2020: [3.4, 4.4, 5.4]}, index=MultiIndex.from_product( [["a"], ["b"], [0, 1, 2]], names=["scen", "mod", "id"] ), ) result = df.add(series, axis=0) tm.assert_frame_equal(result, expected) def test_frame_multiindex_operations_series_index_to_frame_index(self): # GH 43321 df = DataFrame( {2010: [1], 2020: [3]}, index=MultiIndex.from_product([["a"], ["b"]], names=["scen", "mod"]), ) series = Series( [10.0, 20.0, 30.0], index=MultiIndex.from_product( [["a"], ["b"], [0, 1, 2]], names=["scen", "mod", "id"] ), ) expected = DataFrame( {2010: [11.0, 21, 31.0], 2020: [13.0, 23.0, 33.0]}, index=MultiIndex.from_product( [["a"], ["b"], [0, 1, 2]], names=["scen", "mod", "id"] ), ) result = df.add(series, axis=0) tm.assert_frame_equal(result, expected) def test_frame_multiindex_operations_no_align(self): df = DataFrame( {2010: [1, 2, 3], 2020: [3, 4, 5]}, index=MultiIndex.from_product( [["a"], ["b"], [0, 1, 2]], names=["scen", "mod", "id"] ), ) series = Series( [0.4], index=MultiIndex.from_product([["c"], ["a"]], names=["mod", "scen"]), ) expected = DataFrame( {2010: np.nan, 2020: np.nan}, index=MultiIndex.from_tuples( [ ("a", "b", 0), ("a", "b", 1), ("a", "b", 2), ("a", "c", np.nan), ], names=["scen", "mod", "id"], ), ) result = df.add(series, axis=0) tm.assert_frame_equal(result, expected) def test_frame_multiindex_operations_part_align(self): df = DataFrame( {2010: [1, 2, 3], 2020: [3, 4, 5]}, index=MultiIndex.from_tuples( [ ("a", "b", 0), ("a", "b", 1), ("a", "c", 2), ], names=["scen", "mod", "id"], ), ) series = Series( [0.4], index=MultiIndex.from_product([["b"], ["a"]], names=["mod", "scen"]), ) expected = DataFrame( {2010: [1.4, 2.4, np.nan], 2020: [3.4, 4.4, np.nan]}, index=MultiIndex.from_tuples( [ ("a", "b", 0), ("a", "b", 1), ("a", "c", 2), ], names=["scen", "mod", "id"], ), ) result = df.add(series, axis=0) tm.assert_frame_equal(result, expected) class TestFrameArithmetic: def test_td64_op_nat_casting(self): # Make sure we don't accidentally treat timedelta64(NaT) as datetime64 # when calling dispatch_to_series in DataFrame arithmetic ser = Series(["NaT", "NaT"], dtype="timedelta64[ns]") df = DataFrame([[1, 2], [3, 4]]) result = df * ser expected = DataFrame({0: ser, 1: ser}) tm.assert_frame_equal(result, expected) def test_df_add_2d_array_rowlike_broadcasts(self): # GH#23000 arr = np.arange(6).reshape(3, 2) df = DataFrame(arr, columns=[True, False], index=["A", "B", "C"]) rowlike = arr[[1], :] # shape --> (1, ncols) assert rowlike.shape == (1, df.shape[1]) expected = DataFrame( [[2, 4], [4, 6], [6, 8]], columns=df.columns, index=df.index, # specify dtype explicitly to avoid failing # on 32bit builds dtype=arr.dtype, ) result = df + rowlike tm.assert_frame_equal(result, expected) result = rowlike + df tm.assert_frame_equal(result, expected) def test_df_add_2d_array_collike_broadcasts(self): # GH#23000 arr = np.arange(6).reshape(3, 2) df = DataFrame(arr, columns=[True, False], index=["A", "B", "C"]) collike = arr[:, [1]] # shape --> (nrows, 1) assert collike.shape == (df.shape[0], 1) expected = DataFrame( [[1, 2], [5, 6], [9, 10]], columns=df.columns, index=df.index, # specify dtype explicitly to avoid failing # on 32bit builds dtype=arr.dtype, ) result = df + collike tm.assert_frame_equal(result, expected) result = collike + df tm.assert_frame_equal(result, expected) def test_df_arith_2d_array_rowlike_broadcasts( self, request, all_arithmetic_operators, using_array_manager ): # GH#23000 opname = all_arithmetic_operators if using_array_manager and opname in ("__rmod__", "__rfloordiv__"): # TODO(ArrayManager) decide on dtypes td.mark_array_manager_not_yet_implemented(request) arr = np.arange(6).reshape(3, 2) df = DataFrame(arr, columns=[True, False], index=["A", "B", "C"]) rowlike = arr[[1], :] # shape --> (1, ncols) assert rowlike.shape == (1, df.shape[1]) exvals = [ getattr(df.loc["A"], opname)(rowlike.squeeze()), getattr(df.loc["B"], opname)(rowlike.squeeze()), getattr(df.loc["C"], opname)(rowlike.squeeze()), ] expected = DataFrame(exvals, columns=df.columns, index=df.index) result = getattr(df, opname)(rowlike) tm.assert_frame_equal(result, expected) def test_df_arith_2d_array_collike_broadcasts( self, request, all_arithmetic_operators, using_array_manager ): # GH#23000 opname = all_arithmetic_operators if using_array_manager and opname in ("__rmod__", "__rfloordiv__"): # TODO(ArrayManager) decide on dtypes td.mark_array_manager_not_yet_implemented(request) arr = np.arange(6).reshape(3, 2) df = DataFrame(arr, columns=[True, False], index=["A", "B", "C"]) collike = arr[:, [1]] # shape --> (nrows, 1) assert collike.shape == (df.shape[0], 1) exvals = { True: getattr(df[True], opname)(collike.squeeze()), False: getattr(df[False], opname)(collike.squeeze()), } dtype = None if opname in ["__rmod__", "__rfloordiv__"]: # Series ops may return mixed int/float dtypes in cases where # DataFrame op will return all-float. So we upcast `expected` dtype = np.common_type(*(x.values for x in exvals.values())) expected = DataFrame(exvals, columns=df.columns, index=df.index, dtype=dtype) result = getattr(df, opname)(collike) tm.assert_frame_equal(result, expected) def test_df_bool_mul_int(self): # GH 22047, GH 22163 multiplication by 1 should result in int dtype, # not object dtype df = DataFrame([[False, True], [False, False]]) result = df * 1 # On appveyor this comes back as np.int32 instead of np.int64, # so we check dtype.kind instead of just dtype kinds = result.dtypes.apply(lambda x: x.kind) assert (kinds == "i").all() result = 1 * df kinds = result.dtypes.apply(lambda x: x.kind) assert (kinds == "i").all() def test_arith_mixed(self): left = DataFrame({"A": ["a", "b", "c"], "B": [1, 2, 3]}) result = left + left expected = DataFrame({"A": ["aa", "bb", "cc"], "B": [2, 4, 6]}) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("col", ["A", "B"]) def test_arith_getitem_commute(self, all_arithmetic_functions, col): df = DataFrame({"A": [1.1, 3.3], "B": [2.5, -3.9]}) result = all_arithmetic_functions(df, 1)[col] expected = all_arithmetic_functions(df[col], 1) tm.assert_series_equal(result, expected) @pytest.mark.parametrize( "values", [[1, 2], (1, 2), np.array([1, 2]), range(1, 3), deque([1, 2])] ) def test_arith_alignment_non_pandas_object(self, values): # GH#17901 df = DataFrame({"A": [1, 1], "B": [1, 1]}) expected = DataFrame({"A": [2, 2], "B": [3, 3]}) result = df + values tm.assert_frame_equal(result, expected) def test_arith_non_pandas_object(self): df = DataFrame( np.arange(1, 10, dtype="f8").reshape(3, 3), columns=["one", "two", "three"], index=["a", "b", "c"], ) val1 = df.xs("a").values added = DataFrame(df.values + val1, index=df.index, columns=df.columns) tm.assert_frame_equal(df + val1, added) added = DataFrame((df.values.T + val1).T, index=df.index, columns=df.columns) tm.assert_frame_equal(df.add(val1, axis=0), added) val2 = list(df["two"]) added = DataFrame(df.values + val2, index=df.index, columns=df.columns) tm.assert_frame_equal(df + val2, added) added = DataFrame((df.values.T + val2).T, index=df.index, columns=df.columns) tm.assert_frame_equal(df.add(val2, axis="index"), added) val3 = np.random.default_rng(2).random(df.shape) added = DataFrame(df.values + val3, index=df.index, columns=df.columns) tm.assert_frame_equal(df.add(val3), added) def test_operations_with_interval_categories_index(self, all_arithmetic_operators): # GH#27415 op = all_arithmetic_operators ind = pd.CategoricalIndex(pd.interval_range(start=0.0, end=2.0)) data = [1, 2] df = DataFrame([data], columns=ind) num = 10 result = getattr(df, op)(num) expected = DataFrame([[getattr(n, op)(num) for n in data]], columns=ind) tm.assert_frame_equal(result, expected) def test_frame_with_frame_reindex(self): # GH#31623 df = DataFrame( { "foo": [pd.Timestamp("2019"), pd.Timestamp("2020")], "bar": [pd.Timestamp("2018"), pd.Timestamp("2021")], }, columns=["foo", "bar"], dtype="M8[ns]", ) df2 = df[["foo"]] result = df - df2 expected = DataFrame( {"foo": [pd.Timedelta(0), pd.Timedelta(0)], "bar": [np.nan, np.nan]}, columns=["bar", "foo"], ) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( "value, dtype", [ (1, "i8"), (1.0, "f8"), (2**63, "f8"), (1j, "complex128"), (2**63, "complex128"), (True, "bool"), (np.timedelta64(20, "ns"), " b tm.assert_frame_equal(result, expected) result = df.values > b tm.assert_numpy_array_equal(result, expected.values) msg1d = "Unable to coerce to Series, length must be 2: given 3" msg2d = "Unable to coerce to DataFrame, shape must be" msg2db = "operands could not be broadcast together with shapes" with pytest.raises(ValueError, match=msg1d): # wrong shape df > lst with pytest.raises(ValueError, match=msg1d): # wrong shape df > tup # broadcasts like ndarray (GH#23000) result = df > b_r tm.assert_frame_equal(result, expected) result = df.values > b_r tm.assert_numpy_array_equal(result, expected.values) with pytest.raises(ValueError, match=msg2d): df > b_c with pytest.raises(ValueError, match=msg2db): df.values > b_c # == expected = DataFrame([[False, False], [True, False], [False, False]]) result = df == b tm.assert_frame_equal(result, expected) with pytest.raises(ValueError, match=msg1d): df == lst with pytest.raises(ValueError, match=msg1d): df == tup # broadcasts like ndarray (GH#23000) result = df == b_r tm.assert_frame_equal(result, expected) result = df.values == b_r tm.assert_numpy_array_equal(result, expected.values) with pytest.raises(ValueError, match=msg2d): df == b_c assert df.values.shape != b_c.shape # with alignment df = DataFrame( np.arange(6).reshape((3, 2)), columns=list("AB"), index=list("abc") ) expected.index = df.index expected.columns = df.columns with pytest.raises(ValueError, match=msg1d): df == lst with pytest.raises(ValueError, match=msg1d): df == tup def test_inplace_ops_alignment(self): # inplace ops / ops alignment # GH 8511 columns = list("abcdefg") X_orig = DataFrame( np.arange(10 * len(columns)).reshape(-1, len(columns)), columns=columns, index=range(10), ) Z = 100 * X_orig.iloc[:, 1:-1].copy() block1 = list("bedcf") subs = list("bcdef") # add X = X_orig.copy() result1 = (X[block1] + Z).reindex(columns=subs) X[block1] += Z result2 = X.reindex(columns=subs) X = X_orig.copy() result3 = (X[block1] + Z[block1]).reindex(columns=subs) X[block1] += Z[block1] result4 = X.reindex(columns=subs) tm.assert_frame_equal(result1, result2) tm.assert_frame_equal(result1, result3) tm.assert_frame_equal(result1, result4) # sub X = X_orig.copy() result1 = (X[block1] - Z).reindex(columns=subs) X[block1] -= Z result2 = X.reindex(columns=subs) X = X_orig.copy() result3 = (X[block1] - Z[block1]).reindex(columns=subs) X[block1] -= Z[block1] result4 = X.reindex(columns=subs) tm.assert_frame_equal(result1, result2) tm.assert_frame_equal(result1, result3) tm.assert_frame_equal(result1, result4) def test_inplace_ops_identity(self): # GH 5104 # make sure that we are actually changing the object s_orig = Series([1, 2, 3]) df_orig = DataFrame( np.random.default_rng(2).integers(0, 5, size=10).reshape(-1, 5) ) # no dtype change s = s_orig.copy() s2 = s s += 1 tm.assert_series_equal(s, s2) tm.assert_series_equal(s_orig + 1, s) assert s is s2 assert s._mgr is s2._mgr df = df_orig.copy() df2 = df df += 1 tm.assert_frame_equal(df, df2) tm.assert_frame_equal(df_orig + 1, df) assert df is df2 assert df._mgr is df2._mgr # dtype change s = s_orig.copy() s2 = s s += 1.5 tm.assert_series_equal(s, s2) tm.assert_series_equal(s_orig + 1.5, s) df = df_orig.copy() df2 = df df += 1.5 tm.assert_frame_equal(df, df2) tm.assert_frame_equal(df_orig + 1.5, df) assert df is df2 assert df._mgr is df2._mgr # mixed dtype arr = np.random.default_rng(2).integers(0, 10, size=5) df_orig = DataFrame({"A": arr.copy(), "B": "foo"}) df = df_orig.copy() df2 = df df["A"] += 1 expected = DataFrame({"A": arr.copy() + 1, "B": "foo"}) tm.assert_frame_equal(df, expected) tm.assert_frame_equal(df2, expected) assert df._mgr is df2._mgr df = df_orig.copy() df2 = df df["A"] += 1.5 expected = DataFrame({"A": arr.copy() + 1.5, "B": "foo"}) tm.assert_frame_equal(df, expected) tm.assert_frame_equal(df2, expected) assert df._mgr is df2._mgr @pytest.mark.parametrize( "op", [ "add", "and", pytest.param( "div", marks=pytest.mark.xfail( raises=AttributeError, reason="__idiv__ not implemented" ), ), "floordiv", "mod", "mul", "or", "pow", "sub", "truediv", "xor", ], ) def test_inplace_ops_identity2(self, op): df = DataFrame({"a": [1.0, 2.0, 3.0], "b": [1, 2, 3]}) operand = 2 if op in ("and", "or", "xor"): # cannot use floats for boolean ops df["a"] = [True, False, True] df_copy = df.copy() iop = f"__i{op}__" op = f"__{op}__" # no id change and value is correct getattr(df, iop)(operand) expected = getattr(df_copy, op)(operand) tm.assert_frame_equal(df, expected) expected = id(df) assert id(df) == expected @pytest.mark.parametrize( "val", [ [1, 2, 3], (1, 2, 3), np.array([1, 2, 3], dtype=np.int64), range(1, 4), ], ) def test_alignment_non_pandas(self, val): index = ["A", "B", "C"] columns = ["X", "Y", "Z"] df = DataFrame( np.random.default_rng(2).standard_normal((3, 3)), index=index, columns=columns, ) align = DataFrame._align_for_op expected = DataFrame({"X": val, "Y": val, "Z": val}, index=df.index) tm.assert_frame_equal(align(df, val, axis=0)[1], expected) expected = DataFrame( {"X": [1, 1, 1], "Y": [2, 2, 2], "Z": [3, 3, 3]}, index=df.index ) tm.assert_frame_equal(align(df, val, axis=1)[1], expected) @pytest.mark.parametrize("val", [[1, 2], (1, 2), np.array([1, 2]), range(1, 3)]) def test_alignment_non_pandas_length_mismatch(self, val): index = ["A", "B", "C"] columns = ["X", "Y", "Z"] df = DataFrame( np.random.default_rng(2).standard_normal((3, 3)), index=index, columns=columns, ) align = DataFrame._align_for_op # length mismatch msg = "Unable to coerce to Series, length must be 3: given 2" with pytest.raises(ValueError, match=msg): align(df, val, axis=0) with pytest.raises(ValueError, match=msg): align(df, val, axis=1) def test_alignment_non_pandas_index_columns(self): index = ["A", "B", "C"] columns = ["X", "Y", "Z"] df = DataFrame( np.random.default_rng(2).standard_normal((3, 3)), index=index, columns=columns, ) align = DataFrame._align_for_op val = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) tm.assert_frame_equal( align(df, val, axis=0)[1], DataFrame(val, index=df.index, columns=df.columns), ) tm.assert_frame_equal( align(df, val, axis=1)[1], DataFrame(val, index=df.index, columns=df.columns), ) # shape mismatch msg = "Unable to coerce to DataFrame, shape must be" val = np.array([[1, 2, 3], [4, 5, 6]]) with pytest.raises(ValueError, match=msg): align(df, val, axis=0) with pytest.raises(ValueError, match=msg): align(df, val, axis=1) val = np.zeros((3, 3, 3)) msg = re.escape( "Unable to coerce to Series/DataFrame, dimension must be <= 2: (3, 3, 3)" ) with pytest.raises(ValueError, match=msg): align(df, val, axis=0) with pytest.raises(ValueError, match=msg): align(df, val, axis=1) def test_no_warning(self, all_arithmetic_operators): df = DataFrame({"A": [0.0, 0.0], "B": [0.0, None]}) b = df["B"] with tm.assert_produces_warning(None): getattr(df, all_arithmetic_operators)(b) def test_dunder_methods_binary(self, all_arithmetic_operators): # GH#??? frame.__foo__ should only accept one argument df = DataFrame({"A": [0.0, 0.0], "B": [0.0, None]}) b = df["B"] with pytest.raises(TypeError, match="takes 2 positional arguments"): getattr(df, all_arithmetic_operators)(b, 0) def test_align_int_fill_bug(self): # GH#910 X = np.arange(10 * 10, dtype="float64").reshape(10, 10) Y = np.ones((10, 1), dtype=int) df1 = DataFrame(X) df1["0.X"] = Y.squeeze() df2 = df1.astype(float) result = df1 - df1.mean() expected = df2 - df2.mean() tm.assert_frame_equal(result, expected) def test_pow_with_realignment(): # GH#32685 pow has special semantics for operating with null values left = DataFrame({"A": [0, 1, 2]}) right = DataFrame(index=[0, 1, 2]) result = left**right expected = DataFrame({"A": [np.nan, 1.0, np.nan]}) tm.assert_frame_equal(result, expected) def test_dataframe_series_extension_dtypes(): # https://github.com/pandas-dev/pandas/issues/34311 df = DataFrame( np.random.default_rng(2).integers(0, 100, (10, 3)), columns=["a", "b", "c"] ) ser = Series([1, 2, 3], index=["a", "b", "c"]) expected = df.to_numpy("int64") + ser.to_numpy("int64").reshape(-1, 3) expected = DataFrame(expected, columns=df.columns, dtype="Int64") df_ea = df.astype("Int64") result = df_ea + ser tm.assert_frame_equal(result, expected) result = df_ea + ser.astype("Int64") tm.assert_frame_equal(result, expected) def test_dataframe_blockwise_slicelike(): # GH#34367 arr = np.random.default_rng(2).integers(0, 1000, (100, 10)) df1 = DataFrame(arr) # Explicit cast to float to avoid implicit cast when setting nan df2 = df1.copy().astype({1: "float", 3: "float", 7: "float"}) df2.iloc[0, [1, 3, 7]] = np.nan # Explicit cast to float to avoid implicit cast when setting nan df3 = df1.copy().astype({5: "float"}) df3.iloc[0, [5]] = np.nan # Explicit cast to float to avoid implicit cast when setting nan df4 = df1.copy().astype({2: "float", 3: "float", 4: "float"}) df4.iloc[0, np.arange(2, 5)] = np.nan # Explicit cast to float to avoid implicit cast when setting nan df5 = df1.copy().astype({4: "float", 5: "float", 6: "float"}) df5.iloc[0, np.arange(4, 7)] = np.nan for left, right in [(df1, df2), (df2, df3), (df4, df5)]: res = left + right expected = DataFrame({i: left[i] + right[i] for i in left.columns}) tm.assert_frame_equal(res, expected) @pytest.mark.parametrize( "df, col_dtype", [ (DataFrame([[1.0, 2.0], [4.0, 5.0]], columns=list("ab")), "float64"), ( DataFrame([[1.0, "b"], [4.0, "b"]], columns=list("ab")).astype( {"b": object} ), "object", ), ], ) def test_dataframe_operation_with_non_numeric_types(df, col_dtype): # GH #22663 expected = DataFrame([[0.0, np.nan], [3.0, np.nan]], columns=list("ab")) expected = expected.astype({"b": col_dtype}) result = df + Series([-1.0], index=list("a")) tm.assert_frame_equal(result, expected) def test_arith_reindex_with_duplicates(): # https://github.com/pandas-dev/pandas/issues/35194 df1 = DataFrame(data=[[0]], columns=["second"]) df2 = DataFrame(data=[[0, 0, 0]], columns=["first", "second", "second"]) result = df1 + df2 expected = DataFrame([[np.nan, 0, 0]], columns=["first", "second", "second"]) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("to_add", [[Series([1, 1])], [Series([1, 1]), Series([1, 1])]]) def test_arith_list_of_arraylike_raise(to_add): # GH 36702. Raise when trying to add list of array-like to DataFrame df = DataFrame({"x": [1, 2], "y": [1, 2]}) msg = f"Unable to coerce list of {type(to_add[0])} to Series/DataFrame" with pytest.raises(ValueError, match=msg): df + to_add with pytest.raises(ValueError, match=msg): to_add + df def test_inplace_arithmetic_series_update(using_copy_on_write, warn_copy_on_write): # https://github.com/pandas-dev/pandas/issues/36373 df = DataFrame({"A": [1, 2, 3]}) df_orig = df.copy() series = df["A"] vals = series._values with tm.assert_cow_warning(warn_copy_on_write): series += 1 if using_copy_on_write: assert series._values is not vals tm.assert_frame_equal(df, df_orig) else: assert series._values is vals expected = DataFrame({"A": [2, 3, 4]}) tm.assert_frame_equal(df, expected) def test_arithmetic_multiindex_align(): """ Regression test for: https://github.com/pandas-dev/pandas/issues/33765 """ df1 = DataFrame( [[1]], index=["a"], columns=MultiIndex.from_product([[0], [1]], names=["a", "b"]), ) df2 = DataFrame([[1]], index=["a"], columns=Index([0], name="a")) expected = DataFrame( [[0]], index=["a"], columns=MultiIndex.from_product([[0], [1]], names=["a", "b"]), ) result = df1 - df2 tm.assert_frame_equal(result, expected) def test_bool_frame_mult_float(): # GH 18549 df = DataFrame(True, list("ab"), list("cd")) result = df * 1.0 expected = DataFrame(np.ones((2, 2)), list("ab"), list("cd")) tm.assert_frame_equal(result, expected) def test_frame_sub_nullable_int(any_int_ea_dtype): # GH 32822 series1 = Series([1, 2, None], dtype=any_int_ea_dtype) series2 = Series([1, 2, 3], dtype=any_int_ea_dtype) expected = DataFrame([0, 0, None], dtype=any_int_ea_dtype) result = series1.to_frame() - series2.to_frame() tm.assert_frame_equal(result, expected) @pytest.mark.filterwarnings( "ignore:Passing a BlockManager|Passing a SingleBlockManager:DeprecationWarning" ) def test_frame_op_subclass_nonclass_constructor(): # GH#43201 subclass._constructor is a function, not the subclass itself class SubclassedSeries(Series): @property def _constructor(self): return SubclassedSeries @property def _constructor_expanddim(self): return SubclassedDataFrame class SubclassedDataFrame(DataFrame): _metadata = ["my_extra_data"] def __init__(self, my_extra_data, *args, **kwargs) -> None: self.my_extra_data = my_extra_data super().__init__(*args, **kwargs) @property def _constructor(self): return functools.partial(type(self), self.my_extra_data) @property def _constructor_sliced(self): return SubclassedSeries sdf = SubclassedDataFrame("some_data", {"A": [1, 2, 3], "B": [4, 5, 6]}) result = sdf * 2 expected = SubclassedDataFrame("some_data", {"A": [2, 4, 6], "B": [8, 10, 12]}) tm.assert_frame_equal(result, expected) result = sdf + sdf tm.assert_frame_equal(result, expected) def test_enum_column_equality(): Cols = Enum("Cols", "col1 col2") q1 = DataFrame({Cols.col1: [1, 2, 3]}) q2 = DataFrame({Cols.col1: [1, 2, 3]}) result = q1[Cols.col1] == q2[Cols.col1] expected = Series([True, True, True], name=Cols.col1) tm.assert_series_equal(result, expected) def test_mixed_col_index_dtype(): # GH 47382 df1 = DataFrame(columns=list("abc"), data=1.0, index=[0]) df2 = DataFrame(columns=list("abc"), data=0.0, index=[0]) df1.columns = df2.columns.astype("string") result = df1 + df2 expected = DataFrame(columns=list("abc"), data=1.0, index=[0]) tm.assert_frame_equal(result, expected)