simd_csv/
total_reader.rs

1use crate::core::{CoreReader, ReadResult};
2use crate::error;
3use crate::records::{ByteRecord, ByteRecordBuilder};
4
5pub struct TotalReaderBuilder {
6    delimiter: u8,
7    quote: u8,
8}
9
10impl Default for TotalReaderBuilder {
11    fn default() -> Self {
12        Self {
13            delimiter: b',',
14            quote: b'"',
15        }
16    }
17}
18
19impl TotalReaderBuilder {
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    pub fn delimiter(&mut self, delimiter: u8) -> &mut Self {
25        self.delimiter = delimiter;
26        self
27    }
28
29    pub fn quote(&mut self, quote: u8) -> &mut Self {
30        self.quote = quote;
31        self
32    }
33
34    pub fn from_bytes<'b>(&self, bytes: &'b [u8]) -> TotalReader<'b> {
35        TotalReader {
36            inner: CoreReader::new(self.delimiter, self.quote),
37            bytes,
38            pos: 0,
39        }
40    }
41}
42
43// NOTE: a reader to be used when the whole data fits into memory or when using
44// memory maps.
45pub struct TotalReader<'b> {
46    inner: CoreReader,
47    bytes: &'b [u8],
48    pos: usize,
49}
50
51impl<'b> TotalReader<'b> {
52    pub fn from_bytes(bytes: &'b [u8]) -> Self {
53        TotalReaderBuilder::new().from_bytes(bytes)
54    }
55
56    pub fn count_records(&mut self) -> u64 {
57        use ReadResult::*;
58
59        let mut count: u64 = 0;
60
61        loop {
62            let (result, pos) = self.inner.split_record(&self.bytes[self.pos..]);
63
64            self.pos += pos;
65
66            match result {
67                End => break,
68                InputEmpty | Cr | Lf => continue,
69                Record => {
70                    count += 1;
71                }
72            };
73        }
74
75        count
76    }
77
78    pub fn read_byte_record(&mut self, record: &mut ByteRecord) -> error::Result<bool> {
79        use ReadResult::*;
80
81        record.clear();
82
83        let mut record_builder = ByteRecordBuilder::wrap(record);
84
85        loop {
86            let (result, pos) = self
87                .inner
88                .read_record(&self.bytes[self.pos..], &mut record_builder);
89
90            self.pos += pos;
91
92            match result {
93                End => {
94                    return Ok(false);
95                }
96                Cr | Lf | InputEmpty => {
97                    continue;
98                }
99                Record => {
100                    return Ok(true);
101                }
102            };
103        }
104    }
105}