polars_io/parquet/write/
writer.rs

1use std::io::Write;
2use std::sync::Mutex;
3
4use polars_buffer::Buffer;
5use polars_core::frame::chunk_df_for_writing;
6use polars_core::prelude::*;
7use polars_parquet::write::{
8    CompressionOptions, Encoding, FileWriter, StatisticsOptions, Version, WriteOptions,
9    get_dtype_encoding, to_parquet_schema,
10};
11
12use super::batched_writer::BatchedWriter;
13use super::options::ParquetCompression;
14use super::{KeyValueMetadata, ParquetWriteOptions};
15use crate::parquet::encryption::ParquetEncryption;
16use crate::shared::schema_to_arrow_checked;
17
18impl ParquetWriteOptions {
19    pub fn to_writer<F>(&self, f: F) -> ParquetWriter<F>
20    where
21        F: Write,
22    {
23        ParquetWriter::new(f)
24            .with_compression(self.compression)
25            .with_statistics(self.statistics)
26            .with_row_group_size(self.row_group_size)
27            .with_data_page_size(self.data_page_size)
28            .with_key_value_metadata(self.key_value_metadata.clone())
29            .with_encryption(self.encryption.clone())
30    }
31}
32
33/// Write a DataFrame to Parquet format.
34#[must_use]
35pub struct ParquetWriter<W> {
36    writer: W,
37    /// Data page compression
38    compression: CompressionOptions,
39    /// Compute and write column statistics.
40    statistics: StatisticsOptions,
41    /// if `None` will be 512^2 rows
42    row_group_size: Option<usize>,
43    /// if `None` will be 1024^2 bytes
44    data_page_size: Option<usize>,
45    /// Serialize columns in parallel
46    parallel: bool,
47    /// Custom file-level key value metadata
48    key_value_metadata: Option<KeyValueMetadata>,
49    /// Context info for the Parquet file being written.
50    context_info: Option<PlHashMap<String, String>>,
51    /// How to encrypt the file with Parquet modular encryption
52    encryption: Option<ParquetEncryption>,
53}
54
55impl<W> ParquetWriter<W>
56where
57    W: Write,
58{
59    /// Create a new writer
60    pub fn new(writer: W) -> Self
61    where
62        W: Write,
63    {
64        ParquetWriter {
65            writer,
66            compression: ParquetCompression::default().into(),
67            statistics: StatisticsOptions::default(),
68            row_group_size: None,
69            data_page_size: None,
70            parallel: true,
71            key_value_metadata: None,
72            context_info: None,
73            encryption: None,
74        }
75    }
76
77    /// Set the compression used. Defaults to `Zstd`.
78    ///
79    /// The default compression `Zstd` has very good performance, but may not yet been supported
80    /// by older readers. If you want more compatibility guarantees, consider using `Snappy`.
81    pub fn with_compression(mut self, compression: ParquetCompression) -> Self {
82        self.compression = compression.into();
83        self
84    }
85
86    /// Compute and write statistic
87    pub fn with_statistics(mut self, statistics: StatisticsOptions) -> Self {
88        self.statistics = statistics;
89        self
90    }
91
92    /// Set the row group size (in number of rows) during writing. This can reduce memory pressure and improve
93    /// writing performance.
94    pub fn with_row_group_size(mut self, size: Option<usize>) -> Self {
95        self.row_group_size = size;
96        self
97    }
98
99    /// Sets the maximum bytes size of a data page. If `None` will be 1024^2 bytes.
100    pub fn with_data_page_size(mut self, limit: Option<usize>) -> Self {
101        self.data_page_size = limit;
102        self
103    }
104
105    /// Serialize columns in parallel
106    pub fn set_parallel(mut self, parallel: bool) -> Self {
107        self.parallel = parallel;
108        self
109    }
110
111    /// Set custom file-level key value metadata for the Parquet file
112    pub fn with_key_value_metadata(mut self, key_value_metadata: Option<KeyValueMetadata>) -> Self {
113        self.key_value_metadata = key_value_metadata;
114        self
115    }
116
117    /// Set context information for the writer
118    pub fn with_context_info(mut self, context_info: Option<PlHashMap<String, String>>) -> Self {
119        self.context_info = context_info;
120        self
121    }
122
123    /// Set how to encrypt the file with Parquet modular encryption.
124    pub fn with_encryption(mut self, encryption: Option<ParquetEncryption>) -> Self {
125        self.encryption = encryption;
126        self
127    }
128
129    pub fn batched(self, schema: &Schema) -> PolarsResult<BatchedWriter<W>> {
130        let schema = schema_to_arrow_checked(schema, CompatLevel::newest(), "parquet")?;
131        let parquet_schema = to_parquet_schema(&schema)?;
132        let encodings = get_encodings(&schema);
133        let options = self.materialize_options();
134        let encryption_properties = self
135            .encryption
136            .as_ref()
137            .map(ParquetEncryption::file_properties)
138            .transpose()?;
139        let writer = Mutex::new(FileWriter::try_new(
140            self.writer,
141            schema,
142            options,
143            encryption_properties,
144        )?);
145
146        Ok(BatchedWriter {
147            writer,
148            parquet_schema,
149            encodings,
150            options,
151            parallel: self.parallel,
152            key_value_metadata: self.key_value_metadata,
153        })
154    }
155
156    fn materialize_options(&self) -> WriteOptions {
157        WriteOptions {
158            statistics: self.statistics,
159            compression: self.compression,
160            version: Version::V1,
161            data_page_size: self.data_page_size,
162        }
163    }
164
165    /// Write the given DataFrame in the writer `W`.
166    /// Returns the total size of the file.
167    pub fn finish(self, df: &mut DataFrame) -> PolarsResult<u64> {
168        let chunked_df = chunk_df_for_writing(df, self.row_group_size.unwrap_or(512 * 512))?;
169        let mut batched = self.batched(chunked_df.schema())?;
170        batched.write_batch(&chunked_df)?;
171        batched.finish()
172    }
173}
174
175pub fn get_encodings(schema: &ArrowSchema) -> Buffer<Vec<Encoding>> {
176    schema
177        .iter_values()
178        .map(|f| get_dtype_encoding(&f.dtype))
179        .collect()
180}
polars_io/parquet/write/writer.rs

polars_io/parquet/write/
writer.rs
