BUG: bug in head/tail of groupby for returning selected columns (GH6524)

Author: jreback

pandas/core/groupby.py

@@ -607,7 +607,8 @@ def head(self, n=5):
         """
         rng = np.arange(self.grouper._max_groupsize, dtype='int64')
         in_head = self._cumcount_array(rng) < n
-        head = self.obj[in_head]
+        obj = self._obj_with_exclusions
+        head = obj.iloc[in_head]
         if self.as_index:
             head.index = self._index_with_as_index(in_head)
         return head
@@ -636,7 +637,8 @@ def tail(self, n=5):
         """
         rng = np.arange(0, -self.grouper._max_groupsize, -1, dtype='int64')
         in_tail = self._cumcount_array(rng, ascending=False) > -n
-        tail = self.obj[in_tail]
+        obj = self._obj_with_exclusions
+        tail = obj.iloc[in_tail]
         if self.as_index:
             tail.index = self._index_with_as_index(in_tail)
         return tail

pandas/tests/test_groupby.py

@@ -2455,6 +2455,22 @@ def test_getitem_list_of_columns(self):
         assert_frame_equal(result2, expected)
         assert_frame_equal(result3, expected)
 
+    def test_validate_getitem(self):
+
+        # invalid getitems
+        # GH 6524
+        df = DataFrame({"Dummy":[1,2]*6,
+                        "X":[1,3,7]*4,
+                        "Y":[2,3,4]*4,
+                        "group":["A","B"]*6})
+        expected = DataFrame({"X":[1,3],
+                              "Y":[2,3]},
+                             index = MultiIndex.from_tuples([('A',0),('B',1)],names=['group',None]))
+        result = df.groupby('group')[['X', 'Y']].head(1)
+        assert_frame_equal(result, expected)
+
+        self.assertRaises(AttributeError, lambda : df.groupby('group').loc[:,['X', 'Y']].head(1))
+
     def test_agg_multiple_functions_maintain_order(self):
         # GH #610
         funcs = [('mean', np.mean), ('max', np.max), ('min', np.min)]