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::{
99 BoundaryOrder, ColumnOrder, Compression, CompressionCodec, Encoding, EncodingMask, PageType,
100 Type,
101};
102#[cfg(feature = "encryption")]
103use crate::encryption::decrypt::FileDecryptor;
104use crate::errors::{ParquetError, Result};
105#[cfg(feature = "encryption")]
106use crate::file::column_crypto_metadata::ColumnCryptoMetaData;
107pub(crate) use crate::file::metadata::memory::HeapSize;
108#[cfg(feature = "encryption")]
109use crate::file::metadata::thrift::encryption::EncryptionAlgorithm;
110use crate::file::page_index::column_index::{ByteArrayColumnIndex, PrimitiveColumnIndex};
111use crate::file::page_index::{column_index::ColumnIndexMetaData, offset_index::PageLocation};
112use crate::file::statistics::Statistics;
113use crate::geospatial::statistics as geo_statistics;
114use crate::parquet_thrift::{
115 ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
116 WriteThrift, WriteThriftField,
117};
118use crate::schema::types::{
119 ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, SchemaDescriptor,
120 Type as SchemaType,
121};
122use crate::thrift_struct;
123use crate::{
124 data_type::private::ParquetValueType, file::page_index::offset_index::OffsetIndexMetaData,
125};
126
127pub use footer_tail::FooterTail;
128pub use options::{ParquetMetaDataOptions, ParquetStatisticsPolicy};
129pub use push_decoder::ParquetMetaDataPushDecoder;
130pub use reader::{PageIndexPolicy, ParquetMetaDataReader};
131use std::io::Write;
132use std::ops::Range;
133use std::sync::Arc;
134pub use writer::ParquetMetaDataWriter;
135pub(crate) use writer::ThriftMetadataWriter;
136
137#[derive(Debug, Clone, PartialEq)]
250pub struct PageIndex {
251 column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>,
252 offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>,
253}
254
255impl PageIndex {
256 pub(crate) fn new(
257 column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>,
258 offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>,
259 ) -> Self {
260 Self {
261 column_indexes,
262 offset_indexes,
263 }
264 }
265
266 pub fn has_offset_indexes(&self) -> bool {
273 self.offset_indexes.is_some()
274 }
275
276 pub fn has_column_indexes(&self) -> bool {
283 self.column_indexes.is_some()
284 }
285
286 pub fn is_complete(&self) -> bool {
291 self.has_column_indexes() && self.has_offset_indexes()
292 }
293
294 pub fn column_indexes_for_rowgroup(
304 &self,
305 row_group_idx: usize,
306 ) -> Option<&[Option<ColumnIndexMetaData>]> {
307 match self.column_indexes.as_ref() {
308 None => None,
309 Some(indexes) => indexes.get(row_group_idx).map(|ci| ci.as_slice()),
310 }
311 }
312
313 pub fn column_index(
322 &self,
323 row_group_idx: usize,
324 column_idx: usize,
325 ) -> Option<&ColumnIndexMetaData> {
326 if let Some(column_indexes) = self.column_indexes.as_ref() {
327 let rg = column_indexes.get(row_group_idx)?;
328 rg.get(column_idx)?.as_ref()
329 } else {
330 None
331 }
332 }
333
334 pub fn offset_indexes_for_rowgroup(
344 &self,
345 row_group_idx: usize,
346 ) -> Option<&[Option<OffsetIndexMetaData>]> {
347 match self.offset_indexes.as_ref() {
348 None => None,
349 Some(indexes) => indexes.get(row_group_idx).map(|oi| oi.as_slice()),
350 }
351 }
352
353 pub fn offset_index(
364 &self,
365 row_group_idx: usize,
366 column_idx: usize,
367 ) -> Option<&OffsetIndexMetaData> {
368 if let Some(offset_indexes) = self.offset_indexes.as_ref() {
369 let rg = offset_indexes.get(row_group_idx)?;
370 rg.get(column_idx)?.as_ref()
371 } else {
372 None
373 }
374 }
375
376 pub fn num_data_pages(&self, row_group_idx: usize, column_idx: usize) -> Option<usize> {
384 match self.offset_index(row_group_idx, column_idx) {
385 Some(offset_index) => Some(offset_index.page_locations.len()),
386 None => Some(self.column_index(row_group_idx, column_idx)?.num_pages() as usize),
387 }
388 }
389
390 pub fn page_locations(
403 &self,
404 row_group_idx: usize,
405 column_idx: usize,
406 ) -> Option<&Vec<PageLocation>> {
407 if let Some(offset_indexes) = self.offset_indexes.as_ref() {
408 let rg = offset_indexes.get(row_group_idx)?;
409 let off_idx = rg.get(column_idx)?.as_ref()?;
410 Some(off_idx.page_locations())
411 } else {
412 None
413 }
414 }
415}
416
417#[derive(Debug, Clone, PartialEq)]
435pub struct ParquetMetaData {
436 file_metadata: FileMetaData,
438 row_groups: Vec<RowGroupMetaData>,
440 page_index: Option<PageIndex>,
442 #[cfg(feature = "encryption")]
444 file_decryptor: Option<Box<FileDecryptor>>,
445}
446
447impl ParquetMetaData {
448 pub fn new(file_metadata: FileMetaData, row_groups: Vec<RowGroupMetaData>) -> Self {
451 ParquetMetaData {
452 file_metadata,
453 row_groups,
454 page_index: None,
455 #[cfg(feature = "encryption")]
456 file_decryptor: None,
457 }
458 }
459
460 #[cfg(feature = "encryption")]
463 pub(crate) fn with_file_decryptor(&mut self, file_decryptor: Option<FileDecryptor>) {
464 self.file_decryptor = file_decryptor.map(Box::new);
465 }
466
467 pub fn into_builder(self) -> ParquetMetaDataBuilder {
469 self.into()
470 }
471
472 pub fn file_metadata(&self) -> &FileMetaData {
474 &self.file_metadata
475 }
476
477 #[cfg(feature = "encryption")]
479 pub(crate) fn file_decryptor(&self) -> Option<&FileDecryptor> {
480 self.file_decryptor.as_deref()
481 }
482
483 pub fn num_row_groups(&self) -> usize {
485 self.row_groups.len()
486 }
487
488 pub fn row_group(&self, i: usize) -> &RowGroupMetaData {
491 &self.row_groups[i]
492 }
493
494 pub fn row_groups(&self) -> &[RowGroupMetaData] {
496 &self.row_groups
497 }
498
499 pub fn page_index(&self) -> Option<&PageIndex> {
506 self.page_index.as_ref()
507 }
508
509 pub fn memory_size(&self) -> usize {
524 #[cfg(feature = "encryption")]
525 let encryption_size = self.file_decryptor.heap_size();
526 #[cfg(not(feature = "encryption"))]
527 let encryption_size = 0usize;
528
529 std::mem::size_of::<Self>()
530 + self.file_metadata.heap_size()
531 + self.row_groups.heap_size()
532 + self.page_index.heap_size()
533 + encryption_size
534 }
535
536 pub(crate) fn set_page_index(&mut self, index: Option<PageIndex>) {
538 self.page_index = index;
539 }
540}
541
542pub struct ParquetMetaDataBuilder(ParquetMetaData);
580
581impl ParquetMetaDataBuilder {
582 pub fn new(file_meta_data: FileMetaData) -> Self {
584 Self(ParquetMetaData::new(file_meta_data, vec![]))
585 }
586
587 pub fn new_from_metadata(metadata: ParquetMetaData) -> Self {
589 Self(metadata)
590 }
591
592 pub fn add_row_group(mut self, row_group: RowGroupMetaData) -> Self {
594 self.0.row_groups.push(row_group);
595 self
596 }
597
598 pub fn set_row_groups(mut self, row_groups: Vec<RowGroupMetaData>) -> Self {
600 self.0.row_groups = row_groups;
601 self
602 }
603
604 pub fn take_row_groups(&mut self) -> Vec<RowGroupMetaData> {
610 std::mem::take(&mut self.0.row_groups)
611 }
612
613 pub fn row_groups(&self) -> &[RowGroupMetaData] {
615 &self.0.row_groups
616 }
617
618 pub fn set_page_index(mut self, page_index: Option<PageIndex>) -> Self {
620 self.0.page_index = page_index;
621 self
622 }
623
624 pub fn take_page_index(&mut self) -> Option<PageIndex> {
626 std::mem::take(&mut self.0.page_index)
627 }
628
629 pub fn page_index(&self) -> Option<&PageIndex> {
631 self.0.page_index.as_ref()
632 }
633
634 #[cfg(feature = "encryption")]
636 pub(crate) fn set_file_decryptor(mut self, file_decryptor: Option<FileDecryptor>) -> Self {
637 self.0.with_file_decryptor(file_decryptor);
638 self
639 }
640
641 pub fn build(self) -> ParquetMetaData {
643 let Self(metadata) = self;
644 metadata
645 }
646}
647
648impl From<ParquetMetaData> for ParquetMetaDataBuilder {
649 fn from(meta_data: ParquetMetaData) -> Self {
650 Self(meta_data)
651 }
652}
653
654thrift_struct!(
655pub struct KeyValue {
657 1: required string key
658 2: optional string value
659}
660);
661
662impl KeyValue {
663 pub fn new<F2>(key: String, value: F2) -> KeyValue
665 where
666 F2: Into<Option<String>>,
667 {
668 KeyValue {
669 key,
670 value: value.into(),
671 }
672 }
673}
674
675thrift_struct!(
676pub struct PageEncodingStats {
678 1: required PageType page_type;
679 2: required Encoding encoding;
680 3: required i32 count;
681}
682);
683
684#[derive(Debug, Clone, PartialEq)]
687enum ParquetPageEncodingStats {
688 Full(Vec<PageEncodingStats>),
690 Mask(EncodingMask),
692}
693
694pub type FileMetaDataPtr = Arc<FileMetaData>;
696
697#[derive(Debug, Clone, PartialEq)]
701pub struct FileMetaData {
702 version: i32,
703 num_rows: i64,
704 created_by: Option<String>,
705 key_value_metadata: Option<Vec<KeyValue>>,
706 schema_descr: SchemaDescPtr,
707 column_orders: Option<Vec<ColumnOrder>>,
708 #[cfg(feature = "encryption")]
709 encryption_algorithm: Option<Box<EncryptionAlgorithm>>,
710 #[cfg(feature = "encryption")]
711 footer_signing_key_metadata: Option<Vec<u8>>,
712}
713
714impl FileMetaData {
715 pub fn new(
717 version: i32,
718 num_rows: i64,
719 created_by: Option<String>,
720 key_value_metadata: Option<Vec<KeyValue>>,
721 schema_descr: SchemaDescPtr,
722 column_orders: Option<Vec<ColumnOrder>>,
723 ) -> Self {
724 FileMetaData {
725 version,
726 num_rows,
727 created_by,
728 key_value_metadata,
729 schema_descr,
730 column_orders,
731 #[cfg(feature = "encryption")]
732 encryption_algorithm: None,
733 #[cfg(feature = "encryption")]
734 footer_signing_key_metadata: None,
735 }
736 }
737
738 #[cfg(feature = "encryption")]
739 pub(crate) fn with_encryption_algorithm(
740 mut self,
741 encryption_algorithm: Option<EncryptionAlgorithm>,
742 ) -> Self {
743 self.encryption_algorithm = encryption_algorithm.map(Box::new);
744 self
745 }
746
747 #[cfg(feature = "encryption")]
748 pub(crate) fn with_footer_signing_key_metadata(
749 mut self,
750 footer_signing_key_metadata: Option<Vec<u8>>,
751 ) -> Self {
752 self.footer_signing_key_metadata = footer_signing_key_metadata;
753 self
754 }
755
756 pub fn version(&self) -> i32 {
758 self.version
759 }
760
761 pub fn num_rows(&self) -> i64 {
763 self.num_rows
764 }
765
766 pub fn created_by(&self) -> Option<&str> {
775 self.created_by.as_deref()
776 }
777
778 pub fn key_value_metadata(&self) -> Option<&Vec<KeyValue>> {
780 self.key_value_metadata.as_ref()
781 }
782
783 pub fn schema(&self) -> &SchemaType {
787 self.schema_descr.root_schema()
788 }
789
790 pub fn schema_descr(&self) -> &SchemaDescriptor {
792 &self.schema_descr
793 }
794
795 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
797 self.schema_descr.clone()
798 }
799
800 pub fn column_orders(&self) -> Option<&Vec<ColumnOrder>> {
808 self.column_orders.as_ref()
809 }
810
811 pub fn column_order(&self, i: usize) -> ColumnOrder {
814 self.column_orders
815 .as_ref()
816 .map(|data| data[i])
817 .unwrap_or(ColumnOrder::UNDEFINED)
818 }
819}
820
821thrift_struct!(
822pub struct SortingColumn {
824 1: required i32 column_idx
826
827 2: required bool descending
829
830 3: required bool nulls_first
833}
834);
835
836pub type RowGroupMetaDataPtr = Arc<RowGroupMetaData>;
838
839#[derive(Debug, Clone, PartialEq)]
844pub struct RowGroupMetaData {
845 columns: Vec<ColumnChunkMetaData>,
846 num_rows: i64,
847 sorting_columns: Option<Vec<SortingColumn>>,
848 total_byte_size: i64,
849 schema_descr: SchemaDescPtr,
850 file_offset: Option<i64>,
852 ordinal: Option<i32>,
854}
855
856impl RowGroupMetaData {
857 pub fn builder(schema_descr: SchemaDescPtr) -> RowGroupMetaDataBuilder {
859 RowGroupMetaDataBuilder::new(schema_descr)
860 }
861
862 pub fn num_columns(&self) -> usize {
864 self.columns.len()
865 }
866
867 pub fn column(&self, i: usize) -> &ColumnChunkMetaData {
869 &self.columns[i]
870 }
871
872 pub fn columns(&self) -> &[ColumnChunkMetaData] {
874 &self.columns
875 }
876
877 pub fn columns_mut(&mut self) -> &mut [ColumnChunkMetaData] {
879 &mut self.columns
880 }
881
882 pub fn num_rows(&self) -> i64 {
884 self.num_rows
885 }
886
887 pub fn sorting_columns(&self) -> Option<&Vec<SortingColumn>> {
889 self.sorting_columns.as_ref()
890 }
891
892 pub fn total_byte_size(&self) -> i64 {
894 self.total_byte_size
895 }
896
897 pub fn compressed_size(&self) -> i64 {
899 self.columns.iter().map(|c| c.total_compressed_size).sum()
900 }
901
902 pub fn schema_descr(&self) -> &SchemaDescriptor {
904 self.schema_descr.as_ref()
905 }
906
907 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
909 self.schema_descr.clone()
910 }
911
912 #[inline(always)]
917 pub fn ordinal(&self) -> Option<i32> {
918 self.ordinal
919 }
920
921 #[inline(always)]
923 pub fn file_offset(&self) -> Option<i64> {
924 self.file_offset
925 }
926
927 pub fn into_builder(self) -> RowGroupMetaDataBuilder {
929 RowGroupMetaDataBuilder(self)
930 }
931}
932
933pub struct RowGroupMetaDataBuilder(RowGroupMetaData);
935
936impl RowGroupMetaDataBuilder {
937 fn new(schema_descr: SchemaDescPtr) -> Self {
939 Self(RowGroupMetaData {
940 columns: Vec::with_capacity(schema_descr.num_columns()),
941 schema_descr,
942 file_offset: None,
943 num_rows: 0,
944 sorting_columns: None,
945 total_byte_size: 0,
946 ordinal: None,
947 })
948 }
949
950 pub fn set_num_rows(mut self, value: i64) -> Self {
952 self.0.num_rows = value;
953 self
954 }
955
956 pub fn set_sorting_columns(mut self, value: Option<Vec<SortingColumn>>) -> Self {
958 self.0.sorting_columns = value;
959 self
960 }
961
962 pub fn set_total_byte_size(mut self, value: i64) -> Self {
964 self.0.total_byte_size = value;
965 self
966 }
967
968 pub fn take_columns(&mut self) -> Vec<ColumnChunkMetaData> {
974 std::mem::take(&mut self.0.columns)
975 }
976
977 pub fn set_column_metadata(mut self, value: Vec<ColumnChunkMetaData>) -> Self {
979 self.0.columns = value;
980 self
981 }
982
983 pub fn add_column_metadata(mut self, value: ColumnChunkMetaData) -> Self {
985 self.0.columns.push(value);
986 self
987 }
988
989 pub fn set_ordinal(mut self, value: i32) -> Self {
991 self.0.ordinal = Some(value);
992 self
993 }
994
995 pub fn set_file_offset(mut self, value: i64) -> Self {
997 self.0.file_offset = Some(value);
998 self
999 }
1000
1001 pub fn build(self) -> Result<RowGroupMetaData> {
1003 if self.0.schema_descr.num_columns() != self.0.columns.len() {
1004 return Err(general_err!(
1005 "Column length mismatch: {} != {}",
1006 self.0.schema_descr.num_columns(),
1007 self.0.columns.len()
1008 ));
1009 }
1010
1011 Ok(self.0)
1012 }
1013
1014 pub(super) fn build_unchecked(self) -> RowGroupMetaData {
1016 self.0
1017 }
1018}
1019
1020#[derive(Debug, Clone, PartialEq)]
1022pub struct ColumnChunkMetaData {
1023 column_descr: ColumnDescPtr,
1024 encodings: EncodingMask,
1025 file_path: Option<String>,
1026 file_offset: i64,
1027 num_values: i64,
1028 compression: CompressionCodec,
1029 total_compressed_size: i64,
1030 total_uncompressed_size: i64,
1031 data_page_offset: i64,
1032 index_page_offset: Option<i64>,
1033 dictionary_page_offset: Option<i64>,
1034 statistics: Option<Statistics>,
1035 geo_statistics: Option<Box<geo_statistics::GeospatialStatistics>>,
1036 encoding_stats: Option<ParquetPageEncodingStats>,
1037 bloom_filter_offset: Option<i64>,
1038 bloom_filter_length: Option<i32>,
1039 offset_index_offset: Option<i64>,
1040 offset_index_length: Option<i32>,
1041 column_index_offset: Option<i64>,
1042 column_index_length: Option<i32>,
1043 unencoded_byte_array_data_bytes: Option<i64>,
1044 repetition_level_histogram: Option<LevelHistogram>,
1045 definition_level_histogram: Option<LevelHistogram>,
1046 #[cfg(feature = "encryption")]
1047 column_crypto_metadata: Option<Box<ColumnCryptoMetaData>>,
1048 #[cfg(feature = "encryption")]
1049 encrypted_column_metadata: Option<Vec<u8>>,
1050 #[cfg(feature = "encryption")]
1054 plaintext_footer_mode: bool,
1055}
1056
1057#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
1066pub struct LevelHistogram {
1067 inner: Vec<i64>,
1068}
1069
1070impl LevelHistogram {
1071 pub fn try_new(max_level: i16) -> Option<Self> {
1077 if max_level > 0 {
1078 Some(Self {
1079 inner: vec![0; max_level as usize + 1],
1080 })
1081 } else {
1082 None
1083 }
1084 }
1085 pub fn values(&self) -> &[i64] {
1087 &self.inner
1088 }
1089
1090 pub fn into_inner(self) -> Vec<i64> {
1092 self.inner
1093 }
1094
1095 pub fn get(&self, index: usize) -> Option<i64> {
1102 self.inner.get(index).copied()
1103 }
1104
1105 pub fn add(&mut self, other: &Self) {
1110 assert_eq!(self.len(), other.len());
1111 for (dst, src) in self.inner.iter_mut().zip(other.inner.iter()) {
1112 *dst += src;
1113 }
1114 }
1115
1116 pub fn len(&self) -> usize {
1118 self.inner.len()
1119 }
1120
1121 pub fn is_empty(&self) -> bool {
1123 self.inner.is_empty()
1124 }
1125
1126 pub fn reset(&mut self) {
1128 for value in &mut self.inner {
1129 *value = 0;
1130 }
1131 }
1132
1133 #[inline]
1135 pub fn increment_by(&mut self, level: i16, count: i64) {
1136 self.inner[level as usize] += count;
1137 }
1138
1139 #[deprecated(since = "58.2.0", note = "Use `increment_by` instead")]
1145 pub fn update_from_levels(&mut self, levels: &[i16]) {
1146 for &level in levels {
1147 self.increment_by(level, 1);
1148 }
1149 }
1150}
1151
1152impl From<Vec<i64>> for LevelHistogram {
1153 fn from(inner: Vec<i64>) -> Self {
1154 Self { inner }
1155 }
1156}
1157
1158impl From<LevelHistogram> for Vec<i64> {
1159 fn from(value: LevelHistogram) -> Self {
1160 value.into_inner()
1161 }
1162}
1163
1164impl HeapSize for LevelHistogram {
1165 fn heap_size(&self) -> usize {
1166 self.inner.heap_size()
1167 }
1168}
1169
1170impl ColumnChunkMetaData {
1172 pub fn builder(column_descr: ColumnDescPtr) -> ColumnChunkMetaDataBuilder {
1174 ColumnChunkMetaDataBuilder::new(column_descr)
1175 }
1176
1177 pub fn file_path(&self) -> Option<&str> {
1182 self.file_path.as_deref()
1183 }
1184
1185 pub fn file_offset(&self) -> i64 {
1192 self.file_offset
1193 }
1194
1195 pub fn column_type(&self) -> Type {
1197 self.column_descr.physical_type()
1198 }
1199
1200 pub fn column_path(&self) -> &ColumnPath {
1202 self.column_descr.path()
1203 }
1204
1205 pub fn column_descr(&self) -> &ColumnDescriptor {
1207 self.column_descr.as_ref()
1208 }
1209
1210 pub fn column_descr_ptr(&self) -> ColumnDescPtr {
1212 self.column_descr.clone()
1213 }
1214
1215 pub fn encodings(&self) -> impl Iterator<Item = Encoding> {
1217 self.encodings.encodings()
1218 }
1219
1220 pub fn encodings_mask(&self) -> &EncodingMask {
1222 &self.encodings
1223 }
1224
1225 pub fn num_values(&self) -> i64 {
1227 self.num_values
1228 }
1229
1230 pub fn compression(&self) -> Compression {
1237 self.compression.into()
1238 }
1239
1240 pub fn compression_codec(&self) -> CompressionCodec {
1242 self.compression
1243 }
1244
1245 pub fn compressed_size(&self) -> i64 {
1247 self.total_compressed_size
1248 }
1249
1250 pub fn uncompressed_size(&self) -> i64 {
1252 self.total_uncompressed_size
1253 }
1254
1255 pub fn data_page_offset(&self) -> i64 {
1257 self.data_page_offset
1258 }
1259
1260 pub fn index_page_offset(&self) -> Option<i64> {
1262 self.index_page_offset
1263 }
1264
1265 pub fn dictionary_page_offset(&self) -> Option<i64> {
1267 self.dictionary_page_offset
1268 }
1269
1270 pub fn byte_range(&self) -> (u64, u64) {
1276 let col_start = match self.dictionary_page_offset() {
1277 Some(dictionary_page_offset) => dictionary_page_offset,
1278 None => self.data_page_offset(),
1279 };
1280 let col_len = self.compressed_size();
1281 assert!(
1282 col_start >= 0 && col_len >= 0,
1283 "column start and length should not be negative"
1284 );
1285 (col_start as u64, col_len as u64)
1286 }
1287
1288 pub fn statistics(&self) -> Option<&Statistics> {
1291 self.statistics.as_ref()
1292 }
1293
1294 pub fn geo_statistics(&self) -> Option<&geo_statistics::GeospatialStatistics> {
1297 self.geo_statistics.as_deref()
1298 }
1299
1300 pub fn page_encoding_stats(&self) -> Option<&Vec<PageEncodingStats>> {
1307 match self.encoding_stats.as_ref() {
1308 Some(ParquetPageEncodingStats::Full(stats)) => Some(stats),
1309 _ => None,
1310 }
1311 }
1312
1313 pub fn page_encoding_stats_mask(&self) -> Option<&EncodingMask> {
1343 match self.encoding_stats.as_ref() {
1344 Some(ParquetPageEncodingStats::Mask(stats)) => Some(stats),
1345 _ => None,
1346 }
1347 }
1348
1349 pub fn bloom_filter_offset(&self) -> Option<i64> {
1351 self.bloom_filter_offset
1352 }
1353
1354 pub fn bloom_filter_length(&self) -> Option<i32> {
1356 self.bloom_filter_length
1357 }
1358
1359 pub fn column_index_offset(&self) -> Option<i64> {
1361 self.column_index_offset
1362 }
1363
1364 pub fn column_index_length(&self) -> Option<i32> {
1366 self.column_index_length
1367 }
1368
1369 pub(crate) fn column_index_range(&self) -> Option<Range<u64>> {
1371 let offset = u64::try_from(self.column_index_offset?).ok()?;
1372 let length = u64::try_from(self.column_index_length?).ok()?;
1373 Some(offset..(offset + length))
1374 }
1375
1376 pub fn offset_index_offset(&self) -> Option<i64> {
1378 self.offset_index_offset
1379 }
1380
1381 pub fn offset_index_length(&self) -> Option<i32> {
1383 self.offset_index_length
1384 }
1385
1386 pub(crate) fn offset_index_range(&self) -> Option<Range<u64>> {
1388 let offset = u64::try_from(self.offset_index_offset?).ok()?;
1389 let length = u64::try_from(self.offset_index_length?).ok()?;
1390 Some(offset..(offset + length))
1391 }
1392
1393 pub fn unencoded_byte_array_data_bytes(&self) -> Option<i64> {
1398 self.unencoded_byte_array_data_bytes
1399 }
1400
1401 pub fn repetition_level_histogram(&self) -> Option<&LevelHistogram> {
1407 self.repetition_level_histogram.as_ref()
1408 }
1409
1410 pub fn definition_level_histogram(&self) -> Option<&LevelHistogram> {
1416 self.definition_level_histogram.as_ref()
1417 }
1418
1419 #[cfg(feature = "encryption")]
1421 pub fn crypto_metadata(&self) -> Option<&ColumnCryptoMetaData> {
1422 self.column_crypto_metadata.as_deref()
1423 }
1424
1425 pub fn into_builder(self) -> ColumnChunkMetaDataBuilder {
1427 ColumnChunkMetaDataBuilder::from(self)
1428 }
1429}
1430
1431pub struct ColumnChunkMetaDataBuilder(ColumnChunkMetaData);
1450
1451impl ColumnChunkMetaDataBuilder {
1452 fn new(column_descr: ColumnDescPtr) -> Self {
1456 Self(ColumnChunkMetaData {
1457 column_descr,
1458 encodings: Default::default(),
1459 file_path: None,
1460 file_offset: 0,
1461 num_values: 0,
1462 compression: CompressionCodec::UNCOMPRESSED,
1463 total_compressed_size: 0,
1464 total_uncompressed_size: 0,
1465 data_page_offset: 0,
1466 index_page_offset: None,
1467 dictionary_page_offset: None,
1468 statistics: None,
1469 geo_statistics: None,
1470 encoding_stats: None,
1471 bloom_filter_offset: None,
1472 bloom_filter_length: None,
1473 offset_index_offset: None,
1474 offset_index_length: None,
1475 column_index_offset: None,
1476 column_index_length: None,
1477 unencoded_byte_array_data_bytes: None,
1478 repetition_level_histogram: None,
1479 definition_level_histogram: None,
1480 #[cfg(feature = "encryption")]
1481 column_crypto_metadata: None,
1482 #[cfg(feature = "encryption")]
1483 encrypted_column_metadata: None,
1484 #[cfg(feature = "encryption")]
1485 plaintext_footer_mode: false,
1486 })
1487 }
1488
1489 pub fn set_encodings(mut self, encodings: Vec<Encoding>) -> Self {
1491 self.0.encodings = EncodingMask::new_from_encodings(encodings.iter());
1492 self
1493 }
1494
1495 pub fn set_encodings_mask(mut self, encodings: EncodingMask) -> Self {
1497 self.0.encodings = encodings;
1498 self
1499 }
1500
1501 pub fn set_file_path(mut self, value: String) -> Self {
1503 self.0.file_path = Some(value);
1504 self
1505 }
1506
1507 pub fn set_num_values(mut self, value: i64) -> Self {
1509 self.0.num_values = value;
1510 self
1511 }
1512
1513 pub fn set_compression(mut self, value: Compression) -> Self {
1515 self.0.compression = value.into();
1516 self
1517 }
1518
1519 pub fn set_compression_codec(mut self, value: CompressionCodec) -> Self {
1521 self.0.compression = value;
1522 self
1523 }
1524
1525 pub fn set_total_compressed_size(mut self, value: i64) -> Self {
1527 self.0.total_compressed_size = value;
1528 self
1529 }
1530
1531 pub fn set_total_uncompressed_size(mut self, value: i64) -> Self {
1533 self.0.total_uncompressed_size = value;
1534 self
1535 }
1536
1537 pub fn set_data_page_offset(mut self, value: i64) -> Self {
1539 self.0.data_page_offset = value;
1540 self
1541 }
1542
1543 pub fn set_dictionary_page_offset(mut self, value: Option<i64>) -> Self {
1545 self.0.dictionary_page_offset = value;
1546 self
1547 }
1548
1549 pub fn set_index_page_offset(mut self, value: Option<i64>) -> Self {
1551 self.0.index_page_offset = value;
1552 self
1553 }
1554
1555 pub fn set_statistics(mut self, value: Statistics) -> Self {
1557 self.0.statistics = Some(value);
1558 self
1559 }
1560
1561 pub fn set_geo_statistics(mut self, value: Box<geo_statistics::GeospatialStatistics>) -> Self {
1563 self.0.geo_statistics = Some(value);
1564 self
1565 }
1566
1567 pub fn clear_statistics(mut self) -> Self {
1569 self.0.statistics = None;
1570 self
1571 }
1572
1573 pub fn set_page_encoding_stats(mut self, value: Vec<PageEncodingStats>) -> Self {
1577 self.0.encoding_stats = Some(ParquetPageEncodingStats::Full(value));
1578 self
1579 }
1580
1581 pub fn set_page_encoding_stats_mask(mut self, value: EncodingMask) -> Self {
1585 self.0.encoding_stats = Some(ParquetPageEncodingStats::Mask(value));
1586 self
1587 }
1588
1589 pub fn clear_page_encoding_stats(mut self) -> Self {
1591 self.0.encoding_stats = None;
1592 self
1593 }
1594
1595 pub fn set_bloom_filter_offset(mut self, value: Option<i64>) -> Self {
1597 self.0.bloom_filter_offset = value;
1598 self
1599 }
1600
1601 pub fn set_bloom_filter_length(mut self, value: Option<i32>) -> Self {
1603 self.0.bloom_filter_length = value;
1604 self
1605 }
1606
1607 pub fn set_offset_index_offset(mut self, value: Option<i64>) -> Self {
1609 self.0.offset_index_offset = value;
1610 self
1611 }
1612
1613 pub fn set_offset_index_length(mut self, value: Option<i32>) -> Self {
1615 self.0.offset_index_length = value;
1616 self
1617 }
1618
1619 pub fn set_column_index_offset(mut self, value: Option<i64>) -> Self {
1621 self.0.column_index_offset = value;
1622 self
1623 }
1624
1625 pub fn set_column_index_length(mut self, value: Option<i32>) -> Self {
1627 self.0.column_index_length = value;
1628 self
1629 }
1630
1631 pub fn set_unencoded_byte_array_data_bytes(mut self, value: Option<i64>) -> Self {
1633 self.0.unencoded_byte_array_data_bytes = value;
1634 self
1635 }
1636
1637 pub fn set_repetition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1639 self.0.repetition_level_histogram = value;
1640 self
1641 }
1642
1643 pub fn set_definition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1645 self.0.definition_level_histogram = value;
1646 self
1647 }
1648
1649 #[cfg(feature = "encryption")]
1650 pub fn set_column_crypto_metadata(mut self, value: Option<ColumnCryptoMetaData>) -> Self {
1652 self.0.column_crypto_metadata = value.map(Box::new);
1653 self
1654 }
1655
1656 #[cfg(feature = "encryption")]
1657 pub fn set_encrypted_column_metadata(mut self, value: Option<Vec<u8>>) -> Self {
1659 self.0.encrypted_column_metadata = value;
1660 self
1661 }
1662
1663 pub fn build(self) -> Result<ColumnChunkMetaData> {
1665 Ok(self.0)
1666 }
1667}
1668
1669pub struct ColumnIndexBuilder {
1674 column_type: Type,
1675 null_pages: Vec<bool>,
1676 min_values: Vec<Vec<u8>>,
1677 max_values: Vec<Vec<u8>>,
1678 null_counts: Vec<i64>,
1679 nan_counts: Vec<Option<i64>>,
1680 boundary_order: BoundaryOrder,
1681 repetition_level_histograms: Option<Vec<i64>>,
1683 definition_level_histograms: Option<Vec<i64>>,
1685 valid: bool,
1693}
1694
1695impl ColumnIndexBuilder {
1696 pub fn new(column_type: Type) -> Self {
1698 ColumnIndexBuilder {
1699 column_type,
1700 null_pages: Vec::new(),
1701 min_values: Vec::new(),
1702 max_values: Vec::new(),
1703 null_counts: Vec::new(),
1704 nan_counts: Vec::new(),
1705 boundary_order: BoundaryOrder::UNORDERED,
1706 repetition_level_histograms: None,
1707 definition_level_histograms: None,
1708 valid: true,
1709 }
1710 }
1711
1712 pub fn append(
1719 &mut self,
1720 null_page: bool,
1721 min_value: Vec<u8>,
1722 max_value: Vec<u8>,
1723 null_count: i64,
1724 nan_count: Option<i64>,
1725 ) {
1726 self.null_pages.push(null_page);
1727 self.min_values.push(min_value);
1728 self.max_values.push(max_value);
1729 self.null_counts.push(null_count);
1730 self.nan_counts.push(nan_count);
1731 }
1732
1733 pub fn append_histograms(
1738 &mut self,
1739 repetition_level_histogram: &Option<LevelHistogram>,
1740 definition_level_histogram: &Option<LevelHistogram>,
1741 ) {
1742 if !self.valid {
1743 return;
1744 }
1745 if let Some(rep_lvl_hist) = repetition_level_histogram {
1746 let hist = self.repetition_level_histograms.get_or_insert(Vec::new());
1747 hist.reserve(rep_lvl_hist.len());
1748 hist.extend(rep_lvl_hist.values());
1749 }
1750 if let Some(def_lvl_hist) = definition_level_histogram {
1751 let hist = self.definition_level_histograms.get_or_insert(Vec::new());
1752 hist.reserve(def_lvl_hist.len());
1753 hist.extend(def_lvl_hist.values());
1754 }
1755 }
1756
1757 pub fn set_boundary_order(&mut self, boundary_order: BoundaryOrder) {
1759 self.boundary_order = boundary_order;
1760 }
1761
1762 pub fn to_invalid(&mut self) {
1764 self.valid = false;
1765 }
1766
1767 pub fn valid(&self) -> bool {
1769 self.valid
1770 }
1771
1772 pub fn build(self) -> Result<ColumnIndexMetaData> {
1776 Ok(match self.column_type {
1777 Type::BOOLEAN => {
1778 let index = self.build_page_index(false)?;
1779 ColumnIndexMetaData::BOOLEAN(index)
1780 }
1781 Type::INT32 => {
1782 let index = self.build_page_index(false)?;
1783 ColumnIndexMetaData::INT32(index)
1784 }
1785 Type::INT64 => {
1786 let index = self.build_page_index(false)?;
1787 ColumnIndexMetaData::INT64(index)
1788 }
1789 Type::INT96 => {
1790 let index = self.build_page_index(false)?;
1791 ColumnIndexMetaData::INT96(index)
1792 }
1793 Type::FLOAT => {
1794 let index = self.build_page_index(true)?;
1795 ColumnIndexMetaData::FLOAT(index)
1796 }
1797 Type::DOUBLE => {
1798 let index = self.build_page_index(true)?;
1799 ColumnIndexMetaData::DOUBLE(index)
1800 }
1801 Type::BYTE_ARRAY => {
1802 let index = self.build_byte_array_index(false)?;
1803 ColumnIndexMetaData::BYTE_ARRAY(index)
1804 }
1805 Type::FIXED_LEN_BYTE_ARRAY => {
1806 let index = self.build_byte_array_index(true)?;
1807 ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index)
1808 }
1809 })
1810 }
1811
1812 fn build_nan_counts(nan_counts: &[Option<i64>]) -> Option<Vec<i64>> {
1813 let has_some = nan_counts.iter().any(|x| x.is_some());
1814 let has_none = nan_counts.iter().any(|x| x.is_none());
1815
1816 if has_some && !has_none {
1817 Some(nan_counts.iter().map(|x| x.unwrap()).collect())
1818 } else if !has_some && has_none {
1819 None
1820 } else {
1821 debug_assert!(
1822 false,
1823 "Mixed Some/None in nan_counts - caller should provide consistent values"
1824 );
1825 Some(nan_counts.iter().map(|x| x.unwrap_or(0)).collect())
1826 }
1827 }
1828
1829 fn build_page_index<T>(self, may_have_nan: bool) -> Result<PrimitiveColumnIndex<T>>
1830 where
1831 T: ParquetValueType,
1832 {
1833 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1834 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1835
1836 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1841 Self::build_nan_counts(&self.nan_counts)
1842 } else {
1843 None
1844 };
1845
1846 PrimitiveColumnIndex::try_new(
1847 self.null_pages,
1848 self.boundary_order,
1849 Some(self.null_counts),
1850 nan_counts,
1851 self.repetition_level_histograms,
1852 self.definition_level_histograms,
1853 min_values,
1854 max_values,
1855 )
1856 }
1857
1858 fn build_byte_array_index(self, may_have_nan: bool) -> Result<ByteArrayColumnIndex> {
1859 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1860 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1861
1862 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1867 Self::build_nan_counts(&self.nan_counts)
1868 } else {
1869 None
1870 };
1871
1872 ByteArrayColumnIndex::try_new(
1873 self.null_pages,
1874 self.boundary_order,
1875 Some(self.null_counts),
1876 nan_counts,
1877 self.repetition_level_histograms,
1878 self.definition_level_histograms,
1879 min_values,
1880 max_values,
1881 )
1882 }
1883}
1884
1885impl From<ColumnChunkMetaData> for ColumnChunkMetaDataBuilder {
1886 fn from(value: ColumnChunkMetaData) -> Self {
1887 ColumnChunkMetaDataBuilder(value)
1888 }
1889}
1890
1891pub struct OffsetIndexBuilder {
1895 offset_array: Vec<i64>,
1896 compressed_page_size_array: Vec<i32>,
1897 first_row_index_array: Vec<i64>,
1898 unencoded_byte_array_data_bytes_array: Option<Vec<i64>>,
1899 current_first_row_index: i64,
1900}
1901
1902impl Default for OffsetIndexBuilder {
1903 fn default() -> Self {
1904 Self::new()
1905 }
1906}
1907
1908impl OffsetIndexBuilder {
1909 pub fn new() -> Self {
1911 OffsetIndexBuilder {
1912 offset_array: Vec::new(),
1913 compressed_page_size_array: Vec::new(),
1914 first_row_index_array: Vec::new(),
1915 unencoded_byte_array_data_bytes_array: None,
1916 current_first_row_index: 0,
1917 }
1918 }
1919
1920 pub fn append_row_count(&mut self, row_count: i64) {
1922 let current_page_row_index = self.current_first_row_index;
1923 self.first_row_index_array.push(current_page_row_index);
1924 self.current_first_row_index += row_count;
1925 }
1926
1927 pub fn append_offset_and_size(&mut self, offset: i64, compressed_page_size: i32) {
1929 self.offset_array.push(offset);
1930 self.compressed_page_size_array.push(compressed_page_size);
1931 }
1932
1933 pub fn append_unencoded_byte_array_data_bytes(
1935 &mut self,
1936 unencoded_byte_array_data_bytes: Option<i64>,
1937 ) {
1938 if let Some(val) = unencoded_byte_array_data_bytes {
1939 self.unencoded_byte_array_data_bytes_array
1940 .get_or_insert(Vec::new())
1941 .push(val);
1942 }
1943 }
1944
1945 pub fn build(self) -> OffsetIndexMetaData {
1947 let locations = self
1948 .offset_array
1949 .iter()
1950 .zip(self.compressed_page_size_array.iter())
1951 .zip(self.first_row_index_array.iter())
1952 .map(|((offset, size), row_index)| PageLocation {
1953 offset: *offset,
1954 compressed_page_size: *size,
1955 first_row_index: *row_index,
1956 })
1957 .collect::<Vec<_>>();
1958 OffsetIndexMetaData {
1959 page_locations: locations,
1960 unencoded_byte_array_data_bytes: self.unencoded_byte_array_data_bytes_array,
1961 }
1962 }
1963}
1964
1965#[cfg(test)]
1966mod tests {
1967 use super::*;
1968 use crate::basic::{PageType, SortOrder};
1969 use crate::file::metadata::thrift::tests::{
1970 read_column_chunk, read_column_chunk_with_options, read_row_group,
1971 };
1972
1973 #[test]
1974 #[expect(deprecated)]
1975 fn test_level_histogram_update_from_levels_compat() {
1976 let mut histogram = LevelHistogram::try_new(2).unwrap();
1977 histogram.update_from_levels(&[0, 2, 1, 2, 2]);
1978 assert_eq!(histogram.values(), &[1, 1, 3]);
1979 }
1980
1981 #[test]
1982 fn test_row_group_metadata_thrift_conversion() {
1983 let schema_descr = get_test_schema_descr();
1984
1985 let mut columns = vec![];
1986 for ptr in schema_descr.columns() {
1987 let column = ColumnChunkMetaData::builder(ptr.clone()).build().unwrap();
1988 columns.push(column);
1989 }
1990 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
1991 .set_num_rows(1000)
1992 .set_total_byte_size(2000)
1993 .set_column_metadata(columns)
1994 .set_ordinal(1)
1995 .build()
1996 .unwrap();
1997
1998 let mut buf = Vec::new();
1999 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2000 row_group_meta.write_thrift(&mut writer).unwrap();
2001
2002 let row_group_res = read_row_group(&buf, schema_descr).unwrap();
2003
2004 assert_eq!(row_group_res, row_group_meta);
2005 }
2006
2007 #[test]
2008 fn test_row_group_metadata_thrift_conversion_empty() {
2009 let schema_descr = get_test_schema_descr();
2010
2011 let row_group_meta = RowGroupMetaData::builder(schema_descr).build();
2012
2013 assert!(row_group_meta.is_err());
2014 if let Err(e) = row_group_meta {
2015 assert_eq!(
2016 format!("{e}"),
2017 "Parquet error: Column length mismatch: 2 != 0"
2018 );
2019 }
2020 }
2021
2022 #[test]
2024 fn test_row_group_metadata_thrift_corrupted() {
2025 let schema_descr_2cols = Arc::new(SchemaDescriptor::new(Arc::new(
2026 SchemaType::group_type_builder("schema")
2027 .with_fields(vec![
2028 Arc::new(
2029 SchemaType::primitive_type_builder("a", Type::INT32)
2030 .build()
2031 .unwrap(),
2032 ),
2033 Arc::new(
2034 SchemaType::primitive_type_builder("b", Type::INT32)
2035 .build()
2036 .unwrap(),
2037 ),
2038 ])
2039 .build()
2040 .unwrap(),
2041 )));
2042
2043 let schema_descr_3cols = Arc::new(SchemaDescriptor::new(Arc::new(
2044 SchemaType::group_type_builder("schema")
2045 .with_fields(vec![
2046 Arc::new(
2047 SchemaType::primitive_type_builder("a", Type::INT32)
2048 .build()
2049 .unwrap(),
2050 ),
2051 Arc::new(
2052 SchemaType::primitive_type_builder("b", Type::INT32)
2053 .build()
2054 .unwrap(),
2055 ),
2056 Arc::new(
2057 SchemaType::primitive_type_builder("c", Type::INT32)
2058 .build()
2059 .unwrap(),
2060 ),
2061 ])
2062 .build()
2063 .unwrap(),
2064 )));
2065
2066 let row_group_meta_2cols = RowGroupMetaData::builder(schema_descr_2cols.clone())
2067 .set_num_rows(1000)
2068 .set_total_byte_size(2000)
2069 .set_column_metadata(vec![
2070 ColumnChunkMetaData::builder(schema_descr_2cols.column(0))
2071 .build()
2072 .unwrap(),
2073 ColumnChunkMetaData::builder(schema_descr_2cols.column(1))
2074 .build()
2075 .unwrap(),
2076 ])
2077 .set_ordinal(1)
2078 .build()
2079 .unwrap();
2080 let mut buf = Vec::new();
2081 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2082 row_group_meta_2cols.write_thrift(&mut writer).unwrap();
2083
2084 let err = read_row_group(&buf, schema_descr_3cols)
2085 .unwrap_err()
2086 .to_string();
2087 assert_eq!(
2088 err,
2089 "Parquet error: Column count mismatch. Schema has 3 columns while Row Group has 2"
2090 );
2091 }
2092
2093 #[test]
2094 fn test_column_chunk_metadata_thrift_conversion() {
2095 let column_descr = get_test_schema_descr().column(0);
2096 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
2097 .set_encodings_mask(EncodingMask::new_from_encodings(
2098 [Encoding::PLAIN, Encoding::RLE].iter(),
2099 ))
2100 .set_file_path("file_path".to_owned())
2101 .set_num_values(1000)
2102 .set_compression_codec(CompressionCodec::SNAPPY)
2103 .set_total_compressed_size(2000)
2104 .set_total_uncompressed_size(3000)
2105 .set_data_page_offset(4000)
2106 .set_dictionary_page_offset(Some(5000))
2107 .set_page_encoding_stats(vec![
2108 PageEncodingStats {
2109 page_type: PageType::DATA_PAGE,
2110 encoding: Encoding::PLAIN,
2111 count: 3,
2112 },
2113 PageEncodingStats {
2114 page_type: PageType::DATA_PAGE,
2115 encoding: Encoding::RLE,
2116 count: 5,
2117 },
2118 ])
2119 .set_bloom_filter_offset(Some(6000))
2120 .set_bloom_filter_length(Some(25))
2121 .set_offset_index_offset(Some(7000))
2122 .set_offset_index_length(Some(25))
2123 .set_column_index_offset(Some(8000))
2124 .set_column_index_length(Some(25))
2125 .set_unencoded_byte_array_data_bytes(Some(2000))
2126 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
2127 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
2128 .build()
2129 .unwrap();
2130
2131 let mut buf = Vec::new();
2132 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2133 col_metadata.write_thrift(&mut writer).unwrap();
2134 let col_chunk_res = read_column_chunk(&buf, column_descr.clone()).unwrap();
2135
2136 let expected_metadata = ColumnChunkMetaData::builder(column_descr)
2137 .set_encodings_mask(EncodingMask::new_from_encodings(
2138 [Encoding::PLAIN, Encoding::RLE].iter(),
2139 ))
2140 .set_file_path("file_path".to_owned())
2141 .set_num_values(1000)
2142 .set_compression_codec(CompressionCodec::SNAPPY)
2143 .set_total_compressed_size(2000)
2144 .set_total_uncompressed_size(3000)
2145 .set_data_page_offset(4000)
2146 .set_dictionary_page_offset(Some(5000))
2147 .set_page_encoding_stats_mask(EncodingMask::new_from_encodings(
2148 [Encoding::PLAIN, Encoding::RLE].iter(),
2149 ))
2150 .set_bloom_filter_offset(Some(6000))
2151 .set_bloom_filter_length(Some(25))
2152 .set_offset_index_offset(Some(7000))
2153 .set_offset_index_length(Some(25))
2154 .set_column_index_offset(Some(8000))
2155 .set_column_index_length(Some(25))
2156 .set_unencoded_byte_array_data_bytes(Some(2000))
2157 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
2158 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
2159 .build()
2160 .unwrap();
2161
2162 assert_eq!(col_chunk_res, expected_metadata);
2163 }
2164
2165 #[test]
2166 fn test_column_chunk_metadata_thrift_conversion_full_stats() {
2167 let column_descr = get_test_schema_descr().column(0);
2168 let stats = vec![
2169 PageEncodingStats {
2170 page_type: PageType::DATA_PAGE,
2171 encoding: Encoding::PLAIN,
2172 count: 3,
2173 },
2174 PageEncodingStats {
2175 page_type: PageType::DATA_PAGE,
2176 encoding: Encoding::RLE,
2177 count: 5,
2178 },
2179 ];
2180 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
2181 .set_encodings_mask(EncodingMask::new_from_encodings(
2182 [Encoding::PLAIN, Encoding::RLE].iter(),
2183 ))
2184 .set_num_values(1000)
2185 .set_compression_codec(CompressionCodec::SNAPPY)
2186 .set_total_compressed_size(2000)
2187 .set_total_uncompressed_size(3000)
2188 .set_data_page_offset(4000)
2189 .set_page_encoding_stats(stats)
2190 .build()
2191 .unwrap();
2192
2193 let mut buf = Vec::new();
2194 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2195 col_metadata.write_thrift(&mut writer).unwrap();
2196
2197 let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
2198 let col_chunk_res =
2199 read_column_chunk_with_options(&buf, column_descr, Some(&options)).unwrap();
2200
2201 assert_eq!(col_chunk_res, col_metadata);
2202 }
2203
2204 #[test]
2205 fn test_column_chunk_metadata_thrift_conversion_empty() {
2206 let column_descr = get_test_schema_descr().column(0);
2207
2208 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
2209 .build()
2210 .unwrap();
2211
2212 let mut buf = Vec::new();
2213 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2214 col_metadata.write_thrift(&mut writer).unwrap();
2215 let col_chunk_res = read_column_chunk(&buf, column_descr).unwrap();
2216
2217 assert_eq!(col_chunk_res, col_metadata);
2218 }
2219
2220 #[test]
2221 fn test_compressed_size() {
2222 let schema_descr = get_test_schema_descr();
2223
2224 let mut columns = vec![];
2225 for column_descr in schema_descr.columns() {
2226 let column = ColumnChunkMetaData::builder(column_descr.clone())
2227 .set_total_compressed_size(500)
2228 .set_total_uncompressed_size(700)
2229 .build()
2230 .unwrap();
2231 columns.push(column);
2232 }
2233 let row_group_meta = RowGroupMetaData::builder(schema_descr)
2234 .set_num_rows(1000)
2235 .set_column_metadata(columns)
2236 .build()
2237 .unwrap();
2238
2239 let compressed_size_res = row_group_meta.compressed_size();
2240 let compressed_size_exp: i64 = 1000;
2241
2242 assert_eq!(compressed_size_res, compressed_size_exp);
2243 }
2244
2245 #[test]
2246 fn test_memory_size() {
2247 let schema_descr = get_test_schema_descr();
2248
2249 let columns = schema_descr
2250 .columns()
2251 .iter()
2252 .map(|column_descr| {
2253 ColumnChunkMetaData::builder(column_descr.clone())
2254 .set_statistics(Statistics::new::<i32>(None, None, None, None, false))
2255 .build()
2256 })
2257 .collect::<Result<Vec<_>>>()
2258 .unwrap();
2259 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2260 .set_num_rows(1000)
2261 .set_column_metadata(columns)
2262 .build()
2263 .unwrap();
2264 let row_group_meta = vec![row_group_meta];
2265
2266 let version = 2;
2267 let num_rows = 1000;
2268 let created_by = Some(String::from("test harness"));
2269 let key_value_metadata = Some(vec![KeyValue::new(
2270 String::from("Foo"),
2271 Some(String::from("bar")),
2272 )]);
2273 let column_orders = Some(vec![
2274 ColumnOrder::UNDEFINED,
2275 ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED),
2276 ]);
2277 let file_metadata = FileMetaData::new(
2278 version,
2279 num_rows,
2280 created_by,
2281 key_value_metadata,
2282 schema_descr.clone(),
2283 column_orders,
2284 );
2285
2286 let columns_with_stats = schema_descr
2288 .columns()
2289 .iter()
2290 .map(|column_descr| {
2291 ColumnChunkMetaData::builder(column_descr.clone())
2292 .set_statistics(Statistics::new::<i32>(
2293 Some(0),
2294 Some(100),
2295 None,
2296 None,
2297 false,
2298 ))
2299 .build()
2300 })
2301 .collect::<Result<Vec<_>>>()
2302 .unwrap();
2303
2304 let row_group_meta_with_stats = RowGroupMetaData::builder(schema_descr)
2305 .set_num_rows(1000)
2306 .set_column_metadata(columns_with_stats)
2307 .build()
2308 .unwrap();
2309 let row_group_meta_with_stats = vec![row_group_meta_with_stats];
2310
2311 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata.clone())
2312 .set_row_groups(row_group_meta_with_stats)
2313 .build();
2314
2315 #[cfg(not(feature = "encryption"))]
2316 let base_expected_size = 2798;
2317 #[cfg(feature = "encryption")]
2318 let base_expected_size = 2966;
2319
2320 assert_eq!(parquet_meta.memory_size(), base_expected_size);
2321
2322 let mut column_index = ColumnIndexBuilder::new(Type::BOOLEAN);
2323 column_index.append(false, vec![1u8], vec![2u8, 3u8], 4, None);
2324 let column_index = column_index.build().unwrap();
2325 let ColumnIndexMetaData::BOOLEAN(native_index) = column_index else {
2326 panic!("wrong type of column index")
2327 };
2328
2329 let mut offset_index = OffsetIndexBuilder::new();
2331 offset_index.append_row_count(1);
2332 offset_index.append_offset_and_size(2, 3);
2333 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2334 offset_index.append_row_count(1);
2335 offset_index.append_offset_and_size(2, 3);
2336 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2337 let offset_index = Some(offset_index.build());
2338
2339 let page_index = PageIndex::new(
2340 Some(vec![vec![Some(ColumnIndexMetaData::BOOLEAN(native_index))]]),
2341 Some(vec![vec![offset_index]]),
2342 );
2343
2344 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata)
2345 .set_row_groups(row_group_meta)
2346 .set_page_index(Some(page_index))
2347 .build();
2348
2349 #[cfg(not(feature = "encryption"))]
2350 let bigger_expected_size = 3248;
2351 #[cfg(feature = "encryption")]
2352 let bigger_expected_size = 3416;
2353
2354 assert!(bigger_expected_size > base_expected_size);
2356 assert_eq!(parquet_meta.memory_size(), bigger_expected_size);
2357 }
2358
2359 #[test]
2360 #[cfg(feature = "encryption")]
2361 fn test_memory_size_with_decryptor() {
2362 use crate::encryption::decrypt::FileDecryptionProperties;
2363 use crate::file::metadata::thrift::encryption::AesGcmV1;
2364
2365 let schema_descr = get_test_schema_descr();
2366
2367 let columns = schema_descr
2368 .columns()
2369 .iter()
2370 .map(|column_descr| ColumnChunkMetaData::builder(column_descr.clone()).build())
2371 .collect::<Result<Vec<_>>>()
2372 .unwrap();
2373 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2374 .set_num_rows(1000)
2375 .set_column_metadata(columns)
2376 .build()
2377 .unwrap();
2378 let row_group_meta = vec![row_group_meta];
2379
2380 let version = 2;
2381 let num_rows = 1000;
2382 let aad_file_unique = vec![1u8; 8];
2383 let aad_prefix = vec![2u8; 8];
2384 let encryption_algorithm = EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 {
2385 aad_prefix: Some(aad_prefix.clone()),
2386 aad_file_unique: Some(aad_file_unique.clone()),
2387 supply_aad_prefix: Some(true),
2388 });
2389 let footer_key_metadata = Some(vec![3u8; 8]);
2390 let file_metadata =
2391 FileMetaData::new(version, num_rows, None, None, schema_descr.clone(), None)
2392 .with_encryption_algorithm(Some(encryption_algorithm))
2393 .with_footer_signing_key_metadata(footer_key_metadata.clone());
2394
2395 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2396 .set_row_groups(row_group_meta.clone())
2397 .build();
2398
2399 let base_expected_size = 2074;
2400 assert_eq!(parquet_meta_data.memory_size(), base_expected_size);
2401
2402 let footer_key = b"0123456789012345";
2403 let column_key = b"1234567890123450";
2404 let mut decryption_properties_builder =
2405 FileDecryptionProperties::builder(footer_key.to_vec())
2406 .with_aad_prefix(aad_prefix.clone());
2407 for column in schema_descr.columns() {
2408 decryption_properties_builder = decryption_properties_builder
2409 .with_column_key(&column.path().string(), column_key.to_vec());
2410 }
2411 let decryption_properties = decryption_properties_builder.build().unwrap();
2412 let decryptor = FileDecryptor::new(
2413 &decryption_properties,
2414 footer_key_metadata.as_deref(),
2415 aad_file_unique,
2416 aad_prefix,
2417 )
2418 .unwrap();
2419
2420 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2421 .set_row_groups(row_group_meta.clone())
2422 .set_file_decryptor(Some(decryptor))
2423 .build();
2424
2425 let expected_size_with_decryptor = 3088;
2426 assert!(expected_size_with_decryptor > base_expected_size);
2427
2428 assert_eq!(
2429 parquet_meta_data.memory_size(),
2430 expected_size_with_decryptor
2431 );
2432 }
2433
2434 fn get_test_schema_descr() -> SchemaDescPtr {
2436 let schema = SchemaType::group_type_builder("schema")
2437 .with_fields(vec![
2438 Arc::new(
2439 SchemaType::primitive_type_builder("a", Type::INT32)
2440 .build()
2441 .unwrap(),
2442 ),
2443 Arc::new(
2444 SchemaType::primitive_type_builder("b", Type::INT32)
2445 .build()
2446 .unwrap(),
2447 ),
2448 ])
2449 .build()
2450 .unwrap();
2451
2452 Arc::new(SchemaDescriptor::new(Arc::new(schema)))
2453 }
2454}