polars_lazy/scan/
csv.rs

1#[cfg(feature = "csv")]
2use std::num::NonZeroUsize;
3
4#[cfg(feature = "csv")]
5use polars_buffer::Buffer;
6use polars_core::prelude::*;
7use polars_io::cloud::CloudOptions;
8use polars_io::csv::read::{
9    CommentPrefix, CsvEncoding, CsvParseOptions, CsvReadOptions, NullValues,
10};
11use polars_io::path_utils::expand_paths;
12use polars_io::{HiveOptions, RowIndex};
13use polars_utils::mmap::MMapSemaphore;
14use polars_utils::pl_path::PlRefPath;
15use polars_utils::slice_enum::Slice;
16
17use crate::prelude::*;
18
19#[derive(Clone)]
20#[cfg(feature = "csv")]
21pub struct LazyCsvReader {
22    sources: ScanSources,
23    glob: bool,
24    cache: bool,
25    read_options: CsvReadOptions,
26    cloud_options: Option<CloudOptions>,
27    include_file_paths: Option<PlSmallStr>,
28    missing_columns_policy: Option<MissingColumnsPolicy>,
29}
30
31#[cfg(feature = "csv")]
32impl LazyCsvReader {
33    /// Re-export to shorten code.
34    pub fn map_parse_options<F: Fn(CsvParseOptions) -> CsvParseOptions>(
35        mut self,
36        map_func: F,
37    ) -> Self {
38        self.read_options = self.read_options.map_parse_options(map_func);
39        self
40    }
41
42    pub fn new_paths(paths: Buffer<PlRefPath>) -> Self {
43        Self::new_with_sources(ScanSources::Paths(paths))
44    }
45
46    pub fn new_with_sources(sources: ScanSources) -> Self {
47        LazyCsvReader {
48            sources,
49            glob: true,
50            cache: true,
51            read_options: Default::default(),
52            cloud_options: Default::default(),
53            include_file_paths: None,
54            missing_columns_policy: None,
55        }
56    }
57
58    pub fn new(path: PlRefPath) -> Self {
59        Self::new_with_sources(ScanSources::Paths(Buffer::from_iter([path])))
60    }
61
62    /// Skip this number of rows after the header location.
63    #[must_use]
64    pub fn with_skip_rows_after_header(mut self, offset: usize) -> Self {
65        self.read_options.skip_rows_after_header = offset;
66        self
67    }
68
69    /// Add a row index column.
70    #[must_use]
71    pub fn with_row_index(mut self, row_index: Option<RowIndex>) -> Self {
72        self.read_options.row_index = row_index;
73        self
74    }
75
76    /// Try to stop parsing when `n` rows are parsed. During multithreaded parsing the upper bound `n` cannot
77    /// be guaranteed.
78    #[must_use]
79    pub fn with_n_rows(mut self, num_rows: Option<usize>) -> Self {
80        self.read_options.n_rows = num_rows;
81        self
82    }
83
84    /// Sets the number of threads used for CSV parsing.
85    #[must_use]
86    pub fn with_n_threads(mut self, n_threads: Option<usize>) -> Self {
87        self.read_options.n_threads = n_threads;
88        self
89    }
90
91    /// Set the number of rows to use when inferring the csv schema.
92    /// The default is 100 rows.
93    /// Setting to [None] will do a full table scan, which is very slow.
94    #[must_use]
95    pub fn with_infer_schema_length(mut self, num_rows: Option<usize>) -> Self {
96        self.read_options.infer_schema_length = num_rows;
97        self
98    }
99
100    #[must_use]
101    pub fn with_infer_schema_files(mut self, infer_schema_files: NonZeroUsize) -> Self {
102        self.read_options.infer_schema_files = infer_schema_files;
103        self
104    }
105
106    /// Continue with next batch when a ParserError is encountered.
107    #[must_use]
108    pub fn with_ignore_errors(mut self, ignore: bool) -> Self {
109        self.read_options.ignore_errors = ignore;
110        self
111    }
112
113    /// Set the CSV file's schema
114    #[must_use]
115    pub fn with_schema(mut self, schema: Option<SchemaRef>) -> Self {
116        self.read_options.schema = schema;
117        self
118    }
119
120    /// Skip the first `n` rows during parsing. The header will be parsed at row `n`.
121    /// Note that by row we mean valid CSV, encoding and comments are respected.
122    #[must_use]
123    pub fn with_skip_rows(mut self, skip_rows: usize) -> Self {
124        self.read_options.skip_rows = skip_rows;
125        self
126    }
127
128    /// Skip the first `n` lines during parsing. The header will be parsed at line `n`.
129    /// We don't respect CSV escaping when skipping lines.
130    #[must_use]
131    pub fn with_skip_lines(mut self, skip_lines: usize) -> Self {
132        self.read_options.skip_lines = skip_lines;
133        self
134    }
135
136    #[must_use]
137    pub fn with_column_names_overwrite(
138        mut self,
139        column_names_overwrite: Buffer<PlSmallStr>,
140    ) -> Self {
141        self.read_options.column_names_overwrite = Some(column_names_overwrite);
142        self
143    }
144
145    /// Overwrite the schema with the dtypes in this given Schema. The given schema may be a subset
146    /// of the total schema.
147    #[must_use]
148    pub fn with_dtype_overwrite(mut self, schema: Option<SchemaRef>) -> Self {
149        self.read_options.schema_overwrite = schema;
150        self
151    }
152
153    /// Overwrite dtypes by position.
154    #[must_use]
155    pub fn with_dtype_overwrite_by_position(mut self, dtypes: Option<Arc<Vec<DataType>>>) -> Self {
156        self.read_options.dtype_overwrite = dtypes;
157        self
158    }
159
160    /// Set whether the CSV file has headers
161    #[must_use]
162    pub fn with_has_header(mut self, has_header: bool) -> Self {
163        self.read_options.has_header = has_header;
164        self
165    }
166
167    /// Sets the chunk size used by the parser. This influences performance.
168    /// This can be used as a way to reduce memory usage during the parsing at the cost of performance.
169    pub fn with_chunk_size(mut self, chunk_size: usize) -> Self {
170        self.read_options.chunk_size = chunk_size;
171        self
172    }
173
174    /// Set the CSV file's column separator as a byte character
175    #[must_use]
176    pub fn with_separator(self, separator: u8) -> Self {
177        self.map_parse_options(|opts| opts.with_separator(separator))
178    }
179
180    /// Set the comment prefix for this instance. Lines starting with this prefix will be ignored.
181    #[must_use]
182    pub fn with_comment_prefix(self, comment_prefix: Option<PlSmallStr>) -> Self {
183        self.map_parse_options(|opts| {
184            opts.with_comment_prefix(comment_prefix.clone().map(|s| {
185                if s.len() == 1 && s.chars().next().unwrap().is_ascii() {
186                    CommentPrefix::Single(s.as_bytes()[0])
187                } else {
188                    CommentPrefix::Multi(s)
189                }
190            }))
191        })
192    }
193
194    /// Set the `char` used as quote char. The default is `b'"'`. If set to [`None`] quoting is disabled.
195    #[must_use]
196    pub fn with_quote_char(self, quote_char: Option<u8>) -> Self {
197        self.map_parse_options(|opts| opts.with_quote_char(quote_char))
198    }
199
200    /// Set the `char` used as end of line. The default is `b'\n'`.
201    #[must_use]
202    pub fn with_eol_char(self, eol_char: u8) -> Self {
203        self.map_parse_options(|opts| opts.with_eol_char(eol_char))
204    }
205
206    /// Set values that will be interpreted as missing/ null.
207    #[must_use]
208    pub fn with_null_values(self, null_values: Option<NullValues>) -> Self {
209        self.map_parse_options(|opts| opts.with_null_values(null_values.clone()))
210    }
211
212    /// Treat missing fields as null.
213    pub fn with_missing_is_null(self, missing_is_null: bool) -> Self {
214        self.map_parse_options(|opts| opts.with_missing_is_null(missing_is_null))
215    }
216
217    /// Cache the DataFrame after reading.
218    #[must_use]
219    pub fn with_cache(mut self, cache: bool) -> Self {
220        self.cache = cache;
221        self
222    }
223
224    /// Reduce memory usage at the expense of performance
225    #[must_use]
226    pub fn with_low_memory(mut self, low_memory: bool) -> Self {
227        self.read_options.low_memory = low_memory;
228        self
229    }
230
231    /// Set  [`CsvEncoding`]
232    #[must_use]
233    pub fn with_encoding(self, encoding: CsvEncoding) -> Self {
234        self.map_parse_options(|opts| opts.with_encoding(encoding))
235    }
236
237    /// Automatically try to parse dates/datetimes and time.
238    /// If parsing fails, columns remain of dtype [`DataType::String`].
239    #[cfg(feature = "temporal")]
240    pub fn with_try_parse_dates(self, try_parse_dates: bool) -> Self {
241        self.map_parse_options(|opts| opts.with_try_parse_dates(try_parse_dates))
242    }
243
244    /// Raise an error if CSV is empty (otherwise return an empty frame)
245    #[must_use]
246    pub fn with_raise_if_empty(mut self, raise_if_empty: bool) -> Self {
247        self.read_options.raise_if_empty = raise_if_empty;
248        self
249    }
250
251    /// Truncate lines that are longer than the schema.
252    #[must_use]
253    pub fn with_truncate_ragged_lines(self, truncate_ragged_lines: bool) -> Self {
254        self.map_parse_options(|opts| opts.with_truncate_ragged_lines(truncate_ragged_lines))
255    }
256
257    #[must_use]
258    pub fn with_decimal_comma(self, decimal_comma: bool) -> Self {
259        self.map_parse_options(|opts| opts.with_decimal_comma(decimal_comma))
260    }
261
262    #[must_use]
263    /// Expand path given via globbing rules.
264    pub fn with_glob(mut self, toggle: bool) -> Self {
265        self.glob = toggle;
266        self
267    }
268
269    pub fn with_cloud_options(mut self, cloud_options: Option<CloudOptions>) -> Self {
270        self.cloud_options = cloud_options;
271        self
272    }
273
274    /// Modify a schema before we run the lazy scanning.
275    ///
276    /// Important! Run this function latest in the builder!
277    pub fn with_schema_modify<F>(mut self, f: F) -> PolarsResult<Self>
278    where
279        F: Fn(Schema) -> PolarsResult<Schema>,
280    {
281        const ASSUMED_COMPRESSION_RATIO: usize = 4;
282        let n_threads = self.read_options.n_threads;
283
284        let infer_schema = |bytes: Buffer<u8>| {
285            use polars_io::prelude::streaming::read_until_start_and_infer_schema;
286            use polars_io::utils::compression::ByteSourceReader;
287
288            let bytes_len = bytes.len();
289            let mut reader = ByteSourceReader::from_memory(bytes)?;
290            let decompressed_size_hint = Some(
291                bytes_len
292                    * reader
293                        .compression()
294                        .map_or(1, |_| ASSUMED_COMPRESSION_RATIO),
295            );
296
297            let (inferred_schema, _) = read_until_start_and_infer_schema(
298                &self.read_options,
299                None,
300                decompressed_size_hint,
301                None,
302                &mut reader,
303            )?;
304
305            PolarsResult::Ok(inferred_schema)
306        };
307
308        let schema = match self.sources.clone() {
309            ScanSources::Paths(paths) => {
310                // TODO: Path expansion should happen when converting to the IR
311                // https://github.com/pola-rs/polars/issues/17634
312
313                use polars_core::runtime::ASYNC;
314
315                let paths = ASYNC.block_on(expand_paths(
316                    &paths[..],
317                    self.glob(),
318                    &[], // hidden_file_prefix
319                    &mut self.cloud_options,
320                ))?;
321
322                let Some(path) = paths.first() else {
323                    polars_bail!(ComputeError: "no paths specified for this reader");
324                };
325
326                let file = polars_utils::io::open_file(path.as_std_path())?;
327                let mmap = MMapSemaphore::new_from_file(&file)?;
328                infer_schema(Buffer::from_owner(mmap))?
329            },
330            ScanSources::Files(files) => {
331                let Some(file) = files.first() else {
332                    polars_bail!(ComputeError: "no buffers specified for this reader");
333                };
334
335                let mmap = MMapSemaphore::new_from_file(file)?;
336                infer_schema(Buffer::from_owner(mmap))?
337            },
338            ScanSources::Buffers(buffers) => {
339                let Some(buffer) = buffers.first() else {
340                    polars_bail!(ComputeError: "no buffers specified for this reader");
341                };
342
343                infer_schema(buffer.clone())?
344            },
345        };
346
347        self.read_options.n_threads = n_threads;
348        let mut schema = f(schema)?;
349
350        // the dtypes set may be for the new names, so update again
351        if let Some(overwrite_schema) = &self.read_options.schema_overwrite {
352            for (name, dtype) in overwrite_schema.iter() {
353                schema.with_column(name.clone(), dtype.clone());
354            }
355        }
356
357        Ok(self.with_schema(Some(Arc::new(schema))))
358    }
359
360    pub fn with_include_file_paths(mut self, include_file_paths: Option<PlSmallStr>) -> Self {
361        self.include_file_paths = include_file_paths;
362        self
363    }
364
365    #[must_use]
366    pub fn with_missing_columns_policy(mut self, policy: Option<MissingColumnsPolicy>) -> Self {
367        self.missing_columns_policy = policy;
368        self
369    }
370}
371
372impl LazyFileListReader for LazyCsvReader {
373    /// Get the final [LazyFrame].
374    fn finish(self) -> PolarsResult<LazyFrame> {
375        let rechunk = self.rechunk();
376        let row_index = self.row_index().cloned();
377        let pre_slice = self.n_rows().map(|len| Slice::Positive { offset: 0, len });
378
379        let missing_columns_policy = self.missing_columns_policy.unwrap_or_default();
380
381        let lf: LazyFrame = DslBuilder::scan_csv(
382            self.sources,
383            self.read_options,
384            UnifiedScanArgs {
385                schema: None,
386                cloud_options: self.cloud_options,
387                hive_options: HiveOptions::new_disabled(),
388                rechunk,
389                cache: self.cache,
390                glob: self.glob,
391                hidden_file_prefix: None,
392                projection: None,
393                column_mapping: None,
394                default_values: None,
395                row_index,
396                pre_slice,
397                cast_columns_policy: CastColumnsPolicy::ERROR_ON_MISMATCH,
398                missing_columns_policy,
399                extra_columns_policy: ExtraColumnsPolicy::Raise,
400                include_file_paths: self.include_file_paths,
401                deletion_files: None,
402                table_statistics: None,
403                row_count: None,
404            },
405        )?
406        .build()
407        .into();
408        Ok(lf)
409    }
410
411    fn finish_no_glob(self) -> PolarsResult<LazyFrame> {
412        unreachable!();
413    }
414
415    fn glob(&self) -> bool {
416        self.glob
417    }
418
419    fn sources(&self) -> &ScanSources {
420        &self.sources
421    }
422
423    fn with_sources(mut self, sources: ScanSources) -> Self {
424        self.sources = sources;
425        self
426    }
427
428    fn with_n_rows(mut self, n_rows: impl Into<Option<usize>>) -> Self {
429        self.read_options.n_rows = n_rows.into();
430        self
431    }
432
433    fn with_row_index(mut self, row_index: impl Into<Option<RowIndex>>) -> Self {
434        self.read_options.row_index = row_index.into();
435        self
436    }
437
438    fn rechunk(&self) -> bool {
439        self.read_options.rechunk
440    }
441
442    /// Rechunk the memory to contiguous chunks when parsing is done.
443    fn with_rechunk(mut self, rechunk: bool) -> Self {
444        self.read_options.rechunk = rechunk;
445        self
446    }
447
448    /// Try to stop parsing when `n` rows are parsed. During multithreaded parsing the upper bound `n` cannot
449    /// be guaranteed.
450    fn n_rows(&self) -> Option<usize> {
451        self.read_options.n_rows
452    }
453
454    /// Return the row index settings.
455    fn row_index(&self) -> Option<&RowIndex> {
456        self.read_options.row_index.as_ref()
457    }
458
459    fn concat_impl(&self, lfs: Vec<LazyFrame>) -> PolarsResult<LazyFrame> {
460        // set to false, as the csv parser has full thread utilization
461        let args = UnionArgs {
462            rechunk: self.rechunk(),
463            parallel: false,
464            to_supertypes: false,
465            from_partitioned_ds: true,
466            ..Default::default()
467        };
468        concat_impl(&lfs, args)
469    }
470
471    /// [CloudOptions] used to list files.
472    fn cloud_options(&self) -> Option<&CloudOptions> {
473        self.cloud_options.as_ref()
474    }
475}
polars_lazy/scan/csv.rs

polars_lazy/scan/
csv.rs
