def test_groupby_agg(): # Ensure that the result of agg is inferred to be decimal dtype # https://github.com/pandas-dev/pandas/issues/29141 data = make_data()[:5] df = pd.DataFrame( {"id1": [0, 0, 0, 1, 1], "id2": [0, 1, 0, 1, 1], "decimals": DecimalArray(data)} ) # single key, selected column expected = pd.Series(to_decimal([data[0], data[3]])) result = df.groupby("id1")["decimals"].agg(lambda x: x.iloc[0]) tm.assert_series_equal(result, expected, check_names=False) result = df["decimals"].groupby(df["id1"]).agg(lambda x: x.iloc[0]) tm.assert_series_equal(result, expected, check_names=False) # multiple keys, selected column expected = pd.Series( to_decimal([data[0], data[1], data[3]]), index=pd.MultiIndex.from_tuples([(0, 0), (0, 1), (1, 1)]), ) result = df.groupby(["id1", "id2"])["decimals"].agg(lambda x: x.iloc[0]) tm.assert_series_equal(result, expected, check_names=False) result = df["decimals"].groupby([df["id1"], df["id2"]]).agg(lambda x: x.iloc[0]) tm.assert_series_equal(result, expected, check_names=False) # multiple columns expected = pd.DataFrame({"id2": [0, 1], "decimals": to_decimal([data[0], data[3]])}) result = df.groupby("id1").agg(lambda x: x.iloc[0]) tm.assert_frame_equal(result, expected, check_names=False)
def test_array_ufunc_series_defer(): a = to_decimal([1, 2, 3]) s = pd.Series(a) expected = pd.Series(to_decimal([2, 4, 6])) r1 = np.add(s, a) r2 = np.add(a, s) tm.assert_series_equal(r1, expected) tm.assert_series_equal(r2, expected)
def test_divmod_array(reverse, expected_div, expected_mod): # https://github.com/pandas-dev/pandas/issues/22930 arr = to_decimal([1, 2, 3, 4]) if reverse: div, mod = divmod(2, arr) else: div, mod = divmod(arr, 2) expected_div = to_decimal(expected_div) expected_mod = to_decimal(expected_mod) tm.assert_extension_array_equal(div, expected_div) tm.assert_extension_array_equal(mod, expected_mod)
def test_array_ufunc_series_scalar_other(): # check _HANDLED_TYPES a = to_decimal([1, 2, 3]) s = pd.Series(a) result = np.add(s, decimal.Decimal(1)) expected = pd.Series(np.add(a, decimal.Decimal(1))) tm.assert_series_equal(result, expected)
def test_groupby_agg_err_catching(err_cls): # make sure we suppress anything other than TypeError or AssertionError # in _python_agg_general # Use a non-standard EA to make sure we don't go down ndarray paths from pandas.tests.extension.decimal.array import ( DecimalArray, make_data, to_decimal, ) data = make_data()[:5] df = DataFrame( {"id1": [0, 0, 0, 1, 1], "id2": [0, 1, 0, 1, 1], "decimals": DecimalArray(data)} ) expected = Series(to_decimal([data[0], data[3]])) def weird_func(x): # weird function that raise something other than TypeError or IndexError # in _python_agg_general if len(x) == 0: raise err_cls return x.iloc[0] result = df["decimals"].groupby(df["id1"]).agg(weird_func) tm.assert_series_equal(result, expected, check_names=False)
def test_groupby_agg_ea_method(monkeypatch): # Ensure that the result of agg is inferred to be decimal dtype # https://github.com/pandas-dev/pandas/issues/29141 def DecimalArray__my_sum(self): return np.sum(np.array(self)) monkeypatch.setattr(DecimalArray, "my_sum", DecimalArray__my_sum, raising=False) data = make_data()[:5] df = pd.DataFrame({"id": [0, 0, 0, 1, 1], "decimals": DecimalArray(data)}) expected = pd.Series(to_decimal([data[0] + data[1] + data[2], data[3] + data[4]])) result = df.groupby("id")["decimals"].agg(lambda x: x.values.my_sum()) tm.assert_series_equal(result, expected, check_names=False) s = pd.Series(DecimalArray(data)) result = s.groupby(np.array([0, 0, 0, 1, 1])).agg(lambda x: x.values.my_sum()) tm.assert_series_equal(result, expected, check_names=False)
def test_array_ufunc_series(): a = to_decimal([1, 2, 3]) s = pd.Series(a) result = np.exp(s) expected = pd.Series(to_decimal(np.exp(a._data))) tm.assert_series_equal(result, expected)
def test_array_ufunc(): a = to_decimal([1, 2, 3]) result = np.exp(a) expected = to_decimal(np.exp(a._data)) tm.assert_extension_array_equal(result, expected)