polars_io/csv/read/
options.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2use std::num::NonZeroUsize;
3use std::path::PathBuf;
4use std::sync::Arc;
5
6use polars_buffer::Buffer;
7use polars_core::datatypes::{DataType, Field};
8use polars_core::schema::{Schema, SchemaRef};
9use polars_error::PolarsResult;
10use polars_utils::pl_str::PlSmallStr;
11#[cfg(feature = "serde")]
12use serde::{Deserialize, Serialize};
13
14use crate::RowIndex;
15
16#[derive(Clone, Debug, PartialEq, Eq, Hash)]
17#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
18#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
19pub struct CsvReadOptions {
20    pub path: Option<PathBuf>,
21    // Performance related options
22    pub rechunk: bool,
23    pub n_threads: Option<usize>,
24    pub low_memory: bool,
25    // Row-wise options
26    pub n_rows: Option<usize>,
27    pub row_index: Option<RowIndex>,
28    // Column-wise options
29    pub columns: Option<Arc<[PlSmallStr]>>,
30    pub projection: Option<Arc<Vec<usize>>>,
31    pub schema: Option<SchemaRef>,
32    pub schema_overwrite: Option<SchemaRef>,
33    /// Override the names from the file. This is Python `scan_csv(new_columns=...)`
34    pub column_names_overwrite: Option<Buffer<PlSmallStr>>,
35    pub dtype_overwrite: Option<Arc<Vec<DataType>>>,
36    // CSV-specific options
37    pub parse_options: Arc<CsvParseOptions>,
38    pub has_header: bool,
39    pub chunk_size: usize,
40    /// Skip rows according to the CSV spec.
41    pub skip_rows: usize,
42    /// Skip lines according to newline char (e.g. escaping will be ignored)
43    pub skip_lines: usize,
44    pub skip_rows_after_header: usize,
45    pub infer_schema_length: Option<usize>,
46    #[cfg_attr(feature = "serde", serde(default = "nonzero_usize_max"))]
47    pub infer_schema_files: NonZeroUsize,
48    pub raise_if_empty: bool,
49    pub ignore_errors: bool,
50    pub fields_to_cast: Vec<Field>,
51}
52
53const fn nonzero_usize_max() -> NonZeroUsize {
54    NonZeroUsize::MAX
55}
56
57#[derive(Clone, Debug, PartialEq, Eq, Hash)]
58#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
59#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
60pub struct CsvParseOptions {
61    pub separator: u8,
62    pub quote_char: Option<u8>,
63    pub eol_char: u8,
64    pub encoding: CsvEncoding,
65    pub null_values: Option<NullValues>,
66    pub missing_is_null: bool,
67    pub truncate_ragged_lines: bool,
68    pub comment_prefix: Option<CommentPrefix>,
69    pub try_parse_dates: bool,
70    pub decimal_comma: bool,
71}
72
73impl Default for CsvReadOptions {
74    fn default() -> Self {
75        Self {
76            path: None,
77
78            rechunk: false,
79            n_threads: None,
80            low_memory: false,
81
82            n_rows: None,
83            row_index: None,
84
85            columns: None,
86            projection: None,
87            schema: None,
88            schema_overwrite: None,
89            column_names_overwrite: None,
90            dtype_overwrite: None,
91
92            parse_options: Default::default(),
93            has_header: true,
94            chunk_size: 1 << 18,
95            skip_rows: 0,
96            skip_lines: 0,
97            skip_rows_after_header: 0,
98            infer_schema_length: Some(100),
99            infer_schema_files: const { NonZeroUsize::new(10).unwrap() },
100            raise_if_empty: true,
101            ignore_errors: false,
102            fields_to_cast: vec![],
103        }
104    }
105}
106
107/// Options related to parsing the CSV format.
108impl Default for CsvParseOptions {
109    fn default() -> Self {
110        Self {
111            separator: b',',
112            quote_char: Some(b'"'),
113            eol_char: b'\n',
114            encoding: Default::default(),
115            null_values: None,
116            missing_is_null: true,
117            truncate_ragged_lines: false,
118            comment_prefix: None,
119            try_parse_dates: false,
120            decimal_comma: false,
121        }
122    }
123}
124
125impl CsvReadOptions {
126    pub fn get_parse_options(&self) -> Arc<CsvParseOptions> {
127        self.parse_options.clone()
128    }
129
130    pub fn with_path<P: Into<PathBuf>>(mut self, path: Option<P>) -> Self {
131        self.path = path.map(|p| p.into());
132        self
133    }
134
135    /// Whether to makes the columns contiguous in memory.
136    pub fn with_rechunk(mut self, rechunk: bool) -> Self {
137        self.rechunk = rechunk;
138        self
139    }
140
141    /// Number of threads to use for reading. Defaults to the size of the polars
142    /// thread pool.
143    pub fn with_n_threads(mut self, n_threads: Option<usize>) -> Self {
144        self.n_threads = n_threads;
145        self
146    }
147
148    /// Reduce memory consumption at the expense of performance
149    pub fn with_low_memory(mut self, low_memory: bool) -> Self {
150        self.low_memory = low_memory;
151        self
152    }
153
154    /// Limits the number of rows to read.
155    pub fn with_n_rows(mut self, n_rows: Option<usize>) -> Self {
156        self.n_rows = n_rows;
157        self
158    }
159
160    /// Adds a row index column.
161    pub fn with_row_index(mut self, row_index: Option<RowIndex>) -> Self {
162        self.row_index = row_index;
163        self
164    }
165
166    /// Which columns to select.
167    pub fn with_columns(mut self, columns: Option<Arc<[PlSmallStr]>>) -> Self {
168        self.columns = columns;
169        self
170    }
171
172    /// Which columns to select denoted by their index. The index starts from 0
173    /// (i.e. [0, 4] would select the 1st and 5th column).
174    pub fn with_projection(mut self, projection: Option<Arc<Vec<usize>>>) -> Self {
175        self.projection = projection;
176        self
177    }
178
179    /// Set the schema to use for CSV file. The length of the schema must match
180    /// the number of columns in the file. If this is [None], the schema is
181    /// inferred from the file.
182    pub fn with_schema(mut self, schema: Option<SchemaRef>) -> Self {
183        self.schema = schema;
184        self
185    }
186
187    /// Overwrites the data types in the schema by column name.
188    pub fn with_schema_overwrite(mut self, schema_overwrite: Option<SchemaRef>) -> Self {
189        self.schema_overwrite = schema_overwrite;
190        self
191    }
192
193    /// Overwrite the column names inferred from the file.
194    pub fn with_column_names_overwrite(
195        mut self,
196        column_names_overwrite: Buffer<PlSmallStr>,
197    ) -> Self {
198        self.column_names_overwrite = Some(column_names_overwrite);
199        self
200    }
201
202    /// Overwrite the dtypes in the schema in the order of the slice that's given.
203    /// This is useful if you don't know the column names beforehand
204    pub fn with_dtype_overwrite(mut self, dtype_overwrite: Option<Arc<Vec<DataType>>>) -> Self {
205        self.dtype_overwrite = dtype_overwrite;
206        self
207    }
208
209    /// Sets the CSV parsing options. See [map_parse_options][Self::map_parse_options]
210    /// for an easier way to mutate them in-place.
211    pub fn with_parse_options(mut self, parse_options: CsvParseOptions) -> Self {
212        self.parse_options = Arc::new(parse_options);
213        self
214    }
215
216    /// Sets whether the CSV file has a header row.
217    pub fn with_has_header(mut self, has_header: bool) -> Self {
218        self.has_header = has_header;
219        self
220    }
221
222    /// Sets the chunk size used by the parser. This influences performance.
223    pub fn with_chunk_size(mut self, chunk_size: usize) -> Self {
224        self.chunk_size = chunk_size;
225        self
226    }
227
228    /// Start reading after ``skip_rows`` rows. The header will be parsed at this
229    /// offset. Note that we respect CSV escaping/comments when skipping rows.
230    /// If you want to skip by newline char only, use `skip_lines`.
231    pub fn with_skip_rows(mut self, skip_rows: usize) -> Self {
232        self.skip_rows = skip_rows;
233        self
234    }
235
236    /// Start reading after `skip_lines` lines. The header will be parsed at this
237    /// offset. Note that CSV escaping will not be respected when skipping lines.
238    /// If you want to skip valid CSV rows, use ``skip_rows``.
239    pub fn with_skip_lines(mut self, skip_lines: usize) -> Self {
240        self.skip_lines = skip_lines;
241        self
242    }
243
244    /// Number of rows to skip after the header row.
245    pub fn with_skip_rows_after_header(mut self, skip_rows_after_header: usize) -> Self {
246        self.skip_rows_after_header = skip_rows_after_header;
247        self
248    }
249
250    /// Set the number of rows to use when inferring the csv schema.
251    /// The default is 100 rows.
252    /// Setting to [None] will do a full table scan, which is very slow.
253    pub fn with_infer_schema_length(mut self, infer_schema_length: Option<usize>) -> Self {
254        self.infer_schema_length = infer_schema_length;
255        self
256    }
257
258    pub fn with_infer_schema_files(mut self, infer_schema_files: NonZeroUsize) -> Self {
259        self.infer_schema_files = infer_schema_files;
260        self
261    }
262
263    /// Whether to raise an error if the frame is empty. By default an empty
264    /// DataFrame is returned.
265    pub fn with_raise_if_empty(mut self, raise_if_empty: bool) -> Self {
266        self.raise_if_empty = raise_if_empty;
267        self
268    }
269
270    /// Continue with next batch when a ParserError is encountered.
271    pub fn with_ignore_errors(mut self, ignore_errors: bool) -> Self {
272        self.ignore_errors = ignore_errors;
273        self
274    }
275
276    /// Apply a function to the parse options.
277    pub fn map_parse_options<F: Fn(CsvParseOptions) -> CsvParseOptions>(
278        mut self,
279        map_func: F,
280    ) -> Self {
281        let parse_options = Arc::unwrap_or_clone(self.parse_options);
282        self.parse_options = Arc::new(map_func(parse_options));
283        self
284    }
285}
286
287impl CsvParseOptions {
288    /// The character used to separate fields in the CSV file. This
289    /// is most often a comma ','.
290    pub fn with_separator(mut self, separator: u8) -> Self {
291        self.separator = separator;
292        self
293    }
294
295    /// Set the character used for field quoting. This is most often double
296    /// quotes '"'. Set this to [None] to disable quote parsing.
297    pub fn with_quote_char(mut self, quote_char: Option<u8>) -> Self {
298        self.quote_char = quote_char;
299        self
300    }
301
302    /// Set the character used to indicate an end-of-line (eol).
303    pub fn with_eol_char(mut self, eol_char: u8) -> Self {
304        self.eol_char = eol_char;
305        self
306    }
307
308    /// Set the encoding used by the file.
309    pub fn with_encoding(mut self, encoding: CsvEncoding) -> Self {
310        self.encoding = encoding;
311        self
312    }
313
314    /// Set values that will be interpreted as missing/null.
315    ///
316    /// Note: These values are matched before quote-parsing, so if the null values
317    /// are quoted then those quotes also need to be included here.
318    pub fn with_null_values(mut self, null_values: Option<NullValues>) -> Self {
319        self.null_values = null_values;
320        self
321    }
322
323    /// Treat missing fields as null.
324    pub fn with_missing_is_null(mut self, missing_is_null: bool) -> Self {
325        self.missing_is_null = missing_is_null;
326        self
327    }
328
329    /// Truncate lines that are longer than the schema.
330    pub fn with_truncate_ragged_lines(mut self, truncate_ragged_lines: bool) -> Self {
331        self.truncate_ragged_lines = truncate_ragged_lines;
332        self
333    }
334
335    /// Sets the comment prefix for this instance. Lines starting with this
336    /// prefix will be ignored.
337    pub fn with_comment_prefix<T: Into<CommentPrefix>>(
338        mut self,
339        comment_prefix: Option<T>,
340    ) -> Self {
341        self.comment_prefix = comment_prefix.map(Into::into);
342        self
343    }
344
345    /// Automatically try to parse dates/datetimes and time. If parsing fails,
346    /// columns remain of dtype [`DataType::String`].
347    pub fn with_try_parse_dates(mut self, try_parse_dates: bool) -> Self {
348        self.try_parse_dates = try_parse_dates;
349        self
350    }
351
352    /// Parse floats with a comma as decimal separator.
353    pub fn with_decimal_comma(mut self, decimal_comma: bool) -> Self {
354        self.decimal_comma = decimal_comma;
355        self
356    }
357}
358
359#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Hash)]
360#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
361#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
362pub enum CsvEncoding {
363    /// Utf8 encoding.
364    #[default]
365    Utf8,
366    /// Utf8 encoding and unknown bytes are replaced with �.
367    LossyUtf8,
368}
369
370#[derive(Clone, Debug, Eq, PartialEq, Hash)]
371#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
372#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
373pub enum CommentPrefix {
374    /// A single byte character that indicates the start of a comment line.
375    Single(u8),
376    /// A string that indicates the start of a comment line.
377    /// This allows for multiple characters to be used as a comment identifier.
378    Multi(PlSmallStr),
379}
380
381impl CommentPrefix {
382    /// Creates a new `CommentPrefix` for the `Single` variant.
383    pub fn new_single(prefix: u8) -> Self {
384        CommentPrefix::Single(prefix)
385    }
386
387    /// Creates a new `CommentPrefix` for the `Multi` variant.
388    pub fn new_multi(prefix: PlSmallStr) -> Self {
389        CommentPrefix::Multi(prefix)
390    }
391
392    /// Creates a new `CommentPrefix` from a `&str`.
393    pub fn new_from_str(prefix: &str) -> Self {
394        assert!(!prefix.contains("\n"));
395        if prefix.len() == 1 && prefix.chars().next().unwrap().is_ascii() {
396            let c = prefix.as_bytes()[0];
397            CommentPrefix::Single(c)
398        } else {
399            CommentPrefix::Multi(PlSmallStr::from_str(prefix))
400        }
401    }
402}
403
404impl From<&str> for CommentPrefix {
405    fn from(value: &str) -> Self {
406        Self::new_from_str(value)
407    }
408}
409
410#[derive(Clone, Debug, Eq, PartialEq, Hash)]
411#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
412#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
413pub enum NullValues {
414    /// A single value that's used for all columns
415    AllColumnsSingle(PlSmallStr),
416    /// Multiple values that are used for all columns
417    AllColumns(Vec<PlSmallStr>),
418    /// Tuples that map column names to null value of that column
419    Named(Vec<(PlSmallStr, PlSmallStr)>),
420}
421
422impl NullValues {
423    pub fn compile(self, schema: &Schema) -> PolarsResult<NullValuesCompiled> {
424        Ok(match self {
425            NullValues::AllColumnsSingle(v) => NullValuesCompiled::AllColumnsSingle(v),
426            NullValues::AllColumns(v) => NullValuesCompiled::AllColumns(v),
427            NullValues::Named(v) => {
428                let mut null_values = vec![PlSmallStr::from_static(""); schema.len()];
429                for (name, null_value) in v {
430                    let i = schema.try_index_of(&name)?;
431                    null_values[i] = null_value;
432                }
433                NullValuesCompiled::Columns(null_values)
434            },
435        })
436    }
437}
438
439#[derive(Debug, Clone)]
440pub enum NullValuesCompiled {
441    /// A single value that's used for all columns
442    AllColumnsSingle(PlSmallStr),
443    // Multiple null values that are null for all columns
444    AllColumns(Vec<PlSmallStr>),
445    /// A different null value per column, computed from `NullValues::Named`
446    Columns(Vec<PlSmallStr>),
447}
448
449impl NullValuesCompiled {
450    /// # Safety
451    ///
452    /// The caller must ensure that `index` is in bounds
453    pub(super) unsafe fn is_null(&self, field: &[u8], index: usize) -> bool {
454        use NullValuesCompiled::*;
455        match self {
456            AllColumnsSingle(v) => v.as_bytes() == field,
457            AllColumns(v) => v.iter().any(|v| v.as_bytes() == field),
458            Columns(v) => {
459                debug_assert!(index < v.len());
460                v.get_unchecked(index).as_bytes() == field
461            },
462        }
463    }
464}
polars_io/csv/read/options.rs

polars_io/csv/read/
options.rs
