address comments

pola-rs · Sep 24, 2023 · bc7c530 · bc7c530
1 parent c224cde
commit bc7c530
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/crates/polars-core/src/chunked_array/random.rs b/crates/polars-core/src/chunked_array/random.rs
@@ -174,14 +174,14 @@ impl DataFrame {
         shuffle: bool,
         seed: Option<u64>,
     ) -> PolarsResult<Self> {
-        let n = n.cast(&IDX_DTYPE)?;
-        let n = n.idx()?;
-
         polars_ensure!(
         n.len() == 1,
         ComputeError: "Sample size must be a single value."
         );
 
+        let n = n.cast(&IDX_DTYPE)?;
+        let n = n.idx()?;
+
         match n.get(0) {
             Some(n) => self.sample_n_literal(n as usize, with_replacement, shuffle, seed),
             None => {

diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py
@@ -8739,7 +8739,9 @@ def sample(
 
         if n is None:
             n = 1
-        n = _prepare_other_arg(n, 1)
+
+        if not isinstance(n, pl.Series):
+            n = pl.Series("", [n])
 
         return self._from_pydf(self._df.sample_n(n._s, with_replacement, shuffle, seed))
 

diff --git a/py-polars/polars/expr/expr.py b/py-polars/polars/expr/expr.py
@@ -8038,7 +8038,7 @@ def shuffle(self, seed: int | None = None) -> Self:
 
     def sample(
         self,
-        n: IntoExpr = None,
+        n: int | Expr | None = None,
         *,
         fraction: float | None = None,
         with_replacement: bool = False,