polars_io/csv/read/
read_impl.rs

1use std::fmt;
2use std::sync::Mutex;
3
4use polars_buffer::{Buffer, SharedStorage};
5#[cfg(feature = "temporal")]
6use polars_core::chunked_array::temporal::string::StringMethods;
7use polars_core::prelude::*;
8use polars_core::runtime::RAYON;
9use polars_core::utils::{accumulate_dataframes_vertical, handle_casting_failures};
10use polars_utils::relaxed_cell::RelaxedCell;
11use rayon::prelude::*;
12
13use super::CsvParseOptions;
14use super::builder::init_builders;
15use super::options::{CsvEncoding, NullValuesCompiled};
16use super::parser::{CountLines, is_comment_line, parse_lines};
17use super::reader::prepare_csv_schema;
18#[cfg(feature = "decompress")]
19use super::utils::decompress;
20use crate::RowIndex;
21use crate::csv::read::{CsvReadOptions, read_until_start_and_infer_schema_from_compressed_reader};
22use crate::mmap::ReaderBytes;
23use crate::predicates::PhysicalIoExpr;
24use crate::utils::compression::{CompressedReader, SupportedCompression};
25use crate::utils::update_row_counts2;
26
27pub fn cast_columns(
28    df: &mut DataFrame,
29    to_cast: &[Field],
30    parallel: bool,
31    ignore_errors: bool,
32) -> PolarsResult<()> {
33    let cast_fn = |c: &Column, fld: &Field| {
34        let out = match (c.dtype(), fld.dtype()) {
35            #[cfg(feature = "temporal")]
36            (DataType::String, DataType::Date) => c
37                .str()
38                .unwrap()
39                .as_date(None, false)
40                .map(|ca| ca.into_column()),
41            #[cfg(feature = "temporal")]
42            (DataType::String, DataType::Time) => c
43                .str()
44                .unwrap()
45                .as_time(None, false)
46                .map(|ca| ca.into_column()),
47            #[cfg(feature = "temporal")]
48            (DataType::String, DataType::Datetime(tu, _)) => c
49                .str()
50                .unwrap()
51                .as_datetime(
52                    None,
53                    *tu,
54                    false,
55                    false,
56                    None,
57                    &StringChunked::from_iter(std::iter::once("raise")),
58                )
59                .map(|ca| ca.into_column()),
60            (_, dt) => c.cast(dt),
61        }?;
62        if !ignore_errors && c.null_count() != out.null_count() {
63            handle_casting_failures(c.as_materialized_series(), out.as_materialized_series())?;
64        }
65        Ok(out)
66    };
67
68    if parallel {
69        let cols = RAYON.install(|| {
70            df.columns()
71                .into_par_iter()
72                .map(|s| {
73                    if let Some(fld) = to_cast.iter().find(|fld| fld.name() == s.name()) {
74                        cast_fn(s, fld)
75                    } else {
76                        Ok(s.clone())
77                    }
78                })
79                .collect::<PolarsResult<Vec<_>>>()
80        })?;
81        *df = unsafe { DataFrame::new_unchecked(df.height(), cols) }
82    } else {
83        // cast to the original dtypes in the schema
84        for fld in to_cast {
85            // field may not be projected
86            if let Some(idx) = df.get_column_index(fld.name()) {
87                df.try_apply_at_idx(idx, |s| cast_fn(s, fld))?;
88            }
89        }
90    }
91    Ok(())
92}
93
94struct ReaderBytesAndDependents<'a> {
95    // Ensure lifetime dependents are dropped before `reader_bytes`, since their drop impls
96    // could access themselves, this is achieved by placing them before `reader_bytes`.
97    // SAFETY: This is lifetime bound to `reader_bytes`
98    compressed_reader: CompressedReader,
99    // SAFETY: This is lifetime bound to `reader_bytes`
100    leftover: Buffer<u8>,
101    _reader_bytes: ReaderBytes<'a>,
102}
103
104/// CSV file reader
105pub(crate) struct CoreReader<'a> {
106    reader_bytes: Option<ReaderBytesAndDependents<'a>>,
107
108    /// Explicit schema for the CSV file
109    schema: SchemaRef,
110    parse_options: CsvParseOptions,
111    /// Optional projection for which columns to load (zero-based column indices)
112    projection: Option<Vec<usize>>,
113    /// Current line number, used in error reporting
114    current_line: usize,
115    ignore_errors: bool,
116    n_rows: Option<usize>,
117    n_threads: Option<usize>,
118    null_values: Option<NullValuesCompiled>,
119    predicate: Option<Arc<dyn PhysicalIoExpr>>,
120    to_cast: Vec<Field>,
121    row_index: Option<RowIndex>,
122}
123
124impl fmt::Debug for CoreReader<'_> {
125    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
126        f.debug_struct("Reader")
127            .field("schema", &self.schema)
128            .field("projection", &self.projection)
129            .field("current_line", &self.current_line)
130            .finish()
131    }
132}
133
134impl<'a> CoreReader<'a> {
135    #[allow(clippy::too_many_arguments)]
136    pub(crate) fn new(
137        reader_bytes: ReaderBytes<'a>,
138        parse_options: Arc<CsvParseOptions>,
139        n_rows: Option<usize>,
140        skip_rows: usize,
141        skip_lines: usize,
142        mut projection: Option<Vec<usize>>,
143        max_records: Option<usize>,
144        has_header: bool,
145        ignore_errors: bool,
146        schema: Option<SchemaRef>,
147        columns: Option<Arc<[PlSmallStr]>>,
148        n_threads: Option<usize>,
149        schema_overwrite: Option<SchemaRef>,
150        dtype_overwrite: Option<Arc<Vec<DataType>>>,
151        predicate: Option<Arc<dyn PhysicalIoExpr>>,
152        mut to_cast: Vec<Field>,
153        skip_rows_after_header: usize,
154        row_index: Option<RowIndex>,
155        raise_if_empty: bool,
156    ) -> PolarsResult<CoreReader<'a>> {
157        let separator = parse_options.separator;
158
159        #[cfg(feature = "decompress")]
160        let mut reader_bytes = reader_bytes;
161
162        if !cfg!(feature = "decompress") && SupportedCompression::check(&reader_bytes).is_some() {
163            polars_bail!(
164                ComputeError: "cannot read compressed CSV file; \
165                compile with feature 'decompress'"
166            );
167        }
168        // We keep track of the inferred schema bool
169        // In case the file is compressed this schema inference is wrong and has to be done
170        // again after decompression.
171        #[cfg(feature = "decompress")]
172        {
173            let total_n_rows =
174                n_rows.map(|n| skip_rows + (has_header as usize) + skip_rows_after_header + n);
175            if let Some(b) = decompress(
176                &reader_bytes,
177                total_n_rows,
178                separator,
179                parse_options.quote_char,
180                parse_options.eol_char,
181            ) {
182                reader_bytes = ReaderBytes::Owned(b.into());
183            }
184        }
185
186        let reader_slice = match &reader_bytes {
187            ReaderBytes::Borrowed(slice) => {
188                // SAFETY: The produced slice and derived slices MUST not live longer than
189                // `reader_bytes`. TODO use `scan_csv` to implement `read_csv`.
190                let ss = unsafe { SharedStorage::from_slice_unchecked(slice) };
191                Buffer::from_storage(ss)
192            },
193            ReaderBytes::Owned(slice) => slice.clone(),
194        };
195        let mut compressed_reader = CompressedReader::try_new(reader_slice)?;
196
197        let read_options = CsvReadOptions {
198            parse_options: parse_options.clone(),
199            n_rows,
200            skip_rows,
201            skip_lines,
202            projection: projection.clone().map(Arc::new),
203            has_header,
204            ignore_errors,
205            schema: schema.clone(),
206            columns: columns.clone(),
207            n_threads,
208            schema_overwrite,
209            dtype_overwrite: dtype_overwrite.clone(),
210            fields_to_cast: to_cast.clone(),
211            skip_rows_after_header,
212            row_index: row_index.clone(),
213            raise_if_empty,
214            infer_schema_length: max_records,
215            ..Default::default()
216        };
217
218        // Since this is also used to skip to the start, always call it.
219        let (inferred_schema, leftover) = read_until_start_and_infer_schema_from_compressed_reader(
220            &read_options,
221            None,
222            false,
223            false,
224            None,
225            &mut compressed_reader,
226        )?;
227
228        let mut schema = match schema {
229            Some(schema) => schema,
230            None => Arc::new(inferred_schema),
231        };
232        if let Some(dtypes) = dtype_overwrite {
233            polars_ensure!(
234                dtypes.len() <= schema.len(),
235                InvalidOperation: "The number of schema overrides must be less than or equal to the number of fields (got {} overrides for {} fields)",
236                dtypes.len(), schema.len()
237            );
238            let s = Arc::make_mut(&mut schema);
239            for (index, dt) in dtypes.iter().enumerate() {
240                s.set_dtype_at_index(index, dt.clone()).unwrap();
241            }
242        }
243
244        prepare_csv_schema(&mut schema, &mut to_cast)?;
245
246        // Create a null value for every column
247        let null_values = parse_options
248            .null_values
249            .as_ref()
250            .map(|nv| nv.clone().compile(&schema))
251            .transpose()?;
252
253        if let Some(cols) = columns {
254            let mut prj = Vec::with_capacity(cols.len());
255            for col in cols.as_ref() {
256                let i = schema.try_index_of(col)?;
257                prj.push(i);
258            }
259            projection = Some(prj);
260        }
261
262        Ok(CoreReader {
263            reader_bytes: Some(ReaderBytesAndDependents {
264                compressed_reader,
265                leftover,
266                _reader_bytes: reader_bytes,
267            }),
268            parse_options: (*parse_options).clone(),
269            schema,
270            projection,
271            current_line: usize::from(has_header),
272            ignore_errors,
273            n_rows,
274            n_threads,
275            null_values,
276            predicate,
277            to_cast,
278            row_index,
279        })
280    }
281
282    fn get_projection(&mut self) -> PolarsResult<Vec<usize>> {
283        // we also need to sort the projection to have predictable output.
284        // the `parse_lines` function expects this.
285        self.projection
286            .take()
287            .map(|mut v| {
288                v.sort_unstable();
289                if let Some(idx) = v.last() {
290                    polars_ensure!(*idx < self.schema.len(), OutOfBounds: "projection index: {} is out of bounds for csv schema with length: {}", idx, self.schema.len())
291                }
292                Ok(v)
293            })
294            .unwrap_or_else(|| Ok((0..self.schema.len()).collect()))
295    }
296
297    fn read_chunk(
298        &self,
299        bytes: &[u8],
300        projection: &[usize],
301        bytes_offset: usize,
302        capacity: usize,
303        starting_point_offset: Option<usize>,
304        stop_at_nbytes: usize,
305    ) -> PolarsResult<DataFrame> {
306        let mut df = read_chunk(
307            bytes,
308            &self.parse_options,
309            self.schema.as_ref(),
310            self.ignore_errors,
311            projection,
312            bytes_offset,
313            capacity,
314            self.null_values.as_ref(),
315            usize::MAX,
316            stop_at_nbytes,
317            starting_point_offset,
318        )?;
319
320        cast_columns(&mut df, &self.to_cast, false, self.ignore_errors)?;
321        Ok(df)
322    }
323
324    // The code adheres to RFC 4180 in a strict sense, unless explicitly documented otherwise.
325    // Malformed CSV is common, see e.g. the use of lazy_quotes, whitespace and comments.
326    // In case malformed CSV is detected, a warning or an error will be issued.
327    // Not all malformed CSV will be detected, as that would impact performance.
328    fn parse_csv(&mut self, bytes: &[u8]) -> PolarsResult<DataFrame> {
329        let projection = self.get_projection()?;
330
331        // An empty file with a schema should return an empty DataFrame with that schema
332        if bytes.is_empty() {
333            let mut df = if projection.len() == self.schema.len() {
334                DataFrame::empty_with_schema(self.schema.as_ref())
335            } else {
336                DataFrame::empty_with_schema(
337                    &projection
338                        .iter()
339                        .map(|&i| self.schema.get_at_index(i).unwrap())
340                        .map(|(name, dtype)| Field {
341                            name: name.clone(),
342                            dtype: dtype.clone(),
343                        })
344                        .collect::<Schema>(),
345                )
346            };
347
348            cast_columns(&mut df, &self.to_cast, false, self.ignore_errors)?;
349
350            if let Some(ref row_index) = self.row_index {
351                df.insert_column(0, Column::new_empty(row_index.name.clone(), &IDX_DTYPE))?;
352            }
353            return Ok(df);
354        }
355
356        let n_threads = self
357            .n_threads
358            .unwrap_or_else(|| RAYON.current_num_threads());
359
360        // This is chosen by benchmarking on ny city trip csv dataset.
361        // We want small enough chunks such that threads start working as soon as possible
362        // But we also want them large enough, so that we have less chunks related overhead.
363        // We minimize chunks to 16 MB to still fit L3 cache.
364        //
365        // Width-aware adjustment: For wide data (many columns), per-chunk overhead
366        // (allocating column buffers) becomes significant. Each chunk must allocate
367        // O(n_cols) buffers, so total allocation overhead is O(n_chunks * n_cols).
368        // To keep this bounded, we limit n_chunks such that n_chunks * n_cols <= threshold.
369        // With threshold ~500K, this gives:
370        //   - 100 cols: up to 5000 chunks (no practical limit)
371        //   - 1000 cols: up to 500 chunks
372        //   - 10000 cols: up to 50 chunks
373        //   - 30000 cols: up to 16 chunks
374        let n_cols = projection.len();
375        // Empirically determined to balance allocation overhead and parallelism.
376        const ALLOCATION_BUDGET: usize = 500_000;
377        let max_chunks_for_width = ALLOCATION_BUDGET / n_cols.max(1);
378        let n_parts_hint = std::cmp::min(n_threads * 16, max_chunks_for_width.max(n_threads));
379        let chunk_size = std::cmp::min(bytes.len() / n_parts_hint.max(1), 16 * 1024 * 1024);
380
381        // Use a small min chunk size to catch failures in tests.
382        #[cfg(debug_assertions)]
383        let min_chunk_size = 64;
384        #[cfg(not(debug_assertions))]
385        let min_chunk_size = 1024 * 4;
386
387        let mut chunk_size = std::cmp::max(chunk_size, min_chunk_size);
388        let mut total_bytes_offset = 0;
389
390        let results = Arc::new(Mutex::new(vec![]));
391        // We have to do this after parsing as there can be comments.
392        let total_line_count = &RelaxedCell::new_usize(0);
393
394        let counter = CountLines::new(
395            self.parse_options.quote_char,
396            self.parse_options.eol_char,
397            None,
398        );
399        let mut total_offset = 0;
400        let mut previous_total_offset = 0;
401        let check_utf8 = matches!(self.parse_options.encoding, CsvEncoding::Utf8)
402            && self.schema.iter_fields().any(|f| f.dtype().is_string());
403
404        RAYON.scope(|s| {
405            // Pass 1: identify chunks for parallel processing (line parsing).
406            loop {
407                let b = unsafe { bytes.get_unchecked(total_offset..) };
408                if b.is_empty() {
409                    break;
410                }
411                debug_assert!(
412                    total_offset == 0 || bytes[total_offset - 1] == self.parse_options.eol_char
413                );
414
415                // Count is the number of rows for the next chunk. In case of malformed CSV data,
416                // count may not be as expected.
417                let (count, position) = counter.find_next(b, &mut chunk_size);
418                debug_assert!(count == 0 || b[position] == self.parse_options.eol_char);
419
420                let (b, count) = if count == 0
421                    && unsafe {
422                        std::ptr::eq(b.as_ptr().add(b.len()), bytes.as_ptr().add(bytes.len()))
423                    } {
424                    total_offset = bytes.len();
425                    let c = if is_comment_line(bytes, self.parse_options.comment_prefix.as_ref()) {
426                        0
427                    } else {
428                        1
429                    };
430                    (b, c)
431                } else {
432                    let end = total_offset + position + 1;
433                    let b = unsafe { bytes.get_unchecked(total_offset..end) };
434
435                    previous_total_offset = total_offset;
436                    total_offset = end;
437                    (b, count)
438                };
439
440                // Pass 2: process each individual chunk in parallel (field parsing)
441                if !b.is_empty() {
442                    let results = results.clone();
443                    let projection = projection.as_ref();
444                    let slf = &(*self);
445                    s.spawn(move |_| {
446                        if check_utf8 && !super::builder::validate_utf8(b) {
447                            let mut results = results.lock().unwrap();
448                            results.push((
449                                b.as_ptr() as usize,
450                                Err(polars_err!(ComputeError: "invalid utf-8 sequence")),
451                            ));
452                            return;
453                        }
454
455                        let result = slf
456                            .read_chunk(b, projection, 0, count, Some(0), b.len())
457                            .and_then(|mut df| {
458                                // Check malformed
459                                if df.height() > count
460                                    || (df.height() < count
461                                        && slf.parse_options.comment_prefix.is_none())
462                                {
463                                    // Note: in case data is malformed, df.height() is more likely to be correct than count.
464                                    let msg = format!(
465                                        "CSV malformed: expected {} rows, \
466                                        actual {} rows, in chunk starting at \
467                                        byte offset {}, length {}",
468                                        count,
469                                        df.height(),
470                                        previous_total_offset,
471                                        b.len()
472                                    );
473                                    if slf.ignore_errors {
474                                        polars_warn!("{msg}");
475                                    } else {
476                                        polars_bail!(ComputeError: msg)
477                                    }
478                                }
479
480                                if slf.n_rows.is_some() {
481                                    total_line_count.fetch_add(df.height());
482                                }
483
484                                // We cannot use the line count as there can be comments in the lines so we must correct line counts later.
485                                if let Some(rc) = &slf.row_index {
486                                    // is first chunk
487                                    let offset = if std::ptr::eq(b.as_ptr(), bytes.as_ptr()) {
488                                        Some(rc.offset)
489                                    } else {
490                                        None
491                                    };
492
493                                    unsafe { df.with_row_index_mut(rc.name.clone(), offset) };
494                                };
495
496                                if let Some(predicate) = slf.predicate.as_ref() {
497                                    let s = predicate.evaluate_io(&df)?;
498                                    let mask = s.bool()?;
499                                    df = df.filter(mask)?;
500                                }
501                                Ok(df)
502                            });
503
504                        results.lock().unwrap().push((b.as_ptr() as usize, result));
505                    });
506
507                    // Check just after we spawned a chunk. That means we processed all data up until
508                    // row count.
509                    if self.n_rows.is_some() && total_line_count.load() > self.n_rows.unwrap() {
510                        break;
511                    }
512                }
513                total_bytes_offset += b.len();
514            }
515        });
516
517        let mut results = std::mem::take(&mut *results.lock().unwrap());
518        results.sort_unstable_by_key(|k| k.0);
519        let mut dfs = results
520            .into_iter()
521            .map(|k| k.1)
522            .collect::<PolarsResult<Vec<_>>>()?;
523
524        if let Some(rc) = &self.row_index {
525            update_row_counts2(&mut dfs, rc.offset)
526        };
527        accumulate_dataframes_vertical(dfs)
528    }
529
530    /// Read the csv into a DataFrame. The predicate can come from a lazy physical plan.
531    pub fn finish(mut self) -> PolarsResult<DataFrame> {
532        let mut reader_bytes = self.reader_bytes.take().unwrap();
533        let (body_bytes, _) = reader_bytes
534            .compressed_reader
535            .read_next_slice(&reader_bytes.leftover, usize::MAX)?;
536
537        let mut df = self.parse_csv(&body_bytes)?;
538
539        // if multi-threaded the n_rows was probabilistically determined.
540        // Let's slice to correct number of rows if possible.
541        if let Some(n_rows) = self.n_rows {
542            if n_rows < df.height() {
543                df = df.slice(0, n_rows)
544            }
545        }
546        Ok(df)
547    }
548}
549
550#[allow(clippy::too_many_arguments)]
551pub fn read_chunk(
552    bytes: &[u8],
553    parse_options: &CsvParseOptions,
554    schema: &Schema,
555    ignore_errors: bool,
556    projection: &[usize],
557    bytes_offset_thread: usize,
558    capacity: usize,
559    null_values: Option<&NullValuesCompiled>,
560    chunk_size: usize,
561    stop_at_nbytes: usize,
562    starting_point_offset: Option<usize>,
563) -> PolarsResult<DataFrame> {
564    let mut read = bytes_offset_thread;
565    // There's an off-by-one error somewhere in the reading code, where it reads
566    // one more item than the requested capacity. Given the batch sizes are
567    // approximate (sometimes they're smaller), this isn't broken, but it does
568    // mean a bunch of extra allocation and copying. So we allocate a
569    // larger-by-one buffer so the size is more likely to be accurate.
570    let mut buffers = init_builders(
571        projection,
572        capacity + 1,
573        schema,
574        parse_options.quote_char,
575        parse_options.encoding,
576        parse_options.decimal_comma,
577    )?;
578
579    debug_assert!(projection.is_sorted());
580
581    let mut last_read = usize::MAX;
582    loop {
583        if read >= stop_at_nbytes || read == last_read {
584            break;
585        }
586        let local_bytes = &bytes[read..stop_at_nbytes];
587
588        last_read = read;
589        let offset = read + starting_point_offset.unwrap();
590        read += parse_lines(
591            local_bytes,
592            parse_options,
593            offset,
594            ignore_errors,
595            null_values,
596            projection,
597            &mut buffers,
598            chunk_size,
599            schema.len(),
600            schema,
601        )?;
602    }
603
604    let columns = buffers
605        .into_iter()
606        .map(|buf| buf.into_series().map(Column::from))
607        .collect::<PolarsResult<Vec<_>>>()?;
608    Ok(unsafe { DataFrame::new_unchecked_infer_height(columns) })
609}
polars_io/csv/read/read_impl.rs

polars_io/csv/read/
read_impl.rs
