|
| 1 | +import numpy as np |
| 2 | +import pytest |
| 3 | + |
| 4 | +from pandas import DataFrame, NaT, Series, Timedelta, Timestamp |
| 5 | +import pandas._testing as tm |
| 6 | + |
| 7 | + |
| 8 | +def test_group_shift_with_null_key(): |
| 9 | + # This test is designed to replicate the segfault in issue #13813. |
| 10 | + n_rows = 1200 |
| 11 | + |
| 12 | + # Generate a moderately large dataframe with occasional missing |
| 13 | + # values in column `B`, and then group by [`A`, `B`]. This should |
| 14 | + # force `-1` in `labels` array of `g.grouper.group_info` exactly |
| 15 | + # at those places, where the group-by key is partially missing. |
| 16 | + df = DataFrame( |
| 17 | + [(i % 12, i % 3 if i % 3 else np.nan, i) for i in range(n_rows)], |
| 18 | + dtype=float, |
| 19 | + columns=["A", "B", "Z"], |
| 20 | + index=None, |
| 21 | + ) |
| 22 | + g = df.groupby(["A", "B"]) |
| 23 | + |
| 24 | + expected = DataFrame( |
| 25 | + [(i + 12 if i % 3 and i < n_rows - 12 else np.nan) for i in range(n_rows)], |
| 26 | + dtype=float, |
| 27 | + columns=["Z"], |
| 28 | + index=None, |
| 29 | + ) |
| 30 | + result = g.shift(-1) |
| 31 | + |
| 32 | + tm.assert_frame_equal(result, expected) |
| 33 | + |
| 34 | + |
| 35 | +def test_group_shift_with_fill_value(): |
| 36 | + # GH #24128 |
| 37 | + n_rows = 24 |
| 38 | + df = DataFrame( |
| 39 | + [(i % 12, i % 3, i) for i in range(n_rows)], |
| 40 | + dtype=float, |
| 41 | + columns=["A", "B", "Z"], |
| 42 | + index=None, |
| 43 | + ) |
| 44 | + g = df.groupby(["A", "B"]) |
| 45 | + |
| 46 | + expected = DataFrame( |
| 47 | + [(i + 12 if i < n_rows - 12 else 0) for i in range(n_rows)], |
| 48 | + dtype=float, |
| 49 | + columns=["Z"], |
| 50 | + index=None, |
| 51 | + ) |
| 52 | + result = g.shift(-1, fill_value=0)[["Z"]] |
| 53 | + |
| 54 | + tm.assert_frame_equal(result, expected) |
| 55 | + |
| 56 | + |
| 57 | +def test_group_shift_lose_timezone(): |
| 58 | + # GH 30134 |
| 59 | + now_dt = Timestamp.utcnow() |
| 60 | + df = DataFrame({"a": [1, 1], "date": now_dt}) |
| 61 | + result = df.groupby("a").shift(0).iloc[0] |
| 62 | + expected = Series({"date": now_dt}, name=result.name) |
| 63 | + tm.assert_series_equal(result, expected) |
| 64 | + |
| 65 | + |
| 66 | +def test_group_diff_real(any_real_dtype): |
| 67 | + df = DataFrame({"a": [1, 2, 3, 3, 2], "b": [1, 2, 3, 4, 5]}, dtype=any_real_dtype) |
| 68 | + result = df.groupby("a")["b"].diff() |
| 69 | + exp_dtype = "float" |
| 70 | + if any_real_dtype in ["int8", "int16", "float32"]: |
| 71 | + exp_dtype = "float32" |
| 72 | + expected = Series([np.nan, np.nan, np.nan, 1.0, 3.0], dtype=exp_dtype, name="b") |
| 73 | + tm.assert_series_equal(result, expected) |
| 74 | + |
| 75 | + |
| 76 | +@pytest.mark.parametrize( |
| 77 | + "data", |
| 78 | + [ |
| 79 | + [ |
| 80 | + Timestamp("2013-01-01"), |
| 81 | + Timestamp("2013-01-02"), |
| 82 | + Timestamp("2013-01-03"), |
| 83 | + ], |
| 84 | + [Timedelta("5 days"), Timedelta("6 days"), Timedelta("7 days")], |
| 85 | + ], |
| 86 | +) |
| 87 | +def test_group_diff_datetimelike(data): |
| 88 | + df = DataFrame({"a": [1, 2, 2], "b": data}) |
| 89 | + result = df.groupby("a")["b"].diff() |
| 90 | + expected = Series([NaT, NaT, Timedelta("1 days")], name="b") |
| 91 | + tm.assert_series_equal(result, expected) |
| 92 | + |
| 93 | + |
| 94 | +def test_group_diff_bool(): |
| 95 | + df = DataFrame({"a": [1, 2, 3, 3, 2], "b": [True, True, False, False, True]}) |
| 96 | + result = df.groupby("a")["b"].diff() |
| 97 | + expected = Series([np.nan, np.nan, np.nan, False, False], name="b") |
| 98 | + tm.assert_series_equal(result, expected) |
| 99 | + |
| 100 | + |
| 101 | +def test_group_diff_object_raises(object_dtype): |
| 102 | + df = DataFrame( |
| 103 | + {"a": ["foo", "bar", "bar"], "b": ["baz", "foo", "foo"]}, dtype=object_dtype |
| 104 | + ) |
| 105 | + with pytest.raises(TypeError, match=r"unsupported operand type\(s\) for -"): |
| 106 | + df.groupby("a")["b"].diff() |
0 commit comments