From b9928541ae729368ca3e57b9aa9cbe67b4e049b2 Mon Sep 17 00:00:00 2001 From: fjetter Date: Mon, 3 Feb 2020 15:51:07 +0100 Subject: [PATCH] Add test for gh 31605 --- pandas/tests/groupby/test_apply.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pandas/tests/groupby/test_apply.py b/pandas/tests/groupby/test_apply.py index 9c2b045079622..41ec70468aaeb 100644 --- a/pandas/tests/groupby/test_apply.py +++ b/pandas/tests/groupby/test_apply.py @@ -851,3 +851,17 @@ def test_apply_function_returns_non_pandas_non_scalar(function, expected_values) result = df.groupby("groups").apply(function) expected = pd.Series(expected_values, index=pd.Index(["A", "B"], name="groups")) tm.assert_series_equal(result, expected) + + +def test_apply_function_returns_numpy_array(): + # GH 31605 + def fct(group): + return group["B"].values.flatten() + + df = pd.DataFrame({"A": ["a", "a", "b", "none"], "B": [1, 2, 3, np.nan]}) + + result = df.groupby("A").apply(fct) + expected = pd.Series( + [[1.0, 2.0], [3.0], [np.nan]], index=pd.Index(["a", "b", "none"], name="A") + ) + tm.assert_series_equal(result, expected)