ruda_dataset/transform/
composed.rs

1use crate::Dataset;
2
3/// Compose multiple datasets together to create a bigger one.
4#[derive(new)]
5pub struct ComposedDataset<D> {
6    datasets: Vec<D>,
7}
8
9impl<D, I> Dataset<I> for ComposedDataset<D>
10where
11    D: Dataset<I>,
12    I: Clone,
13{
14    fn get(&self, index: usize) -> Option<I> {
15        let mut current_index = 0;
16        for dataset in self.datasets.iter() {
17            if index < dataset.len() + current_index {
18                return dataset.get(index - current_index);
19            }
20            current_index += dataset.len();
21        }
22        None
23    }
24
25    fn get_many(&self, indices: &[usize]) -> Option<Vec<I>> {
26        if indices.is_empty() {
27            return Some(Vec::new());
28        }
29
30        let mut offset = 0;
31        let ends: Vec<_> = self.datasets
32            .iter()
33            .map(|dataset| {
34                offset += dataset.len();
35                offset
36            })
37            .collect();
38        if indices.iter().any(|&index| index >= offset) {
39            return None;
40        }
41        let mut items = Vec::with_capacity(indices.len());
42        let mut local_indices = Vec::new();
43        let mut start = 0;
44        while start < indices.len() {
45            let dataset = ends.partition_point(|&end| end <= indices[start]);
46            let offset = if dataset == 0 { 0 } else { ends[dataset - 1] };
47            let mut end = start + 1;
48            while end < indices.len() && indices[end] >= offset && indices[end] < ends[dataset] {
49                end += 1;
50            }
51            local_indices.clear();
52            local_indices.extend(indices[start..end].iter().map(|&index| index - offset));
53            let batch = self.datasets[dataset].get_many(&local_indices)?;
54            if batch.len() != local_indices.len() {
55                return None;
56            }
57            items.extend(batch);
58            start = end;
59        }
60        Some(items)
61    }
62
63    fn len(&self) -> usize {
64        let mut total = 0;
65        for dataset in self.datasets.iter() {
66            total += dataset.len();
67        }
68        total
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75    use crate::FakeDataset;
76
77    #[test]
78    fn test_composed_dataset() {
79        let dataset1 = FakeDataset::<String>::new(10);
80        let dataset2 = FakeDataset::<String>::new(5);
81
82        let items1 = dataset1.iter().collect::<Vec<_>>();
83        let items2 = dataset2.iter().collect::<Vec<_>>();
84
85        let composed = ComposedDataset::new(vec![dataset1, dataset2]);
86
87        assert_eq!(composed.len(), 15);
88
89        let expected_items: Vec<String> = items1.iter().chain(items2.iter()).cloned().collect();
90
91        let items = composed.iter().collect::<Vec<_>>();
92
93        assert_eq!(items, expected_items);
94    }
95}
ruda_dataset/transform/composed.rs

ruda_dataset/transform/
composed.rs
