From 21fc53c350cd5c77621370e9a5fc2d270e27cc6a Mon Sep 17 00:00:00 2001 From: pranavb Date: Fri, 9 Dec 2022 10:55:06 -0500 Subject: [PATCH 1/2] TST: Added test for groupby dropping NA when grouping on columns GH21755 --- pandas/tests/groupby/test_groupby.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/pandas/tests/groupby/test_groupby.py b/pandas/tests/groupby/test_groupby.py index 667656cb4de02..81cb9974a91a8 100644 --- a/pandas/tests/groupby/test_groupby.py +++ b/pandas/tests/groupby/test_groupby.py @@ -2840,3 +2840,15 @@ def test_groupby_index_name_in_index_content(val_in, index, val_out): result = series.to_frame().groupby("blah").sum() expected = expected.to_frame() tm.assert_frame_equal(result, expected) + + +def test_groupby_on_column_drops_na(): + # GH 21755 + df = DataFrame( + {"a": [np.nan, "foo", "bar", np.nan, np.nan, "baz"], "b": [0, 1, 2, 3, 4, 5]} + ) + result = df.groupby("a").head(1) + expected = DataFrame({"a": ["foo", "bar", "baz"], "b": [1, 2, 5]}) + expected.index = [1, 2, 5] + + tm.assert_frame_equal(result, expected) From a8dc3d671d8a14aae26ebc0b0ab8618893bc1652 Mon Sep 17 00:00:00 2001 From: pranavb Date: Sun, 11 Dec 2022 18:13:45 -0500 Subject: [PATCH 2/2] TST: Added index list inside of expected DataFrame constructor --- pandas/tests/groupby/test_groupby.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pandas/tests/groupby/test_groupby.py b/pandas/tests/groupby/test_groupby.py index 81cb9974a91a8..58c5e5155e4e8 100644 --- a/pandas/tests/groupby/test_groupby.py +++ b/pandas/tests/groupby/test_groupby.py @@ -2848,7 +2848,6 @@ def test_groupby_on_column_drops_na(): {"a": [np.nan, "foo", "bar", np.nan, np.nan, "baz"], "b": [0, 1, 2, 3, 4, 5]} ) result = df.groupby("a").head(1) - expected = DataFrame({"a": ["foo", "bar", "baz"], "b": [1, 2, 5]}) - expected.index = [1, 2, 5] + expected = DataFrame({"a": ["foo", "bar", "baz"], "b": [1, 2, 5]}, index=[1, 2, 5]) tm.assert_frame_equal(result, expected)