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#[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 .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#[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 states.sort_by_key(|(_, id)| *id);
151
152 ends_in_unterminated_row(&last_slice, eol_char, comment_prefix)
155 } else {
156 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 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#[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#[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
259pub(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
269pub(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 expected_fields.wrapping_sub(count) <= 1
302 }
303
304 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 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 (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 (Some(_), None) => return Some(total_pos + pos),
361 _ => 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#[inline(always)]
374pub(super) fn could_be_whitespace_fast(b: u8) -> bool {
375 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#[inline]
401pub(super) fn skip_whitespace(input: &[u8]) -> &[u8] {
402 skip_condition(input, is_whitespace)
403}
404
405pub 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 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 in_field = !in_field;
493 }
494 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 let ret = Some(
512 self.v
513 .get_unchecked(..(self.total_index + pos as usize - 1)),
514 );
515 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 let ret = Some(self.v.get_unchecked(..(pos - 1)));
526 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 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 let ret = Some(self.v.get_unchecked(..pos));
560 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 let ret = Some(self.v.get_unchecked(..pos));
614 self.v = self.v.get_unchecked(pos + 1..);
616 return ret;
617 }
618 } else {
619 self.total_index += SIMD_SIZE;
620 }
621 } else {
622 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 in_field = !in_field;
636 }
637 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 let ret = Some(
655 self.v
656 .get_unchecked(..(self.total_index + pos as usize - 1)),
657 );
658 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 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 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 #[allow(unused_assignments)]
745 let mut global_quote_parity = false;
746 let mut scan_offset = 0;
747
748 #[cfg(feature = "simd")]
749 {
750 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 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 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 #[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 in_field = !in_field;
978 }
979 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
991fn ends_in_unterminated_row(
992 bytes: &[u8],
993 eol_char: u8,
994 comment_prefix: Option<&CommentPrefix>,
995) -> bool {
996 if !bytes.is_empty() && bytes.last().copied().unwrap() != eol_char {
997 let last_new_line_post = memchr::memrchr(eol_char, bytes).unwrap_or(0);
1000 let last_line_is_comment_line = bytes
1001 .get(last_new_line_post + 1..)
1002 .map(|line| is_comment_line(line, comment_prefix))
1003 .unwrap_or(false);
1004
1005 return !last_line_is_comment_line;
1006 }
1007
1008 false
1009}
1010
1011#[inline]
1012fn find_quoted(bytes: &[u8], quote_char: u8, needle: u8) -> Option<usize> {
1013 let mut in_field = false;
1014
1015 let mut idx = 0u32;
1016 #[allow(clippy::explicit_counter_loop)]
1018 for &c in bytes.iter() {
1019 if c == quote_char {
1020 in_field = !in_field;
1024 }
1025
1026 if !in_field && c == needle {
1027 return Some(idx as usize);
1028 }
1029 idx += 1;
1030 }
1031 None
1032}
1033
1034#[inline]
1035pub(super) fn skip_this_line(bytes: &[u8], quote: Option<u8>, eol_char: u8) -> &[u8] {
1036 let pos = match quote {
1037 Some(quote) => find_quoted(bytes, quote, eol_char),
1038 None => bytes.iter().position(|x| *x == eol_char),
1039 };
1040 match pos {
1041 None => &[],
1042 Some(pos) => &bytes[pos + 1..],
1043 }
1044}
1045
1046#[inline]
1047pub(super) fn skip_this_line_naive(input: &[u8], eol_char: u8) -> &[u8] {
1048 if let Some(pos) = next_line_position_naive(input, eol_char) {
1049 unsafe { input.get_unchecked(pos..) }
1050 } else {
1051 &[]
1052 }
1053}
1054
1055#[allow(clippy::too_many_arguments)]
1067pub(super) fn parse_lines(
1068 mut bytes: &[u8],
1069 parse_options: &CsvParseOptions,
1070 offset: usize,
1071 ignore_errors: bool,
1072 null_values: Option<&NullValuesCompiled>,
1073 projection: &[usize],
1074 buffers: &mut [Builder],
1075 n_lines: usize,
1076 schema_len: usize,
1078 schema: &Schema,
1079) -> PolarsResult<usize> {
1080 assert!(
1081 !projection.is_empty(),
1082 "at least one column should be projected"
1083 );
1084 let mut truncate_ragged_lines = parse_options.truncate_ragged_lines;
1085 if projection.len() != schema_len {
1089 truncate_ragged_lines = true
1090 }
1091
1092 let start = bytes.as_ptr() as usize;
1094 let original_bytes_len = bytes.len();
1095 let n_lines = n_lines as u32;
1096
1097 let mut line_count = 0u32;
1098 loop {
1099 if line_count > n_lines {
1100 let end = bytes.as_ptr() as usize;
1101 return Ok(end - start);
1102 }
1103
1104 if bytes.is_empty() {
1105 return Ok(original_bytes_len);
1106 } else if is_comment_line(bytes, parse_options.comment_prefix.as_ref()) {
1107 let bytes_rem = skip_this_line_naive(bytes, parse_options.eol_char);
1109 bytes = bytes_rem;
1110 continue;
1111 }
1112
1113 let mut projection_iter = projection.iter().copied();
1117 let mut next_projected = unsafe { projection_iter.next().unwrap_unchecked() };
1118 let mut processed_fields = 0;
1119
1120 let mut iter = SplitFields::new(
1121 bytes,
1122 parse_options.separator,
1123 parse_options.quote_char,
1124 parse_options.eol_char,
1125 );
1126 let mut idx = 0u32;
1127 let mut read_sol = 0;
1128 loop {
1129 match iter.next() {
1130 None => {
1132 bytes = unsafe { bytes.get_unchecked(std::cmp::min(read_sol, bytes.len())..) };
1133 break;
1134 },
1135 Some((mut field, needs_escaping)) => {
1136 let field_len = field.len();
1137
1138 read_sol += field_len + 1;
1140
1141 if idx == next_projected as u32 {
1142 unsafe {
1145 if field_len > 0 && *field.get_unchecked(field_len - 1) == b'\r' {
1146 field = field.get_unchecked(..field_len - 1);
1147 }
1148 }
1149
1150 debug_assert!(processed_fields < buffers.len());
1151 let buf = unsafe {
1152 buffers.get_unchecked_mut(processed_fields)
1154 };
1155 let mut add_null = false;
1156
1157 if let Some(null_values) = null_values {
1159 let field = if needs_escaping && !field.is_empty() {
1160 unsafe { field.get_unchecked(1..field.len() - 1) }
1161 } else {
1162 field
1163 };
1164
1165 add_null = unsafe { null_values.is_null(field, idx as usize) }
1168 }
1169 if add_null {
1170 buf.add_null(!parse_options.missing_is_null && field.is_empty())
1171 } else {
1172 buf.add(field, ignore_errors, needs_escaping, parse_options.missing_is_null)
1173 .map_err(|e| {
1174 let bytes_offset = offset + field.as_ptr() as usize - start;
1175 let unparsable = String::from_utf8_lossy(field);
1176 let column_name = schema.get_at_index(idx as usize).unwrap().0;
1177 polars_err!(
1178 ComputeError:
1179 "could not parse `{}` as dtype `{}` at column '{}' (column number {})\n\n\
1180 The current offset in the file is {} bytes.\n\
1181 \n\
1182 You might want to try:\n\
1183 - increasing `infer_schema_length` (e.g. `infer_schema_length=10000`),\n\
1184 - increasing `infer_schema_files` (e.g. `infer_schema_files=50`),\n\
1185 - specifying correct dtype with the `schema_overrides` argument\n\
1186 - setting `ignore_errors` to `True`,\n\
1187 - adding `{}` to the `null_values` list.\n\n\
1188 Original error: ```{}```",
1189 &unparsable,
1190 buf.dtype(),
1191 column_name,
1192 idx + 1,
1193 bytes_offset,
1194 &unparsable,
1195 e
1196 )
1197 })?;
1198 }
1199 processed_fields += 1;
1200
1201 match projection_iter.next() {
1203 Some(p) => next_projected = p,
1204 None => {
1205 if bytes.get(read_sol - 1) == Some(&parse_options.eol_char) {
1206 bytes = unsafe { bytes.get_unchecked(read_sol..) };
1207 } else {
1208 if !truncate_ragged_lines && read_sol < bytes.len() {
1209 polars_bail!(ComputeError: r#"found more fields than defined in 'Schema'
1210
1211Consider setting 'truncate_ragged_lines={}'."#, polars_error::constants::TRUE)
1212 }
1213 let bytes_rem = skip_this_line(
1214 unsafe { bytes.get_unchecked(read_sol - 1..) },
1215 parse_options.quote_char,
1216 parse_options.eol_char,
1217 );
1218 bytes = bytes_rem;
1219 }
1220 break;
1221 },
1222 }
1223 }
1224 idx += 1;
1225 },
1226 }
1227 }
1228
1229 while processed_fields < projection.len() {
1233 debug_assert!(processed_fields < buffers.len());
1234 let buf = unsafe {
1235 buffers.get_unchecked_mut(processed_fields)
1237 };
1238 buf.add_null(!parse_options.missing_is_null);
1239 processed_fields += 1;
1240 }
1241 line_count += 1;
1242 }
1243}
1244
1245#[cfg(test)]
1246mod test {
1247 use super::SplitLines;
1248
1249 #[test]
1250 fn test_splitlines() {
1251 let input = "1,\"foo\n\"\n2,\"foo\n\"\n";
1252 let mut lines = SplitLines::new(input.as_bytes(), Some(b'"'), b'\n', None);
1253 assert_eq!(lines.next(), Some("1,\"foo\n\"".as_bytes()));
1254 assert_eq!(lines.next(), Some("2,\"foo\n\"".as_bytes()));
1255 assert_eq!(lines.next(), None);
1256
1257 let input2 = "1,'foo\n'\n2,'foo\n'\n";
1258 let mut lines2 = SplitLines::new(input2.as_bytes(), Some(b'\''), b'\n', None);
1259 assert_eq!(lines2.next(), Some("1,'foo\n'".as_bytes()));
1260 assert_eq!(lines2.next(), Some("2,'foo\n'".as_bytes()));
1261 assert_eq!(lines2.next(), None);
1262 }
1263}