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

polars_lazy/scan/
csv.rs
