change sample method to work for pandas

narwhals-dev · Mar 19, 2024 · ce542f2 · ce542f2
1 parent e0f26d6
commit ce542f2
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 8 deletions.
diff --git a/narwhals/pandas_like/expr.py b/narwhals/pandas_like/expr.py
@@ -180,8 +180,16 @@ def n_unique(self) -> Self:
     def unique(self) -> Self:
         return register_expression_call(self, "unique")
 
-    def sample(self, n: int, fraction: float, *, with_replacement: bool) -> Self:
-        return register_expression_call(self, "sample", n, fraction, with_replacement)
+    def sample(
+        self,
+        n: int | None = None,
+        fraction: float | None = None,
+        *,
+        with_replacement: bool = False,
+    ) -> Self:
+        return register_expression_call(
+            self, "sample", n, fraction=fraction, with_replacement=with_replacement
+        )
 
     def alias(self, name: str) -> Self:
         # Define this one manually, so that we can

diff --git a/narwhals/pandas_like/series.py b/narwhals/pandas_like/series.py
@@ -291,11 +291,15 @@ def zip_with(self, mask: PandasSeries, other: PandasSeries) -> PandasSeries:
         ser = self._series
         return self._from_series(ser.where(mask, other))
 
-    def sample(self, n: int, fraction: float, *, with_replacement: bool) -> PandasSeries:
+    def sample(
+        self,
+        n: int | None = None,
+        fraction: float | None = None,
+        *,
+        with_replacement: bool = False,
+    ) -> PandasSeries:
         ser = self._series
-        return self._from_series(
-            ser.sample(n=n, frac=fraction, with_replacement=with_replacement)
-        )
+        return self._from_series(ser.sample(n=n, frac=fraction, replace=with_replacement))
 
     def unique(self) -> PandasSeries:
         if self._implementation != "pandas":

diff --git a/narwhals/series.py b/narwhals/series.py
@@ -118,9 +118,15 @@ def zip_with(self, mask: Self, other: Self) -> Self:
             self._series.zip_with(self._extract_native(mask), self._extract_native(other))
         )
 
-    def sample(self, n: int, fraction: float, *, with_replacement: bool) -> Self:
+    def sample(
+        self,
+        n: int | None = None,
+        fraction: float | None = None,
+        *,
+        with_replacement: bool = False,
+    ) -> Self:
         return self._from_series(
-            self._series.sample(n, fraction=fraction, with_replacement=with_replacement)
+            self._series.sample(n=n, fraction=fraction, with_replacement=with_replacement)
         )
 
     def to_numpy(self) -> Any: