simd_csv/select/
headers.rs

1use std::collections::BTreeMap;
2use std::fmt;
3use std::ops::Index;
4
5use crate::debug;
6use crate::records::ByteRecord;
7
8#[derive(Debug, PartialEq, Clone)]
9pub enum ColumIndexationBy<'b> {
10    Name(&'b [u8]),
11    NameAndNth(&'b [u8], isize),
12    Pos(isize),
13}
14
15impl ColumIndexationBy<'_> {
16    pub fn has_name(&self) -> bool {
17        matches!(self, Self::Name(_) | Self::NameAndNth(_, _))
18    }
19}
20
21/// A [`ByteRecord`] wrapper for CSV headers that can be used to efficiently find
22/// columns by name, position or even both.
23pub struct ByteHeadersIndex {
24    inner: ByteRecord,
25    map: Option<BTreeMap<Vec<u8>, Vec<usize>>>,
26}
27
28impl fmt::Debug for ByteHeadersIndex {
29    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
30        write!(f, "ByteHeadersIndex(")?;
31        f.debug_list()
32            .entries(self.inner.iter().map(debug::Bytes))
33            .finish()?;
34        write!(f, ")")?;
35        Ok(())
36    }
37}
38
39impl ByteHeadersIndex {
40    pub fn new(record: ByteRecord, has_names: bool) -> Self {
41        let map = if !has_names {
42            None
43        } else {
44            let mut map = BTreeMap::new();
45
46            for (i, name) in record.iter().enumerate() {
47                let indices = map.entry(name.to_vec()).or_insert_with(Vec::new);
48                indices.push(i);
49            }
50
51            Some(map)
52        };
53
54        Self { inner: record, map }
55    }
56
57    #[inline]
58    pub fn len(&self) -> usize {
59        self.inner.len()
60    }
61
62    #[inline]
63    pub fn is_empty(&self) -> bool {
64        self.inner.is_empty()
65    }
66
67    #[inline]
68    pub fn has_names(&self) -> bool {
69        self.map.is_some()
70    }
71
72    pub fn first_column_index_by_name(&self, name: impl AsRef<[u8]>) -> Option<usize> {
73        self.map
74            .as_ref()?
75            .get(name.as_ref())
76            .map(|indices| indices[0])
77    }
78
79    pub fn find_column_index(&self, indexation: ColumIndexationBy) -> Option<usize> {
80        match indexation {
81            ColumIndexationBy::Name(name) => self
82                .map
83                .as_ref()?
84                .get(name)
85                .and_then(|positions| positions.first())
86                .copied(),
87            ColumIndexationBy::Pos(pos) => {
88                let len = self.inner.len();
89
90                if pos < 0 {
91                    // Negative indexing
92                    let pos = pos.unsigned_abs();
93
94                    if pos > len {
95                        None
96                    } else {
97                        Some(len - pos)
98                    }
99                } else {
100                    let pos = pos as usize;
101
102                    if pos >= len {
103                        None
104                    } else {
105                        Some(pos)
106                    }
107                }
108            }
109            ColumIndexationBy::NameAndNth(name, pos) => self
110                .map
111                .as_ref()?
112                .get(name)
113                .and_then(|positions| {
114                    if pos < 0 {
115                        let pos = pos.unsigned_abs();
116                        let len = positions.len();
117
118                        if pos > len {
119                            None
120                        } else {
121                            positions.get(len - pos)
122                        }
123                    } else {
124                        positions.get(pos as usize)
125                    }
126                })
127                .copied(),
128        }
129    }
130}
131
132impl AsRef<ByteRecord> for ByteHeadersIndex {
133    fn as_ref(&self) -> &ByteRecord {
134        &self.inner
135    }
136}
137
138impl Index<usize> for ByteHeadersIndex {
139    type Output = [u8];
140
141    #[inline(always)]
142    fn index(&self, index: usize) -> &Self::Output {
143        &self.inner[index]
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[test]
152    fn test_byte_headers_index() {
153        let headers = brec!["name", "surname", "age", "name"];
154        let index = ByteHeadersIndex::new(headers.clone(), true);
155
156        assert_eq!(&headers, index.as_ref());
157        assert_eq!(index.has_names(), true);
158        assert_eq!(index.len(), 4);
159        assert_eq!(index.is_empty(), false);
160        assert_eq!(index.first_column_index_by_name("surname"), Some(1));
161    }
162}
simd_csv/select/headers.rs

simd_csv/select/
headers.rs
