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

polars_core/chunked_array/
random.rs
