Skip to main content

rudb_csv/
scan.rs

1//! Bytes into fields.
2//!
3//! Two ways of doing it, which agree. [`record`] reads one record into owned strings, and it is
4//! what the sniffer runs: working out the delimiter means splitting the sample under each candidate
5//! and seeing which one gives every line the same number of fields, and that has to be the same
6//! splitting the reader will do or the sniffer answers a question nobody asked. [`records`] reads
7//! a chunk's worth of records at once and writes down where each field is rather than copying it,
8//! and it is what the reader runs. The tests at the bottom hold the two to the same answer on a
9//! few thousand generated files, which is what lets the sniffer keep the simple one.
10//!
11//! A record ends at a newline that is not inside a quoted field, and the three line endings are all
12//! accepted whatever the file mostly uses. DuckDB reports the one it found and reads either, and a
13//! file that a text editor has half converted is a real thing.
14//!
15//! [`records`] finds the delimiters, quotes and line endings sixty four bytes at a time, as one bit
16//! a byte in a `u64`, and works out which of them are inside a quoted field with a prefix XOR over
17//! the quote bits, the way simdjson and simdcsv do. A field only has to be looked at byte by byte
18//! when it has a quote in it, and then only to check that the quotes are where a quoted field puts
19//! them. A quote anywhere else, which is a literal character in the middle of a bare field, hands
20//! that one record to the byte loop, since the parity trick has no way to know the quote was not
21//! meant. A dialect whose escape is not its quote is read by the byte loop throughout.
22//!
23//! Fields used to arrive as owned strings, a copy per field before the conversion made a second
24//! one. The borrowed ranges were the change to make when there was a number saying it mattered, and
25//! the number is this: on a 200,000 row file shaped like TPC-H lineitem, 27 MB, one core of an
26//! Apple M4 read 95 MB/s a record at a time and 613 MB/s a chunk at a time, with the scan itself
27//! running at about 3 GB/s and the rest going to the column conversions. The ignored test
28//! `reads_lineitem_faster_a_chunk_at_a_time` in the reader measures it again.
29
30use std::borrow::Cow;
31
32use rudb_common::{Error, Result};
33
34use crate::dialect::Dialect;
35
36/// Reads the record starting at `from` into `out`, and answers where the next one starts.
37///
38/// `None` means the buffer does not hold a whole record: either it ran out mid line and more bytes
39/// may follow, or a quoted field was left open. With `eof` set there are no more bytes, so a last
40/// line with no newline on the end is a record like any other and only an open quote is short.
41///
42/// # Errors
43///
44/// When a quoted field has something other than a delimiter or a line ending after its closing
45/// quote, which is a file that is not the file it claims to be and reading past it would invent
46/// data.
47pub fn record(
48    bytes: &[u8],
49    from: usize,
50    dialect: Dialect,
51    eof: bool,
52    out: &mut Vec<String>,
53) -> Result<Option<usize>> {
54    if from >= bytes.len() {
55        // No bytes left is no record, and it is not a record of one empty field. The difference
56        // matters at the end of a file, where a reader that took the second reading would hand back
57        // an empty row forever, and it matters for the last line of a file that ends in a newline,
58        // which is not a row of nothing.
59        return Ok(None);
60    }
61    let quote = dialect.quote_byte();
62    let escape = dialect.escape_byte();
63    let mut at = from;
64    let mut count = 0;
65    let mut field = Vec::new();
66    loop {
67        field.clear();
68        if bytes.get(at) == Some(&quote) {
69            at += 1;
70            loop {
71                let Some(&byte) = bytes.get(at) else { return Ok(None) };
72                if byte == escape && escape != quote {
73                    // An escape of its own, such as the backslash the Join Order Benchmark's
74                    // files are written with, keeps whatever byte follows it, so `\"` is a quote
75                    // in the value, `\\` is one backslash and `\n` is the letter n. That last one
76                    // looks like a line break and is not one, which is DuckDB's reading.
77                    let Some(&next) = bytes.get(at + 1) else { return Ok(None) };
78                    field.push(next);
79                    at += 2;
80                    continue;
81                }
82                if byte == escape && bytes.get(at + 1) == Some(&quote) {
83                    field.push(quote);
84                    at += 2;
85                    continue;
86                }
87                if byte == quote {
88                    at += 1;
89                    break;
90                }
91                field.push(byte);
92                at += 1;
93            }
94            match bytes.get(at) {
95                None if !eof => return Ok(None),
96                None => {}
97                Some(&byte) if byte == dialect.delimiter || byte == b'\n' || byte == b'\r' => {}
98                Some(&byte) => return Err(after_quote(byte)),
99            }
100        } else {
101            while let Some(&byte) = bytes.get(at) {
102                if byte == dialect.delimiter || byte == b'\n' || byte == b'\r' {
103                    break;
104                }
105                field.push(byte);
106                at += 1;
107            }
108            if at >= bytes.len() && !eof {
109                return Ok(None);
110            }
111        }
112        place(out, count, &field);
113        count += 1;
114        match bytes.get(at) {
115            Some(&byte) if byte == dialect.delimiter => at += 1,
116            Some(b'\r') => {
117                at += 1;
118                if bytes.get(at) == Some(&b'\n') {
119                    at += 1;
120                } else if at >= bytes.len() && !eof {
121                    // A trailing carriage return may be the first half of a `\r\n` that has not
122                    // arrived, and guessing wrong here splits one record into two.
123                    return Ok(None);
124                }
125                break;
126            }
127            Some(b'\n') => {
128                at += 1;
129                break;
130            }
131            Some(_) => unreachable!("a field stops at a delimiter, a line ending or the end"),
132            None => break,
133        }
134    }
135    out.truncate(count);
136    Ok(Some(at))
137}
138
139/// Writes the field at `at`, reusing the string that is already there.
140///
141/// A CSV file is bytes and the encoding is not stated anywhere in it. UTF-8 comes through unchanged
142/// because no byte of a multi byte sequence can be a delimiter or a quote. A byte that is not valid
143/// UTF-8 becomes the replacement character rather than an error, which is what one bad byte in a
144/// text file deserves, and which is also what DuckDB does when it is not told to check.
145fn place(out: &mut Vec<String>, at: usize, field: &[u8]) {
146    let text = String::from_utf8_lossy(field);
147    match out.get_mut(at) {
148        Some(held) => {
149            held.clear();
150            held.push_str(&text);
151        }
152        None => out.push(text.into_owned()),
153    }
154}
155
156/// The error for a quoted field with something after its closing quote, shared by both readers so
157/// that they cannot come to say it differently.
158fn after_quote(byte: u8) -> Error {
159    Error::io(format!(
160        "a quoted value is followed by '{}' rather than by a delimiter or the end of the line",
161        byte as char
162    ))
163}
164
165/// Where one field's bytes are in the buffer it was read from.
166///
167/// A quoted field's range is what is between its quotes. Whether a doubled quote or an escaped one
168/// is in there is kept too, since then the bytes are not the value yet and [`Span::text`] has to
169/// take the escapes out, which is the one case that copies before the value is used.
170///
171/// Eight bytes, two `u32`s with the escape in the top bit of the end, because a chunk holds one of
172/// these for every field of eight thousand records and the converters walk them once per column.
173/// That caps a chunk's buffer at [`Span::MOST`] bytes, which the reader keeps to.
174#[derive(Debug, Clone, Copy, PartialEq, Eq)]
175pub struct Span {
176    start: u32,
177    end: u32,
178}
179
180/// The bit of [`Span::end`] that says the field holds escapes.
181const ESCAPED: u32 = 1 << 31;
182
183impl Span {
184    /// The furthest into its buffer a range can reach.
185    pub const MOST: usize = (ESCAPED - 1) as usize;
186
187    /// The field from `start` to `end`, with escapes in it or not.
188    ///
189    /// # Panics
190    ///
191    /// In a debug build, if `end` is past [`Self::MOST`], which the reader never lets a buffer be.
192    #[must_use]
193    pub fn new(start: usize, end: usize, escaped: bool) -> Self {
194        debug_assert!(start <= end && end <= Self::MOST);
195        #[allow(clippy::cast_possible_truncation)]
196        let (start, end) = (start as u32, end as u32);
197        Self { start, end: end | if escaped { ESCAPED } else { 0 } }
198    }
199
200    /// The first byte of the field.
201    #[must_use]
202    pub const fn start(self) -> usize {
203        self.start as usize
204    }
205
206    /// One past the last byte of the field.
207    #[must_use]
208    pub const fn end(self) -> usize {
209        (self.end & !ESCAPED) as usize
210    }
211
212    /// Whether the bytes still hold escapes.
213    #[must_use]
214    pub const fn escaped(self) -> bool {
215        self.end & ESCAPED != 0
216    }
217
218    /// Whether the field is empty, which the reader turns into a null unless it was told a
219    /// different null string.
220    ///
221    /// A field with an escape in it is never empty once the escape is taken out, since every escape
222    /// leaves a byte behind, so the range answers for the value.
223    #[must_use]
224    pub const fn is_empty(self) -> bool {
225        self.start() == self.end()
226    }
227
228    /// How many bytes the field has in the buffer, escapes and all.
229    #[must_use]
230    pub const fn len(self) -> usize {
231        self.end() - self.start()
232    }
233
234    /// The bytes of the field as they are in the buffer, escapes and all.
235    #[must_use]
236    pub fn raw(self, bytes: &[u8]) -> &[u8] {
237        &bytes[self.start()..self.end()]
238    }
239
240    /// The field as text, borrowed from the buffer unless it had an escape or a byte that is not
241    /// UTF-8 in it.
242    ///
243    /// The bytes that are not UTF-8 become the replacement character, exactly as
244    /// `String::from_utf8_lossy` makes them, because that is what [`record`] does and the two have
245    /// to agree.
246    #[must_use]
247    pub fn text(self, bytes: &[u8], dialect: Dialect) -> Cow<'_, str> {
248        let raw = self.raw(bytes);
249        if !self.escaped() {
250            return String::from_utf8_lossy(raw);
251        }
252        Cow::Owned(String::from_utf8_lossy(&unescape(raw, dialect)).into_owned())
253    }
254}
255
256/// The bytes of a quoted field with its escapes taken out.
257///
258/// This pairs the bytes up exactly the way the loop that found the field did, from the left, an
259/// escape byte followed by a quote being one quote and anything else being itself. The range ends
260/// before the closing quote, so there is no closing quote in it to be mistaken for the second half
261/// of a pair, and reading it again from the same end gives the same pairs.
262///
263/// An escape that is not the quote pairs with whatever follows it, the same as in [`record`], and
264/// the loop that found the field never ends one on a lone escape, so every escape here has a byte
265/// after it.
266fn unescape(raw: &[u8], dialect: Dialect) -> Vec<u8> {
267    let quote = dialect.quote_byte();
268    let escape = dialect.escape_byte();
269    let mut out = Vec::with_capacity(raw.len());
270    let mut at = 0;
271    while let Some(&byte) = raw.get(at) {
272        if byte == escape
273            && escape != quote
274            && let Some(&next) = raw.get(at + 1)
275        {
276            out.push(next);
277            at += 2;
278            continue;
279        }
280        if byte == escape && raw.get(at + 1) == Some(&quote) {
281            out.push(quote);
282            at += 2;
283        } else {
284            out.push(byte);
285            at += 1;
286        }
287    }
288    out
289}
290
291/// A chunk's worth of records, as ranges into the buffer they were read from.
292///
293/// One list of fields for all of the records and one list of where each record's fields end, which
294/// is two allocations that are reused from chunk to chunk rather than one per row.
295#[derive(Debug, Clone, Default)]
296pub struct Records {
297    spans: Vec<Span>,
298    ends: Vec<usize>,
299}
300
301impl Records {
302    /// No records, with the room the last chunk needed kept.
303    pub fn clear(&mut self) {
304        self.spans.clear();
305        self.ends.clear();
306    }
307
308    /// How many whole records have been read.
309    #[must_use]
310    pub fn len(&self) -> usize {
311        self.ends.len()
312    }
313
314    /// Whether no whole record has been read.
315    #[must_use]
316    pub fn is_empty(&self) -> bool {
317        self.ends.is_empty()
318    }
319
320    /// The fields of record `at`, in order.
321    ///
322    /// # Panics
323    ///
324    /// If there is no record `at`.
325    #[must_use]
326    pub fn fields(&self, at: usize) -> &[Span] {
327        let start = if at == 0 { 0 } else { self.ends[at - 1] };
328        &self.spans[start..self.ends[at]]
329    }
330
331    /// Field `column` of record `row`, or `None` when the record is shorter than that.
332    #[must_use]
333    pub fn field(&self, row: usize, column: usize) -> Option<Span> {
334        let start = if row == 0 { 0 } else { self.ends[row - 1] };
335        let at = start + column;
336        if at < self.ends[row] { Some(self.spans[at]) } else { None }
337    }
338
339    /// Moves every range down by `by` bytes, for a buffer that has just had that many bytes taken
340    /// off its front.
341    pub fn shift(&mut self, by: usize) {
342        if by == 0 {
343            return;
344        }
345        // A shift is by less than the buffer is long, so it fits in the `u32` every range does.
346        #[allow(clippy::cast_possible_truncation)]
347        let by = by as u32;
348        for span in &mut self.spans {
349            span.start -= by;
350            span.end -= by;
351        }
352    }
353}
354
355/// Reads whole records starting at `from` into `out`, until it holds `limit` of them or the buffer
356/// has no whole record left, and answers where the first record it did not take starts.
357///
358/// The rules are [`record`]'s, including `eof`: without it a record the buffer ends in the middle
359/// of is left for when there are more bytes, and with it a last line with no newline is a record
360/// and only an open quote is short. Fewer than `limit` records and not at the end of the file means
361/// the caller should read more and ask again from the answer.
362///
363/// # Errors
364///
365/// The one [`record`] reports, for a quoted field with something after its closing quote. The
366/// records before it are in `out` and the one it was found in is not.
367pub fn records(
368    bytes: &[u8],
369    from: usize,
370    dialect: Dialect,
371    eof: bool,
372    limit: usize,
373    out: &mut Records,
374) -> Result<usize> {
375    let quote = dialect.quote_byte();
376    let delimiter = dialect.delimiter;
377    let structural = |byte: u8| byte == quote || byte == b'\n' || byte == b'\r';
378    if dialect.escape_byte() == quote && !structural(delimiter) {
379        blocks(bytes, from, dialect, eof, limit, out)
380    } else {
381        let mut at = from;
382        while out.len() < limit {
383            let Some(next) = spans(bytes, at, dialect, eof, &mut out.spans)? else { break };
384            out.ends.push(out.spans.len());
385            at = next;
386        }
387        Ok(at)
388    }
389}
390
391/// [`records`] for a dialect whose escape is its quote, sixty four bytes at a time.
392///
393/// Each block becomes three masks: the quotes, the bytes that can end a field, and the line endings
394/// among those. The prefix XOR of the quote mask has a bit set for every byte after an odd number of
395/// quotes, which is every byte inside a quoted field, and a doubled quote flips it twice between two
396/// adjacent bytes and so changes nothing. Whether the block ended inside a quote is carried into the next one. The field
397/// endings left once those bits are taken away are walked in order with `trailing_zeros`.
398///
399/// That is only right while every quote is where a quoted field puts it: first and last in its
400/// field, with any others in pairs. A field with no quote in it cannot be wrong, because the bits
401/// were outside a quote at its start and nothing in it flips them. A field with a quote in it is
402/// checked, and one that is not shaped like a quoted field goes to the byte loop from the start of
403/// its record, after which the blocks start again behind that record with the parity cleared.
404fn blocks(
405    bytes: &[u8],
406    from: usize,
407    dialect: Dialect,
408    eof: bool,
409    limit: usize,
410    out: &mut Records,
411) -> Result<usize> {
412    let delimiter = dialect.delimiter;
413    let quote = dialect.quote_byte();
414    let len = bytes.len();
415    let mut record = from;
416    loop {
417        // Where the fields of the record being read start in `out`, so that a record that turns
418        // out to be short or to need the byte loop can be taken back out.
419        let mut mark = out.spans.len();
420        let mut field = record;
421        let mut quoted = false;
422        let mut inside = 0u64;
423        let mut block = record;
424        'blocks: while block < len && out.len() < limit {
425            let end = len.min(block + 64);
426            let (quotes, ends, lines) = masks(&bytes[block..end], delimiter, quote);
427            let prefix = prefix_xor(quotes) ^ inside;
428            inside = 0u64.wrapping_sub(prefix >> 63);
429            let mut structural = ends & !prefix;
430            while structural != 0 {
431                let bit = structural.trailing_zeros() as usize;
432                structural &= structural - 1;
433                let at = block + bit;
434                if at < field {
435                    // The `\n` of a `\r\n`, which the `\r` already dealt with.
436                    continue;
437                }
438                let low = field.saturating_sub(block);
439                if quoted || quotes & below(bit) & !below(low) != 0 {
440                    let Some(escaped) = enclosed(&bytes[field..at], quote) else {
441                        break 'blocks;
442                    };
443                    out.spans.push(Span::new(field + 1, at - 1, escaped));
444                    quoted = false;
445                } else {
446                    out.spans.push(Span::new(field, at, false));
447                }
448                if lines >> bit & 1 == 0 {
449                    field = at + 1;
450                    continue;
451                }
452                let mut next = at + 1;
453                if bytes[at] == b'\r' {
454                    match bytes.get(next) {
455                        Some(b'\n') => next += 1,
456                        Some(_) => {}
457                        None if eof => {}
458                        // The first half of a `\r\n` whose second half has not been read, the
459                        // same wait `record` makes.
460                        None => {
461                            out.spans.truncate(mark);
462                            return Ok(record);
463                        }
464                    }
465                }
466                out.ends.push(out.spans.len());
467                mark = out.spans.len();
468                record = next;
469                field = next;
470                if out.len() == limit {
471                    return Ok(record);
472                }
473            }
474            if field < end && quotes >> field.saturating_sub(block) != 0 {
475                quoted = true;
476            }
477            block = end;
478        }
479        if out.len() >= limit || record >= len {
480            return Ok(record);
481        }
482        // Either a field whose quotes the blocks cannot vouch for, or the end of the buffer in the
483        // middle of a record. The byte loop settles both, the second one being where a last line
484        // with no newline, a record that is not all here yet and an open quote are told apart.
485        out.spans.truncate(mark);
486        let Some(next) = spans(bytes, record, dialect, eof, &mut out.spans)? else {
487            return Ok(record);
488        };
489        out.ends.push(out.spans.len());
490        record = next;
491    }
492}
493
494/// The quotes in a block of up to sixty four bytes, the bytes that can end a field, and of those
495/// the ones that can end a line, one bit a byte with the first byte in the lowest bit.
496///
497/// A short block at the end of the buffer is padded and the bits for the padding are cleared
498/// afterwards, so the padding byte does not have to be one that cannot be a delimiter.
499#[inline]
500fn masks(block: &[u8], delimiter: u8, quote: u8) -> (u64, u64, u64) {
501    let needles = [quote, delimiter, b'\n', b'\r'];
502    let [quotes, delimiters, newlines, returns] = if let Ok(full) = block.try_into() {
503        rudb_vector::bytes::masks(full, needles)
504    } else {
505        let mut padded = [0u8; 64];
506        padded[..block.len()].copy_from_slice(block);
507        let live = below(block.len());
508        rudb_vector::bytes::masks(&padded, needles).map(|mask| mask & live)
509    };
510    let lines = newlines | returns;
511    (quotes, delimiters | lines, lines)
512}
513
514/// Every bit set that has an odd number of set bits at or below it in `bits`.
515const fn prefix_xor(mut bits: u64) -> u64 {
516    bits ^= bits << 1;
517    bits ^= bits << 2;
518    bits ^= bits << 4;
519    bits ^= bits << 8;
520    bits ^= bits << 16;
521    bits ^= bits << 32;
522    bits
523}
524
525/// The bits below bit `count`, for a count under sixty four.
526const fn below(count: usize) -> u64 {
527    (1u64 << count) - 1
528}
529
530/// Whether a field is shaped like a quoted field, and if it is, whether it has an escape in it.
531///
532/// Shaped means a quote first and last and every quote between them one of a doubled pair, which is
533/// exactly the field the byte loop would read as quoted and close on the last byte. Pairs are taken
534/// from the left the way the byte loop takes them, so `"a"""` is `a"` and `"a""` is not a field the
535/// byte loop would have closed where it stops.
536fn enclosed(field: &[u8], quote: u8) -> Option<bool> {
537    let [first, inner @ .., last] = field else { return None };
538    if *first != quote || *last != quote {
539        return None;
540    }
541    let mut escaped = false;
542    let mut at = 0;
543    while let Some(found) = inner[at..].iter().position(|&byte| byte == quote) {
544        if inner.get(at + found + 1) != Some(&quote) {
545            return None;
546        }
547        escaped = true;
548        at += found + 2;
549    }
550    Some(escaped)
551}
552
553/// [`record`] writing ranges rather than strings, which is the byte loop [`records`] falls back on.
554///
555/// The same loop step for step, with the pushes of bytes taken out and a range recorded in their
556/// place. On `None` or an error nothing this call added is left in `out`.
557fn spans(
558    bytes: &[u8],
559    from: usize,
560    dialect: Dialect,
561    eof: bool,
562    out: &mut Vec<Span>,
563) -> Result<Option<usize>> {
564    if from >= bytes.len() {
565        return Ok(None);
566    }
567    let mark = out.len();
568    let quote = dialect.quote_byte();
569    let escape = dialect.escape_byte();
570    let mut at = from;
571    loop {
572        let span;
573        if bytes.get(at) == Some(&quote) {
574            at += 1;
575            let start = at;
576            let mut escaped = false;
577            loop {
578                let Some(&byte) = bytes.get(at) else {
579                    out.truncate(mark);
580                    return Ok(None);
581                };
582                if byte == escape && escape != quote {
583                    if at + 1 >= bytes.len() {
584                        out.truncate(mark);
585                        return Ok(None);
586                    }
587                    escaped = true;
588                    at += 2;
589                    continue;
590                }
591                if byte == escape && bytes.get(at + 1) == Some(&quote) {
592                    escaped = true;
593                    at += 2;
594                    continue;
595                }
596                if byte == quote {
597                    break;
598                }
599                at += 1;
600            }
601            span = Span::new(start, at, escaped);
602            at += 1;
603            match bytes.get(at) {
604                None if !eof => {
605                    out.truncate(mark);
606                    return Ok(None);
607                }
608                None => {}
609                Some(&byte) if byte == dialect.delimiter || byte == b'\n' || byte == b'\r' => {}
610                Some(&byte) => {
611                    out.truncate(mark);
612                    return Err(after_quote(byte));
613                }
614            }
615        } else {
616            let start = at;
617            while let Some(&byte) = bytes.get(at) {
618                if byte == dialect.delimiter || byte == b'\n' || byte == b'\r' {
619                    break;
620                }
621                at += 1;
622            }
623            if at >= bytes.len() && !eof {
624                out.truncate(mark);
625                return Ok(None);
626            }
627            span = Span::new(start, at, false);
628        }
629        out.push(span);
630        match bytes.get(at) {
631            Some(&byte) if byte == dialect.delimiter => at += 1,
632            Some(b'\r') => {
633                at += 1;
634                if bytes.get(at) == Some(&b'\n') {
635                    at += 1;
636                } else if at >= bytes.len() && !eof {
637                    out.truncate(mark);
638                    return Ok(None);
639                }
640                break;
641            }
642            Some(b'\n') => {
643                at += 1;
644                break;
645            }
646            Some(_) => unreachable!("a field stops at a delimiter, a line ending or the end"),
647            None => break,
648        }
649    }
650    Ok(Some(at))
651}
652
653#[cfg(test)]
654mod tests {
655    use super::*;
656
657    fn split(bytes: &[u8], dialect: Dialect) -> Vec<Vec<String>> {
658        let mut rows = Vec::new();
659        let mut fields = Vec::new();
660        let mut at = 0;
661        while at < bytes.len() {
662            let next = record(bytes, at, dialect, true, &mut fields)
663                .expect("splits")
664                .expect("a whole record");
665            rows.push(fields.clone());
666            at = next;
667        }
668        rows
669    }
670
671    fn comma() -> Dialect {
672        Dialect { delimiter: b',', quote: Some(b'"'), escape: Some(b'"'), header: true }
673    }
674
675    #[test]
676    fn a_line_of_fields_is_the_fields_of_that_line() {
677        assert_eq!(split(b"a,b,c\n1,2,3\n", comma()), [["a", "b", "c"], ["1", "2", "3"]]);
678    }
679
680    #[test]
681    fn the_last_line_does_not_need_a_newline_on_it() {
682        assert_eq!(split(b"a,b\n1,2", comma()), [["a", "b"], ["1", "2"]]);
683    }
684
685    #[test]
686    fn all_three_line_endings_end_a_line() {
687        assert_eq!(split(b"a\r\nb\rc\n", comma()), [["a"], ["b"], ["c"]]);
688    }
689
690    #[test]
691    fn a_quoted_field_may_hold_the_delimiter_and_a_newline() {
692        assert_eq!(split(b"1,\"x,y\"\n", comma()), [["1", "x,y"]]);
693        assert_eq!(split(b"1,\"x\ny\"\n", comma()), [["1", "x\ny"]]);
694    }
695
696    #[test]
697    fn a_doubled_quote_inside_a_quoted_field_is_one_quote() {
698        assert_eq!(split(b"1,\"say \"\"hi\"\"\"\n", comma()), [["1", "say \"hi\""]]);
699    }
700
701    #[test]
702    fn an_empty_field_is_an_empty_string_here_and_becomes_a_null_above() {
703        assert_eq!(split(b"1,,3\n", comma()), [["1", "", "3"]]);
704        assert_eq!(split(b"1,\"\",3\n", comma()), [["1", "", "3"]]);
705    }
706
707    #[test]
708    fn a_trailing_delimiter_makes_a_last_empty_field() {
709        assert_eq!(split(b"1|x|\n", Dialect { delimiter: b'|', ..comma() }), [["1", "x", ""]]);
710    }
711
712    #[test]
713    fn a_quote_in_the_middle_of_a_bare_field_is_just_a_character() {
714        assert_eq!(split(b"1,he said \"hi\"\n", comma()), [["1", "he said \"hi\""]]);
715    }
716
717    #[test]
718    fn a_record_that_the_buffer_does_not_hold_all_of_is_not_a_record_yet() {
719        let mut fields = Vec::new();
720        assert_eq!(record(b"a,b", 0, comma(), false, &mut fields).unwrap(), None);
721        assert_eq!(record(b"a,\"b", 0, comma(), true, &mut fields).unwrap(), None);
722        assert_eq!(record(b"a,b\n", 0, comma(), false, &mut fields).unwrap(), Some(4));
723    }
724
725    #[test]
726    fn rubbish_after_a_closing_quote_is_an_error_rather_than_a_guess() {
727        let mut fields = Vec::new();
728        let error = record(b"\"x\"y,2\n", 0, comma(), true, &mut fields).unwrap_err();
729        assert!(error.message().contains("quoted value"), "{error}");
730    }
731
732    #[test]
733    fn utf8_survives_being_read_one_byte_at_a_time() {
734        assert_eq!(split("a,héllo\n".as_bytes(), comma()), [["a", "héllo"]]);
735    }
736
737    #[test]
738    fn a_backslash_escape_inside_quotes_keeps_whatever_byte_follows_it() {
739        // The Join Order Benchmark's files, which say `ESCAPE '\'`. DuckDB `v2.0.0-dev84237` reads
740        // `\"` as a quote, `\\` as one backslash and `\n` as the letter n, and a backslash
741        // outside quotes as a backslash.
742        let backslash = Dialect { escape: Some(b'\\'), ..comma() };
743        assert_eq!(
744            split(b"1,\"O\\\"Brien, Pat\"\n2,\"a\\\\b\",\"x\\ny\"\n3,c\\d\n", backslash),
745            vec![vec!["1", "O\"Brien, Pat"], vec!["2", "a\\b", "xny"], vec!["3", "c\\d"]]
746        );
747        // A doubled quote is not an escaped quote any more once the escape is something else.
748        let mut fields = Vec::new();
749        assert!(record(b"\"a\"\"b\"\n", 0, backslash, true, &mut fields).is_err());
750    }
751}
752
753/// The two readers held to the same answers.
754///
755/// [`record`] is the reference, because it is the loop the sniffer runs and the one every rule in
756/// this file was first written down in. Everything [`records`] reports is compared with it: the
757/// fields of every record, where the reading stopped, and the error if there was one.
758#[cfg(test)]
759mod agree {
760    use super::*;
761
762    /// What a whole scan found: the records, where it stopped, and the error it stopped on.
763    type Outcome = (Vec<Vec<String>>, usize, Option<String>);
764
765    fn by_record(bytes: &[u8], from: usize, dialect: Dialect, eof: bool) -> Outcome {
766        let mut rows = Vec::new();
767        let mut fields = Vec::new();
768        let mut at = from;
769        loop {
770            match record(bytes, at, dialect, eof, &mut fields) {
771                Ok(Some(next)) => {
772                    rows.push(fields.clone());
773                    at = next;
774                }
775                Ok(None) => return (rows, at, None),
776                Err(error) => return (rows, at, Some(error.to_string())),
777            }
778        }
779    }
780
781    fn by_chunk(bytes: &[u8], from: usize, dialect: Dialect, eof: bool, limit: usize) -> Outcome {
782        let mut rows = Vec::new();
783        let mut at = from;
784        let mut out = Records::default();
785        loop {
786            out.clear();
787            let result = records(bytes, at, dialect, eof, limit, &mut out);
788            for row in 0..out.len() {
789                rows.push(
790                    out.fields(row)
791                        .iter()
792                        .map(|span| span.text(bytes, dialect).into_owned())
793                        .collect(),
794                );
795            }
796            match result {
797                Ok(next) => {
798                    at = next;
799                    if out.len() < limit {
800                        return (rows, at, None);
801                    }
802                }
803                Err(error) => {
804                    // The error leaves the records before it in `out`, so the position is the
805                    // start of the record it was found in, which is where `record` stopped too.
806                    let mut from = at;
807                    let mut fields = Vec::new();
808                    for _ in 0..out.len() {
809                        from = record(bytes, from, dialect, eof, &mut fields).unwrap().unwrap();
810                    }
811                    return (rows, from, Some(error.to_string()));
812                }
813            }
814        }
815    }
816
817    /// Both readers over `bytes`, whole and cut at every one of `cuts` the way a refill cuts it,
818    /// with a few chunk sizes.
819    fn check(bytes: &[u8], dialect: Dialect, cuts: &[usize]) {
820        for eof in [true, false] {
821            let expected = by_record(bytes, 0, dialect, eof);
822            for limit in [1, 2, 3, 7, 8192] {
823                assert_eq!(
824                    by_chunk(bytes, 0, dialect, eof, limit),
825                    expected,
826                    "{:?} under {dialect:?}, eof {eof}, {limit} at a time",
827                    String::from_utf8_lossy(bytes),
828                );
829            }
830        }
831        let whole = by_record(bytes, 0, dialect, true);
832        for &cut in cuts {
833            let cut = cut.min(bytes.len());
834            // The front of the file as a buffer that more bytes will follow, and then the rest of
835            // it from wherever that stopped, which is what a refill does.
836            let (mut rows, at, error) = by_chunk(&bytes[..cut], 0, dialect, false, 8192);
837            if error.is_some() {
838                assert_eq!(error, whole.2, "an error in the front is the error in the whole");
839                continue;
840            }
841            let (rest, end, error) = by_chunk(bytes, at, dialect, true, 8192);
842            rows.extend(rest);
843            assert_eq!(
844                (rows, end, error),
845                whole,
846                "{:?} under {dialect:?} cut at {cut}",
847                String::from_utf8_lossy(bytes),
848            );
849        }
850    }
851
852    fn dialects() -> [Dialect; 5] {
853        let comma =
854            Dialect { delimiter: b',', quote: Some(b'"'), escape: Some(b'"'), header: false };
855        [
856            comma,
857            Dialect { quote: None, escape: None, ..comma },
858            Dialect { delimiter: b'|', ..comma },
859            Dialect { delimiter: b'\t', quote: Some(b'\''), escape: Some(b'\''), ..comma },
860            Dialect { escape: Some(b'\\'), ..comma },
861        ]
862    }
863
864    #[test]
865    fn the_tricky_ones_split_the_same_both_ways() {
866        let long = "x".repeat(61);
867        let cases: Vec<Vec<u8>> = vec![
868            b"".to_vec(),
869            b"\n".to_vec(),
870            b"\n\n".to_vec(),
871            b"a".to_vec(),
872            b"a,b,c\n1,2,3\n".to_vec(),
873            b"a,b\n1,2".to_vec(),
874            b"a\r\nb\rc\n".to_vec(),
875            b"a\r".to_vec(),
876            b"a\r\r\n\n".to_vec(),
877            b"1,\"x,y\"\n2,\"x\ny\"\n".to_vec(),
878            b"1,\"say \"\"hi\"\"\"\n".to_vec(),
879            b"1,,3\n1,\"\",3\n".to_vec(),
880            b"1|x|\n1,x,\n".to_vec(),
881            b"1,he said \"hi\"\n2,x\n".to_vec(),
882            b"1,he said \"hi\n2,x\n".to_vec(),
883            b"\"x\"y,2\n".to_vec(),
884            b"\"x\"\"\n".to_vec(),
885            b"\"x\"\"".to_vec(),
886            b"a,\"b".to_vec(),
887            b"a,\"b\"".to_vec(),
888            b"\"".to_vec(),
889            b"\"\"\"\"\n".to_vec(),
890            b"\"a\"\"\",b\n".to_vec(),
891            b"\"a\\\"b\",c\n\"a\\\\\"\n".to_vec(),
892            b"\"a\\\\b\",\"x\\ny\",\"\\\n\"\n\"a\\".to_vec(),
893            b"'a,b'\t'c''d'\n".to_vec(),
894            b"x\xffy,\"\xfe\"\"\"\n".to_vec(),
895            "h\u{e9}llo,w\u{f6}rld\n".as_bytes().to_vec(),
896            format!("{long},\"a\nb\",c\n{long}\r\n\"{long}\"\"{long}\",d\n").into_bytes(),
897            format!("{long}ab\r\n{long}abc\r\n").into_bytes(),
898            format!("\"{long}\"\"\",\"\n\"\n{long},x\"\n").into_bytes(),
899        ];
900        for bytes in &cases {
901            let cuts: Vec<usize> = (0..=bytes.len()).collect();
902            for dialect in dialects() {
903                check(bytes, dialect, &cuts);
904            }
905        }
906    }
907
908    /// A small generator, since the crate has no dependencies to take one from.
909    struct Rng(u64);
910
911    impl Rng {
912        fn next(&mut self) -> u64 {
913            self.0 ^= self.0 << 13;
914            self.0 ^= self.0 >> 7;
915            self.0 ^= self.0 << 17;
916            self.0
917        }
918
919        fn below(&mut self, n: usize) -> usize {
920            (self.next() % n as u64) as usize
921        }
922    }
923
924    /// A file that is mostly a CSV file, with the kinds of damage real ones have.
925    fn generate(rng: &mut Rng) -> Vec<u8> {
926        const PIECES: [&[u8]; 20] = [
927            b"a",
928            b"1",
929            b"-2.5",
930            b"xyz",
931            b",",
932            b"|",
933            b"\t",
934            b"\"",
935            b"\"\"",
936            b"'",
937            b"\\",
938            b"\n",
939            b"\r",
940            b"\r\n",
941            b" ",
942            b"\xc3\xa9",
943            b"\xff",
944            b"2020-01-02",
945            b"",
946            b"0123456789abcdef",
947        ];
948        let mut out = Vec::new();
949        let rows = rng.below(40);
950        for _ in 0..rows {
951            let fields = 1 + rng.below(5);
952            for field in 0..fields {
953                if field > 0 {
954                    out.push(b",,,|\t"[rng.below(5)]);
955                }
956                let mut body = Vec::new();
957                for _ in 0..rng.below(6) {
958                    let piece = PIECES[rng.below(PIECES.len())];
959                    // Long runs of plain bytes now and then, so that fields straddle the blocks.
960                    if rng.below(10) == 0 {
961                        body.extend(std::iter::repeat_n(b'q', rng.below(90)));
962                    }
963                    body.extend_from_slice(piece);
964                }
965                match rng.below(4) {
966                    0 => {
967                        out.push(b'"');
968                        for &byte in &body {
969                            if byte == b'"' {
970                                out.push(b'"');
971                            }
972                            out.push(byte);
973                        }
974                        out.push(b'"');
975                    }
976                    1 => {
977                        out.push(b'"');
978                        out.extend_from_slice(&body);
979                        out.push(b'"');
980                    }
981                    _ => out.extend_from_slice(&body),
982                }
983            }
984            out.extend_from_slice([&b"\n"[..], b"\r\n", b"\r"][rng.below(3)]);
985        }
986        if rng.below(3) == 0 {
987            out.truncate(out.len().saturating_sub(1 + rng.below(3)));
988        }
989        out
990    }
991
992    #[test]
993    fn thousands_of_generated_files_split_the_same_both_ways() {
994        let mut rng = Rng(0x9e37_79b9_7f4a_7c15);
995        for _ in 0..3000 {
996            let bytes = generate(&mut rng);
997            let cuts: Vec<usize> = (0..4).map(|_| rng.below(bytes.len() + 1)).collect();
998            for dialect in dialects() {
999                check(&bytes, dialect, &cuts);
1000            }
1001        }
1002    }
1003}