Repository: narwhals-dev/narwhals
Pull request #975 (2024-09-14 18:07:57)

feat: add `DataFrame.sample` method
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #973 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.

I mean... why not? 😂


Patch:
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index f78b4e3d..fe0c4025 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -29,6 +29,7 @@
         - rename
         - row
         - rows
+        - sample
         - schema
         - select
         - shape
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index f409ef73..7029cc0e 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -572,3 +572,25 @@ class ArrowDataFrame:
 
     def to_arrow(self: Self) -> Any:
         return self._native_frame
+
+    def sample(
+        self: Self,
+        n: int | None = None,
+        *,
+        fraction: float | None = None,
+        with_replacement: bool = False,
+        seed: int | None = None,
+    ) -> Self:
+        import numpy as np  # ignore-banned-import
+        import pyarrow.compute as pc  # ignore-banned-import()
+
+        frame = self._native_frame
+        num_rows = len(self)
+        if n is None and fraction is not None:
+            n = int(num_rows * fraction)
+
+        rng = np.random.default_rng(seed=seed)
+        idx = np.arange(0, num_rows)
+        mask = rng.choice(idx, size=n, replace=with_replacement)
+
+        return self._from_native_frame(pc.take(frame, mask))
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 24e4fe5c..b324be2b 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -249,12 +249,18 @@ class ArrowExpr:
     def sample(
         self: Self,
         n: int | None = None,
-        fraction: float | None = None,
         *,
+        fraction: float | None = None,
         with_replacement: bool = False,
+        seed: int | None = None,
     ) -> Self:
         return reuse_series_implementation(
-            self, "sample", n=n, fraction=fraction, with_replacement=with_replacement
+            self,
+            "sample",
+            n=n,
+            fraction=fraction,
+            with_replacement=with_replacement,
+            seed=seed,
         )
 
     def fill_null(self: Self, value: Any) -> Self:
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 1f98ad80..4f53a3f0 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -534,9 +534,10 @@ class ArrowSeries:
     def sample(
         self: Self,
         n: int | None = None,
-        fraction: float | None = None,
         *,
+        fraction: float | None = None,
         with_replacement: bool = False,
+        seed: int | None = None,
     ) -> Self:
         import numpy as np  # ignore-banned-import
         import pyarrow.compute as pc  # ignore-banned-import()
@@ -547,8 +548,10 @@ class ArrowSeries:
         if n is None and fraction is not None:
             n = int(num_rows * fraction)
 
+        rng = np.random.default_rng(seed=seed)
         idx = np.arange(0, num_rows)
-        mask = np.random.choice(idx, size=n, replace=with_replacement)
+        mask = rng.choice(idx, size=n, replace=with_replacement)
+
         return self._from_native_series(pc.take(ser, mask))
 
     def fill_null(self: Self, value: Any) -> Self:
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 71a65999..62067069 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -693,3 +693,17 @@ class PandasLikeDataFrame:
         import pyarrow as pa  # ignore-banned-import()
 
         return pa.Table.from_pandas(self._native_frame)
+
+    def sample(
+        self: Self,
+        n: int | None = None,
+        *,
+        fraction: float | None = None,
+        with_replacement: bool = False,
+        seed: int | None = None,
+    ) -> Self:
+        return self._from_native_frame(
+            self._native_frame.sample(
+                n=n, frac=fraction, replace=with_replacement, random_state=seed
+            )
+        )
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 74a2ee31..409d1ab0 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -254,14 +254,20 @@ class PandasLikeExpr:
         return reuse_series_implementation(self, "shift", n=n)
 
     def sample(
-        self,
+        self: Self,
         n: int | None = None,
-        fraction: float | None = None,
         *,
+        fraction: float | None = None,
         with_replacement: bool = False,
+        seed: int | None = None,
     ) -> Self:
         return reuse_series_implementation(
-            self, "sample", n=n, fraction=fraction, with_replacement=with_replacement
+            self,
+            "sample",
+            n=n,
+            fraction=fraction,
+            with_replacement=with_replacement,
+            seed=seed,
         )
 
     def alias(self, name: str) -> Self:
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 8288be26..14d7a128 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -436,15 +436,16 @@ class PandasLikeSeries:
         return ser.nunique(dropna=False)  # type: ignore[no-any-return]
 
     def sample(
-        self,
+        self: Self,
         n: int | None = None,
-        fraction: float | None = None,
         *,
+        fraction: float | None = None,
         with_replacement: bool = False,
-    ) -> PandasLikeSeries:
+        seed: int | None = None,
+    ) -> Self:
         ser = self._native_series
         return self._from_native_series(
-            ser.sample(n=n, frac=fraction, replace=with_replacement)
+            ser.sample(n=n, frac=fraction, replace=with_replacement, random_state=seed)
         )
 
     def abs(self) -> PandasLikeSeries:
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 980e6694..ce5999aa 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -2459,6 +2459,66 @@ class DataFrame(BaseFrame[FrameT]):
         """
         return self._compliant_frame.to_arrow()
 
+    def sample(
+        self: Self,
+        n: int | None = None,
+        *,
+        fraction: float | None = None,
+        with_replacement: bool = False,
+        seed: int | None = None,
+    ) -> Self:
+        r"""
+        Sample from this DataFrame.
+
+        Arguments:
+            n: Number of items to return. Cannot be used with fraction.
+            fraction: Fraction of items to return. Cannot be used with n.
+            with_replacement: Allow values to be sampled more than once.
+            seed: Seed for the random number generator. If set to None (default), a random
+                seed is generated for each sample operation.
+
+        Notes:
+            The results may not be consistent across libraries.
+
+        Examples:
+            >>> import narwhals as nw
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> data = {"a": [1, 2, 3, 4], "b": ["x", "y", "x", "y"]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.sample(n=2, seed=123)
+
+            We can then pass either pandas or Polars to `func`:
+            >>> func(df_pd)
+               a  b
+            3  4  y
+            0  1  x
+            >>> func(df_pl)
+            shape: (2, 2)
+            ┌─────┬─────┐
+            │ a   ┆ b   │
+            │ --- ┆ --- │
+            │ i64 ┆ str │
+            ╞═════╪═════╡
+            │ 2   ┆ y   │
+            │ 3   ┆ x   │
+            └─────┴─────┘
+
+            As you can see, by using the same seed, the result will be consistent within
+            the same backend, but not necessarely across different backends.
+        """
+        return self._from_compliant_dataframe(
+            self._compliant_frame.sample(
+                n=n, fraction=fraction, with_replacement=with_replacement, seed=seed
+            )
+        )
+
 
 class LazyFrame(BaseFrame[FrameT]):
     """
diff --git a/narwhals/expr.py b/narwhals/expr.py
index 4d54c288..b39a5981 100644
--- a/narwhals/expr.py
+++ b/narwhals/expr.py
@@ -1228,21 +1228,22 @@ class Expr:
         return self.__class__(lambda plx: self._call(plx).drop_nulls())
 
     def sample(
-        self,
+        self: Self,
         n: int | None = None,
-        fraction: float | None = None,
         *,
+        fraction: float | None = None,
         with_replacement: bool = False,
+        seed: int | None = None,
     ) -> Self:
         """
         Sample randomly from this expression.
 
         Arguments:
             n: Number of items to return. Cannot be used with fraction.
-
             fraction: Fraction of items to return. Cannot be used with n.
-
             with_replacement: Allow values to be sampled more than once.
+            seed: Seed for the random number generator. If set to None (default), a random
+                seed is generated for each sample operation.
 
         Examples:
             >>> import narwhals as nw
@@ -1279,7 +1280,7 @@ class Expr:
         """
         return self.__class__(
             lambda plx: self._call(plx).sample(
-                n, fraction=fraction, with_replacement=with_replacement
+                n, fraction=fraction, with_replacement=with_replacement, seed=seed
             )
         )
 
diff --git a/narwhals/series.py b/narwhals/series.py
index 9fcb07a2..9d21058f 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -1076,21 +1076,22 @@ class Series:
         return self._from_compliant_series(self._compliant_series.shift(n))
 
     def sample(
-        self,
+        self: Self,
         n: int | None = None,
-        fraction: float | None = None,
         *,
+        fraction: float | None = None,
         with_replacement: bool = False,
+        seed: int | None = None,
     ) -> Self:
         """
         Sample randomly from this Series.
 
         Arguments:
             n: Number of items to return. Cannot be used with fraction.
-
             fraction: Fraction of items to return. Cannot be used with n.
-
             with_replacement: Allow values to be sampled more than once.
+            seed: Seed for the random number generator. If set to None (default), a random
+                seed is generated for each sample operation.
 
         Notes:
             The `sample` method returns a Series with a specified number of
@@ -1131,7 +1132,7 @@ class Series:
         """
         return self._from_compliant_series(
             self._compliant_series.sample(
-                n=n, fraction=fraction, with_replacement=with_replacement
+                n=n, fraction=fraction, with_replacement=with_replacement, seed=seed
             )
         )
 


Test patch:
diff --git a/tests/expr_and_series/sample_test.py b/tests/expr_and_series/sample_test.py
index 039bd31b..eb6d853e 100644
--- a/tests/expr_and_series/sample_test.py
+++ b/tests/expr_and_series/sample_test.py
@@ -2,6 +2,7 @@ import pytest
 
 import narwhals.stable.v1 as nw
 from tests.utils import Constructor
+from tests.utils import compare_dicts
 
 
 def test_expr_sample(constructor: Constructor, request: pytest.FixtureRequest) -> None:
@@ -32,3 +33,37 @@ def test_expr_sample_fraction(
     result_series = df.collect()["a"].sample(fraction=0.1).shape
     expected_series = (3,)
     assert result_series == expected_series
+
+
+def test_sample_with_seed(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if "dask" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+
+    size, n = 100, 10
+    df = nw.from_native(constructor({"a": list(range(size))})).lazy()
+    expected = {"res1": [True], "res2": [False]}
+    result = (
+        df.select(
+            seed1=nw.col("a").sample(n=n, seed=123),
+            seed2=nw.col("a").sample(n=n, seed=123),
+            seed3=nw.col("a").sample(n=n, seed=42),
+        )
+        .select(
+            res1=(nw.col("seed1") == nw.col("seed2")).all(),
+            res2=(nw.col("seed1") == nw.col("seed3")).all(),
+        )
+        .collect()
+    )
+
+    compare_dicts(result, expected)
+
+    series = df.collect()["a"]
+    seed1 = series.sample(n=n, seed=123)
+    seed2 = series.sample(n=n, seed=123)
+    seed3 = series.sample(n=n, seed=42)
+
+    compare_dicts(
+        {"res1": [(seed1 == seed2).all()], "res2": [(seed1 == seed3).all()]}, expected
+    )
diff --git a/tests/frame/sample_test.py b/tests/frame/sample_test.py
new file mode 100644
index 00000000..88d5969c
--- /dev/null
+++ b/tests/frame/sample_test.py
@@ -0,0 +1,34 @@
+import narwhals.stable.v1 as nw
+from tests.utils import Constructor
+
+
+def test_sample_n(constructor_eager: Constructor) -> None:
+    df = nw.from_native(
+        constructor_eager({"a": [1, 2, 3, 4], "b": ["x", "y", "x", "y"]}), eager_only=True
+    )
+
+    result_expr = df.sample(n=2).shape
+    expected_expr = (2, 2)
+    assert result_expr == expected_expr
+
+
+def test_sample_fraction(constructor_eager: Constructor) -> None:
+    df = nw.from_native(
+        constructor_eager({"a": [1, 2, 3, 4], "b": ["x", "y", "x", "y"]}), eager_only=True
+    )
+
+    result_expr = df.sample(fraction=0.5).shape
+    expected_expr = (2, 2)
+    assert result_expr == expected_expr
+
+
+def test_sample_with_seed(constructor_eager: Constructor) -> None:
+    size, n = 100, 10
+    df = nw.from_native(constructor_eager({"a": range(size)}), eager_only=True)
+
+    r1 = nw.to_native(df.sample(n=n, seed=123))
+    r2 = nw.to_native(df.sample(n=n, seed=123))
+    r3 = nw.to_native(df.sample(n=n, seed=42))
+
+    assert r1.equals(r2)  # type: ignore[union-attr]
+    assert not r1.equals(r3)  # type: ignore[union-attr]