polars_lazy/scan/
csv.rs

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

polars_lazy/scan/
csv.rs
