Skip to main content

polars_io/parquet/write/
options.rs

1use polars_arrow::datatypes::ArrowSchemaRef;
2use polars_core::prelude::CompatLevel;
3use polars_parquet::write::{
4    BrotliLevel, CompressionOptions, GzipLevel, StatisticsOptions, ZstdLevel,
5};
6#[cfg(feature = "serde")]
7use serde::{Deserialize, Serialize};
8
9use super::KeyValueMetadata;
10use crate::parquet::encryption::ParquetEncryption;
11
12#[derive(Default, Clone, Debug, PartialEq, Hash, Eq)]
13#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
14#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
15pub struct ParquetWriteOptions {
16    /// Data page compression
17    pub compression: ParquetCompression,
18    /// Compute and write column statistics.
19    pub statistics: StatisticsOptions,
20    /// If `None` will be all written to a single row group.
21    pub row_group_size: Option<usize>,
22    /// if `None` will be 1024^2 bytes
23    pub data_page_size: Option<usize>,
24    /// Custom file-level key value metadata
25    pub key_value_metadata: Option<KeyValueMetadata>,
26    pub arrow_schema: Option<ArrowSchemaRef>,
27    #[cfg_attr(feature = "serde", serde(default))]
28    pub compat_level: Option<CompatLevel>,
29    /// How to encrypt files with Parquet modular encryption
30    #[cfg_attr(feature = "serde", serde(default))]
31    pub encryption: Option<ParquetEncryption>,
32}
33
34impl ParquetWriteOptions {
35    pub fn compat_level(&self) -> CompatLevel {
36        self.compat_level.unwrap_or(CompatLevel::oldest())
37    }
38}
39
40/// The compression strategy to use for writing Parquet files.
41#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
42#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
43#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
44pub enum ParquetCompression {
45    Uncompressed,
46    Snappy,
47    Gzip(Option<GzipLevel>),
48    Brotli(Option<BrotliLevel>),
49    Zstd(Option<ZstdLevel>),
50    Lz4Raw,
51}
52
53impl Default for ParquetCompression {
54    fn default() -> Self {
55        Self::Zstd(None)
56    }
57}
58
59impl From<ParquetCompression> for CompressionOptions {
60    fn from(value: ParquetCompression) -> Self {
61        use ParquetCompression::*;
62        match value {
63            Uncompressed => CompressionOptions::Uncompressed,
64            Snappy => CompressionOptions::Snappy,
65            Gzip(level) => CompressionOptions::Gzip(level),
66            Brotli(level) => CompressionOptions::Brotli(level),
67            Lz4Raw => CompressionOptions::Lz4Raw,
68            Zstd(level) => CompressionOptions::Zstd(level),
69        }
70    }
71}