1use 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
51pub struct TrackedWrite<W: Write> {
55 inner: BufWriter<W>,
56 bytes_written: usize,
57}
58
59impl<W: Write> TrackedWrite<W> {
60 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 pub fn bytes_written(&self) -> usize {
71 self.bytes_written
72 }
73
74 pub fn inner(&self) -> &W {
76 self.inner.get_ref()
77 }
78
79 pub fn inner_mut(&mut self) -> &mut W {
84 self.inner.get_mut()
85 }
86
87 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
120pub type OnCloseColumnChunk<'a> = Box<dyn FnOnce(ColumnCloseResult) -> Result<()> + 'a>;
122
123pub 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
140pub 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: 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 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 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 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 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 #[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 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 pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] {
305 &self.row_groups
306 }
307
308 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 pub fn close(mut self) -> Result<ParquetMetaData> {
322 self.finish()
323 }
324
325 #[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 #[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 fn write_metadata(&mut self) -> Result<ParquetMetaData> {
344 self.finished = true;
345
346 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 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 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 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 pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
421 self.kv_metadatas.push(kv_metadata);
422 }
423
424 pub fn schema_descr(&self) -> &SchemaDescriptor {
426 &self.descr
427 }
428
429 #[cfg(feature = "arrow")]
431 pub(crate) fn schema_descr_ptr(&self) -> &SchemaDescPtr {
432 &self.descr
433 }
434
435 pub fn properties(&self) -> &WriterPropertiesPtr {
437 &self.props
438 }
439
440 pub fn inner(&self) -> &W {
442 self.buf.inner()
443 }
444
445 pub fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
454 self.buf.write_all(buf)
455 }
456
457 pub fn flush(&mut self) -> std::io::Result<()> {
459 self.buf.flush()
460 }
461
462 pub fn inner_mut(&mut self) -> &mut W {
471 self.buf.inner_mut()
472 }
473
474 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 pub fn bytes_written(&self) -> usize {
484 self.buf.bytes_written()
485 }
486
487 #[cfg(feature = "encryption")]
489 pub(crate) fn file_encryptor(&self) -> Option<Arc<FileEncryptor>> {
490 self.file_encryptor.clone()
491 }
492}
493
494fn 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 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 *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
533pub 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 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 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 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 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 *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 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 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 pub fn append_column<R: ChunkReader>(
717 &mut self,
718 reader: &R,
719 close: ColumnCloseResult,
720 ) -> Result<()> {
721 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 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 #[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 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 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 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 #[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 #[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 #[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 #[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 #[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#[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
1003pub struct SerializedColumnWriter<'a> {
1005 inner: ColumnWriter<'a>,
1006 on_close: Option<OnCloseColumnChunk<'a>>,
1007}
1008
1009impl<'a> SerializedColumnWriter<'a> {
1010 pub fn new(inner: ColumnWriter<'a>, on_close: Option<OnCloseColumnChunk<'a>>) -> Self {
1013 Self { inner, on_close }
1014 }
1015
1016 pub fn untyped(&mut self) -> &mut ColumnWriter<'a> {
1018 &mut self.inner
1019 }
1020
1021 pub fn typed<T: DataType>(&mut self) -> &mut ColumnWriterImpl<'a, T> {
1023 get_typed_column_writer_mut(&mut self.inner)
1024 }
1025
1026 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
1037pub 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 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 #[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 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#[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#[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 let expected = vec![
1365 ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED),
1367 ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED),
1369 ColumnOrder::IEEE_754_TOTAL_ORDER,
1371 ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED),
1373 ColumnOrder::IEEE_754_TOTAL_ORDER,
1375 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 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 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 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 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 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 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 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 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 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 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 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 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 for writer in column_writers {
2124 writer.close().unwrap()
2125 }
2126
2127 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 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) .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 assert!(row_group.column(0).offset_index_offset().is_some());
2201 assert!(row_group.column(0).column_index_offset().is_some());
2202 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 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 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 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 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 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"; 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 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)] 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 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 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 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 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 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 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)] #[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 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 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 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)] 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)] fn test_rewrite_missing_column_index() {
2723 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 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 fn retrieve_stats(metadata: &ParquetMetaData) -> (&[u8], &[u8], &Int96, &Int96) {
2769 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 let (exp_min, exp_max, exp_col_min, exp_col_max) = retrieve_stats(reader.metadata());
2802
2803 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 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 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}