BUG: Cannot sample on DataFrameGroupBy with weights when index is specified

Author: wjsi

pandas/core/groupby/groupby.py

@@ -3076,7 +3076,7 @@ def sample(
 
         if weights is not None:
             weights = Series(weights, index=self._selected_obj.index)
-            ws = [weights[idx] for idx in self.indices.values()]
+            ws = [weights.iloc[idx] for idx in self.indices.values()]
         else:
             ws = [None] * self.ngroups
 

pandas/tests/groupby/test_sample.py

@@ -116,14 +116,18 @@ def test_groupby_sample_without_n_or_frac():
     tm.assert_series_equal(result, expected)
 
 
-def test_groupby_sample_with_weights():
+@pytest.mark.parametrize(
+    "index, expect_index",
+    [(["w", "x", "y", "z"], ["w", "w", "y", "y"]), ([3, 4, 5, 6], [3, 3, 5, 5])],
+)
+def test_groupby_sample_with_weights(index, expect_index):
     values = [1] * 2 + [2] * 2
-    df = DataFrame({"a": values, "b": values}, index=Index(["w", "x", "y", "z"]))
+    df = DataFrame({"a": values, "b": values}, index=Index(index))
 
     result = df.groupby("a").sample(n=2, replace=True, weights=[1, 0, 1, 0])
-    expected = DataFrame({"a": values, "b": values}, index=Index(["w", "w", "y", "y"]))
+    expected = DataFrame({"a": values, "b": values}, index=Index(expect_index))
     tm.assert_frame_equal(result, expected)
 
     result = df.groupby("a")["b"].sample(n=2, replace=True, weights=[1, 0, 1, 0])
-    expected = Series(values, name="b", index=Index(["w", "w", "y", "y"]))
+    expected = Series(values, name="b", index=Index(expect_index))
     tm.assert_series_equal(result, expected)