From 01dce9d82e7fa751d24355b9c6eb9aa6a548cb88 Mon Sep 17 00:00:00 2001 From: Matt Roeschke Date: Wed, 16 Nov 2016 22:06:30 -0800 Subject: [PATCH] TST: Test aggregation over arrays (#3788) --- pandas/tests/test_groupby.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/pandas/tests/test_groupby.py b/pandas/tests/test_groupby.py index dc326aeaa88ac..52d1c5c3681e0 100644 --- a/pandas/tests/test_groupby.py +++ b/pandas/tests/test_groupby.py @@ -6813,6 +6813,23 @@ def test_group_shift_with_null_key(self): assert_frame_equal(result, expected) + def test_agg_over_numpy_arrays(self): + # GH 3788 + df = pd.DataFrame([[1, np.array([10, 20, 30])], + [1, np.array([40, 50, 60])], + [2, np.array([20, 30, 40])]], + columns=['category', 'arraydata']) + result = df.groupby('category').agg(sum) + + expected_data = [[np.array([50, 70, 90])], [np.array([20, 30, 40])]] + expected_index = pd.Index([1, 2], name='category') + expected_column = ['arraydata'] + expected = pd.DataFrame(expected_data, + index=expected_index, + columns=expected_column) + + assert_frame_equal(result, expected) + def assert_fp_equal(a, b): assert (np.abs(a - b) < 1e-12).all()