PERF/ENH: allow Generator in sampling methods

doc/source/whatsnew/v1.4.0.rst

@@ -29,7 +29,7 @@ enhancement2
 
 Other enhancements
 ^^^^^^^^^^^^^^^^^^
--
+- :meth:`Series.sample`, :meth:`DataFrame.sample`, and :meth:`.GroupBy.sample` now accept a ``np.random.Generator`` as input to ``random_state``. A generator will be more performant, especially with ``replace=False`` (:issue:`38100`)
 -
 
 .. ---------------------------------------------------------------------------

pandas/_typing.py

@@ -125,7 +125,17 @@
 JSONSerializable = Optional[Union[PythonScalar, List, Dict]]
 Frequency = Union[str, "DateOffset"]
 Axes = Collection[Any]
-RandomState = Union[int, ArrayLike, np.random.Generator, np.random.RandomState]
+
+if TYPE_CHECKING:
+    RandomState = Union[
+        int,
+        ArrayLike,
+        np.random.Generator,
+        np.random.BitGenerator,
+        np.random.RandomState,
+    ]
+else:
+    RandomState = Union[int, ArrayLike, np.random.Generator, np.random.RandomState]
 
 # dtypes
 NpDtype = Union[str, np.dtype]

pandas/core/common.py

@@ -21,6 +21,7 @@
     Iterable,
     Iterator,
     cast,
+    overload,
 )
 import warnings
 
@@ -29,6 +30,7 @@
 from pandas._libs import lib
 from pandas._typing import (
     AnyArrayLike,
+    ArrayLike,
     NpDtype,
     RandomState,
     Scalar,
@@ -389,16 +391,28 @@ def standardize_mapping(into):
     return into
 
 
-def random_state(state: RandomState | None = None) -> np.random.RandomState:
+@overload
+def random_state(state: np.random.Generator) -> np.random.Generator:
+    ...
+
+
+@overload
+def random_state(
+    state: int | ArrayLike | np.random.BitGenerator | np.random.RandomState | None,
+) -> np.random.RandomState:
+    ...
+
+
+def random_state(state: RandomState | None = None):
     """
     Helper function for processing random_state arguments.
 
     Parameters
     ----------
-    state : int, array-like, BitGenerator, np.random.RandomState, None.
+    state : int, array-like, BitGenerator, Generator, np.random.RandomState, None.
         If receives an int, array-like, or BitGenerator, passes to
         np.random.RandomState() as seed.
-        If receives an np.random.RandomState object, just returns object.
+        If receives an np.random RandomState or Generator, just returns that unchanged.
         If receives `None`, returns np.random.
         If receives anything else, raises an informative ValueError.
 
@@ -411,7 +425,7 @@ def random_state(state: RandomState | None = None) -> np.random.RandomState:
 
     Returns
     -------
-    np.random.RandomState or np.random if state is None
+    np.random.RandomState or np.random.Generator. If state is None, returns np.random
 
     """
     if (
@@ -434,12 +448,13 @@ def random_state(state: RandomState | None = None) -> np.random.RandomState:
         return np.random.RandomState(state)  # type: ignore[arg-type]
     elif isinstance(state, np.random.RandomState):
         return state
+    elif isinstance(state, np.random.Generator):
+        return state
     elif state is None:
-        # error: Incompatible return value type (got Module, expected "RandomState")
-        return np.random  # type: ignore[return-value]
+        return np.random
     else:
         raise ValueError(
-            "random_state must be an integer, array-like, a BitGenerator, "
+            "random_state must be an integer, array-like, a BitGenerator, Generator, "
             "a numpy RandomState, or None"
         )
 

pandas/core/generic.py

@@ -5178,14 +5178,19 @@ def sample(
             If weights do not sum to 1, they will be normalized to sum to 1.
             Missing values in the weights column will be treated as zero.
             Infinite values not allowed.
-        random_state : int, array-like, BitGenerator, np.random.RandomState, optional
-            If int, array-like, or BitGenerator, seed for random number generator.
-            If np.random.RandomState, use as numpy RandomState object.
+        random_state : int, array-like, BitGenerator, np.random.RandomState,
+            np.random.Generator, optional. If int, array-like, or BitGenerator, seed for
+            random number generator. If np.random.RandomState or np.random.Generator,
+            use as given.
 
             .. versionchanged:: 1.1.0
 
-                array-like and BitGenerator (for NumPy>=1.17) object now passed to
-                np.random.RandomState() as seed
+                array-like and BitGenerator object now passed to np.random.RandomState()
+                as seed
+
+            .. versionchanged:: 1.4.0
+
+                np.random.Generator objects now accepted
 
         axis : {0 or ‘index’, 1 or ‘columns’, None}, default None
             Axis to sample. Accepts axis number or name. Default is stat axis

pandas/core/groupby/groupby.py

@@ -3211,9 +3211,14 @@ def sample(
             sampling probabilities after normalization within each group.
             Values must be non-negative with at least one positive element
             within each group.
-        random_state : int, array-like, BitGenerator, np.random.RandomState, optional
-            If int, array-like, or BitGenerator, seed for random number generator.
-            If np.random.RandomState, use as numpy RandomState object.
+        random_state : int, array-like, BitGenerator, np.random.RandomState,
+            np.random.Generator, optional. If int, array-like, or BitGenerator, seed for
+            random number generator. If np.random.RandomState or np.random.Generator,
+            use as given.
+
+            .. versionchanged:: 1.4.0
+
+                np.random.Generator objects now accepted
 
         Returns
         -------

pandas/tests/frame/methods/test_sample.py

@@ -71,8 +71,8 @@ def test_sample_lengths(self, obj):
     def test_sample_invalid_random_state(self, obj):
         # Check for error when random_state argument invalid.
         msg = (
-            "random_state must be an integer, array-like, a BitGenerator, a numpy "
-            "RandomState, or None"
+            "random_state must be an integer, array-like, a BitGenerator, Generator, "
+            "a numpy RandomState, or None"
         )
         with pytest.raises(ValueError, match=msg):
             obj.sample(random_state="a_string")
@@ -177,6 +177,22 @@ def test_sample_random_state(self, func_str, arg, frame_or_series):
         expected = obj.sample(n=3, random_state=com.random_state(eval(func_str)(arg)))
         tm.assert_equal(result, expected)
 
+    def test_sample_generator(self, frame_or_series):
+        # GH#38100
+        obj = frame_or_series(np.arange(100))
+        rng = np.random.default_rng()
+
+        # Consecutive calls should advance the seed
+        result1 = obj.sample(n=50, random_state=rng)
+        result2 = obj.sample(n=50, random_state=rng)
+        assert not (result1.index.values == result2.index.values).all()
+
+        # Matching generator initialization must give same result
+        # Consecutive calls should advance the seed
+        result1 = obj.sample(n=50, random_state=np.random.default_rng(11))
+        result2 = obj.sample(n=50, random_state=np.random.default_rng(11))
+        tm.assert_equal(result1, result2)
+
     def test_sample_upsampling_without_replacement(self, frame_or_series):
         # GH#27451
 

pandas/tests/test_common.py

@@ -84,7 +84,7 @@ def test_random_state():
 
     # Error for floats or strings
     msg = (
-        "random_state must be an integer, array-like, a BitGenerator, "
+        "random_state must be an integer, array-like, a BitGenerator, Generator, "
         "a numpy RandomState, or None"
     )
     with pytest.raises(ValueError, match=msg):