code: def test_groupby_sample_balanced_groups_shape(n, frac):
      values = [1] * 10 + [2] * 10
      df = DataFrame({"a": values, "b": values})
  
      result = df.groupby("a").sample(n=n, frac=frac)
      values = [1] * 2 + [2] * 2
      expected = DataFrame(
          {"a": values, "b": values}, index=result.index
      )
      tm.assert_frame_equal(result, expected)
  
      result = df.groupby("a")["b"].sample(n=n, frac=frac)
      expected = Series(values, name="b", index=result.index)
      tm.assert_series_equal(result, expected)
  
file_path: pandas/pandas/tests/groupby/methods/test_sample.py
imports: import pytest
  from pandas import DataFrame, Index, Series
  import pandas._testing
label_desc: Use image node_id 1 for calling a global function with example usage: test_groupby_sample_balanced_groups_shape(n, frac) without return types