polars_core/chunked_array/
random.rs

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

polars_core/chunked_array/
random.rs
