1mod footer_tail;
53mod memory;
54mod options;
55mod parser;
56mod push_decoder;
57pub(crate) mod reader;
58pub(crate) mod thrift;
59mod writer;
60
61use crate::basic::{
62 BoundaryOrder, ColumnOrder, Compression, CompressionCodec, Encoding, EncodingMask, PageType,
63 Type,
64};
65#[cfg(feature = "encryption")]
66use crate::encryption::decrypt::FileDecryptor;
67use crate::errors::{ParquetError, Result};
68#[cfg(feature = "encryption")]
69use crate::file::column_crypto_metadata::ColumnCryptoMetaData;
70pub(crate) use crate::file::metadata::memory::HeapSize;
71#[cfg(feature = "encryption")]
72use crate::file::metadata::thrift::encryption::EncryptionAlgorithm;
73use crate::file::page_index::column_index::{ByteArrayColumnIndex, PrimitiveColumnIndex};
74use crate::file::page_index::{column_index::ColumnIndexMetaData, offset_index::PageLocation};
75use crate::file::statistics::Statistics;
76use crate::geospatial::statistics as geo_statistics;
77use crate::parquet_thrift::{
78 ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
79 WriteThrift, WriteThriftField,
80};
81use crate::schema::types::{
82 ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, SchemaDescriptor,
83 Type as SchemaType,
84};
85use crate::thrift_struct;
86use crate::{
87 data_type::private::ParquetValueType, file::page_index::offset_index::OffsetIndexMetaData,
88};
89
90pub use footer_tail::FooterTail;
91pub use options::{ParquetMetaDataOptions, ParquetStatisticsPolicy};
92pub use push_decoder::ParquetMetaDataPushDecoder;
93pub use reader::{PageIndexPolicy, ParquetMetaDataReader};
94use std::io::Write;
95use std::ops::Range;
96use std::sync::Arc;
97pub use writer::ParquetMetaDataWriter;
98pub(crate) use writer::ThriftMetadataWriter;
99
100#[derive(Debug, Clone, PartialEq)]
213pub struct PageIndex {
214 column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>,
215 offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>,
216}
217
218impl PageIndex {
219 pub(crate) fn new(
220 column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>,
221 offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>,
222 ) -> Self {
223 Self {
224 column_indexes,
225 offset_indexes,
226 }
227 }
228
229 pub fn has_offset_indexes(&self) -> bool {
236 self.offset_indexes.is_some()
237 }
238
239 pub fn has_column_indexes(&self) -> bool {
246 self.column_indexes.is_some()
247 }
248
249 pub fn is_complete(&self) -> bool {
254 self.has_column_indexes() && self.has_offset_indexes()
255 }
256
257 pub fn column_indexes_for_rowgroup(
267 &self,
268 row_group_idx: usize,
269 ) -> Option<&[Option<ColumnIndexMetaData>]> {
270 match self.column_indexes.as_ref() {
271 None => None,
272 Some(indexes) => indexes.get(row_group_idx).map(|ci| ci.as_slice()),
273 }
274 }
275
276 pub fn column_index(
285 &self,
286 row_group_idx: usize,
287 column_idx: usize,
288 ) -> Option<&ColumnIndexMetaData> {
289 if let Some(column_indexes) = self.column_indexes.as_ref() {
290 let rg = column_indexes.get(row_group_idx)?;
291 rg.get(column_idx)?.as_ref()
292 } else {
293 None
294 }
295 }
296
297 pub fn offset_indexes_for_rowgroup(
307 &self,
308 row_group_idx: usize,
309 ) -> Option<&[Option<OffsetIndexMetaData>]> {
310 match self.offset_indexes.as_ref() {
311 None => None,
312 Some(indexes) => indexes.get(row_group_idx).map(|oi| oi.as_slice()),
313 }
314 }
315
316 pub fn offset_index(
327 &self,
328 row_group_idx: usize,
329 column_idx: usize,
330 ) -> Option<&OffsetIndexMetaData> {
331 if let Some(offset_indexes) = self.offset_indexes.as_ref() {
332 let rg = offset_indexes.get(row_group_idx)?;
333 rg.get(column_idx)?.as_ref()
334 } else {
335 None
336 }
337 }
338
339 pub fn num_data_pages(&self, row_group_idx: usize, column_idx: usize) -> Option<usize> {
347 match self.offset_index(row_group_idx, column_idx) {
348 Some(offset_index) => Some(offset_index.page_locations.len()),
349 None => Some(self.column_index(row_group_idx, column_idx)?.num_pages() as usize),
350 }
351 }
352
353 pub fn page_locations(
366 &self,
367 row_group_idx: usize,
368 column_idx: usize,
369 ) -> Option<&Vec<PageLocation>> {
370 if let Some(offset_indexes) = self.offset_indexes.as_ref() {
371 let rg = offset_indexes.get(row_group_idx)?;
372 let off_idx = rg.get(column_idx)?.as_ref()?;
373 Some(off_idx.page_locations())
374 } else {
375 None
376 }
377 }
378}
379
380#[derive(Debug, Clone, PartialEq)]
398pub struct ParquetMetaData {
399 file_metadata: FileMetaData,
401 row_groups: Vec<RowGroupMetaData>,
403 page_index: Option<PageIndex>,
405 #[cfg(feature = "encryption")]
407 file_decryptor: Option<Box<FileDecryptor>>,
408}
409
410impl ParquetMetaData {
411 pub fn new(file_metadata: FileMetaData, row_groups: Vec<RowGroupMetaData>) -> Self {
414 ParquetMetaData {
415 file_metadata,
416 row_groups,
417 page_index: None,
418 #[cfg(feature = "encryption")]
419 file_decryptor: None,
420 }
421 }
422
423 #[cfg(feature = "encryption")]
426 pub(crate) fn with_file_decryptor(&mut self, file_decryptor: Option<FileDecryptor>) {
427 self.file_decryptor = file_decryptor.map(Box::new);
428 }
429
430 pub fn into_builder(self) -> ParquetMetaDataBuilder {
432 self.into()
433 }
434
435 pub fn file_metadata(&self) -> &FileMetaData {
437 &self.file_metadata
438 }
439
440 #[cfg(feature = "encryption")]
442 pub(crate) fn file_decryptor(&self) -> Option<&FileDecryptor> {
443 self.file_decryptor.as_deref()
444 }
445
446 pub fn num_row_groups(&self) -> usize {
448 self.row_groups.len()
449 }
450
451 pub fn row_group(&self, i: usize) -> &RowGroupMetaData {
454 &self.row_groups[i]
455 }
456
457 pub fn row_groups(&self) -> &[RowGroupMetaData] {
459 &self.row_groups
460 }
461
462 pub fn row_group_num_rows(&self, row_group_idx: usize) -> Result<usize> {
467 self.row_groups
468 .get(row_group_idx)
469 .ok_or_else(|| {
470 ParquetError::General(format!(
471 "Row group index {row_group_idx} out of bounds for file with {} row groups",
472 self.num_row_groups()
473 ))
474 })?
475 .num_rows()
476 .try_into()
477 .map_err(|e| ParquetError::General(format!("Row count overflow: {e}")))
478 }
479
480 pub fn page_index(&self) -> Option<&PageIndex> {
487 self.page_index.as_ref()
488 }
489
490 pub fn memory_size(&self) -> usize {
505 #[cfg(feature = "encryption")]
506 let encryption_size = self.file_decryptor.heap_size();
507 #[cfg(not(feature = "encryption"))]
508 let encryption_size = 0usize;
509
510 std::mem::size_of::<Self>()
511 + self.file_metadata.heap_size()
512 + self.row_groups.heap_size()
513 + self.page_index.heap_size()
514 + encryption_size
515 }
516
517 pub(crate) fn set_page_index(&mut self, index: Option<PageIndex>) {
519 self.page_index = index;
520 }
521}
522
523pub struct ParquetMetaDataBuilder(ParquetMetaData);
561
562impl ParquetMetaDataBuilder {
563 pub fn new(file_meta_data: FileMetaData) -> Self {
565 Self(ParquetMetaData::new(file_meta_data, vec![]))
566 }
567
568 pub fn new_from_metadata(metadata: ParquetMetaData) -> Self {
570 Self(metadata)
571 }
572
573 pub fn add_row_group(mut self, row_group: RowGroupMetaData) -> Self {
575 self.0.row_groups.push(row_group);
576 self
577 }
578
579 pub fn set_row_groups(mut self, row_groups: Vec<RowGroupMetaData>) -> Self {
581 self.0.row_groups = row_groups;
582 self
583 }
584
585 pub fn take_row_groups(&mut self) -> Vec<RowGroupMetaData> {
591 std::mem::take(&mut self.0.row_groups)
592 }
593
594 pub fn row_groups(&self) -> &[RowGroupMetaData] {
596 &self.0.row_groups
597 }
598
599 pub fn set_page_index(mut self, page_index: Option<PageIndex>) -> Self {
601 self.0.page_index = page_index;
602 self
603 }
604
605 pub fn take_page_index(&mut self) -> Option<PageIndex> {
607 std::mem::take(&mut self.0.page_index)
608 }
609
610 pub fn page_index(&self) -> Option<&PageIndex> {
612 self.0.page_index.as_ref()
613 }
614
615 #[cfg(feature = "encryption")]
617 pub(crate) fn set_file_decryptor(mut self, file_decryptor: Option<FileDecryptor>) -> Self {
618 self.0.with_file_decryptor(file_decryptor);
619 self
620 }
621
622 pub fn build(self) -> ParquetMetaData {
624 let Self(metadata) = self;
625 metadata
626 }
627}
628
629impl From<ParquetMetaData> for ParquetMetaDataBuilder {
630 fn from(meta_data: ParquetMetaData) -> Self {
631 Self(meta_data)
632 }
633}
634
635thrift_struct!(
636pub struct KeyValue {
638 1: required string key
639 2: optional string value
640}
641);
642
643impl KeyValue {
644 pub fn new<F2>(key: String, value: F2) -> KeyValue
646 where
647 F2: Into<Option<String>>,
648 {
649 KeyValue {
650 key,
651 value: value.into(),
652 }
653 }
654}
655
656thrift_struct!(
657pub struct PageEncodingStats {
659 1: required PageType page_type;
660 2: required Encoding encoding;
661 3: required i32 count;
662}
663);
664
665#[derive(Debug, Clone, PartialEq)]
668enum ParquetPageEncodingStats {
669 Full(Vec<PageEncodingStats>),
671 Mask(EncodingMask),
673}
674
675pub type FileMetaDataPtr = Arc<FileMetaData>;
677
678#[derive(Debug, Clone, PartialEq)]
682pub struct FileMetaData {
683 version: i32,
684 num_rows: i64,
685 created_by: Option<String>,
686 key_value_metadata: Option<Vec<KeyValue>>,
687 schema_descr: SchemaDescPtr,
688 column_orders: Option<Vec<ColumnOrder>>,
689 #[cfg(feature = "encryption")]
690 encryption_algorithm: Option<Box<EncryptionAlgorithm>>,
691 #[cfg(feature = "encryption")]
692 footer_signing_key_metadata: Option<Vec<u8>>,
693}
694
695impl FileMetaData {
696 pub fn new(
698 version: i32,
699 num_rows: i64,
700 created_by: Option<String>,
701 key_value_metadata: Option<Vec<KeyValue>>,
702 schema_descr: SchemaDescPtr,
703 column_orders: Option<Vec<ColumnOrder>>,
704 ) -> Self {
705 FileMetaData {
706 version,
707 num_rows,
708 created_by,
709 key_value_metadata,
710 schema_descr,
711 column_orders,
712 #[cfg(feature = "encryption")]
713 encryption_algorithm: None,
714 #[cfg(feature = "encryption")]
715 footer_signing_key_metadata: None,
716 }
717 }
718
719 #[cfg(feature = "encryption")]
720 pub(crate) fn with_encryption_algorithm(
721 mut self,
722 encryption_algorithm: Option<EncryptionAlgorithm>,
723 ) -> Self {
724 self.encryption_algorithm = encryption_algorithm.map(Box::new);
725 self
726 }
727
728 #[cfg(feature = "encryption")]
729 pub(crate) fn with_footer_signing_key_metadata(
730 mut self,
731 footer_signing_key_metadata: Option<Vec<u8>>,
732 ) -> Self {
733 self.footer_signing_key_metadata = footer_signing_key_metadata;
734 self
735 }
736
737 pub fn version(&self) -> i32 {
739 self.version
740 }
741
742 pub fn num_rows(&self) -> i64 {
744 self.num_rows
745 }
746
747 pub fn created_by(&self) -> Option<&str> {
756 self.created_by.as_deref()
757 }
758
759 pub fn key_value_metadata(&self) -> Option<&Vec<KeyValue>> {
761 self.key_value_metadata.as_ref()
762 }
763
764 pub fn schema(&self) -> &SchemaType {
768 self.schema_descr.root_schema()
769 }
770
771 pub fn schema_descr(&self) -> &SchemaDescriptor {
773 &self.schema_descr
774 }
775
776 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
778 self.schema_descr.clone()
779 }
780
781 pub fn column_orders(&self) -> Option<&Vec<ColumnOrder>> {
789 self.column_orders.as_ref()
790 }
791
792 pub fn column_order(&self, i: usize) -> ColumnOrder {
795 self.column_orders
796 .as_ref()
797 .map(|data| data[i])
798 .unwrap_or(ColumnOrder::UNDEFINED)
799 }
800}
801
802thrift_struct!(
803pub struct SortingColumn {
805 1: required i32 column_idx
807
808 2: required bool descending
810
811 3: required bool nulls_first
814}
815);
816
817pub type RowGroupMetaDataPtr = Arc<RowGroupMetaData>;
819
820#[derive(Debug, Clone, PartialEq)]
825pub struct RowGroupMetaData {
826 columns: Vec<ColumnChunkMetaData>,
827 num_rows: i64,
828 sorting_columns: Option<Vec<SortingColumn>>,
829 total_byte_size: i64,
830 schema_descr: SchemaDescPtr,
831 file_offset: Option<i64>,
833 ordinal: Option<i32>,
835}
836
837impl RowGroupMetaData {
838 pub fn builder(schema_descr: SchemaDescPtr) -> RowGroupMetaDataBuilder {
840 RowGroupMetaDataBuilder::new(schema_descr)
841 }
842
843 pub fn num_columns(&self) -> usize {
845 self.columns.len()
846 }
847
848 pub fn column(&self, i: usize) -> &ColumnChunkMetaData {
850 &self.columns[i]
851 }
852
853 pub fn columns(&self) -> &[ColumnChunkMetaData] {
855 &self.columns
856 }
857
858 pub fn columns_mut(&mut self) -> &mut [ColumnChunkMetaData] {
860 &mut self.columns
861 }
862
863 pub fn num_rows(&self) -> i64 {
865 self.num_rows
866 }
867
868 pub fn sorting_columns(&self) -> Option<&Vec<SortingColumn>> {
870 self.sorting_columns.as_ref()
871 }
872
873 pub fn total_byte_size(&self) -> i64 {
875 self.total_byte_size
876 }
877
878 pub fn compressed_size(&self) -> i64 {
880 self.columns.iter().map(|c| c.total_compressed_size).sum()
881 }
882
883 pub fn schema_descr(&self) -> &SchemaDescriptor {
885 self.schema_descr.as_ref()
886 }
887
888 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
890 self.schema_descr.clone()
891 }
892
893 #[inline(always)]
898 pub fn ordinal(&self) -> Option<i32> {
899 self.ordinal
900 }
901
902 #[inline(always)]
904 pub fn file_offset(&self) -> Option<i64> {
905 self.file_offset
906 }
907
908 pub fn into_builder(self) -> RowGroupMetaDataBuilder {
910 RowGroupMetaDataBuilder(self)
911 }
912}
913
914pub struct RowGroupMetaDataBuilder(RowGroupMetaData);
916
917impl RowGroupMetaDataBuilder {
918 fn new(schema_descr: SchemaDescPtr) -> Self {
920 Self(RowGroupMetaData {
921 columns: Vec::with_capacity(schema_descr.num_columns()),
922 schema_descr,
923 file_offset: None,
924 num_rows: 0,
925 sorting_columns: None,
926 total_byte_size: 0,
927 ordinal: None,
928 })
929 }
930
931 pub fn set_num_rows(mut self, value: i64) -> Self {
933 self.0.num_rows = value;
934 self
935 }
936
937 pub fn set_sorting_columns(mut self, value: Option<Vec<SortingColumn>>) -> Self {
939 self.0.sorting_columns = value;
940 self
941 }
942
943 pub fn set_total_byte_size(mut self, value: i64) -> Self {
945 self.0.total_byte_size = value;
946 self
947 }
948
949 pub fn take_columns(&mut self) -> Vec<ColumnChunkMetaData> {
955 std::mem::take(&mut self.0.columns)
956 }
957
958 pub fn set_column_metadata(mut self, value: Vec<ColumnChunkMetaData>) -> Self {
960 self.0.columns = value;
961 self
962 }
963
964 pub fn add_column_metadata(mut self, value: ColumnChunkMetaData) -> Self {
966 self.0.columns.push(value);
967 self
968 }
969
970 pub fn set_ordinal(mut self, value: i32) -> Self {
972 self.0.ordinal = Some(value);
973 self
974 }
975
976 pub fn set_file_offset(mut self, value: i64) -> Self {
978 self.0.file_offset = Some(value);
979 self
980 }
981
982 pub fn build(self) -> Result<RowGroupMetaData> {
984 if self.0.schema_descr.num_columns() != self.0.columns.len() {
985 return Err(general_err!(
986 "Column length mismatch: {} != {}",
987 self.0.schema_descr.num_columns(),
988 self.0.columns.len()
989 ));
990 }
991
992 Ok(self.0)
993 }
994
995 pub(super) fn build_unchecked(self) -> RowGroupMetaData {
997 self.0
998 }
999}
1000
1001#[derive(Debug, Clone, PartialEq)]
1003pub struct ColumnChunkMetaData {
1004 column_descr: ColumnDescPtr,
1005 encodings: EncodingMask,
1006 file_path: Option<String>,
1007 file_offset: i64,
1008 num_values: i64,
1009 compression: CompressionCodec,
1010 total_compressed_size: i64,
1011 total_uncompressed_size: i64,
1012 data_page_offset: i64,
1013 index_page_offset: Option<i64>,
1014 dictionary_page_offset: Option<i64>,
1015 statistics: Option<Statistics>,
1016 geo_statistics: Option<Box<geo_statistics::GeospatialStatistics>>,
1017 encoding_stats: Option<ParquetPageEncodingStats>,
1018 bloom_filter_offset: Option<i64>,
1019 bloom_filter_length: Option<i32>,
1020 offset_index_offset: Option<i64>,
1021 offset_index_length: Option<i32>,
1022 column_index_offset: Option<i64>,
1023 column_index_length: Option<i32>,
1024 unencoded_byte_array_data_bytes: Option<i64>,
1025 repetition_level_histogram: Option<LevelHistogram>,
1026 definition_level_histogram: Option<LevelHistogram>,
1027 #[cfg(feature = "encryption")]
1028 column_crypto_metadata: Option<Box<ColumnCryptoMetaData>>,
1029 #[cfg(feature = "encryption")]
1030 encrypted_column_metadata: Option<Vec<u8>>,
1031 #[cfg(feature = "encryption")]
1035 plaintext_footer_mode: bool,
1036}
1037
1038#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
1047pub struct LevelHistogram {
1048 inner: Vec<i64>,
1049}
1050
1051impl LevelHistogram {
1052 pub fn try_new(max_level: i16) -> Option<Self> {
1058 if max_level > 0 {
1059 Some(Self {
1060 inner: vec![0; max_level as usize + 1],
1061 })
1062 } else {
1063 None
1064 }
1065 }
1066 pub fn values(&self) -> &[i64] {
1068 &self.inner
1069 }
1070
1071 pub fn into_inner(self) -> Vec<i64> {
1073 self.inner
1074 }
1075
1076 pub fn get(&self, index: usize) -> Option<i64> {
1083 self.inner.get(index).copied()
1084 }
1085
1086 pub fn add(&mut self, other: &Self) {
1091 assert_eq!(self.len(), other.len());
1092 for (dst, src) in self.inner.iter_mut().zip(other.inner.iter()) {
1093 *dst += src;
1094 }
1095 }
1096
1097 pub fn len(&self) -> usize {
1099 self.inner.len()
1100 }
1101
1102 pub fn is_empty(&self) -> bool {
1104 self.inner.is_empty()
1105 }
1106
1107 pub fn reset(&mut self) {
1109 for value in &mut self.inner {
1110 *value = 0;
1111 }
1112 }
1113
1114 #[inline]
1116 pub fn increment_by(&mut self, level: i16, count: i64) {
1117 self.inner[level as usize] += count;
1118 }
1119
1120 #[deprecated(since = "58.2.0", note = "Use `increment_by` instead")]
1126 pub fn update_from_levels(&mut self, levels: &[i16]) {
1127 for &level in levels {
1128 self.increment_by(level, 1);
1129 }
1130 }
1131}
1132
1133impl From<Vec<i64>> for LevelHistogram {
1134 fn from(inner: Vec<i64>) -> Self {
1135 Self { inner }
1136 }
1137}
1138
1139impl From<LevelHistogram> for Vec<i64> {
1140 fn from(value: LevelHistogram) -> Self {
1141 value.into_inner()
1142 }
1143}
1144
1145impl HeapSize for LevelHistogram {
1146 fn heap_size(&self) -> usize {
1147 self.inner.heap_size()
1148 }
1149}
1150
1151impl ColumnChunkMetaData {
1153 pub fn builder(column_descr: ColumnDescPtr) -> ColumnChunkMetaDataBuilder {
1155 ColumnChunkMetaDataBuilder::new(column_descr)
1156 }
1157
1158 pub fn file_path(&self) -> Option<&str> {
1163 self.file_path.as_deref()
1164 }
1165
1166 pub fn file_offset(&self) -> i64 {
1173 self.file_offset
1174 }
1175
1176 pub fn column_type(&self) -> Type {
1178 self.column_descr.physical_type()
1179 }
1180
1181 pub fn column_path(&self) -> &ColumnPath {
1183 self.column_descr.path()
1184 }
1185
1186 pub fn column_descr(&self) -> &ColumnDescriptor {
1188 self.column_descr.as_ref()
1189 }
1190
1191 pub fn column_descr_ptr(&self) -> ColumnDescPtr {
1193 self.column_descr.clone()
1194 }
1195
1196 pub fn encodings(&self) -> impl Iterator<Item = Encoding> {
1198 self.encodings.encodings()
1199 }
1200
1201 pub fn encodings_mask(&self) -> &EncodingMask {
1203 &self.encodings
1204 }
1205
1206 pub fn num_values(&self) -> i64 {
1208 self.num_values
1209 }
1210
1211 pub fn compression(&self) -> Compression {
1218 self.compression.into()
1219 }
1220
1221 pub fn compression_codec(&self) -> CompressionCodec {
1223 self.compression
1224 }
1225
1226 pub fn compressed_size(&self) -> i64 {
1228 self.total_compressed_size
1229 }
1230
1231 pub fn uncompressed_size(&self) -> i64 {
1233 self.total_uncompressed_size
1234 }
1235
1236 pub fn data_page_offset(&self) -> i64 {
1238 self.data_page_offset
1239 }
1240
1241 pub fn index_page_offset(&self) -> Option<i64> {
1243 self.index_page_offset
1244 }
1245
1246 pub fn dictionary_page_offset(&self) -> Option<i64> {
1248 self.dictionary_page_offset
1249 }
1250
1251 pub fn byte_range(&self) -> (u64, u64) {
1257 let col_start = match self.dictionary_page_offset() {
1258 Some(dictionary_page_offset) => dictionary_page_offset,
1259 None => self.data_page_offset(),
1260 };
1261 let col_len = self.compressed_size();
1262 assert!(
1263 col_start >= 0 && col_len >= 0,
1264 "column start and length should not be negative"
1265 );
1266 (col_start as u64, col_len as u64)
1267 }
1268
1269 pub fn statistics(&self) -> Option<&Statistics> {
1272 self.statistics.as_ref()
1273 }
1274
1275 pub fn geo_statistics(&self) -> Option<&geo_statistics::GeospatialStatistics> {
1278 self.geo_statistics.as_deref()
1279 }
1280
1281 pub fn page_encoding_stats(&self) -> Option<&Vec<PageEncodingStats>> {
1288 match self.encoding_stats.as_ref() {
1289 Some(ParquetPageEncodingStats::Full(stats)) => Some(stats),
1290 _ => None,
1291 }
1292 }
1293
1294 pub fn page_encoding_stats_mask(&self) -> Option<&EncodingMask> {
1324 match self.encoding_stats.as_ref() {
1325 Some(ParquetPageEncodingStats::Mask(stats)) => Some(stats),
1326 _ => None,
1327 }
1328 }
1329
1330 pub fn bloom_filter_offset(&self) -> Option<i64> {
1332 self.bloom_filter_offset
1333 }
1334
1335 pub fn bloom_filter_length(&self) -> Option<i32> {
1337 self.bloom_filter_length
1338 }
1339
1340 pub fn column_index_offset(&self) -> Option<i64> {
1342 self.column_index_offset
1343 }
1344
1345 pub fn column_index_length(&self) -> Option<i32> {
1347 self.column_index_length
1348 }
1349
1350 pub(crate) fn column_index_range(&self) -> Option<Range<u64>> {
1352 let offset = u64::try_from(self.column_index_offset?).ok()?;
1353 let length = u64::try_from(self.column_index_length?).ok()?;
1354 Some(offset..(offset + length))
1355 }
1356
1357 pub fn offset_index_offset(&self) -> Option<i64> {
1359 self.offset_index_offset
1360 }
1361
1362 pub fn offset_index_length(&self) -> Option<i32> {
1364 self.offset_index_length
1365 }
1366
1367 pub(crate) fn offset_index_range(&self) -> Option<Range<u64>> {
1369 let offset = u64::try_from(self.offset_index_offset?).ok()?;
1370 let length = u64::try_from(self.offset_index_length?).ok()?;
1371 Some(offset..(offset + length))
1372 }
1373
1374 pub fn unencoded_byte_array_data_bytes(&self) -> Option<i64> {
1379 self.unencoded_byte_array_data_bytes
1380 }
1381
1382 pub fn repetition_level_histogram(&self) -> Option<&LevelHistogram> {
1388 self.repetition_level_histogram.as_ref()
1389 }
1390
1391 pub fn definition_level_histogram(&self) -> Option<&LevelHistogram> {
1397 self.definition_level_histogram.as_ref()
1398 }
1399
1400 #[cfg(feature = "encryption")]
1402 pub fn crypto_metadata(&self) -> Option<&ColumnCryptoMetaData> {
1403 self.column_crypto_metadata.as_deref()
1404 }
1405
1406 pub fn into_builder(self) -> ColumnChunkMetaDataBuilder {
1408 ColumnChunkMetaDataBuilder::from(self)
1409 }
1410}
1411
1412pub struct ColumnChunkMetaDataBuilder(ColumnChunkMetaData);
1431
1432impl ColumnChunkMetaDataBuilder {
1433 fn new(column_descr: ColumnDescPtr) -> Self {
1437 Self(ColumnChunkMetaData {
1438 column_descr,
1439 encodings: Default::default(),
1440 file_path: None,
1441 file_offset: 0,
1442 num_values: 0,
1443 compression: CompressionCodec::UNCOMPRESSED,
1444 total_compressed_size: 0,
1445 total_uncompressed_size: 0,
1446 data_page_offset: 0,
1447 index_page_offset: None,
1448 dictionary_page_offset: None,
1449 statistics: None,
1450 geo_statistics: None,
1451 encoding_stats: None,
1452 bloom_filter_offset: None,
1453 bloom_filter_length: None,
1454 offset_index_offset: None,
1455 offset_index_length: None,
1456 column_index_offset: None,
1457 column_index_length: None,
1458 unencoded_byte_array_data_bytes: None,
1459 repetition_level_histogram: None,
1460 definition_level_histogram: None,
1461 #[cfg(feature = "encryption")]
1462 column_crypto_metadata: None,
1463 #[cfg(feature = "encryption")]
1464 encrypted_column_metadata: None,
1465 #[cfg(feature = "encryption")]
1466 plaintext_footer_mode: false,
1467 })
1468 }
1469
1470 pub fn set_encodings(mut self, encodings: Vec<Encoding>) -> Self {
1472 self.0.encodings = EncodingMask::new_from_encodings(encodings.iter());
1473 self
1474 }
1475
1476 pub fn set_encodings_mask(mut self, encodings: EncodingMask) -> Self {
1478 self.0.encodings = encodings;
1479 self
1480 }
1481
1482 pub fn set_file_path(mut self, value: String) -> Self {
1484 self.0.file_path = Some(value);
1485 self
1486 }
1487
1488 pub fn set_num_values(mut self, value: i64) -> Self {
1490 self.0.num_values = value;
1491 self
1492 }
1493
1494 pub fn set_compression(mut self, value: Compression) -> Self {
1496 self.0.compression = value.into();
1497 self
1498 }
1499
1500 pub fn set_compression_codec(mut self, value: CompressionCodec) -> Self {
1502 self.0.compression = value;
1503 self
1504 }
1505
1506 pub fn set_total_compressed_size(mut self, value: i64) -> Self {
1508 self.0.total_compressed_size = value;
1509 self
1510 }
1511
1512 pub fn set_total_uncompressed_size(mut self, value: i64) -> Self {
1514 self.0.total_uncompressed_size = value;
1515 self
1516 }
1517
1518 pub fn set_data_page_offset(mut self, value: i64) -> Self {
1520 self.0.data_page_offset = value;
1521 self
1522 }
1523
1524 pub fn set_dictionary_page_offset(mut self, value: Option<i64>) -> Self {
1526 self.0.dictionary_page_offset = value;
1527 self
1528 }
1529
1530 pub fn set_index_page_offset(mut self, value: Option<i64>) -> Self {
1532 self.0.index_page_offset = value;
1533 self
1534 }
1535
1536 pub fn set_statistics(mut self, value: Statistics) -> Self {
1538 self.0.statistics = Some(value);
1539 self
1540 }
1541
1542 pub fn set_geo_statistics(mut self, value: Box<geo_statistics::GeospatialStatistics>) -> Self {
1544 self.0.geo_statistics = Some(value);
1545 self
1546 }
1547
1548 pub fn clear_statistics(mut self) -> Self {
1550 self.0.statistics = None;
1551 self
1552 }
1553
1554 pub fn set_page_encoding_stats(mut self, value: Vec<PageEncodingStats>) -> Self {
1558 self.0.encoding_stats = Some(ParquetPageEncodingStats::Full(value));
1559 self
1560 }
1561
1562 pub fn set_page_encoding_stats_mask(mut self, value: EncodingMask) -> Self {
1566 self.0.encoding_stats = Some(ParquetPageEncodingStats::Mask(value));
1567 self
1568 }
1569
1570 pub fn clear_page_encoding_stats(mut self) -> Self {
1572 self.0.encoding_stats = None;
1573 self
1574 }
1575
1576 pub fn set_bloom_filter_offset(mut self, value: Option<i64>) -> Self {
1578 self.0.bloom_filter_offset = value;
1579 self
1580 }
1581
1582 pub fn set_bloom_filter_length(mut self, value: Option<i32>) -> Self {
1584 self.0.bloom_filter_length = value;
1585 self
1586 }
1587
1588 pub fn set_offset_index_offset(mut self, value: Option<i64>) -> Self {
1590 self.0.offset_index_offset = value;
1591 self
1592 }
1593
1594 pub fn set_offset_index_length(mut self, value: Option<i32>) -> Self {
1596 self.0.offset_index_length = value;
1597 self
1598 }
1599
1600 pub fn set_column_index_offset(mut self, value: Option<i64>) -> Self {
1602 self.0.column_index_offset = value;
1603 self
1604 }
1605
1606 pub fn set_column_index_length(mut self, value: Option<i32>) -> Self {
1608 self.0.column_index_length = value;
1609 self
1610 }
1611
1612 pub fn set_unencoded_byte_array_data_bytes(mut self, value: Option<i64>) -> Self {
1614 self.0.unencoded_byte_array_data_bytes = value;
1615 self
1616 }
1617
1618 pub fn set_repetition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1620 self.0.repetition_level_histogram = value;
1621 self
1622 }
1623
1624 pub fn set_definition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1626 self.0.definition_level_histogram = value;
1627 self
1628 }
1629
1630 #[cfg(feature = "encryption")]
1631 pub fn set_column_crypto_metadata(mut self, value: Option<ColumnCryptoMetaData>) -> Self {
1633 self.0.column_crypto_metadata = value.map(Box::new);
1634 self
1635 }
1636
1637 #[cfg(feature = "encryption")]
1638 pub fn set_encrypted_column_metadata(mut self, value: Option<Vec<u8>>) -> Self {
1640 self.0.encrypted_column_metadata = value;
1641 self
1642 }
1643
1644 pub fn build(self) -> Result<ColumnChunkMetaData> {
1646 Ok(self.0)
1647 }
1648}
1649
1650pub struct ColumnIndexBuilder {
1655 column_type: Type,
1656 null_pages: Vec<bool>,
1657 min_values: Vec<Vec<u8>>,
1658 max_values: Vec<Vec<u8>>,
1659 null_counts: Vec<i64>,
1660 nan_counts: Vec<Option<i64>>,
1661 boundary_order: BoundaryOrder,
1662 repetition_level_histograms: Option<Vec<i64>>,
1664 definition_level_histograms: Option<Vec<i64>>,
1666 valid: bool,
1674}
1675
1676impl ColumnIndexBuilder {
1677 pub fn new(column_type: Type) -> Self {
1679 ColumnIndexBuilder {
1680 column_type,
1681 null_pages: Vec::new(),
1682 min_values: Vec::new(),
1683 max_values: Vec::new(),
1684 null_counts: Vec::new(),
1685 nan_counts: Vec::new(),
1686 boundary_order: BoundaryOrder::UNORDERED,
1687 repetition_level_histograms: None,
1688 definition_level_histograms: None,
1689 valid: true,
1690 }
1691 }
1692
1693 pub fn append(
1700 &mut self,
1701 null_page: bool,
1702 min_value: Vec<u8>,
1703 max_value: Vec<u8>,
1704 null_count: i64,
1705 nan_count: Option<i64>,
1706 ) {
1707 self.null_pages.push(null_page);
1708 self.min_values.push(min_value);
1709 self.max_values.push(max_value);
1710 self.null_counts.push(null_count);
1711 self.nan_counts.push(nan_count);
1712 }
1713
1714 pub fn append_histograms(
1719 &mut self,
1720 repetition_level_histogram: &Option<LevelHistogram>,
1721 definition_level_histogram: &Option<LevelHistogram>,
1722 ) {
1723 if !self.valid {
1724 return;
1725 }
1726 if let Some(rep_lvl_hist) = repetition_level_histogram {
1727 let hist = self.repetition_level_histograms.get_or_insert(Vec::new());
1728 hist.reserve(rep_lvl_hist.len());
1729 hist.extend(rep_lvl_hist.values());
1730 }
1731 if let Some(def_lvl_hist) = definition_level_histogram {
1732 let hist = self.definition_level_histograms.get_or_insert(Vec::new());
1733 hist.reserve(def_lvl_hist.len());
1734 hist.extend(def_lvl_hist.values());
1735 }
1736 }
1737
1738 pub fn set_boundary_order(&mut self, boundary_order: BoundaryOrder) {
1740 self.boundary_order = boundary_order;
1741 }
1742
1743 pub fn to_invalid(&mut self) {
1745 self.valid = false;
1746 }
1747
1748 pub fn valid(&self) -> bool {
1750 self.valid
1751 }
1752
1753 pub fn build(self) -> Result<ColumnIndexMetaData> {
1757 Ok(match self.column_type {
1758 Type::BOOLEAN => {
1759 let index = self.build_page_index(false)?;
1760 ColumnIndexMetaData::BOOLEAN(index)
1761 }
1762 Type::INT32 => {
1763 let index = self.build_page_index(false)?;
1764 ColumnIndexMetaData::INT32(index)
1765 }
1766 Type::INT64 => {
1767 let index = self.build_page_index(false)?;
1768 ColumnIndexMetaData::INT64(index)
1769 }
1770 Type::INT96 => {
1771 let index = self.build_page_index(false)?;
1772 ColumnIndexMetaData::INT96(index)
1773 }
1774 Type::FLOAT => {
1775 let index = self.build_page_index(true)?;
1776 ColumnIndexMetaData::FLOAT(index)
1777 }
1778 Type::DOUBLE => {
1779 let index = self.build_page_index(true)?;
1780 ColumnIndexMetaData::DOUBLE(index)
1781 }
1782 Type::BYTE_ARRAY => {
1783 let index = self.build_byte_array_index(false)?;
1784 ColumnIndexMetaData::BYTE_ARRAY(index)
1785 }
1786 Type::FIXED_LEN_BYTE_ARRAY => {
1787 let index = self.build_byte_array_index(true)?;
1788 ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index)
1789 }
1790 })
1791 }
1792
1793 fn build_nan_counts(nan_counts: &[Option<i64>]) -> Option<Vec<i64>> {
1794 let has_some = nan_counts.iter().any(|x| x.is_some());
1795 let has_none = nan_counts.iter().any(|x| x.is_none());
1796
1797 if has_some && !has_none {
1798 Some(nan_counts.iter().map(|x| x.unwrap()).collect())
1799 } else if !has_some && has_none {
1800 None
1801 } else {
1802 debug_assert!(
1803 false,
1804 "Mixed Some/None in nan_counts - caller should provide consistent values"
1805 );
1806 Some(nan_counts.iter().map(|x| x.unwrap_or(0)).collect())
1807 }
1808 }
1809
1810 fn build_page_index<T>(self, may_have_nan: bool) -> Result<PrimitiveColumnIndex<T>>
1811 where
1812 T: ParquetValueType,
1813 {
1814 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1815 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1816
1817 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1822 Self::build_nan_counts(&self.nan_counts)
1823 } else {
1824 None
1825 };
1826
1827 PrimitiveColumnIndex::try_new(
1828 self.null_pages,
1829 self.boundary_order,
1830 Some(self.null_counts),
1831 nan_counts,
1832 self.repetition_level_histograms,
1833 self.definition_level_histograms,
1834 min_values,
1835 max_values,
1836 )
1837 }
1838
1839 fn build_byte_array_index(self, may_have_nan: bool) -> Result<ByteArrayColumnIndex> {
1840 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1841 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1842
1843 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1848 Self::build_nan_counts(&self.nan_counts)
1849 } else {
1850 None
1851 };
1852
1853 ByteArrayColumnIndex::try_new(
1854 self.null_pages,
1855 self.boundary_order,
1856 Some(self.null_counts),
1857 nan_counts,
1858 self.repetition_level_histograms,
1859 self.definition_level_histograms,
1860 min_values,
1861 max_values,
1862 )
1863 }
1864}
1865
1866impl From<ColumnChunkMetaData> for ColumnChunkMetaDataBuilder {
1867 fn from(value: ColumnChunkMetaData) -> Self {
1868 ColumnChunkMetaDataBuilder(value)
1869 }
1870}
1871
1872pub struct OffsetIndexBuilder {
1876 offset_array: Vec<i64>,
1877 compressed_page_size_array: Vec<i32>,
1878 first_row_index_array: Vec<i64>,
1879 unencoded_byte_array_data_bytes_array: Option<Vec<i64>>,
1880 current_first_row_index: i64,
1881}
1882
1883impl Default for OffsetIndexBuilder {
1884 fn default() -> Self {
1885 Self::new()
1886 }
1887}
1888
1889impl OffsetIndexBuilder {
1890 pub fn new() -> Self {
1892 OffsetIndexBuilder {
1893 offset_array: Vec::new(),
1894 compressed_page_size_array: Vec::new(),
1895 first_row_index_array: Vec::new(),
1896 unencoded_byte_array_data_bytes_array: None,
1897 current_first_row_index: 0,
1898 }
1899 }
1900
1901 pub fn append_row_count(&mut self, row_count: i64) {
1903 let current_page_row_index = self.current_first_row_index;
1904 self.first_row_index_array.push(current_page_row_index);
1905 self.current_first_row_index += row_count;
1906 }
1907
1908 pub fn append_offset_and_size(&mut self, offset: i64, compressed_page_size: i32) {
1910 self.offset_array.push(offset);
1911 self.compressed_page_size_array.push(compressed_page_size);
1912 }
1913
1914 pub fn append_unencoded_byte_array_data_bytes(
1916 &mut self,
1917 unencoded_byte_array_data_bytes: Option<i64>,
1918 ) {
1919 if let Some(val) = unencoded_byte_array_data_bytes {
1920 self.unencoded_byte_array_data_bytes_array
1921 .get_or_insert(Vec::new())
1922 .push(val);
1923 }
1924 }
1925
1926 pub fn build(self) -> OffsetIndexMetaData {
1928 let locations = self
1929 .offset_array
1930 .iter()
1931 .zip(self.compressed_page_size_array.iter())
1932 .zip(self.first_row_index_array.iter())
1933 .map(|((offset, size), row_index)| PageLocation {
1934 offset: *offset,
1935 compressed_page_size: *size,
1936 first_row_index: *row_index,
1937 })
1938 .collect::<Vec<_>>();
1939 OffsetIndexMetaData {
1940 page_locations: locations,
1941 unencoded_byte_array_data_bytes: self.unencoded_byte_array_data_bytes_array,
1942 }
1943 }
1944}
1945
1946#[cfg(test)]
1947mod tests {
1948 use super::*;
1949 use crate::basic::{PageType, SortOrder};
1950 use crate::file::metadata::thrift::tests::{
1951 read_column_chunk, read_column_chunk_with_options, read_row_group,
1952 };
1953
1954 #[test]
1955 #[expect(deprecated)]
1956 fn test_level_histogram_update_from_levels_compat() {
1957 let mut histogram = LevelHistogram::try_new(2).unwrap();
1958 histogram.update_from_levels(&[0, 2, 1, 2, 2]);
1959 assert_eq!(histogram.values(), &[1, 1, 3]);
1960 }
1961
1962 #[test]
1963 fn test_row_group_metadata_thrift_conversion() {
1964 let schema_descr = get_test_schema_descr();
1965
1966 let mut columns = vec![];
1967 for ptr in schema_descr.columns() {
1968 let column = ColumnChunkMetaData::builder(ptr.clone()).build().unwrap();
1969 columns.push(column);
1970 }
1971 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
1972 .set_num_rows(1000)
1973 .set_total_byte_size(2000)
1974 .set_column_metadata(columns)
1975 .set_ordinal(1)
1976 .build()
1977 .unwrap();
1978
1979 let mut buf = Vec::new();
1980 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1981 row_group_meta.write_thrift(&mut writer).unwrap();
1982
1983 let row_group_res = read_row_group(&buf, schema_descr).unwrap();
1984
1985 assert_eq!(row_group_res, row_group_meta);
1986 }
1987
1988 #[test]
1989 fn test_row_group_metadata_thrift_conversion_empty() {
1990 let schema_descr = get_test_schema_descr();
1991
1992 let row_group_meta = RowGroupMetaData::builder(schema_descr).build();
1993
1994 assert!(row_group_meta.is_err());
1995 if let Err(e) = row_group_meta {
1996 assert_eq!(
1997 format!("{e}"),
1998 "Parquet error: Column length mismatch: 2 != 0"
1999 );
2000 }
2001 }
2002
2003 #[test]
2005 fn test_row_group_metadata_thrift_corrupted() {
2006 let schema_descr_2cols = Arc::new(SchemaDescriptor::new(Arc::new(
2007 SchemaType::group_type_builder("schema")
2008 .with_fields(vec![
2009 Arc::new(
2010 SchemaType::primitive_type_builder("a", Type::INT32)
2011 .build()
2012 .unwrap(),
2013 ),
2014 Arc::new(
2015 SchemaType::primitive_type_builder("b", Type::INT32)
2016 .build()
2017 .unwrap(),
2018 ),
2019 ])
2020 .build()
2021 .unwrap(),
2022 )));
2023
2024 let schema_descr_3cols = Arc::new(SchemaDescriptor::new(Arc::new(
2025 SchemaType::group_type_builder("schema")
2026 .with_fields(vec![
2027 Arc::new(
2028 SchemaType::primitive_type_builder("a", Type::INT32)
2029 .build()
2030 .unwrap(),
2031 ),
2032 Arc::new(
2033 SchemaType::primitive_type_builder("b", Type::INT32)
2034 .build()
2035 .unwrap(),
2036 ),
2037 Arc::new(
2038 SchemaType::primitive_type_builder("c", Type::INT32)
2039 .build()
2040 .unwrap(),
2041 ),
2042 ])
2043 .build()
2044 .unwrap(),
2045 )));
2046
2047 let row_group_meta_2cols = RowGroupMetaData::builder(schema_descr_2cols.clone())
2048 .set_num_rows(1000)
2049 .set_total_byte_size(2000)
2050 .set_column_metadata(vec![
2051 ColumnChunkMetaData::builder(schema_descr_2cols.column(0))
2052 .build()
2053 .unwrap(),
2054 ColumnChunkMetaData::builder(schema_descr_2cols.column(1))
2055 .build()
2056 .unwrap(),
2057 ])
2058 .set_ordinal(1)
2059 .build()
2060 .unwrap();
2061 let mut buf = Vec::new();
2062 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2063 row_group_meta_2cols.write_thrift(&mut writer).unwrap();
2064
2065 let err = read_row_group(&buf, schema_descr_3cols)
2066 .unwrap_err()
2067 .to_string();
2068 assert_eq!(
2069 err,
2070 "Parquet error: Column count mismatch. Schema has 3 columns while Row Group has 2"
2071 );
2072 }
2073
2074 #[test]
2075 fn test_column_chunk_metadata_thrift_conversion() {
2076 let column_descr = get_test_schema_descr().column(0);
2077 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
2078 .set_encodings_mask(EncodingMask::new_from_encodings(
2079 [Encoding::PLAIN, Encoding::RLE].iter(),
2080 ))
2081 .set_file_path("file_path".to_owned())
2082 .set_num_values(1000)
2083 .set_compression_codec(CompressionCodec::SNAPPY)
2084 .set_total_compressed_size(2000)
2085 .set_total_uncompressed_size(3000)
2086 .set_data_page_offset(4000)
2087 .set_dictionary_page_offset(Some(5000))
2088 .set_page_encoding_stats(vec![
2089 PageEncodingStats {
2090 page_type: PageType::DATA_PAGE,
2091 encoding: Encoding::PLAIN,
2092 count: 3,
2093 },
2094 PageEncodingStats {
2095 page_type: PageType::DATA_PAGE,
2096 encoding: Encoding::RLE,
2097 count: 5,
2098 },
2099 ])
2100 .set_bloom_filter_offset(Some(6000))
2101 .set_bloom_filter_length(Some(25))
2102 .set_offset_index_offset(Some(7000))
2103 .set_offset_index_length(Some(25))
2104 .set_column_index_offset(Some(8000))
2105 .set_column_index_length(Some(25))
2106 .set_unencoded_byte_array_data_bytes(Some(2000))
2107 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
2108 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
2109 .build()
2110 .unwrap();
2111
2112 let mut buf = Vec::new();
2113 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2114 col_metadata.write_thrift(&mut writer).unwrap();
2115 let col_chunk_res = read_column_chunk(&buf, column_descr.clone()).unwrap();
2116
2117 let expected_metadata = ColumnChunkMetaData::builder(column_descr)
2118 .set_encodings_mask(EncodingMask::new_from_encodings(
2119 [Encoding::PLAIN, Encoding::RLE].iter(),
2120 ))
2121 .set_file_path("file_path".to_owned())
2122 .set_num_values(1000)
2123 .set_compression_codec(CompressionCodec::SNAPPY)
2124 .set_total_compressed_size(2000)
2125 .set_total_uncompressed_size(3000)
2126 .set_data_page_offset(4000)
2127 .set_dictionary_page_offset(Some(5000))
2128 .set_page_encoding_stats_mask(EncodingMask::new_from_encodings(
2129 [Encoding::PLAIN, Encoding::RLE].iter(),
2130 ))
2131 .set_bloom_filter_offset(Some(6000))
2132 .set_bloom_filter_length(Some(25))
2133 .set_offset_index_offset(Some(7000))
2134 .set_offset_index_length(Some(25))
2135 .set_column_index_offset(Some(8000))
2136 .set_column_index_length(Some(25))
2137 .set_unencoded_byte_array_data_bytes(Some(2000))
2138 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
2139 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
2140 .build()
2141 .unwrap();
2142
2143 assert_eq!(col_chunk_res, expected_metadata);
2144 }
2145
2146 #[test]
2147 fn test_column_chunk_metadata_thrift_conversion_full_stats() {
2148 let column_descr = get_test_schema_descr().column(0);
2149 let stats = vec![
2150 PageEncodingStats {
2151 page_type: PageType::DATA_PAGE,
2152 encoding: Encoding::PLAIN,
2153 count: 3,
2154 },
2155 PageEncodingStats {
2156 page_type: PageType::DATA_PAGE,
2157 encoding: Encoding::RLE,
2158 count: 5,
2159 },
2160 ];
2161 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
2162 .set_encodings_mask(EncodingMask::new_from_encodings(
2163 [Encoding::PLAIN, Encoding::RLE].iter(),
2164 ))
2165 .set_num_values(1000)
2166 .set_compression_codec(CompressionCodec::SNAPPY)
2167 .set_total_compressed_size(2000)
2168 .set_total_uncompressed_size(3000)
2169 .set_data_page_offset(4000)
2170 .set_page_encoding_stats(stats)
2171 .build()
2172 .unwrap();
2173
2174 let mut buf = Vec::new();
2175 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2176 col_metadata.write_thrift(&mut writer).unwrap();
2177
2178 let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
2179 let col_chunk_res =
2180 read_column_chunk_with_options(&buf, column_descr, Some(&options)).unwrap();
2181
2182 assert_eq!(col_chunk_res, col_metadata);
2183 }
2184
2185 #[test]
2186 fn test_column_chunk_metadata_thrift_conversion_empty() {
2187 let column_descr = get_test_schema_descr().column(0);
2188
2189 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
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 let col_chunk_res = read_column_chunk(&buf, column_descr).unwrap();
2197
2198 assert_eq!(col_chunk_res, col_metadata);
2199 }
2200
2201 #[test]
2202 fn test_compressed_size() {
2203 let schema_descr = get_test_schema_descr();
2204
2205 let mut columns = vec![];
2206 for column_descr in schema_descr.columns() {
2207 let column = ColumnChunkMetaData::builder(column_descr.clone())
2208 .set_total_compressed_size(500)
2209 .set_total_uncompressed_size(700)
2210 .build()
2211 .unwrap();
2212 columns.push(column);
2213 }
2214 let row_group_meta = RowGroupMetaData::builder(schema_descr)
2215 .set_num_rows(1000)
2216 .set_column_metadata(columns)
2217 .build()
2218 .unwrap();
2219
2220 let compressed_size_res = row_group_meta.compressed_size();
2221 let compressed_size_exp: i64 = 1000;
2222
2223 assert_eq!(compressed_size_res, compressed_size_exp);
2224 }
2225
2226 #[test]
2227 fn test_memory_size() {
2228 let schema_descr = get_test_schema_descr();
2229
2230 let columns = schema_descr
2231 .columns()
2232 .iter()
2233 .map(|column_descr| {
2234 ColumnChunkMetaData::builder(column_descr.clone())
2235 .set_statistics(Statistics::new::<i32>(None, None, None, None, false))
2236 .build()
2237 })
2238 .collect::<Result<Vec<_>>>()
2239 .unwrap();
2240 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2241 .set_num_rows(1000)
2242 .set_column_metadata(columns)
2243 .build()
2244 .unwrap();
2245 let row_group_meta = vec![row_group_meta];
2246
2247 let version = 2;
2248 let num_rows = 1000;
2249 let created_by = Some(String::from("test harness"));
2250 let key_value_metadata = Some(vec![KeyValue::new(
2251 String::from("Foo"),
2252 Some(String::from("bar")),
2253 )]);
2254 let column_orders = Some(vec![
2255 ColumnOrder::UNDEFINED,
2256 ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED),
2257 ]);
2258 let file_metadata = FileMetaData::new(
2259 version,
2260 num_rows,
2261 created_by,
2262 key_value_metadata,
2263 schema_descr.clone(),
2264 column_orders,
2265 );
2266
2267 let columns_with_stats = schema_descr
2269 .columns()
2270 .iter()
2271 .map(|column_descr| {
2272 ColumnChunkMetaData::builder(column_descr.clone())
2273 .set_statistics(Statistics::new::<i32>(
2274 Some(0),
2275 Some(100),
2276 None,
2277 None,
2278 false,
2279 ))
2280 .build()
2281 })
2282 .collect::<Result<Vec<_>>>()
2283 .unwrap();
2284
2285 let row_group_meta_with_stats = RowGroupMetaData::builder(schema_descr)
2286 .set_num_rows(1000)
2287 .set_column_metadata(columns_with_stats)
2288 .build()
2289 .unwrap();
2290 let row_group_meta_with_stats = vec![row_group_meta_with_stats];
2291
2292 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata.clone())
2293 .set_row_groups(row_group_meta_with_stats)
2294 .build();
2295
2296 #[cfg(not(feature = "encryption"))]
2297 let base_expected_size = 2798;
2298 #[cfg(feature = "encryption")]
2299 let base_expected_size = 2966;
2300
2301 assert_eq!(parquet_meta.memory_size(), base_expected_size);
2302
2303 let mut column_index = ColumnIndexBuilder::new(Type::BOOLEAN);
2304 column_index.append(false, vec![1u8], vec![2u8, 3u8], 4, None);
2305 let column_index = column_index.build().unwrap();
2306 let ColumnIndexMetaData::BOOLEAN(native_index) = column_index else {
2307 panic!("wrong type of column index")
2308 };
2309
2310 let mut offset_index = OffsetIndexBuilder::new();
2312 offset_index.append_row_count(1);
2313 offset_index.append_offset_and_size(2, 3);
2314 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2315 offset_index.append_row_count(1);
2316 offset_index.append_offset_and_size(2, 3);
2317 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2318 let offset_index = Some(offset_index.build());
2319
2320 let page_index = PageIndex::new(
2321 Some(vec![vec![Some(ColumnIndexMetaData::BOOLEAN(native_index))]]),
2322 Some(vec![vec![offset_index]]),
2323 );
2324
2325 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata)
2326 .set_row_groups(row_group_meta)
2327 .set_page_index(Some(page_index))
2328 .build();
2329
2330 #[cfg(not(feature = "encryption"))]
2331 let bigger_expected_size = 3248;
2332 #[cfg(feature = "encryption")]
2333 let bigger_expected_size = 3416;
2334
2335 assert!(bigger_expected_size > base_expected_size);
2337 assert_eq!(parquet_meta.memory_size(), bigger_expected_size);
2338 }
2339
2340 #[test]
2341 #[cfg(feature = "encryption")]
2342 fn test_memory_size_with_decryptor() {
2343 use crate::encryption::decrypt::FileDecryptionProperties;
2344 use crate::file::metadata::thrift::encryption::AesGcmV1;
2345
2346 let schema_descr = get_test_schema_descr();
2347
2348 let columns = schema_descr
2349 .columns()
2350 .iter()
2351 .map(|column_descr| ColumnChunkMetaData::builder(column_descr.clone()).build())
2352 .collect::<Result<Vec<_>>>()
2353 .unwrap();
2354 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2355 .set_num_rows(1000)
2356 .set_column_metadata(columns)
2357 .build()
2358 .unwrap();
2359 let row_group_meta = vec![row_group_meta];
2360
2361 let version = 2;
2362 let num_rows = 1000;
2363 let aad_file_unique = vec![1u8; 8];
2364 let aad_prefix = vec![2u8; 8];
2365 let encryption_algorithm = EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 {
2366 aad_prefix: Some(aad_prefix.clone()),
2367 aad_file_unique: Some(aad_file_unique.clone()),
2368 supply_aad_prefix: Some(true),
2369 });
2370 let footer_key_metadata = Some(vec![3u8; 8]);
2371 let file_metadata =
2372 FileMetaData::new(version, num_rows, None, None, schema_descr.clone(), None)
2373 .with_encryption_algorithm(Some(encryption_algorithm))
2374 .with_footer_signing_key_metadata(footer_key_metadata.clone());
2375
2376 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2377 .set_row_groups(row_group_meta.clone())
2378 .build();
2379
2380 let base_expected_size = 2074;
2381 assert_eq!(parquet_meta_data.memory_size(), base_expected_size);
2382
2383 let footer_key = b"0123456789012345";
2384 let column_key = b"1234567890123450";
2385 let mut decryption_properties_builder =
2386 FileDecryptionProperties::builder(footer_key.to_vec())
2387 .with_aad_prefix(aad_prefix.clone());
2388 for column in schema_descr.columns() {
2389 decryption_properties_builder = decryption_properties_builder
2390 .with_column_key(&column.path().string(), column_key.to_vec());
2391 }
2392 let decryption_properties = decryption_properties_builder.build().unwrap();
2393 let decryptor = FileDecryptor::new(
2394 &decryption_properties,
2395 footer_key_metadata.as_deref(),
2396 aad_file_unique,
2397 aad_prefix,
2398 )
2399 .unwrap();
2400
2401 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2402 .set_row_groups(row_group_meta.clone())
2403 .set_file_decryptor(Some(decryptor))
2404 .build();
2405
2406 let expected_size_with_decryptor = 3088;
2407 assert!(expected_size_with_decryptor > base_expected_size);
2408
2409 assert_eq!(
2410 parquet_meta_data.memory_size(),
2411 expected_size_with_decryptor
2412 );
2413 }
2414
2415 fn get_test_schema_descr() -> SchemaDescPtr {
2417 let schema = SchemaType::group_type_builder("schema")
2418 .with_fields(vec![
2419 Arc::new(
2420 SchemaType::primitive_type_builder("a", Type::INT32)
2421 .build()
2422 .unwrap(),
2423 ),
2424 Arc::new(
2425 SchemaType::primitive_type_builder("b", Type::INT32)
2426 .build()
2427 .unwrap(),
2428 ),
2429 ])
2430 .build()
2431 .unwrap();
2432
2433 Arc::new(SchemaDescriptor::new(Arc::new(schema)))
2434 }
2435}