BUG: Sampling over selected groupbys does not reflect the selection (#40039)

Author: wjsi

doc/source/whatsnew/v1.3.0.rst

@@ -514,6 +514,7 @@ Groupby/resample/rolling
 - Bug in :meth:`DataFrameGroupBy.apply` where a :class:`MultiIndex` would be created instead of an :class:`Index` if a :class:`:meth:`core.window.rolling.RollingGroupby` object was created (:issue:`39732`)
 - Bug in :meth:`DataFrameGroupBy.sample` where error was raised when ``weights`` was specified and the index was an :class:`Int64Index` (:issue:`39927`)
 - Bug in :meth:`DataFrameGroupBy.aggregate` and :meth:`.Resampler.aggregate` would sometimes raise ``SpecificationError`` when passed a dictionary and columns were missing; will now always raise a ``KeyError`` instead (:issue:`40004`)
+- Bug in :meth:`DataFrameGroupBy.sample` where column selection was not applied to sample result (:issue:`39928`)
 -
 
 Reshaping

pandas/core/groupby/groupby.py

@@ -3098,11 +3098,12 @@ def sample(
         if random_state is not None:
             random_state = com.random_state(random_state)
 
+        group_iterator = self.grouper.get_iterator(self._selected_obj, self.axis)
         samples = [
             obj.sample(
                 n=n, frac=frac, replace=replace, weights=w, random_state=random_state
             )
-            for (_, obj), w in zip(self, ws)
+            for (_, obj), w in zip(group_iterator, ws)
         ]
 
         return concat(samples, axis=self.axis)

pandas/tests/groupby/test_sample.py

@@ -132,3 +132,13 @@ def test_groupby_sample_with_weights(index, expected_index):
     result = df.groupby("a")["b"].sample(n=2, replace=True, weights=[1, 0, 1, 0])
     expected = Series(values, name="b", index=Index(expected_index))
     tm.assert_series_equal(result, expected)
+
+
+def test_groupby_sample_with_selections():
+    # GH 39928
+    values = [1] * 10 + [2] * 10
+    df = DataFrame({"a": values, "b": values, "c": values})
+
+    result = df.groupby("a")[["b", "c"]].sample(n=None, frac=None)
+    expected = DataFrame({"b": [1, 2], "c": [1, 2]}, index=result.index)
+    tm.assert_frame_equal(result, expected)