polars_io/csv/read/
options.rs

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

polars_io/csv/read/
options.rs
