ai_dataloader/collate/default_collate/
map.rs

1use super::super::Collate;
2use super::DefaultCollate;
3use std::{
4    cmp::Eq,
5    collections::{BTreeMap, HashMap},
6    hash::{BuildHasher, Hash},
7};
8
9impl<K, V, H> Collate<HashMap<K, V, H>> for DefaultCollate
10where
11    K: Eq + Hash + Clone,
12    V: Clone,
13    Self: Collate<V>,
14    H: BuildHasher,
15{
16    type Output = HashMap<K, <Self as Collate<V>>::Output>;
17    fn collate(&self, batch: Vec<HashMap<K, V, H>>) -> Self::Output {
18        let mut collated = HashMap::with_capacity(batch[0].keys().len());
19        for key in batch[0].keys() {
20            let vec: Vec<_> = batch.iter().map(|hash_map| hash_map[key].clone()).collect();
21            collated.insert(key.clone(), self.collate(vec));
22        }
23        collated
24    }
25}
26impl<K, V> Collate<BTreeMap<K, V>> for DefaultCollate
27where
28    K: Ord + Clone,
29    V: Clone,
30    Self: Collate<V>,
31{
32    type Output = BTreeMap<K, <Self as Collate<V>>::Output>;
33    fn collate(&self, batch: Vec<BTreeMap<K, V>>) -> Self::Output {
34        let mut collated = BTreeMap::new();
35        for key in batch[0].keys() {
36            let vec: Vec<_> = batch.iter().map(|hash_map| hash_map[key].clone()).collect();
37            collated.insert(key.clone(), self.collate(vec));
38        }
39        collated
40    }
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46    use ndarray::array;
47
48    #[test]
49    fn vec_of_hash_map() {
50        let map1 = HashMap::from([("A", 0), ("B", 1)]);
51        let map2 = HashMap::from([("A", 100), ("B", 100)]);
52        let expected_result = HashMap::from([("A", array![0, 100]), ("B", array![1, 100])]);
53        assert_eq!(DefaultCollate.collate(vec![map1, map2]), expected_result);
54
55        // Same value type but different key
56        let map1 = HashMap::from([(1, 0), (2, 1)]);
57        let map2 = HashMap::from([(1, 100), (2, 100)]);
58        let expected_result = HashMap::from([(1, array![0, 100]), (2, array![1, 100])]);
59        assert_eq!(DefaultCollate.collate(vec![map1, map2]), expected_result);
60
61        let map1 = HashMap::from([("A", 0.0), ("B", 1.0)]);
62        let map2 = HashMap::from([("A", 100.0), ("B", 100.0)]);
63        let expected_result = HashMap::from([("A", array![0.0, 100.0]), ("B", array![1.0, 100.0])]);
64        assert_eq!(DefaultCollate.collate(vec![map1, map2]), expected_result);
65    }
66
67    #[test]
68    fn specialized() {
69        let map1 = HashMap::from([("A", String::from("0")), ("B", String::from("1"))]);
70        let map2 = HashMap::from([("A", String::from("100")), ("B", String::from("100"))]);
71        let expected_result = HashMap::from([
72            ("A", vec![String::from("0"), String::from("100")]),
73            ("B", vec![String::from("1"), String::from("100")]),
74        ]);
75        assert_eq!(DefaultCollate.collate(vec![map1, map2]), expected_result);
76
77        let map1 = HashMap::from([("A", "0"), ("B", "1")]);
78        let map2 = HashMap::from([("A", "100"), ("B", "100")]);
79        let expected_result = HashMap::from([("A", vec!["0", "100"]), ("B", vec!["1", "100"])]);
80        assert_eq!(DefaultCollate.collate(vec![map1, map2]), expected_result);
81    }
82}
ai_dataloader/collate/default_collate/map.rs

ai_dataloader/collate/default_collate/
map.rs
