lumen_dataset/transform/
selection.rs

1use std::sync::Arc;
2use rand::seq::SliceRandom;
3use crate::Dataset;
4
5pub struct SubsetDataset<D> 
6where 
7    D: Dataset,
8{
9    pub wrapped: Arc<D>,
10    pub indices: Vec<usize>,
11}
12
13impl<D: Dataset> SubsetDataset<D> {
14    /// Creates a new selection dataset with the given dataset and indices without checking bounds.
15    ///
16    /// ## Arguments
17    ///
18    /// * `dataset` - The original dataset to select from.
19    /// * `indices` - A vector of indices to select from the dataset.
20    pub fn new(dataset: impl Into<Arc<D>>, indices: Vec<usize>) -> Self {
21        Self { wrapped: dataset.into(), indices }
22    }
23
24    /// Creates a new selection dataset with the given dataset and indices.
25    ///
26    /// Checks that all indices are within the bounds of the dataset.
27    ///
28    /// ## Arguments
29    ///
30    /// * `dataset` - The original dataset to select from.
31    /// * `indices` - A slice of indices to select from the dataset.
32    ///   These indices must be within the bounds of the dataset.
33    pub fn from_indices(dataset: impl Into<Arc<D>>, indices: Vec<usize>) -> SubsetDatasetResult<Self> {
34        let dataset = dataset.into();
35        
36        // check index
37        let size = dataset.len();
38        if let Some(&idx) = indices.iter().find(|&i| *i >= size) {
39            Err(SubsetDatasetError::IndexOutOfRangeWhenSelectDataset(idx, size, "from indices"))?;
40        }
41
42        Ok(Self::new(dataset, indices))
43    }
44
45    /// Creates a new selection dataset that selects all indices from the dataset.
46    ///
47    /// This allocates a 1-to-1 mapping of indices to the dataset size,
48    /// essentially functioning as a no-op selection. This is only useful
49    /// when the dataset will later be shuffled or transformed in place.
50    ///
51    /// ## Arguments
52    ///
53    /// * `dataset` - The original dataset to select from.
54    ///
55    /// ## Returns
56    ///
57    /// A new `SubsetDataset` that selects all indices from the dataset.
58    pub fn select_all(dataset: impl Into<Arc<D>>) -> Self {
59        let dataset = dataset.into();
60        let size = dataset.len();
61        Self::new(dataset, iota(size))
62    }
63
64    /// Creates a new dataset that is a slice of the current selection dataset.
65    ///
66    /// Slices the *selection indices* from ``[start..end]``.
67    ///
68    /// Independent of future shuffles on the parent, but shares the same wrapped dataset.
69    ///
70    /// ## Arguments
71    ///
72    /// * `start` - The start of the range.
73    /// * `end` - The end of the range (exclusive).
74    pub fn slice(&self, start: usize, end: usize) -> SubsetDatasetResult<Self> {
75        if start >= self.len() {
76            Err(SubsetDatasetError::IndexOutOfRangeWhenSelectDataset(start, self.len(), "slice"))?;
77        }
78        if end >= self.len() {
79            Err(SubsetDatasetError::IndexOutOfRangeWhenSelectDataset(end, self.len(), "slice"))?;
80        }
81        Ok(Self::new(self.wrapped.clone(), self.indices[start..end].to_vec()))
82    }
83
84    /// Split into `num` datasets by slicing the selection indices evenly.
85    ///
86    /// Split is done via `slice`, so the datasets share the same wrapped dataset.
87    ///
88    /// Independent of future shuffles on the parent, but shares the same wrapped dataset.
89    ///
90    /// ## Arguments
91    ///
92    /// * `num` - The number of datasets to split into.
93    ///
94    /// ## Returns
95    ///
96    /// A vector of `SubsetDataset` instances, each containing a subset of the indices.
97    pub fn split(&self, num: usize) -> SubsetDatasetResult<Vec<Self>> {
98        if num == 0 {
99            Err(SubsetDatasetError::NumSplitZeroWhenSelectDataset)?;
100        }
101        
102        let n = self.indices.len();
103        let mut datasets = Vec::with_capacity(num);
104        
105        let base_size = n / num;
106        let remainder = n % num;
107
108        let mut start = 0;
109        for i in 0..num {
110            let size = base_size + (if i < remainder { 1 } else { 0 });
111            let end = start + size;
112            
113            datasets.push(self.slice(start, end)?);
114            
115            start = end;
116        }
117
118        Ok(datasets)
119    }
120}
121
122impl<D: Dataset> Dataset for SubsetDataset<D> {
123    type Item = D::Item;
124    type Error = D::Error;
125
126    fn get(&self, index: usize) -> Result<Option<Self::Item>, Self::Error> {
127        let index = match self.indices.get(index) {
128            Some(index) => index.clone(),
129            None => return Ok(None),
130        };
131        self.wrapped.get(index)
132    }
133
134    fn len(&self) -> usize {
135        self.indices.len()
136    }
137}
138
139pub fn random_split<D: Dataset>(dataset: D, ratio: f64) -> (SubsetDataset<D>, SubsetDataset<D>) {
140    let length = dataset.len();
141    let mut indices = iota(length);
142    
143    let mut rng = rand::rng();
144    indices.shuffle(&mut rng);
145
146    let split_idx = (length as f64 * ratio).floor() as usize;
147    let (indices1, indices2) = indices.split_at(split_idx);
148
149    let dataset = Arc::new(dataset);
150    let subset1 = SubsetDataset::new(dataset.clone(), indices1.to_vec());
151    let subset2 = SubsetDataset::new(dataset.clone(), indices2.to_vec());
152
153    (subset1, subset2)
154}
155
156#[inline(always)]
157fn iota(size: usize) -> Vec<usize> {
158    (0..size).collect()
159}
160
161#[derive(Debug, thiserror::Error)]
162pub enum SubsetDatasetError {
163    #[error("Index out of bounds for wrapped dataset size: {0} >= {1} when {2}")]
164    IndexOutOfRangeWhenSelectDataset(usize, usize, &'static str),
165
166    #[error("Use 0 to split dataset")]
167    NumSplitZeroWhenSelectDataset,
168}
169
170pub type SubsetDatasetResult<T> = std::result::Result<T, SubsetDatasetError>;
lumen_dataset/transform/selection.rs

lumen_dataset/transform/
selection.rs
