1#[cfg(feature = "arrow")]
54mod dictionary;
55mod footer_tail;
56mod memory;
57mod options;
58pub mod page_index;
59mod parser;
60mod push_decoder;
61pub(crate) mod reader;
62pub(crate) mod thrift;
63mod writer;
64
65use crate::basic::{
66 BoundaryOrder, ColumnOrder, Compression, CompressionCodec, Encoding, EncodingMask, PageType,
67 Type,
68};
69#[cfg(feature = "encryption")]
70use crate::encryption::decrypt::FileDecryptor;
71use crate::errors::{ParquetError, Result};
72#[cfg(feature = "encryption")]
73use crate::file::column_crypto_metadata::ColumnCryptoMetaData;
74pub(crate) use crate::file::metadata::memory::HeapSize;
75use crate::file::metadata::page_index::{PageIndex, PageIndexProvider, RowGroupPageIndex};
76#[cfg(feature = "encryption")]
77use crate::file::metadata::thrift::encryption::EncryptionAlgorithm;
78use crate::file::page_index::column_index::{ByteArrayColumnIndex, PrimitiveColumnIndex};
79use crate::file::page_index::{column_index::ColumnIndexMetaData, offset_index::PageLocation};
80use crate::file::statistics::Statistics;
81use crate::geospatial::statistics as geo_statistics;
82use crate::parquet_thrift::{
83 ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
84 WriteThrift, WriteThriftField,
85};
86use crate::schema::types::{
87 ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, SchemaDescriptor,
88 Type as SchemaType,
89};
90use crate::thrift_struct;
91use crate::{
92 data_type::private::ParquetValueType, file::page_index::offset_index::OffsetIndexMetaData,
93};
94
95pub use footer_tail::FooterTail;
96pub use options::{ParquetMetaDataOptions, ParquetStatisticsPolicy};
97pub use push_decoder::ParquetMetaDataPushDecoder;
98pub use reader::{ColumnChunkMask, PageIndexPolicy, ParquetMetaDataReader};
99use std::io::Write;
100use std::ops::Range;
101use std::sync::Arc;
102pub use writer::ParquetMetaDataWriter;
103pub(crate) use writer::ThriftMetadataWriter;
104
105#[derive(Debug, Clone)]
126pub struct ParquetMetaData {
127 file_metadata: Arc<FileMetaData>,
129 row_groups: Arc<[RowGroupMetaData]>,
131 page_index: Option<Arc<dyn PageIndexProvider>>,
133 #[cfg(feature = "encryption")]
135 file_decryptor: Option<Arc<FileDecryptor>>,
136}
137
138impl ParquetMetaData {
139 pub fn new(file_metadata: FileMetaData, row_groups: Vec<RowGroupMetaData>) -> Self {
142 ParquetMetaData {
143 file_metadata: Arc::new(file_metadata),
144 row_groups: row_groups.into(),
145 page_index: None,
146 #[cfg(feature = "encryption")]
147 file_decryptor: None,
148 }
149 }
150
151 pub fn into_builder(self) -> ParquetMetaDataBuilder {
156 self.into()
157 }
158
159 pub fn file_metadata(&self) -> &FileMetaData {
161 &self.file_metadata
162 }
163
164 #[cfg(feature = "encryption")]
166 pub(crate) fn file_decryptor(&self) -> Option<&FileDecryptor> {
167 self.file_decryptor.as_deref()
168 }
169
170 pub fn num_row_groups(&self) -> usize {
172 self.row_groups.len()
173 }
174
175 pub fn row_group(&self, i: usize) -> &RowGroupMetaData {
178 &self.row_groups[i]
179 }
180
181 pub fn row_groups(&self) -> &[RowGroupMetaData] {
183 &self.row_groups
184 }
185
186 pub fn row_group_num_rows(&self, row_group_idx: usize) -> Result<usize> {
191 self.row_groups
192 .get(row_group_idx)
193 .ok_or_else(|| {
194 ParquetError::General(format!(
195 "Row group index {row_group_idx} out of bounds for file with {} row groups",
196 self.num_row_groups()
197 ))
198 })?
199 .num_rows()
200 .try_into()
201 .map_err(|e| ParquetError::General(format!("Row count overflow: {e}")))
202 }
203
204 pub fn page_index(&self) -> Option<&Arc<dyn PageIndexProvider>> {
211 self.page_index.as_ref()
212 }
213
214 pub fn page_index_for_row_group(&self, row_group_idx: usize) -> RowGroupPageIndex {
216 RowGroupPageIndex::new(row_group_idx, self.page_index.clone())
217 }
218
219 pub fn memory_size(&self) -> usize {
234 #[cfg(feature = "encryption")]
235 let encryption_size = self.file_decryptor.heap_size();
236 #[cfg(not(feature = "encryption"))]
237 let encryption_size = 0usize;
238
239 let page_index_size = if let Some(page_index) = self.page_index.as_ref() {
244 if let Some(page_index) = page_index.as_any().downcast_ref::<PageIndex>() {
245 2 * std::mem::size_of::<usize>() + std::mem::size_of::<PageIndex>() + page_index.heap_size()
249 } else {
250 0
251 }
252 } else {
253 0
254 };
255
256 std::mem::size_of::<Self>()
257 + self.file_metadata.heap_size()
258 + self.row_groups.heap_size()
259 + page_index_size
260 + encryption_size
261 }
262
263 pub(crate) fn set_page_index(&mut self, index: Option<Arc<dyn PageIndexProvider>>) {
265 self.page_index = index;
266 }
267}
268
269impl PartialEq for ParquetMetaData {
270 fn eq(&self, other: &Self) -> bool {
271 if self.file_metadata != other.file_metadata || self.row_groups != other.row_groups {
273 return false;
274 }
275
276 #[cfg(feature = "encryption")]
277 if self.file_decryptor != other.file_decryptor {
278 return false;
279 }
280
281 match (&self.page_index, &other.page_index) {
283 (None, None) => true,
284 (Some(a), Some(b)) => {
285 let a_page_index = a.as_any().downcast_ref::<PageIndex>();
287 let b_page_index = b.as_any().downcast_ref::<PageIndex>();
288 match (a_page_index, b_page_index) {
289 (Some(a_idx), Some(b_idx)) => a_idx == b_idx,
290 _ => Arc::ptr_eq(a, b), }
292 }
293 _ => false,
294 }
295 }
296}
297
298enum RowGroups {
300 Owned(Vec<RowGroupMetaData>),
302 Shared(Arc<[RowGroupMetaData]>),
304}
305
306impl RowGroups {
307 fn as_slice(&self) -> &[RowGroupMetaData] {
309 match self {
310 RowGroups::Owned(v) => v.as_slice(),
311 RowGroups::Shared(arc) => arc.as_ref(),
312 }
313 }
314
315 fn into_vec(self) -> Vec<RowGroupMetaData> {
317 match self {
318 RowGroups::Owned(v) => v,
319 RowGroups::Shared(arc) => arc.to_vec(),
320 }
321 }
322
323 fn convert_to_owned(&mut self) {
325 if matches!(self, RowGroups::Shared(_)) {
326 let old = std::mem::replace(self, RowGroups::Owned(Vec::new()));
327 *self = RowGroups::Owned(old.into_vec());
328 }
329 }
330
331 fn as_mut_vec(&mut self) -> &mut Vec<RowGroupMetaData> {
333 self.convert_to_owned();
334 match self {
335 RowGroups::Owned(v) => v,
336 RowGroups::Shared(_) => unreachable!(),
337 }
338 }
339}
340
341pub struct ParquetMetaDataBuilder {
379 file_metadata: Arc<FileMetaData>,
380 row_groups: RowGroups,
381 page_index: Option<Arc<dyn PageIndexProvider>>,
382 #[cfg(feature = "encryption")]
383 file_decryptor: Option<Arc<FileDecryptor>>,
384}
385
386impl ParquetMetaDataBuilder {
387 pub fn new(file_meta_data: FileMetaData) -> Self {
389 Self {
390 file_metadata: Arc::new(file_meta_data),
391 row_groups: RowGroups::Owned(Vec::new()),
392 page_index: None,
393 #[cfg(feature = "encryption")]
394 file_decryptor: None,
395 }
396 }
397
398 pub fn new_from_metadata(metadata: ParquetMetaData) -> Self {
400 Self {
401 file_metadata: metadata.file_metadata,
402 row_groups: RowGroups::Shared(metadata.row_groups),
403 page_index: metadata.page_index,
404 #[cfg(feature = "encryption")]
405 file_decryptor: metadata.file_decryptor,
406 }
407 }
408
409 pub fn add_row_group(mut self, row_group: RowGroupMetaData) -> Self {
414 self.row_groups.as_mut_vec().push(row_group);
415 self
416 }
417
418 pub fn set_row_groups(mut self, row_groups: Vec<RowGroupMetaData>) -> Self {
420 self.row_groups = RowGroups::Owned(row_groups);
421 self
422 }
423
424 pub fn take_row_groups(&mut self) -> Vec<RowGroupMetaData> {
433 std::mem::replace(&mut self.row_groups, RowGroups::Owned(Vec::new())).into_vec()
434 }
435
436 pub fn row_groups(&self) -> &[RowGroupMetaData] {
438 self.row_groups.as_slice()
439 }
440
441 pub fn set_page_index(mut self, page_index: Option<Arc<dyn PageIndexProvider>>) -> Self {
450 self.page_index = page_index;
451 self
452 }
453
454 pub fn take_page_index(&mut self) -> Option<Arc<dyn PageIndexProvider>> {
456 std::mem::take(&mut self.page_index)
457 }
458
459 pub fn page_index(&self) -> Option<&dyn PageIndexProvider> {
461 self.page_index.as_ref().map(|arc| arc.as_ref())
462 }
463
464 #[cfg(feature = "encryption")]
466 pub(crate) fn set_file_decryptor(mut self, file_decryptor: Option<FileDecryptor>) -> Self {
467 self.file_decryptor = file_decryptor.map(Arc::new);
468 self
469 }
470
471 pub fn build(self) -> ParquetMetaData {
473 ParquetMetaData {
474 file_metadata: self.file_metadata,
475 row_groups: match self.row_groups {
476 RowGroups::Owned(v) => v.into(),
477 RowGroups::Shared(arc) => arc,
478 },
479 page_index: self.page_index,
480 #[cfg(feature = "encryption")]
481 file_decryptor: self.file_decryptor,
482 }
483 }
484}
485
486impl From<ParquetMetaData> for ParquetMetaDataBuilder {
487 fn from(meta_data: ParquetMetaData) -> Self {
488 Self::new_from_metadata(meta_data)
489 }
490}
491
492thrift_struct!(
493pub struct KeyValue {
495 1: required string key
496 2: optional string value
497}
498);
499
500impl KeyValue {
501 pub fn new<F2>(key: String, value: F2) -> KeyValue
503 where
504 F2: Into<Option<String>>,
505 {
506 KeyValue {
507 key,
508 value: value.into(),
509 }
510 }
511}
512
513thrift_struct!(
514pub struct PageEncodingStats {
516 1: required PageType page_type;
517 2: required Encoding encoding;
518 3: required i32 count;
519}
520);
521
522#[derive(Debug, Clone, PartialEq)]
525enum ParquetPageEncodingStats {
526 Full(Vec<PageEncodingStats>),
528 Mask(EncodingMask),
530}
531
532pub type FileMetaDataPtr = Arc<FileMetaData>;
534
535#[derive(Debug, Clone, PartialEq)]
539pub struct FileMetaData {
540 version: i32,
541 num_rows: i64,
542 created_by: Option<String>,
543 key_value_metadata: Option<Vec<KeyValue>>,
544 schema_descr: SchemaDescPtr,
545 column_orders: Option<Vec<ColumnOrder>>,
546 #[cfg(feature = "encryption")]
547 encryption_algorithm: Option<Box<EncryptionAlgorithm>>,
548 #[cfg(feature = "encryption")]
549 footer_signing_key_metadata: Option<Vec<u8>>,
550}
551
552impl FileMetaData {
553 pub fn new(
555 version: i32,
556 num_rows: i64,
557 created_by: Option<String>,
558 key_value_metadata: Option<Vec<KeyValue>>,
559 schema_descr: SchemaDescPtr,
560 column_orders: Option<Vec<ColumnOrder>>,
561 ) -> Self {
562 FileMetaData {
563 version,
564 num_rows,
565 created_by,
566 key_value_metadata,
567 schema_descr,
568 column_orders,
569 #[cfg(feature = "encryption")]
570 encryption_algorithm: None,
571 #[cfg(feature = "encryption")]
572 footer_signing_key_metadata: None,
573 }
574 }
575
576 #[cfg(feature = "encryption")]
577 pub(crate) fn with_encryption_algorithm(
578 mut self,
579 encryption_algorithm: Option<EncryptionAlgorithm>,
580 ) -> Self {
581 self.encryption_algorithm = encryption_algorithm.map(Box::new);
582 self
583 }
584
585 #[cfg(feature = "encryption")]
586 pub(crate) fn with_footer_signing_key_metadata(
587 mut self,
588 footer_signing_key_metadata: Option<Vec<u8>>,
589 ) -> Self {
590 self.footer_signing_key_metadata = footer_signing_key_metadata;
591 self
592 }
593
594 pub fn version(&self) -> i32 {
596 self.version
597 }
598
599 pub fn num_rows(&self) -> i64 {
601 self.num_rows
602 }
603
604 pub fn created_by(&self) -> Option<&str> {
613 self.created_by.as_deref()
614 }
615
616 pub fn key_value_metadata(&self) -> Option<&Vec<KeyValue>> {
618 self.key_value_metadata.as_ref()
619 }
620
621 pub fn schema(&self) -> &SchemaType {
625 self.schema_descr.root_schema()
626 }
627
628 pub fn schema_descr(&self) -> &SchemaDescriptor {
630 &self.schema_descr
631 }
632
633 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
635 self.schema_descr.clone()
636 }
637
638 pub fn column_orders(&self) -> Option<&Vec<ColumnOrder>> {
646 self.column_orders.as_ref()
647 }
648
649 pub fn column_order(&self, i: usize) -> ColumnOrder {
652 self.column_orders
653 .as_ref()
654 .map(|data| data[i])
655 .unwrap_or(ColumnOrder::UNDEFINED)
656 }
657}
658
659thrift_struct!(
660pub struct SortingColumn {
662 1: required i32 column_idx
664
665 2: required bool descending
667
668 3: required bool nulls_first
671}
672);
673
674pub type RowGroupMetaDataPtr = Arc<RowGroupMetaData>;
676
677#[derive(Debug, Clone, PartialEq)]
682pub struct RowGroupMetaData {
683 columns: Vec<ColumnChunkMetaData>,
684 num_rows: i64,
685 sorting_columns: Option<Vec<SortingColumn>>,
686 total_byte_size: i64,
687 schema_descr: SchemaDescPtr,
688 file_offset: Option<i64>,
690 ordinal: Option<i32>,
692}
693
694impl RowGroupMetaData {
695 pub fn builder(schema_descr: SchemaDescPtr) -> RowGroupMetaDataBuilder {
697 RowGroupMetaDataBuilder::new(schema_descr)
698 }
699
700 pub fn num_columns(&self) -> usize {
702 self.columns.len()
703 }
704
705 pub fn column(&self, i: usize) -> &ColumnChunkMetaData {
707 &self.columns[i]
708 }
709
710 pub fn columns(&self) -> &[ColumnChunkMetaData] {
712 &self.columns
713 }
714
715 pub fn columns_mut(&mut self) -> &mut [ColumnChunkMetaData] {
717 &mut self.columns
718 }
719
720 pub fn num_rows(&self) -> i64 {
722 self.num_rows
723 }
724
725 pub fn sorting_columns(&self) -> Option<&Vec<SortingColumn>> {
727 self.sorting_columns.as_ref()
728 }
729
730 pub fn total_byte_size(&self) -> i64 {
732 self.total_byte_size
733 }
734
735 pub fn compressed_size(&self) -> i64 {
737 self.columns.iter().map(|c| c.total_compressed_size).sum()
738 }
739
740 pub fn schema_descr(&self) -> &SchemaDescriptor {
742 self.schema_descr.as_ref()
743 }
744
745 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
747 self.schema_descr.clone()
748 }
749
750 #[inline(always)]
755 pub fn ordinal(&self) -> Option<i32> {
756 self.ordinal
757 }
758
759 #[inline(always)]
761 pub fn file_offset(&self) -> Option<i64> {
762 self.file_offset
763 }
764
765 pub fn into_builder(self) -> RowGroupMetaDataBuilder {
767 RowGroupMetaDataBuilder(self)
768 }
769}
770
771impl HeapSize for Arc<[RowGroupMetaData]> {
772 fn heap_size(&self) -> usize {
773 let rg_size = self.iter().map(|t| t.heap_size()).sum::<usize>();
775 2 * std::mem::size_of::<usize>() + std::mem::size_of_val(self.as_ref()) + rg_size
778 }
779}
780
781pub struct RowGroupMetaDataBuilder(RowGroupMetaData);
783
784impl RowGroupMetaDataBuilder {
785 fn new(schema_descr: SchemaDescPtr) -> Self {
787 Self(RowGroupMetaData {
788 columns: Vec::with_capacity(schema_descr.num_columns()),
789 schema_descr,
790 file_offset: None,
791 num_rows: 0,
792 sorting_columns: None,
793 total_byte_size: 0,
794 ordinal: None,
795 })
796 }
797
798 pub fn set_num_rows(mut self, value: i64) -> Self {
800 self.0.num_rows = value;
801 self
802 }
803
804 pub fn set_sorting_columns(mut self, value: Option<Vec<SortingColumn>>) -> Self {
806 self.0.sorting_columns = value;
807 self
808 }
809
810 pub fn set_total_byte_size(mut self, value: i64) -> Self {
812 self.0.total_byte_size = value;
813 self
814 }
815
816 pub fn take_columns(&mut self) -> Vec<ColumnChunkMetaData> {
822 std::mem::take(&mut self.0.columns)
823 }
824
825 pub fn set_column_metadata(mut self, value: Vec<ColumnChunkMetaData>) -> Self {
827 self.0.columns = value;
828 self
829 }
830
831 pub fn add_column_metadata(mut self, value: ColumnChunkMetaData) -> Self {
833 self.0.columns.push(value);
834 self
835 }
836
837 pub fn set_ordinal(mut self, value: i32) -> Self {
839 self.0.ordinal = Some(value);
840 self
841 }
842
843 pub fn set_file_offset(mut self, value: i64) -> Self {
845 self.0.file_offset = Some(value);
846 self
847 }
848
849 pub fn build(self) -> Result<RowGroupMetaData> {
851 if self.0.schema_descr.num_columns() != self.0.columns.len() {
852 return Err(general_err!(
853 "Column length mismatch: {} != {}",
854 self.0.schema_descr.num_columns(),
855 self.0.columns.len()
856 ));
857 }
858
859 Ok(self.0)
860 }
861
862 pub(super) fn build_unchecked(self) -> RowGroupMetaData {
864 self.0
865 }
866}
867
868#[derive(Debug, Clone, PartialEq)]
870pub struct ColumnChunkMetaData {
871 column_descr: ColumnDescPtr,
872 encodings: EncodingMask,
873 file_path: Option<String>,
874 file_offset: i64,
875 num_values: i64,
876 compression: CompressionCodec,
877 total_compressed_size: i64,
878 total_uncompressed_size: i64,
879 data_page_offset: i64,
880 index_page_offset: Option<i64>,
881 dictionary_page_offset: Option<i64>,
882 statistics: Option<Statistics>,
883 geo_statistics: Option<Box<geo_statistics::GeospatialStatistics>>,
884 encoding_stats: Option<ParquetPageEncodingStats>,
885 bloom_filter_offset: Option<i64>,
886 bloom_filter_length: Option<i32>,
887 offset_index_offset: Option<i64>,
888 offset_index_length: Option<i32>,
889 column_index_offset: Option<i64>,
890 column_index_length: Option<i32>,
891 unencoded_byte_array_data_bytes: Option<i64>,
892 repetition_level_histogram: Option<LevelHistogram>,
893 definition_level_histogram: Option<LevelHistogram>,
894 #[cfg(feature = "encryption")]
895 column_crypto_metadata: Option<Box<ColumnCryptoMetaData>>,
896 #[cfg(feature = "encryption")]
897 encrypted_column_metadata: Option<Vec<u8>>,
898 #[cfg(feature = "encryption")]
902 plaintext_footer_mode: bool,
903}
904
905#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
914pub struct LevelHistogram {
915 inner: Vec<i64>,
916}
917
918impl LevelHistogram {
919 pub fn try_new(max_level: i16) -> Option<Self> {
925 if max_level > 0 {
926 Some(Self {
927 inner: vec![0; max_level as usize + 1],
928 })
929 } else {
930 None
931 }
932 }
933 pub fn values(&self) -> &[i64] {
935 &self.inner
936 }
937
938 pub fn into_inner(self) -> Vec<i64> {
940 self.inner
941 }
942
943 pub fn get(&self, index: usize) -> Option<i64> {
950 self.inner.get(index).copied()
951 }
952
953 pub fn add(&mut self, other: &Self) {
958 assert_eq!(self.len(), other.len());
959 for (dst, src) in self.inner.iter_mut().zip(other.inner.iter()) {
960 *dst += src;
961 }
962 }
963
964 pub fn len(&self) -> usize {
966 self.inner.len()
967 }
968
969 pub fn is_empty(&self) -> bool {
971 self.inner.is_empty()
972 }
973
974 pub fn reset(&mut self) {
976 self.inner.fill(0);
977 }
978
979 #[inline]
981 pub fn increment_by(&mut self, level: i16, count: i64) {
982 self.inner[level as usize] += count;
983 }
984
985 #[deprecated(since = "58.2.0", note = "Use `increment_by` instead")]
991 pub fn update_from_levels(&mut self, levels: &[i16]) {
992 for &level in levels {
993 self.increment_by(level, 1);
994 }
995 }
996}
997
998impl From<Vec<i64>> for LevelHistogram {
999 fn from(inner: Vec<i64>) -> Self {
1000 Self { inner }
1001 }
1002}
1003
1004impl From<LevelHistogram> for Vec<i64> {
1005 fn from(value: LevelHistogram) -> Self {
1006 value.into_inner()
1007 }
1008}
1009
1010impl HeapSize for LevelHistogram {
1011 fn heap_size(&self) -> usize {
1012 self.inner.heap_size()
1013 }
1014}
1015
1016impl ColumnChunkMetaData {
1018 pub fn builder(column_descr: ColumnDescPtr) -> ColumnChunkMetaDataBuilder {
1020 ColumnChunkMetaDataBuilder::new(column_descr)
1021 }
1022
1023 pub fn file_path(&self) -> Option<&str> {
1028 self.file_path.as_deref()
1029 }
1030
1031 pub fn file_offset(&self) -> i64 {
1038 self.file_offset
1039 }
1040
1041 pub fn column_type(&self) -> Type {
1043 self.column_descr.physical_type()
1044 }
1045
1046 pub fn column_path(&self) -> &ColumnPath {
1048 self.column_descr.path()
1049 }
1050
1051 pub fn column_descr(&self) -> &ColumnDescriptor {
1053 self.column_descr.as_ref()
1054 }
1055
1056 pub fn column_descr_ptr(&self) -> ColumnDescPtr {
1058 self.column_descr.clone()
1059 }
1060
1061 pub fn encodings(&self) -> impl Iterator<Item = Encoding> {
1063 self.encodings.encodings()
1064 }
1065
1066 pub fn encodings_mask(&self) -> &EncodingMask {
1068 &self.encodings
1069 }
1070
1071 pub fn num_values(&self) -> i64 {
1073 self.num_values
1074 }
1075
1076 pub fn compression(&self) -> Compression {
1083 self.compression.into()
1084 }
1085
1086 pub fn compression_codec(&self) -> CompressionCodec {
1088 self.compression
1089 }
1090
1091 pub fn compressed_size(&self) -> i64 {
1093 self.total_compressed_size
1094 }
1095
1096 pub fn uncompressed_size(&self) -> i64 {
1098 self.total_uncompressed_size
1099 }
1100
1101 pub fn data_page_offset(&self) -> i64 {
1103 self.data_page_offset
1104 }
1105
1106 pub fn index_page_offset(&self) -> Option<i64> {
1108 self.index_page_offset
1109 }
1110
1111 pub fn dictionary_page_offset(&self) -> Option<i64> {
1113 self.dictionary_page_offset
1114 }
1115
1116 pub fn byte_range(&self) -> (u64, u64) {
1122 let col_start = match self.dictionary_page_offset() {
1123 Some(dictionary_page_offset) => dictionary_page_offset,
1124 None => self.data_page_offset(),
1125 };
1126 let col_len = self.compressed_size();
1127 assert!(
1128 col_start >= 0 && col_len >= 0,
1129 "column start and length should not be negative"
1130 );
1131 (col_start as u64, col_len as u64)
1132 }
1133
1134 pub fn statistics(&self) -> Option<&Statistics> {
1137 self.statistics.as_ref()
1138 }
1139
1140 pub fn geo_statistics(&self) -> Option<&geo_statistics::GeospatialStatistics> {
1143 self.geo_statistics.as_deref()
1144 }
1145
1146 pub fn page_encoding_stats(&self) -> Option<&Vec<PageEncodingStats>> {
1153 match self.encoding_stats.as_ref() {
1154 Some(ParquetPageEncodingStats::Full(stats)) => Some(stats),
1155 _ => None,
1156 }
1157 }
1158
1159 pub fn page_encoding_stats_mask(&self) -> Option<&EncodingMask> {
1189 match self.encoding_stats.as_ref() {
1190 Some(ParquetPageEncodingStats::Mask(stats)) => Some(stats),
1191 _ => None,
1192 }
1193 }
1194
1195 pub fn bloom_filter_offset(&self) -> Option<i64> {
1197 self.bloom_filter_offset
1198 }
1199
1200 pub fn bloom_filter_length(&self) -> Option<i32> {
1202 self.bloom_filter_length
1203 }
1204
1205 pub fn column_index_offset(&self) -> Option<i64> {
1207 self.column_index_offset
1208 }
1209
1210 pub fn column_index_length(&self) -> Option<i32> {
1212 self.column_index_length
1213 }
1214
1215 pub fn column_index_range(&self) -> Option<Range<u64>> {
1221 let offset = u64::try_from(self.column_index_offset?).ok()?;
1222 let length = u64::try_from(self.column_index_length?).ok()?;
1223 Some(offset..(offset + length))
1224 }
1225
1226 pub fn offset_index_offset(&self) -> Option<i64> {
1228 self.offset_index_offset
1229 }
1230
1231 pub fn offset_index_length(&self) -> Option<i32> {
1233 self.offset_index_length
1234 }
1235
1236 pub fn offset_index_range(&self) -> Option<Range<u64>> {
1242 let offset = u64::try_from(self.offset_index_offset?).ok()?;
1243 let length = u64::try_from(self.offset_index_length?).ok()?;
1244 Some(offset..(offset + length))
1245 }
1246
1247 pub fn unencoded_byte_array_data_bytes(&self) -> Option<i64> {
1252 self.unencoded_byte_array_data_bytes
1253 }
1254
1255 pub fn repetition_level_histogram(&self) -> Option<&LevelHistogram> {
1261 self.repetition_level_histogram.as_ref()
1262 }
1263
1264 pub fn definition_level_histogram(&self) -> Option<&LevelHistogram> {
1270 self.definition_level_histogram.as_ref()
1271 }
1272
1273 #[cfg(feature = "encryption")]
1275 pub fn crypto_metadata(&self) -> Option<&ColumnCryptoMetaData> {
1276 self.column_crypto_metadata.as_deref()
1277 }
1278
1279 pub fn into_builder(self) -> ColumnChunkMetaDataBuilder {
1281 ColumnChunkMetaDataBuilder::from(self)
1282 }
1283}
1284
1285pub struct ColumnChunkMetaDataBuilder(ColumnChunkMetaData);
1304
1305impl ColumnChunkMetaDataBuilder {
1306 fn new(column_descr: ColumnDescPtr) -> Self {
1310 Self(ColumnChunkMetaData {
1311 column_descr,
1312 encodings: Default::default(),
1313 file_path: None,
1314 file_offset: 0,
1315 num_values: 0,
1316 compression: CompressionCodec::UNCOMPRESSED,
1317 total_compressed_size: 0,
1318 total_uncompressed_size: 0,
1319 data_page_offset: 0,
1320 index_page_offset: None,
1321 dictionary_page_offset: None,
1322 statistics: None,
1323 geo_statistics: None,
1324 encoding_stats: None,
1325 bloom_filter_offset: None,
1326 bloom_filter_length: None,
1327 offset_index_offset: None,
1328 offset_index_length: None,
1329 column_index_offset: None,
1330 column_index_length: None,
1331 unencoded_byte_array_data_bytes: None,
1332 repetition_level_histogram: None,
1333 definition_level_histogram: None,
1334 #[cfg(feature = "encryption")]
1335 column_crypto_metadata: None,
1336 #[cfg(feature = "encryption")]
1337 encrypted_column_metadata: None,
1338 #[cfg(feature = "encryption")]
1339 plaintext_footer_mode: false,
1340 })
1341 }
1342
1343 pub fn set_encodings(mut self, encodings: Vec<Encoding>) -> Self {
1345 self.0.encodings = EncodingMask::new_from_encodings(encodings.iter());
1346 self
1347 }
1348
1349 pub fn set_encodings_mask(mut self, encodings: EncodingMask) -> Self {
1351 self.0.encodings = encodings;
1352 self
1353 }
1354
1355 pub fn set_file_path(mut self, value: String) -> Self {
1357 self.0.file_path = Some(value);
1358 self
1359 }
1360
1361 pub fn set_num_values(mut self, value: i64) -> Self {
1363 self.0.num_values = value;
1364 self
1365 }
1366
1367 pub fn set_compression(mut self, value: Compression) -> Self {
1369 self.0.compression = value.into();
1370 self
1371 }
1372
1373 pub fn set_compression_codec(mut self, value: CompressionCodec) -> Self {
1375 self.0.compression = value;
1376 self
1377 }
1378
1379 pub fn set_total_compressed_size(mut self, value: i64) -> Self {
1381 self.0.total_compressed_size = value;
1382 self
1383 }
1384
1385 pub fn set_total_uncompressed_size(mut self, value: i64) -> Self {
1387 self.0.total_uncompressed_size = value;
1388 self
1389 }
1390
1391 pub fn set_data_page_offset(mut self, value: i64) -> Self {
1393 self.0.data_page_offset = value;
1394 self
1395 }
1396
1397 pub fn set_dictionary_page_offset(mut self, value: Option<i64>) -> Self {
1399 self.0.dictionary_page_offset = value;
1400 self
1401 }
1402
1403 pub fn set_index_page_offset(mut self, value: Option<i64>) -> Self {
1405 self.0.index_page_offset = value;
1406 self
1407 }
1408
1409 pub fn set_statistics(mut self, value: Statistics) -> Self {
1411 self.0.statistics = Some(value);
1412 self
1413 }
1414
1415 pub fn set_geo_statistics(mut self, value: Box<geo_statistics::GeospatialStatistics>) -> Self {
1417 self.0.geo_statistics = Some(value);
1418 self
1419 }
1420
1421 pub fn clear_statistics(mut self) -> Self {
1423 self.0.statistics = None;
1424 self
1425 }
1426
1427 pub fn set_page_encoding_stats(mut self, value: Vec<PageEncodingStats>) -> Self {
1431 self.0.encoding_stats = Some(ParquetPageEncodingStats::Full(value));
1432 self
1433 }
1434
1435 pub fn set_page_encoding_stats_mask(mut self, value: EncodingMask) -> Self {
1439 self.0.encoding_stats = Some(ParquetPageEncodingStats::Mask(value));
1440 self
1441 }
1442
1443 pub fn clear_page_encoding_stats(mut self) -> Self {
1445 self.0.encoding_stats = None;
1446 self
1447 }
1448
1449 pub fn set_bloom_filter_offset(mut self, value: Option<i64>) -> Self {
1451 self.0.bloom_filter_offset = value;
1452 self
1453 }
1454
1455 pub fn set_bloom_filter_length(mut self, value: Option<i32>) -> Self {
1457 self.0.bloom_filter_length = value;
1458 self
1459 }
1460
1461 pub fn set_offset_index_offset(mut self, value: Option<i64>) -> Self {
1463 self.0.offset_index_offset = value;
1464 self
1465 }
1466
1467 pub fn set_offset_index_length(mut self, value: Option<i32>) -> Self {
1469 self.0.offset_index_length = value;
1470 self
1471 }
1472
1473 pub fn set_column_index_offset(mut self, value: Option<i64>) -> Self {
1475 self.0.column_index_offset = value;
1476 self
1477 }
1478
1479 pub fn set_column_index_length(mut self, value: Option<i32>) -> Self {
1481 self.0.column_index_length = value;
1482 self
1483 }
1484
1485 pub fn set_unencoded_byte_array_data_bytes(mut self, value: Option<i64>) -> Self {
1487 self.0.unencoded_byte_array_data_bytes = value;
1488 self
1489 }
1490
1491 pub fn set_repetition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1493 self.0.repetition_level_histogram = value;
1494 self
1495 }
1496
1497 pub fn set_definition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1499 self.0.definition_level_histogram = value;
1500 self
1501 }
1502
1503 #[cfg(feature = "encryption")]
1504 pub fn set_column_crypto_metadata(mut self, value: Option<ColumnCryptoMetaData>) -> Self {
1506 self.0.column_crypto_metadata = value.map(Box::new);
1507 self
1508 }
1509
1510 #[cfg(feature = "encryption")]
1511 pub fn set_encrypted_column_metadata(mut self, value: Option<Vec<u8>>) -> Self {
1513 self.0.encrypted_column_metadata = value;
1514 self
1515 }
1516
1517 pub fn build(self) -> Result<ColumnChunkMetaData> {
1519 Ok(self.0)
1520 }
1521}
1522
1523pub struct ColumnIndexBuilder {
1528 column_type: Type,
1529 null_pages: Vec<bool>,
1530 min_values: Vec<Vec<u8>>,
1531 max_values: Vec<Vec<u8>>,
1532 null_counts: Vec<i64>,
1533 nan_counts: Vec<Option<i64>>,
1534 boundary_order: BoundaryOrder,
1535 repetition_level_histograms: Option<Vec<i64>>,
1537 definition_level_histograms: Option<Vec<i64>>,
1539 valid: bool,
1547}
1548
1549impl ColumnIndexBuilder {
1550 pub fn new(column_type: Type) -> Self {
1552 ColumnIndexBuilder {
1553 column_type,
1554 null_pages: Vec::new(),
1555 min_values: Vec::new(),
1556 max_values: Vec::new(),
1557 null_counts: Vec::new(),
1558 nan_counts: Vec::new(),
1559 boundary_order: BoundaryOrder::UNORDERED,
1560 repetition_level_histograms: None,
1561 definition_level_histograms: None,
1562 valid: true,
1563 }
1564 }
1565
1566 pub fn append(
1573 &mut self,
1574 null_page: bool,
1575 min_value: Vec<u8>,
1576 max_value: Vec<u8>,
1577 null_count: i64,
1578 nan_count: Option<i64>,
1579 ) {
1580 self.null_pages.push(null_page);
1581 self.min_values.push(min_value);
1582 self.max_values.push(max_value);
1583 self.null_counts.push(null_count);
1584 self.nan_counts.push(nan_count);
1585 }
1586
1587 pub fn append_histograms(
1592 &mut self,
1593 repetition_level_histogram: &Option<LevelHistogram>,
1594 definition_level_histogram: &Option<LevelHistogram>,
1595 ) {
1596 if !self.valid {
1597 return;
1598 }
1599 if let Some(rep_lvl_hist) = repetition_level_histogram {
1600 let hist = self.repetition_level_histograms.get_or_insert(Vec::new());
1601 hist.reserve(rep_lvl_hist.len());
1602 hist.extend(rep_lvl_hist.values());
1603 }
1604 if let Some(def_lvl_hist) = definition_level_histogram {
1605 let hist = self.definition_level_histograms.get_or_insert(Vec::new());
1606 hist.reserve(def_lvl_hist.len());
1607 hist.extend(def_lvl_hist.values());
1608 }
1609 }
1610
1611 pub fn set_boundary_order(&mut self, boundary_order: BoundaryOrder) {
1613 self.boundary_order = boundary_order;
1614 }
1615
1616 pub fn to_invalid(&mut self) {
1618 self.valid = false;
1619 }
1620
1621 pub fn valid(&self) -> bool {
1623 self.valid
1624 }
1625
1626 pub fn build(self) -> Result<ColumnIndexMetaData> {
1630 Ok(match self.column_type {
1631 Type::BOOLEAN => {
1632 let index = self.build_page_index(false)?;
1633 ColumnIndexMetaData::BOOLEAN(index)
1634 }
1635 Type::INT32 => {
1636 let index = self.build_page_index(false)?;
1637 ColumnIndexMetaData::INT32(index)
1638 }
1639 Type::INT64 => {
1640 let index = self.build_page_index(false)?;
1641 ColumnIndexMetaData::INT64(index)
1642 }
1643 Type::INT96 => {
1644 let index = self.build_page_index(false)?;
1645 ColumnIndexMetaData::INT96(index)
1646 }
1647 Type::FLOAT => {
1648 let index = self.build_page_index(true)?;
1649 ColumnIndexMetaData::FLOAT(index)
1650 }
1651 Type::DOUBLE => {
1652 let index = self.build_page_index(true)?;
1653 ColumnIndexMetaData::DOUBLE(index)
1654 }
1655 Type::BYTE_ARRAY => {
1656 let index = self.build_byte_array_index(false)?;
1657 ColumnIndexMetaData::BYTE_ARRAY(index)
1658 }
1659 Type::FIXED_LEN_BYTE_ARRAY => {
1660 let index = self.build_byte_array_index(true)?;
1661 ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index)
1662 }
1663 })
1664 }
1665
1666 fn build_nan_counts(nan_counts: &[Option<i64>]) -> Option<Vec<i64>> {
1667 let has_some = nan_counts.iter().any(|x| x.is_some());
1668 let has_none = nan_counts.iter().any(|x| x.is_none());
1669
1670 if has_some && !has_none {
1671 Some(nan_counts.iter().map(|x| x.unwrap()).collect())
1672 } else if !has_some && has_none {
1673 None
1674 } else {
1675 debug_assert!(
1676 false,
1677 "Mixed Some/None in nan_counts - caller should provide consistent values"
1678 );
1679 Some(nan_counts.iter().map(|x| x.unwrap_or(0)).collect())
1680 }
1681 }
1682
1683 fn build_page_index<T>(self, may_have_nan: bool) -> Result<PrimitiveColumnIndex<T>>
1684 where
1685 T: ParquetValueType,
1686 {
1687 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1688 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1689
1690 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1695 Self::build_nan_counts(&self.nan_counts)
1696 } else {
1697 None
1698 };
1699
1700 PrimitiveColumnIndex::try_new(
1701 self.null_pages,
1702 self.boundary_order,
1703 Some(self.null_counts),
1704 nan_counts,
1705 self.repetition_level_histograms,
1706 self.definition_level_histograms,
1707 min_values,
1708 max_values,
1709 )
1710 }
1711
1712 fn build_byte_array_index(self, may_have_nan: bool) -> Result<ByteArrayColumnIndex> {
1713 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1714 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1715
1716 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1721 Self::build_nan_counts(&self.nan_counts)
1722 } else {
1723 None
1724 };
1725
1726 ByteArrayColumnIndex::try_new(
1727 self.null_pages,
1728 self.boundary_order,
1729 Some(self.null_counts),
1730 nan_counts,
1731 self.repetition_level_histograms,
1732 self.definition_level_histograms,
1733 min_values,
1734 max_values,
1735 )
1736 }
1737}
1738
1739impl From<ColumnChunkMetaData> for ColumnChunkMetaDataBuilder {
1740 fn from(value: ColumnChunkMetaData) -> Self {
1741 ColumnChunkMetaDataBuilder(value)
1742 }
1743}
1744
1745pub struct OffsetIndexBuilder {
1749 offset_array: Vec<i64>,
1750 compressed_page_size_array: Vec<i32>,
1751 first_row_index_array: Vec<i64>,
1752 unencoded_byte_array_data_bytes_array: Option<Vec<i64>>,
1753 current_first_row_index: i64,
1754}
1755
1756impl Default for OffsetIndexBuilder {
1757 fn default() -> Self {
1758 Self::new()
1759 }
1760}
1761
1762impl OffsetIndexBuilder {
1763 pub fn new() -> Self {
1765 OffsetIndexBuilder {
1766 offset_array: Vec::new(),
1767 compressed_page_size_array: Vec::new(),
1768 first_row_index_array: Vec::new(),
1769 unencoded_byte_array_data_bytes_array: None,
1770 current_first_row_index: 0,
1771 }
1772 }
1773
1774 pub fn append_row_count(&mut self, row_count: i64) {
1776 let current_page_row_index = self.current_first_row_index;
1777 self.first_row_index_array.push(current_page_row_index);
1778 self.current_first_row_index += row_count;
1779 }
1780
1781 pub fn append_offset_and_size(&mut self, offset: i64, compressed_page_size: i32) {
1783 self.offset_array.push(offset);
1784 self.compressed_page_size_array.push(compressed_page_size);
1785 }
1786
1787 pub fn append_unencoded_byte_array_data_bytes(
1789 &mut self,
1790 unencoded_byte_array_data_bytes: Option<i64>,
1791 ) {
1792 if let Some(val) = unencoded_byte_array_data_bytes {
1793 self.unencoded_byte_array_data_bytes_array
1794 .get_or_insert(Vec::new())
1795 .push(val);
1796 }
1797 }
1798
1799 pub fn build(self) -> OffsetIndexMetaData {
1801 let locations = self
1802 .offset_array
1803 .iter()
1804 .zip(self.compressed_page_size_array.iter())
1805 .zip(self.first_row_index_array.iter())
1806 .map(|((offset, size), row_index)| PageLocation {
1807 offset: *offset,
1808 compressed_page_size: *size,
1809 first_row_index: *row_index,
1810 })
1811 .collect::<Vec<_>>();
1812 OffsetIndexMetaData {
1813 page_locations: locations,
1814 unencoded_byte_array_data_bytes: self.unencoded_byte_array_data_bytes_array,
1815 }
1816 }
1817}
1818
1819#[cfg(test)]
1820mod tests {
1821 use super::*;
1822 use crate::basic::{PageType, SortOrder};
1823 use crate::file::metadata::page_index::PageIndexBuilder;
1824 use crate::file::metadata::thrift::tests::{
1825 read_column_chunk, read_column_chunk_with_options, read_row_group,
1826 };
1827
1828 #[test]
1829 #[expect(deprecated)]
1830 fn test_level_histogram_update_from_levels_compat() {
1831 let mut histogram = LevelHistogram::try_new(2).unwrap();
1832 histogram.update_from_levels(&[0, 2, 1, 2, 2]);
1833 assert_eq!(histogram.values(), &[1, 1, 3]);
1834 }
1835
1836 #[test]
1837 fn test_row_group_metadata_thrift_conversion() {
1838 let schema_descr = get_test_schema_descr();
1839
1840 let mut columns = vec![];
1841 for ptr in schema_descr.columns() {
1842 let column = ColumnChunkMetaData::builder(ptr.clone()).build().unwrap();
1843 columns.push(column);
1844 }
1845 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
1846 .set_num_rows(1000)
1847 .set_total_byte_size(2000)
1848 .set_column_metadata(columns)
1849 .set_ordinal(1)
1850 .build()
1851 .unwrap();
1852
1853 let mut buf = Vec::new();
1854 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1855 row_group_meta.write_thrift(&mut writer).unwrap();
1856
1857 let row_group_res = read_row_group(&buf, schema_descr).unwrap();
1858
1859 assert_eq!(row_group_res, row_group_meta);
1860 }
1861
1862 #[test]
1863 fn test_row_group_metadata_thrift_conversion_empty() {
1864 let schema_descr = get_test_schema_descr();
1865
1866 let row_group_meta = RowGroupMetaData::builder(schema_descr).build();
1867
1868 assert!(row_group_meta.is_err());
1869 if let Err(e) = row_group_meta {
1870 assert_eq!(
1871 format!("{e}"),
1872 "Parquet error: Column length mismatch: 2 != 0"
1873 );
1874 }
1875 }
1876
1877 #[test]
1879 fn test_row_group_metadata_thrift_corrupted() {
1880 let schema_descr_2cols = Arc::new(SchemaDescriptor::new(Arc::new(
1881 SchemaType::group_type_builder("schema")
1882 .with_fields(vec![
1883 Arc::new(
1884 SchemaType::primitive_type_builder("a", Type::INT32)
1885 .build()
1886 .unwrap(),
1887 ),
1888 Arc::new(
1889 SchemaType::primitive_type_builder("b", Type::INT32)
1890 .build()
1891 .unwrap(),
1892 ),
1893 ])
1894 .build()
1895 .unwrap(),
1896 )));
1897
1898 let schema_descr_3cols = Arc::new(SchemaDescriptor::new(Arc::new(
1899 SchemaType::group_type_builder("schema")
1900 .with_fields(vec![
1901 Arc::new(
1902 SchemaType::primitive_type_builder("a", Type::INT32)
1903 .build()
1904 .unwrap(),
1905 ),
1906 Arc::new(
1907 SchemaType::primitive_type_builder("b", Type::INT32)
1908 .build()
1909 .unwrap(),
1910 ),
1911 Arc::new(
1912 SchemaType::primitive_type_builder("c", Type::INT32)
1913 .build()
1914 .unwrap(),
1915 ),
1916 ])
1917 .build()
1918 .unwrap(),
1919 )));
1920
1921 let row_group_meta_2cols = RowGroupMetaData::builder(schema_descr_2cols.clone())
1922 .set_num_rows(1000)
1923 .set_total_byte_size(2000)
1924 .set_column_metadata(vec![
1925 ColumnChunkMetaData::builder(schema_descr_2cols.column(0))
1926 .build()
1927 .unwrap(),
1928 ColumnChunkMetaData::builder(schema_descr_2cols.column(1))
1929 .build()
1930 .unwrap(),
1931 ])
1932 .set_ordinal(1)
1933 .build()
1934 .unwrap();
1935 let mut buf = Vec::new();
1936 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1937 row_group_meta_2cols.write_thrift(&mut writer).unwrap();
1938
1939 let err = read_row_group(&buf, schema_descr_3cols)
1940 .unwrap_err()
1941 .to_string();
1942 assert_eq!(
1943 err,
1944 "Parquet error: Column count mismatch. Schema has 3 columns while Row Group has 2"
1945 );
1946 }
1947
1948 #[test]
1949 fn test_column_chunk_metadata_thrift_conversion() {
1950 let column_descr = get_test_schema_descr().column(0);
1951 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1952 .set_encodings_mask(EncodingMask::new_from_encodings(
1953 [Encoding::PLAIN, Encoding::RLE].iter(),
1954 ))
1955 .set_file_path("file_path".to_owned())
1956 .set_num_values(1000)
1957 .set_compression_codec(CompressionCodec::SNAPPY)
1958 .set_total_compressed_size(2000)
1959 .set_total_uncompressed_size(3000)
1960 .set_data_page_offset(4000)
1961 .set_dictionary_page_offset(Some(5000))
1962 .set_page_encoding_stats(vec![
1963 PageEncodingStats {
1964 page_type: PageType::DATA_PAGE,
1965 encoding: Encoding::PLAIN,
1966 count: 3,
1967 },
1968 PageEncodingStats {
1969 page_type: PageType::DATA_PAGE,
1970 encoding: Encoding::RLE,
1971 count: 5,
1972 },
1973 ])
1974 .set_bloom_filter_offset(Some(6000))
1975 .set_bloom_filter_length(Some(25))
1976 .set_offset_index_offset(Some(7000))
1977 .set_offset_index_length(Some(25))
1978 .set_column_index_offset(Some(8000))
1979 .set_column_index_length(Some(25))
1980 .set_unencoded_byte_array_data_bytes(Some(2000))
1981 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
1982 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
1983 .build()
1984 .unwrap();
1985
1986 let mut buf = Vec::new();
1987 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1988 col_metadata.write_thrift(&mut writer).unwrap();
1989 let col_chunk_res = read_column_chunk(&buf, column_descr.clone()).unwrap();
1990
1991 let expected_metadata = ColumnChunkMetaData::builder(column_descr)
1992 .set_encodings_mask(EncodingMask::new_from_encodings(
1993 [Encoding::PLAIN, Encoding::RLE].iter(),
1994 ))
1995 .set_file_path("file_path".to_owned())
1996 .set_num_values(1000)
1997 .set_compression_codec(CompressionCodec::SNAPPY)
1998 .set_total_compressed_size(2000)
1999 .set_total_uncompressed_size(3000)
2000 .set_data_page_offset(4000)
2001 .set_dictionary_page_offset(Some(5000))
2002 .set_page_encoding_stats_mask(EncodingMask::new_from_encodings(
2003 [Encoding::PLAIN, Encoding::RLE].iter(),
2004 ))
2005 .set_bloom_filter_offset(Some(6000))
2006 .set_bloom_filter_length(Some(25))
2007 .set_offset_index_offset(Some(7000))
2008 .set_offset_index_length(Some(25))
2009 .set_column_index_offset(Some(8000))
2010 .set_column_index_length(Some(25))
2011 .set_unencoded_byte_array_data_bytes(Some(2000))
2012 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
2013 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
2014 .build()
2015 .unwrap();
2016
2017 assert_eq!(col_chunk_res, expected_metadata);
2018 }
2019
2020 #[test]
2021 fn test_column_chunk_metadata_thrift_conversion_full_stats() {
2022 let column_descr = get_test_schema_descr().column(0);
2023 let stats = vec![
2024 PageEncodingStats {
2025 page_type: PageType::DATA_PAGE,
2026 encoding: Encoding::PLAIN,
2027 count: 3,
2028 },
2029 PageEncodingStats {
2030 page_type: PageType::DATA_PAGE,
2031 encoding: Encoding::RLE,
2032 count: 5,
2033 },
2034 ];
2035 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
2036 .set_encodings_mask(EncodingMask::new_from_encodings(
2037 [Encoding::PLAIN, Encoding::RLE].iter(),
2038 ))
2039 .set_num_values(1000)
2040 .set_compression_codec(CompressionCodec::SNAPPY)
2041 .set_total_compressed_size(2000)
2042 .set_total_uncompressed_size(3000)
2043 .set_data_page_offset(4000)
2044 .set_page_encoding_stats(stats)
2045 .build()
2046 .unwrap();
2047
2048 let mut buf = Vec::new();
2049 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2050 col_metadata.write_thrift(&mut writer).unwrap();
2051
2052 let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
2053 let col_chunk_res =
2054 read_column_chunk_with_options(&buf, column_descr, Some(&options)).unwrap();
2055
2056 assert_eq!(col_chunk_res, col_metadata);
2057 }
2058
2059 #[test]
2060 fn test_column_chunk_metadata_thrift_conversion_empty() {
2061 let column_descr = get_test_schema_descr().column(0);
2062
2063 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
2064 .build()
2065 .unwrap();
2066
2067 let mut buf = Vec::new();
2068 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
2069 col_metadata.write_thrift(&mut writer).unwrap();
2070 let col_chunk_res = read_column_chunk(&buf, column_descr).unwrap();
2071
2072 assert_eq!(col_chunk_res, col_metadata);
2073 }
2074
2075 #[test]
2076 fn test_compressed_size() {
2077 let schema_descr = get_test_schema_descr();
2078
2079 let mut columns = vec![];
2080 for column_descr in schema_descr.columns() {
2081 let column = ColumnChunkMetaData::builder(column_descr.clone())
2082 .set_total_compressed_size(500)
2083 .set_total_uncompressed_size(700)
2084 .build()
2085 .unwrap();
2086 columns.push(column);
2087 }
2088 let row_group_meta = RowGroupMetaData::builder(schema_descr)
2089 .set_num_rows(1000)
2090 .set_column_metadata(columns)
2091 .build()
2092 .unwrap();
2093
2094 let compressed_size_res = row_group_meta.compressed_size();
2095 let compressed_size_exp: i64 = 1000;
2096
2097 assert_eq!(compressed_size_res, compressed_size_exp);
2098 }
2099
2100 #[test]
2101 fn test_memory_size() {
2102 let schema_descr = get_test_schema_descr();
2103
2104 let columns = schema_descr
2105 .columns()
2106 .iter()
2107 .map(|column_descr| {
2108 ColumnChunkMetaData::builder(column_descr.clone())
2109 .set_statistics(Statistics::new::<i32>(None, None, None, None, false))
2110 .build()
2111 })
2112 .collect::<Result<Vec<_>>>()
2113 .unwrap();
2114 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2115 .set_num_rows(1000)
2116 .set_column_metadata(columns)
2117 .build()
2118 .unwrap();
2119 let row_group_meta = vec![row_group_meta];
2120
2121 let version = 2;
2122 let num_rows = 1000;
2123 let created_by = Some(String::from("test harness"));
2124 let key_value_metadata = Some(vec![KeyValue::new(
2125 String::from("Foo"),
2126 Some(String::from("bar")),
2127 )]);
2128 let column_orders = Some(vec![
2129 ColumnOrder::UNDEFINED,
2130 ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED),
2131 ]);
2132 let file_metadata = FileMetaData::new(
2133 version,
2134 num_rows,
2135 created_by,
2136 key_value_metadata,
2137 schema_descr.clone(),
2138 column_orders,
2139 );
2140
2141 let columns_with_stats = schema_descr
2143 .columns()
2144 .iter()
2145 .map(|column_descr| {
2146 ColumnChunkMetaData::builder(column_descr.clone())
2147 .set_statistics(Statistics::new::<i32>(
2148 Some(0),
2149 Some(100),
2150 None,
2151 None,
2152 false,
2153 ))
2154 .build()
2155 })
2156 .collect::<Result<Vec<_>>>()
2157 .unwrap();
2158
2159 let row_group_meta_with_stats = RowGroupMetaData::builder(schema_descr)
2160 .set_num_rows(1000)
2161 .set_column_metadata(columns_with_stats)
2162 .build()
2163 .unwrap();
2164 let row_group_meta_with_stats = vec![row_group_meta_with_stats];
2165
2166 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata.clone())
2167 .set_row_groups(row_group_meta_with_stats)
2168 .build();
2169
2170 #[cfg(not(feature = "encryption"))]
2172 let base_expected_size = 2798;
2173 #[cfg(feature = "encryption")]
2174 let base_expected_size = 2966;
2175
2176 assert_eq!(parquet_meta.memory_size(), base_expected_size);
2177
2178 let mut page_index = PageIndexBuilder::new(1, 1);
2179
2180 let mut column_index = ColumnIndexBuilder::new(Type::BOOLEAN);
2181 column_index.append(false, vec![1u8], vec![2u8, 3u8], 4, None);
2182 let column_index = column_index.build().unwrap();
2183 {
2184 let ColumnIndexMetaData::BOOLEAN(_) = column_index else {
2185 panic!("wrong type of column index")
2186 };
2187 }
2188
2189 page_index.put_column_index(column_index, 0, 0);
2190
2191 let mut offset_index = OffsetIndexBuilder::new();
2193 offset_index.append_row_count(1);
2194 offset_index.append_offset_and_size(2, 3);
2195 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2196 offset_index.append_row_count(1);
2197 offset_index.append_offset_and_size(2, 3);
2198 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2199 let offset_index = offset_index.build();
2200 page_index.put_offset_index(offset_index, 0, 0);
2201
2202 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata)
2203 .set_row_groups(row_group_meta)
2204 .set_page_index(Some(Arc::new(page_index.build())))
2205 .build();
2206
2207 #[cfg(not(feature = "encryption"))]
2209 let bigger_expected_size = 3312;
2210 #[cfg(feature = "encryption")]
2211 let bigger_expected_size = 3480;
2212
2213 assert!(bigger_expected_size > base_expected_size);
2215 assert_eq!(parquet_meta.memory_size(), bigger_expected_size);
2216 }
2217
2218 #[test]
2219 #[cfg(feature = "encryption")]
2220 fn test_memory_size_with_decryptor() {
2221 use crate::encryption::decrypt::FileDecryptionProperties;
2222 use crate::file::metadata::thrift::encryption::AesGcmV1;
2223
2224 let schema_descr = get_test_schema_descr();
2225
2226 let columns = schema_descr
2227 .columns()
2228 .iter()
2229 .map(|column_descr| ColumnChunkMetaData::builder(column_descr.clone()).build())
2230 .collect::<Result<Vec<_>>>()
2231 .unwrap();
2232 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2233 .set_num_rows(1000)
2234 .set_column_metadata(columns)
2235 .build()
2236 .unwrap();
2237 let row_group_meta = vec![row_group_meta];
2238
2239 let version = 2;
2240 let num_rows = 1000;
2241 let aad_file_unique = vec![1u8; 8];
2242 let aad_prefix = vec![2u8; 8];
2243 let encryption_algorithm = EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 {
2244 aad_prefix: Some(aad_prefix.clone()),
2245 aad_file_unique: Some(aad_file_unique.clone()),
2246 supply_aad_prefix: Some(true),
2247 });
2248 let footer_key_metadata = Some(vec![3u8; 8]);
2249 let file_metadata =
2250 FileMetaData::new(version, num_rows, None, None, schema_descr.clone(), None)
2251 .with_encryption_algorithm(Some(encryption_algorithm))
2252 .with_footer_signing_key_metadata(footer_key_metadata.clone());
2253
2254 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2255 .set_row_groups(row_group_meta.clone())
2256 .build();
2257
2258 let base_expected_size = 2074;
2259 assert_eq!(parquet_meta_data.memory_size(), base_expected_size);
2260
2261 let footer_key = b"0123456789012345";
2262 let column_key = b"1234567890123450";
2263 let mut decryption_properties_builder =
2264 FileDecryptionProperties::builder(footer_key.to_vec())
2265 .with_aad_prefix(aad_prefix.clone());
2266 for column in schema_descr.columns() {
2267 decryption_properties_builder = decryption_properties_builder
2268 .with_column_key(&column.path().string(), column_key.to_vec());
2269 }
2270 let decryption_properties = decryption_properties_builder.build().unwrap();
2271 let decryptor = FileDecryptor::new(
2272 &decryption_properties,
2273 footer_key_metadata.as_deref(),
2274 aad_file_unique,
2275 aad_prefix,
2276 )
2277 .unwrap();
2278
2279 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2280 .set_row_groups(row_group_meta.clone())
2281 .set_file_decryptor(Some(decryptor))
2282 .build();
2283
2284 let expected_size_with_decryptor = 3104;
2285 assert!(expected_size_with_decryptor > base_expected_size);
2286
2287 assert_eq!(
2288 parquet_meta_data.memory_size(),
2289 expected_size_with_decryptor
2290 );
2291 }
2292
2293 fn get_test_schema_descr() -> SchemaDescPtr {
2295 let schema = SchemaType::group_type_builder("schema")
2296 .with_fields(vec![
2297 Arc::new(
2298 SchemaType::primitive_type_builder("a", Type::INT32)
2299 .build()
2300 .unwrap(),
2301 ),
2302 Arc::new(
2303 SchemaType::primitive_type_builder("b", Type::INT32)
2304 .build()
2305 .unwrap(),
2306 ),
2307 ])
2308 .build()
2309 .unwrap();
2310
2311 Arc::new(SchemaDescriptor::new(Arc::new(schema)))
2312 }
2313
2314 #[test]
2315 fn test_page_index_builder_partial_population() {
2316 let mut builder = PageIndexBuilder::new(1, 1);
2317
2318 let mut offset_idx_builder = OffsetIndexBuilder::new();
2320 offset_idx_builder.append_row_count(50);
2321 offset_idx_builder.append_offset_and_size(1000, 500);
2322 let offset_idx = offset_idx_builder.build();
2323 builder.put_offset_index(offset_idx, 0, 0);
2324
2325 let page_index = builder.build();
2326
2327 assert!(!page_index.has_column_indexes());
2329 assert!(page_index.has_offset_indexes());
2331 }
2332
2333 #[test]
2334 fn test_page_index_builder_empty_to_none() {
2335 let builder = PageIndexBuilder::new(2, 2);
2337
2338 let page_index = builder.build();
2339
2340 assert!(!page_index.has_column_indexes());
2342 assert!(!page_index.has_offset_indexes());
2343 }
2344
2345 #[test]
2346 fn test_replace_page_index_preserves_shared_metadata() {
2347 let schema_descr = get_test_schema_descr();
2348 let file_metadata = FileMetaData::new(1, 0, None, None, schema_descr, None);
2349 let cached = Arc::new(ParquetMetaData::new(file_metadata, vec![]));
2350
2351 let page_index = Arc::new(PageIndexBuilder::new(0, 0).build());
2352 let rebuilt = cached
2353 .as_ref()
2354 .clone()
2355 .into_builder()
2356 .set_page_index(Some(page_index))
2357 .build();
2358
2359 assert!(Arc::ptr_eq(&cached.file_metadata, &rebuilt.file_metadata));
2360 assert!(Arc::ptr_eq(&cached.row_groups, &rebuilt.row_groups));
2361 }
2362
2363 #[test]
2364 fn test_rebuild_page_index() {
2365 let mut builder = PageIndexBuilder::new(1, 2);
2367
2368 let mut col_idx_builder = ColumnIndexBuilder::new(Type::INT32);
2370 col_idx_builder.append(false, vec![1, 0, 0, 0], vec![100, 0, 0, 0], 10, None);
2371 let col_idx = col_idx_builder.build().unwrap();
2372 builder.put_column_index(col_idx.clone(), 0, 0);
2373
2374 let mut offset_idx_builder = OffsetIndexBuilder::new();
2376 offset_idx_builder.append_row_count(50);
2377 offset_idx_builder.append_offset_and_size(1000, 500);
2378 let offset_idx = offset_idx_builder.build();
2379 builder.put_offset_index(offset_idx.clone(), 0, 0);
2380
2381 let page_index = builder.build();
2382
2383 assert!(page_index.is_complete());
2385
2386 assert!(page_index.column_index(0, 0).is_some());
2388 assert!(page_index.offset_index(0, 0).is_some());
2389 assert!(page_index.column_index(0, 1).is_none());
2391 assert!(page_index.offset_index(0, 1).is_none());
2392
2393 let mut builder = PageIndexBuilder::new_from(page_index);
2394
2395 builder.put_column_index(col_idx, 0, 1);
2397 builder.put_offset_index(offset_idx, 0, 1);
2398 let page_index = builder.build();
2399
2400 assert!(page_index.column_index(0, 0).is_some());
2402 assert!(page_index.offset_index(0, 0).is_some());
2403 assert!(page_index.column_index(0, 1).is_some());
2404 assert!(page_index.offset_index(0, 1).is_some());
2405 }
2406}