polars_lazy/scan/
parquet.rs

1use polars_buffer::Buffer;
2use polars_core::prelude::*;
3use polars_io::cloud::CloudOptions;
4use polars_io::parquet::read::ParallelStrategy;
5use polars_io::prelude::ParquetOptions;
6use polars_io::{HiveOptions, RowIndex};
7use polars_utils::pl_path::PlRefPath;
8use polars_utils::slice_enum::Slice;
9
10use crate::prelude::*;
11
12#[derive(Clone)]
13pub struct ScanArgsParquet {
14    pub n_rows: Option<usize>,
15    pub parallel: ParallelStrategy,
16    pub row_index: Option<RowIndex>,
17    pub cloud_options: Option<CloudOptions>,
18    pub hive_options: HiveOptions,
19    pub use_statistics: bool,
20    pub schema: Option<SchemaRef>,
21    pub low_memory: bool,
22    pub rechunk: bool,
23    pub cache: bool,
24    /// Expand path given via globbing rules.
25    pub glob: bool,
26    pub include_file_paths: Option<PlSmallStr>,
27    pub allow_missing_columns: bool,
28}
29
30impl Default for ScanArgsParquet {
31    fn default() -> Self {
32        Self {
33            n_rows: None,
34            parallel: Default::default(),
35            row_index: None,
36            cloud_options: None,
37            hive_options: Default::default(),
38            use_statistics: true,
39            schema: None,
40            rechunk: false,
41            low_memory: false,
42            cache: true,
43            glob: true,
44            include_file_paths: None,
45            allow_missing_columns: false,
46        }
47    }
48}
49
50#[derive(Clone)]
51struct LazyParquetReader {
52    args: ScanArgsParquet,
53    sources: ScanSources,
54}
55
56impl LazyParquetReader {
57    fn new(args: ScanArgsParquet) -> Self {
58        Self {
59            args,
60            sources: ScanSources::default(),
61        }
62    }
63}
64
65impl LazyFileListReader for LazyParquetReader {
66    /// Get the final [LazyFrame].
67    fn finish(self) -> PolarsResult<LazyFrame> {
68        let parquet_options = ParquetOptions {
69            schema: self.args.schema,
70            parallel: self.args.parallel,
71            low_memory: self.args.low_memory,
72            use_statistics: self.args.use_statistics,
73            decryption: None,
74        };
75
76        let unified_scan_args = UnifiedScanArgs {
77            schema: None,
78            cloud_options: self.args.cloud_options,
79            hive_options: self.args.hive_options,
80            rechunk: self.args.rechunk,
81            cache: self.args.cache,
82            glob: self.args.glob,
83            expand_paths: true,
84            hidden_file_prefix: None,
85            projection: None,
86            column_mapping: None,
87            default_values: None,
88            // Note: We call `with_row_index()` on the LazyFrame below
89            row_index: None,
90            pre_slice: self
91                .args
92                .n_rows
93                .map(|len| Slice::Positive { offset: 0, len }),
94            cast_columns_policy: CastColumnsPolicy::ERROR_ON_MISMATCH,
95            missing_columns_policy: if self.args.allow_missing_columns {
96                MissingColumnsPolicy::Insert
97            } else {
98                MissingColumnsPolicy::Raise
99            },
100            extra_columns_policy: ExtraColumnsPolicy::Raise,
101            include_file_paths: self.args.include_file_paths,
102            deletion_files: None,
103            table_statistics: None,
104            hive_partitions: None,
105            row_count: None,
106            source_sizes: None,
107            resolve_heavy_sources: None,
108        };
109
110        let mut lf: LazyFrame =
111            DslBuilder::scan_parquet(self.sources, parquet_options, unified_scan_args)?
112                .build()
113                .into();
114
115        // It's a bit hacky, but this row_index function updates the schema.
116        if let Some(row_index) = self.args.row_index {
117            lf = lf.with_row_index(row_index.name, Some(row_index.offset))
118        }
119
120        Ok(lf)
121    }
122
123    fn glob(&self) -> bool {
124        self.args.glob
125    }
126
127    fn finish_no_glob(self) -> PolarsResult<LazyFrame> {
128        unreachable!();
129    }
130
131    fn sources(&self) -> &ScanSources {
132        &self.sources
133    }
134
135    fn with_sources(mut self, sources: ScanSources) -> Self {
136        self.sources = sources;
137        self
138    }
139
140    fn with_n_rows(mut self, n_rows: impl Into<Option<usize>>) -> Self {
141        self.args.n_rows = n_rows.into();
142        self
143    }
144
145    fn with_row_index(mut self, row_index: impl Into<Option<RowIndex>>) -> Self {
146        self.args.row_index = row_index.into();
147        self
148    }
149
150    fn rechunk(&self) -> bool {
151        self.args.rechunk
152    }
153
154    fn with_rechunk(mut self, toggle: bool) -> Self {
155        self.args.rechunk = toggle;
156        self
157    }
158
159    fn cloud_options(&self) -> Option<&CloudOptions> {
160        self.args.cloud_options.as_ref()
161    }
162
163    fn n_rows(&self) -> Option<usize> {
164        self.args.n_rows
165    }
166
167    fn row_index(&self) -> Option<&RowIndex> {
168        self.args.row_index.as_ref()
169    }
170}
171
172impl LazyFrame {
173    /// Create a LazyFrame directly from a parquet scan.
174    pub fn scan_parquet(path: PlRefPath, args: ScanArgsParquet) -> PolarsResult<Self> {
175        Self::scan_parquet_sources(ScanSources::Paths(Buffer::from_iter([path])), args)
176    }
177
178    /// Create a LazyFrame directly from a parquet scan.
179    pub fn scan_parquet_sources(sources: ScanSources, args: ScanArgsParquet) -> PolarsResult<Self> {
180        LazyParquetReader::new(args).with_sources(sources).finish()
181    }
182
183    /// Create a LazyFrame directly from a parquet scan.
184    pub fn scan_parquet_files(
185        paths: Buffer<PlRefPath>,
186        args: ScanArgsParquet,
187    ) -> PolarsResult<Self> {
188        Self::scan_parquet_sources(ScanSources::Paths(paths), args)
189    }
190}
polars_lazy/scan/parquet.rs

polars_lazy/scan/
parquet.rs
