Skip to main content

polars_io/csv/read/
streaming.rs

1use std::cmp;
2use std::iter::Iterator;
3use std::num::NonZeroUsize;
4
5use polars_buffer::Buffer;
6use polars_core::prelude::Schema;
7use polars_core::schema::SchemaRef;
8use polars_error::{PolarsResult, polars_bail, polars_ensure};
9
10use crate::csv::read::schema_inference::infer_file_schema_impl;
11use crate::prelude::_csv_read_internal::{SplitLines, is_comment_line};
12use crate::prelude::{CsvParseOptions, CsvReadOptions};
13use crate::utils::compression::{ByteSourceReader, CompressedReader};
14use crate::utils::stream_buf_reader::ReaderSource;
15
16pub type InspectContentFn<'a> = Box<dyn FnMut(&[u8]) + 'a>;
17
18/// Reads bytes from `reader` until the CSV starting point is reached depending on the options.
19///
20/// Returns the inferred schema and leftover bytes not yet consumed, which may be empty. The
21/// leftover bytes + `reader.read_next_slice` is guaranteed to start at first real content row.
22///
23/// `inspect_first_content_row_fn` allows looking at the first content row, this is where parsing
24/// will start. Beware even if the function is provided it's *not* guaranteed that the returned
25/// value will be `Some`, since the CSV may be incomplete.
26///
27/// The reading is done in an iterative streaming fashion
28///
29/// This function isn't perf critical but would increase binary-size so don't inline it.
30#[inline(never)]
31pub fn read_until_start_and_infer_schema_from_compressed_reader(
32    options: &CsvReadOptions,
33    projected_schema: Option<SchemaRef>,
34    ignore_extra_columns: bool,
35    insert_missing_columns: bool,
36    mut inspect_first_content_row_fn: Option<InspectContentFn<'_>>,
37    reader: &mut CompressedReader,
38) -> PolarsResult<(Schema, Buffer<u8>)> {
39    // It's better to be above than below here.
40    const ESTIMATED_BYTES_PER_ROW: usize = 200;
41
42    #[derive(Copy, Clone)]
43    enum State {
44        // Ordered so that all states only happen after the ones before it.
45        SkipEmpty,
46        SkipRowsBeforeHeader(usize),
47        SkipHeader(bool),
48        SkipRowsAfterHeader(usize),
49        ContentInspect,
50        InferCollect,
51        Done,
52    }
53
54    polars_ensure!(
55        !(options.skip_lines != 0 && options.skip_rows != 0),
56        InvalidOperation: "only one of 'skip_rows'/'skip_lines' may be set"
57    );
58
59    // We have to treat skip_lines differently since the lines it skips may not follow regular CSV
60    // quote escape rules.
61    let prev_leftover = skip_lines_naive_from_compressed_reader(
62        options.parse_options.eol_char,
63        options.skip_lines,
64        options.raise_if_empty,
65        reader,
66    )?;
67
68    let mut state = if options.has_header {
69        State::SkipEmpty
70    } else if options.skip_lines != 0 {
71        // skip_lines shouldn't skip extra comments before the header, so directly go to SkipHeader
72        // state.
73        State::SkipHeader(false)
74    } else {
75        State::SkipRowsBeforeHeader(options.skip_rows)
76    };
77
78    let comment_prefix = options.parse_options.comment_prefix.as_ref();
79    let infer_schema_length = if options.schema.is_some() {
80        // Don't actually infer if the schema is set.
81        Some(0)
82    } else {
83        options.infer_schema_length
84    };
85
86    let mut header_line = None;
87    let mut content_lines = Vec::with_capacity(infer_schema_length.unwrap_or_else(|| {
88        reader
89            .total_len_estimate()
90            .saturating_div(ESTIMATED_BYTES_PER_ROW)
91    }));
92
93    // In the compressed case `reader.read_next_slice` has to copy the previous leftover into a new
94    // `Vec` which would lead to quadratic copying if we don't factor in `infer_schema_length` into
95    // the initial read size. We have to retain the row memory for schema inference and also for
96    // actual morsel generation. If `infer_schema_length` is set to `None` we will have to read the
97    // full input anyway so we can do so once and avoid re-copying.
98    let initial_read_size = infer_schema_length
99        .map(|isl| {
100            cmp::max(
101                CompressedReader::initial_read_size(),
102                isl.saturating_mul(ESTIMATED_BYTES_PER_ROW),
103            )
104        })
105        .unwrap_or(usize::MAX);
106
107    let leftover = for_each_line_from_reader_from_compressed_reader(
108        &options.parse_options,
109        true,
110        prev_leftover,
111        initial_read_size,
112        reader,
113        |mem_slice_line| {
114            let line = &*mem_slice_line;
115
116            let done = loop {
117                match &mut state {
118                    State::SkipEmpty => {
119                        if line.is_empty() || line == b"\r" {
120                            break LineUse::ConsumeDiscard;
121                        }
122
123                        state = State::SkipRowsBeforeHeader(options.skip_rows);
124                    },
125                    State::SkipRowsBeforeHeader(remaining) => {
126                        let is_comment = is_comment_line(line, comment_prefix);
127
128                        if *remaining == 0 && !is_comment {
129                            state = State::SkipHeader(false);
130                            continue;
131                        }
132
133                        *remaining -= !is_comment as usize;
134                        break LineUse::ConsumeDiscard;
135                    },
136                    State::SkipHeader(did_skip) => {
137                        if !options.has_header || *did_skip {
138                            state = State::SkipRowsAfterHeader(options.skip_rows_after_header);
139                            continue;
140                        }
141
142                        header_line = Some(mem_slice_line.clone());
143                        *did_skip = true;
144                        break LineUse::ConsumeDiscard;
145                    },
146                    State::SkipRowsAfterHeader(remaining) => {
147                        let is_comment = is_comment_line(line, comment_prefix);
148
149                        if *remaining == 0 && !is_comment {
150                            state = State::ContentInspect;
151                            continue;
152                        }
153
154                        *remaining -= !is_comment as usize;
155                        break LineUse::ConsumeDiscard;
156                    },
157                    State::ContentInspect => {
158                        if let Some(func) = &mut inspect_first_content_row_fn {
159                            func(line);
160                        }
161
162                        state = State::InferCollect;
163                    },
164                    State::InferCollect => {
165                        if !is_comment_line(line, comment_prefix) {
166                            content_lines.push(mem_slice_line.clone());
167                            if content_lines.len() >= infer_schema_length.unwrap_or(usize::MAX) {
168                                state = State::Done;
169                                continue;
170                            }
171                        }
172
173                        break LineUse::ConsumeKeep;
174                    },
175                    State::Done => {
176                        break LineUse::Done;
177                    },
178                }
179            };
180
181            Ok(done)
182        },
183    )?;
184
185    let infer_all_as_str = infer_schema_length == Some(0);
186
187    let inferred_schema = infer_schema(
188        &header_line,
189        &content_lines,
190        infer_all_as_str,
191        options,
192        projected_schema,
193        ignore_extra_columns,
194        insert_missing_columns,
195    )?;
196
197    Ok((inferred_schema, leftover))
198}
199
200/// Reads bytes from `reader` until the CSV starting point is reached depending on the options.
201///
202/// Returns the inferred schema and leftover bytes not yet consumed, which may be empty. The
203/// leftover bytes + `reader.read_next_slice` is guaranteed to start at first real content row.
204///
205/// `inspect_first_content_row_fn` allows looking at the first content row, this is where parsing
206/// will start. Beware even if the function is provided it's *not* guaranteed that the returned
207/// value will be `Some`, since the CSV may be incomplete.
208///
209/// The reading is done in an iterative streaming fashion
210///
211/// This function isn't perf critical but would increase binary-size so don't inline it.
212#[inline(never)]
213pub fn read_until_start_and_infer_schema(
214    options: &CsvReadOptions,
215    projected_schema: Option<SchemaRef>,
216    ignore_extra_columns: bool,
217    insert_missing_columns: bool,
218    decompressed_file_size_hint: Option<usize>,
219    mut inspect_first_content_row_fn: Option<InspectContentFn<'_>>,
220    reader: &mut ByteSourceReader<ReaderSource>,
221) -> PolarsResult<(Schema, Buffer<u8>)> {
222    // It's better to be above than below here.
223    const ESTIMATED_BYTES_PER_ROW: usize = 200;
224
225    #[derive(Copy, Clone)]
226    enum State {
227        // Ordered so that all states only happen after the ones before it.
228        SkipEmpty,
229        SkipRowsBeforeHeader(usize),
230        SkipHeader(bool),
231        SkipRowsAfterHeader(usize),
232        ContentInspect,
233        InferCollect,
234        Done,
235    }
236
237    polars_ensure!(
238        !(options.skip_lines != 0 && options.skip_rows != 0),
239        InvalidOperation: "only one of 'skip_rows'/'skip_lines' may be set"
240    );
241
242    // We have to treat skip_lines differently since the lines it skips may not follow regular CSV
243    // quote escape rules.
244    let prev_leftover = skip_lines_naive(
245        options.parse_options.eol_char,
246        options.skip_lines,
247        options.raise_if_empty,
248        decompressed_file_size_hint,
249        reader,
250    )?;
251
252    let mut state = if options.has_header {
253        State::SkipEmpty
254    } else if options.skip_lines != 0 {
255        // skip_lines shouldn't skip extra comments before the header, so directly go to SkipHeader
256        // state.
257        State::SkipHeader(false)
258    } else {
259        State::SkipRowsBeforeHeader(options.skip_rows)
260    };
261
262    let comment_prefix = options.parse_options.comment_prefix.as_ref();
263    let infer_schema_length = if options.schema.is_some() {
264        // Don't actually infer if the schema is set.
265        Some(0)
266    } else {
267        options.infer_schema_length
268    };
269
270    let mut header_line = None;
271    let mut content_lines = Vec::with_capacity(infer_schema_length.unwrap_or_else(|| {
272        decompressed_file_size_hint
273            .map(|size| size.saturating_div(ESTIMATED_BYTES_PER_ROW))
274            .unwrap_or(100)
275    }));
276
277    // In the compressed case `reader.read_next_slice` has to copy the previous leftover into a new
278    // `Vec` which would lead to quadratic copying if we don't factor in `infer_schema_length` into
279    // the initial read size. We have to retain the row memory for schema inference and also for
280    // actual morsel generation. If `infer_schema_length` is set to `None` we will have to read the
281    // full input anyway so we can do so once and avoid re-copying.
282    let initial_read_size = infer_schema_length
283        .map(|isl| {
284            cmp::max(
285                CompressedReader::initial_read_size(),
286                isl.saturating_mul(ESTIMATED_BYTES_PER_ROW),
287            )
288        })
289        .unwrap_or(usize::MAX);
290
291    let leftover = for_each_line_from_reader(
292        &options.parse_options,
293        true,
294        prev_leftover,
295        initial_read_size,
296        decompressed_file_size_hint,
297        reader,
298        |mem_slice_line| {
299            let line = &*mem_slice_line;
300
301            let done = loop {
302                match &mut state {
303                    State::SkipEmpty => {
304                        if line.is_empty() || line == b"\r" {
305                            break LineUse::ConsumeDiscard;
306                        }
307
308                        state = State::SkipRowsBeforeHeader(options.skip_rows);
309                    },
310                    State::SkipRowsBeforeHeader(remaining) => {
311                        let is_comment = is_comment_line(line, comment_prefix);
312
313                        if *remaining == 0 && !is_comment {
314                            state = State::SkipHeader(false);
315                            continue;
316                        }
317
318                        *remaining -= !is_comment as usize;
319                        break LineUse::ConsumeDiscard;
320                    },
321                    State::SkipHeader(did_skip) => {
322                        if !options.has_header || *did_skip {
323                            state = State::SkipRowsAfterHeader(options.skip_rows_after_header);
324                            continue;
325                        }
326
327                        header_line = Some(mem_slice_line.clone());
328                        *did_skip = true;
329                        break LineUse::ConsumeDiscard;
330                    },
331                    State::SkipRowsAfterHeader(remaining) => {
332                        let is_comment = is_comment_line(line, comment_prefix);
333
334                        if *remaining == 0 && !is_comment {
335                            state = State::ContentInspect;
336                            continue;
337                        }
338
339                        *remaining -= !is_comment as usize;
340                        break LineUse::ConsumeDiscard;
341                    },
342                    State::ContentInspect => {
343                        if let Some(func) = &mut inspect_first_content_row_fn {
344                            func(line);
345                        }
346
347                        state = State::InferCollect;
348                    },
349                    State::InferCollect => {
350                        if !is_comment_line(line, comment_prefix) {
351                            content_lines.push(mem_slice_line.clone());
352                            if content_lines.len() >= infer_schema_length.unwrap_or(usize::MAX) {
353                                state = State::Done;
354                                continue;
355                            }
356                        }
357
358                        break LineUse::ConsumeKeep;
359                    },
360                    State::Done => {
361                        break LineUse::Done;
362                    },
363                }
364            };
365
366            Ok(done)
367        },
368    )?;
369
370    let infer_all_as_str = infer_schema_length == Some(0);
371
372    let inferred_schema = infer_schema(
373        &header_line,
374        &content_lines,
375        infer_all_as_str,
376        options,
377        projected_schema,
378        ignore_extra_columns,
379        insert_missing_columns,
380    )?;
381
382    Ok((inferred_schema, leftover))
383}
384
385enum LineUse {
386    ConsumeDiscard,
387    ConsumeKeep,
388    Done,
389}
390
391/// Iterate over valid CSV lines produced by reader.
392///
393/// Returning `ConsumeDiscard` after `ConsumeKeep` is a logic error, since a segmented `Buffer`
394/// can't be constructed.
395fn for_each_line_from_reader_from_compressed_reader(
396    parse_options: &CsvParseOptions,
397    is_file_start: bool,
398    mut prev_leftover: Buffer<u8>,
399    initial_read_size: usize,
400    reader: &mut CompressedReader,
401    mut line_fn: impl FnMut(Buffer<u8>) -> PolarsResult<LineUse>,
402) -> PolarsResult<Buffer<u8>> {
403    let mut is_first_line = is_file_start;
404
405    let fixed_read_size = std::env::var("POLARS_FORCE_CSV_INFER_READ_SIZE")
406        .map(|x| {
407            x.parse::<NonZeroUsize>()
408                .unwrap_or_else(|_| {
409                    panic!("invalid value for POLARS_FORCE_CSV_INFER_READ_SIZE: {x}")
410                })
411                .get()
412        })
413        .ok();
414
415    let mut read_size = fixed_read_size.unwrap_or(initial_read_size);
416    let mut retain_offset = None;
417
418    loop {
419        let (mut slice, bytes_read) = reader.read_next_slice(&prev_leftover, read_size)?;
420        if slice.is_empty() {
421            return Ok(Buffer::new());
422        }
423
424        if is_first_line {
425            is_first_line = false;
426            const UTF8_BOM_MARKER: Option<&[u8]> = Some(b"\xef\xbb\xbf");
427            if slice.get(0..3) == UTF8_BOM_MARKER {
428                slice = slice.sliced(3..);
429            }
430        }
431
432        let line_to_sub_slice = |line: &[u8]| {
433            let start = line.as_ptr() as usize - slice.as_ptr() as usize;
434            slice.clone().sliced(start..(start + line.len()))
435        };
436
437        // When reading a CSV with `has_header=False` we need to read up to `infer_schema_length` lines, but we only want to decompress the input once, so we grow a `Buffer` that will be returned as leftover.
438        let effective_slice = if let Some(offset) = retain_offset {
439            slice.clone().sliced(offset..)
440        } else {
441            slice.clone()
442        };
443
444        let mut lines = SplitLines::new(
445            &effective_slice,
446            parse_options.quote_char,
447            parse_options.eol_char,
448            parse_options.comment_prefix.as_ref(),
449        );
450        let Some(mut prev_line) = lines.next() else {
451            read_size = read_size.saturating_mul(2);
452            prev_leftover = slice;
453            continue;
454        };
455
456        let mut should_ret = false;
457
458        // The last line in `SplitLines` may be incomplete if `slice` ends before the file does, so
459        // we iterate everything except the last line.
460        for next_line in lines {
461            match line_fn(line_to_sub_slice(prev_line))? {
462                LineUse::ConsumeDiscard => debug_assert!(retain_offset.is_none()),
463                LineUse::ConsumeKeep => {
464                    if retain_offset.is_none() {
465                        let retain_start_offset =
466                            prev_line.as_ptr() as usize - slice.as_ptr() as usize;
467                        prev_leftover = slice.clone().sliced(retain_start_offset..);
468                        retain_offset = Some(0);
469                    }
470                },
471                LineUse::Done => {
472                    should_ret = true;
473                    break;
474                },
475            }
476            prev_line = next_line;
477        }
478
479        let mut unconsumed_offset = prev_line.as_ptr() as usize - effective_slice.as_ptr() as usize;
480
481        // EOF file reached, the last line will have no continuation on the next call to
482        // `read_next_slice`.
483        if bytes_read < read_size {
484            match line_fn(line_to_sub_slice(prev_line))? {
485                LineUse::ConsumeDiscard => {
486                    debug_assert!(retain_offset.is_none());
487                    unconsumed_offset += prev_line.len();
488                    if effective_slice.get(unconsumed_offset) == Some(&parse_options.eol_char) {
489                        unconsumed_offset += 1;
490                    }
491                },
492                LineUse::ConsumeKeep | LineUse::Done => (),
493            }
494            should_ret = true;
495        }
496
497        if let Some(offset) = &mut retain_offset {
498            if *offset == 0 {
499                // `unconsumed_offset` was computed with the full `slice` as base reference
500                // compensate retained offset.
501                *offset = unconsumed_offset - (slice.len() - prev_leftover.len());
502            } else {
503                prev_leftover = slice;
504                *offset += unconsumed_offset;
505            }
506        } else {
507            // Since `read_next_slice` has to copy the leftover bytes in the decompression case,
508            // it's more efficient to hand in as little as possible.
509            prev_leftover = slice.sliced(unconsumed_offset..);
510        }
511
512        if should_ret {
513            return Ok(prev_leftover);
514        }
515
516        if fixed_read_size.is_none() {
517            if read_size < CompressedReader::ideal_read_size() {
518                read_size *= 4;
519            }
520            // Copying readers re-copy `prev_leftover` on every read and kept lines pin each copy.
521            read_size = read_size.max(prev_leftover.len());
522        }
523    }
524}
525
526/// Iterate over valid CSV lines produced by reader.
527///
528/// Returning `ConsumeDiscard` after `ConsumeKeep` is a logic error, since a segmented `Buffer`
529/// can't be constructed.
530fn for_each_line_from_reader(
531    parse_options: &CsvParseOptions,
532    is_file_start: bool,
533    mut prev_leftover: Buffer<u8>,
534    initial_read_size: usize,
535    decompressed_file_size_hint: Option<usize>,
536    reader: &mut ByteSourceReader<ReaderSource>,
537    mut line_fn: impl FnMut(Buffer<u8>) -> PolarsResult<LineUse>,
538) -> PolarsResult<Buffer<u8>> {
539    let mut is_first_line = is_file_start;
540
541    let fixed_read_size = std::env::var("POLARS_FORCE_CSV_INFER_READ_SIZE")
542        .map(|x| {
543            x.parse::<NonZeroUsize>()
544                .unwrap_or_else(|_| {
545                    panic!("invalid value for POLARS_FORCE_CSV_INFER_READ_SIZE: {x}")
546                })
547                .get()
548        })
549        .ok();
550
551    let mut read_size = fixed_read_size.unwrap_or(initial_read_size);
552    let mut retain_offset = None;
553
554    loop {
555        let (mut slice, bytes_read) =
556            reader.read_next_slice(&prev_leftover, read_size, decompressed_file_size_hint)?;
557        if slice.is_empty() {
558            return Ok(Buffer::new());
559        }
560
561        if is_first_line {
562            is_first_line = false;
563            const UTF8_BOM_MARKER: Option<&[u8]> = Some(b"\xef\xbb\xbf");
564            if slice.get(0..3) == UTF8_BOM_MARKER {
565                slice = slice.sliced(3..);
566            }
567        }
568
569        let line_to_sub_slice = |line: &[u8]| {
570            let start = line.as_ptr() as usize - slice.as_ptr() as usize;
571            slice.clone().sliced(start..(start + line.len()))
572        };
573
574        // When reading a CSV with `has_header=False` we need to read up to `infer_schema_length` lines, but we only want to decompress the input once, so we grow a `Buffer` that will be returned as leftover.
575        let effective_slice = if let Some(offset) = retain_offset {
576            slice.clone().sliced(offset..)
577        } else {
578            slice.clone()
579        };
580
581        let mut lines = SplitLines::new(
582            &effective_slice,
583            parse_options.quote_char,
584            parse_options.eol_char,
585            parse_options.comment_prefix.as_ref(),
586        );
587        let Some(mut prev_line) = lines.next() else {
588            read_size = read_size.saturating_mul(2);
589            prev_leftover = slice;
590            continue;
591        };
592
593        let mut should_ret = false;
594
595        // The last line in `SplitLines` may be incomplete if `slice` ends before the file does, so
596        // we iterate everything except the last line.
597        for next_line in lines {
598            match line_fn(line_to_sub_slice(prev_line))? {
599                LineUse::ConsumeDiscard => debug_assert!(retain_offset.is_none()),
600                LineUse::ConsumeKeep => {
601                    if retain_offset.is_none() {
602                        let retain_start_offset =
603                            prev_line.as_ptr() as usize - slice.as_ptr() as usize;
604                        prev_leftover = slice.clone().sliced(retain_start_offset..);
605                        retain_offset = Some(0);
606                    }
607                },
608                LineUse::Done => {
609                    should_ret = true;
610                    break;
611                },
612            }
613            prev_line = next_line;
614        }
615
616        let mut unconsumed_offset = prev_line.as_ptr() as usize - effective_slice.as_ptr() as usize;
617
618        // EOF file reached, the last line will have no continuation on the next call to
619        // `read_next_slice`.
620        if bytes_read < read_size {
621            match line_fn(line_to_sub_slice(prev_line))? {
622                LineUse::ConsumeDiscard => {
623                    debug_assert!(retain_offset.is_none());
624                    unconsumed_offset += prev_line.len();
625                    if effective_slice.get(unconsumed_offset) == Some(&parse_options.eol_char) {
626                        unconsumed_offset += 1;
627                    }
628                },
629                LineUse::ConsumeKeep | LineUse::Done => (),
630            }
631            should_ret = true;
632        }
633
634        if let Some(offset) = &mut retain_offset {
635            if *offset == 0 {
636                // `unconsumed_offset` was computed with the full `slice` as base reference
637                // compensate retained offset.
638                *offset = unconsumed_offset - (slice.len() - prev_leftover.len());
639            } else {
640                prev_leftover = slice;
641                *offset += unconsumed_offset;
642            }
643        } else {
644            // Since `read_next_slice` has to copy the leftover bytes in the decompression case,
645            // it's more efficient to hand in as little as possible.
646            prev_leftover = slice.sliced(unconsumed_offset..);
647        }
648
649        if should_ret {
650            return Ok(prev_leftover);
651        }
652
653        if fixed_read_size.is_none() {
654            if read_size < ByteSourceReader::<ReaderSource>::ideal_read_size() {
655                read_size *= 4;
656            }
657            // Copying readers re-copy `prev_leftover` on every read and kept lines pin each copy.
658            read_size = read_size.max(prev_leftover.len());
659        }
660    }
661}
662
663fn skip_lines_naive_from_compressed_reader(
664    eol_char: u8,
665    skip_lines: usize,
666    raise_if_empty: bool,
667    reader: &mut CompressedReader,
668) -> PolarsResult<Buffer<u8>> {
669    let mut prev_leftover = Buffer::new();
670
671    if skip_lines == 0 {
672        return Ok(prev_leftover);
673    }
674
675    let mut remaining = skip_lines;
676    let mut read_size = CompressedReader::initial_read_size();
677
678    loop {
679        let (slice, bytes_read) = reader.read_next_slice(&prev_leftover, read_size)?;
680        let mut bytes: &[u8] = &slice;
681
682        'inner: loop {
683            let Some(mut pos) = memchr::memchr(eol_char, bytes) else {
684                read_size = read_size.saturating_mul(2);
685                break 'inner;
686            };
687            pos = cmp::min(pos + 1, bytes.len());
688
689            bytes = &bytes[pos..];
690            remaining -= 1;
691
692            if remaining == 0 {
693                let unconsumed_offset = bytes.as_ptr() as usize - slice.as_ptr() as usize;
694                prev_leftover = slice.sliced(unconsumed_offset..);
695                return Ok(prev_leftover);
696            }
697        }
698
699        if bytes_read == 0 {
700            if raise_if_empty {
701                polars_bail!(NoData: "specified skip_lines is larger than total number of lines.");
702            } else {
703                return Ok(Buffer::new());
704            }
705        }
706
707        // No need to search for naive eol twice in the leftover.
708        prev_leftover = Buffer::new();
709
710        if read_size < CompressedReader::ideal_read_size() {
711            read_size *= 4;
712        }
713    }
714}
715
716fn skip_lines_naive(
717    eol_char: u8,
718    skip_lines: usize,
719    raise_if_empty: bool,
720    decompressed_file_size_hint: Option<usize>,
721    reader: &mut ByteSourceReader<ReaderSource>,
722) -> PolarsResult<Buffer<u8>> {
723    let mut prev_leftover = Buffer::new();
724
725    if skip_lines == 0 {
726        return Ok(prev_leftover);
727    }
728
729    let mut remaining = skip_lines;
730    let mut read_size = CompressedReader::initial_read_size();
731
732    loop {
733        let (slice, bytes_read) =
734            reader.read_next_slice(&prev_leftover, read_size, decompressed_file_size_hint)?;
735        let mut bytes: &[u8] = &slice;
736
737        'inner: loop {
738            let Some(mut pos) = memchr::memchr(eol_char, bytes) else {
739                read_size = read_size.saturating_mul(2);
740                break 'inner;
741            };
742            pos = cmp::min(pos + 1, bytes.len());
743
744            bytes = &bytes[pos..];
745            remaining -= 1;
746
747            if remaining == 0 {
748                let unconsumed_offset = bytes.as_ptr() as usize - slice.as_ptr() as usize;
749                prev_leftover = slice.sliced(unconsumed_offset..);
750                return Ok(prev_leftover);
751            }
752        }
753
754        if bytes_read == 0 {
755            if raise_if_empty {
756                polars_bail!(NoData: "specified skip_lines is larger than total number of lines.");
757            } else {
758                return Ok(Buffer::new());
759            }
760        }
761
762        // No need to search for naive eol twice in the leftover.
763        prev_leftover = Buffer::new();
764
765        if read_size < CompressedReader::ideal_read_size() {
766            read_size *= 4;
767        }
768    }
769}
770
771fn infer_schema(
772    header_line: &Option<Buffer<u8>>,
773    content_lines: &[Buffer<u8>],
774    infer_all_as_str: bool,
775    options: &CsvReadOptions,
776    projected_schema: Option<SchemaRef>,
777    ignore_extra_columns: bool,
778    insert_missing_columns: bool,
779) -> PolarsResult<Schema> {
780    let has_no_inference_data = if options.has_header {
781        header_line.is_none()
782    } else {
783        content_lines.is_empty()
784    };
785
786    if options.raise_if_empty && has_no_inference_data {
787        polars_bail!(NoData: "empty CSV");
788    }
789
790    let mut inferred_schema = if has_no_inference_data {
791        Schema::default()
792    } else {
793        infer_file_schema_impl(
794            header_line,
795            content_lines,
796            infer_all_as_str,
797            &options.parse_options,
798            options.column_names_overwrite.as_deref(),
799            options.schema_overwrite.as_deref(),
800            ignore_extra_columns,
801            insert_missing_columns,
802        )?
803    };
804
805    if let Some(schema) = &options.schema {
806        if !has_no_inference_data {
807            if !ignore_extra_columns {
808                let mut extra_names = vec![];
809
810                let num_extra_names = if options.has_header {
811                    extra_names.extend(
812                        inferred_schema
813                            .iter_names()
814                            .filter(|name| !schema.contains(name))
815                            .collect::<Vec<_>>(),
816                    );
817                    extra_names.len()
818                } else {
819                    inferred_schema.len().saturating_sub(schema.len())
820                };
821
822                if num_extra_names != 0 {
823                    let mut names = String::new();
824
825                    if !extra_names.is_empty() {
826                        names = format!(" (extra names: {extra_names:?})");
827                    }
828
829                    polars_bail!(
830                        SchemaMismatch:
831                        "CSV file contained column names not specified in schema (n_extra = {num_extra_names}). \
832                        Specify these names in the schema, or pass `extra_columns='ignore'` to \
833                        ignore these columns.{names}"
834                    )
835                }
836            }
837
838            if !insert_missing_columns {
839                let mut missing_names = vec![];
840
841                let num_missing_names = if options.has_header {
842                    missing_names.extend(
843                        schema
844                            .iter_names()
845                            .filter(|name| !inferred_schema.contains(name))
846                            .collect::<Vec<_>>(),
847                    );
848                    missing_names.len()
849                } else {
850                    schema.len().saturating_sub(inferred_schema.len())
851                };
852
853                if num_missing_names != 0 {
854                    let mut names = String::new();
855
856                    if !missing_names.is_empty() {
857                        names = format!(" (missing names: {missing_names:?})");
858                    }
859
860                    polars_bail!(
861                        SchemaMismatch:
862                        "column names specified in schema not found in CSV file (n_missing = {num_missing_names}). \
863                        Remove these names from the schema, or pass `missing_columns='insert'` to \
864                        insert these columns with NULL row values.{names}"
865                    )
866                }
867            }
868        }
869
870        if !options.has_header {
871            inferred_schema = schema.as_ref().clone();
872        } else {
873            for (name, dtype) in schema.iter() {
874                inferred_schema.insert(name.clone(), dtype.clone());
875            }
876        }
877    }
878
879    if let Some(dtypes) = options.dtype_overwrite.as_deref() {
880        polars_ensure!(
881            dtypes.len() == inferred_schema.len(),
882            SchemaMismatch:
883            "The number of dtypes in schema override must be equal to the number \
884            of fields in the file ({} != {}).",
885            dtypes.len(), inferred_schema.len()
886        );
887        for (i, dtype) in dtypes.iter().enumerate() {
888            inferred_schema.set_dtype_at_index(i, dtype.clone());
889        }
890    }
891
892    // TODO: We currently always override with the projected dtype, but this may cause issues e.g.
893    // with temporal types. This can be improved to better choose between the 2 dtypes.
894    if let Some(projected_schema) = projected_schema {
895        for (name, inferred_dtype) in inferred_schema.iter_mut() {
896            if let Some(projected_dtype) = projected_schema.get(name) {
897                *inferred_dtype = projected_dtype.clone();
898            }
899        }
900    }
901
902    Ok(inferred_schema)
903}