Skip to main content

ruda_dataset/transform/
shuffle.rs

1use crate::Dataset;
2use crate::transform::{RngSource, SelectionDataset};
3
4/// A Shuffled a dataset.
5///
6/// This is a thin wrapper around a [SelectionDataset] which selects and shuffles
7/// the full indices of the original dataset.
8///
9/// Consider using [SelectionDataset] if you are only interested in
10/// shuffling mechanisms.
11///
12/// Consider using [sampler dataset](crate::transform::SamplerDataset) if you
13/// want a probability distribution which is computed lazily.
14pub struct ShuffledDataset<D, I>
15where
16    D: Dataset<I>,
17    I: Clone + Send + Sync,
18{
19    wrapped: SelectionDataset<D, I>,
20}
21
22impl<D, I> ShuffledDataset<D, I>
23where
24    D: Dataset<I>,
25    I: Clone + Send + Sync,
26{
27    /// Creates a new selection dataset with shuffled indices.
28    ///
29    /// This is a thin wrapper around `SelectionDataset::new_shuffled`.
30    ///
31    /// # Arguments
32    ///
33    /// * `dataset` - The original dataset to select from.
34    /// * `rng_source` - The source of the random number generator.
35    ///
36    /// # Returns
37    ///
38    /// A new `ShuffledDataset`.
39    pub fn new<R>(dataset: D, rng_source: R) -> Self
40    where
41        R: Into<RngSource>,
42    {
43        Self {
44            wrapped: SelectionDataset::new_shuffled(dataset, rng_source),
45        }
46    }
47
48    /// Creates a new selection dataset with shuffled indices using a fixed seed.
49    ///
50    /// This is a thin wrapper around `SelectionDataset::new_shuffled_with_seed`.
51    ///
52    /// # Arguments
53    ///
54    /// * `dataset` - The original dataset to select from.
55    /// * `seed` - A fixed seed for the random number generator.
56    ///
57    /// # Returns
58    ///
59    /// A new `ShuffledDataset`.
60    #[deprecated(since = "0.19.0", note = "Use `new(dataset, seed)` instead`")]
61    pub fn with_seed(dataset: D, seed: u64) -> Self {
62        Self::new(dataset, seed)
63    }
64}
65
66impl<D, I> Dataset<I> for ShuffledDataset<D, I>
67where
68    D: Dataset<I>,
69    I: Clone + Send + Sync,
70{
71    fn get(&self, index: usize) -> Option<I> {
72        self.wrapped.get(index)
73    }
74
75    fn get_many(&self, indices: &[usize]) -> Option<Vec<I>> {
76        self.wrapped.get_many(indices)
77    }
78
79    fn len(&self) -> usize {
80        self.wrapped.len()
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87    use crate::FakeDataset;
88    use crate::transform::selection::shuffled_indices;
89    use rand::SeedableRng;
90    use rand::prelude::StdRng;
91
92    #[test]
93    fn test_shuffled_dataset() {
94        let dataset = FakeDataset::<String>::new(27);
95        let source_items = dataset.iter().collect::<Vec<_>>();
96
97        let seed = 42;
98
99        #[allow(deprecated)]
100        let shuffled = ShuffledDataset::with_seed(dataset, seed);
101
102        let mut rng = StdRng::seed_from_u64(seed);
103        let indices = shuffled_indices(source_items.len(), &mut rng);
104
105        assert_eq!(shuffled.len(), source_items.len());
106
107        let expected_items: Vec<_> = indices
108            .iter()
109            .map(|&i| source_items[i].to_string())
110            .collect();
111        assert_eq!(&shuffled.iter().collect::<Vec<_>>(), &expected_items);
112    }
113}