Skip to main content

polars_io/csv/read/
parser.rs

1use std::cmp;
2
3use memchr::memchr2_iter;
4use polars_buffer::Buffer;
5use polars_core::prelude::*;
6use polars_core::runtime::RAYON;
7use polars_error::feature_gated;
8use polars_utils::mmap::MMapSemaphore;
9use polars_utils::pl_path::PlRefPath;
10use polars_utils::select::select_unpredictable;
11use rayon::prelude::*;
12
13use super::CsvParseOptions;
14use super::builder::Builder;
15use super::options::{CommentPrefix, NullValuesCompiled};
16use super::splitfields::SplitFields;
17use crate::prelude::CsvReadOptions;
18use crate::prelude::streaming::read_until_start_and_infer_schema;
19use crate::utils::compression::ByteSourceReader;
20use crate::utils::stream_buf_reader::ReaderSource;
21
22/// Read the number of rows without parsing columns
23/// useful for count(*) queries
24#[allow(clippy::too_many_arguments)]
25pub fn count_rows(
26    path: PlRefPath,
27    quote_char: Option<u8>,
28    comment_prefix: Option<&CommentPrefix>,
29    eol_char: u8,
30    has_header: bool,
31    skip_lines: usize,
32    skip_rows_before_header: usize,
33    skip_rows_after_header: usize,
34    raise_if_empty: bool,
35) -> PolarsResult<usize> {
36    let file = if path.has_scheme() || polars_config::config().force_async() {
37        feature_gated!("cloud", {
38            crate::file_cache::FILE_CACHE
39                .get_entry(path)
40                // Safety: This was initialized by schema inference.
41                .unwrap()
42                .try_open_assume_latest()?
43        })
44    } else {
45        polars_utils::io::open_file(path.as_std_path())?
46    };
47
48    let mmap = MMapSemaphore::new_from_file(&file).unwrap();
49
50    count_rows_from_slice_par(
51        Buffer::from_owner(mmap),
52        quote_char,
53        comment_prefix,
54        eol_char,
55        has_header,
56        skip_lines,
57        skip_rows_before_header,
58        skip_rows_after_header,
59        raise_if_empty,
60    )
61}
62
63/// Read the number of rows without parsing columns.
64/// Useful for count(*) queries.
65/// Supports transparent decompression. Does not support truncated compressed files.
66#[allow(clippy::too_many_arguments)]
67pub fn count_rows_from_reader_par(
68    mut reader: ByteSourceReader<ReaderSource>,
69    quote_char: Option<u8>,
70    comment_prefix: Option<&CommentPrefix>,
71    eol_char: u8,
72    has_header: bool,
73    skip_lines: usize,
74    skip_rows_before_header: usize,
75    skip_rows_after_header: usize,
76    raise_if_empty: bool,
77    decompressed_size_hint: Option<usize>,
78) -> PolarsResult<usize> {
79    let reader_options = CsvReadOptions {
80        parse_options: Arc::new(CsvParseOptions {
81            quote_char,
82            comment_prefix: comment_prefix.cloned(),
83            eol_char,
84            ..Default::default()
85        }),
86        has_header,
87        skip_lines,
88        skip_rows: skip_rows_before_header,
89        skip_rows_after_header,
90        raise_if_empty,
91        ..Default::default()
92    };
93
94    let (_, mut leftover) = read_until_start_and_infer_schema(
95        &reader_options,
96        None,
97        true,
98        true,
99        decompressed_size_hint,
100        None,
101        &mut reader,
102    )?;
103
104    const BYTES_PER_CHUNK: usize = if cfg!(debug_assertions) {
105        128
106    } else {
107        512 * 1024
108    };
109
110    let count = CountLines::new(quote_char, eol_char, comment_prefix.cloned());
111    RAYON.install(|| {
112        let mut states = Vec::new();
113
114        let eof_unterminated_row = if comment_prefix.is_none() {
115            let mut last_slice = Buffer::new();
116            let mut err = None;
117
118            let streaming_iter = std::iter::from_fn(|| {
119                let (slice, read_n) =
120                    match reader.read_next_slice(&leftover, BYTES_PER_CHUNK, Some(BYTES_PER_CHUNK))
121                    {
122                        Ok(tup) => tup,
123                        Err(e) => {
124                            err = Some(e);
125                            return None;
126                        },
127                    };
128
129                leftover = Buffer::new();
130                if slice.is_empty() && read_n == 0 {
131                    return None;
132                }
133
134                last_slice = slice.clone();
135                Some(slice)
136            });
137
138            states = streaming_iter
139                .enumerate()
140                .par_bridge()
141                .map(|(id, slice)| (count.analyze_chunk(&slice), id))
142                .collect::<Vec<_>>();
143
144            if let Some(e) = err {
145                return Err(e.into());
146            }
147
148            // par_bridge does not guarantee order, but is mostly sorted so `slice::sort` is a
149            // decent fit.
150            states.sort_by_key(|(_, id)| *id);
151
152            // Technically this is broken if the input has a comment line at the end that is longer
153            // than `BYTES_PER_CHUNK`, but in practice this ought to be fine.
154            ends_in_unterminated_row(&last_slice, eol_char, comment_prefix)
155        } else {
156            // For the non-compressed case this is a zero-copy op.
157            // TODO: Implement streaming chunk logic.
158            let (bytes, _) =
159                reader.read_next_slice(&leftover, usize::MAX, decompressed_size_hint)?;
160
161            let num_chunks = bytes.len().div_ceil(BYTES_PER_CHUNK);
162            (0..num_chunks)
163                .into_par_iter()
164                .map(|chunk_idx| {
165                    let mut start_offset = chunk_idx * BYTES_PER_CHUNK;
166                    let next_start_offset = (start_offset + BYTES_PER_CHUNK).min(bytes.len());
167
168                    if start_offset != 0 {
169                        // Ensure we start at the start of a line.
170                        if let Some(nl_off) = bytes[start_offset..next_start_offset]
171                            .iter()
172                            .position(|b| *b == eol_char)
173                        {
174                            start_offset += nl_off + 1;
175                        } else {
176                            return (count.analyze_chunk(&[]), 0);
177                        }
178                    }
179
180                    let stop_offset = if let Some(nl_off) = bytes[next_start_offset..]
181                        .iter()
182                        .position(|b| *b == eol_char)
183                    {
184                        next_start_offset + nl_off + 1
185                    } else {
186                        bytes.len()
187                    };
188
189                    (count.analyze_chunk(&bytes[start_offset..stop_offset]), 0)
190                })
191                .collect_into_vec(&mut states);
192
193            ends_in_unterminated_row(&bytes, eol_char, comment_prefix)
194        };
195
196        let mut n = 0;
197        let mut in_string = false;
198        for (pair, _) in states {
199            n += pair[in_string as usize].newline_count;
200            in_string = pair[in_string as usize].end_inside_string;
201        }
202        n += eof_unterminated_row as usize;
203        Ok(n)
204    })
205}
206
207/// Read the number of rows without parsing columns.
208/// Useful for count(*) queries.
209/// Supports transparent decompression.
210#[allow(clippy::too_many_arguments)]
211pub fn count_rows_from_slice_par(
212    buffer: Buffer<u8>,
213    quote_char: Option<u8>,
214    comment_prefix: Option<&CommentPrefix>,
215    eol_char: u8,
216    has_header: bool,
217    skip_lines: usize,
218    skip_rows_before_header: usize,
219    skip_rows_after_header: usize,
220    raise_if_empty: bool,
221) -> PolarsResult<usize> {
222    const ASSUMED_COMPRESSION_RATIO: usize = 4;
223
224    let buffer_len = buffer.len();
225    let reader = ByteSourceReader::from_memory(buffer)?;
226    let decompressed_size_hint = Some(
227        buffer_len
228            * reader
229                .compression()
230                .map_or(1, |_| ASSUMED_COMPRESSION_RATIO),
231    );
232
233    count_rows_from_reader_par(
234        reader,
235        quote_char,
236        comment_prefix,
237        eol_char,
238        has_header,
239        skip_lines,
240        skip_rows_before_header,
241        skip_rows_after_header,
242        raise_if_empty,
243        decompressed_size_hint,
244    )
245}
246
247/// Checks if a line in a CSV file is a comment based on the given comment prefix configuration.
248///
249/// This function is used during CSV parsing to determine whether a line should be ignored based on its starting characters.
250#[inline]
251pub fn is_comment_line(line: &[u8], comment_prefix: Option<&CommentPrefix>) -> bool {
252    match comment_prefix {
253        Some(CommentPrefix::Single(c)) => line.first() == Some(c),
254        Some(CommentPrefix::Multi(s)) => line.starts_with(s.as_bytes()),
255        None => false,
256    }
257}
258
259/// Find the nearest next line position.
260/// Does not check for new line characters embedded in String fields.
261pub(super) fn next_line_position_naive(input: &[u8], eol_char: u8) -> Option<usize> {
262    let pos = memchr::memchr(eol_char, input)? + 1;
263    if input.len() - pos == 0 {
264        return None;
265    }
266    Some(pos)
267}
268
269/// Find the nearest next line position that is not embedded in a String field.
270pub(super) fn next_line_position(
271    mut input: &[u8],
272    mut expected_fields: Option<usize>,
273    separator: u8,
274    quote_char: Option<u8>,
275    eol_char: u8,
276) -> Option<usize> {
277    fn accept_line(
278        line: &[u8],
279        expected_fields: usize,
280        separator: u8,
281        eol_char: u8,
282        quote_char: Option<u8>,
283    ) -> bool {
284        let mut count = 0usize;
285        for (field, _) in SplitFields::new(line, separator, quote_char, eol_char) {
286            if memchr2_iter(separator, eol_char, field).count() >= expected_fields {
287                return false;
288            }
289            count += 1;
290        }
291
292        // if the latest field is missing
293        // e.g.:
294        // a,b,c
295        // vala,valb,
296        // SplitFields returns a count that is 1 less
297        // Therefore we accept:
298        // expected == count
299        // and
300        // expected == count - 1
301        expected_fields.wrapping_sub(count) <= 1
302    }
303
304    // we check 3 subsequent lines for `accept_line` before we accept
305    // if 3 groups are rejected we reject completely
306    let mut rejected_line_groups = 0u8;
307
308    let mut total_pos = 0;
309    if input.is_empty() {
310        return None;
311    }
312    let mut lines_checked = 0u8;
313    loop {
314        if rejected_line_groups >= 3 {
315            return None;
316        }
317        lines_checked = lines_checked.wrapping_add(1);
318        // headers might have an extra value
319        // So if we have churned through enough lines
320        // we try one field less.
321        if lines_checked == u8::MAX {
322            if let Some(ef) = expected_fields {
323                expected_fields = Some(ef.saturating_sub(1))
324            }
325        };
326        let pos = memchr::memchr(eol_char, input)? + 1;
327        if input.len() - pos == 0 {
328            return None;
329        }
330        debug_assert!(pos <= input.len());
331        let new_input = unsafe { input.get_unchecked(pos..) };
332        let mut lines = SplitLines::new(new_input, quote_char, eol_char, None);
333        let line = lines.next();
334
335        match (line, expected_fields) {
336            // count the fields, and determine if they are equal to what we expect from the schema
337            (Some(line), Some(expected_fields)) => {
338                if accept_line(line, expected_fields, separator, eol_char, quote_char) {
339                    let mut valid = true;
340                    for line in lines.take(2) {
341                        if !accept_line(line, expected_fields, separator, eol_char, quote_char) {
342                            valid = false;
343                            break;
344                        }
345                    }
346                    if valid {
347                        return Some(total_pos + pos);
348                    } else {
349                        rejected_line_groups += 1;
350                    }
351                } else {
352                    debug_assert!(pos < input.len());
353                    unsafe {
354                        input = input.get_unchecked(pos + 1..);
355                    }
356                    total_pos += pos + 1;
357                }
358            },
359            // don't count the fields
360            (Some(_), None) => return Some(total_pos + pos),
361            // // no new line found, check latest line (without eol) for number of fields
362            _ => return None,
363        }
364    }
365}
366
367#[inline(always)]
368pub(super) fn is_whitespace(b: u8) -> bool {
369    b == b' ' || b == b'\t'
370}
371
372/// May have false-positives, but not false negatives.
373#[inline(always)]
374pub(super) fn could_be_whitespace_fast(b: u8) -> bool {
375    // We're interested in \t (ASCII 9) and " " (ASCII 32), both of which are
376    // <= 32. In that range there aren't a lot of other common symbols (besides
377    // newline), so this is a quick test which can be worth doing to avoid the
378    // exact test.
379    b <= 32
380}
381
382#[inline]
383fn skip_condition<F>(input: &[u8], f: F) -> &[u8]
384where
385    F: Fn(u8) -> bool,
386{
387    if input.is_empty() {
388        return input;
389    }
390
391    let read = input.iter().position(|b| !f(*b)).unwrap_or(input.len());
392    &input[read..]
393}
394
395/// Remove whitespace from the start of buffer.
396/// Makes sure that the bytes stream starts with
397///     'field_1,field_2'
398/// and not with
399///     '\nfield_1,field_1'
400#[inline]
401pub(super) fn skip_whitespace(input: &[u8]) -> &[u8] {
402    skip_condition(input, is_whitespace)
403}
404
405/// An adapted version of std::iter::Split.
406/// This exists solely because we cannot split the file in lines naively as
407///
408/// ```text
409///    for line in bytes.split(b'\n') {
410/// ```
411///
412/// This will fail when string fields have embedded end line characters.
413/// For instance: "This is a valid field\nI have multiples lines" is a valid string field, that contains multiple lines.
414pub struct SplitLines<'a> {
415    v: &'a [u8],
416    quote_char: u8,
417    eol_char: u8,
418    #[cfg(feature = "simd")]
419    simd_eol_char: SimdVec,
420    #[cfg(feature = "simd")]
421    simd_quote_char: SimdVec,
422    #[cfg(feature = "simd")]
423    previous_valid_eols: u64,
424    total_index: usize,
425    quoting: bool,
426    comment_prefix: Option<&'a CommentPrefix>,
427}
428
429#[cfg(feature = "simd")]
430const SIMD_SIZE: usize = 64;
431#[cfg(feature = "simd")]
432use std::simd::prelude::*;
433
434#[cfg(feature = "simd")]
435use polars_utils::clmul::prefix_xorsum_inclusive;
436
437#[cfg(feature = "simd")]
438type SimdVec = u8x64;
439
440impl<'a> SplitLines<'a> {
441    pub fn new(
442        slice: &'a [u8],
443        quote_char: Option<u8>,
444        eol_char: u8,
445        comment_prefix: Option<&'a CommentPrefix>,
446    ) -> Self {
447        let quoting = quote_char.is_some();
448        let quote_char = quote_char.unwrap_or(b'\"');
449        #[cfg(feature = "simd")]
450        let simd_eol_char = SimdVec::splat(eol_char);
451        #[cfg(feature = "simd")]
452        let simd_quote_char = SimdVec::splat(quote_char);
453        Self {
454            v: slice,
455            quote_char,
456            eol_char,
457            #[cfg(feature = "simd")]
458            simd_eol_char,
459            #[cfg(feature = "simd")]
460            simd_quote_char,
461            #[cfg(feature = "simd")]
462            previous_valid_eols: 0,
463            total_index: 0,
464            quoting,
465            comment_prefix,
466        }
467    }
468}
469
470impl<'a> SplitLines<'a> {
471    // scalar as in non-simd
472    fn next_scalar(&mut self) -> Option<&'a [u8]> {
473        if self.v.is_empty() {
474            return None;
475        }
476        if is_comment_line(self.v, self.comment_prefix) {
477            return self.next_comment_line();
478        }
479        {
480            let mut pos = 0u32;
481            let mut iter = self.v.iter();
482            let mut in_field = false;
483            loop {
484                match iter.next() {
485                    Some(&c) => {
486                        pos += 1;
487
488                        if self.quoting && c == self.quote_char {
489                            // toggle between string field enclosure
490                            //      if we encounter a starting '"' -> in_field = true;
491                            //      if we encounter a closing '"' -> in_field = false;
492                            in_field = !in_field;
493                        }
494                        // if we are not in a string and we encounter '\n' we can stop at this position.
495                        else if c == self.eol_char && !in_field {
496                            break;
497                        }
498                    },
499                    None => {
500                        let remainder = self.v;
501                        self.v = &[];
502                        return Some(remainder);
503                    },
504                }
505            }
506
507            unsafe {
508                debug_assert!((pos as usize) <= self.v.len());
509
510                // return line up to this position
511                let ret = Some(
512                    self.v
513                        .get_unchecked(..(self.total_index + pos as usize - 1)),
514                );
515                // skip the '\n' token and update slice.
516                self.v = self.v.get_unchecked(self.total_index + pos as usize..);
517                ret
518            }
519        }
520    }
521    fn next_comment_line(&mut self) -> Option<&'a [u8]> {
522        if let Some(pos) = next_line_position_naive(self.v, self.eol_char) {
523            unsafe {
524                // return line up to this position
525                let ret = Some(self.v.get_unchecked(..(pos - 1)));
526                // skip the '\n' token and update slice.
527                self.v = self.v.get_unchecked(pos..);
528                ret
529            }
530        } else {
531            let remainder = self.v;
532            self.v = &[];
533            Some(remainder)
534        }
535    }
536}
537
538impl<'a> Iterator for SplitLines<'a> {
539    type Item = &'a [u8];
540
541    #[inline]
542    #[cfg(not(feature = "simd"))]
543    fn next(&mut self) -> Option<&'a [u8]> {
544        self.next_scalar()
545    }
546
547    #[inline]
548    #[cfg(feature = "simd")]
549    fn next(&mut self) -> Option<&'a [u8]> {
550        // First check cached value
551        if self.previous_valid_eols != 0 {
552            let pos = self.previous_valid_eols.trailing_zeros() as usize;
553            self.previous_valid_eols >>= (pos + 1) as u64;
554
555            unsafe {
556                debug_assert!((pos) <= self.v.len());
557
558                // return line up to this position
559                let ret = Some(self.v.get_unchecked(..pos));
560                // skip the '\n' token and update slice.
561                self.v = self.v.get_unchecked(pos + 1..);
562                return ret;
563            }
564        }
565        if self.v.is_empty() {
566            return None;
567        }
568        if self.comment_prefix.is_some() {
569            return self.next_scalar();
570        }
571
572        self.total_index = 0;
573        let mut not_in_field_previous_iter = true;
574
575        loop {
576            let bytes = unsafe { self.v.get_unchecked(self.total_index..) };
577            if bytes.len() > SIMD_SIZE {
578                let lane: [u8; SIMD_SIZE] = unsafe {
579                    bytes
580                        .get_unchecked(0..SIMD_SIZE)
581                        .try_into()
582                        .unwrap_unchecked()
583                };
584                let simd_bytes = SimdVec::from(lane);
585                let eol_mask = simd_bytes.simd_eq(self.simd_eol_char).to_bitmask();
586
587                let valid_eols = if self.quoting {
588                    let quote_mask = simd_bytes.simd_eq(self.simd_quote_char).to_bitmask();
589                    let mut not_in_quote_field = prefix_xorsum_inclusive(quote_mask);
590
591                    if not_in_field_previous_iter {
592                        not_in_quote_field = !not_in_quote_field;
593                    }
594                    not_in_field_previous_iter = (not_in_quote_field & (1 << (SIMD_SIZE - 1))) > 0;
595                    eol_mask & not_in_quote_field
596                } else {
597                    eol_mask
598                };
599
600                if valid_eols != 0 {
601                    let pos = valid_eols.trailing_zeros() as usize;
602                    if pos == SIMD_SIZE - 1 {
603                        self.previous_valid_eols = 0;
604                    } else {
605                        self.previous_valid_eols = valid_eols >> (pos + 1) as u64;
606                    }
607
608                    unsafe {
609                        let pos = self.total_index + pos;
610                        debug_assert!((pos) <= self.v.len());
611
612                        // return line up to this position
613                        let ret = Some(self.v.get_unchecked(..pos));
614                        // skip the '\n' token and update slice.
615                        self.v = self.v.get_unchecked(pos + 1..);
616                        return ret;
617                    }
618                } else {
619                    self.total_index += SIMD_SIZE;
620                }
621            } else {
622                // Denotes if we are in a string field, started with a quote
623                let mut in_field = !not_in_field_previous_iter;
624                let mut pos = 0u32;
625                let mut iter = bytes.iter();
626                loop {
627                    match iter.next() {
628                        Some(&c) => {
629                            pos += 1;
630
631                            if self.quoting && c == self.quote_char {
632                                // toggle between string field enclosure
633                                //      if we encounter a starting '"' -> in_field = true;
634                                //      if we encounter a closing '"' -> in_field = false;
635                                in_field = !in_field;
636                            }
637                            // if we are not in a string and we encounter '\n' we can stop at this position.
638                            else if c == self.eol_char && !in_field {
639                                break;
640                            }
641                        },
642                        None => {
643                            let remainder = self.v;
644                            self.v = &[];
645                            return Some(remainder);
646                        },
647                    }
648                }
649
650                unsafe {
651                    debug_assert!((pos as usize) <= self.v.len());
652
653                    // return line up to this position
654                    let ret = Some(
655                        self.v
656                            .get_unchecked(..(self.total_index + pos as usize - 1)),
657                    );
658                    // skip the '\n' token and update slice.
659                    self.v = self.v.get_unchecked(self.total_index + pos as usize..);
660                    return ret;
661                }
662            }
663        }
664    }
665}
666
667#[derive(Clone)]
668pub struct CountLines {
669    quote_char: u8,
670    eol_char: u8,
671    #[cfg(feature = "simd")]
672    simd_eol_char: SimdVec,
673    #[cfg(feature = "simd")]
674    simd_quote_char: SimdVec,
675    quoting: bool,
676    comment_prefix: Option<CommentPrefix>,
677}
678
679#[derive(Copy, Clone, Debug, Default)]
680pub struct LineStats {
681    pub newline_count: usize,
682    pub last_newline_offset: usize,
683    pub end_inside_string: bool,
684}
685
686impl CountLines {
687    pub fn new(
688        quote_char: Option<u8>,
689        eol_char: u8,
690        comment_prefix: Option<CommentPrefix>,
691    ) -> Self {
692        let quoting = quote_char.is_some();
693        let quote_char = quote_char.unwrap_or(b'\"');
694        #[cfg(feature = "simd")]
695        let simd_eol_char = SimdVec::splat(eol_char);
696        #[cfg(feature = "simd")]
697        let simd_quote_char = SimdVec::splat(quote_char);
698        Self {
699            quote_char,
700            eol_char,
701            #[cfg(feature = "simd")]
702            simd_eol_char,
703            #[cfg(feature = "simd")]
704            simd_quote_char,
705            quoting,
706            comment_prefix,
707        }
708    }
709
710    pub fn has_comment_prefix(&self) -> bool {
711        self.comment_prefix.is_some()
712    }
713
714    /// Analyzes a chunk of CSV data.
715    ///
716    /// Returns (newline_count, last_newline_offset, end_inside_string) twice,
717    /// the first is assuming the start of the chunk is *not* inside a string,
718    /// the second assuming the start is inside a string.
719    ///
720    /// If comment_prefix is not None the start of bytes must be at the start of
721    /// a line (and thus not in the middle of a comment).
722    pub fn analyze_chunk(&self, bytes: &[u8]) -> [LineStats; 2] {
723        let mut states = [
724            LineStats {
725                newline_count: 0,
726                last_newline_offset: 0,
727                end_inside_string: false,
728            },
729            LineStats {
730                newline_count: 0,
731                last_newline_offset: 0,
732                end_inside_string: false,
733            },
734        ];
735
736        // If we have to deal with comments we can't use SIMD and have to explicitly do two passes.
737        if self.comment_prefix.is_some() {
738            states[0] = self.analyze_chunk_with_comment(bytes, false);
739            states[1] = self.analyze_chunk_with_comment(bytes, true);
740            return states;
741        }
742
743        // False if even number of quotes seen so far, true otherwise.
744        #[allow(unused_assignments)]
745        let mut global_quote_parity = false;
746        let mut scan_offset = 0;
747
748        #[cfg(feature = "simd")]
749        {
750            // 0 if even number of quotes seen so far, u64::MAX otherwise.
751            let mut global_quote_parity_mask = 0;
752            while scan_offset + 64 <= bytes.len() {
753                let block: [u8; 64] = unsafe {
754                    bytes
755                        .get_unchecked(scan_offset..scan_offset + 64)
756                        .try_into()
757                        .unwrap_unchecked()
758                };
759                let simd_bytes = SimdVec::from(block);
760                let eol_mask = simd_bytes.simd_eq(self.simd_eol_char).to_bitmask();
761                if self.quoting {
762                    let quote_mask = simd_bytes.simd_eq(self.simd_quote_char).to_bitmask();
763                    let quote_parity =
764                        prefix_xorsum_inclusive(quote_mask) ^ global_quote_parity_mask;
765                    global_quote_parity_mask = ((quote_parity as i64) >> 63) as u64;
766
767                    let start_outside_string_eol_mask = eol_mask & !quote_parity;
768                    states[0].newline_count += start_outside_string_eol_mask.count_ones() as usize;
769                    states[0].last_newline_offset = select_unpredictable(
770                        start_outside_string_eol_mask != 0,
771                        (scan_offset + 63)
772                            .wrapping_sub(start_outside_string_eol_mask.leading_zeros() as usize),
773                        states[0].last_newline_offset,
774                    );
775
776                    let start_inside_string_eol_mask = eol_mask & quote_parity;
777                    states[1].newline_count += start_inside_string_eol_mask.count_ones() as usize;
778                    states[1].last_newline_offset = select_unpredictable(
779                        start_inside_string_eol_mask != 0,
780                        (scan_offset + 63)
781                            .wrapping_sub(start_inside_string_eol_mask.leading_zeros() as usize),
782                        states[1].last_newline_offset,
783                    );
784                } else {
785                    states[0].newline_count += eol_mask.count_ones() as usize;
786                    states[0].last_newline_offset = select_unpredictable(
787                        eol_mask != 0,
788                        (scan_offset + 63).wrapping_sub(eol_mask.leading_zeros() as usize),
789                        states[0].last_newline_offset,
790                    );
791                }
792
793                scan_offset += 64;
794            }
795
796            global_quote_parity = global_quote_parity_mask > 0;
797        }
798
799        while scan_offset < bytes.len() {
800            let c = unsafe { *bytes.get_unchecked(scan_offset) };
801            global_quote_parity ^= (c == self.quote_char) & self.quoting;
802
803            let state = &mut states[global_quote_parity as usize];
804            state.newline_count += (c == self.eol_char) as usize;
805            state.last_newline_offset =
806                select_unpredictable(c == self.eol_char, scan_offset, state.last_newline_offset);
807
808            scan_offset += 1;
809        }
810
811        states[0].end_inside_string = global_quote_parity;
812        states[1].end_inside_string = !global_quote_parity;
813        states
814    }
815
816    // bytes must begin at the start of a line.
817    fn analyze_chunk_with_comment(&self, bytes: &[u8], mut in_string: bool) -> LineStats {
818        let pre_s = match self.comment_prefix.as_ref().unwrap() {
819            CommentPrefix::Single(pc) => core::slice::from_ref(pc),
820            CommentPrefix::Multi(ps) => ps.as_bytes(),
821        };
822
823        let mut state = LineStats::default();
824        let mut scan_offset = 0;
825        while scan_offset < bytes.len() {
826            // Skip comment line if needed.
827            while bytes[scan_offset..].starts_with(pre_s) {
828                scan_offset += pre_s.len();
829                let Some(nl_off) = bytes[scan_offset..]
830                    .iter()
831                    .position(|c| *c == self.eol_char)
832                else {
833                    break;
834                };
835                scan_offset += nl_off + 1;
836            }
837
838            while scan_offset < bytes.len() {
839                let c = unsafe { *bytes.get_unchecked(scan_offset) };
840                in_string ^= (c == self.quote_char) & self.quoting;
841
842                if c == self.eol_char && !in_string {
843                    state.newline_count += 1;
844                    state.last_newline_offset = scan_offset;
845                    scan_offset += 1;
846                    break;
847                } else {
848                    scan_offset += 1;
849                }
850            }
851        }
852
853        state.end_inside_string = in_string;
854        state
855    }
856
857    pub fn find_next(&self, bytes: &[u8], chunk_size: &mut usize) -> (usize, usize) {
858        loop {
859            let b = unsafe { bytes.get_unchecked(..(*chunk_size).min(bytes.len())) };
860
861            let (count, offset) = if self.comment_prefix.is_some() {
862                let stats = self.analyze_chunk_with_comment(b, false);
863                (stats.newline_count, stats.last_newline_offset)
864            } else {
865                self.count(b)
866            };
867
868            if count > 0 || b.len() == bytes.len() {
869                return (count, offset);
870            }
871
872            *chunk_size = chunk_size.saturating_mul(2);
873        }
874    }
875
876    pub fn count_rows(&self, bytes: &[u8], is_eof: bool) -> (usize, usize) {
877        let stats = if self.comment_prefix.is_some() {
878            self.analyze_chunk_with_comment(bytes, false)
879        } else {
880            self.analyze_chunk(bytes)[0]
881        };
882
883        let mut count = stats.newline_count;
884        let mut offset = stats.last_newline_offset;
885
886        if count > 0 {
887            offset = cmp::min(offset + 1, bytes.len());
888        } else {
889            debug_assert!(offset == 0);
890        }
891
892        if is_eof {
893            count += ends_in_unterminated_row(bytes, self.eol_char, self.comment_prefix.as_ref())
894                as usize;
895            offset = bytes.len();
896        }
897
898        (count, offset)
899    }
900
901    /// Returns count and offset to split for remainder in slice.
902    #[cfg(feature = "simd")]
903    pub fn count(&self, bytes: &[u8]) -> (usize, usize) {
904        let mut total_idx = 0;
905        let original_bytes = bytes;
906        let mut count = 0;
907        let mut position = 0;
908        let mut not_in_field_previous_iter = true;
909
910        loop {
911            let bytes = unsafe { original_bytes.get_unchecked(total_idx..) };
912
913            if bytes.len() > SIMD_SIZE {
914                let lane: [u8; SIMD_SIZE] = unsafe {
915                    bytes
916                        .get_unchecked(0..SIMD_SIZE)
917                        .try_into()
918                        .unwrap_unchecked()
919                };
920                let simd_bytes = SimdVec::from(lane);
921                let eol_mask = simd_bytes.simd_eq(self.simd_eol_char).to_bitmask();
922
923                let valid_eols = if self.quoting {
924                    let quote_mask = simd_bytes.simd_eq(self.simd_quote_char).to_bitmask();
925                    let mut not_in_quote_field = prefix_xorsum_inclusive(quote_mask);
926
927                    if not_in_field_previous_iter {
928                        not_in_quote_field = !not_in_quote_field;
929                    }
930                    not_in_field_previous_iter = (not_in_quote_field & (1 << (SIMD_SIZE - 1))) > 0;
931                    eol_mask & not_in_quote_field
932                } else {
933                    eol_mask
934                };
935
936                if valid_eols != 0 {
937                    count += valid_eols.count_ones() as usize;
938                    position = total_idx + 63 - valid_eols.leading_zeros() as usize;
939                    debug_assert_eq!(original_bytes[position], self.eol_char)
940                }
941                total_idx += SIMD_SIZE;
942            } else if bytes.is_empty() {
943                debug_assert!(count == 0 || original_bytes[position] == self.eol_char);
944                return (count, position);
945            } else {
946                let (c, o) = self.count_no_simd(bytes, !not_in_field_previous_iter);
947
948                let (count, position) = if c > 0 {
949                    (count + c, total_idx + o)
950                } else {
951                    (count, position)
952                };
953                debug_assert!(count == 0 || original_bytes[position] == self.eol_char);
954
955                return (count, position);
956            }
957        }
958    }
959
960    #[cfg(not(feature = "simd"))]
961    pub fn count(&self, bytes: &[u8]) -> (usize, usize) {
962        self.count_no_simd(bytes, false)
963    }
964
965    fn count_no_simd(&self, bytes: &[u8], in_field: bool) -> (usize, usize) {
966        let iter = bytes.iter();
967        let mut in_field = in_field;
968        let mut count = 0;
969        let mut position = 0;
970
971        for b in iter {
972            let c = *b;
973            if self.quoting && c == self.quote_char {
974                // toggle between string field enclosure
975                //      if we encounter a starting '"' -> in_field = true;
976                //      if we encounter a closing '"' -> in_field = false;
977                in_field = !in_field;
978            }
979            // If we are not in a string and we encounter '\n' we can stop at this position.
980            else if c == self.eol_char && !in_field {
981                position = (b as *const _ as usize) - (bytes.as_ptr() as usize);
982                count += 1;
983            }
984        }
985        debug_assert!(count == 0 || bytes[position] == self.eol_char);
986
987        (count, position)
988    }
989}
990
991/// Returns the length of the longest prefix of `bytes` that only holds complete rows, i.e. up to
992/// and including the last `eol_char` outside of a quoted field. Like `skip_lines` itself, the
993/// first `skip_lines` lines are split without regard to quoting. Like the CSV reader, a UTF-8 BOM
994/// after them is skipped, so that a comment prefix on the next line is found.
995pub fn complete_rows_len(
996    bytes: &[u8],
997    quote_char: Option<u8>,
998    comment_prefix: Option<&CommentPrefix>,
999    eol_char: u8,
1000    skip_lines: usize,
1001) -> usize {
1002    let mut start = 0;
1003    for _ in 0..skip_lines {
1004        match memchr::memchr(eol_char, &bytes[start..]) {
1005            Some(pos) => start += pos + 1,
1006            None => return start,
1007        }
1008    }
1009
1010    const UTF8_BOM_MARKER: &[u8] = b"\xef\xbb\xbf";
1011    if bytes[start..].starts_with(UTF8_BOM_MARKER) {
1012        start += UTF8_BOM_MARKER.len();
1013    }
1014
1015    let count = CountLines::new(quote_char, eol_char, comment_prefix.cloned());
1016    let (_, offset) = count.count_rows(&bytes[start..], false);
1017    start + offset
1018}
1019
1020fn ends_in_unterminated_row(
1021    bytes: &[u8],
1022    eol_char: u8,
1023    comment_prefix: Option<&CommentPrefix>,
1024) -> bool {
1025    if !bytes.is_empty() && bytes.last().copied().unwrap() != eol_char {
1026        // We can do a simple backwards-scan to find the start of last line if it is a
1027        // comment line, since comment lines can't escape new-lines.
1028        let last_new_line_pos = memchr::memrchr(eol_char, bytes).unwrap_or(0);
1029        let last_line_is_comment_line = bytes
1030            .get(last_new_line_pos + 1..)
1031            .map(|line| is_comment_line(line, comment_prefix))
1032            .unwrap_or(false);
1033
1034        return !last_line_is_comment_line;
1035    }
1036
1037    false
1038}
1039
1040#[inline]
1041fn find_quoted(bytes: &[u8], quote_char: u8, needle: u8) -> Option<usize> {
1042    let mut in_field = false;
1043
1044    let mut idx = 0u32;
1045    // micro optimizations
1046    #[allow(clippy::explicit_counter_loop)]
1047    for &c in bytes.iter() {
1048        if c == quote_char {
1049            // toggle between string field enclosure
1050            //      if we encounter a starting '"' -> in_field = true;
1051            //      if we encounter a closing '"' -> in_field = false;
1052            in_field = !in_field;
1053        }
1054
1055        if !in_field && c == needle {
1056            return Some(idx as usize);
1057        }
1058        idx += 1;
1059    }
1060    None
1061}
1062
1063#[inline]
1064pub(super) fn skip_this_line(bytes: &[u8], quote: Option<u8>, eol_char: u8) -> &[u8] {
1065    let pos = match quote {
1066        Some(quote) => find_quoted(bytes, quote, eol_char),
1067        None => bytes.iter().position(|x| *x == eol_char),
1068    };
1069    match pos {
1070        None => &[],
1071        Some(pos) => &bytes[pos + 1..],
1072    }
1073}
1074
1075#[inline]
1076pub(super) fn skip_this_line_naive(input: &[u8], eol_char: u8) -> &[u8] {
1077    if let Some(pos) = next_line_position_naive(input, eol_char) {
1078        unsafe { input.get_unchecked(pos..) }
1079    } else {
1080        &[]
1081    }
1082}
1083
1084/// Parse CSV.
1085///
1086/// # Arguments
1087/// * `bytes` - input to parse
1088/// * `offset` - offset in bytes in total input. This is 0 if single threaded. If multi-threaded every
1089///   thread has a different offset.
1090/// * `projection` - Indices of the columns to project.
1091/// * `buffers` - Parsed output will be written to these buffers. Except for UTF8 data. The offsets of the
1092///   fields are written to the buffers. The UTF8 data will be parsed later.
1093///
1094/// Returns the number of bytes parsed successfully.
1095#[allow(clippy::too_many_arguments)]
1096pub(super) fn parse_lines(
1097    mut bytes: &[u8],
1098    parse_options: &CsvParseOptions,
1099    offset: usize,
1100    ignore_errors: bool,
1101    null_values: Option<&NullValuesCompiled>,
1102    projection: &[usize],
1103    buffers: &mut [Builder],
1104    n_lines: usize,
1105    // length of original schema
1106    schema_len: usize,
1107    schema: &Schema,
1108) -> PolarsResult<usize> {
1109    assert!(
1110        !projection.is_empty(),
1111        "at least one column should be projected"
1112    );
1113    let mut truncate_ragged_lines = parse_options.truncate_ragged_lines;
1114    // During projection pushdown we are not checking other csv fields.
1115    // This would be very expensive and we don't care as we only want
1116    // the projected columns.
1117    if projection.len() != schema_len {
1118        truncate_ragged_lines = true
1119    }
1120
1121    // we use the pointers to track the no of bytes read.
1122    let start = bytes.as_ptr() as usize;
1123    let original_bytes_len = bytes.len();
1124    let n_lines = n_lines as u32;
1125
1126    let mut iter = SplitFields::new(
1127        bytes,
1128        parse_options.separator,
1129        parse_options.quote_char,
1130        parse_options.eol_char,
1131    );
1132    let mut line_count = 0u32;
1133    loop {
1134        if line_count > n_lines {
1135            let end = bytes.as_ptr() as usize;
1136            return Ok(end - start);
1137        }
1138
1139        if bytes.is_empty() {
1140            return Ok(original_bytes_len);
1141        } else if is_comment_line(bytes, parse_options.comment_prefix.as_ref()) {
1142            // deal with comments
1143            let bytes_rem = skip_this_line_naive(bytes, parse_options.eol_char);
1144            bytes = bytes_rem;
1145            continue;
1146        }
1147
1148        // Every line we only need to parse the columns that are projected.
1149        // Therefore we check if the idx of the field is in our projected columns.
1150        // If it is not, we skip the field.
1151        let mut projection_iter = projection.iter().copied();
1152        let mut next_projected = unsafe { projection_iter.next().unwrap_unchecked() };
1153        let mut processed_fields = 0;
1154
1155        iter.next_line(bytes);
1156        let mut idx = 0u32;
1157        let mut read_sol = 0;
1158        loop {
1159            match iter.next() {
1160                // end of line
1161                None => {
1162                    bytes = unsafe { bytes.get_unchecked(std::cmp::min(read_sol, bytes.len())..) };
1163                    break;
1164                },
1165                Some((mut field, needs_escaping)) => {
1166                    let field_len = field.len();
1167
1168                    // +1 is the split character that is consumed by the iterator.
1169                    read_sol += field_len + 1;
1170
1171                    if idx == next_projected as u32 {
1172                        // the iterator is finished when it encounters a `\n`
1173                        // this could be preceded by a '\r'
1174                        unsafe {
1175                            if field_len > 0 && *field.get_unchecked(field_len - 1) == b'\r' {
1176                                field = field.get_unchecked(..field_len - 1);
1177                            }
1178                        }
1179
1180                        debug_assert!(processed_fields < buffers.len());
1181                        let buf = unsafe {
1182                            // SAFETY: processed fields index can never exceed the projection indices.
1183                            buffers.get_unchecked_mut(processed_fields)
1184                        };
1185                        let mut add_null = false;
1186
1187                        // if we have null values argument, check if this field equal null value
1188                        if let Some(null_values) = null_values {
1189                            let field = if needs_escaping && !field.is_empty() {
1190                                unsafe { field.get_unchecked(1..field.len() - 1) }
1191                            } else {
1192                                field
1193                            };
1194
1195                            // SAFETY:
1196                            // process fields is in bounds
1197                            add_null = unsafe { null_values.is_null(field, idx as usize) }
1198                        }
1199                        if add_null {
1200                            buf.add_null(!parse_options.missing_is_null && field.is_empty())
1201                        } else {
1202                            buf.add(field, ignore_errors, needs_escaping, parse_options.missing_is_null)
1203                                .map_err(|e| {
1204                                    let bytes_offset = offset + field.as_ptr() as usize - start;
1205                                    let unparsable = String::from_utf8_lossy(field);
1206                                    let column_name = schema.get_at_index(idx as usize).unwrap().0;
1207                                    polars_err!(
1208                                        ComputeError:
1209                                        "could not parse `{}` as dtype `{}` at column '{}' (column number {})\n\n\
1210                                        The current offset in the file is {} bytes.\n\
1211                                        \n\
1212                                        You might want to try:\n\
1213                                        - increasing `infer_schema_length` (e.g. `infer_schema_length=10000`),\n\
1214                                        - increasing `infer_schema_files` (e.g. `infer_schema_files=50`),\n\
1215                                        - specifying correct dtype with the `schema_overrides` argument\n\
1216                                        - setting `ignore_errors` to `True`,\n\
1217                                        - adding `{}` to the `null_values` list.\n\n\
1218                                        Original error: ```{}```",
1219                                        &unparsable,
1220                                        buf.dtype(),
1221                                        column_name,
1222                                        idx + 1,
1223                                        bytes_offset,
1224                                        &unparsable,
1225                                        e
1226                                    )
1227                                })?;
1228                        }
1229                        processed_fields += 1;
1230
1231                        // if we have all projected columns we are done with this line
1232                        match projection_iter.next() {
1233                            Some(p) => next_projected = p,
1234                            None => {
1235                                if bytes.get(read_sol - 1) == Some(&parse_options.eol_char) {
1236                                    bytes = unsafe { bytes.get_unchecked(read_sol..) };
1237                                } else {
1238                                    if !truncate_ragged_lines && read_sol < bytes.len() {
1239                                        polars_bail!(ComputeError: r#"found more fields than defined in 'Schema'
1240
1241Consider setting 'truncate_ragged_lines={}'."#, polars_error::constants::TRUE)
1242                                    }
1243                                    let bytes_rem = skip_this_line(
1244                                        unsafe { bytes.get_unchecked(read_sol - 1..) },
1245                                        parse_options.quote_char,
1246                                        parse_options.eol_char,
1247                                    );
1248                                    bytes = bytes_rem;
1249                                }
1250                                break;
1251                            },
1252                        }
1253                    }
1254                    idx += 1;
1255                },
1256            }
1257        }
1258
1259        // there can be lines that miss fields (also the comma values)
1260        // this means the splitter won't process them.
1261        // We traverse them to read them as null values.
1262        while processed_fields < projection.len() {
1263            debug_assert!(processed_fields < buffers.len());
1264            let buf = unsafe {
1265                // SAFETY: processed fields index can never exceed the projection indices.
1266                buffers.get_unchecked_mut(processed_fields)
1267            };
1268            buf.add_null(!parse_options.missing_is_null);
1269            processed_fields += 1;
1270        }
1271        line_count += 1;
1272    }
1273}
1274
1275#[cfg(test)]
1276mod test {
1277    use super::SplitLines;
1278
1279    #[test]
1280    fn test_splitlines() {
1281        let input = "1,\"foo\n\"\n2,\"foo\n\"\n";
1282        let mut lines = SplitLines::new(input.as_bytes(), Some(b'"'), b'\n', None);
1283        assert_eq!(lines.next(), Some("1,\"foo\n\"".as_bytes()));
1284        assert_eq!(lines.next(), Some("2,\"foo\n\"".as_bytes()));
1285        assert_eq!(lines.next(), None);
1286
1287        let input2 = "1,'foo\n'\n2,'foo\n'\n";
1288        let mut lines2 = SplitLines::new(input2.as_bytes(), Some(b'\''), b'\n', None);
1289        assert_eq!(lines2.next(), Some("1,'foo\n'".as_bytes()));
1290        assert_eq!(lines2.next(), Some("2,'foo\n'".as_bytes()));
1291        assert_eq!(lines2.next(), None);
1292    }
1293}