polars_core/chunked_array/
random.rs

1use num_traits::{Float, NumCast};
2use polars_error::to_compute_err;
3use rand::distr::Bernoulli;
4use rand::prelude::*;
5use rand::seq::index::IndexVec;
6use rand_distr::{Normal, StandardNormal, StandardUniform, Uniform};
7
8use crate::prelude::DataType::Float64;
9use crate::prelude::*;
10use crate::random::get_global_random_u64;
11use crate::utils::NoNull;
12
13fn create_rand_index_with_replacement(
14    n: usize,
15    len: usize,
16    seed: Option<u64>,
17    shuffle: Option<bool>,
18) -> IdxCa {
19    if len == 0 {
20        return IdxCa::new_vec(PlSmallStr::EMPTY, vec![]);
21    }
22    let mut rng = SmallRng::seed_from_u64(seed.unwrap_or_else(get_global_random_u64));
23    let dist = Uniform::new(0, len as IdxSize).unwrap();
24    let idxs = (0..n as IdxSize)
25        .map(move |_| dist.sample(&mut rng))
26        .collect_trusted::<NoNull<IdxCa>>()
27        .into_inner();
28    if shuffle == Some(false) {
29        idxs.sort(false)
30    } else {
31        idxs
32    }
33}
34
35fn create_rand_index_no_replacement(
36    n: usize,
37    len: usize,
38    seed: Option<u64>,
39    shuffle: Option<bool>,
40) -> IdxCa {
41    let mut rng = SmallRng::seed_from_u64(seed.unwrap_or_else(get_global_random_u64));
42    let mut buf: Vec<IdxSize>;
43    if n == len {
44        buf = (0..len as IdxSize).collect();
45        // None and Some(false) coincide here because the natural output is already ordered and
46        // forcing a shuffle would violate the fastest algorithm contract for None
47        if let Some(true) = shuffle {
48            buf.shuffle(&mut rng);
49        }
50    } else {
51        // TODO: avoid extra potential copy by vendoring rand::seq::index::sample,
52        // or genericize take over slices over any unsigned type. The optimizer
53        // should get rid of the extra copy already if IdxSize matches the IndexVec
54        // size returned.
55        buf = match rand::seq::index::sample(&mut rng, len, n) {
56            IndexVec::U32(v) => v.into_iter().map(|x| x as IdxSize).collect(),
57            #[cfg(target_pointer_width = "64")]
58            IndexVec::U64(v) => v.into_iter().map(|x| x as IdxSize).collect(),
59        };
60        // None and Some(true) coincide here because the rand::seq::index::sample
61        // already returns indices in an unspecified order so neither needs additional work
62        if let Some(false) = shuffle {
63            buf.sort_unstable();
64        }
65    }
66    IdxCa::new_vec(PlSmallStr::EMPTY, buf)
67}
68
69impl<T> ChunkedArray<T>
70where
71    T: PolarsNumericType,
72    StandardUniform: Distribution<T::Native>,
73{
74    pub fn init_rand(size: usize, null_density: f32, seed: Option<u64>) -> Self {
75        let mut rng = SmallRng::seed_from_u64(seed.unwrap_or_else(get_global_random_u64));
76        (0..size)
77            .map(|_| {
78                if rng.random::<f32>() < null_density {
79                    None
80                } else {
81                    Some(rng.random())
82                }
83            })
84            .collect()
85    }
86}
87
88fn ensure_shape(n: usize, len: usize, with_replacement: bool) -> PolarsResult<()> {
89    polars_ensure!(
90        with_replacement || n <= len,
91        ShapeMismatch:
92        "cannot take a larger sample than the total population when `with_replacement=false`"
93    );
94    Ok(())
95}
96
97impl Series {
98    pub fn sample_n(
99        &self,
100        n: usize,
101        with_replacement: bool,
102        shuffle: Option<bool>,
103        seed: Option<u64>,
104    ) -> PolarsResult<Self> {
105        ensure_shape(n, self.len(), with_replacement)?;
106        if n == 0 {
107            return Ok(self.clear());
108        }
109        let len = self.len();
110
111        match with_replacement {
112            true => {
113                let idx = create_rand_index_with_replacement(n, len, seed, shuffle);
114                debug_assert_eq!(len, self.len());
115                // SAFETY: we know that we never go out of bounds.
116                unsafe { Ok(self.take_unchecked(&idx)) }
117            },
118            false => {
119                let idx = create_rand_index_no_replacement(n, len, seed, shuffle);
120                debug_assert_eq!(len, self.len());
121                // SAFETY: we know that we never go out of bounds.
122                unsafe { Ok(self.take_unchecked(&idx)) }
123            },
124        }
125    }
126
127    /// Sample a fraction between 0.0-1.0 of this [`ChunkedArray`].
128    pub fn sample_frac(
129        &self,
130        frac: f64,
131        with_replacement: bool,
132        shuffle: Option<bool>,
133        seed: Option<u64>,
134    ) -> PolarsResult<Self> {
135        let n = (self.len() as f64 * frac) as usize;
136        self.sample_n(n, with_replacement, shuffle, seed)
137    }
138
139    pub fn shuffle(&self, seed: Option<u64>) -> Self {
140        let len = self.len();
141        let n = len;
142        let idx = create_rand_index_no_replacement(n, len, seed, Some(true));
143        debug_assert_eq!(len, self.len());
144        // SAFETY: we know that we never go out of bounds.
145        unsafe { self.take_unchecked(&idx) }
146    }
147}
148
149impl<T> ChunkedArray<T>
150where
151    T: PolarsDataType,
152    ChunkedArray<T>: ChunkTake<IdxCa>,
153{
154    /// Sample n datapoints from this [`ChunkedArray`].
155    pub fn sample_n(
156        &self,
157        n: usize,
158        with_replacement: bool,
159        shuffle: Option<bool>,
160        seed: Option<u64>,
161    ) -> PolarsResult<Self> {
162        ensure_shape(n, self.len(), with_replacement)?;
163        let len = self.len();
164
165        match with_replacement {
166            true => {
167                let idx = create_rand_index_with_replacement(n, len, seed, shuffle);
168                debug_assert_eq!(len, self.len());
169                // SAFETY: we know that we never go out of bounds.
170                unsafe { Ok(self.take_unchecked(&idx)) }
171            },
172            false => {
173                let idx = create_rand_index_no_replacement(n, len, seed, shuffle);
174                debug_assert_eq!(len, self.len());
175                // SAFETY: we know that we never go out of bounds.
176                unsafe { Ok(self.take_unchecked(&idx)) }
177            },
178        }
179    }
180
181    /// Sample a fraction between 0.0-1.0 of this [`ChunkedArray`].
182    pub fn sample_frac(
183        &self,
184        frac: f64,
185        with_replacement: bool,
186        shuffle: Option<bool>,
187        seed: Option<u64>,
188    ) -> PolarsResult<Self> {
189        let n = (self.len() as f64 * frac) as usize;
190        self.sample_n(n, with_replacement, shuffle, seed)
191    }
192}
193
194impl DataFrame {
195    /// Sample n datapoints from this [`DataFrame`].
196    pub fn sample_n(
197        &self,
198        n: &Series,
199        with_replacement: bool,
200        shuffle: Option<bool>,
201        seed: Option<u64>,
202    ) -> PolarsResult<Self> {
203        polars_ensure!(
204        n.len() == 1,
205        ComputeError: "Sample size must be a single value."
206        );
207
208        let n = n.strict_cast(&IDX_DTYPE)?;
209        let n = n.idx()?;
210
211        match n.get(0) {
212            Some(n) => self.sample_n_literal(n as usize, with_replacement, shuffle, seed),
213            None => Ok(self.clear()),
214        }
215    }
216
217    pub fn sample_n_literal(
218        &self,
219        n: usize,
220        with_replacement: bool,
221        shuffle: Option<bool>,
222        seed: Option<u64>,
223    ) -> PolarsResult<Self> {
224        ensure_shape(n, self.height(), with_replacement)?;
225        // All columns should used the same indices. So we first create the indices.
226        let idx = match with_replacement {
227            true => create_rand_index_with_replacement(n, self.height(), seed, shuffle),
228            false => create_rand_index_no_replacement(n, self.height(), seed, shuffle),
229        };
230        // SAFETY: the indices are within bounds.
231        Ok(unsafe { self.take_unchecked(&idx) })
232    }
233
234    /// Sample a fraction between 0.0-1.0 of this [`DataFrame`].
235    pub fn sample_frac(
236        &self,
237        frac: &Series,
238        with_replacement: bool,
239        shuffle: Option<bool>,
240        seed: Option<u64>,
241    ) -> PolarsResult<Self> {
242        polars_ensure!(
243        frac.len() == 1,
244        ComputeError: "Sample fraction must be a single value."
245        );
246
247        let frac = frac.cast(&Float64)?;
248        let frac = frac.f64()?;
249
250        match frac.get(0) {
251            Some(frac) => {
252                let n = (self.height() as f64 * frac) as usize;
253                self.sample_n_literal(n, with_replacement, shuffle, seed)
254            },
255            None => Ok(self.clear()),
256        }
257    }
258}
259
260impl<T> ChunkedArray<T>
261where
262    T: PolarsNumericType,
263    T::Native: Float,
264{
265    /// Create [`ChunkedArray`] with samples from a Normal distribution.
266    pub fn rand_normal(
267        name: PlSmallStr,
268        length: usize,
269        mean: f64,
270        std_dev: f64,
271    ) -> PolarsResult<Self> {
272        let normal = Normal::new(mean, std_dev).map_err(to_compute_err)?;
273        let mut builder = PrimitiveChunkedBuilder::<T>::new(name, length);
274        let mut rng = rand::rng();
275        for _ in 0..length {
276            let smpl = normal.sample(&mut rng);
277            let smpl = NumCast::from(smpl).unwrap();
278            builder.append_value(smpl)
279        }
280        Ok(builder.finish())
281    }
282
283    /// Create [`ChunkedArray`] with samples from a Standard Normal distribution.
284    pub fn rand_standard_normal(name: PlSmallStr, length: usize) -> Self {
285        let mut builder = PrimitiveChunkedBuilder::<T>::new(name, length);
286        let mut rng = rand::rng();
287        for _ in 0..length {
288            let smpl: f64 = rng.sample(StandardNormal);
289            let smpl = NumCast::from(smpl).unwrap();
290            builder.append_value(smpl)
291        }
292        builder.finish()
293    }
294
295    /// Create [`ChunkedArray`] with samples from a Uniform distribution.
296    pub fn rand_uniform(name: PlSmallStr, length: usize, low: f64, high: f64) -> Self {
297        let uniform = Uniform::new(low, high).unwrap();
298        let mut builder = PrimitiveChunkedBuilder::<T>::new(name, length);
299        let mut rng = rand::rng();
300        for _ in 0..length {
301            let smpl = uniform.sample(&mut rng);
302            let smpl = NumCast::from(smpl).unwrap();
303            builder.append_value(smpl)
304        }
305        builder.finish()
306    }
307}
308
309impl BooleanChunked {
310    /// Create [`ChunkedArray`] with samples from a Bernoulli distribution.
311    pub fn rand_bernoulli(name: PlSmallStr, length: usize, p: f64) -> PolarsResult<Self> {
312        let dist = Bernoulli::new(p).map_err(to_compute_err)?;
313        let mut rng = rand::rng();
314        let mut builder = BooleanChunkedBuilder::new(name, length);
315        for _ in 0..length {
316            let smpl = dist.sample(&mut rng);
317            builder.append_value(smpl)
318        }
319        Ok(builder.finish())
320    }
321}
322
323#[cfg(test)]
324mod test {
325    use super::*;
326
327    #[test]
328    fn test_sample() {
329        let df = df![
330            "foo" => &[1, 2, 3, 4, 5]
331        ]
332        .unwrap();
333
334        // Default samples are random and don't require seeds.
335        assert!(
336            df.sample_n(
337                &Series::new(PlSmallStr::from_static("s"), &[3]),
338                false,
339                None,
340                None
341            )
342            .is_ok()
343        );
344        assert!(
345            df.sample_frac(
346                &Series::new(PlSmallStr::from_static("frac"), &[0.4]),
347                false,
348                None,
349                None
350            )
351            .is_ok()
352        );
353        // With seeding.
354        assert!(
355            df.sample_n(
356                &Series::new(PlSmallStr::from_static("s"), &[3]),
357                false,
358                None,
359                Some(0)
360            )
361            .is_ok()
362        );
363        assert!(
364            df.sample_frac(
365                &Series::new(PlSmallStr::from_static("frac"), &[0.4]),
366                false,
367                None,
368                Some(0)
369            )
370            .is_ok()
371        );
372        // Without replacement can not sample more than 100%.
373        assert!(
374            df.sample_frac(
375                &Series::new(PlSmallStr::from_static("frac"), &[2.0]),
376                false,
377                None,
378                Some(0)
379            )
380            .is_err()
381        );
382        assert!(
383            df.sample_n(
384                &Series::new(PlSmallStr::from_static("s"), &[3]),
385                true,
386                None,
387                Some(0)
388            )
389            .is_ok()
390        );
391        assert!(
392            df.sample_frac(
393                &Series::new(PlSmallStr::from_static("frac"), &[0.4]),
394                true,
395                None,
396                Some(0)
397            )
398            .is_ok()
399        );
400        // With replacement can sample more than 100%.
401        assert!(
402            df.sample_frac(
403                &Series::new(PlSmallStr::from_static("frac"), &[2.0]),
404                true,
405                None,
406                Some(0)
407            )
408            .is_ok()
409        );
410    }
411}
polars_core/chunked_array/random.rs

polars_core/chunked_array/
random.rs
