ensure before cast

pola-rs · Sep 25, 2023 · eeffe3a · eeffe3a
1 parent a1b0860
commit eeffe3a
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/crates/polars-plan/src/dsl/function_expr/random.rs b/crates/polars-plan/src/dsl/function_expr/random.rs
@@ -49,14 +49,14 @@ pub(super) fn sample_n(
     let src = &s[0];
     let n_s = &s[1];
 
-    let n_s = n_s.cast(&IDX_DTYPE)?;
-    let n = n_s.idx()?;
-
     polars_ensure!(
-        n.len() == 1,
+        n_s.len() == 1,
         ComputeError: "Sample size must be a single value."
     );
 
+    let n_s = n_s.cast(&IDX_DTYPE)?;
+    let n = n_s.idx()?;
+
     match n.get(0) {
         Some(n) => src.sample_n(n as usize, with_replacement, shuffle, seed),
         None => Ok(Series::new_empty(src.name(), src.dtype())),