Skip to main content

parquet/file/
writer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! [`SerializedFileWriter`]: Low level Parquet writer API
6
7use crate::bloom_filter::Sbbf;
8use crate::file::metadata::page_index::PageIndex;
9use crate::file::metadata::thrift::PageHeader;
10use crate::file::page_index::column_index::ColumnIndexMetaData;
11use crate::file::page_index::offset_index::OffsetIndexMetaData;
12use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift};
13#[cfg(feature = "arrow")]
14use bytes::Bytes;
15use std::fmt::Debug;
16use std::io::{BufWriter, IoSlice, Read};
17use std::{io::Write, sync::Arc};
18
19use crate::column::page_encryption::PageEncryptor;
20use crate::column::writer::{ColumnCloseResult, ColumnWriterImpl, get_typed_column_writer_mut};
21use crate::column::{
22    page::{CompressedPage, PageWriteSpec, PageWriter},
23    writer::{ColumnWriter, get_column_writer},
24};
25use crate::data_type::DataType;
26#[cfg(feature = "encryption")]
27use crate::encryption::encrypt::{
28    FileEncryptionProperties, FileEncryptor, get_column_crypto_metadata,
29};
30use crate::errors::{ParquetError, Result};
31#[cfg(feature = "encryption")]
32use crate::file::PARQUET_MAGIC_ENCR_FOOTER;
33use crate::file::properties::{BloomFilterPosition, WriterPropertiesPtr};
34use crate::file::reader::ChunkReader;
35use crate::file::{PARQUET_MAGIC, metadata::*};
36use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, SchemaDescriptor, TypePtr};
37
38/// A wrapper around a [`Write`] that keeps track of the number
39/// of bytes that have been written. The given [`Write`] is wrapped
40/// with a [`BufWriter`] to optimize writing performance.
41pub struct TrackedWrite<W: Write> {
42    inner: BufWriter<W>,
43    bytes_written: usize,
44}
45
46impl<W: Write> TrackedWrite<W> {
47    /// Create a new [`TrackedWrite`] from a [`Write`]
48    pub fn new(inner: W) -> Self {
49        let buf_write = BufWriter::new(inner);
50        Self {
51            inner: buf_write,
52            bytes_written: 0,
53        }
54    }
55
56    /// Returns the number of bytes written to this instance
57    pub fn bytes_written(&self) -> usize {
58        self.bytes_written
59    }
60
61    /// Returns a reference to the underlying writer.
62    pub fn inner(&self) -> &W {
63        self.inner.get_ref()
64    }
65
66    /// Returns a mutable reference to the underlying writer.
67    ///
68    /// It is inadvisable to directly write to the underlying writer, doing so
69    /// will likely result in data corruption
70    pub fn inner_mut(&mut self) -> &mut W {
71        self.inner.get_mut()
72    }
73
74    /// Returns the underlying writer.
75    pub fn into_inner(self) -> Result<W> {
76        self.inner.into_inner().map_err(|err| {
77            ParquetError::General(format!("fail to get inner writer: {:?}", err.to_string()))
78        })
79    }
80}
81
82impl<W: Write> Write for TrackedWrite<W> {
83    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
84        let bytes = self.inner.write(buf)?;
85        self.bytes_written += bytes;
86        Ok(bytes)
87    }
88
89    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> std::io::Result<usize> {
90        let bytes = self.inner.write_vectored(bufs)?;
91        self.bytes_written += bytes;
92        Ok(bytes)
93    }
94
95    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
96        self.inner.write_all(buf)?;
97        self.bytes_written += buf.len();
98
99        Ok(())
100    }
101
102    fn flush(&mut self) -> std::io::Result<()> {
103        self.inner.flush()
104    }
105}
106
107/// Callback invoked on closing a column chunk
108pub type OnCloseColumnChunk<'a> = Box<dyn FnOnce(ColumnCloseResult) -> Result<()> + 'a>;
109
110/// Callback invoked on closing a row group, arguments are:
111///
112/// - the row group metadata
113/// - the column index for each column chunk
114/// - the offset index for each column chunk
115pub type OnCloseRowGroup<'a, W> = Box<
116    dyn FnOnce(
117            &'a mut TrackedWrite<W>,
118            RowGroupMetaData,
119            Vec<Option<Sbbf>>,
120            Vec<Option<ColumnIndexMetaData>>,
121            Vec<Option<OffsetIndexMetaData>>,
122        ) -> Result<()>
123        + 'a
124        + Send,
125>;
126
127// ----------------------------------------------------------------------
128// Serialized impl for file & row group writers
129
130/// Parquet file writer API.
131///
132/// This is a low level API for writing Parquet files directly, and handles
133/// tracking the location of file structures such as row groups and column
134/// chunks, and writing the metadata and file footer.
135///
136/// Data is written to row groups using  [`SerializedRowGroupWriter`] and
137/// columns using [`SerializedColumnWriter`]. The `SerializedFileWriter` tracks
138/// where all the data is written, and assembles the final file metadata.
139///
140/// The main workflow should be as following:
141/// - Create file writer, this will open a new file and potentially write some metadata.
142/// - Request a new row group writer by calling `next_row_group`.
143/// - Once finished writing row group, close row group writer by calling `close`
144/// - Write subsequent row groups, if necessary.
145/// - After all row groups have been written, close the file writer using `close` method.
146pub struct SerializedFileWriter<W: Write> {
147    buf: TrackedWrite<W>,
148    descr: SchemaDescPtr,
149    props: WriterPropertiesPtr,
150    row_groups: Vec<RowGroupMetaData>,
151    bloom_filters: Vec<Vec<Option<Sbbf>>>,
152    column_indexes: Vec<Vec<Option<ColumnIndexMetaData>>>,
153    offset_indexes: Vec<Vec<Option<OffsetIndexMetaData>>>,
154    row_group_index: usize,
155    // kv_metadatas will be appended to `props` when `write_metadata`
156    kv_metadatas: Vec<KeyValue>,
157    finished: bool,
158    #[cfg(feature = "encryption")]
159    file_encryptor: Option<Arc<FileEncryptor>>,
160}
161
162impl<W: Write> Debug for SerializedFileWriter<W> {
163    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
164        // implement Debug so this can be used with #[derive(Debug)]
165        // in client code rather than actually listing all the fields
166        f.debug_struct("SerializedFileWriter")
167            .field("descr", &self.descr)
168            .field("row_group_index", &self.row_group_index)
169            .field("kv_metadatas", &self.kv_metadatas)
170            .finish_non_exhaustive()
171    }
172}
173
174impl<W: Write + Send> SerializedFileWriter<W> {
175    /// Creates new file writer.
176    pub fn new(buf: W, schema: TypePtr, properties: WriterPropertiesPtr) -> Result<Self> {
177        let mut buf = TrackedWrite::new(buf);
178
179        let schema_descriptor = SchemaDescriptor::new(schema.clone());
180
181        #[cfg(feature = "encryption")]
182        let file_encryptor = Self::get_file_encryptor(&properties, &schema_descriptor)?;
183
184        Self::start_file(&properties, &mut buf)?;
185        Ok(Self {
186            buf,
187            descr: Arc::new(schema_descriptor),
188            props: properties,
189            row_groups: vec![],
190            bloom_filters: vec![],
191            column_indexes: Vec::new(),
192            offset_indexes: Vec::new(),
193            row_group_index: 0,
194            kv_metadatas: Vec::new(),
195            finished: false,
196            #[cfg(feature = "encryption")]
197            file_encryptor,
198        })
199    }
200
201    #[cfg(feature = "encryption")]
202    fn get_file_encryptor(
203        properties: &WriterPropertiesPtr,
204        schema_descriptor: &SchemaDescriptor,
205    ) -> Result<Option<Arc<FileEncryptor>>> {
206        if let Some(file_encryption_properties) = properties.file_encryption_properties() {
207            file_encryption_properties.validate_encrypted_column_names(schema_descriptor)?;
208
209            Ok(Some(Arc::new(FileEncryptor::new(Arc::clone(
210                file_encryption_properties,
211            ))?)))
212        } else {
213            Ok(None)
214        }
215    }
216
217    /// Creates new row group from this file writer.
218    ///
219    /// Note: Parquet files are limited to at most 2^31 row groups in a file. If encryption is
220    /// enabled, this is reduced to 2^15, and row groups must be written sequentially.
221    ///
222    /// Every time the next row group is requested, the previous row group must
223    /// be finalised and closed using the [`SerializedRowGroupWriter::close`]
224    /// method or an error will be returned.
225    pub fn next_row_group(&mut self) -> Result<SerializedRowGroupWriter<'_, W>> {
226        self.assert_previous_writer_closed()?;
227        let ordinal = self.row_group_index;
228
229        // Thrift cannot encode lists with more than i32::MAX elements
230        let ordinal: i32 = ordinal.try_into().map_err(|_| {
231            ParquetError::General(format!(
232                "Parquet does not support more than {} row groups per file (currently: {})",
233                i32::MAX,
234                ordinal
235            ))
236        })?;
237
238        // If encryption is enabled, the max is 32767
239        #[cfg(feature = "encryption")]
240        if self.file_encryptor.is_some() && ordinal > i16::MAX as i32 {
241            return Err(ParquetError::General(format!(
242                "Parquet with encryption does not support more than {} row groups per file (currently: {})",
243                i16::MAX,
244                ordinal
245            )));
246        }
247
248        self.row_group_index = self
249            .row_group_index
250            .checked_add(1)
251            .ok_or_else(|| ParquetError::General("Row group index overflowed".to_string()))?;
252
253        let bloom_filter_position = self.properties().bloom_filter_position();
254        let row_groups = &mut self.row_groups;
255        let row_bloom_filters = &mut self.bloom_filters;
256        let row_column_indexes = &mut self.column_indexes;
257        let row_offset_indexes = &mut self.offset_indexes;
258        let on_close = move |buf,
259                             mut metadata,
260                             row_group_bloom_filter,
261                             row_group_column_index,
262                             row_group_offset_index| {
263            row_bloom_filters.push(row_group_bloom_filter);
264            row_column_indexes.push(row_group_column_index);
265            row_offset_indexes.push(row_group_offset_index);
266            // write bloom filters out immediately after the row group if requested
267            match bloom_filter_position {
268                BloomFilterPosition::AfterRowGroup => {
269                    write_bloom_filters(buf, row_bloom_filters, &mut metadata)?
270                }
271                BloomFilterPosition::End => (),
272            }
273            row_groups.push(metadata);
274            Ok(())
275        };
276
277        let row_group_writer = SerializedRowGroupWriter::new(
278            self.descr.clone(),
279            self.props.clone(),
280            &mut self.buf,
281            ordinal,
282            Some(Box::new(on_close)),
283        );
284        #[cfg(feature = "encryption")]
285        let row_group_writer = row_group_writer.with_file_encryptor(self.file_encryptor.clone());
286
287        Ok(row_group_writer)
288    }
289
290    /// Returns metadata for any flushed row groups
291    pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] {
292        &self.row_groups
293    }
294
295    /// Close and finalize the underlying Parquet writer
296    ///
297    /// Unlike [`Self::close`] this does not consume self
298    ///
299    /// Attempting to write after calling finish will result in an error
300    pub fn finish(&mut self) -> Result<ParquetMetaData> {
301        self.assert_previous_writer_closed()?;
302        let metadata = self.write_metadata()?;
303        self.buf.flush()?;
304        Ok(metadata)
305    }
306
307    /// Closes and finalises file writer, returning the file metadata.
308    pub fn close(mut self) -> Result<ParquetMetaData> {
309        self.finish()
310    }
311
312    /// Writes magic bytes at the beginning of the file.
313    #[cfg(not(feature = "encryption"))]
314    fn start_file(_properties: &WriterPropertiesPtr, buf: &mut TrackedWrite<W>) -> Result<()> {
315        buf.write_all(get_file_magic())?;
316        Ok(())
317    }
318
319    /// Writes magic bytes at the beginning of the file.
320    #[cfg(feature = "encryption")]
321    fn start_file(properties: &WriterPropertiesPtr, buf: &mut TrackedWrite<W>) -> Result<()> {
322        let magic = get_file_magic(properties.file_encryption_properties.as_ref());
323
324        buf.write_all(magic)?;
325        Ok(())
326    }
327
328    /// Assembles and writes metadata at the end of the file. This will take ownership
329    /// of `row_groups` and the page index structures.
330    fn write_metadata(&mut self) -> Result<ParquetMetaData> {
331        self.finished = true;
332
333        // write out any remaining bloom filters after all row groups
334        for row_group in &mut self.row_groups {
335            write_bloom_filters(&mut self.buf, &mut self.bloom_filters, row_group)?;
336        }
337
338        let key_value_metadata = match self.props.key_value_metadata() {
339            Some(kv) => Some(kv.iter().chain(&self.kv_metadatas).cloned().collect()),
340            None if self.kv_metadatas.is_empty() => None,
341            None => Some(self.kv_metadatas.clone()),
342        };
343
344        // take ownership of metadata
345        let row_groups = std::mem::take(&mut self.row_groups);
346        let column_indexes = std::mem::take(&mut self.column_indexes);
347        let offset_indexes = std::mem::take(&mut self.offset_indexes);
348
349        let write_path_in_schema = self.props.write_path_in_schema();
350        let mut encoder = ThriftMetadataWriter::new(
351            &mut self.buf,
352            &self.descr,
353            row_groups,
354            Some(self.props.created_by().to_string()),
355            self.props.writer_version().as_num(),
356            write_path_in_schema,
357            false,
358        );
359
360        #[cfg(feature = "encryption")]
361        {
362            encoder = encoder.with_file_encryptor(self.file_encryptor.clone());
363        }
364
365        if let Some(key_value_metadata) = key_value_metadata {
366            encoder = encoder.with_key_value_metadata(key_value_metadata)
367        }
368
369        // check for empty column index
370        let column_indexes = if column_indexes.is_empty()
371            || column_indexes
372                .iter()
373                .all(|cis| cis.iter().all(|ci| ci.is_none()))
374        {
375            None
376        } else {
377            Some(column_indexes)
378        };
379        // offset index will always be created unless explicitly disabled
380        let offset_indexes = if self.props.offset_index_disabled() {
381            None
382        } else {
383            Some(offset_indexes)
384        };
385        if column_indexes.is_some() || offset_indexes.is_some() {
386            let page_index = PageIndex::new(column_indexes, offset_indexes);
387            encoder = encoder.with_page_index(Arc::new(page_index));
388        }
389
390        encoder.finish()
391    }
392
393    #[inline]
394    fn assert_previous_writer_closed(&self) -> Result<()> {
395        if self.finished {
396            return Err(general_err!("SerializedFileWriter already finished"));
397        }
398
399        if self.row_group_index != self.row_groups.len() {
400            Err(general_err!("Previous row group writer was not closed"))
401        } else {
402            Ok(())
403        }
404    }
405
406    /// Add a [`KeyValue`] to the file writer's metadata
407    pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
408        self.kv_metadatas.push(kv_metadata);
409    }
410
411    /// Returns a reference to schema descriptor.
412    pub fn schema_descr(&self) -> &SchemaDescriptor {
413        &self.descr
414    }
415
416    /// Returns a reference to schema descriptor Arc.
417    #[cfg(feature = "arrow")]
418    pub(crate) fn schema_descr_ptr(&self) -> &SchemaDescPtr {
419        &self.descr
420    }
421
422    /// Returns a reference to the writer properties
423    pub fn properties(&self) -> &WriterPropertiesPtr {
424        &self.props
425    }
426
427    /// Returns a reference to the underlying writer.
428    pub fn inner(&self) -> &W {
429        self.buf.inner()
430    }
431
432    /// Writes the given buf bytes to the internal buffer.
433    ///
434    /// This can be used to write raw data to an in-progress Parquet file, for
435    /// example, custom index structures or other payloads. Other Parquet readers
436    /// will skip this data when reading the files.
437    ///
438    /// It's safe to use this method to write data to the underlying writer,
439    /// because it will ensure that the buffering and byte‐counting layers are used.
440    pub fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
441        self.buf.write_all(buf)
442    }
443
444    /// Flushes underlying writer
445    pub fn flush(&mut self) -> std::io::Result<()> {
446        self.buf.flush()
447    }
448
449    /// Returns a mutable reference to the underlying writer.
450    ///
451    /// **Warning**: if you write directly to this writer, you will skip
452    /// the `TrackedWrite` buffering and byte‐counting layers, which can cause
453    /// the file footer’s recorded offsets and sizes to diverge from reality,
454    /// resulting in an unreadable or corrupted Parquet file.
455    ///
456    /// If you want to write safely to the underlying writer, use [`Self::write_all`].
457    pub fn inner_mut(&mut self) -> &mut W {
458        self.buf.inner_mut()
459    }
460
461    /// Writes the file footer and returns the underlying writer.
462    pub fn into_inner(mut self) -> Result<W> {
463        self.assert_previous_writer_closed()?;
464        let _ = self.write_metadata()?;
465
466        self.buf.into_inner()
467    }
468
469    /// Returns the number of bytes written to this instance
470    pub fn bytes_written(&self) -> usize {
471        self.buf.bytes_written()
472    }
473
474    /// Get the file encryptor used by this instance to encrypt data
475    #[cfg(feature = "encryption")]
476    pub(crate) fn file_encryptor(&self) -> Option<Arc<FileEncryptor>> {
477        self.file_encryptor.clone()
478    }
479}
480
481/// Serialize all the bloom filters of the given row group to the given buffer,
482/// and returns the updated row group metadata.
483fn write_bloom_filters<W: Write + Send>(
484    buf: &mut TrackedWrite<W>,
485    bloom_filters: &mut [Vec<Option<Sbbf>>],
486    row_group: &mut RowGroupMetaData,
487) -> Result<()> {
488    // iter row group
489    // iter each column
490    // write bloom filter to the file
491
492    let row_group_idx: u32 = row_group
493        .ordinal()
494        .expect("Missing row group ordinal")
495        .try_into()
496        .map_err(|_| {
497            ParquetError::General(format!(
498                "Negative row group ordinal: {})",
499                row_group.ordinal().unwrap()
500            ))
501        })?;
502    let row_group_idx = row_group_idx as usize;
503    for (column_idx, column_chunk) in row_group.columns_mut().iter_mut().enumerate() {
504        if let Some(bloom_filter) = bloom_filters[row_group_idx][column_idx].take() {
505            let start_offset = buf.bytes_written();
506            bloom_filter.write(&mut *buf)?;
507            let end_offset = buf.bytes_written();
508            // set offset and index for bloom filter
509            *column_chunk = column_chunk
510                .clone()
511                .into_builder()
512                .set_bloom_filter_offset(Some(start_offset as i64))
513                .set_bloom_filter_length(Some((end_offset - start_offset) as i32))
514                .build()?;
515        }
516    }
517    Ok(())
518}
519
520/// Parquet row group writer API.
521///
522/// Provides methods to access column writers in an iterator-like fashion, order is
523/// guaranteed to match the order of schema leaves (column descriptors).
524///
525/// All columns should be written sequentially; the main workflow is:
526/// - Request the next column using `next_column` method - this will return `None` if no
527///   more columns are available to write.
528/// - Once done writing a column, close column writer with `close`
529/// - Once all columns have been written, close row group writer with `close`
530///   method. The close method will return row group metadata and is no-op
531///   on already closed row group.
532pub struct SerializedRowGroupWriter<'a, W: Write> {
533    descr: SchemaDescPtr,
534    props: WriterPropertiesPtr,
535    buf: &'a mut TrackedWrite<W>,
536    total_rows_written: Option<u64>,
537    total_bytes_written: u64,
538    total_uncompressed_bytes: i64,
539    column_index: usize,
540    row_group_metadata: Option<RowGroupMetaDataPtr>,
541    column_chunks: Vec<ColumnChunkMetaData>,
542    bloom_filters: Vec<Option<Sbbf>>,
543    column_indexes: Vec<Option<ColumnIndexMetaData>>,
544    offset_indexes: Vec<Option<OffsetIndexMetaData>>,
545    row_group_index: i32,
546    file_offset: i64,
547    on_close: Option<OnCloseRowGroup<'a, W>>,
548    #[cfg(feature = "encryption")]
549    file_encryptor: Option<Arc<FileEncryptor>>,
550}
551
552impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> {
553    /// Creates a new `SerializedRowGroupWriter` with:
554    ///
555    /// - `schema_descr` - the schema to write
556    /// - `properties` - writer properties
557    /// - `buf` - the buffer to write data to
558    /// - `row_group_index` - row group index in this parquet file.
559    /// - `file_offset` - file offset of this row group in this parquet file.
560    /// - `on_close` - an optional callback that will invoked on [`Self::close`]
561    pub fn new(
562        schema_descr: SchemaDescPtr,
563        properties: WriterPropertiesPtr,
564        buf: &'a mut TrackedWrite<W>,
565        row_group_index: i32,
566        on_close: Option<OnCloseRowGroup<'a, W>>,
567    ) -> Self {
568        let num_columns = schema_descr.num_columns();
569        let file_offset = buf.bytes_written() as i64;
570        Self {
571            buf,
572            row_group_index,
573            file_offset,
574            on_close,
575            total_rows_written: None,
576            descr: schema_descr,
577            props: properties,
578            column_index: 0,
579            row_group_metadata: None,
580            column_chunks: Vec::with_capacity(num_columns),
581            bloom_filters: Vec::with_capacity(num_columns),
582            column_indexes: Vec::with_capacity(num_columns),
583            offset_indexes: Vec::with_capacity(num_columns),
584            total_bytes_written: 0,
585            total_uncompressed_bytes: 0,
586            #[cfg(feature = "encryption")]
587            file_encryptor: None,
588        }
589    }
590
591    #[cfg(feature = "encryption")]
592    /// Set the file encryptor to use for encrypting row group data and metadata
593    pub(crate) fn with_file_encryptor(
594        mut self,
595        file_encryptor: Option<Arc<FileEncryptor>>,
596    ) -> Self {
597        self.file_encryptor = file_encryptor;
598        self
599    }
600
601    /// Advance `self.column_index` returning the next [`ColumnDescPtr`] if any
602    fn next_column_desc(&mut self) -> Option<ColumnDescPtr> {
603        let ret = self.descr.columns().get(self.column_index)?.clone();
604        self.column_index += 1;
605        Some(ret)
606    }
607
608    /// Returns [`OnCloseColumnChunk`] for the next writer
609    fn get_on_close(&mut self) -> (&mut TrackedWrite<W>, OnCloseColumnChunk<'_>) {
610        let total_bytes_written = &mut self.total_bytes_written;
611        let total_uncompressed_bytes = &mut self.total_uncompressed_bytes;
612        let total_rows_written = &mut self.total_rows_written;
613        let column_chunks = &mut self.column_chunks;
614        let column_indexes = &mut self.column_indexes;
615        let offset_indexes = &mut self.offset_indexes;
616        let bloom_filters = &mut self.bloom_filters;
617
618        let on_close = |r: ColumnCloseResult| {
619            // Update row group writer metrics
620            *total_bytes_written += r.bytes_written;
621            *total_uncompressed_bytes += r.metadata.uncompressed_size();
622            column_chunks.push(r.metadata);
623            bloom_filters.push(r.bloom_filter);
624            column_indexes.push(r.column_index);
625            offset_indexes.push(r.offset_index);
626
627            if let Some(rows) = *total_rows_written {
628                if rows != r.rows_written {
629                    return Err(general_err!(
630                        "Incorrect number of rows, expected {} != {} rows",
631                        rows,
632                        r.rows_written
633                    ));
634                }
635            } else {
636                *total_rows_written = Some(r.rows_written);
637            }
638
639            Ok(())
640        };
641        (self.buf, Box::new(on_close))
642    }
643
644    /// Returns the next column writer, if available, using the factory function;
645    /// otherwise returns `None`.
646    pub(crate) fn next_column_with_factory<'b, F, C>(&'b mut self, factory: F) -> Result<Option<C>>
647    where
648        F: FnOnce(
649            ColumnDescPtr,
650            WriterPropertiesPtr,
651            Box<dyn PageWriter + 'b>,
652            OnCloseColumnChunk<'b>,
653        ) -> Result<C>,
654    {
655        self.assert_previous_writer_closed()?;
656
657        let encryptor_context = self.get_page_encryptor_context();
658
659        Ok(match self.next_column_desc() {
660            Some(column) => {
661                let props = self.props.clone();
662                let (buf, on_close) = self.get_on_close();
663
664                let page_writer = SerializedPageWriter::new(buf);
665                let page_writer =
666                    Self::set_page_writer_encryptor(&column, encryptor_context, page_writer)?;
667
668                Some(factory(
669                    column,
670                    props,
671                    Box::new(page_writer),
672                    Box::new(on_close),
673                )?)
674            }
675            None => None,
676        })
677    }
678
679    /// Returns the next column writer, if available; otherwise returns `None`.
680    /// In case of any IO error or Thrift error, or if row group writer has already been
681    /// closed returns `Err`.
682    pub fn next_column(&mut self) -> Result<Option<SerializedColumnWriter<'_>>> {
683        self.next_column_with_factory(|descr, props, page_writer, on_close| {
684            let column_writer = get_column_writer(descr, props, page_writer);
685            Ok(SerializedColumnWriter::new(column_writer, Some(on_close)))
686        })
687    }
688
689    /// Append an encoded column chunk from `reader` directly to the underlying
690    /// writer.
691    ///
692    /// This method can be used for efficiently concatenating or projecting
693    /// Parquet data, or encoding Parquet data to temporary in-memory buffers.
694    ///
695    /// Arguments:
696    /// - `reader`: a [`ChunkReader`] containing the encoded column data
697    /// - `close`: the [`ColumnCloseResult`] metadata returned from closing
698    ///   the column writer that wrote the data in `reader`.
699    ///
700    /// See Also:
701    /// 1. [`get_column_writer`]  for creating writers that can encode data.
702    /// 2. [`Self::next_column`] for writing data that isn't already encoded
703    pub fn append_column<R: ChunkReader>(
704        &mut self,
705        reader: &R,
706        close: ColumnCloseResult,
707    ) -> Result<()> {
708        // Position a reader at the start of the buffered chunk, then splice the
709        // bytes through the shared streaming path.
710        let metadata = &close.metadata;
711        let src_offset = metadata
712            .dictionary_page_offset()
713            .unwrap_or_else(|| metadata.data_page_offset());
714        let read = reader.get_read(src_offset as _)?;
715        self.append_column_from_read(read, close)
716    }
717
718    /// Splice an already-encoded column chunk into the row group, reading its
719    /// bytes sequentially from `read`.
720    ///
721    /// `read` must be positioned at the start of the chunk (the dictionary page
722    /// if present, otherwise the first data page — i.e. `src_offset` below) and
723    /// yield exactly the chunk's compressed bytes. Unlike [`Self::append_column`]
724    /// this consumes an owned [`Read`], which lets the caller stream the bytes
725    /// back from a [`PageStore`](crate::column::page_store::PageStore) one page
726    /// at a time without materializing the whole chunk in memory.
727    pub(crate) fn append_column_from_read<R: Read>(
728        &mut self,
729        read: R,
730        close: ColumnCloseResult,
731    ) -> Result<()> {
732        let (src_offset, src_length, write_offset) = self.begin_appended_column(&close)?;
733
734        let mut read = read.take(src_length as _);
735        let write_length = std::io::copy(&mut read, &mut *self.buf)?;
736
737        if src_length as u64 != write_length {
738            return Err(general_err!(
739                "Failed to splice column data, expected {src_length} got {write_length}"
740            ));
741        }
742
743        self.finish_appended_column(close, src_offset, write_offset)
744    }
745
746    /// Splice an already-encoded column chunk into the row group from an
747    /// in-order sequence of byte buffers (typically its serialized pages).
748    ///
749    /// This is a lower-overhead alternative to [`Self::append_column`] /
750    /// [`Self::append_column_from_read`] for callers that already hold the
751    /// chunk as owned [`Bytes`]: each buffer is written straight to the output
752    /// with a single `write_all`, skipping the intermediate copy through
753    /// [`std::io::copy`]'s fixed-size buffer.
754    ///
755    /// `pages` must yield the chunk's compressed bytes in final file order
756    /// (the dictionary page, if any, first) and together total exactly the
757    /// compressed size recorded in `close`.
758    #[cfg(feature = "arrow")]
759    pub(crate) fn append_column_from_pages<I>(
760        &mut self,
761        pages: I,
762        close: ColumnCloseResult,
763    ) -> Result<()>
764    where
765        I: IntoIterator<Item = Result<Bytes>>,
766    {
767        let (src_offset, src_length, write_offset) = self.begin_appended_column(&close)?;
768
769        let mut write_length = 0u64;
770        for page in pages {
771            let page = page?;
772            self.buf.write_all(&page)?;
773            write_length += page.len() as u64;
774        }
775
776        if src_length as u64 != write_length {
777            return Err(general_err!(
778                "Failed to splice column data, expected {src_length} got {write_length}"
779            ));
780        }
781
782        self.finish_appended_column(close, src_offset, write_offset)
783    }
784
785    /// [`Self::append_column_from_read`] / [`Self::append_column_from_pages`]
786    /// preamble: validates the writer state and that `close` matches the next
787    /// expected column.
788    ///
789    /// Returns `(src_offset, src_length, write_offset)`: the chunk's start
790    /// offset and length in the source buffer, and the offset at which it will
791    /// land in the output file.
792    fn begin_appended_column(&mut self, close: &ColumnCloseResult) -> Result<(i64, i64, usize)> {
793        self.assert_previous_writer_closed()?;
794        let desc = self
795            .next_column_desc()
796            .ok_or_else(|| general_err!("exhausted columns in SerializedRowGroupWriter"))?;
797
798        let metadata = &close.metadata;
799
800        if metadata.column_descr() != desc.as_ref() {
801            return Err(general_err!(
802                "column descriptor mismatch, expected {:?} got {:?}",
803                desc,
804                metadata.column_descr()
805            ));
806        }
807
808        let src_offset = metadata
809            .dictionary_page_offset()
810            .unwrap_or_else(|| metadata.data_page_offset());
811        let src_length = metadata.compressed_size();
812        let write_offset = self.buf.bytes_written();
813        Ok((src_offset, src_length, write_offset))
814    }
815
816    /// [`Self::append_column_from_read`] / [`Self::append_column_from_pages`]
817    /// epilogue: rewrites the buffer-relative page offsets recorded in `close`
818    /// to their final positions in the output file and closes the column.
819    fn finish_appended_column(
820        &mut self,
821        mut close: ColumnCloseResult,
822        src_offset: i64,
823        write_offset: usize,
824    ) -> Result<()> {
825        let metadata = close.metadata;
826        let src_dictionary_offset = metadata.dictionary_page_offset();
827        let src_data_offset = metadata.data_page_offset();
828
829        let map_offset = |x| x - src_offset + write_offset as i64;
830        let mut builder = ColumnChunkMetaData::builder(metadata.column_descr_ptr())
831            .set_compression_codec(metadata.compression_codec())
832            .set_encodings_mask(*metadata.encodings_mask())
833            .set_total_compressed_size(metadata.compressed_size())
834            .set_total_uncompressed_size(metadata.uncompressed_size())
835            .set_num_values(metadata.num_values())
836            .set_data_page_offset(map_offset(src_data_offset))
837            .set_dictionary_page_offset(src_dictionary_offset.map(map_offset))
838            .set_unencoded_byte_array_data_bytes(metadata.unencoded_byte_array_data_bytes());
839
840        if let Some(rep_hist) = metadata.repetition_level_histogram() {
841            builder = builder.set_repetition_level_histogram(Some(rep_hist.clone()))
842        }
843        if let Some(def_hist) = metadata.definition_level_histogram() {
844            builder = builder.set_definition_level_histogram(Some(def_hist.clone()))
845        }
846        if let Some(statistics) = metadata.statistics() {
847            builder = builder.set_statistics(statistics.clone())
848        }
849        if let Some(geo_statistics) = metadata.geo_statistics() {
850            builder = builder.set_geo_statistics(Box::new(geo_statistics.clone()))
851        }
852        if let Some(page_encoding_stats) = metadata.page_encoding_stats() {
853            builder = builder.set_page_encoding_stats(page_encoding_stats.clone())
854        }
855        builder = self.set_column_crypto_metadata(builder, &metadata);
856        close.metadata = builder.build()?;
857
858        if let Some(offsets) = close.offset_index.as_mut() {
859            for location in &mut offsets.page_locations {
860                location.offset = map_offset(location.offset)
861            }
862        }
863
864        let (_, on_close) = self.get_on_close();
865        on_close(close)
866    }
867
868    /// Closes this row group writer and returns row group metadata.
869    pub fn close(mut self) -> Result<RowGroupMetaDataPtr> {
870        if self.row_group_metadata.is_none() {
871            self.assert_previous_writer_closed()?;
872
873            let column_chunks = std::mem::take(&mut self.column_chunks);
874            let row_group_metadata = RowGroupMetaData::builder(self.descr.clone())
875                .set_column_metadata(column_chunks)
876                .set_total_byte_size(self.total_uncompressed_bytes)
877                .set_num_rows(self.total_rows_written.unwrap_or(0) as i64)
878                .set_sorting_columns(self.props.sorting_columns().cloned())
879                .set_ordinal(self.row_group_index)
880                .set_file_offset(self.file_offset)
881                .build()?;
882
883            self.row_group_metadata = Some(Arc::new(row_group_metadata.clone()));
884
885            if let Some(on_close) = self.on_close.take() {
886                on_close(
887                    self.buf,
888                    row_group_metadata,
889                    self.bloom_filters,
890                    self.column_indexes,
891                    self.offset_indexes,
892                )?
893            }
894        }
895
896        let metadata = self.row_group_metadata.as_ref().unwrap().clone();
897        Ok(metadata)
898    }
899
900    /// Set the column crypto metadata for a column chunk
901    #[cfg(feature = "encryption")]
902    fn set_column_crypto_metadata(
903        &self,
904        builder: ColumnChunkMetaDataBuilder,
905        metadata: &ColumnChunkMetaData,
906    ) -> ColumnChunkMetaDataBuilder {
907        if let Some(file_encryptor) = self.file_encryptor.as_ref() {
908            builder.set_column_crypto_metadata(get_column_crypto_metadata(
909                file_encryptor.properties(),
910                &metadata.column_descr_ptr(),
911            ))
912        } else {
913            builder
914        }
915    }
916
917    /// Get context required to create a [`PageEncryptor`] for a column
918    #[cfg(feature = "encryption")]
919    fn get_page_encryptor_context(&self) -> PageEncryptorContext {
920        PageEncryptorContext {
921            file_encryptor: self.file_encryptor.clone(),
922            row_group_index: self.row_group_index as usize,
923            column_index: self.column_index,
924        }
925    }
926
927    /// Set the [`PageEncryptor`] on a page writer if a column is encrypted
928    #[cfg(feature = "encryption")]
929    fn set_page_writer_encryptor<'b>(
930        column: &ColumnDescPtr,
931        context: PageEncryptorContext,
932        page_writer: SerializedPageWriter<'b, W>,
933    ) -> Result<SerializedPageWriter<'b, W>> {
934        let page_encryptor = PageEncryptor::create_if_column_encrypted(
935            context.file_encryptor.as_ref(),
936            context.row_group_index,
937            context.column_index,
938            &column.path().string(),
939        )?;
940
941        Ok(page_writer.with_page_encryptor(page_encryptor))
942    }
943
944    /// No-op implementation of setting the column crypto metadata for a column chunk
945    #[cfg(not(feature = "encryption"))]
946    fn set_column_crypto_metadata(
947        &self,
948        builder: ColumnChunkMetaDataBuilder,
949        _metadata: &ColumnChunkMetaData,
950    ) -> ColumnChunkMetaDataBuilder {
951        builder
952    }
953
954    #[cfg(not(feature = "encryption"))]
955    fn get_page_encryptor_context(&self) -> PageEncryptorContext {
956        PageEncryptorContext {}
957    }
958
959    /// No-op implementation of setting a [`PageEncryptor`] for when encryption is disabled
960    #[cfg(not(feature = "encryption"))]
961    fn set_page_writer_encryptor<'b>(
962        _column: &ColumnDescPtr,
963        _context: PageEncryptorContext,
964        page_writer: SerializedPageWriter<'b, W>,
965    ) -> Result<SerializedPageWriter<'b, W>> {
966        Ok(page_writer)
967    }
968
969    #[inline]
970    fn assert_previous_writer_closed(&self) -> Result<()> {
971        if self.column_index != self.column_chunks.len() {
972            Err(general_err!("Previous column writer was not closed"))
973        } else {
974            Ok(())
975        }
976    }
977}
978
979/// Context required to create a [`PageEncryptor`] for a column
980#[cfg(feature = "encryption")]
981struct PageEncryptorContext {
982    file_encryptor: Option<Arc<FileEncryptor>>,
983    row_group_index: usize,
984    column_index: usize,
985}
986
987#[cfg(not(feature = "encryption"))]
988struct PageEncryptorContext {}
989
990/// A wrapper around a [`ColumnWriter`] that invokes a callback on [`Self::close`]
991pub struct SerializedColumnWriter<'a> {
992    inner: ColumnWriter<'a>,
993    on_close: Option<OnCloseColumnChunk<'a>>,
994}
995
996impl<'a> SerializedColumnWriter<'a> {
997    /// Create a new [`SerializedColumnWriter`] from a [`ColumnWriter`] and an
998    /// optional callback to be invoked on [`Self::close`]
999    pub fn new(inner: ColumnWriter<'a>, on_close: Option<OnCloseColumnChunk<'a>>) -> Self {
1000        Self { inner, on_close }
1001    }
1002
1003    /// Returns a reference to an untyped [`ColumnWriter`]
1004    pub fn untyped(&mut self) -> &mut ColumnWriter<'a> {
1005        &mut self.inner
1006    }
1007
1008    /// Returns a reference to a typed [`ColumnWriterImpl`]
1009    pub fn typed<T: DataType>(&mut self) -> &mut ColumnWriterImpl<'a, T> {
1010        get_typed_column_writer_mut(&mut self.inner)
1011    }
1012
1013    /// Close this [`SerializedColumnWriter`]
1014    pub fn close(mut self) -> Result<()> {
1015        let r = self.inner.close()?;
1016        if let Some(on_close) = self.on_close.take() {
1017            on_close(r)?
1018        }
1019
1020        Ok(())
1021    }
1022}
1023
1024/// A serialized implementation for Parquet [`PageWriter`].
1025/// Writes and serializes pages and metadata into output stream.
1026///
1027/// `SerializedPageWriter` should not be used after calling `close()`.
1028pub struct SerializedPageWriter<'a, W: Write> {
1029    sink: &'a mut TrackedWrite<W>,
1030    #[cfg(feature = "encryption")]
1031    page_encryptor: Option<PageEncryptor>,
1032}
1033
1034impl<'a, W: Write> SerializedPageWriter<'a, W> {
1035    /// Creates new page writer.
1036    pub fn new(sink: &'a mut TrackedWrite<W>) -> Self {
1037        Self {
1038            sink,
1039            #[cfg(feature = "encryption")]
1040            page_encryptor: None,
1041        }
1042    }
1043
1044    /// Serializes page header into Thrift.
1045    /// Returns number of bytes that have been written into the sink.
1046    #[inline]
1047    fn serialize_page_header(&mut self, header: PageHeader) -> Result<usize> {
1048        let start_pos = self.sink.bytes_written();
1049        match self.page_encryptor_and_sink_mut() {
1050            Some((page_encryptor, sink)) => {
1051                page_encryptor.encrypt_page_header(&header, sink)?;
1052            }
1053            None => {
1054                let mut protocol = ThriftCompactOutputProtocol::new(&mut *self.sink);
1055                header.write_thrift(&mut protocol)?;
1056            }
1057        }
1058        Ok(self.sink.bytes_written() - start_pos)
1059    }
1060}
1061
1062#[cfg(feature = "encryption")]
1063impl<W: Write> SerializedPageWriter<'_, W> {
1064    /// Set the encryptor to use to encrypt page data
1065    fn with_page_encryptor(mut self, page_encryptor: Option<PageEncryptor>) -> Self {
1066        self.page_encryptor = page_encryptor;
1067        self
1068    }
1069
1070    fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
1071        self.page_encryptor.as_mut()
1072    }
1073
1074    fn page_encryptor_and_sink_mut(
1075        &mut self,
1076    ) -> Option<(&mut PageEncryptor, &mut TrackedWrite<W>)> {
1077        self.page_encryptor.as_mut().map(|pe| (pe, &mut *self.sink))
1078    }
1079}
1080
1081// These mirror the signatures of the encryption-enabled versions above, so that the
1082// callers do not need a `cfg` of their own.
1083#[cfg(not(feature = "encryption"))]
1084#[expect(
1085    clippy::needless_pass_by_ref_mut,
1086    reason = "mirrors the encryption-enabled signatures"
1087)]
1088impl<W: Write> SerializedPageWriter<'_, W> {
1089    fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
1090        None
1091    }
1092
1093    fn page_encryptor_and_sink_mut(
1094        &mut self,
1095    ) -> Option<(&mut PageEncryptor, &mut TrackedWrite<W>)> {
1096        None
1097    }
1098}
1099
1100impl<W: Write + Send> PageWriter for SerializedPageWriter<'_, W> {
1101    fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
1102        let page = match self.page_encryptor_mut() {
1103            Some(page_encryptor) => page_encryptor.encrypt_compressed_page(page)?,
1104            None => page,
1105        };
1106
1107        let page_type = page.page_type();
1108        let start_pos = self.sink.bytes_written() as u64;
1109
1110        let page_header = page.to_thrift_header()?;
1111        let header_size = self.serialize_page_header(page_header)?;
1112
1113        self.sink.write_all(page.data())?;
1114
1115        let mut spec = PageWriteSpec::new();
1116        spec.page_type = page_type;
1117        spec.uncompressed_size = page.uncompressed_size() + header_size;
1118        spec.compressed_size = page.compressed_size() + header_size;
1119        spec.offset = start_pos;
1120        spec.bytes_written = self.sink.bytes_written() as u64 - start_pos;
1121        spec.num_values = page.num_values();
1122
1123        if let Some(page_encryptor) = self.page_encryptor_mut()
1124            && page.compressed_page().is_data_page()
1125        {
1126            page_encryptor.increment_page();
1127        }
1128        Ok(spec)
1129    }
1130
1131    fn close(&mut self) -> Result<()> {
1132        self.sink.flush()?;
1133        Ok(())
1134    }
1135}
1136
1137/// Get the magic bytes at the start and end of the file that identify this
1138/// as a Parquet file.
1139#[cfg(feature = "encryption")]
1140pub(crate) fn get_file_magic(
1141    file_encryption_properties: Option<&Arc<FileEncryptionProperties>>,
1142) -> &'static [u8; 4] {
1143    match file_encryption_properties.as_ref() {
1144        Some(encryption_properties) if encryption_properties.encrypt_footer() => {
1145            &PARQUET_MAGIC_ENCR_FOOTER
1146        }
1147        _ => &PARQUET_MAGIC,
1148    }
1149}
1150
1151#[cfg(not(feature = "encryption"))]
1152pub(crate) fn get_file_magic() -> &'static [u8; 4] {
1153    &PARQUET_MAGIC
1154}
1155
1156#[cfg(test)]
1157mod tests {
1158    use super::*;
1159
1160    #[cfg(feature = "arrow")]
1161    use arrow_array::RecordBatchReader;
1162    use bytes::Bytes;
1163    use std::fs::File;
1164
1165    #[cfg(feature = "arrow")]
1166    use crate::arrow::ArrowWriter;
1167    #[cfg(feature = "arrow")]
1168    use crate::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
1169    use crate::basic::{
1170        ColumnOrder, Compression, ConvertedType, Encoding, LogicalType, Repetition, SortOrder, Type,
1171    };
1172    use crate::column::page::{Page, PageReader};
1173    use crate::column::reader::get_typed_column_reader;
1174    use crate::compression::{Codec, CodecOptionsBuilder, create_codec};
1175    use crate::data_type::{BoolType, ByteArrayType, Int32Type, Int96, Int96Type};
1176    use crate::file::page_index::column_index::ColumnIndexMetaData;
1177    use crate::file::properties::EnabledStatistics;
1178    use crate::file::serialized_reader::ReadOptionsBuilder;
1179    use crate::file::statistics::{from_thrift_page_stats, page_stats_to_thrift};
1180    use crate::file::{
1181        properties::{ReaderProperties, WriterProperties, WriterVersion},
1182        reader::{FileReader, SerializedFileReader, SerializedPageReader},
1183        statistics::Statistics,
1184    };
1185    use crate::record::{Row, RowAccessor};
1186    use crate::schema::parser::parse_message_type;
1187    use crate::schema::types;
1188    use crate::schema::types::{ColumnDescriptor, ColumnPath};
1189    use crate::util::test_common::file_util::get_test_file;
1190    use crate::util::test_common::rand_gen::RandGen;
1191
1192    #[test]
1193    fn test_row_group_writer_error_not_all_columns_written() {
1194        let file = tempfile::tempfile().unwrap();
1195        let schema = Arc::new(
1196            types::Type::group_type_builder("schema")
1197                .with_fields(vec![Arc::new(
1198                    types::Type::primitive_type_builder("col1", Type::INT32)
1199                        .build()
1200                        .unwrap(),
1201                )])
1202                .build()
1203                .unwrap(),
1204        );
1205        let props = Default::default();
1206        let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
1207        let row_group_writer = writer.next_row_group().unwrap();
1208        let res = row_group_writer.close();
1209        assert!(res.is_err());
1210        if let Err(err) = res {
1211            assert_eq!(
1212                format!("{err}"),
1213                "Parquet error: Column length mismatch: 1 != 0"
1214            );
1215        }
1216    }
1217
1218    #[test]
1219    fn test_row_group_writer_num_records_mismatch() {
1220        let file = tempfile::tempfile().unwrap();
1221        let schema = Arc::new(
1222            types::Type::group_type_builder("schema")
1223                .with_fields(vec![
1224                    Arc::new(
1225                        types::Type::primitive_type_builder("col1", Type::INT32)
1226                            .with_repetition(Repetition::REQUIRED)
1227                            .build()
1228                            .unwrap(),
1229                    ),
1230                    Arc::new(
1231                        types::Type::primitive_type_builder("col2", Type::INT32)
1232                            .with_repetition(Repetition::REQUIRED)
1233                            .build()
1234                            .unwrap(),
1235                    ),
1236                ])
1237                .build()
1238                .unwrap(),
1239        );
1240        let props = Default::default();
1241        let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
1242        let mut row_group_writer = writer.next_row_group().unwrap();
1243
1244        let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
1245        col_writer
1246            .typed::<Int32Type>()
1247            .write_batch(&[1, 2, 3], None, None)
1248            .unwrap();
1249        col_writer.close().unwrap();
1250
1251        let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
1252        col_writer
1253            .typed::<Int32Type>()
1254            .write_batch(&[1, 2], None, None)
1255            .unwrap();
1256
1257        let err = col_writer.close().unwrap_err();
1258        assert_eq!(
1259            err.to_string(),
1260            "Parquet error: Incorrect number of rows, expected 3 != 2 rows"
1261        );
1262    }
1263
1264    #[test]
1265    fn test_file_writer_empty_file() {
1266        let file = tempfile::tempfile().unwrap();
1267
1268        let schema = Arc::new(
1269            types::Type::group_type_builder("schema")
1270                .with_fields(vec![Arc::new(
1271                    types::Type::primitive_type_builder("col1", Type::INT32)
1272                        .build()
1273                        .unwrap(),
1274                )])
1275                .build()
1276                .unwrap(),
1277        );
1278        let props = Default::default();
1279        let writer = SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap();
1280        writer.close().unwrap();
1281
1282        let reader = SerializedFileReader::new(file).unwrap();
1283        assert_eq!(reader.get_row_iter(None).unwrap().count(), 0);
1284    }
1285
1286    #[test]
1287    fn test_file_writer_column_orders_populated() {
1288        let file = tempfile::tempfile().unwrap();
1289
1290        let schema = Arc::new(
1291            types::Type::group_type_builder("schema")
1292                .with_fields(vec![
1293                    Arc::new(
1294                        types::Type::primitive_type_builder("col1", Type::INT32)
1295                            .build()
1296                            .unwrap(),
1297                    ),
1298                    Arc::new(
1299                        types::Type::primitive_type_builder("col2", Type::FIXED_LEN_BYTE_ARRAY)
1300                            .with_converted_type(ConvertedType::INTERVAL)
1301                            .with_length(12)
1302                            .build()
1303                            .unwrap(),
1304                    ),
1305                    Arc::new(
1306                        types::Type::group_type_builder("nested")
1307                            .with_repetition(Repetition::REQUIRED)
1308                            .with_fields(vec![
1309                                Arc::new(
1310                                    types::Type::primitive_type_builder(
1311                                        "col3",
1312                                        Type::FIXED_LEN_BYTE_ARRAY,
1313                                    )
1314                                    .with_logical_type(Some(LogicalType::Float16))
1315                                    .with_length(2)
1316                                    .build()
1317                                    .unwrap(),
1318                                ),
1319                                Arc::new(
1320                                    types::Type::primitive_type_builder("col4", Type::BYTE_ARRAY)
1321                                        .with_logical_type(Some(LogicalType::String))
1322                                        .build()
1323                                        .unwrap(),
1324                                ),
1325                            ])
1326                            .build()
1327                            .unwrap(),
1328                    ),
1329                    Arc::new(
1330                        types::Type::primitive_type_builder("col5", Type::FLOAT)
1331                            .build()
1332                            .unwrap(),
1333                    ),
1334                    Arc::new(
1335                        types::Type::primitive_type_builder("col6", Type::DOUBLE)
1336                            .build()
1337                            .unwrap(),
1338                    ),
1339                ])
1340                .build()
1341                .unwrap(),
1342        );
1343
1344        let props = Default::default();
1345        let writer = SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap();
1346        writer.close().unwrap();
1347
1348        let reader = SerializedFileReader::new(file).unwrap();
1349
1350        // only leaves
1351        let expected = vec![
1352            // INT32
1353            ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED),
1354            // INTERVAL
1355            ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED),
1356            // Float16
1357            ColumnOrder::IEEE_754_TOTAL_ORDER,
1358            // String
1359            ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED),
1360            // FLOAT
1361            ColumnOrder::IEEE_754_TOTAL_ORDER,
1362            // DOUBLE
1363            ColumnOrder::IEEE_754_TOTAL_ORDER,
1364        ];
1365        let actual = reader.metadata().file_metadata().column_orders();
1366
1367        assert!(actual.is_some());
1368        let actual = actual.unwrap();
1369        assert_eq!(*actual, expected);
1370    }
1371
1372    #[test]
1373    fn test_file_writer_with_metadata() {
1374        let file = tempfile::tempfile().unwrap();
1375
1376        let schema = Arc::new(
1377            types::Type::group_type_builder("schema")
1378                .with_fields(vec![Arc::new(
1379                    types::Type::primitive_type_builder("col1", Type::INT32)
1380                        .build()
1381                        .unwrap(),
1382                )])
1383                .build()
1384                .unwrap(),
1385        );
1386        let props = Arc::new(
1387            WriterProperties::builder()
1388                .set_key_value_metadata(Some(vec![KeyValue::new(
1389                    "key".to_string(),
1390                    "value".to_string(),
1391                )]))
1392                .build(),
1393        );
1394        let writer = SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap();
1395        writer.close().unwrap();
1396
1397        let reader = SerializedFileReader::new(file).unwrap();
1398        assert_eq!(
1399            reader
1400                .metadata()
1401                .file_metadata()
1402                .key_value_metadata()
1403                .unwrap()
1404                .len(),
1405            1
1406        );
1407    }
1408
1409    #[test]
1410    fn test_file_writer_v2_with_metadata() {
1411        let file = tempfile::tempfile().unwrap();
1412        let field_logical_type = Some(LogicalType::integer(8, false));
1413        let field = Arc::new(
1414            types::Type::primitive_type_builder("col1", Type::INT32)
1415                .with_logical_type(field_logical_type.clone())
1416                .with_converted_type(field_logical_type.into())
1417                .build()
1418                .unwrap(),
1419        );
1420        let schema = Arc::new(
1421            types::Type::group_type_builder("schema")
1422                .with_fields(vec![field.clone()])
1423                .build()
1424                .unwrap(),
1425        );
1426        let props = Arc::new(
1427            WriterProperties::builder()
1428                .set_key_value_metadata(Some(vec![KeyValue::new(
1429                    "key".to_string(),
1430                    "value".to_string(),
1431                )]))
1432                .set_writer_version(WriterVersion::PARQUET_2_0)
1433                .build(),
1434        );
1435        let writer = SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap();
1436        writer.close().unwrap();
1437
1438        let reader = SerializedFileReader::new(file).unwrap();
1439
1440        assert_eq!(
1441            reader
1442                .metadata()
1443                .file_metadata()
1444                .key_value_metadata()
1445                .unwrap()
1446                .len(),
1447            1
1448        );
1449
1450        // ARROW-11803: Test that the converted and logical types have been populated
1451        let fields = reader.metadata().file_metadata().schema().get_fields();
1452        assert_eq!(fields.len(), 1);
1453        assert_eq!(fields[0], field);
1454    }
1455
1456    #[test]
1457    fn test_file_writer_with_sorting_columns_metadata() {
1458        let file = tempfile::tempfile().unwrap();
1459
1460        let schema = Arc::new(
1461            types::Type::group_type_builder("schema")
1462                .with_fields(vec![
1463                    Arc::new(
1464                        types::Type::primitive_type_builder("col1", Type::INT32)
1465                            .build()
1466                            .unwrap(),
1467                    ),
1468                    Arc::new(
1469                        types::Type::primitive_type_builder("col2", Type::INT32)
1470                            .build()
1471                            .unwrap(),
1472                    ),
1473                ])
1474                .build()
1475                .unwrap(),
1476        );
1477        let expected_result = Some(vec![SortingColumn {
1478            column_idx: 0,
1479            descending: false,
1480            nulls_first: true,
1481        }]);
1482        let props = Arc::new(
1483            WriterProperties::builder()
1484                .set_key_value_metadata(Some(vec![KeyValue::new(
1485                    "key".to_string(),
1486                    "value".to_string(),
1487                )]))
1488                .set_sorting_columns(expected_result.clone())
1489                .build(),
1490        );
1491        let mut writer =
1492            SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap();
1493        let mut row_group_writer = writer.next_row_group().expect("get row group writer");
1494
1495        let col_writer = row_group_writer.next_column().unwrap().unwrap();
1496        col_writer.close().unwrap();
1497
1498        let col_writer = row_group_writer.next_column().unwrap().unwrap();
1499        col_writer.close().unwrap();
1500
1501        row_group_writer.close().unwrap();
1502        writer.close().unwrap();
1503
1504        let reader = SerializedFileReader::new(file).unwrap();
1505        let result: Vec<Option<&Vec<SortingColumn>>> = reader
1506            .metadata()
1507            .row_groups()
1508            .iter()
1509            .map(|f| f.sorting_columns())
1510            .collect();
1511        // validate the sorting column read match the one written above
1512        assert_eq!(expected_result.as_ref(), result[0]);
1513    }
1514
1515    #[test]
1516    fn test_file_writer_empty_row_groups() {
1517        let file = tempfile::tempfile().unwrap();
1518        test_file_roundtrip(file, vec![]);
1519    }
1520
1521    #[test]
1522    fn test_file_writer_single_row_group() {
1523        let file = tempfile::tempfile().unwrap();
1524        test_file_roundtrip(file, vec![vec![1, 2, 3, 4, 5]]);
1525    }
1526
1527    #[test]
1528    fn test_file_writer_multiple_row_groups() {
1529        let file = tempfile::tempfile().unwrap();
1530        test_file_roundtrip(
1531            file,
1532            vec![
1533                vec![1, 2, 3, 4, 5],
1534                vec![1, 2, 3],
1535                vec![1],
1536                vec![1, 2, 3, 4, 5, 6],
1537            ],
1538        );
1539    }
1540
1541    #[test]
1542    fn test_file_writer_multiple_large_row_groups() {
1543        let file = tempfile::tempfile().unwrap();
1544        test_file_roundtrip(
1545            file,
1546            vec![vec![123; 1024], vec![124; 1000], vec![125; 15], vec![]],
1547        );
1548    }
1549
1550    #[test]
1551    fn test_page_writer_data_pages() {
1552        let pages = [
1553            Page::DataPage {
1554                buf: Bytes::from(vec![1, 2, 3, 4, 5, 6, 7, 8]),
1555                num_values: 10,
1556                encoding: Encoding::DELTA_BINARY_PACKED,
1557                def_level_encoding: Encoding::RLE,
1558                rep_level_encoding: Encoding::RLE,
1559                statistics: Some(Statistics::int32(Some(1), Some(3), None, Some(7), true)),
1560            },
1561            Page::DataPageV2 {
1562                buf: Bytes::from(vec![4; 128]),
1563                num_values: 10,
1564                encoding: Encoding::DELTA_BINARY_PACKED,
1565                num_nulls: 2,
1566                num_rows: 12,
1567                def_levels_byte_len: 24,
1568                rep_levels_byte_len: 32,
1569                is_compressed: false,
1570                statistics: Some(Statistics::int32(Some(1), Some(3), None, Some(7), true)),
1571            },
1572        ];
1573
1574        test_page_roundtrip(&pages[..], Compression::SNAPPY, Type::INT32);
1575        test_page_roundtrip(&pages[..], Compression::UNCOMPRESSED, Type::INT32);
1576    }
1577
1578    #[test]
1579    fn test_page_writer_dict_pages() {
1580        let pages = [
1581            Page::DictionaryPage {
1582                buf: Bytes::from(vec![1, 2, 3, 4, 5]),
1583                num_values: 5,
1584                encoding: Encoding::RLE_DICTIONARY,
1585                is_sorted: false,
1586            },
1587            Page::DataPage {
1588                buf: Bytes::from(vec![1, 2, 3, 4, 5, 6, 7, 8]),
1589                num_values: 10,
1590                encoding: Encoding::DELTA_BINARY_PACKED,
1591                def_level_encoding: Encoding::RLE,
1592                rep_level_encoding: Encoding::RLE,
1593                statistics: Some(Statistics::int32(Some(1), Some(3), None, Some(7), true)),
1594            },
1595            Page::DataPageV2 {
1596                buf: Bytes::from(vec![4; 128]),
1597                num_values: 10,
1598                encoding: Encoding::DELTA_BINARY_PACKED,
1599                num_nulls: 2,
1600                num_rows: 12,
1601                def_levels_byte_len: 24,
1602                rep_levels_byte_len: 32,
1603                is_compressed: false,
1604                statistics: None,
1605            },
1606        ];
1607
1608        test_page_roundtrip(&pages[..], Compression::SNAPPY, Type::INT32);
1609        test_page_roundtrip(&pages[..], Compression::UNCOMPRESSED, Type::INT32);
1610    }
1611
1612    /// Tests writing and reading pages.
1613    /// Physical type is for statistics only, should match any defined statistics type in
1614    /// pages.
1615    fn test_page_roundtrip(pages: &[Page], codec: Compression, physical_type: Type) {
1616        let mut compressed_pages = vec![];
1617        let mut total_num_values = 0i64;
1618        let codec_options = CodecOptionsBuilder::default()
1619            .set_backward_compatible_lz4(false)
1620            .build();
1621        let mut compressor = create_codec(codec, &codec_options).unwrap();
1622
1623        for page in pages {
1624            let uncompressed_len = page.buffer().len();
1625
1626            let compressed_page = match *page {
1627                Page::DataPage {
1628                    ref buf,
1629                    num_values,
1630                    encoding,
1631                    def_level_encoding,
1632                    rep_level_encoding,
1633                    ref statistics,
1634                } => {
1635                    total_num_values += num_values as i64;
1636                    let output_buf = compress_helper(compressor.as_mut(), buf);
1637
1638                    Page::DataPage {
1639                        buf: Bytes::from(output_buf),
1640                        num_values,
1641                        encoding,
1642                        def_level_encoding,
1643                        rep_level_encoding,
1644                        statistics: from_thrift_page_stats(
1645                            physical_type,
1646                            page_stats_to_thrift(statistics.as_ref()),
1647                        )
1648                        .unwrap(),
1649                    }
1650                }
1651                Page::DataPageV2 {
1652                    ref buf,
1653                    num_values,
1654                    encoding,
1655                    num_nulls,
1656                    num_rows,
1657                    def_levels_byte_len,
1658                    rep_levels_byte_len,
1659                    ref statistics,
1660                    ..
1661                } => {
1662                    total_num_values += num_values as i64;
1663                    let offset = (def_levels_byte_len + rep_levels_byte_len) as usize;
1664                    let cmp_buf = compress_helper(compressor.as_mut(), &buf[offset..]);
1665                    let mut output_buf = Vec::from(&buf[..offset]);
1666                    output_buf.extend_from_slice(&cmp_buf[..]);
1667
1668                    Page::DataPageV2 {
1669                        buf: Bytes::from(output_buf),
1670                        num_values,
1671                        encoding,
1672                        num_nulls,
1673                        num_rows,
1674                        def_levels_byte_len,
1675                        rep_levels_byte_len,
1676                        is_compressed: compressor.is_some(),
1677                        statistics: from_thrift_page_stats(
1678                            physical_type,
1679                            page_stats_to_thrift(statistics.as_ref()),
1680                        )
1681                        .unwrap(),
1682                    }
1683                }
1684                Page::DictionaryPage {
1685                    ref buf,
1686                    num_values,
1687                    encoding,
1688                    is_sorted,
1689                } => {
1690                    let output_buf = compress_helper(compressor.as_mut(), buf);
1691
1692                    Page::DictionaryPage {
1693                        buf: Bytes::from(output_buf),
1694                        num_values,
1695                        encoding,
1696                        is_sorted,
1697                    }
1698                }
1699            };
1700
1701            let compressed_page = CompressedPage::new(compressed_page, uncompressed_len);
1702            compressed_pages.push(compressed_page);
1703        }
1704
1705        let mut buffer: Vec<u8> = vec![];
1706        let mut result_pages: Vec<Page> = vec![];
1707        {
1708            let mut writer = TrackedWrite::new(&mut buffer);
1709            let mut page_writer = SerializedPageWriter::new(&mut writer);
1710
1711            for page in compressed_pages {
1712                page_writer.write_page(page).unwrap();
1713            }
1714            page_writer.close().unwrap();
1715        }
1716        {
1717            let reader = bytes::Bytes::from(buffer);
1718
1719            let t = types::Type::primitive_type_builder("t", physical_type)
1720                .build()
1721                .unwrap();
1722
1723            let desc = ColumnDescriptor::new(Arc::new(t), 0, 0, ColumnPath::new(vec![]));
1724            let meta = ColumnChunkMetaData::builder(Arc::new(desc))
1725                .set_compression_codec(codec.into())
1726                .set_total_compressed_size(reader.len() as i64)
1727                .set_num_values(total_num_values)
1728                .build()
1729                .unwrap();
1730
1731            let props = ReaderProperties::builder()
1732                .set_backward_compatible_lz4(false)
1733                .set_read_page_statistics(true)
1734                .build();
1735            let mut page_reader = SerializedPageReader::new_with_properties(
1736                Arc::new(reader),
1737                &meta,
1738                total_num_values as usize,
1739                None,
1740                Arc::new(props),
1741            )
1742            .unwrap();
1743
1744            while let Some(page) = page_reader.get_next_page().unwrap() {
1745                result_pages.push(page);
1746            }
1747        }
1748
1749        assert_eq!(result_pages.len(), pages.len());
1750        for i in 0..result_pages.len() {
1751            assert_page(&result_pages[i], &pages[i]);
1752        }
1753    }
1754
1755    /// Helper function to compress a slice
1756    fn compress_helper(compressor: Option<&mut Box<dyn Codec>>, data: &[u8]) -> Vec<u8> {
1757        let mut output_buf = vec![];
1758        if let Some(cmpr) = compressor {
1759            cmpr.compress(data, &mut output_buf).unwrap();
1760        } else {
1761            output_buf.extend_from_slice(data);
1762        }
1763        output_buf
1764    }
1765
1766    /// Check if pages match.
1767    fn assert_page(left: &Page, right: &Page) {
1768        assert_eq!(left.page_type(), right.page_type());
1769        assert_eq!(&left.buffer(), &right.buffer());
1770        assert_eq!(left.num_values(), right.num_values());
1771        assert_eq!(left.encoding(), right.encoding());
1772        assert_eq!(
1773            page_stats_to_thrift(left.statistics()),
1774            page_stats_to_thrift(right.statistics())
1775        );
1776    }
1777
1778    /// Tests roundtrip of i32 data written using `W` and read using `R`
1779    fn test_roundtrip_i32<W, R>(
1780        file: W,
1781        data: Vec<Vec<i32>>,
1782        compression: Compression,
1783    ) -> ParquetMetaData
1784    where
1785        W: Write + Send,
1786        R: ChunkReader + From<W> + 'static,
1787    {
1788        test_roundtrip::<W, R, Int32Type, _>(file, data, |r| r.get_int(0).unwrap(), compression)
1789    }
1790
1791    /// Tests roundtrip of data of type `D` written using `W` and read using `R`
1792    /// and the provided `values` function
1793    fn test_roundtrip<W, R, D, F>(
1794        mut file: W,
1795        data: Vec<Vec<D::T>>,
1796        value: F,
1797        compression: Compression,
1798    ) -> ParquetMetaData
1799    where
1800        W: Write + Send,
1801        R: ChunkReader + From<W> + 'static,
1802        D: DataType,
1803        F: Fn(Row) -> D::T,
1804    {
1805        let schema = Arc::new(
1806            types::Type::group_type_builder("schema")
1807                .with_fields(vec![Arc::new(
1808                    types::Type::primitive_type_builder("col1", D::get_physical_type())
1809                        .with_repetition(Repetition::REQUIRED)
1810                        .build()
1811                        .unwrap(),
1812                )])
1813                .build()
1814                .unwrap(),
1815        );
1816        let props = Arc::new(
1817            WriterProperties::builder()
1818                .set_compression(compression)
1819                .build(),
1820        );
1821        let mut file_writer = SerializedFileWriter::new(&mut file, schema, props).unwrap();
1822        let mut rows: i64 = 0;
1823
1824        for (idx, subset) in data.iter().enumerate() {
1825            let row_group_file_offset = file_writer.buf.bytes_written();
1826            let mut row_group_writer = file_writer.next_row_group().unwrap();
1827            if let Some(mut writer) = row_group_writer.next_column().unwrap() {
1828                rows += writer
1829                    .typed::<D>()
1830                    .write_batch(&subset[..], None, None)
1831                    .unwrap() as i64;
1832                writer.close().unwrap();
1833            }
1834            let last_group = row_group_writer.close().unwrap();
1835            let flushed = file_writer.flushed_row_groups();
1836            assert_eq!(flushed.len(), idx + 1);
1837            assert_eq!(Some(idx as i32), last_group.ordinal());
1838            assert_eq!(Some(row_group_file_offset as i64), last_group.file_offset());
1839            assert_eq!(&flushed[idx], last_group.as_ref());
1840        }
1841        let file_metadata = file_writer.close().unwrap();
1842
1843        let reader = SerializedFileReader::new(R::from(file)).unwrap();
1844        assert_eq!(reader.num_row_groups(), data.len());
1845        assert_eq!(
1846            reader.metadata().file_metadata().num_rows(),
1847            rows,
1848            "row count in metadata not equal to number of rows written"
1849        );
1850        for (i, item) in data.iter().enumerate().take(reader.num_row_groups()) {
1851            let row_group_reader = reader.get_row_group(i).unwrap();
1852            let iter = row_group_reader.get_row_iter(None).unwrap();
1853            let res: Vec<_> = iter.map(|row| row.unwrap()).map(&value).collect();
1854            let row_group_size = row_group_reader.metadata().total_byte_size();
1855            let uncompressed_size: i64 = row_group_reader
1856                .metadata()
1857                .columns()
1858                .iter()
1859                .map(|v| v.uncompressed_size())
1860                .sum();
1861            assert_eq!(row_group_size, uncompressed_size);
1862            assert_eq!(res, *item);
1863        }
1864        file_metadata
1865    }
1866
1867    /// File write-read roundtrip.
1868    /// `data` consists of arrays of values for each row group.
1869    fn test_file_roundtrip(file: File, data: Vec<Vec<i32>>) -> ParquetMetaData {
1870        test_roundtrip_i32::<File, File>(file, data, Compression::UNCOMPRESSED)
1871    }
1872
1873    #[test]
1874    fn test_bytes_writer_empty_row_groups() {
1875        test_bytes_roundtrip(vec![], Compression::UNCOMPRESSED);
1876    }
1877
1878    #[test]
1879    fn test_bytes_writer_single_row_group() {
1880        test_bytes_roundtrip(vec![vec![1, 2, 3, 4, 5]], Compression::UNCOMPRESSED);
1881    }
1882
1883    #[test]
1884    fn test_bytes_writer_multiple_row_groups() {
1885        test_bytes_roundtrip(
1886            vec![
1887                vec![1, 2, 3, 4, 5],
1888                vec![1, 2, 3],
1889                vec![1],
1890                vec![1, 2, 3, 4, 5, 6],
1891            ],
1892            Compression::UNCOMPRESSED,
1893        );
1894    }
1895
1896    #[test]
1897    fn test_bytes_writer_single_row_group_compressed() {
1898        test_bytes_roundtrip(vec![vec![1, 2, 3, 4, 5]], Compression::SNAPPY);
1899    }
1900
1901    #[test]
1902    fn test_bytes_writer_multiple_row_groups_compressed() {
1903        test_bytes_roundtrip(
1904            vec![
1905                vec![1, 2, 3, 4, 5],
1906                vec![1, 2, 3],
1907                vec![1],
1908                vec![1, 2, 3, 4, 5, 6],
1909            ],
1910            Compression::SNAPPY,
1911        );
1912    }
1913
1914    fn test_bytes_roundtrip(data: Vec<Vec<i32>>, compression: Compression) {
1915        test_roundtrip_i32::<Vec<u8>, Bytes>(Vec::with_capacity(1024), data, compression);
1916    }
1917
1918    #[test]
1919    fn test_boolean_roundtrip() {
1920        let my_bool_values: Vec<_> = (0..2049).map(|idx| idx % 2 == 0).collect();
1921        test_roundtrip::<Vec<u8>, Bytes, BoolType, _>(
1922            Vec::with_capacity(1024),
1923            vec![my_bool_values],
1924            |r| r.get_bool(0).unwrap(),
1925            Compression::UNCOMPRESSED,
1926        );
1927    }
1928
1929    #[test]
1930    fn test_boolean_compressed_roundtrip() {
1931        let my_bool_values: Vec<_> = (0..2049).map(|idx| idx % 2 == 0).collect();
1932        test_roundtrip::<Vec<u8>, Bytes, BoolType, _>(
1933            Vec::with_capacity(1024),
1934            vec![my_bool_values],
1935            |r| r.get_bool(0).unwrap(),
1936            Compression::SNAPPY,
1937        );
1938    }
1939
1940    #[test]
1941    fn test_column_offset_index_file() {
1942        let file = tempfile::tempfile().unwrap();
1943        let file_metadata = test_file_roundtrip(file, vec![vec![1, 2, 3, 4, 5]]);
1944        file_metadata.row_groups().iter().for_each(|row_group| {
1945            row_group.columns().iter().for_each(|column_chunk| {
1946                assert!(column_chunk.column_index_offset().is_some());
1947                assert!(column_chunk.column_index_length().is_some());
1948                assert!(column_chunk.offset_index_offset().is_some());
1949                assert!(column_chunk.offset_index_length().is_some());
1950            })
1951        });
1952    }
1953
1954    fn test_kv_metadata(initial_kv: Option<Vec<KeyValue>>, final_kv: Option<Vec<KeyValue>>) {
1955        let schema = Arc::new(
1956            types::Type::group_type_builder("schema")
1957                .with_fields(vec![Arc::new(
1958                    types::Type::primitive_type_builder("col1", Type::INT32)
1959                        .with_repetition(Repetition::REQUIRED)
1960                        .build()
1961                        .unwrap(),
1962                )])
1963                .build()
1964                .unwrap(),
1965        );
1966        let mut out = Vec::with_capacity(1024);
1967        let props = Arc::new(
1968            WriterProperties::builder()
1969                .set_key_value_metadata(initial_kv.clone())
1970                .build(),
1971        );
1972        let mut writer = SerializedFileWriter::new(&mut out, schema, props).unwrap();
1973        let mut row_group_writer = writer.next_row_group().unwrap();
1974        let column = row_group_writer.next_column().unwrap().unwrap();
1975        column.close().unwrap();
1976        row_group_writer.close().unwrap();
1977        if let Some(kvs) = &final_kv {
1978            for kv in kvs {
1979                writer.append_key_value_metadata(kv.clone())
1980            }
1981        }
1982        writer.close().unwrap();
1983
1984        let reader = SerializedFileReader::new(Bytes::from(out)).unwrap();
1985        let metadata = reader.metadata().file_metadata();
1986        let keys = metadata.key_value_metadata();
1987
1988        match (initial_kv, final_kv) {
1989            (Some(a), Some(b)) => {
1990                let keys = keys.unwrap();
1991                assert_eq!(keys.len(), a.len() + b.len());
1992                assert_eq!(&keys[..a.len()], a.as_slice());
1993                assert_eq!(&keys[a.len()..], b.as_slice());
1994            }
1995            (Some(v), None) => assert_eq!(keys.unwrap(), &v),
1996            (None, Some(v)) if !v.is_empty() => assert_eq!(keys.unwrap(), &v),
1997            _ => assert!(keys.is_none()),
1998        }
1999    }
2000
2001    #[test]
2002    fn test_append_metadata() {
2003        let kv1 = KeyValue::new("cupcakes".to_string(), "awesome".to_string());
2004        let kv2 = KeyValue::new("bingo".to_string(), "bongo".to_string());
2005
2006        test_kv_metadata(None, None);
2007        test_kv_metadata(Some(vec![kv1.clone()]), None);
2008        test_kv_metadata(None, Some(vec![kv2.clone()]));
2009        test_kv_metadata(Some(vec![kv1.clone()]), Some(vec![kv2.clone()]));
2010        test_kv_metadata(Some(vec![]), Some(vec![kv2]));
2011        test_kv_metadata(Some(vec![]), Some(vec![]));
2012        test_kv_metadata(Some(vec![kv1]), Some(vec![]));
2013        test_kv_metadata(None, Some(vec![]));
2014    }
2015
2016    #[test]
2017    fn test_backwards_compatible_statistics() {
2018        let message_type = "
2019            message test_schema {
2020                REQUIRED INT32 decimal1 (DECIMAL(8,2));
2021                REQUIRED INT32 i32 (INTEGER(32,true));
2022                REQUIRED INT32 u32 (INTEGER(32,false));
2023            }
2024        ";
2025
2026        let schema = Arc::new(parse_message_type(message_type).unwrap());
2027        let props = Default::default();
2028        let mut writer = SerializedFileWriter::new(vec![], schema, props).unwrap();
2029        let mut row_group_writer = writer.next_row_group().unwrap();
2030
2031        for _ in 0..3 {
2032            let mut writer = row_group_writer.next_column().unwrap().unwrap();
2033            writer
2034                .typed::<Int32Type>()
2035                .write_batch(&[1, 2, 3], None, None)
2036                .unwrap();
2037            writer.close().unwrap();
2038        }
2039        let metadata = row_group_writer.close().unwrap();
2040        writer.close().unwrap();
2041
2042        // decimal
2043        let s = page_stats_to_thrift(metadata.column(0).statistics()).unwrap();
2044        assert_eq!(s.min.as_deref(), Some(1_i32.to_le_bytes().as_ref()));
2045        assert_eq!(s.max.as_deref(), Some(3_i32.to_le_bytes().as_ref()));
2046        assert_eq!(s.min_value.as_deref(), Some(1_i32.to_le_bytes().as_ref()));
2047        assert_eq!(s.max_value.as_deref(), Some(3_i32.to_le_bytes().as_ref()));
2048
2049        // i32
2050        let s = page_stats_to_thrift(metadata.column(1).statistics()).unwrap();
2051        assert_eq!(s.min.as_deref(), Some(1_i32.to_le_bytes().as_ref()));
2052        assert_eq!(s.max.as_deref(), Some(3_i32.to_le_bytes().as_ref()));
2053        assert_eq!(s.min_value.as_deref(), Some(1_i32.to_le_bytes().as_ref()));
2054        assert_eq!(s.max_value.as_deref(), Some(3_i32.to_le_bytes().as_ref()));
2055
2056        // u32
2057        let s = page_stats_to_thrift(metadata.column(2).statistics()).unwrap();
2058        assert_eq!(s.min.as_deref(), None);
2059        assert_eq!(s.max.as_deref(), None);
2060        assert_eq!(s.min_value.as_deref(), Some(1_i32.to_le_bytes().as_ref()));
2061        assert_eq!(s.max_value.as_deref(), Some(3_i32.to_le_bytes().as_ref()));
2062    }
2063
2064    #[test]
2065    fn test_spliced_write() {
2066        let message_type = "
2067            message test_schema {
2068                REQUIRED INT32 i32 (INTEGER(32,true));
2069                REQUIRED INT32 u32 (INTEGER(32,false));
2070            }
2071        ";
2072        let schema = Arc::new(parse_message_type(message_type).unwrap());
2073        let props = Arc::new(WriterProperties::builder().build());
2074
2075        let mut file = Vec::with_capacity(1024);
2076        let mut file_writer = SerializedFileWriter::new(&mut file, schema, props.clone()).unwrap();
2077
2078        let columns = file_writer.descr.columns();
2079        let mut column_state: Vec<(_, Option<ColumnCloseResult>)> = columns
2080            .iter()
2081            .map(|_| (TrackedWrite::new(Vec::with_capacity(1024)), None))
2082            .collect();
2083
2084        let mut column_state_slice = column_state.as_mut_slice();
2085        let mut column_writers = Vec::with_capacity(columns.len());
2086        for c in columns {
2087            let ((buf, out), tail) = column_state_slice.split_first_mut().unwrap();
2088            column_state_slice = tail;
2089
2090            let page_writer = Box::new(SerializedPageWriter::new(buf));
2091            let col_writer = get_column_writer(c.clone(), props.clone(), page_writer);
2092            column_writers.push(SerializedColumnWriter::new(
2093                col_writer,
2094                Some(Box::new(|on_close| {
2095                    *out = Some(on_close);
2096                    Ok(())
2097                })),
2098            ));
2099        }
2100
2101        let column_data = [[1, 2, 3, 4], [7, 3, 7, 3]];
2102
2103        // Interleaved writing to the column writers
2104        for (writer, batch) in column_writers.iter_mut().zip(column_data) {
2105            let writer = writer.typed::<Int32Type>();
2106            writer.write_batch(&batch, None, None).unwrap();
2107        }
2108
2109        // Close the column writers
2110        for writer in column_writers {
2111            writer.close().unwrap()
2112        }
2113
2114        // Splice column data into a row group
2115        let mut row_group_writer = file_writer.next_row_group().unwrap();
2116        for (write, close) in column_state {
2117            let buf = Bytes::from(write.into_inner().unwrap());
2118            row_group_writer
2119                .append_column(&buf, close.unwrap())
2120                .unwrap();
2121        }
2122        row_group_writer.close().unwrap();
2123        file_writer.close().unwrap();
2124
2125        // Check data was written correctly
2126        let file = Bytes::from(file);
2127        let test_read = |reader: SerializedFileReader<Bytes>| {
2128            let row_group = reader.get_row_group(0).unwrap();
2129
2130            let mut out = Vec::with_capacity(4);
2131            let c1 = row_group.get_column_reader(0).unwrap();
2132            let mut c1 = get_typed_column_reader::<Int32Type>(c1);
2133            c1.read_records(4, None, None, &mut out).unwrap();
2134            assert_eq!(out, column_data[0]);
2135
2136            out.clear();
2137
2138            let c2 = row_group.get_column_reader(1).unwrap();
2139            let mut c2 = get_typed_column_reader::<Int32Type>(c2);
2140            c2.read_records(4, None, None, &mut out).unwrap();
2141            assert_eq!(out, column_data[1]);
2142        };
2143
2144        let reader = SerializedFileReader::new(file.clone()).unwrap();
2145        test_read(reader);
2146
2147        let options = ReadOptionsBuilder::new().with_page_index().build();
2148        let reader = SerializedFileReader::new_with_options(file, options).unwrap();
2149        test_read(reader);
2150    }
2151
2152    #[test]
2153    fn test_disabled_statistics() {
2154        let message_type = "
2155            message test_schema {
2156                REQUIRED INT32 a;
2157                REQUIRED INT32 b;
2158            }
2159        ";
2160        let schema = Arc::new(parse_message_type(message_type).unwrap());
2161        let props = WriterProperties::builder()
2162            .set_statistics_enabled(EnabledStatistics::None)
2163            .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
2164            .set_offset_index_disabled(true) // this should be ignored because of the line above
2165            .build();
2166        let mut file = Vec::with_capacity(1024);
2167        let mut file_writer =
2168            SerializedFileWriter::new(&mut file, schema, Arc::new(props)).unwrap();
2169
2170        let mut row_group_writer = file_writer.next_row_group().unwrap();
2171        let mut a_writer = row_group_writer.next_column().unwrap().unwrap();
2172        let col_writer = a_writer.typed::<Int32Type>();
2173        col_writer.write_batch(&[1, 2, 3], None, None).unwrap();
2174        a_writer.close().unwrap();
2175
2176        let mut b_writer = row_group_writer.next_column().unwrap().unwrap();
2177        let col_writer = b_writer.typed::<Int32Type>();
2178        col_writer.write_batch(&[4, 5, 6], None, None).unwrap();
2179        b_writer.close().unwrap();
2180        row_group_writer.close().unwrap();
2181
2182        let metadata = file_writer.finish().unwrap();
2183        assert_eq!(metadata.num_row_groups(), 1);
2184        let row_group = metadata.row_group(0);
2185        assert_eq!(row_group.num_columns(), 2);
2186        // Column "a" has both offset and column index, as requested
2187        assert!(row_group.column(0).offset_index_offset().is_some());
2188        assert!(row_group.column(0).column_index_offset().is_some());
2189        // Column "b" should only have offset index
2190        assert!(row_group.column(1).offset_index_offset().is_some());
2191        assert!(row_group.column(1).column_index_offset().is_none());
2192
2193        let err = file_writer.next_row_group().err().unwrap().to_string();
2194        assert_eq!(err, "Parquet error: SerializedFileWriter already finished");
2195
2196        drop(file_writer);
2197
2198        let options = ReadOptionsBuilder::new().with_page_index().build();
2199        let reader = SerializedFileReader::new_with_options(Bytes::from(file), options).unwrap();
2200
2201        let a_idx = reader.metadata().page_index().unwrap().column_index(0, 0);
2202        assert!(
2203            matches!(a_idx, Some(ColumnIndexMetaData::INT32(_))),
2204            "{a_idx:?}"
2205        );
2206        let b_idx = reader.metadata().page_index().unwrap().column_index(0, 1);
2207        assert!(b_idx.is_none(), "{b_idx:?}");
2208    }
2209
2210    #[test]
2211    fn test_offset_index_disabled() {
2212        let message_type = "
2213            message test_schema {
2214                REQUIRED INT32 a;
2215                REQUIRED INT32 b;
2216            }
2217        ";
2218        // write file with indexes disabled (including offset indexes)
2219        let schema = Arc::new(parse_message_type(message_type).unwrap());
2220        let props = WriterProperties::builder()
2221            .set_statistics_enabled(EnabledStatistics::None)
2222            .set_offset_index_disabled(true)
2223            .build();
2224        let mut file = Vec::with_capacity(1024);
2225        let mut file_writer =
2226            SerializedFileWriter::new(&mut file, schema, Arc::new(props)).unwrap();
2227
2228        let mut row_group_writer = file_writer.next_row_group().unwrap();
2229        let mut a_writer = row_group_writer.next_column().unwrap().unwrap();
2230        let col_writer = a_writer.typed::<Int32Type>();
2231        col_writer.write_batch(&[1, 2, 3], None, None).unwrap();
2232        a_writer.close().unwrap();
2233
2234        let mut b_writer = row_group_writer.next_column().unwrap().unwrap();
2235        let col_writer = b_writer.typed::<Int32Type>();
2236        col_writer.write_batch(&[4, 5, 6], None, None).unwrap();
2237        b_writer.close().unwrap();
2238        row_group_writer.close().unwrap();
2239
2240        let metadata = file_writer.finish().unwrap();
2241        assert_eq!(metadata.num_row_groups(), 1);
2242        let row_group = metadata.row_group(0);
2243        assert_eq!(row_group.num_columns(), 2);
2244        // no page indexes should exist
2245        assert!(row_group.column(0).offset_index_offset().is_none());
2246        assert!(row_group.column(0).column_index_offset().is_none());
2247        assert!(row_group.column(1).offset_index_offset().is_none());
2248        assert!(row_group.column(1).column_index_offset().is_none());
2249
2250        drop(file_writer);
2251
2252        // read file and request page index...should be `None`
2253        let options = ReadOptionsBuilder::new().with_page_index().build();
2254        let reader = SerializedFileReader::new_with_options(Bytes::from(file), options).unwrap();
2255
2256        assert!(reader.metadata().page_index().is_none());
2257    }
2258
2259    #[test]
2260    fn test_byte_array_size_statistics() {
2261        let message_type = "
2262            message test_schema {
2263                OPTIONAL BYTE_ARRAY a (UTF8);
2264            }
2265        ";
2266        let schema = Arc::new(parse_message_type(message_type).unwrap());
2267        let data = ByteArrayType::gen_vec(32, 7);
2268        let def_levels = [1, 1, 1, 1, 0, 1, 0, 1, 0, 1];
2269        let unenc_size: i64 = data.iter().map(|x| x.len() as i64).sum();
2270        let file: File = tempfile::tempfile().unwrap();
2271        let props = Arc::new(
2272            WriterProperties::builder()
2273                .set_statistics_enabled(EnabledStatistics::Page)
2274                .build(),
2275        );
2276
2277        let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
2278        let mut row_group_writer = writer.next_row_group().unwrap();
2279
2280        let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
2281        col_writer
2282            .typed::<ByteArrayType>()
2283            .write_batch(&data, Some(&def_levels), None)
2284            .unwrap();
2285        col_writer.close().unwrap();
2286        row_group_writer.close().unwrap();
2287        let file_metadata = writer.close().unwrap();
2288
2289        assert_eq!(file_metadata.num_row_groups(), 1);
2290        assert_eq!(file_metadata.row_group(0).num_columns(), 1);
2291
2292        let check_def_hist = |def_hist: &[i64]| {
2293            assert_eq!(def_hist.len(), 2);
2294            assert_eq!(def_hist[0], 3);
2295            assert_eq!(def_hist[1], 7);
2296        };
2297
2298        let meta_data = file_metadata.row_group(0).column(0);
2299
2300        assert!(meta_data.repetition_level_histogram().is_none());
2301        assert!(meta_data.definition_level_histogram().is_some());
2302        assert!(meta_data.unencoded_byte_array_data_bytes().is_some());
2303        assert_eq!(
2304            unenc_size,
2305            meta_data.unencoded_byte_array_data_bytes().unwrap()
2306        );
2307        check_def_hist(meta_data.definition_level_histogram().unwrap().values());
2308
2309        // check that the read metadata is also correct
2310        let options = ReadOptionsBuilder::new().with_page_index().build();
2311        let reader = SerializedFileReader::new_with_options(file, options).unwrap();
2312
2313        let rfile_metadata = reader.metadata().file_metadata();
2314        assert_eq!(
2315            rfile_metadata.num_rows(),
2316            file_metadata.file_metadata().num_rows()
2317        );
2318        assert_eq!(reader.num_row_groups(), 1);
2319        let rowgroup = reader.get_row_group(0).unwrap();
2320        assert_eq!(rowgroup.num_columns(), 1);
2321        let column = rowgroup.metadata().column(0);
2322        assert!(column.definition_level_histogram().is_some());
2323        assert!(column.repetition_level_histogram().is_none());
2324        assert!(column.unencoded_byte_array_data_bytes().is_some());
2325        check_def_hist(column.definition_level_histogram().unwrap().values());
2326        assert_eq!(
2327            unenc_size,
2328            column.unencoded_byte_array_data_bytes().unwrap()
2329        );
2330
2331        // check histogram in column index as well
2332        assert!(reader.metadata().page_index().is_some());
2333        let page_index = reader.metadata().page_index().unwrap();
2334        let col_idx =
2335            if let Some(ColumnIndexMetaData::BYTE_ARRAY(index)) = page_index.column_index(0, 0) {
2336                assert_eq!(index.num_pages(), 1);
2337                index
2338            } else {
2339                unreachable!()
2340            };
2341
2342        assert!(col_idx.repetition_level_histogram(0).is_none());
2343        assert!(col_idx.definition_level_histogram(0).is_some());
2344        check_def_hist(col_idx.definition_level_histogram(0).unwrap());
2345
2346        assert!(page_index.offset_index(0, 0).is_some());
2347        assert!(
2348            page_index
2349                .offset_index(0, 0)
2350                .unwrap()
2351                .unencoded_byte_array_data_bytes
2352                .is_some()
2353        );
2354        let page_sizes = page_index
2355            .offset_index(0, 0)
2356            .unwrap()
2357            .unencoded_byte_array_data_bytes
2358            .as_ref()
2359            .unwrap();
2360        assert_eq!(page_sizes.len(), 1);
2361        assert_eq!(page_sizes[0], unenc_size);
2362    }
2363
2364    #[cfg(feature = "encryption")]
2365    #[test]
2366    fn test_too_many_rowgroups() {
2367        let message_type = "
2368            message test_schema {
2369                REQUIRED BYTE_ARRAY a (UTF8);
2370            }
2371        ";
2372        let schema = Arc::new(parse_message_type(message_type).unwrap());
2373        let file: File = tempfile::tempfile().unwrap();
2374
2375        const AES_128_FOOTER_KEY: &[u8; 16] = b"0123456789012345"; // 128bit/16
2376        let footer_key = AES_128_FOOTER_KEY;
2377        let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec())
2378            .build()
2379            .unwrap();
2380        let props = Arc::new(
2381            WriterProperties::builder()
2382                .set_statistics_enabled(EnabledStatistics::None)
2383                .set_max_row_group_row_count(Some(1))
2384                .with_file_encryption_properties(file_encryption_properties)
2385                .build(),
2386        );
2387        let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
2388
2389        // Create 32k empty rowgroups. Should error when i == 32768.
2390        for i in 0..0x8001 {
2391            match writer.next_row_group() {
2392                Ok(mut row_group_writer) => {
2393                    assert_ne!(i, 0x8000);
2394                    let col_writer = row_group_writer.next_column().unwrap().unwrap();
2395                    col_writer.close().unwrap();
2396                    row_group_writer.close().unwrap();
2397                }
2398                Err(e) => {
2399                    assert_eq!(i, 0x8000);
2400                    assert_eq!(
2401                        e.to_string(),
2402                        "Parquet error: Parquet with encryption does not support more than 32767 row groups per file (currently: 32768)"
2403                    );
2404                }
2405            }
2406        }
2407        writer.close().unwrap();
2408    }
2409
2410    #[test]
2411    #[cfg_attr(miri, ignore)] // Takes too long
2412    fn test_32k_rowgroups() {
2413        let message_type = "
2414            message test_schema {
2415                REQUIRED BYTE_ARRAY a (UTF8);
2416            }
2417        ";
2418        let schema = Arc::new(parse_message_type(message_type).unwrap());
2419        let file: File = tempfile::tempfile().unwrap();
2420        let props = Arc::new(
2421            WriterProperties::builder()
2422                .set_statistics_enabled(EnabledStatistics::None)
2423                .set_max_row_group_row_count(Some(1))
2424                .build(),
2425        );
2426        let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
2427
2428        // Create 32k + 1 empty rowgroups. No row group ordinals should be written (but we can't
2429        // test for that).
2430        for _ in 0..0x8001 {
2431            let mut row_group_writer = writer.next_row_group().unwrap();
2432            let col_writer = row_group_writer.next_column().unwrap().unwrap();
2433            col_writer.close().unwrap();
2434            row_group_writer.close().unwrap();
2435        }
2436        writer.close().unwrap();
2437
2438        // Parse the written metadata and check that ordinals were replaced.
2439        let reader = SerializedFileReader::new(file).unwrap();
2440        let metadata = reader.metadata();
2441
2442        for (i, rg) in metadata.row_groups().iter().enumerate() {
2443            assert_eq!(i as i32, rg.ordinal().unwrap());
2444        }
2445    }
2446
2447    #[test]
2448    fn test_size_statistics_with_repetition_and_nulls() {
2449        let message_type = "
2450            message test_schema {
2451                OPTIONAL group i32_list (LIST) {
2452                    REPEATED group list {
2453                        OPTIONAL INT32 element;
2454                    }
2455                }
2456            }
2457        ";
2458        // column is:
2459        // row 0: [1, 2]
2460        // row 1: NULL
2461        // row 2: [4, NULL]
2462        // row 3: []
2463        // row 4: [7, 8, 9, 10]
2464        let schema = Arc::new(parse_message_type(message_type).unwrap());
2465        let data = [1, 2, 4, 7, 8, 9, 10];
2466        let def_levels = [3, 3, 0, 3, 2, 1, 3, 3, 3, 3];
2467        let rep_levels = [0, 1, 0, 0, 1, 0, 0, 1, 1, 1];
2468        let file = tempfile::tempfile().unwrap();
2469        let props = Arc::new(
2470            WriterProperties::builder()
2471                .set_statistics_enabled(EnabledStatistics::Page)
2472                .build(),
2473        );
2474        let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
2475        let mut row_group_writer = writer.next_row_group().unwrap();
2476
2477        let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
2478        col_writer
2479            .typed::<Int32Type>()
2480            .write_batch(&data, Some(&def_levels), Some(&rep_levels))
2481            .unwrap();
2482        col_writer.close().unwrap();
2483        row_group_writer.close().unwrap();
2484        let file_metadata = writer.close().unwrap();
2485
2486        assert_eq!(file_metadata.num_row_groups(), 1);
2487        assert_eq!(file_metadata.row_group(0).num_columns(), 1);
2488
2489        let check_def_hist = |def_hist: &[i64]| {
2490            assert_eq!(def_hist.len(), 4);
2491            assert_eq!(def_hist[0], 1);
2492            assert_eq!(def_hist[1], 1);
2493            assert_eq!(def_hist[2], 1);
2494            assert_eq!(def_hist[3], 7);
2495        };
2496
2497        let check_rep_hist = |rep_hist: &[i64]| {
2498            assert_eq!(rep_hist.len(), 2);
2499            assert_eq!(rep_hist[0], 5);
2500            assert_eq!(rep_hist[1], 5);
2501        };
2502
2503        // check that histograms are set properly in the write and read metadata
2504        // also check that unencoded_byte_array_data_bytes is not set
2505        let meta_data = file_metadata.row_group(0).column(0);
2506        assert!(meta_data.repetition_level_histogram().is_some());
2507        assert!(meta_data.definition_level_histogram().is_some());
2508        assert!(meta_data.unencoded_byte_array_data_bytes().is_none());
2509        check_def_hist(meta_data.definition_level_histogram().unwrap().values());
2510        check_rep_hist(meta_data.repetition_level_histogram().unwrap().values());
2511
2512        // check that the read metadata is also correct
2513        let options = ReadOptionsBuilder::new().with_page_index().build();
2514        let reader = SerializedFileReader::new_with_options(file, options).unwrap();
2515
2516        let rfile_metadata = reader.metadata().file_metadata();
2517        assert_eq!(
2518            rfile_metadata.num_rows(),
2519            file_metadata.file_metadata().num_rows()
2520        );
2521        assert_eq!(reader.num_row_groups(), 1);
2522        let rowgroup = reader.get_row_group(0).unwrap();
2523        assert_eq!(rowgroup.num_columns(), 1);
2524        let column = rowgroup.metadata().column(0);
2525        assert!(column.definition_level_histogram().is_some());
2526        assert!(column.repetition_level_histogram().is_some());
2527        assert!(column.unencoded_byte_array_data_bytes().is_none());
2528        check_def_hist(column.definition_level_histogram().unwrap().values());
2529        check_rep_hist(column.repetition_level_histogram().unwrap().values());
2530
2531        assert!(
2532            reader
2533                .metadata()
2534                .page_index()
2535                .is_some_and(|idx| idx.is_complete())
2536        );
2537        let page_index = reader.metadata().page_index().unwrap();
2538
2539        // check histogram in column index as well
2540        let col_idx = if let Some(ColumnIndexMetaData::INT32(index)) = page_index.column_index(0, 0)
2541        {
2542            assert_eq!(index.num_pages(), 1);
2543            index
2544        } else {
2545            unreachable!()
2546        };
2547
2548        check_def_hist(col_idx.definition_level_histogram(0).unwrap());
2549        check_rep_hist(col_idx.repetition_level_histogram(0).unwrap());
2550
2551        assert!(
2552            page_index
2553                .offset_index(0, 0)
2554                .unwrap()
2555                .unencoded_byte_array_data_bytes
2556                .is_none()
2557        );
2558    }
2559
2560    #[test]
2561    #[cfg_attr(miri, ignore)] // Takes too long
2562    #[cfg(feature = "arrow")]
2563    fn test_byte_stream_split_extended_roundtrip() {
2564        let path = format!(
2565            "{}/byte_stream_split_extended.gzip.parquet",
2566            arrow::util::test_util::parquet_test_data(),
2567        );
2568        let file = File::open(path).unwrap();
2569
2570        // Read in test file and rewrite to tmp
2571        let parquet_reader = ParquetRecordBatchReaderBuilder::try_new(file)
2572            .expect("parquet open")
2573            .build()
2574            .expect("parquet open");
2575
2576        let file = tempfile::tempfile().unwrap();
2577        let props = WriterProperties::builder()
2578            .set_dictionary_enabled(false)
2579            .set_column_encoding(
2580                ColumnPath::from("float16_byte_stream_split"),
2581                Encoding::BYTE_STREAM_SPLIT,
2582            )
2583            .set_column_encoding(
2584                ColumnPath::from("float_byte_stream_split"),
2585                Encoding::BYTE_STREAM_SPLIT,
2586            )
2587            .set_column_encoding(
2588                ColumnPath::from("double_byte_stream_split"),
2589                Encoding::BYTE_STREAM_SPLIT,
2590            )
2591            .set_column_encoding(
2592                ColumnPath::from("int32_byte_stream_split"),
2593                Encoding::BYTE_STREAM_SPLIT,
2594            )
2595            .set_column_encoding(
2596                ColumnPath::from("int64_byte_stream_split"),
2597                Encoding::BYTE_STREAM_SPLIT,
2598            )
2599            .set_column_encoding(
2600                ColumnPath::from("flba5_byte_stream_split"),
2601                Encoding::BYTE_STREAM_SPLIT,
2602            )
2603            .set_column_encoding(
2604                ColumnPath::from("decimal_byte_stream_split"),
2605                Encoding::BYTE_STREAM_SPLIT,
2606            )
2607            .build();
2608
2609        let mut parquet_writer = ArrowWriter::try_new(
2610            file.try_clone().expect("cannot open file"),
2611            parquet_reader.schema(),
2612            Some(props),
2613        )
2614        .expect("create arrow writer");
2615
2616        for maybe_batch in parquet_reader {
2617            let batch = maybe_batch.expect("reading batch");
2618            parquet_writer.write(&batch).expect("writing data");
2619        }
2620
2621        parquet_writer.close().expect("finalizing file");
2622
2623        let reader = SerializedFileReader::new(file).expect("Failed to create reader");
2624        let filemeta = reader.metadata();
2625
2626        // Make sure byte_stream_split encoding was used
2627        let check_encoding = |x: usize, filemeta: &ParquetMetaData| {
2628            assert!(
2629                filemeta
2630                    .row_group(0)
2631                    .column(x)
2632                    .encodings()
2633                    .collect::<Vec<_>>()
2634                    .contains(&Encoding::BYTE_STREAM_SPLIT)
2635            );
2636        };
2637
2638        check_encoding(1, filemeta);
2639        check_encoding(3, filemeta);
2640        check_encoding(5, filemeta);
2641        check_encoding(7, filemeta);
2642        check_encoding(9, filemeta);
2643        check_encoding(11, filemeta);
2644        check_encoding(13, filemeta);
2645
2646        // Read back tmpfile and make sure all values are correct
2647        let mut iter = reader
2648            .get_row_iter(None)
2649            .expect("Failed to create row iterator");
2650
2651        let mut start = 0;
2652        let end = reader.metadata().file_metadata().num_rows();
2653
2654        let check_row = |row: Result<Row, ParquetError>| {
2655            assert!(row.is_ok());
2656            let r = row.unwrap();
2657            assert_eq!(r.get_float16(0).unwrap(), r.get_float16(1).unwrap());
2658            assert_eq!(r.get_float(2).unwrap(), r.get_float(3).unwrap());
2659            assert_eq!(r.get_double(4).unwrap(), r.get_double(5).unwrap());
2660            assert_eq!(r.get_int(6).unwrap(), r.get_int(7).unwrap());
2661            assert_eq!(r.get_long(8).unwrap(), r.get_long(9).unwrap());
2662            assert_eq!(r.get_bytes(10).unwrap(), r.get_bytes(11).unwrap());
2663            assert_eq!(r.get_decimal(12).unwrap(), r.get_decimal(13).unwrap());
2664        };
2665
2666        while start < end {
2667            match iter.next() {
2668                Some(row) => check_row(row),
2669                None => break,
2670            }
2671            start += 1;
2672        }
2673    }
2674
2675    #[test]
2676    #[cfg_attr(miri, ignore)] // Takes too long
2677    fn test_rewrite_no_page_indexes() {
2678        let file = get_test_file("alltypes_tiny_pages.parquet");
2679        let metadata = ParquetMetaDataReader::new()
2680            .with_page_index_policy(PageIndexPolicy::Optional)
2681            .parse_and_finish(&file)
2682            .unwrap();
2683
2684        let props = Arc::new(WriterProperties::builder().build());
2685        let schema = metadata.file_metadata().schema_descr().root_schema_ptr();
2686        let output = Vec::<u8>::new();
2687        let mut writer = SerializedFileWriter::new(output, schema, props).unwrap();
2688
2689        for rg in metadata.row_groups() {
2690            let mut rg_out = writer.next_row_group().unwrap();
2691            for column in rg.columns() {
2692                let result = ColumnCloseResult {
2693                    bytes_written: column.compressed_size() as _,
2694                    rows_written: rg.num_rows() as _,
2695                    metadata: column.clone(),
2696                    bloom_filter: None,
2697                    column_index: None,
2698                    offset_index: None,
2699                };
2700                rg_out.append_column(&file, result).unwrap();
2701            }
2702            rg_out.close().unwrap();
2703        }
2704        writer.close().unwrap();
2705    }
2706
2707    #[test]
2708    #[cfg_attr(miri, ignore)] // Takes too long
2709    fn test_rewrite_missing_column_index() {
2710        // this file has an INT96 column that lacks a column index entry
2711        let file = get_test_file("alltypes_tiny_pages.parquet");
2712        let metadata = ParquetMetaDataReader::new()
2713            .with_page_index_policy(PageIndexPolicy::Optional)
2714            .parse_and_finish(&file)
2715            .unwrap();
2716
2717        let props = Arc::new(WriterProperties::builder().build());
2718        let schema = metadata.file_metadata().schema_descr().root_schema_ptr();
2719        let output = Vec::<u8>::new();
2720        let mut writer = SerializedFileWriter::new(output, schema, props).unwrap();
2721
2722        for (rg_idx, rg) in metadata.row_groups().iter().enumerate() {
2723            let rg_page_index = metadata.page_index_for_row_group(rg_idx);
2724            let mut rg_out = writer.next_row_group().unwrap();
2725            for (col_idx, column) in rg.columns().iter().enumerate() {
2726                let column_index = rg_page_index.column_index(col_idx).cloned();
2727                let offset_index = rg_page_index.offset_index(col_idx).cloned();
2728
2729                let result = ColumnCloseResult {
2730                    bytes_written: column.compressed_size() as _,
2731                    rows_written: rg.num_rows() as _,
2732                    metadata: column.clone(),
2733                    bloom_filter: None,
2734                    column_index,
2735                    offset_index,
2736                };
2737                rg_out.append_column(&file, result).unwrap();
2738            }
2739            rg_out.close().unwrap();
2740        }
2741        writer.close().unwrap();
2742    }
2743
2744    #[test]
2745    fn test_int96_interop() {
2746        // this file has an INT96 column. rewrite it with min/max statistics sorted per
2747        // recent changes to the spec. (see https://github.com/apache/parquet-format/pull/584)
2748        let file = get_test_file("int96_timestamp_order.parquet");
2749        let read_opts = ReadOptionsBuilder::new().with_page_index().build();
2750        let reader = SerializedFileReader::new_with_options(file, read_opts).unwrap();
2751        let file_metadata = reader.metadata().file_metadata();
2752        let schema = file_metadata.schema_descr().root_schema_ptr();
2753
2754        // helper function to extract Int96 min/max from column metadata and the column index
2755        fn retrieve_stats(metadata: &ParquetMetaData) -> (&[u8], &[u8], &Int96, &Int96) {
2756            // sanity check that the proper column order is specified
2757            let column_orders = metadata
2758                .file_metadata()
2759                .column_orders()
2760                .expect("column_orders is missing");
2761            assert_eq!(column_orders[0], ColumnOrder::INT96_TIMESTAMP_ORDER);
2762
2763            let stats = metadata
2764                .row_group(0)
2765                .column(0)
2766                .statistics()
2767                .expect("statistics missing");
2768            let min = stats.min_bytes_opt().expect("min stats missing");
2769            let max = stats.max_bytes_opt().expect("max stats missing");
2770
2771            assert!(
2772                metadata
2773                    .page_index()
2774                    .is_some_and(|pi| pi.has_column_indexes())
2775            );
2776            let Some(ColumnIndexMetaData::INT96(col0)) =
2777                metadata.page_index().unwrap().column_index(0, 0)
2778            else {
2779                panic!("expected INT96 stats")
2780            };
2781            let col_min = col0.min_value(0).expect("ColumnIndex min not present");
2782            let col_max = col0.max_value(0).expect("ColumnIndex max not present");
2783
2784            (min, max, col_min, col_max)
2785        }
2786
2787        // save read stats for later
2788        let (exp_min, exp_max, exp_col_min, exp_col_max) = retrieve_stats(reader.metadata());
2789
2790        // write file back out again
2791        let props = Arc::new(WriterProperties::builder().build());
2792        let output = Vec::<u8>::new();
2793        let mut writer = SerializedFileWriter::new(output, schema, props).unwrap();
2794
2795        let mut rg_out = writer.next_row_group().unwrap();
2796        let rg_in = reader.get_row_group(0).unwrap();
2797
2798        // int96 is column 0
2799        let col_in = rg_in.get_column_reader(0).unwrap();
2800        let mut typed_in = get_typed_column_reader::<Int96Type>(col_in);
2801
2802        let mut values = Vec::new();
2803        typed_in.read_records(4, None, None, &mut values).unwrap();
2804
2805        let mut col_out = rg_out.next_column().unwrap().unwrap();
2806        col_out
2807            .typed::<Int96Type>()
2808            .write_batch(&values, None, None)
2809            .unwrap();
2810        col_out.close().unwrap();
2811        rg_out.close().unwrap();
2812
2813        let new_metadata = writer.close().unwrap();
2814
2815        // check that new stats match the original stats
2816        let (new_min, new_max, new_col_min, new_col_max) = retrieve_stats(&new_metadata);
2817        assert_eq!(new_min, exp_min);
2818        assert_eq!(new_max, exp_max);
2819        assert_eq!(new_col_min, exp_col_min);
2820        assert_eq!(new_col_max, exp_col_max);
2821    }
2822
2823    #[test]
2824    fn test_empty_row_group_data_page_offset() {
2825        // Write a parquet file with one row group and one column but no data written.
2826        // With no rows written no pages are flushed, so sizes should be zero.
2827        let schema = Arc::new(
2828            types::Type::group_type_builder("schema")
2829                .with_fields(vec![Arc::new(
2830                    types::Type::primitive_type_builder("col1", Type::INT32)
2831                        .build()
2832                        .unwrap(),
2833                )])
2834                .build()
2835                .unwrap(),
2836        );
2837        let props = Default::default();
2838        let mut buf = Vec::new();
2839        let mut writer =
2840            SerializedFileWriter::new(std::io::Cursor::new(&mut buf), schema, props).unwrap();
2841
2842        let mut row_group = writer.next_row_group().unwrap();
2843        let col = row_group.next_column().unwrap().unwrap();
2844        col.close().unwrap();
2845        row_group.close().unwrap();
2846        let metadata = writer.close().unwrap();
2847
2848        assert_eq!(metadata.num_row_groups(), 1);
2849        let col_meta = metadata.row_group(0).column(0);
2850        // No pages written: sizes are zero.
2851        assert_eq!(col_meta.compressed_size(), 0);
2852        assert_eq!(col_meta.uncompressed_size(), 0);
2853        assert_eq!(metadata.row_group(0).compressed_size(), 0);
2854        // No dictionary page written.
2855        assert_eq!(col_meta.dictionary_page_offset(), None);
2856        // data_page_offset is 0. The field is required,
2857        // but 0 would be invalid since it points to the PAR1 magic.
2858        assert_eq!(col_meta.data_page_offset(), 0);
2859        // No pages written: no encodings recorded.
2860        assert_eq!(col_meta.encodings().collect::<Vec<_>>(), vec![]);
2861    }
2862}