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