1mod footer_tail;
53mod memory;
54mod options;
55pub mod page_index;
56mod parser;
57mod push_decoder;
58pub(crate) mod reader;
59pub(crate) mod thrift;
60mod writer;
61
62use crate::basic::{
63 BoundaryOrder, ColumnOrder, Compression, CompressionCodec, Encoding, EncodingMask, PageType,
64 Type,
65};
66#[cfg(feature = "encryption")]
67use crate::encryption::decrypt::FileDecryptor;
68use crate::errors::{ParquetError, Result};
69#[cfg(feature = "encryption")]
70use crate::file::column_crypto_metadata::ColumnCryptoMetaData;
71pub(crate) use crate::file::metadata::memory::HeapSize;
72use crate::file::metadata::page_index::{PageIndex, PageIndexProvider, RowGroupPageIndex};
73#[cfg(feature = "encryption")]
74use crate::file::metadata::thrift::encryption::EncryptionAlgorithm;
75use crate::file::page_index::column_index::{ByteArrayColumnIndex, PrimitiveColumnIndex};
76use crate::file::page_index::{column_index::ColumnIndexMetaData, offset_index::PageLocation};
77use crate::file::statistics::Statistics;
78use crate::geospatial::statistics as geo_statistics;
79use crate::parquet_thrift::{
80 ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
81 WriteThrift, WriteThriftField,
82};
83use crate::schema::types::{
84 ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, SchemaDescriptor,
85 Type as SchemaType,
86};
87use crate::thrift_struct;
88use crate::{
89 data_type::private::ParquetValueType, file::page_index::offset_index::OffsetIndexMetaData,
90};
91
92pub use footer_tail::FooterTail;
93pub use options::{ParquetMetaDataOptions, ParquetStatisticsPolicy};
94pub use push_decoder::ParquetMetaDataPushDecoder;
95pub use reader::{PageIndexPolicy, ParquetMetaDataReader};
96use std::io::Write;
97use std::ops::Range;
98use std::sync::Arc;
99pub use writer::ParquetMetaDataWriter;
100pub(crate) use writer::ThriftMetadataWriter;
101
102#[derive(Debug, Clone)]
120pub struct ParquetMetaData {
121 file_metadata: FileMetaData,
123 row_groups: Vec<RowGroupMetaData>,
125 page_index: Option<Arc<dyn PageIndexProvider>>,
127 #[cfg(feature = "encryption")]
129 file_decryptor: Option<Box<FileDecryptor>>,
130}
131
132impl ParquetMetaData {
133 pub fn new(file_metadata: FileMetaData, row_groups: Vec<RowGroupMetaData>) -> Self {
136 ParquetMetaData {
137 file_metadata,
138 row_groups,
139 page_index: None,
140 #[cfg(feature = "encryption")]
141 file_decryptor: None,
142 }
143 }
144
145 #[cfg(feature = "encryption")]
148 pub(crate) fn with_file_decryptor(&mut self, file_decryptor: Option<FileDecryptor>) {
149 self.file_decryptor = file_decryptor.map(Box::new);
150 }
151
152 pub fn into_builder(self) -> ParquetMetaDataBuilder {
154 self.into()
155 }
156
157 pub fn file_metadata(&self) -> &FileMetaData {
159 &self.file_metadata
160 }
161
162 #[cfg(feature = "encryption")]
164 pub(crate) fn file_decryptor(&self) -> Option<&FileDecryptor> {
165 self.file_decryptor.as_deref()
166 }
167
168 pub fn num_row_groups(&self) -> usize {
170 self.row_groups.len()
171 }
172
173 pub fn row_group(&self, i: usize) -> &RowGroupMetaData {
176 &self.row_groups[i]
177 }
178
179 pub fn row_groups(&self) -> &[RowGroupMetaData] {
181 &self.row_groups
182 }
183
184 pub fn row_group_num_rows(&self, row_group_idx: usize) -> Result<usize> {
189 self.row_groups
190 .get(row_group_idx)
191 .ok_or_else(|| {
192 ParquetError::General(format!(
193 "Row group index {row_group_idx} out of bounds for file with {} row groups",
194 self.num_row_groups()
195 ))
196 })?
197 .num_rows()
198 .try_into()
199 .map_err(|e| ParquetError::General(format!("Row count overflow: {e}")))
200 }
201
202 pub fn page_index(&self) -> Option<&Arc<dyn PageIndexProvider>> {
209 self.page_index.as_ref()
210 }
211
212 pub fn page_index_for_row_group(&self, row_group_idx: usize) -> RowGroupPageIndex {
214 RowGroupPageIndex::new(row_group_idx, self.page_index.clone())
215 }
216
217 pub fn memory_size(&self) -> usize {
232 #[cfg(feature = "encryption")]
233 let encryption_size = self.file_decryptor.heap_size();
234 #[cfg(not(feature = "encryption"))]
235 let encryption_size = 0usize;
236
237 let page_index_size = if let Some(page_index) = self.page_index.as_ref() {
242 if let Some(page_index) = page_index.as_any().downcast_ref::<PageIndex>() {
243 2 * std::mem::size_of::<usize>() + std::mem::size_of::<PageIndex>() + page_index.heap_size()
247 } else {
248 0
249 }
250 } else {
251 0
252 };
253
254 std::mem::size_of::<Self>()
255 + self.file_metadata.heap_size()
256 + self.row_groups.heap_size()
257 + page_index_size
258 + encryption_size
259 }
260
261 pub(crate) fn set_page_index(&mut self, index: Option<Arc<dyn PageIndexProvider>>) {
263 self.page_index = index;
264 }
265}
266
267impl PartialEq for ParquetMetaData {
268 fn eq(&self, other: &Self) -> bool {
269 if self.file_metadata != other.file_metadata || self.row_groups != other.row_groups {
271 return false;
272 }
273
274 #[cfg(feature = "encryption")]
275 if self.file_decryptor != other.file_decryptor {
276 return false;
277 }
278
279 match (&self.page_index, &other.page_index) {
281 (None, None) => true,
282 (Some(a), Some(b)) => {
283 let a_page_index = a.as_any().downcast_ref::<PageIndex>();
285 let b_page_index = b.as_any().downcast_ref::<PageIndex>();
286 match (a_page_index, b_page_index) {
287 (Some(a_idx), Some(b_idx)) => a_idx == b_idx,
288 _ => Arc::ptr_eq(a, b), }
290 }
291 _ => false,
292 }
293 }
294}
295
296pub struct ParquetMetaDataBuilder(ParquetMetaData);
334
335impl ParquetMetaDataBuilder {
336 pub fn new(file_meta_data: FileMetaData) -> Self {
338 Self(ParquetMetaData::new(file_meta_data, vec![]))
339 }
340
341 pub fn new_from_metadata(metadata: ParquetMetaData) -> Self {
343 Self(metadata)
344 }
345
346 pub fn add_row_group(mut self, row_group: RowGroupMetaData) -> Self {
348 self.0.row_groups.push(row_group);
349 self
350 }
351
352 pub fn set_row_groups(mut self, row_groups: Vec<RowGroupMetaData>) -> Self {
354 self.0.row_groups = row_groups;
355 self
356 }
357
358 pub fn take_row_groups(&mut self) -> Vec<RowGroupMetaData> {
364 std::mem::take(&mut self.0.row_groups)
365 }
366
367 pub fn row_groups(&self) -> &[RowGroupMetaData] {
369 &self.0.row_groups
370 }
371
372 pub fn set_page_index(mut self, page_index: Option<Arc<dyn PageIndexProvider>>) -> Self {
381 self.0.page_index = page_index;
382 self
383 }
384
385 pub fn take_page_index(&mut self) -> Option<Arc<dyn PageIndexProvider>> {
387 std::mem::take(&mut self.0.page_index)
388 }
389
390 pub fn page_index(&self) -> Option<&dyn PageIndexProvider> {
392 self.0.page_index.as_ref().map(|arc| arc.as_ref())
393 }
394
395 #[cfg(feature = "encryption")]
397 pub(crate) fn set_file_decryptor(mut self, file_decryptor: Option<FileDecryptor>) -> Self {
398 self.0.with_file_decryptor(file_decryptor);
399 self
400 }
401
402 pub fn build(self) -> ParquetMetaData {
404 let Self(metadata) = self;
405 metadata
406 }
407}
408
409impl From<ParquetMetaData> for ParquetMetaDataBuilder {
410 fn from(meta_data: ParquetMetaData) -> Self {
411 Self(meta_data)
412 }
413}
414
415thrift_struct!(
416pub struct KeyValue {
418 1: required string key
419 2: optional string value
420}
421);
422
423impl KeyValue {
424 pub fn new<F2>(key: String, value: F2) -> KeyValue
426 where
427 F2: Into<Option<String>>,
428 {
429 KeyValue {
430 key,
431 value: value.into(),
432 }
433 }
434}
435
436thrift_struct!(
437pub struct PageEncodingStats {
439 1: required PageType page_type;
440 2: required Encoding encoding;
441 3: required i32 count;
442}
443);
444
445#[derive(Debug, Clone, PartialEq)]
448enum ParquetPageEncodingStats {
449 Full(Vec<PageEncodingStats>),
451 Mask(EncodingMask),
453}
454
455pub type FileMetaDataPtr = Arc<FileMetaData>;
457
458#[derive(Debug, Clone, PartialEq)]
462pub struct FileMetaData {
463 version: i32,
464 num_rows: i64,
465 created_by: Option<String>,
466 key_value_metadata: Option<Vec<KeyValue>>,
467 schema_descr: SchemaDescPtr,
468 column_orders: Option<Vec<ColumnOrder>>,
469 #[cfg(feature = "encryption")]
470 encryption_algorithm: Option<Box<EncryptionAlgorithm>>,
471 #[cfg(feature = "encryption")]
472 footer_signing_key_metadata: Option<Vec<u8>>,
473}
474
475impl FileMetaData {
476 pub fn new(
478 version: i32,
479 num_rows: i64,
480 created_by: Option<String>,
481 key_value_metadata: Option<Vec<KeyValue>>,
482 schema_descr: SchemaDescPtr,
483 column_orders: Option<Vec<ColumnOrder>>,
484 ) -> Self {
485 FileMetaData {
486 version,
487 num_rows,
488 created_by,
489 key_value_metadata,
490 schema_descr,
491 column_orders,
492 #[cfg(feature = "encryption")]
493 encryption_algorithm: None,
494 #[cfg(feature = "encryption")]
495 footer_signing_key_metadata: None,
496 }
497 }
498
499 #[cfg(feature = "encryption")]
500 pub(crate) fn with_encryption_algorithm(
501 mut self,
502 encryption_algorithm: Option<EncryptionAlgorithm>,
503 ) -> Self {
504 self.encryption_algorithm = encryption_algorithm.map(Box::new);
505 self
506 }
507
508 #[cfg(feature = "encryption")]
509 pub(crate) fn with_footer_signing_key_metadata(
510 mut self,
511 footer_signing_key_metadata: Option<Vec<u8>>,
512 ) -> Self {
513 self.footer_signing_key_metadata = footer_signing_key_metadata;
514 self
515 }
516
517 pub fn version(&self) -> i32 {
519 self.version
520 }
521
522 pub fn num_rows(&self) -> i64 {
524 self.num_rows
525 }
526
527 pub fn created_by(&self) -> Option<&str> {
536 self.created_by.as_deref()
537 }
538
539 pub fn key_value_metadata(&self) -> Option<&Vec<KeyValue>> {
541 self.key_value_metadata.as_ref()
542 }
543
544 pub fn schema(&self) -> &SchemaType {
548 self.schema_descr.root_schema()
549 }
550
551 pub fn schema_descr(&self) -> &SchemaDescriptor {
553 &self.schema_descr
554 }
555
556 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
558 self.schema_descr.clone()
559 }
560
561 pub fn column_orders(&self) -> Option<&Vec<ColumnOrder>> {
569 self.column_orders.as_ref()
570 }
571
572 pub fn column_order(&self, i: usize) -> ColumnOrder {
575 self.column_orders
576 .as_ref()
577 .map(|data| data[i])
578 .unwrap_or(ColumnOrder::UNDEFINED)
579 }
580}
581
582thrift_struct!(
583pub struct SortingColumn {
585 1: required i32 column_idx
587
588 2: required bool descending
590
591 3: required bool nulls_first
594}
595);
596
597pub type RowGroupMetaDataPtr = Arc<RowGroupMetaData>;
599
600#[derive(Debug, Clone, PartialEq)]
605pub struct RowGroupMetaData {
606 columns: Vec<ColumnChunkMetaData>,
607 num_rows: i64,
608 sorting_columns: Option<Vec<SortingColumn>>,
609 total_byte_size: i64,
610 schema_descr: SchemaDescPtr,
611 file_offset: Option<i64>,
613 ordinal: Option<i32>,
615}
616
617impl RowGroupMetaData {
618 pub fn builder(schema_descr: SchemaDescPtr) -> RowGroupMetaDataBuilder {
620 RowGroupMetaDataBuilder::new(schema_descr)
621 }
622
623 pub fn num_columns(&self) -> usize {
625 self.columns.len()
626 }
627
628 pub fn column(&self, i: usize) -> &ColumnChunkMetaData {
630 &self.columns[i]
631 }
632
633 pub fn columns(&self) -> &[ColumnChunkMetaData] {
635 &self.columns
636 }
637
638 pub fn columns_mut(&mut self) -> &mut [ColumnChunkMetaData] {
640 &mut self.columns
641 }
642
643 pub fn num_rows(&self) -> i64 {
645 self.num_rows
646 }
647
648 pub fn sorting_columns(&self) -> Option<&Vec<SortingColumn>> {
650 self.sorting_columns.as_ref()
651 }
652
653 pub fn total_byte_size(&self) -> i64 {
655 self.total_byte_size
656 }
657
658 pub fn compressed_size(&self) -> i64 {
660 self.columns.iter().map(|c| c.total_compressed_size).sum()
661 }
662
663 pub fn schema_descr(&self) -> &SchemaDescriptor {
665 self.schema_descr.as_ref()
666 }
667
668 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
670 self.schema_descr.clone()
671 }
672
673 #[inline(always)]
678 pub fn ordinal(&self) -> Option<i32> {
679 self.ordinal
680 }
681
682 #[inline(always)]
684 pub fn file_offset(&self) -> Option<i64> {
685 self.file_offset
686 }
687
688 pub fn into_builder(self) -> RowGroupMetaDataBuilder {
690 RowGroupMetaDataBuilder(self)
691 }
692}
693
694pub struct RowGroupMetaDataBuilder(RowGroupMetaData);
696
697impl RowGroupMetaDataBuilder {
698 fn new(schema_descr: SchemaDescPtr) -> Self {
700 Self(RowGroupMetaData {
701 columns: Vec::with_capacity(schema_descr.num_columns()),
702 schema_descr,
703 file_offset: None,
704 num_rows: 0,
705 sorting_columns: None,
706 total_byte_size: 0,
707 ordinal: None,
708 })
709 }
710
711 pub fn set_num_rows(mut self, value: i64) -> Self {
713 self.0.num_rows = value;
714 self
715 }
716
717 pub fn set_sorting_columns(mut self, value: Option<Vec<SortingColumn>>) -> Self {
719 self.0.sorting_columns = value;
720 self
721 }
722
723 pub fn set_total_byte_size(mut self, value: i64) -> Self {
725 self.0.total_byte_size = value;
726 self
727 }
728
729 pub fn take_columns(&mut self) -> Vec<ColumnChunkMetaData> {
735 std::mem::take(&mut self.0.columns)
736 }
737
738 pub fn set_column_metadata(mut self, value: Vec<ColumnChunkMetaData>) -> Self {
740 self.0.columns = value;
741 self
742 }
743
744 pub fn add_column_metadata(mut self, value: ColumnChunkMetaData) -> Self {
746 self.0.columns.push(value);
747 self
748 }
749
750 pub fn set_ordinal(mut self, value: i32) -> Self {
752 self.0.ordinal = Some(value);
753 self
754 }
755
756 pub fn set_file_offset(mut self, value: i64) -> Self {
758 self.0.file_offset = Some(value);
759 self
760 }
761
762 pub fn build(self) -> Result<RowGroupMetaData> {
764 if self.0.schema_descr.num_columns() != self.0.columns.len() {
765 return Err(general_err!(
766 "Column length mismatch: {} != {}",
767 self.0.schema_descr.num_columns(),
768 self.0.columns.len()
769 ));
770 }
771
772 Ok(self.0)
773 }
774
775 pub(super) fn build_unchecked(self) -> RowGroupMetaData {
777 self.0
778 }
779}
780
781#[derive(Debug, Clone, PartialEq)]
783pub struct ColumnChunkMetaData {
784 column_descr: ColumnDescPtr,
785 encodings: EncodingMask,
786 file_path: Option<String>,
787 file_offset: i64,
788 num_values: i64,
789 compression: CompressionCodec,
790 total_compressed_size: i64,
791 total_uncompressed_size: i64,
792 data_page_offset: i64,
793 index_page_offset: Option<i64>,
794 dictionary_page_offset: Option<i64>,
795 statistics: Option<Statistics>,
796 geo_statistics: Option<Box<geo_statistics::GeospatialStatistics>>,
797 encoding_stats: Option<ParquetPageEncodingStats>,
798 bloom_filter_offset: Option<i64>,
799 bloom_filter_length: Option<i32>,
800 offset_index_offset: Option<i64>,
801 offset_index_length: Option<i32>,
802 column_index_offset: Option<i64>,
803 column_index_length: Option<i32>,
804 unencoded_byte_array_data_bytes: Option<i64>,
805 repetition_level_histogram: Option<LevelHistogram>,
806 definition_level_histogram: Option<LevelHistogram>,
807 #[cfg(feature = "encryption")]
808 column_crypto_metadata: Option<Box<ColumnCryptoMetaData>>,
809 #[cfg(feature = "encryption")]
810 encrypted_column_metadata: Option<Vec<u8>>,
811 #[cfg(feature = "encryption")]
815 plaintext_footer_mode: bool,
816}
817
818#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
827pub struct LevelHistogram {
828 inner: Vec<i64>,
829}
830
831impl LevelHistogram {
832 pub fn try_new(max_level: i16) -> Option<Self> {
838 if max_level > 0 {
839 Some(Self {
840 inner: vec![0; max_level as usize + 1],
841 })
842 } else {
843 None
844 }
845 }
846 pub fn values(&self) -> &[i64] {
848 &self.inner
849 }
850
851 pub fn into_inner(self) -> Vec<i64> {
853 self.inner
854 }
855
856 pub fn get(&self, index: usize) -> Option<i64> {
863 self.inner.get(index).copied()
864 }
865
866 pub fn add(&mut self, other: &Self) {
871 assert_eq!(self.len(), other.len());
872 for (dst, src) in self.inner.iter_mut().zip(other.inner.iter()) {
873 *dst += src;
874 }
875 }
876
877 pub fn len(&self) -> usize {
879 self.inner.len()
880 }
881
882 pub fn is_empty(&self) -> bool {
884 self.inner.is_empty()
885 }
886
887 pub fn reset(&mut self) {
889 self.inner.fill(0);
890 }
891
892 #[inline]
894 pub fn increment_by(&mut self, level: i16, count: i64) {
895 self.inner[level as usize] += count;
896 }
897
898 #[deprecated(since = "58.2.0", note = "Use `increment_by` instead")]
904 pub fn update_from_levels(&mut self, levels: &[i16]) {
905 for &level in levels {
906 self.increment_by(level, 1);
907 }
908 }
909}
910
911impl From<Vec<i64>> for LevelHistogram {
912 fn from(inner: Vec<i64>) -> Self {
913 Self { inner }
914 }
915}
916
917impl From<LevelHistogram> for Vec<i64> {
918 fn from(value: LevelHistogram) -> Self {
919 value.into_inner()
920 }
921}
922
923impl HeapSize for LevelHistogram {
924 fn heap_size(&self) -> usize {
925 self.inner.heap_size()
926 }
927}
928
929impl ColumnChunkMetaData {
931 pub fn builder(column_descr: ColumnDescPtr) -> ColumnChunkMetaDataBuilder {
933 ColumnChunkMetaDataBuilder::new(column_descr)
934 }
935
936 pub fn file_path(&self) -> Option<&str> {
941 self.file_path.as_deref()
942 }
943
944 pub fn file_offset(&self) -> i64 {
951 self.file_offset
952 }
953
954 pub fn column_type(&self) -> Type {
956 self.column_descr.physical_type()
957 }
958
959 pub fn column_path(&self) -> &ColumnPath {
961 self.column_descr.path()
962 }
963
964 pub fn column_descr(&self) -> &ColumnDescriptor {
966 self.column_descr.as_ref()
967 }
968
969 pub fn column_descr_ptr(&self) -> ColumnDescPtr {
971 self.column_descr.clone()
972 }
973
974 pub fn encodings(&self) -> impl Iterator<Item = Encoding> {
976 self.encodings.encodings()
977 }
978
979 pub fn encodings_mask(&self) -> &EncodingMask {
981 &self.encodings
982 }
983
984 pub fn num_values(&self) -> i64 {
986 self.num_values
987 }
988
989 pub fn compression(&self) -> Compression {
996 self.compression.into()
997 }
998
999 pub fn compression_codec(&self) -> CompressionCodec {
1001 self.compression
1002 }
1003
1004 pub fn compressed_size(&self) -> i64 {
1006 self.total_compressed_size
1007 }
1008
1009 pub fn uncompressed_size(&self) -> i64 {
1011 self.total_uncompressed_size
1012 }
1013
1014 pub fn data_page_offset(&self) -> i64 {
1016 self.data_page_offset
1017 }
1018
1019 pub fn index_page_offset(&self) -> Option<i64> {
1021 self.index_page_offset
1022 }
1023
1024 pub fn dictionary_page_offset(&self) -> Option<i64> {
1026 self.dictionary_page_offset
1027 }
1028
1029 pub fn byte_range(&self) -> (u64, u64) {
1035 let col_start = match self.dictionary_page_offset() {
1036 Some(dictionary_page_offset) => dictionary_page_offset,
1037 None => self.data_page_offset(),
1038 };
1039 let col_len = self.compressed_size();
1040 assert!(
1041 col_start >= 0 && col_len >= 0,
1042 "column start and length should not be negative"
1043 );
1044 (col_start as u64, col_len as u64)
1045 }
1046
1047 pub fn statistics(&self) -> Option<&Statistics> {
1050 self.statistics.as_ref()
1051 }
1052
1053 pub fn geo_statistics(&self) -> Option<&geo_statistics::GeospatialStatistics> {
1056 self.geo_statistics.as_deref()
1057 }
1058
1059 pub fn page_encoding_stats(&self) -> Option<&Vec<PageEncodingStats>> {
1066 match self.encoding_stats.as_ref() {
1067 Some(ParquetPageEncodingStats::Full(stats)) => Some(stats),
1068 _ => None,
1069 }
1070 }
1071
1072 pub fn page_encoding_stats_mask(&self) -> Option<&EncodingMask> {
1102 match self.encoding_stats.as_ref() {
1103 Some(ParquetPageEncodingStats::Mask(stats)) => Some(stats),
1104 _ => None,
1105 }
1106 }
1107
1108 pub fn bloom_filter_offset(&self) -> Option<i64> {
1110 self.bloom_filter_offset
1111 }
1112
1113 pub fn bloom_filter_length(&self) -> Option<i32> {
1115 self.bloom_filter_length
1116 }
1117
1118 pub fn column_index_offset(&self) -> Option<i64> {
1120 self.column_index_offset
1121 }
1122
1123 pub fn column_index_length(&self) -> Option<i32> {
1125 self.column_index_length
1126 }
1127
1128 pub fn column_index_range(&self) -> Option<Range<u64>> {
1134 let offset = u64::try_from(self.column_index_offset?).ok()?;
1135 let length = u64::try_from(self.column_index_length?).ok()?;
1136 Some(offset..(offset + length))
1137 }
1138
1139 pub fn offset_index_offset(&self) -> Option<i64> {
1141 self.offset_index_offset
1142 }
1143
1144 pub fn offset_index_length(&self) -> Option<i32> {
1146 self.offset_index_length
1147 }
1148
1149 pub fn offset_index_range(&self) -> Option<Range<u64>> {
1155 let offset = u64::try_from(self.offset_index_offset?).ok()?;
1156 let length = u64::try_from(self.offset_index_length?).ok()?;
1157 Some(offset..(offset + length))
1158 }
1159
1160 pub fn unencoded_byte_array_data_bytes(&self) -> Option<i64> {
1165 self.unencoded_byte_array_data_bytes
1166 }
1167
1168 pub fn repetition_level_histogram(&self) -> Option<&LevelHistogram> {
1174 self.repetition_level_histogram.as_ref()
1175 }
1176
1177 pub fn definition_level_histogram(&self) -> Option<&LevelHistogram> {
1183 self.definition_level_histogram.as_ref()
1184 }
1185
1186 #[cfg(feature = "encryption")]
1188 pub fn crypto_metadata(&self) -> Option<&ColumnCryptoMetaData> {
1189 self.column_crypto_metadata.as_deref()
1190 }
1191
1192 pub fn into_builder(self) -> ColumnChunkMetaDataBuilder {
1194 ColumnChunkMetaDataBuilder::from(self)
1195 }
1196}
1197
1198pub struct ColumnChunkMetaDataBuilder(ColumnChunkMetaData);
1217
1218impl ColumnChunkMetaDataBuilder {
1219 fn new(column_descr: ColumnDescPtr) -> Self {
1223 Self(ColumnChunkMetaData {
1224 column_descr,
1225 encodings: Default::default(),
1226 file_path: None,
1227 file_offset: 0,
1228 num_values: 0,
1229 compression: CompressionCodec::UNCOMPRESSED,
1230 total_compressed_size: 0,
1231 total_uncompressed_size: 0,
1232 data_page_offset: 0,
1233 index_page_offset: None,
1234 dictionary_page_offset: None,
1235 statistics: None,
1236 geo_statistics: None,
1237 encoding_stats: None,
1238 bloom_filter_offset: None,
1239 bloom_filter_length: None,
1240 offset_index_offset: None,
1241 offset_index_length: None,
1242 column_index_offset: None,
1243 column_index_length: None,
1244 unencoded_byte_array_data_bytes: None,
1245 repetition_level_histogram: None,
1246 definition_level_histogram: None,
1247 #[cfg(feature = "encryption")]
1248 column_crypto_metadata: None,
1249 #[cfg(feature = "encryption")]
1250 encrypted_column_metadata: None,
1251 #[cfg(feature = "encryption")]
1252 plaintext_footer_mode: false,
1253 })
1254 }
1255
1256 pub fn set_encodings(mut self, encodings: Vec<Encoding>) -> Self {
1258 self.0.encodings = EncodingMask::new_from_encodings(encodings.iter());
1259 self
1260 }
1261
1262 pub fn set_encodings_mask(mut self, encodings: EncodingMask) -> Self {
1264 self.0.encodings = encodings;
1265 self
1266 }
1267
1268 pub fn set_file_path(mut self, value: String) -> Self {
1270 self.0.file_path = Some(value);
1271 self
1272 }
1273
1274 pub fn set_num_values(mut self, value: i64) -> Self {
1276 self.0.num_values = value;
1277 self
1278 }
1279
1280 pub fn set_compression(mut self, value: Compression) -> Self {
1282 self.0.compression = value.into();
1283 self
1284 }
1285
1286 pub fn set_compression_codec(mut self, value: CompressionCodec) -> Self {
1288 self.0.compression = value;
1289 self
1290 }
1291
1292 pub fn set_total_compressed_size(mut self, value: i64) -> Self {
1294 self.0.total_compressed_size = value;
1295 self
1296 }
1297
1298 pub fn set_total_uncompressed_size(mut self, value: i64) -> Self {
1300 self.0.total_uncompressed_size = value;
1301 self
1302 }
1303
1304 pub fn set_data_page_offset(mut self, value: i64) -> Self {
1306 self.0.data_page_offset = value;
1307 self
1308 }
1309
1310 pub fn set_dictionary_page_offset(mut self, value: Option<i64>) -> Self {
1312 self.0.dictionary_page_offset = value;
1313 self
1314 }
1315
1316 pub fn set_index_page_offset(mut self, value: Option<i64>) -> Self {
1318 self.0.index_page_offset = value;
1319 self
1320 }
1321
1322 pub fn set_statistics(mut self, value: Statistics) -> Self {
1324 self.0.statistics = Some(value);
1325 self
1326 }
1327
1328 pub fn set_geo_statistics(mut self, value: Box<geo_statistics::GeospatialStatistics>) -> Self {
1330 self.0.geo_statistics = Some(value);
1331 self
1332 }
1333
1334 pub fn clear_statistics(mut self) -> Self {
1336 self.0.statistics = None;
1337 self
1338 }
1339
1340 pub fn set_page_encoding_stats(mut self, value: Vec<PageEncodingStats>) -> Self {
1344 self.0.encoding_stats = Some(ParquetPageEncodingStats::Full(value));
1345 self
1346 }
1347
1348 pub fn set_page_encoding_stats_mask(mut self, value: EncodingMask) -> Self {
1352 self.0.encoding_stats = Some(ParquetPageEncodingStats::Mask(value));
1353 self
1354 }
1355
1356 pub fn clear_page_encoding_stats(mut self) -> Self {
1358 self.0.encoding_stats = None;
1359 self
1360 }
1361
1362 pub fn set_bloom_filter_offset(mut self, value: Option<i64>) -> Self {
1364 self.0.bloom_filter_offset = value;
1365 self
1366 }
1367
1368 pub fn set_bloom_filter_length(mut self, value: Option<i32>) -> Self {
1370 self.0.bloom_filter_length = value;
1371 self
1372 }
1373
1374 pub fn set_offset_index_offset(mut self, value: Option<i64>) -> Self {
1376 self.0.offset_index_offset = value;
1377 self
1378 }
1379
1380 pub fn set_offset_index_length(mut self, value: Option<i32>) -> Self {
1382 self.0.offset_index_length = value;
1383 self
1384 }
1385
1386 pub fn set_column_index_offset(mut self, value: Option<i64>) -> Self {
1388 self.0.column_index_offset = value;
1389 self
1390 }
1391
1392 pub fn set_column_index_length(mut self, value: Option<i32>) -> Self {
1394 self.0.column_index_length = value;
1395 self
1396 }
1397
1398 pub fn set_unencoded_byte_array_data_bytes(mut self, value: Option<i64>) -> Self {
1400 self.0.unencoded_byte_array_data_bytes = value;
1401 self
1402 }
1403
1404 pub fn set_repetition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1406 self.0.repetition_level_histogram = value;
1407 self
1408 }
1409
1410 pub fn set_definition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1412 self.0.definition_level_histogram = value;
1413 self
1414 }
1415
1416 #[cfg(feature = "encryption")]
1417 pub fn set_column_crypto_metadata(mut self, value: Option<ColumnCryptoMetaData>) -> Self {
1419 self.0.column_crypto_metadata = value.map(Box::new);
1420 self
1421 }
1422
1423 #[cfg(feature = "encryption")]
1424 pub fn set_encrypted_column_metadata(mut self, value: Option<Vec<u8>>) -> Self {
1426 self.0.encrypted_column_metadata = value;
1427 self
1428 }
1429
1430 pub fn build(self) -> Result<ColumnChunkMetaData> {
1432 Ok(self.0)
1433 }
1434}
1435
1436pub struct ColumnIndexBuilder {
1441 column_type: Type,
1442 null_pages: Vec<bool>,
1443 min_values: Vec<Vec<u8>>,
1444 max_values: Vec<Vec<u8>>,
1445 null_counts: Vec<i64>,
1446 nan_counts: Vec<Option<i64>>,
1447 boundary_order: BoundaryOrder,
1448 repetition_level_histograms: Option<Vec<i64>>,
1450 definition_level_histograms: Option<Vec<i64>>,
1452 valid: bool,
1460}
1461
1462impl ColumnIndexBuilder {
1463 pub fn new(column_type: Type) -> Self {
1465 ColumnIndexBuilder {
1466 column_type,
1467 null_pages: Vec::new(),
1468 min_values: Vec::new(),
1469 max_values: Vec::new(),
1470 null_counts: Vec::new(),
1471 nan_counts: Vec::new(),
1472 boundary_order: BoundaryOrder::UNORDERED,
1473 repetition_level_histograms: None,
1474 definition_level_histograms: None,
1475 valid: true,
1476 }
1477 }
1478
1479 pub fn append(
1486 &mut self,
1487 null_page: bool,
1488 min_value: Vec<u8>,
1489 max_value: Vec<u8>,
1490 null_count: i64,
1491 nan_count: Option<i64>,
1492 ) {
1493 self.null_pages.push(null_page);
1494 self.min_values.push(min_value);
1495 self.max_values.push(max_value);
1496 self.null_counts.push(null_count);
1497 self.nan_counts.push(nan_count);
1498 }
1499
1500 pub fn append_histograms(
1505 &mut self,
1506 repetition_level_histogram: &Option<LevelHistogram>,
1507 definition_level_histogram: &Option<LevelHistogram>,
1508 ) {
1509 if !self.valid {
1510 return;
1511 }
1512 if let Some(rep_lvl_hist) = repetition_level_histogram {
1513 let hist = self.repetition_level_histograms.get_or_insert(Vec::new());
1514 hist.reserve(rep_lvl_hist.len());
1515 hist.extend(rep_lvl_hist.values());
1516 }
1517 if let Some(def_lvl_hist) = definition_level_histogram {
1518 let hist = self.definition_level_histograms.get_or_insert(Vec::new());
1519 hist.reserve(def_lvl_hist.len());
1520 hist.extend(def_lvl_hist.values());
1521 }
1522 }
1523
1524 pub fn set_boundary_order(&mut self, boundary_order: BoundaryOrder) {
1526 self.boundary_order = boundary_order;
1527 }
1528
1529 pub fn to_invalid(&mut self) {
1531 self.valid = false;
1532 }
1533
1534 pub fn valid(&self) -> bool {
1536 self.valid
1537 }
1538
1539 pub fn build(self) -> Result<ColumnIndexMetaData> {
1543 Ok(match self.column_type {
1544 Type::BOOLEAN => {
1545 let index = self.build_page_index(false)?;
1546 ColumnIndexMetaData::BOOLEAN(index)
1547 }
1548 Type::INT32 => {
1549 let index = self.build_page_index(false)?;
1550 ColumnIndexMetaData::INT32(index)
1551 }
1552 Type::INT64 => {
1553 let index = self.build_page_index(false)?;
1554 ColumnIndexMetaData::INT64(index)
1555 }
1556 Type::INT96 => {
1557 let index = self.build_page_index(false)?;
1558 ColumnIndexMetaData::INT96(index)
1559 }
1560 Type::FLOAT => {
1561 let index = self.build_page_index(true)?;
1562 ColumnIndexMetaData::FLOAT(index)
1563 }
1564 Type::DOUBLE => {
1565 let index = self.build_page_index(true)?;
1566 ColumnIndexMetaData::DOUBLE(index)
1567 }
1568 Type::BYTE_ARRAY => {
1569 let index = self.build_byte_array_index(false)?;
1570 ColumnIndexMetaData::BYTE_ARRAY(index)
1571 }
1572 Type::FIXED_LEN_BYTE_ARRAY => {
1573 let index = self.build_byte_array_index(true)?;
1574 ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index)
1575 }
1576 })
1577 }
1578
1579 fn build_nan_counts(nan_counts: &[Option<i64>]) -> Option<Vec<i64>> {
1580 let has_some = nan_counts.iter().any(|x| x.is_some());
1581 let has_none = nan_counts.iter().any(|x| x.is_none());
1582
1583 if has_some && !has_none {
1584 Some(nan_counts.iter().map(|x| x.unwrap()).collect())
1585 } else if !has_some && has_none {
1586 None
1587 } else {
1588 debug_assert!(
1589 false,
1590 "Mixed Some/None in nan_counts - caller should provide consistent values"
1591 );
1592 Some(nan_counts.iter().map(|x| x.unwrap_or(0)).collect())
1593 }
1594 }
1595
1596 fn build_page_index<T>(self, may_have_nan: bool) -> Result<PrimitiveColumnIndex<T>>
1597 where
1598 T: ParquetValueType,
1599 {
1600 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1601 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1602
1603 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1608 Self::build_nan_counts(&self.nan_counts)
1609 } else {
1610 None
1611 };
1612
1613 PrimitiveColumnIndex::try_new(
1614 self.null_pages,
1615 self.boundary_order,
1616 Some(self.null_counts),
1617 nan_counts,
1618 self.repetition_level_histograms,
1619 self.definition_level_histograms,
1620 min_values,
1621 max_values,
1622 )
1623 }
1624
1625 fn build_byte_array_index(self, may_have_nan: bool) -> Result<ByteArrayColumnIndex> {
1626 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1627 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1628
1629 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1634 Self::build_nan_counts(&self.nan_counts)
1635 } else {
1636 None
1637 };
1638
1639 ByteArrayColumnIndex::try_new(
1640 self.null_pages,
1641 self.boundary_order,
1642 Some(self.null_counts),
1643 nan_counts,
1644 self.repetition_level_histograms,
1645 self.definition_level_histograms,
1646 min_values,
1647 max_values,
1648 )
1649 }
1650}
1651
1652impl From<ColumnChunkMetaData> for ColumnChunkMetaDataBuilder {
1653 fn from(value: ColumnChunkMetaData) -> Self {
1654 ColumnChunkMetaDataBuilder(value)
1655 }
1656}
1657
1658pub struct OffsetIndexBuilder {
1662 offset_array: Vec<i64>,
1663 compressed_page_size_array: Vec<i32>,
1664 first_row_index_array: Vec<i64>,
1665 unencoded_byte_array_data_bytes_array: Option<Vec<i64>>,
1666 current_first_row_index: i64,
1667}
1668
1669impl Default for OffsetIndexBuilder {
1670 fn default() -> Self {
1671 Self::new()
1672 }
1673}
1674
1675impl OffsetIndexBuilder {
1676 pub fn new() -> Self {
1678 OffsetIndexBuilder {
1679 offset_array: Vec::new(),
1680 compressed_page_size_array: Vec::new(),
1681 first_row_index_array: Vec::new(),
1682 unencoded_byte_array_data_bytes_array: None,
1683 current_first_row_index: 0,
1684 }
1685 }
1686
1687 pub fn append_row_count(&mut self, row_count: i64) {
1689 let current_page_row_index = self.current_first_row_index;
1690 self.first_row_index_array.push(current_page_row_index);
1691 self.current_first_row_index += row_count;
1692 }
1693
1694 pub fn append_offset_and_size(&mut self, offset: i64, compressed_page_size: i32) {
1696 self.offset_array.push(offset);
1697 self.compressed_page_size_array.push(compressed_page_size);
1698 }
1699
1700 pub fn append_unencoded_byte_array_data_bytes(
1702 &mut self,
1703 unencoded_byte_array_data_bytes: Option<i64>,
1704 ) {
1705 if let Some(val) = unencoded_byte_array_data_bytes {
1706 self.unencoded_byte_array_data_bytes_array
1707 .get_or_insert(Vec::new())
1708 .push(val);
1709 }
1710 }
1711
1712 pub fn build(self) -> OffsetIndexMetaData {
1714 let locations = self
1715 .offset_array
1716 .iter()
1717 .zip(self.compressed_page_size_array.iter())
1718 .zip(self.first_row_index_array.iter())
1719 .map(|((offset, size), row_index)| PageLocation {
1720 offset: *offset,
1721 compressed_page_size: *size,
1722 first_row_index: *row_index,
1723 })
1724 .collect::<Vec<_>>();
1725 OffsetIndexMetaData {
1726 page_locations: locations,
1727 unencoded_byte_array_data_bytes: self.unencoded_byte_array_data_bytes_array,
1728 }
1729 }
1730}
1731
1732#[cfg(test)]
1733mod tests {
1734 use super::*;
1735 use crate::basic::{PageType, SortOrder};
1736 use crate::file::metadata::page_index::PageIndexBuilder;
1737 use crate::file::metadata::thrift::tests::{
1738 read_column_chunk, read_column_chunk_with_options, read_row_group,
1739 };
1740
1741 #[test]
1742 #[expect(deprecated)]
1743 fn test_level_histogram_update_from_levels_compat() {
1744 let mut histogram = LevelHistogram::try_new(2).unwrap();
1745 histogram.update_from_levels(&[0, 2, 1, 2, 2]);
1746 assert_eq!(histogram.values(), &[1, 1, 3]);
1747 }
1748
1749 #[test]
1750 fn test_row_group_metadata_thrift_conversion() {
1751 let schema_descr = get_test_schema_descr();
1752
1753 let mut columns = vec![];
1754 for ptr in schema_descr.columns() {
1755 let column = ColumnChunkMetaData::builder(ptr.clone()).build().unwrap();
1756 columns.push(column);
1757 }
1758 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
1759 .set_num_rows(1000)
1760 .set_total_byte_size(2000)
1761 .set_column_metadata(columns)
1762 .set_ordinal(1)
1763 .build()
1764 .unwrap();
1765
1766 let mut buf = Vec::new();
1767 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1768 row_group_meta.write_thrift(&mut writer).unwrap();
1769
1770 let row_group_res = read_row_group(&buf, schema_descr).unwrap();
1771
1772 assert_eq!(row_group_res, row_group_meta);
1773 }
1774
1775 #[test]
1776 fn test_row_group_metadata_thrift_conversion_empty() {
1777 let schema_descr = get_test_schema_descr();
1778
1779 let row_group_meta = RowGroupMetaData::builder(schema_descr).build();
1780
1781 assert!(row_group_meta.is_err());
1782 if let Err(e) = row_group_meta {
1783 assert_eq!(
1784 format!("{e}"),
1785 "Parquet error: Column length mismatch: 2 != 0"
1786 );
1787 }
1788 }
1789
1790 #[test]
1792 fn test_row_group_metadata_thrift_corrupted() {
1793 let schema_descr_2cols = Arc::new(SchemaDescriptor::new(Arc::new(
1794 SchemaType::group_type_builder("schema")
1795 .with_fields(vec![
1796 Arc::new(
1797 SchemaType::primitive_type_builder("a", Type::INT32)
1798 .build()
1799 .unwrap(),
1800 ),
1801 Arc::new(
1802 SchemaType::primitive_type_builder("b", Type::INT32)
1803 .build()
1804 .unwrap(),
1805 ),
1806 ])
1807 .build()
1808 .unwrap(),
1809 )));
1810
1811 let schema_descr_3cols = Arc::new(SchemaDescriptor::new(Arc::new(
1812 SchemaType::group_type_builder("schema")
1813 .with_fields(vec![
1814 Arc::new(
1815 SchemaType::primitive_type_builder("a", Type::INT32)
1816 .build()
1817 .unwrap(),
1818 ),
1819 Arc::new(
1820 SchemaType::primitive_type_builder("b", Type::INT32)
1821 .build()
1822 .unwrap(),
1823 ),
1824 Arc::new(
1825 SchemaType::primitive_type_builder("c", Type::INT32)
1826 .build()
1827 .unwrap(),
1828 ),
1829 ])
1830 .build()
1831 .unwrap(),
1832 )));
1833
1834 let row_group_meta_2cols = RowGroupMetaData::builder(schema_descr_2cols.clone())
1835 .set_num_rows(1000)
1836 .set_total_byte_size(2000)
1837 .set_column_metadata(vec![
1838 ColumnChunkMetaData::builder(schema_descr_2cols.column(0))
1839 .build()
1840 .unwrap(),
1841 ColumnChunkMetaData::builder(schema_descr_2cols.column(1))
1842 .build()
1843 .unwrap(),
1844 ])
1845 .set_ordinal(1)
1846 .build()
1847 .unwrap();
1848 let mut buf = Vec::new();
1849 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1850 row_group_meta_2cols.write_thrift(&mut writer).unwrap();
1851
1852 let err = read_row_group(&buf, schema_descr_3cols)
1853 .unwrap_err()
1854 .to_string();
1855 assert_eq!(
1856 err,
1857 "Parquet error: Column count mismatch. Schema has 3 columns while Row Group has 2"
1858 );
1859 }
1860
1861 #[test]
1862 fn test_column_chunk_metadata_thrift_conversion() {
1863 let column_descr = get_test_schema_descr().column(0);
1864 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1865 .set_encodings_mask(EncodingMask::new_from_encodings(
1866 [Encoding::PLAIN, Encoding::RLE].iter(),
1867 ))
1868 .set_file_path("file_path".to_owned())
1869 .set_num_values(1000)
1870 .set_compression_codec(CompressionCodec::SNAPPY)
1871 .set_total_compressed_size(2000)
1872 .set_total_uncompressed_size(3000)
1873 .set_data_page_offset(4000)
1874 .set_dictionary_page_offset(Some(5000))
1875 .set_page_encoding_stats(vec![
1876 PageEncodingStats {
1877 page_type: PageType::DATA_PAGE,
1878 encoding: Encoding::PLAIN,
1879 count: 3,
1880 },
1881 PageEncodingStats {
1882 page_type: PageType::DATA_PAGE,
1883 encoding: Encoding::RLE,
1884 count: 5,
1885 },
1886 ])
1887 .set_bloom_filter_offset(Some(6000))
1888 .set_bloom_filter_length(Some(25))
1889 .set_offset_index_offset(Some(7000))
1890 .set_offset_index_length(Some(25))
1891 .set_column_index_offset(Some(8000))
1892 .set_column_index_length(Some(25))
1893 .set_unencoded_byte_array_data_bytes(Some(2000))
1894 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
1895 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
1896 .build()
1897 .unwrap();
1898
1899 let mut buf = Vec::new();
1900 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1901 col_metadata.write_thrift(&mut writer).unwrap();
1902 let col_chunk_res = read_column_chunk(&buf, column_descr.clone()).unwrap();
1903
1904 let expected_metadata = ColumnChunkMetaData::builder(column_descr)
1905 .set_encodings_mask(EncodingMask::new_from_encodings(
1906 [Encoding::PLAIN, Encoding::RLE].iter(),
1907 ))
1908 .set_file_path("file_path".to_owned())
1909 .set_num_values(1000)
1910 .set_compression_codec(CompressionCodec::SNAPPY)
1911 .set_total_compressed_size(2000)
1912 .set_total_uncompressed_size(3000)
1913 .set_data_page_offset(4000)
1914 .set_dictionary_page_offset(Some(5000))
1915 .set_page_encoding_stats_mask(EncodingMask::new_from_encodings(
1916 [Encoding::PLAIN, Encoding::RLE].iter(),
1917 ))
1918 .set_bloom_filter_offset(Some(6000))
1919 .set_bloom_filter_length(Some(25))
1920 .set_offset_index_offset(Some(7000))
1921 .set_offset_index_length(Some(25))
1922 .set_column_index_offset(Some(8000))
1923 .set_column_index_length(Some(25))
1924 .set_unencoded_byte_array_data_bytes(Some(2000))
1925 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
1926 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
1927 .build()
1928 .unwrap();
1929
1930 assert_eq!(col_chunk_res, expected_metadata);
1931 }
1932
1933 #[test]
1934 fn test_column_chunk_metadata_thrift_conversion_full_stats() {
1935 let column_descr = get_test_schema_descr().column(0);
1936 let stats = vec![
1937 PageEncodingStats {
1938 page_type: PageType::DATA_PAGE,
1939 encoding: Encoding::PLAIN,
1940 count: 3,
1941 },
1942 PageEncodingStats {
1943 page_type: PageType::DATA_PAGE,
1944 encoding: Encoding::RLE,
1945 count: 5,
1946 },
1947 ];
1948 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1949 .set_encodings_mask(EncodingMask::new_from_encodings(
1950 [Encoding::PLAIN, Encoding::RLE].iter(),
1951 ))
1952 .set_num_values(1000)
1953 .set_compression_codec(CompressionCodec::SNAPPY)
1954 .set_total_compressed_size(2000)
1955 .set_total_uncompressed_size(3000)
1956 .set_data_page_offset(4000)
1957 .set_page_encoding_stats(stats)
1958 .build()
1959 .unwrap();
1960
1961 let mut buf = Vec::new();
1962 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1963 col_metadata.write_thrift(&mut writer).unwrap();
1964
1965 let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
1966 let col_chunk_res =
1967 read_column_chunk_with_options(&buf, column_descr, Some(&options)).unwrap();
1968
1969 assert_eq!(col_chunk_res, col_metadata);
1970 }
1971
1972 #[test]
1973 fn test_column_chunk_metadata_thrift_conversion_empty() {
1974 let column_descr = get_test_schema_descr().column(0);
1975
1976 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1977 .build()
1978 .unwrap();
1979
1980 let mut buf = Vec::new();
1981 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1982 col_metadata.write_thrift(&mut writer).unwrap();
1983 let col_chunk_res = read_column_chunk(&buf, column_descr).unwrap();
1984
1985 assert_eq!(col_chunk_res, col_metadata);
1986 }
1987
1988 #[test]
1989 fn test_compressed_size() {
1990 let schema_descr = get_test_schema_descr();
1991
1992 let mut columns = vec![];
1993 for column_descr in schema_descr.columns() {
1994 let column = ColumnChunkMetaData::builder(column_descr.clone())
1995 .set_total_compressed_size(500)
1996 .set_total_uncompressed_size(700)
1997 .build()
1998 .unwrap();
1999 columns.push(column);
2000 }
2001 let row_group_meta = RowGroupMetaData::builder(schema_descr)
2002 .set_num_rows(1000)
2003 .set_column_metadata(columns)
2004 .build()
2005 .unwrap();
2006
2007 let compressed_size_res = row_group_meta.compressed_size();
2008 let compressed_size_exp: i64 = 1000;
2009
2010 assert_eq!(compressed_size_res, compressed_size_exp);
2011 }
2012
2013 #[test]
2014 fn test_memory_size() {
2015 let schema_descr = get_test_schema_descr();
2016
2017 let columns = schema_descr
2018 .columns()
2019 .iter()
2020 .map(|column_descr| {
2021 ColumnChunkMetaData::builder(column_descr.clone())
2022 .set_statistics(Statistics::new::<i32>(None, None, None, None, false))
2023 .build()
2024 })
2025 .collect::<Result<Vec<_>>>()
2026 .unwrap();
2027 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2028 .set_num_rows(1000)
2029 .set_column_metadata(columns)
2030 .build()
2031 .unwrap();
2032 let row_group_meta = vec![row_group_meta];
2033
2034 let version = 2;
2035 let num_rows = 1000;
2036 let created_by = Some(String::from("test harness"));
2037 let key_value_metadata = Some(vec![KeyValue::new(
2038 String::from("Foo"),
2039 Some(String::from("bar")),
2040 )]);
2041 let column_orders = Some(vec![
2042 ColumnOrder::UNDEFINED,
2043 ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED),
2044 ]);
2045 let file_metadata = FileMetaData::new(
2046 version,
2047 num_rows,
2048 created_by,
2049 key_value_metadata,
2050 schema_descr.clone(),
2051 column_orders,
2052 );
2053
2054 let columns_with_stats = schema_descr
2056 .columns()
2057 .iter()
2058 .map(|column_descr| {
2059 ColumnChunkMetaData::builder(column_descr.clone())
2060 .set_statistics(Statistics::new::<i32>(
2061 Some(0),
2062 Some(100),
2063 None,
2064 None,
2065 false,
2066 ))
2067 .build()
2068 })
2069 .collect::<Result<Vec<_>>>()
2070 .unwrap();
2071
2072 let row_group_meta_with_stats = RowGroupMetaData::builder(schema_descr)
2073 .set_num_rows(1000)
2074 .set_column_metadata(columns_with_stats)
2075 .build()
2076 .unwrap();
2077 let row_group_meta_with_stats = vec![row_group_meta_with_stats];
2078
2079 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata.clone())
2080 .set_row_groups(row_group_meta_with_stats)
2081 .build();
2082
2083 #[cfg(not(feature = "encryption"))]
2085 let base_expected_size = 2766;
2086 #[cfg(feature = "encryption")]
2087 let base_expected_size = 2934;
2088
2089 assert_eq!(parquet_meta.memory_size(), base_expected_size);
2090
2091 let mut page_index = PageIndexBuilder::new(1, 1);
2092
2093 let mut column_index = ColumnIndexBuilder::new(Type::BOOLEAN);
2094 column_index.append(false, vec![1u8], vec![2u8, 3u8], 4, None);
2095 let column_index = column_index.build().unwrap();
2096 {
2097 let ColumnIndexMetaData::BOOLEAN(_) = column_index else {
2098 panic!("wrong type of column index")
2099 };
2100 }
2101
2102 page_index.put_column_index(column_index, 0, 0);
2103
2104 let mut offset_index = OffsetIndexBuilder::new();
2106 offset_index.append_row_count(1);
2107 offset_index.append_offset_and_size(2, 3);
2108 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2109 offset_index.append_row_count(1);
2110 offset_index.append_offset_and_size(2, 3);
2111 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2112 let offset_index = offset_index.build();
2113 page_index.put_offset_index(offset_index, 0, 0);
2114
2115 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata)
2116 .set_row_groups(row_group_meta)
2117 .set_page_index(Some(Arc::new(page_index.build())))
2118 .build();
2119
2120 #[cfg(not(feature = "encryption"))]
2122 let bigger_expected_size = 3280;
2123 #[cfg(feature = "encryption")]
2124 let bigger_expected_size = 3448;
2125
2126 assert!(bigger_expected_size > base_expected_size);
2128 assert_eq!(parquet_meta.memory_size(), bigger_expected_size);
2129 }
2130
2131 #[test]
2132 #[cfg(feature = "encryption")]
2133 fn test_memory_size_with_decryptor() {
2134 use crate::encryption::decrypt::FileDecryptionProperties;
2135 use crate::file::metadata::thrift::encryption::AesGcmV1;
2136
2137 let schema_descr = get_test_schema_descr();
2138
2139 let columns = schema_descr
2140 .columns()
2141 .iter()
2142 .map(|column_descr| ColumnChunkMetaData::builder(column_descr.clone()).build())
2143 .collect::<Result<Vec<_>>>()
2144 .unwrap();
2145 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2146 .set_num_rows(1000)
2147 .set_column_metadata(columns)
2148 .build()
2149 .unwrap();
2150 let row_group_meta = vec![row_group_meta];
2151
2152 let version = 2;
2153 let num_rows = 1000;
2154 let aad_file_unique = vec![1u8; 8];
2155 let aad_prefix = vec![2u8; 8];
2156 let encryption_algorithm = EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 {
2157 aad_prefix: Some(aad_prefix.clone()),
2158 aad_file_unique: Some(aad_file_unique.clone()),
2159 supply_aad_prefix: Some(true),
2160 });
2161 let footer_key_metadata = Some(vec![3u8; 8]);
2162 let file_metadata =
2163 FileMetaData::new(version, num_rows, None, None, schema_descr.clone(), None)
2164 .with_encryption_algorithm(Some(encryption_algorithm))
2165 .with_footer_signing_key_metadata(footer_key_metadata.clone());
2166
2167 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2168 .set_row_groups(row_group_meta.clone())
2169 .build();
2170
2171 let base_expected_size = 2042;
2172 assert_eq!(parquet_meta_data.memory_size(), base_expected_size);
2173
2174 let footer_key = b"0123456789012345";
2175 let column_key = b"1234567890123450";
2176 let mut decryption_properties_builder =
2177 FileDecryptionProperties::builder(footer_key.to_vec())
2178 .with_aad_prefix(aad_prefix.clone());
2179 for column in schema_descr.columns() {
2180 decryption_properties_builder = decryption_properties_builder
2181 .with_column_key(&column.path().string(), column_key.to_vec());
2182 }
2183 let decryption_properties = decryption_properties_builder.build().unwrap();
2184 let decryptor = FileDecryptor::new(
2185 &decryption_properties,
2186 footer_key_metadata.as_deref(),
2187 aad_file_unique,
2188 aad_prefix,
2189 )
2190 .unwrap();
2191
2192 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2193 .set_row_groups(row_group_meta.clone())
2194 .set_file_decryptor(Some(decryptor))
2195 .build();
2196
2197 let expected_size_with_decryptor = 3056;
2198 assert!(expected_size_with_decryptor > base_expected_size);
2199
2200 assert_eq!(
2201 parquet_meta_data.memory_size(),
2202 expected_size_with_decryptor
2203 );
2204 }
2205
2206 fn get_test_schema_descr() -> SchemaDescPtr {
2208 let schema = SchemaType::group_type_builder("schema")
2209 .with_fields(vec![
2210 Arc::new(
2211 SchemaType::primitive_type_builder("a", Type::INT32)
2212 .build()
2213 .unwrap(),
2214 ),
2215 Arc::new(
2216 SchemaType::primitive_type_builder("b", Type::INT32)
2217 .build()
2218 .unwrap(),
2219 ),
2220 ])
2221 .build()
2222 .unwrap();
2223
2224 Arc::new(SchemaDescriptor::new(Arc::new(schema)))
2225 }
2226
2227 #[test]
2228 fn test_page_index_builder_partial_population() {
2229 let mut builder = PageIndexBuilder::new(1, 1);
2230
2231 let mut offset_idx_builder = OffsetIndexBuilder::new();
2233 offset_idx_builder.append_row_count(50);
2234 offset_idx_builder.append_offset_and_size(1000, 500);
2235 let offset_idx = offset_idx_builder.build();
2236 builder.put_offset_index(offset_idx, 0, 0);
2237
2238 let page_index = builder.build();
2239
2240 assert!(!page_index.has_column_indexes());
2242 assert!(page_index.has_offset_indexes());
2244 }
2245
2246 #[test]
2247 fn test_page_index_builder_empty_to_none() {
2248 let builder = PageIndexBuilder::new(2, 2);
2250
2251 let page_index = builder.build();
2252
2253 assert!(!page_index.has_column_indexes());
2255 assert!(!page_index.has_offset_indexes());
2256 }
2257
2258 #[test]
2259 fn test_rebuild_page_index() {
2260 let mut builder = PageIndexBuilder::new(1, 2);
2262
2263 let mut col_idx_builder = ColumnIndexBuilder::new(Type::INT32);
2265 col_idx_builder.append(false, vec![1, 0, 0, 0], vec![100, 0, 0, 0], 10, None);
2266 let col_idx = col_idx_builder.build().unwrap();
2267 builder.put_column_index(col_idx.clone(), 0, 0);
2268
2269 let mut offset_idx_builder = OffsetIndexBuilder::new();
2271 offset_idx_builder.append_row_count(50);
2272 offset_idx_builder.append_offset_and_size(1000, 500);
2273 let offset_idx = offset_idx_builder.build();
2274 builder.put_offset_index(offset_idx.clone(), 0, 0);
2275
2276 let page_index = builder.build();
2277
2278 assert!(page_index.is_complete());
2280
2281 assert!(page_index.column_index(0, 0).is_some());
2283 assert!(page_index.offset_index(0, 0).is_some());
2284 assert!(page_index.column_index(0, 1).is_none());
2286 assert!(page_index.offset_index(0, 1).is_none());
2287
2288 let mut builder = PageIndexBuilder::new_from(page_index);
2289
2290 builder.put_column_index(col_idx, 0, 1);
2292 builder.put_offset_index(offset_idx, 0, 1);
2293 let page_index = builder.build();
2294
2295 assert!(page_index.column_index(0, 0).is_some());
2297 assert!(page_index.offset_index(0, 0).is_some());
2298 assert!(page_index.column_index(0, 1).is_some());
2299 assert!(page_index.offset_index(0, 1).is_some());
2300 }
2301}