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` (sample size: {}, population: {})",
71        n, len
72    );
73    Ok(())
74}
75
76impl Series {
77    pub fn sample_n(
78        &self,
79        n: usize,
80        with_replacement: bool,
81        shuffle: Option<bool>,
82        seed: Option<u64>,
83    ) -> PolarsResult<Self> {
84        ensure_shape(n, self.len(), with_replacement)?;
85        if n == 0 {
86            return Ok(self.clear());
87        }
88        let len = self.len();
89
90        match with_replacement {
91            true => {
92                let idx = create_rand_index_with_replacement(n, len, seed, shuffle);
93                debug_assert_eq!(len, self.len());
94                // SAFETY: we know that we never go out of bounds.
95                unsafe { Ok(self.take_unchecked(&idx)) }
96            },
97            false => {
98                let idx = create_rand_index_no_replacement(n, len, seed, shuffle);
99                debug_assert_eq!(len, self.len());
100                // SAFETY: we know that we never go out of bounds.
101                unsafe { Ok(self.take_unchecked(&idx)) }
102            },
103        }
104    }
105
106    /// Sample a fraction between 0.0-1.0 of this [`ChunkedArray`].
107    pub fn sample_frac(
108        &self,
109        frac: f64,
110        with_replacement: bool,
111        shuffle: Option<bool>,
112        seed: Option<u64>,
113    ) -> PolarsResult<Self> {
114        let n = (self.len() as f64 * frac) as usize;
115        self.sample_n(n, with_replacement, shuffle, seed)
116    }
117
118    pub fn shuffle(&self, seed: Option<u64>) -> Self {
119        let len = self.len();
120        let n = len;
121        let idx = create_rand_index_no_replacement(n, len, seed, Some(true));
122        debug_assert_eq!(len, self.len());
123        // SAFETY: we know that we never go out of bounds.
124        unsafe { self.take_unchecked(&idx) }
125    }
126}
127
128impl<T> ChunkedArray<T>
129where
130    T: PolarsDataType,
131    ChunkedArray<T>: ChunkTake<IdxCa>,
132{
133    /// Sample n datapoints from this [`ChunkedArray`].
134    pub fn sample_n(
135        &self,
136        n: usize,
137        with_replacement: bool,
138        shuffle: Option<bool>,
139        seed: Option<u64>,
140    ) -> PolarsResult<Self> {
141        ensure_shape(n, self.len(), with_replacement)?;
142        let len = self.len();
143
144        match with_replacement {
145            true => {
146                let idx = create_rand_index_with_replacement(n, len, seed, shuffle);
147                debug_assert_eq!(len, self.len());
148                // SAFETY: we know that we never go out of bounds.
149                unsafe { Ok(self.take_unchecked(&idx)) }
150            },
151            false => {
152                let idx = create_rand_index_no_replacement(n, len, seed, shuffle);
153                debug_assert_eq!(len, self.len());
154                // SAFETY: we know that we never go out of bounds.
155                unsafe { Ok(self.take_unchecked(&idx)) }
156            },
157        }
158    }
159
160    /// Sample a fraction between 0.0-1.0 of this [`ChunkedArray`].
161    pub fn sample_frac(
162        &self,
163        frac: f64,
164        with_replacement: bool,
165        shuffle: Option<bool>,
166        seed: Option<u64>,
167    ) -> PolarsResult<Self> {
168        let n = (self.len() as f64 * frac) as usize;
169        self.sample_n(n, with_replacement, shuffle, seed)
170    }
171}
172
173impl DataFrame {
174    /// Sample n datapoints from this [`DataFrame`].
175    pub fn sample_n(
176        &self,
177        n: &Series,
178        with_replacement: bool,
179        shuffle: Option<bool>,
180        seed: Option<u64>,
181    ) -> PolarsResult<Self> {
182        polars_ensure!(
183        n.len() == 1,
184        ComputeError: "Sample size must be a single value."
185        );
186
187        let n = n.strict_cast(&IDX_DTYPE)?;
188        let n = n.idx()?;
189
190        match n.get(0) {
191            Some(n) => self.sample_n_literal(n as usize, with_replacement, shuffle, seed),
192            None => Ok(self.clear()),
193        }
194    }
195
196    pub fn sample_n_literal(
197        &self,
198        n: usize,
199        with_replacement: bool,
200        shuffle: Option<bool>,
201        seed: Option<u64>,
202    ) -> PolarsResult<Self> {
203        ensure_shape(n, self.height(), with_replacement)?;
204        // All columns should used the same indices. So we first create the indices.
205        let idx = match with_replacement {
206            true => create_rand_index_with_replacement(n, self.height(), seed, shuffle),
207            false => create_rand_index_no_replacement(n, self.height(), seed, shuffle),
208        };
209        // SAFETY: the indices are within bounds.
210        Ok(unsafe { self.take_unchecked(&idx) })
211    }
212
213    /// Sample a fraction between 0.0-1.0 of this [`DataFrame`].
214    pub fn sample_frac(
215        &self,
216        frac: &Series,
217        with_replacement: bool,
218        shuffle: Option<bool>,
219        seed: Option<u64>,
220    ) -> PolarsResult<Self> {
221        polars_ensure!(
222        frac.len() == 1,
223        ComputeError: "Sample fraction must be a single value."
224        );
225
226        let frac = frac.cast(&Float64)?;
227        let frac = frac.f64()?;
228
229        match frac.get(0) {
230            Some(frac) => {
231                let n = (self.height() as f64 * frac) as usize;
232                self.sample_n_literal(n, with_replacement, shuffle, seed)
233            },
234            None => Ok(self.clear()),
235        }
236    }
237}
238
239#[cfg(test)]
240mod test {
241    use super::*;
242
243    #[test]
244    fn test_sample() {
245        let df = df![
246            "foo" => &[1, 2, 3, 4, 5]
247        ]
248        .unwrap();
249
250        // Default samples are random and don't require seeds.
251        assert!(
252            df.sample_n(
253                &Series::new(PlSmallStr::from_static("s"), &[3]),
254                false,
255                None,
256                None
257            )
258            .is_ok()
259        );
260        assert!(
261            df.sample_frac(
262                &Series::new(PlSmallStr::from_static("frac"), &[0.4]),
263                false,
264                None,
265                None
266            )
267            .is_ok()
268        );
269        // With seeding.
270        assert!(
271            df.sample_n(
272                &Series::new(PlSmallStr::from_static("s"), &[3]),
273                false,
274                None,
275                Some(0)
276            )
277            .is_ok()
278        );
279        assert!(
280            df.sample_frac(
281                &Series::new(PlSmallStr::from_static("frac"), &[0.4]),
282                false,
283                None,
284                Some(0)
285            )
286            .is_ok()
287        );
288        // Without replacement can not sample more than 100%.
289        assert!(
290            df.sample_frac(
291                &Series::new(PlSmallStr::from_static("frac"), &[2.0]),
292                false,
293                None,
294                Some(0)
295            )
296            .is_err()
297        );
298        assert!(
299            df.sample_n(
300                &Series::new(PlSmallStr::from_static("s"), &[3]),
301                true,
302                None,
303                Some(0)
304            )
305            .is_ok()
306        );
307        assert!(
308            df.sample_frac(
309                &Series::new(PlSmallStr::from_static("frac"), &[0.4]),
310                true,
311                None,
312                Some(0)
313            )
314            .is_ok()
315        );
316        // With replacement can sample more than 100%.
317        assert!(
318            df.sample_frac(
319                &Series::new(PlSmallStr::from_static("frac"), &[2.0]),
320                true,
321                None,
322                Some(0)
323            )
324            .is_ok()
325        );
326    }
327}
polars_core/chunked_array/random.rs

polars_core/chunked_array/
random.rs
