1mod footer_tail;
90mod memory;
91mod options;
92mod parser;
93mod push_decoder;
94pub(crate) mod reader;
95pub(crate) mod thrift;
96mod writer;
97
98use crate::basic::{EncodingMask, PageType};
99#[cfg(feature = "encryption")]
100use crate::encryption::decrypt::FileDecryptor;
101#[cfg(feature = "encryption")]
102use crate::file::column_crypto_metadata::ColumnCryptoMetaData;
103pub(crate) use crate::file::metadata::memory::HeapSize;
104#[cfg(feature = "encryption")]
105use crate::file::metadata::thrift::encryption::EncryptionAlgorithm;
106use crate::file::page_index::column_index::{ByteArrayColumnIndex, PrimitiveColumnIndex};
107use crate::file::page_index::{column_index::ColumnIndexMetaData, offset_index::PageLocation};
108use crate::file::statistics::Statistics;
109use crate::geospatial::statistics as geo_statistics;
110use crate::schema::types::{
111 ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, SchemaDescriptor,
112 Type as SchemaType,
113};
114use crate::thrift_struct;
115use crate::{
116 basic::BoundaryOrder,
117 errors::{ParquetError, Result},
118};
119use crate::{
120 basic::{ColumnOrder, Compression, Encoding, Type},
121 parquet_thrift::{
122 ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol,
123 ThriftCompactOutputProtocol, WriteThrift, WriteThriftField,
124 },
125};
126use crate::{
127 data_type::private::ParquetValueType, file::page_index::offset_index::OffsetIndexMetaData,
128};
129
130pub use footer_tail::FooterTail;
131pub use options::{ParquetMetaDataOptions, ParquetStatisticsPolicy};
132pub use push_decoder::ParquetMetaDataPushDecoder;
133pub use reader::{PageIndexPolicy, ParquetMetaDataReader};
134use std::io::Write;
135use std::ops::Range;
136use std::sync::Arc;
137pub use writer::ParquetMetaDataWriter;
138pub(crate) use writer::ThriftMetadataWriter;
139
140pub type ParquetColumnIndex = Vec<Vec<ColumnIndexMetaData>>;
157
158pub type ParquetOffsetIndex = Vec<Vec<OffsetIndexMetaData>>;
170
171#[derive(Debug, Clone, PartialEq)]
189pub struct ParquetMetaData {
190 file_metadata: FileMetaData,
192 row_groups: Vec<RowGroupMetaData>,
194 column_index: Option<ParquetColumnIndex>,
196 offset_index: Option<ParquetOffsetIndex>,
198 #[cfg(feature = "encryption")]
200 file_decryptor: Option<Box<FileDecryptor>>,
201}
202
203impl ParquetMetaData {
204 pub fn new(file_metadata: FileMetaData, row_groups: Vec<RowGroupMetaData>) -> Self {
207 ParquetMetaData {
208 file_metadata,
209 row_groups,
210 column_index: None,
211 offset_index: None,
212 #[cfg(feature = "encryption")]
213 file_decryptor: None,
214 }
215 }
216
217 #[cfg(feature = "encryption")]
220 pub(crate) fn with_file_decryptor(&mut self, file_decryptor: Option<FileDecryptor>) {
221 self.file_decryptor = file_decryptor.map(Box::new);
222 }
223
224 pub fn into_builder(self) -> ParquetMetaDataBuilder {
226 self.into()
227 }
228
229 pub fn file_metadata(&self) -> &FileMetaData {
231 &self.file_metadata
232 }
233
234 #[cfg(feature = "encryption")]
236 pub(crate) fn file_decryptor(&self) -> Option<&FileDecryptor> {
237 self.file_decryptor.as_deref()
238 }
239
240 pub fn num_row_groups(&self) -> usize {
242 self.row_groups.len()
243 }
244
245 pub fn row_group(&self, i: usize) -> &RowGroupMetaData {
248 &self.row_groups[i]
249 }
250
251 pub fn row_groups(&self) -> &[RowGroupMetaData] {
253 &self.row_groups
254 }
255
256 pub fn column_index(&self) -> Option<&ParquetColumnIndex> {
263 self.column_index.as_ref()
264 }
265
266 pub fn offset_index(&self) -> Option<&ParquetOffsetIndex> {
273 self.offset_index.as_ref()
274 }
275
276 pub fn memory_size(&self) -> usize {
291 #[cfg(feature = "encryption")]
292 let encryption_size = self.file_decryptor.heap_size();
293 #[cfg(not(feature = "encryption"))]
294 let encryption_size = 0usize;
295
296 std::mem::size_of::<Self>()
297 + self.file_metadata.heap_size()
298 + self.row_groups.heap_size()
299 + self.column_index.heap_size()
300 + self.offset_index.heap_size()
301 + encryption_size
302 }
303
304 pub(crate) fn set_column_index(&mut self, index: Option<ParquetColumnIndex>) {
306 self.column_index = index;
307 }
308
309 pub(crate) fn set_offset_index(&mut self, index: Option<ParquetOffsetIndex>) {
311 self.offset_index = index;
312 }
313}
314
315pub struct ParquetMetaDataBuilder(ParquetMetaData);
353
354impl ParquetMetaDataBuilder {
355 pub fn new(file_meta_data: FileMetaData) -> Self {
357 Self(ParquetMetaData::new(file_meta_data, vec![]))
358 }
359
360 pub fn new_from_metadata(metadata: ParquetMetaData) -> Self {
362 Self(metadata)
363 }
364
365 pub fn add_row_group(mut self, row_group: RowGroupMetaData) -> Self {
367 self.0.row_groups.push(row_group);
368 self
369 }
370
371 pub fn set_row_groups(mut self, row_groups: Vec<RowGroupMetaData>) -> Self {
373 self.0.row_groups = row_groups;
374 self
375 }
376
377 pub fn take_row_groups(&mut self) -> Vec<RowGroupMetaData> {
383 std::mem::take(&mut self.0.row_groups)
384 }
385
386 pub fn row_groups(&self) -> &[RowGroupMetaData] {
388 &self.0.row_groups
389 }
390
391 pub fn set_column_index(mut self, column_index: Option<ParquetColumnIndex>) -> Self {
393 self.0.column_index = column_index;
394 self
395 }
396
397 pub fn take_column_index(&mut self) -> Option<ParquetColumnIndex> {
399 std::mem::take(&mut self.0.column_index)
400 }
401
402 pub fn column_index(&self) -> Option<&ParquetColumnIndex> {
404 self.0.column_index.as_ref()
405 }
406
407 pub fn set_offset_index(mut self, offset_index: Option<ParquetOffsetIndex>) -> Self {
409 self.0.offset_index = offset_index;
410 self
411 }
412
413 pub fn take_offset_index(&mut self) -> Option<ParquetOffsetIndex> {
415 std::mem::take(&mut self.0.offset_index)
416 }
417
418 pub fn offset_index(&self) -> Option<&ParquetOffsetIndex> {
420 self.0.offset_index.as_ref()
421 }
422
423 #[cfg(feature = "encryption")]
425 pub(crate) fn set_file_decryptor(mut self, file_decryptor: Option<FileDecryptor>) -> Self {
426 self.0.with_file_decryptor(file_decryptor);
427 self
428 }
429
430 pub fn build(self) -> ParquetMetaData {
432 let Self(metadata) = self;
433 metadata
434 }
435}
436
437impl From<ParquetMetaData> for ParquetMetaDataBuilder {
438 fn from(meta_data: ParquetMetaData) -> Self {
439 Self(meta_data)
440 }
441}
442
443thrift_struct!(
444pub struct KeyValue {
446 1: required string key
447 2: optional string value
448}
449);
450
451impl KeyValue {
452 pub fn new<F2>(key: String, value: F2) -> KeyValue
454 where
455 F2: Into<Option<String>>,
456 {
457 KeyValue {
458 key,
459 value: value.into(),
460 }
461 }
462}
463
464thrift_struct!(
465pub struct PageEncodingStats {
467 1: required PageType page_type;
468 2: required Encoding encoding;
469 3: required i32 count;
470}
471);
472
473#[derive(Debug, Clone, PartialEq)]
476enum ParquetPageEncodingStats {
477 Full(Vec<PageEncodingStats>),
479 Mask(EncodingMask),
481}
482
483pub type FileMetaDataPtr = Arc<FileMetaData>;
485
486#[derive(Debug, Clone, PartialEq)]
490pub struct FileMetaData {
491 version: i32,
492 num_rows: i64,
493 created_by: Option<String>,
494 key_value_metadata: Option<Vec<KeyValue>>,
495 schema_descr: SchemaDescPtr,
496 column_orders: Option<Vec<ColumnOrder>>,
497 #[cfg(feature = "encryption")]
498 encryption_algorithm: Option<Box<EncryptionAlgorithm>>,
499 #[cfg(feature = "encryption")]
500 footer_signing_key_metadata: Option<Vec<u8>>,
501}
502
503impl FileMetaData {
504 pub fn new(
506 version: i32,
507 num_rows: i64,
508 created_by: Option<String>,
509 key_value_metadata: Option<Vec<KeyValue>>,
510 schema_descr: SchemaDescPtr,
511 column_orders: Option<Vec<ColumnOrder>>,
512 ) -> Self {
513 FileMetaData {
514 version,
515 num_rows,
516 created_by,
517 key_value_metadata,
518 schema_descr,
519 column_orders,
520 #[cfg(feature = "encryption")]
521 encryption_algorithm: None,
522 #[cfg(feature = "encryption")]
523 footer_signing_key_metadata: None,
524 }
525 }
526
527 #[cfg(feature = "encryption")]
528 pub(crate) fn with_encryption_algorithm(
529 mut self,
530 encryption_algorithm: Option<EncryptionAlgorithm>,
531 ) -> Self {
532 self.encryption_algorithm = encryption_algorithm.map(Box::new);
533 self
534 }
535
536 #[cfg(feature = "encryption")]
537 pub(crate) fn with_footer_signing_key_metadata(
538 mut self,
539 footer_signing_key_metadata: Option<Vec<u8>>,
540 ) -> Self {
541 self.footer_signing_key_metadata = footer_signing_key_metadata;
542 self
543 }
544
545 pub fn version(&self) -> i32 {
547 self.version
548 }
549
550 pub fn num_rows(&self) -> i64 {
552 self.num_rows
553 }
554
555 pub fn created_by(&self) -> Option<&str> {
564 self.created_by.as_deref()
565 }
566
567 pub fn key_value_metadata(&self) -> Option<&Vec<KeyValue>> {
569 self.key_value_metadata.as_ref()
570 }
571
572 pub fn schema(&self) -> &SchemaType {
576 self.schema_descr.root_schema()
577 }
578
579 pub fn schema_descr(&self) -> &SchemaDescriptor {
581 &self.schema_descr
582 }
583
584 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
586 self.schema_descr.clone()
587 }
588
589 pub fn column_orders(&self) -> Option<&Vec<ColumnOrder>> {
597 self.column_orders.as_ref()
598 }
599
600 pub fn column_order(&self, i: usize) -> ColumnOrder {
603 self.column_orders
604 .as_ref()
605 .map(|data| data[i])
606 .unwrap_or(ColumnOrder::UNDEFINED)
607 }
608}
609
610thrift_struct!(
611pub struct SortingColumn {
613 1: required i32 column_idx
615
616 2: required bool descending
618
619 3: required bool nulls_first
622}
623);
624
625pub type RowGroupMetaDataPtr = Arc<RowGroupMetaData>;
627
628#[derive(Debug, Clone, PartialEq)]
633pub struct RowGroupMetaData {
634 columns: Vec<ColumnChunkMetaData>,
635 num_rows: i64,
636 sorting_columns: Option<Vec<SortingColumn>>,
637 total_byte_size: i64,
638 schema_descr: SchemaDescPtr,
639 file_offset: Option<i64>,
641 ordinal: Option<i16>,
643}
644
645impl RowGroupMetaData {
646 pub fn builder(schema_descr: SchemaDescPtr) -> RowGroupMetaDataBuilder {
648 RowGroupMetaDataBuilder::new(schema_descr)
649 }
650
651 pub fn num_columns(&self) -> usize {
653 self.columns.len()
654 }
655
656 pub fn column(&self, i: usize) -> &ColumnChunkMetaData {
658 &self.columns[i]
659 }
660
661 pub fn columns(&self) -> &[ColumnChunkMetaData] {
663 &self.columns
664 }
665
666 pub fn columns_mut(&mut self) -> &mut [ColumnChunkMetaData] {
668 &mut self.columns
669 }
670
671 pub fn num_rows(&self) -> i64 {
673 self.num_rows
674 }
675
676 pub fn sorting_columns(&self) -> Option<&Vec<SortingColumn>> {
678 self.sorting_columns.as_ref()
679 }
680
681 pub fn total_byte_size(&self) -> i64 {
683 self.total_byte_size
684 }
685
686 pub fn compressed_size(&self) -> i64 {
688 self.columns.iter().map(|c| c.total_compressed_size).sum()
689 }
690
691 pub fn schema_descr(&self) -> &SchemaDescriptor {
693 self.schema_descr.as_ref()
694 }
695
696 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
698 self.schema_descr.clone()
699 }
700
701 #[inline(always)]
706 pub fn ordinal(&self) -> Option<i16> {
707 self.ordinal
708 }
709
710 #[inline(always)]
712 pub fn file_offset(&self) -> Option<i64> {
713 self.file_offset
714 }
715
716 pub fn into_builder(self) -> RowGroupMetaDataBuilder {
718 RowGroupMetaDataBuilder(self)
719 }
720}
721
722pub struct RowGroupMetaDataBuilder(RowGroupMetaData);
724
725impl RowGroupMetaDataBuilder {
726 fn new(schema_descr: SchemaDescPtr) -> Self {
728 Self(RowGroupMetaData {
729 columns: Vec::with_capacity(schema_descr.num_columns()),
730 schema_descr,
731 file_offset: None,
732 num_rows: 0,
733 sorting_columns: None,
734 total_byte_size: 0,
735 ordinal: None,
736 })
737 }
738
739 pub fn set_num_rows(mut self, value: i64) -> Self {
741 self.0.num_rows = value;
742 self
743 }
744
745 pub fn set_sorting_columns(mut self, value: Option<Vec<SortingColumn>>) -> Self {
747 self.0.sorting_columns = value;
748 self
749 }
750
751 pub fn set_total_byte_size(mut self, value: i64) -> Self {
753 self.0.total_byte_size = value;
754 self
755 }
756
757 pub fn take_columns(&mut self) -> Vec<ColumnChunkMetaData> {
763 std::mem::take(&mut self.0.columns)
764 }
765
766 pub fn set_column_metadata(mut self, value: Vec<ColumnChunkMetaData>) -> Self {
768 self.0.columns = value;
769 self
770 }
771
772 pub fn add_column_metadata(mut self, value: ColumnChunkMetaData) -> Self {
774 self.0.columns.push(value);
775 self
776 }
777
778 pub fn set_ordinal(mut self, value: i16) -> Self {
780 self.0.ordinal = Some(value);
781 self
782 }
783
784 pub fn set_file_offset(mut self, value: i64) -> Self {
786 self.0.file_offset = Some(value);
787 self
788 }
789
790 pub fn build(self) -> Result<RowGroupMetaData> {
792 if self.0.schema_descr.num_columns() != self.0.columns.len() {
793 return Err(general_err!(
794 "Column length mismatch: {} != {}",
795 self.0.schema_descr.num_columns(),
796 self.0.columns.len()
797 ));
798 }
799
800 Ok(self.0)
801 }
802
803 pub(super) fn build_unchecked(self) -> RowGroupMetaData {
805 self.0
806 }
807}
808
809#[derive(Debug, Clone, PartialEq)]
811pub struct ColumnChunkMetaData {
812 column_descr: ColumnDescPtr,
813 encodings: EncodingMask,
814 file_path: Option<String>,
815 file_offset: i64,
816 num_values: i64,
817 compression: Compression,
818 total_compressed_size: i64,
819 total_uncompressed_size: i64,
820 data_page_offset: i64,
821 index_page_offset: Option<i64>,
822 dictionary_page_offset: Option<i64>,
823 statistics: Option<Statistics>,
824 geo_statistics: Option<Box<geo_statistics::GeospatialStatistics>>,
825 encoding_stats: Option<ParquetPageEncodingStats>,
826 bloom_filter_offset: Option<i64>,
827 bloom_filter_length: Option<i32>,
828 offset_index_offset: Option<i64>,
829 offset_index_length: Option<i32>,
830 column_index_offset: Option<i64>,
831 column_index_length: Option<i32>,
832 unencoded_byte_array_data_bytes: Option<i64>,
833 repetition_level_histogram: Option<LevelHistogram>,
834 definition_level_histogram: Option<LevelHistogram>,
835 #[cfg(feature = "encryption")]
836 column_crypto_metadata: Option<Box<ColumnCryptoMetaData>>,
837 #[cfg(feature = "encryption")]
838 encrypted_column_metadata: Option<Vec<u8>>,
839 #[cfg(feature = "encryption")]
843 plaintext_footer_mode: bool,
844}
845
846#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
855pub struct LevelHistogram {
856 inner: Vec<i64>,
857}
858
859impl LevelHistogram {
860 pub fn try_new(max_level: i16) -> Option<Self> {
866 if max_level > 0 {
867 Some(Self {
868 inner: vec![0; max_level as usize + 1],
869 })
870 } else {
871 None
872 }
873 }
874 pub fn values(&self) -> &[i64] {
876 &self.inner
877 }
878
879 pub fn into_inner(self) -> Vec<i64> {
881 self.inner
882 }
883
884 pub fn get(&self, index: usize) -> Option<i64> {
891 self.inner.get(index).copied()
892 }
893
894 pub fn add(&mut self, other: &Self) {
899 assert_eq!(self.len(), other.len());
900 for (dst, src) in self.inner.iter_mut().zip(other.inner.iter()) {
901 *dst += src;
902 }
903 }
904
905 pub fn len(&self) -> usize {
907 self.inner.len()
908 }
909
910 pub fn is_empty(&self) -> bool {
912 self.inner.is_empty()
913 }
914
915 pub fn reset(&mut self) {
917 for value in self.inner.iter_mut() {
918 *value = 0;
919 }
920 }
921
922 #[inline]
924 pub fn increment_by(&mut self, level: i16, count: i64) {
925 self.inner[level as usize] += count;
926 }
927
928 #[deprecated(since = "58.2.0", note = "Use `increment_by` instead")]
934 pub fn update_from_levels(&mut self, levels: &[i16]) {
935 for &level in levels {
936 self.increment_by(level, 1);
937 }
938 }
939}
940
941impl From<Vec<i64>> for LevelHistogram {
942 fn from(inner: Vec<i64>) -> Self {
943 Self { inner }
944 }
945}
946
947impl From<LevelHistogram> for Vec<i64> {
948 fn from(value: LevelHistogram) -> Self {
949 value.into_inner()
950 }
951}
952
953impl HeapSize for LevelHistogram {
954 fn heap_size(&self) -> usize {
955 self.inner.heap_size()
956 }
957}
958
959impl ColumnChunkMetaData {
961 pub fn builder(column_descr: ColumnDescPtr) -> ColumnChunkMetaDataBuilder {
963 ColumnChunkMetaDataBuilder::new(column_descr)
964 }
965
966 pub fn file_path(&self) -> Option<&str> {
971 self.file_path.as_deref()
972 }
973
974 pub fn file_offset(&self) -> i64 {
981 self.file_offset
982 }
983
984 pub fn column_type(&self) -> Type {
986 self.column_descr.physical_type()
987 }
988
989 pub fn column_path(&self) -> &ColumnPath {
991 self.column_descr.path()
992 }
993
994 pub fn column_descr(&self) -> &ColumnDescriptor {
996 self.column_descr.as_ref()
997 }
998
999 pub fn column_descr_ptr(&self) -> ColumnDescPtr {
1001 self.column_descr.clone()
1002 }
1003
1004 pub fn encodings(&self) -> impl Iterator<Item = Encoding> {
1006 self.encodings.encodings()
1007 }
1008
1009 pub fn encodings_mask(&self) -> &EncodingMask {
1011 &self.encodings
1012 }
1013
1014 pub fn num_values(&self) -> i64 {
1016 self.num_values
1017 }
1018
1019 pub fn compression(&self) -> Compression {
1021 self.compression
1022 }
1023
1024 pub fn compressed_size(&self) -> i64 {
1026 self.total_compressed_size
1027 }
1028
1029 pub fn uncompressed_size(&self) -> i64 {
1031 self.total_uncompressed_size
1032 }
1033
1034 pub fn data_page_offset(&self) -> i64 {
1036 self.data_page_offset
1037 }
1038
1039 pub fn index_page_offset(&self) -> Option<i64> {
1041 self.index_page_offset
1042 }
1043
1044 pub fn dictionary_page_offset(&self) -> Option<i64> {
1046 self.dictionary_page_offset
1047 }
1048
1049 pub fn byte_range(&self) -> (u64, u64) {
1051 let col_start = match self.dictionary_page_offset() {
1052 Some(dictionary_page_offset) => dictionary_page_offset,
1053 None => self.data_page_offset(),
1054 };
1055 let col_len = self.compressed_size();
1056 assert!(
1057 col_start >= 0 && col_len >= 0,
1058 "column start and length should not be negative"
1059 );
1060 (col_start as u64, col_len as u64)
1061 }
1062
1063 pub fn statistics(&self) -> Option<&Statistics> {
1066 self.statistics.as_ref()
1067 }
1068
1069 pub fn geo_statistics(&self) -> Option<&geo_statistics::GeospatialStatistics> {
1072 self.geo_statistics.as_deref()
1073 }
1074
1075 pub fn page_encoding_stats(&self) -> Option<&Vec<PageEncodingStats>> {
1082 match self.encoding_stats.as_ref() {
1083 Some(ParquetPageEncodingStats::Full(stats)) => Some(stats),
1084 _ => None,
1085 }
1086 }
1087
1088 pub fn page_encoding_stats_mask(&self) -> Option<&EncodingMask> {
1118 match self.encoding_stats.as_ref() {
1119 Some(ParquetPageEncodingStats::Mask(stats)) => Some(stats),
1120 _ => None,
1121 }
1122 }
1123
1124 pub fn bloom_filter_offset(&self) -> Option<i64> {
1126 self.bloom_filter_offset
1127 }
1128
1129 pub fn bloom_filter_length(&self) -> Option<i32> {
1131 self.bloom_filter_length
1132 }
1133
1134 pub fn column_index_offset(&self) -> Option<i64> {
1136 self.column_index_offset
1137 }
1138
1139 pub fn column_index_length(&self) -> Option<i32> {
1141 self.column_index_length
1142 }
1143
1144 pub(crate) fn column_index_range(&self) -> Option<Range<u64>> {
1146 let offset = u64::try_from(self.column_index_offset?).ok()?;
1147 let length = u64::try_from(self.column_index_length?).ok()?;
1148 Some(offset..(offset + length))
1149 }
1150
1151 pub fn offset_index_offset(&self) -> Option<i64> {
1153 self.offset_index_offset
1154 }
1155
1156 pub fn offset_index_length(&self) -> Option<i32> {
1158 self.offset_index_length
1159 }
1160
1161 pub(crate) fn offset_index_range(&self) -> Option<Range<u64>> {
1163 let offset = u64::try_from(self.offset_index_offset?).ok()?;
1164 let length = u64::try_from(self.offset_index_length?).ok()?;
1165 Some(offset..(offset + length))
1166 }
1167
1168 pub fn unencoded_byte_array_data_bytes(&self) -> Option<i64> {
1173 self.unencoded_byte_array_data_bytes
1174 }
1175
1176 pub fn repetition_level_histogram(&self) -> Option<&LevelHistogram> {
1182 self.repetition_level_histogram.as_ref()
1183 }
1184
1185 pub fn definition_level_histogram(&self) -> Option<&LevelHistogram> {
1191 self.definition_level_histogram.as_ref()
1192 }
1193
1194 #[cfg(feature = "encryption")]
1196 pub fn crypto_metadata(&self) -> Option<&ColumnCryptoMetaData> {
1197 self.column_crypto_metadata.as_deref()
1198 }
1199
1200 pub fn into_builder(self) -> ColumnChunkMetaDataBuilder {
1202 ColumnChunkMetaDataBuilder::from(self)
1203 }
1204}
1205
1206pub struct ColumnChunkMetaDataBuilder(ColumnChunkMetaData);
1225
1226impl ColumnChunkMetaDataBuilder {
1227 fn new(column_descr: ColumnDescPtr) -> Self {
1231 Self(ColumnChunkMetaData {
1232 column_descr,
1233 encodings: Default::default(),
1234 file_path: None,
1235 file_offset: 0,
1236 num_values: 0,
1237 compression: Compression::UNCOMPRESSED,
1238 total_compressed_size: 0,
1239 total_uncompressed_size: 0,
1240 data_page_offset: 0,
1241 index_page_offset: None,
1242 dictionary_page_offset: None,
1243 statistics: None,
1244 geo_statistics: None,
1245 encoding_stats: None,
1246 bloom_filter_offset: None,
1247 bloom_filter_length: None,
1248 offset_index_offset: None,
1249 offset_index_length: None,
1250 column_index_offset: None,
1251 column_index_length: None,
1252 unencoded_byte_array_data_bytes: None,
1253 repetition_level_histogram: None,
1254 definition_level_histogram: None,
1255 #[cfg(feature = "encryption")]
1256 column_crypto_metadata: None,
1257 #[cfg(feature = "encryption")]
1258 encrypted_column_metadata: None,
1259 #[cfg(feature = "encryption")]
1260 plaintext_footer_mode: false,
1261 })
1262 }
1263
1264 pub fn set_encodings(mut self, encodings: Vec<Encoding>) -> Self {
1266 self.0.encodings = EncodingMask::new_from_encodings(encodings.iter());
1267 self
1268 }
1269
1270 pub fn set_encodings_mask(mut self, encodings: EncodingMask) -> Self {
1272 self.0.encodings = encodings;
1273 self
1274 }
1275
1276 pub fn set_file_path(mut self, value: String) -> Self {
1278 self.0.file_path = Some(value);
1279 self
1280 }
1281
1282 pub fn set_num_values(mut self, value: i64) -> Self {
1284 self.0.num_values = value;
1285 self
1286 }
1287
1288 pub fn set_compression(mut self, value: Compression) -> Self {
1290 self.0.compression = value;
1291 self
1292 }
1293
1294 pub fn set_total_compressed_size(mut self, value: i64) -> Self {
1296 self.0.total_compressed_size = value;
1297 self
1298 }
1299
1300 pub fn set_total_uncompressed_size(mut self, value: i64) -> Self {
1302 self.0.total_uncompressed_size = value;
1303 self
1304 }
1305
1306 pub fn set_data_page_offset(mut self, value: i64) -> Self {
1308 self.0.data_page_offset = value;
1309 self
1310 }
1311
1312 pub fn set_dictionary_page_offset(mut self, value: Option<i64>) -> Self {
1314 self.0.dictionary_page_offset = value;
1315 self
1316 }
1317
1318 pub fn set_index_page_offset(mut self, value: Option<i64>) -> Self {
1320 self.0.index_page_offset = value;
1321 self
1322 }
1323
1324 pub fn set_statistics(mut self, value: Statistics) -> Self {
1326 self.0.statistics = Some(value);
1327 self
1328 }
1329
1330 pub fn set_geo_statistics(mut self, value: Box<geo_statistics::GeospatialStatistics>) -> Self {
1332 self.0.geo_statistics = Some(value);
1333 self
1334 }
1335
1336 pub fn clear_statistics(mut self) -> Self {
1338 self.0.statistics = None;
1339 self
1340 }
1341
1342 pub fn set_page_encoding_stats(mut self, value: Vec<PageEncodingStats>) -> Self {
1346 self.0.encoding_stats = Some(ParquetPageEncodingStats::Full(value));
1347 self
1348 }
1349
1350 pub fn set_page_encoding_stats_mask(mut self, value: EncodingMask) -> Self {
1354 self.0.encoding_stats = Some(ParquetPageEncodingStats::Mask(value));
1355 self
1356 }
1357
1358 pub fn clear_page_encoding_stats(mut self) -> Self {
1360 self.0.encoding_stats = None;
1361 self
1362 }
1363
1364 pub fn set_bloom_filter_offset(mut self, value: Option<i64>) -> Self {
1366 self.0.bloom_filter_offset = value;
1367 self
1368 }
1369
1370 pub fn set_bloom_filter_length(mut self, value: Option<i32>) -> Self {
1372 self.0.bloom_filter_length = value;
1373 self
1374 }
1375
1376 pub fn set_offset_index_offset(mut self, value: Option<i64>) -> Self {
1378 self.0.offset_index_offset = value;
1379 self
1380 }
1381
1382 pub fn set_offset_index_length(mut self, value: Option<i32>) -> Self {
1384 self.0.offset_index_length = value;
1385 self
1386 }
1387
1388 pub fn set_column_index_offset(mut self, value: Option<i64>) -> Self {
1390 self.0.column_index_offset = value;
1391 self
1392 }
1393
1394 pub fn set_column_index_length(mut self, value: Option<i32>) -> Self {
1396 self.0.column_index_length = value;
1397 self
1398 }
1399
1400 pub fn set_unencoded_byte_array_data_bytes(mut self, value: Option<i64>) -> Self {
1402 self.0.unencoded_byte_array_data_bytes = value;
1403 self
1404 }
1405
1406 pub fn set_repetition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1408 self.0.repetition_level_histogram = value;
1409 self
1410 }
1411
1412 pub fn set_definition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1414 self.0.definition_level_histogram = value;
1415 self
1416 }
1417
1418 #[cfg(feature = "encryption")]
1419 pub fn set_column_crypto_metadata(mut self, value: Option<ColumnCryptoMetaData>) -> Self {
1421 self.0.column_crypto_metadata = value.map(Box::new);
1422 self
1423 }
1424
1425 #[cfg(feature = "encryption")]
1426 pub fn set_encrypted_column_metadata(mut self, value: Option<Vec<u8>>) -> Self {
1428 self.0.encrypted_column_metadata = value;
1429 self
1430 }
1431
1432 pub fn build(self) -> Result<ColumnChunkMetaData> {
1434 Ok(self.0)
1435 }
1436}
1437
1438pub struct ColumnIndexBuilder {
1443 column_type: Type,
1444 null_pages: Vec<bool>,
1445 min_values: Vec<Vec<u8>>,
1446 max_values: Vec<Vec<u8>>,
1447 null_counts: Vec<i64>,
1448 boundary_order: BoundaryOrder,
1449 repetition_level_histograms: Option<Vec<i64>>,
1451 definition_level_histograms: Option<Vec<i64>>,
1453 valid: bool,
1461}
1462
1463impl ColumnIndexBuilder {
1464 pub fn new(column_type: Type) -> Self {
1466 ColumnIndexBuilder {
1467 column_type,
1468 null_pages: Vec::new(),
1469 min_values: Vec::new(),
1470 max_values: Vec::new(),
1471 null_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(
1481 &mut self,
1482 null_page: bool,
1483 min_value: Vec<u8>,
1484 max_value: Vec<u8>,
1485 null_count: i64,
1486 ) {
1487 self.null_pages.push(null_page);
1488 self.min_values.push(min_value);
1489 self.max_values.push(max_value);
1490 self.null_counts.push(null_count);
1491 }
1492
1493 pub fn append_histograms(
1498 &mut self,
1499 repetition_level_histogram: &Option<LevelHistogram>,
1500 definition_level_histogram: &Option<LevelHistogram>,
1501 ) {
1502 if !self.valid {
1503 return;
1504 }
1505 if let Some(rep_lvl_hist) = repetition_level_histogram {
1506 let hist = self.repetition_level_histograms.get_or_insert(Vec::new());
1507 hist.reserve(rep_lvl_hist.len());
1508 hist.extend(rep_lvl_hist.values());
1509 }
1510 if let Some(def_lvl_hist) = definition_level_histogram {
1511 let hist = self.definition_level_histograms.get_or_insert(Vec::new());
1512 hist.reserve(def_lvl_hist.len());
1513 hist.extend(def_lvl_hist.values());
1514 }
1515 }
1516
1517 pub fn set_boundary_order(&mut self, boundary_order: BoundaryOrder) {
1519 self.boundary_order = boundary_order;
1520 }
1521
1522 pub fn to_invalid(&mut self) {
1524 self.valid = false;
1525 }
1526
1527 pub fn valid(&self) -> bool {
1529 self.valid
1530 }
1531
1532 pub fn build(self) -> Result<ColumnIndexMetaData> {
1536 Ok(match self.column_type {
1537 Type::BOOLEAN => {
1538 let index = self.build_page_index()?;
1539 ColumnIndexMetaData::BOOLEAN(index)
1540 }
1541 Type::INT32 => {
1542 let index = self.build_page_index()?;
1543 ColumnIndexMetaData::INT32(index)
1544 }
1545 Type::INT64 => {
1546 let index = self.build_page_index()?;
1547 ColumnIndexMetaData::INT64(index)
1548 }
1549 Type::INT96 => {
1550 let index = self.build_page_index()?;
1551 ColumnIndexMetaData::INT96(index)
1552 }
1553 Type::FLOAT => {
1554 let index = self.build_page_index()?;
1555 ColumnIndexMetaData::FLOAT(index)
1556 }
1557 Type::DOUBLE => {
1558 let index = self.build_page_index()?;
1559 ColumnIndexMetaData::DOUBLE(index)
1560 }
1561 Type::BYTE_ARRAY => {
1562 let index = self.build_byte_array_index()?;
1563 ColumnIndexMetaData::BYTE_ARRAY(index)
1564 }
1565 Type::FIXED_LEN_BYTE_ARRAY => {
1566 let index = self.build_byte_array_index()?;
1567 ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index)
1568 }
1569 })
1570 }
1571
1572 fn build_page_index<T>(self) -> Result<PrimitiveColumnIndex<T>>
1573 where
1574 T: ParquetValueType,
1575 {
1576 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1577 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1578
1579 PrimitiveColumnIndex::try_new(
1580 self.null_pages,
1581 self.boundary_order,
1582 Some(self.null_counts),
1583 self.repetition_level_histograms,
1584 self.definition_level_histograms,
1585 min_values,
1586 max_values,
1587 )
1588 }
1589
1590 fn build_byte_array_index(self) -> Result<ByteArrayColumnIndex> {
1591 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1592 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1593
1594 ByteArrayColumnIndex::try_new(
1595 self.null_pages,
1596 self.boundary_order,
1597 Some(self.null_counts),
1598 self.repetition_level_histograms,
1599 self.definition_level_histograms,
1600 min_values,
1601 max_values,
1602 )
1603 }
1604}
1605
1606impl From<ColumnChunkMetaData> for ColumnChunkMetaDataBuilder {
1607 fn from(value: ColumnChunkMetaData) -> Self {
1608 ColumnChunkMetaDataBuilder(value)
1609 }
1610}
1611
1612pub struct OffsetIndexBuilder {
1616 offset_array: Vec<i64>,
1617 compressed_page_size_array: Vec<i32>,
1618 first_row_index_array: Vec<i64>,
1619 unencoded_byte_array_data_bytes_array: Option<Vec<i64>>,
1620 current_first_row_index: i64,
1621}
1622
1623impl Default for OffsetIndexBuilder {
1624 fn default() -> Self {
1625 Self::new()
1626 }
1627}
1628
1629impl OffsetIndexBuilder {
1630 pub fn new() -> Self {
1632 OffsetIndexBuilder {
1633 offset_array: Vec::new(),
1634 compressed_page_size_array: Vec::new(),
1635 first_row_index_array: Vec::new(),
1636 unencoded_byte_array_data_bytes_array: None,
1637 current_first_row_index: 0,
1638 }
1639 }
1640
1641 pub fn append_row_count(&mut self, row_count: i64) {
1643 let current_page_row_index = self.current_first_row_index;
1644 self.first_row_index_array.push(current_page_row_index);
1645 self.current_first_row_index += row_count;
1646 }
1647
1648 pub fn append_offset_and_size(&mut self, offset: i64, compressed_page_size: i32) {
1650 self.offset_array.push(offset);
1651 self.compressed_page_size_array.push(compressed_page_size);
1652 }
1653
1654 pub fn append_unencoded_byte_array_data_bytes(
1656 &mut self,
1657 unencoded_byte_array_data_bytes: Option<i64>,
1658 ) {
1659 if let Some(val) = unencoded_byte_array_data_bytes {
1660 self.unencoded_byte_array_data_bytes_array
1661 .get_or_insert(Vec::new())
1662 .push(val);
1663 }
1664 }
1665
1666 pub fn build(self) -> OffsetIndexMetaData {
1668 let locations = self
1669 .offset_array
1670 .iter()
1671 .zip(self.compressed_page_size_array.iter())
1672 .zip(self.first_row_index_array.iter())
1673 .map(|((offset, size), row_index)| PageLocation {
1674 offset: *offset,
1675 compressed_page_size: *size,
1676 first_row_index: *row_index,
1677 })
1678 .collect::<Vec<_>>();
1679 OffsetIndexMetaData {
1680 page_locations: locations,
1681 unencoded_byte_array_data_bytes: self.unencoded_byte_array_data_bytes_array,
1682 }
1683 }
1684}
1685
1686#[cfg(test)]
1687mod tests {
1688 use super::*;
1689 use crate::basic::{PageType, SortOrder};
1690 use crate::file::metadata::thrift::tests::{
1691 read_column_chunk, read_column_chunk_with_options, read_row_group,
1692 };
1693
1694 #[test]
1695 #[allow(deprecated)]
1696 fn test_level_histogram_update_from_levels_compat() {
1697 let mut histogram = LevelHistogram::try_new(2).unwrap();
1698 histogram.update_from_levels(&[0, 2, 1, 2, 2]);
1699 assert_eq!(histogram.values(), &[1, 1, 3]);
1700 }
1701
1702 #[test]
1703 fn test_row_group_metadata_thrift_conversion() {
1704 let schema_descr = get_test_schema_descr();
1705
1706 let mut columns = vec![];
1707 for ptr in schema_descr.columns() {
1708 let column = ColumnChunkMetaData::builder(ptr.clone()).build().unwrap();
1709 columns.push(column);
1710 }
1711 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
1712 .set_num_rows(1000)
1713 .set_total_byte_size(2000)
1714 .set_column_metadata(columns)
1715 .set_ordinal(1)
1716 .build()
1717 .unwrap();
1718
1719 let mut buf = Vec::new();
1720 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1721 row_group_meta.write_thrift(&mut writer).unwrap();
1722
1723 let row_group_res = read_row_group(&mut buf, schema_descr).unwrap();
1724
1725 assert_eq!(row_group_res, row_group_meta);
1726 }
1727
1728 #[test]
1729 fn test_row_group_metadata_thrift_conversion_empty() {
1730 let schema_descr = get_test_schema_descr();
1731
1732 let row_group_meta = RowGroupMetaData::builder(schema_descr).build();
1733
1734 assert!(row_group_meta.is_err());
1735 if let Err(e) = row_group_meta {
1736 assert_eq!(
1737 format!("{e}"),
1738 "Parquet error: Column length mismatch: 2 != 0"
1739 );
1740 }
1741 }
1742
1743 #[test]
1745 fn test_row_group_metadata_thrift_corrupted() {
1746 let schema_descr_2cols = Arc::new(SchemaDescriptor::new(Arc::new(
1747 SchemaType::group_type_builder("schema")
1748 .with_fields(vec![
1749 Arc::new(
1750 SchemaType::primitive_type_builder("a", Type::INT32)
1751 .build()
1752 .unwrap(),
1753 ),
1754 Arc::new(
1755 SchemaType::primitive_type_builder("b", Type::INT32)
1756 .build()
1757 .unwrap(),
1758 ),
1759 ])
1760 .build()
1761 .unwrap(),
1762 )));
1763
1764 let schema_descr_3cols = Arc::new(SchemaDescriptor::new(Arc::new(
1765 SchemaType::group_type_builder("schema")
1766 .with_fields(vec![
1767 Arc::new(
1768 SchemaType::primitive_type_builder("a", Type::INT32)
1769 .build()
1770 .unwrap(),
1771 ),
1772 Arc::new(
1773 SchemaType::primitive_type_builder("b", Type::INT32)
1774 .build()
1775 .unwrap(),
1776 ),
1777 Arc::new(
1778 SchemaType::primitive_type_builder("c", Type::INT32)
1779 .build()
1780 .unwrap(),
1781 ),
1782 ])
1783 .build()
1784 .unwrap(),
1785 )));
1786
1787 let row_group_meta_2cols = RowGroupMetaData::builder(schema_descr_2cols.clone())
1788 .set_num_rows(1000)
1789 .set_total_byte_size(2000)
1790 .set_column_metadata(vec![
1791 ColumnChunkMetaData::builder(schema_descr_2cols.column(0))
1792 .build()
1793 .unwrap(),
1794 ColumnChunkMetaData::builder(schema_descr_2cols.column(1))
1795 .build()
1796 .unwrap(),
1797 ])
1798 .set_ordinal(1)
1799 .build()
1800 .unwrap();
1801 let mut buf = Vec::new();
1802 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1803 row_group_meta_2cols.write_thrift(&mut writer).unwrap();
1804
1805 let err = read_row_group(&mut buf, schema_descr_3cols)
1806 .unwrap_err()
1807 .to_string();
1808 assert_eq!(
1809 err,
1810 "Parquet error: Column count mismatch. Schema has 3 columns while Row Group has 2"
1811 );
1812 }
1813
1814 #[test]
1815 fn test_column_chunk_metadata_thrift_conversion() {
1816 let column_descr = get_test_schema_descr().column(0);
1817 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1818 .set_encodings_mask(EncodingMask::new_from_encodings(
1819 [Encoding::PLAIN, Encoding::RLE].iter(),
1820 ))
1821 .set_file_path("file_path".to_owned())
1822 .set_num_values(1000)
1823 .set_compression(Compression::SNAPPY)
1824 .set_total_compressed_size(2000)
1825 .set_total_uncompressed_size(3000)
1826 .set_data_page_offset(4000)
1827 .set_dictionary_page_offset(Some(5000))
1828 .set_page_encoding_stats(vec![
1829 PageEncodingStats {
1830 page_type: PageType::DATA_PAGE,
1831 encoding: Encoding::PLAIN,
1832 count: 3,
1833 },
1834 PageEncodingStats {
1835 page_type: PageType::DATA_PAGE,
1836 encoding: Encoding::RLE,
1837 count: 5,
1838 },
1839 ])
1840 .set_bloom_filter_offset(Some(6000))
1841 .set_bloom_filter_length(Some(25))
1842 .set_offset_index_offset(Some(7000))
1843 .set_offset_index_length(Some(25))
1844 .set_column_index_offset(Some(8000))
1845 .set_column_index_length(Some(25))
1846 .set_unencoded_byte_array_data_bytes(Some(2000))
1847 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
1848 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
1849 .build()
1850 .unwrap();
1851
1852 let mut buf = Vec::new();
1853 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1854 col_metadata.write_thrift(&mut writer).unwrap();
1855 let col_chunk_res = read_column_chunk(&mut buf, column_descr.clone()).unwrap();
1856
1857 let expected_metadata = ColumnChunkMetaData::builder(column_descr)
1858 .set_encodings_mask(EncodingMask::new_from_encodings(
1859 [Encoding::PLAIN, Encoding::RLE].iter(),
1860 ))
1861 .set_file_path("file_path".to_owned())
1862 .set_num_values(1000)
1863 .set_compression(Compression::SNAPPY)
1864 .set_total_compressed_size(2000)
1865 .set_total_uncompressed_size(3000)
1866 .set_data_page_offset(4000)
1867 .set_dictionary_page_offset(Some(5000))
1868 .set_page_encoding_stats_mask(EncodingMask::new_from_encodings(
1869 [Encoding::PLAIN, Encoding::RLE].iter(),
1870 ))
1871 .set_bloom_filter_offset(Some(6000))
1872 .set_bloom_filter_length(Some(25))
1873 .set_offset_index_offset(Some(7000))
1874 .set_offset_index_length(Some(25))
1875 .set_column_index_offset(Some(8000))
1876 .set_column_index_length(Some(25))
1877 .set_unencoded_byte_array_data_bytes(Some(2000))
1878 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
1879 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
1880 .build()
1881 .unwrap();
1882
1883 assert_eq!(col_chunk_res, expected_metadata);
1884 }
1885
1886 #[test]
1887 fn test_column_chunk_metadata_thrift_conversion_full_stats() {
1888 let column_descr = get_test_schema_descr().column(0);
1889 let stats = vec![
1890 PageEncodingStats {
1891 page_type: PageType::DATA_PAGE,
1892 encoding: Encoding::PLAIN,
1893 count: 3,
1894 },
1895 PageEncodingStats {
1896 page_type: PageType::DATA_PAGE,
1897 encoding: Encoding::RLE,
1898 count: 5,
1899 },
1900 ];
1901 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1902 .set_encodings_mask(EncodingMask::new_from_encodings(
1903 [Encoding::PLAIN, Encoding::RLE].iter(),
1904 ))
1905 .set_num_values(1000)
1906 .set_compression(Compression::SNAPPY)
1907 .set_total_compressed_size(2000)
1908 .set_total_uncompressed_size(3000)
1909 .set_data_page_offset(4000)
1910 .set_page_encoding_stats(stats)
1911 .build()
1912 .unwrap();
1913
1914 let mut buf = Vec::new();
1915 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1916 col_metadata.write_thrift(&mut writer).unwrap();
1917
1918 let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
1919 let col_chunk_res =
1920 read_column_chunk_with_options(&mut buf, column_descr, Some(&options)).unwrap();
1921
1922 assert_eq!(col_chunk_res, col_metadata);
1923 }
1924
1925 #[test]
1926 fn test_column_chunk_metadata_thrift_conversion_empty() {
1927 let column_descr = get_test_schema_descr().column(0);
1928
1929 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1930 .build()
1931 .unwrap();
1932
1933 let mut buf = Vec::new();
1934 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1935 col_metadata.write_thrift(&mut writer).unwrap();
1936 let col_chunk_res = read_column_chunk(&mut buf, column_descr).unwrap();
1937
1938 assert_eq!(col_chunk_res, col_metadata);
1939 }
1940
1941 #[test]
1942 fn test_compressed_size() {
1943 let schema_descr = get_test_schema_descr();
1944
1945 let mut columns = vec![];
1946 for column_descr in schema_descr.columns() {
1947 let column = ColumnChunkMetaData::builder(column_descr.clone())
1948 .set_total_compressed_size(500)
1949 .set_total_uncompressed_size(700)
1950 .build()
1951 .unwrap();
1952 columns.push(column);
1953 }
1954 let row_group_meta = RowGroupMetaData::builder(schema_descr)
1955 .set_num_rows(1000)
1956 .set_column_metadata(columns)
1957 .build()
1958 .unwrap();
1959
1960 let compressed_size_res: i64 = row_group_meta.compressed_size();
1961 let compressed_size_exp: i64 = 1000;
1962
1963 assert_eq!(compressed_size_res, compressed_size_exp);
1964 }
1965
1966 #[test]
1967 fn test_memory_size() {
1968 let schema_descr = get_test_schema_descr();
1969
1970 let columns = schema_descr
1971 .columns()
1972 .iter()
1973 .map(|column_descr| {
1974 ColumnChunkMetaData::builder(column_descr.clone())
1975 .set_statistics(Statistics::new::<i32>(None, None, None, None, false))
1976 .build()
1977 })
1978 .collect::<Result<Vec<_>>>()
1979 .unwrap();
1980 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
1981 .set_num_rows(1000)
1982 .set_column_metadata(columns)
1983 .build()
1984 .unwrap();
1985 let row_group_meta = vec![row_group_meta];
1986
1987 let version = 2;
1988 let num_rows = 1000;
1989 let created_by = Some(String::from("test harness"));
1990 let key_value_metadata = Some(vec![KeyValue::new(
1991 String::from("Foo"),
1992 Some(String::from("bar")),
1993 )]);
1994 let column_orders = Some(vec![
1995 ColumnOrder::UNDEFINED,
1996 ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED),
1997 ]);
1998 let file_metadata = FileMetaData::new(
1999 version,
2000 num_rows,
2001 created_by,
2002 key_value_metadata,
2003 schema_descr.clone(),
2004 column_orders,
2005 );
2006
2007 let columns_with_stats = schema_descr
2009 .columns()
2010 .iter()
2011 .map(|column_descr| {
2012 ColumnChunkMetaData::builder(column_descr.clone())
2013 .set_statistics(Statistics::new::<i32>(
2014 Some(0),
2015 Some(100),
2016 None,
2017 None,
2018 false,
2019 ))
2020 .build()
2021 })
2022 .collect::<Result<Vec<_>>>()
2023 .unwrap();
2024
2025 let row_group_meta_with_stats = RowGroupMetaData::builder(schema_descr)
2026 .set_num_rows(1000)
2027 .set_column_metadata(columns_with_stats)
2028 .build()
2029 .unwrap();
2030 let row_group_meta_with_stats = vec![row_group_meta_with_stats];
2031
2032 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata.clone())
2033 .set_row_groups(row_group_meta_with_stats)
2034 .build();
2035
2036 #[cfg(not(feature = "encryption"))]
2037 let base_expected_size = 2766;
2038 #[cfg(feature = "encryption")]
2039 let base_expected_size = 2934;
2040
2041 assert_eq!(parquet_meta.memory_size(), base_expected_size);
2042
2043 let mut column_index = ColumnIndexBuilder::new(Type::BOOLEAN);
2044 column_index.append(false, vec![1u8], vec![2u8, 3u8], 4);
2045 let column_index = column_index.build().unwrap();
2046 let native_index = match column_index {
2047 ColumnIndexMetaData::BOOLEAN(index) => index,
2048 _ => panic!("wrong type of column index"),
2049 };
2050
2051 let mut offset_index = OffsetIndexBuilder::new();
2053 offset_index.append_row_count(1);
2054 offset_index.append_offset_and_size(2, 3);
2055 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2056 offset_index.append_row_count(1);
2057 offset_index.append_offset_and_size(2, 3);
2058 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2059 let offset_index = offset_index.build();
2060
2061 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata)
2062 .set_row_groups(row_group_meta)
2063 .set_column_index(Some(vec![vec![ColumnIndexMetaData::BOOLEAN(native_index)]]))
2064 .set_offset_index(Some(vec![vec![offset_index]]))
2065 .build();
2066
2067 #[cfg(not(feature = "encryption"))]
2068 let bigger_expected_size = 3192;
2069 #[cfg(feature = "encryption")]
2070 let bigger_expected_size = 3360;
2071
2072 assert!(bigger_expected_size > base_expected_size);
2074 assert_eq!(parquet_meta.memory_size(), bigger_expected_size);
2075 }
2076
2077 #[test]
2078 #[cfg(feature = "encryption")]
2079 fn test_memory_size_with_decryptor() {
2080 use crate::encryption::decrypt::FileDecryptionProperties;
2081 use crate::file::metadata::thrift::encryption::AesGcmV1;
2082
2083 let schema_descr = get_test_schema_descr();
2084
2085 let columns = schema_descr
2086 .columns()
2087 .iter()
2088 .map(|column_descr| ColumnChunkMetaData::builder(column_descr.clone()).build())
2089 .collect::<Result<Vec<_>>>()
2090 .unwrap();
2091 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2092 .set_num_rows(1000)
2093 .set_column_metadata(columns)
2094 .build()
2095 .unwrap();
2096 let row_group_meta = vec![row_group_meta];
2097
2098 let version = 2;
2099 let num_rows = 1000;
2100 let aad_file_unique = vec![1u8; 8];
2101 let aad_prefix = vec![2u8; 8];
2102 let encryption_algorithm = EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 {
2103 aad_prefix: Some(aad_prefix.clone()),
2104 aad_file_unique: Some(aad_file_unique.clone()),
2105 supply_aad_prefix: Some(true),
2106 });
2107 let footer_key_metadata = Some(vec![3u8; 8]);
2108 let file_metadata =
2109 FileMetaData::new(version, num_rows, None, None, schema_descr.clone(), None)
2110 .with_encryption_algorithm(Some(encryption_algorithm))
2111 .with_footer_signing_key_metadata(footer_key_metadata.clone());
2112
2113 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2114 .set_row_groups(row_group_meta.clone())
2115 .build();
2116
2117 let base_expected_size = 2058;
2118 assert_eq!(parquet_meta_data.memory_size(), base_expected_size);
2119
2120 let footer_key = "0123456789012345".as_bytes();
2121 let column_key = "1234567890123450".as_bytes();
2122 let mut decryption_properties_builder =
2123 FileDecryptionProperties::builder(footer_key.to_vec())
2124 .with_aad_prefix(aad_prefix.clone());
2125 for column in schema_descr.columns() {
2126 decryption_properties_builder = decryption_properties_builder
2127 .with_column_key(&column.path().string(), column_key.to_vec());
2128 }
2129 let decryption_properties = decryption_properties_builder.build().unwrap();
2130 let decryptor = FileDecryptor::new(
2131 &decryption_properties,
2132 footer_key_metadata.as_deref(),
2133 aad_file_unique,
2134 aad_prefix,
2135 )
2136 .unwrap();
2137
2138 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2139 .set_row_groups(row_group_meta.clone())
2140 .set_file_decryptor(Some(decryptor))
2141 .build();
2142
2143 let expected_size_with_decryptor = 3072;
2144 assert!(expected_size_with_decryptor > base_expected_size);
2145
2146 assert_eq!(
2147 parquet_meta_data.memory_size(),
2148 expected_size_with_decryptor
2149 );
2150 }
2151
2152 fn get_test_schema_descr() -> SchemaDescPtr {
2154 let schema = SchemaType::group_type_builder("schema")
2155 .with_fields(vec![
2156 Arc::new(
2157 SchemaType::primitive_type_builder("a", Type::INT32)
2158 .build()
2159 .unwrap(),
2160 ),
2161 Arc::new(
2162 SchemaType::primitive_type_builder("b", Type::INT32)
2163 .build()
2164 .unwrap(),
2165 ),
2166 ])
2167 .build()
2168 .unwrap();
2169
2170 Arc::new(SchemaDescriptor::new(Arc::new(schema)))
2171 }
2172}