Skip to main content

parquet/file/
writer.rs

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