Repository: narwhals-dev/narwhals
Pull request #550 (2024-07-18 20:53:08)

feat: arrow series `sample`
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #372 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes




Patch:
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 576645b6..f8bfbffc 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -201,6 +201,17 @@ class ArrowExpr:
     def is_in(self, other: ArrowExpr | Any) -> Self:
         return reuse_series_implementation(self, "is_in", other)
 
+    def sample(
+        self: Self,
+        n: int | None = None,
+        fraction: float | None = None,
+        *,
+        with_replacement: bool = False,
+    ) -> Self:
+        return reuse_series_implementation(
+            self, "sample", n=n, fraction=fraction, with_replacement=with_replacement
+        )
+
     @property
     def dt(self) -> ArrowExprDateTimeNamespace:
         return ArrowExprDateTimeNamespace(self)
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index e02b89d3..c287dd81 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -8,6 +8,7 @@ from typing import Sequence
 from narwhals._arrow.utils import reverse_translate_dtype
 from narwhals._arrow.utils import translate_dtype
 from narwhals._arrow.utils import validate_column_comparand
+from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pyarrow
 from narwhals.dependencies import get_pyarrow_compute
 
@@ -294,6 +295,25 @@ class ArrowSeries:
             )
         )
 
+    def sample(
+        self: Self,
+        n: int | None = None,
+        fraction: float | None = None,
+        *,
+        with_replacement: bool = False,
+    ) -> Self:
+        np = get_numpy()
+        pc = get_pyarrow_compute()
+        ser = self._native_series
+        num_rows = len(self)
+
+        if n is None and fraction is not None:
+            n = int(num_rows * fraction)
+
+        idx = np.arange(0, num_rows)
+        mask = np.random.choice(idx, size=n, replace=with_replacement)
+        return self._from_native_series(pc.take(ser, mask))
+
     @property
     def shape(self) -> tuple[int]:
         return (len(self._native_series),)
diff --git a/utils/check_backend_completeness.py b/utils/check_backend_completeness.py
index 827eae5c..90506dc4 100644
--- a/utils/check_backend_completeness.py
+++ b/utils/check_backend_completeness.py
@@ -29,7 +29,6 @@ MISSING = [
     "Series.n_unique",
     "Series.quantile",
     "Series.round",
-    "Series.sample",
     "Series.shift",
     "Series.sort",
     "Series.to_frame",


Test patch:
diff --git a/tests/expr_and_series/sample_test.py b/tests/expr_and_series/sample_test.py
index 4fb54837..a19c686e 100644
--- a/tests/expr_and_series/sample_test.py
+++ b/tests/expr_and_series/sample_test.py
@@ -1,17 +1,27 @@
 from typing import Any
 
-import pytest
-
 import narwhals.stable.v1 as nw
 
 
-def test_expr_sample(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_expr_sample(constructor: Any) -> None:
     df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]})).lazy()
-    result_shape = nw.to_native(df.select(nw.col("a").sample(n=2)).collect()).shape
-    expected = (2, 1)
-    assert result_shape == expected
-    result_shape = nw.to_native(df.collect()["a"].sample(n=2)).shape
-    expected = (2,)  # type: ignore[assignment]
-    assert result_shape == expected
+
+    result_expr = df.select(nw.col("a").sample(n=2)).collect().shape
+    expected_expr = (2, 1)
+    assert result_expr == expected_expr
+
+    result_series = df.collect()["a"].sample(n=2).shape
+    expected_series = (2,)
+    assert result_series == expected_series
+
+
+def test_expr_sample_fraction(constructor: Any) -> None:
+    df = nw.from_native(constructor({"a": [1, 2, 3] * 10, "b": [4, 5, 6] * 10})).lazy()
+
+    result_expr = df.select(nw.col("a").sample(fraction=0.1)).collect().shape
+    expected_expr = (3, 1)
+    assert result_expr == expected_expr
+
+    result_series = df.collect()["a"].sample(fraction=0.1).shape
+    expected_series = (3,)
+    assert result_series == expected_series