dataloader-rs 0.3.0

High-performance DataLoader for Rust with a PyTorch-like interface
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

use std::collections::{BTreeMap, HashMap};
use std::hash::{BuildHasher, Hash};

use crate::{collator::default_collator::DefaultCollate, error::Result};

const DIFFERENT_KEYS: &str = "all maps in batch must have the same keys";

// Both impls transpose the batch into one column of values per key, then
// collate each column; they differ only in the map type.

impl<K, V, H> DefaultCollate for HashMap<K, V, H>
where
    K: Eq + Hash + Send + 'static,
    V: DefaultCollate,
    H: BuildHasher + Default + Send + 'static,
{
    type Batch = HashMap<K, V::Batch, H>;

    fn collate_items(items: Vec<Self>) -> Result<Self::Batch> {
        let rows = items.len();
        let mut maps = items.into_iter();
        let mut columns: HashMap<K, Vec<V>, H> = match maps.next() {
            Some(first) => first
                .into_iter()
                .map(|(k, v)| (k, column(rows, v)))
                .collect(),
            None => HashMap::default(),
        };
        for map in maps {
            if map.len() != columns.len() {
                return Err(DIFFERENT_KEYS.into());
            }
            for (k, v) in map {
                columns.get_mut(&k).ok_or(DIFFERENT_KEYS)?.push(v);
            }
        }
        columns
            .into_iter()
            .map(|(k, values)| Ok((k, V::collate_items(values)?)))
            .collect()
    }
}

impl<K, V> DefaultCollate for BTreeMap<K, V>
where
    K: Ord + Send + 'static,
    V: DefaultCollate,
{
    type Batch = BTreeMap<K, V::Batch>;

    fn collate_items(items: Vec<Self>) -> Result<Self::Batch> {
        let rows = items.len();
        let mut maps = items.into_iter();
        let mut columns: BTreeMap<K, Vec<V>> = match maps.next() {
            Some(first) => first
                .into_iter()
                .map(|(k, v)| (k, column(rows, v)))
                .collect(),
            None => BTreeMap::new(),
        };
        for map in maps {
            if map.len() != columns.len() {
                return Err(DIFFERENT_KEYS.into());
            }
            for (k, v) in map {
                columns.get_mut(&k).ok_or(DIFFERENT_KEYS)?.push(v);
            }
        }
        columns
            .into_iter()
            .map(|(k, values)| Ok((k, V::collate_items(values)?)))
            .collect()
    }
}

/// A key's column, sized for the whole batch, starting with `first`.
fn column<V>(rows: usize, first: V) -> Vec<V> {
    let mut column = Vec::with_capacity(rows);
    column.push(first);
    column
}

#[cfg(test)]
mod tests {
    use std::collections::{BTreeMap, HashMap};

    use crate::collator::{Collator, DefaultCollator};

    #[test]
    fn collates_hashmaps_by_key() {
        let c = DefaultCollator;
        let m1 = HashMap::from([("x", 1u32), ("y", 2u32)]);
        let m2 = HashMap::from([("x", 10u32), ("y", 20u32)]);
        let out = c.collate(vec![m1, m2]).unwrap();
        assert_eq!(out["x"], vec![1, 10]);
        assert_eq!(out["y"], vec![2, 20]);
    }

    #[test]
    fn hashmaps_with_different_keys_fail() {
        let c = DefaultCollator;
        let m1 = HashMap::from([("x", 1u32)]);
        let m2 = HashMap::from([("y", 2u32)]);
        assert!(c.collate(vec![m1, m2]).is_err());
    }

    #[test]
    fn collates_btreemaps_by_key() {
        let c = DefaultCollator;
        let m1 = BTreeMap::from([("x", 1u32), ("y", 2u32)]);
        let m2 = BTreeMap::from([("x", 10u32), ("y", 20u32)]);
        let out = c.collate(vec![m1, m2]).unwrap();
        assert_eq!(out["x"], vec![1, 10]);
        assert_eq!(out["y"], vec![2, 20]);
    }

    #[test]
    fn btreemaps_with_other_keys_fail() {
        let c = DefaultCollator;
        let m1 = BTreeMap::from([("x", 1u32), ("y", 2u32)]);
        let m2 = BTreeMap::from([("x", 10u32), ("z", 20u32)]);
        assert!(c.collate(vec![m1, m2]).is_err());
        let m3 = BTreeMap::from([("x", 1u32)]);
        let m4 = BTreeMap::from([("x", 10u32), ("y", 20u32)]);
        assert!(c.collate(vec![m3, m4]).is_err());
    }
}