laminar_connectors/schema/parquet/encoder/
mod.rs

1//! Parquet format encoder implementing [`FormatEncoder`].
2//!
3//! Encodes Arrow `RecordBatch`es into Parquet file bytes using
4//! `ArrowWriter<Vec<u8>>` with configurable compression and row-group sizing.
5
6use arrow_array::RecordBatch;
7use arrow_schema::SchemaRef;
8use parquet::arrow::ArrowWriter;
9use parquet::basic::Compression;
10use parquet::file::properties::WriterProperties;
11
12use crate::schema::error::{SchemaError, SchemaResult};
13use crate::schema::traits::FormatEncoder;
14
15/// Configuration for the Parquet encoder.
16#[derive(Debug, Clone)]
17pub struct ParquetEncoderConfig {
18    /// Compression codec (default: Snappy).
19    pub compression: Compression,
20
21    /// Parquet writer version (1 or 2, default: 2).
22    pub writer_version: i32,
23
24    /// Maximum rows per row group (default: `1_000_000`).
25    pub max_row_group_size: usize,
26
27    /// Whether to write column statistics (default: true).
28    pub write_statistics: bool,
29}
30
31impl Default for ParquetEncoderConfig {
32    fn default() -> Self {
33        Self {
34            compression: Compression::SNAPPY,
35            writer_version: 2,
36            max_row_group_size: 1_000_000,
37            write_statistics: true,
38        }
39    }
40}
41
42impl ParquetEncoderConfig {
43    /// Sets the compression codec.
44    #[must_use]
45    pub fn with_compression(mut self, compression: Compression) -> Self {
46        self.compression = compression;
47        self
48    }
49
50    /// Sets the writer version.
51    #[must_use]
52    pub fn with_writer_version(mut self, version: i32) -> Self {
53        self.writer_version = version;
54        self
55    }
56
57    /// Sets the maximum rows per row group.
58    #[must_use]
59    pub fn with_max_row_group_size(mut self, size: usize) -> Self {
60        self.max_row_group_size = size;
61        self
62    }
63
64    /// Enables or disables column statistics.
65    #[must_use]
66    pub fn with_statistics(mut self, enabled: bool) -> Self {
67        self.write_statistics = enabled;
68        self
69    }
70}
71
72/// Encodes Arrow `RecordBatch`es into Parquet file bytes.
73///
74/// Each call to `encode_batch` produces a single Parquet file (as `Vec<u8>`)
75/// containing the entire batch. The file includes footer metadata and is
76/// self-contained.
77#[derive(Debug)]
78pub struct ParquetEncoder {
79    schema: SchemaRef,
80    config: ParquetEncoderConfig,
81}
82
83impl ParquetEncoder {
84    /// Creates a new Parquet encoder for the given schema.
85    #[must_use]
86    pub fn new(schema: SchemaRef) -> Self {
87        Self::with_config(schema, ParquetEncoderConfig::default())
88    }
89
90    /// Creates a new Parquet encoder with custom configuration.
91    #[must_use]
92    pub fn with_config(schema: SchemaRef, config: ParquetEncoderConfig) -> Self {
93        Self { schema, config }
94    }
95}
96
97impl FormatEncoder for ParquetEncoder {
98    fn input_schema(&self) -> SchemaRef {
99        self.schema.clone()
100    }
101
102    fn encode_batch(&self, batch: &RecordBatch) -> SchemaResult<Vec<Vec<u8>>> {
103        if batch.num_rows() == 0 {
104            return Ok(vec![]);
105        }
106
107        let mut props_builder = WriterProperties::builder()
108            .set_compression(self.config.compression)
109            .set_max_row_group_row_count(Some(self.config.max_row_group_size));
110
111        if !self.config.write_statistics {
112            props_builder = props_builder
113                .set_statistics_enabled(parquet::file::properties::EnabledStatistics::None);
114        }
115
116        let props = props_builder.build();
117
118        let mut buf = Vec::new();
119        let mut writer = ArrowWriter::try_new(&mut buf, self.schema.clone(), Some(props))
120            .map_err(|e| SchemaError::DecodeError(format!("Parquet writer init: {e}")))?;
121
122        writer
123            .write(batch)
124            .map_err(|e| SchemaError::DecodeError(format!("Parquet write error: {e}")))?;
125
126        writer
127            .close()
128            .map_err(|e| SchemaError::DecodeError(format!("Parquet close error: {e}")))?;
129
130        // Single Parquet file containing the full batch.
131        Ok(vec![buf])
132    }
133
134    fn format_name(&self) -> &'static str {
135        "parquet"
136    }
137}
138
139#[cfg(test)]
140mod tests;
laminar_connectors/schema/parquet/encoder/mod.rs

laminar_connectors/schema/parquet/encoder/
mod.rs
