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)]
123pub struct ParquetMetaData {
124 file_metadata: FileMetaData,
126 row_groups: Vec<RowGroupMetaData>,
128 page_index: Option<Arc<dyn PageIndexProvider>>,
130 #[cfg(feature = "encryption")]
132 file_decryptor: Option<Box<FileDecryptor>>,
133}
134
135impl ParquetMetaData {
136 pub fn new(file_metadata: FileMetaData, row_groups: Vec<RowGroupMetaData>) -> Self {
139 ParquetMetaData {
140 file_metadata,
141 row_groups,
142 page_index: None,
143 #[cfg(feature = "encryption")]
144 file_decryptor: None,
145 }
146 }
147
148 #[cfg(feature = "encryption")]
151 pub(crate) fn with_file_decryptor(&mut self, file_decryptor: Option<FileDecryptor>) {
152 self.file_decryptor = file_decryptor.map(Box::new);
153 }
154
155 pub fn into_builder(self) -> ParquetMetaDataBuilder {
157 self.into()
158 }
159
160 pub fn file_metadata(&self) -> &FileMetaData {
162 &self.file_metadata
163 }
164
165 #[cfg(feature = "encryption")]
167 pub(crate) fn file_decryptor(&self) -> Option<&FileDecryptor> {
168 self.file_decryptor.as_deref()
169 }
170
171 pub fn num_row_groups(&self) -> usize {
173 self.row_groups.len()
174 }
175
176 pub fn row_group(&self, i: usize) -> &RowGroupMetaData {
179 &self.row_groups[i]
180 }
181
182 pub fn row_groups(&self) -> &[RowGroupMetaData] {
184 &self.row_groups
185 }
186
187 pub fn row_group_num_rows(&self, row_group_idx: usize) -> Result<usize> {
192 self.row_groups
193 .get(row_group_idx)
194 .ok_or_else(|| {
195 ParquetError::General(format!(
196 "Row group index {row_group_idx} out of bounds for file with {} row groups",
197 self.num_row_groups()
198 ))
199 })?
200 .num_rows()
201 .try_into()
202 .map_err(|e| ParquetError::General(format!("Row count overflow: {e}")))
203 }
204
205 pub fn page_index(&self) -> Option<&Arc<dyn PageIndexProvider>> {
212 self.page_index.as_ref()
213 }
214
215 pub fn page_index_for_row_group(&self, row_group_idx: usize) -> RowGroupPageIndex {
217 RowGroupPageIndex::new(row_group_idx, self.page_index.clone())
218 }
219
220 pub fn memory_size(&self) -> usize {
235 #[cfg(feature = "encryption")]
236 let encryption_size = self.file_decryptor.heap_size();
237 #[cfg(not(feature = "encryption"))]
238 let encryption_size = 0usize;
239
240 let page_index_size = if let Some(page_index) = self.page_index.as_ref() {
245 if let Some(page_index) = page_index.as_any().downcast_ref::<PageIndex>() {
246 2 * std::mem::size_of::<usize>() + std::mem::size_of::<PageIndex>() + page_index.heap_size()
250 } else {
251 0
252 }
253 } else {
254 0
255 };
256
257 std::mem::size_of::<Self>()
258 + self.file_metadata.heap_size()
259 + self.row_groups.heap_size()
260 + page_index_size
261 + encryption_size
262 }
263
264 pub(crate) fn set_page_index(&mut self, index: Option<Arc<dyn PageIndexProvider>>) {
266 self.page_index = index;
267 }
268}
269
270impl PartialEq for ParquetMetaData {
271 fn eq(&self, other: &Self) -> bool {
272 if self.file_metadata != other.file_metadata || self.row_groups != other.row_groups {
274 return false;
275 }
276
277 #[cfg(feature = "encryption")]
278 if self.file_decryptor != other.file_decryptor {
279 return false;
280 }
281
282 match (&self.page_index, &other.page_index) {
284 (None, None) => true,
285 (Some(a), Some(b)) => {
286 let a_page_index = a.as_any().downcast_ref::<PageIndex>();
288 let b_page_index = b.as_any().downcast_ref::<PageIndex>();
289 match (a_page_index, b_page_index) {
290 (Some(a_idx), Some(b_idx)) => a_idx == b_idx,
291 _ => Arc::ptr_eq(a, b), }
293 }
294 _ => false,
295 }
296 }
297}
298
299pub struct ParquetMetaDataBuilder(ParquetMetaData);
337
338impl ParquetMetaDataBuilder {
339 pub fn new(file_meta_data: FileMetaData) -> Self {
341 Self(ParquetMetaData::new(file_meta_data, vec![]))
342 }
343
344 pub fn new_from_metadata(metadata: ParquetMetaData) -> Self {
346 Self(metadata)
347 }
348
349 pub fn add_row_group(mut self, row_group: RowGroupMetaData) -> Self {
351 self.0.row_groups.push(row_group);
352 self
353 }
354
355 pub fn set_row_groups(mut self, row_groups: Vec<RowGroupMetaData>) -> Self {
357 self.0.row_groups = row_groups;
358 self
359 }
360
361 pub fn take_row_groups(&mut self) -> Vec<RowGroupMetaData> {
367 std::mem::take(&mut self.0.row_groups)
368 }
369
370 pub fn row_groups(&self) -> &[RowGroupMetaData] {
372 &self.0.row_groups
373 }
374
375 pub fn set_page_index(mut self, page_index: Option<Arc<dyn PageIndexProvider>>) -> Self {
384 self.0.page_index = page_index;
385 self
386 }
387
388 pub fn take_page_index(&mut self) -> Option<Arc<dyn PageIndexProvider>> {
390 std::mem::take(&mut self.0.page_index)
391 }
392
393 pub fn page_index(&self) -> Option<&dyn PageIndexProvider> {
395 self.0.page_index.as_ref().map(|arc| arc.as_ref())
396 }
397
398 #[cfg(feature = "encryption")]
400 pub(crate) fn set_file_decryptor(mut self, file_decryptor: Option<FileDecryptor>) -> Self {
401 self.0.with_file_decryptor(file_decryptor);
402 self
403 }
404
405 pub fn build(self) -> ParquetMetaData {
407 let Self(metadata) = self;
408 metadata
409 }
410}
411
412impl From<ParquetMetaData> for ParquetMetaDataBuilder {
413 fn from(meta_data: ParquetMetaData) -> Self {
414 Self(meta_data)
415 }
416}
417
418thrift_struct!(
419pub struct KeyValue {
421 1: required string key
422 2: optional string value
423}
424);
425
426impl KeyValue {
427 pub fn new<F2>(key: String, value: F2) -> KeyValue
429 where
430 F2: Into<Option<String>>,
431 {
432 KeyValue {
433 key,
434 value: value.into(),
435 }
436 }
437}
438
439thrift_struct!(
440pub struct PageEncodingStats {
442 1: required PageType page_type;
443 2: required Encoding encoding;
444 3: required i32 count;
445}
446);
447
448#[derive(Debug, Clone, PartialEq)]
451enum ParquetPageEncodingStats {
452 Full(Vec<PageEncodingStats>),
454 Mask(EncodingMask),
456}
457
458pub type FileMetaDataPtr = Arc<FileMetaData>;
460
461#[derive(Debug, Clone, PartialEq)]
465pub struct FileMetaData {
466 version: i32,
467 num_rows: i64,
468 created_by: Option<String>,
469 key_value_metadata: Option<Vec<KeyValue>>,
470 schema_descr: SchemaDescPtr,
471 column_orders: Option<Vec<ColumnOrder>>,
472 #[cfg(feature = "encryption")]
473 encryption_algorithm: Option<Box<EncryptionAlgorithm>>,
474 #[cfg(feature = "encryption")]
475 footer_signing_key_metadata: Option<Vec<u8>>,
476}
477
478impl FileMetaData {
479 pub fn new(
481 version: i32,
482 num_rows: i64,
483 created_by: Option<String>,
484 key_value_metadata: Option<Vec<KeyValue>>,
485 schema_descr: SchemaDescPtr,
486 column_orders: Option<Vec<ColumnOrder>>,
487 ) -> Self {
488 FileMetaData {
489 version,
490 num_rows,
491 created_by,
492 key_value_metadata,
493 schema_descr,
494 column_orders,
495 #[cfg(feature = "encryption")]
496 encryption_algorithm: None,
497 #[cfg(feature = "encryption")]
498 footer_signing_key_metadata: None,
499 }
500 }
501
502 #[cfg(feature = "encryption")]
503 pub(crate) fn with_encryption_algorithm(
504 mut self,
505 encryption_algorithm: Option<EncryptionAlgorithm>,
506 ) -> Self {
507 self.encryption_algorithm = encryption_algorithm.map(Box::new);
508 self
509 }
510
511 #[cfg(feature = "encryption")]
512 pub(crate) fn with_footer_signing_key_metadata(
513 mut self,
514 footer_signing_key_metadata: Option<Vec<u8>>,
515 ) -> Self {
516 self.footer_signing_key_metadata = footer_signing_key_metadata;
517 self
518 }
519
520 pub fn version(&self) -> i32 {
522 self.version
523 }
524
525 pub fn num_rows(&self) -> i64 {
527 self.num_rows
528 }
529
530 pub fn created_by(&self) -> Option<&str> {
539 self.created_by.as_deref()
540 }
541
542 pub fn key_value_metadata(&self) -> Option<&Vec<KeyValue>> {
544 self.key_value_metadata.as_ref()
545 }
546
547 pub fn schema(&self) -> &SchemaType {
551 self.schema_descr.root_schema()
552 }
553
554 pub fn schema_descr(&self) -> &SchemaDescriptor {
556 &self.schema_descr
557 }
558
559 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
561 self.schema_descr.clone()
562 }
563
564 pub fn column_orders(&self) -> Option<&Vec<ColumnOrder>> {
572 self.column_orders.as_ref()
573 }
574
575 pub fn column_order(&self, i: usize) -> ColumnOrder {
578 self.column_orders
579 .as_ref()
580 .map(|data| data[i])
581 .unwrap_or(ColumnOrder::UNDEFINED)
582 }
583}
584
585thrift_struct!(
586pub struct SortingColumn {
588 1: required i32 column_idx
590
591 2: required bool descending
593
594 3: required bool nulls_first
597}
598);
599
600pub type RowGroupMetaDataPtr = Arc<RowGroupMetaData>;
602
603#[derive(Debug, Clone, PartialEq)]
608pub struct RowGroupMetaData {
609 columns: Vec<ColumnChunkMetaData>,
610 num_rows: i64,
611 sorting_columns: Option<Vec<SortingColumn>>,
612 total_byte_size: i64,
613 schema_descr: SchemaDescPtr,
614 file_offset: Option<i64>,
616 ordinal: Option<i32>,
618}
619
620impl RowGroupMetaData {
621 pub fn builder(schema_descr: SchemaDescPtr) -> RowGroupMetaDataBuilder {
623 RowGroupMetaDataBuilder::new(schema_descr)
624 }
625
626 pub fn num_columns(&self) -> usize {
628 self.columns.len()
629 }
630
631 pub fn column(&self, i: usize) -> &ColumnChunkMetaData {
633 &self.columns[i]
634 }
635
636 pub fn columns(&self) -> &[ColumnChunkMetaData] {
638 &self.columns
639 }
640
641 pub fn columns_mut(&mut self) -> &mut [ColumnChunkMetaData] {
643 &mut self.columns
644 }
645
646 pub fn num_rows(&self) -> i64 {
648 self.num_rows
649 }
650
651 pub fn sorting_columns(&self) -> Option<&Vec<SortingColumn>> {
653 self.sorting_columns.as_ref()
654 }
655
656 pub fn total_byte_size(&self) -> i64 {
658 self.total_byte_size
659 }
660
661 pub fn compressed_size(&self) -> i64 {
663 self.columns.iter().map(|c| c.total_compressed_size).sum()
664 }
665
666 pub fn schema_descr(&self) -> &SchemaDescriptor {
668 self.schema_descr.as_ref()
669 }
670
671 pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
673 self.schema_descr.clone()
674 }
675
676 #[inline(always)]
681 pub fn ordinal(&self) -> Option<i32> {
682 self.ordinal
683 }
684
685 #[inline(always)]
687 pub fn file_offset(&self) -> Option<i64> {
688 self.file_offset
689 }
690
691 pub fn into_builder(self) -> RowGroupMetaDataBuilder {
693 RowGroupMetaDataBuilder(self)
694 }
695}
696
697pub struct RowGroupMetaDataBuilder(RowGroupMetaData);
699
700impl RowGroupMetaDataBuilder {
701 fn new(schema_descr: SchemaDescPtr) -> Self {
703 Self(RowGroupMetaData {
704 columns: Vec::with_capacity(schema_descr.num_columns()),
705 schema_descr,
706 file_offset: None,
707 num_rows: 0,
708 sorting_columns: None,
709 total_byte_size: 0,
710 ordinal: None,
711 })
712 }
713
714 pub fn set_num_rows(mut self, value: i64) -> Self {
716 self.0.num_rows = value;
717 self
718 }
719
720 pub fn set_sorting_columns(mut self, value: Option<Vec<SortingColumn>>) -> Self {
722 self.0.sorting_columns = value;
723 self
724 }
725
726 pub fn set_total_byte_size(mut self, value: i64) -> Self {
728 self.0.total_byte_size = value;
729 self
730 }
731
732 pub fn take_columns(&mut self) -> Vec<ColumnChunkMetaData> {
738 std::mem::take(&mut self.0.columns)
739 }
740
741 pub fn set_column_metadata(mut self, value: Vec<ColumnChunkMetaData>) -> Self {
743 self.0.columns = value;
744 self
745 }
746
747 pub fn add_column_metadata(mut self, value: ColumnChunkMetaData) -> Self {
749 self.0.columns.push(value);
750 self
751 }
752
753 pub fn set_ordinal(mut self, value: i32) -> Self {
755 self.0.ordinal = Some(value);
756 self
757 }
758
759 pub fn set_file_offset(mut self, value: i64) -> Self {
761 self.0.file_offset = Some(value);
762 self
763 }
764
765 pub fn build(self) -> Result<RowGroupMetaData> {
767 if self.0.schema_descr.num_columns() != self.0.columns.len() {
768 return Err(general_err!(
769 "Column length mismatch: {} != {}",
770 self.0.schema_descr.num_columns(),
771 self.0.columns.len()
772 ));
773 }
774
775 Ok(self.0)
776 }
777
778 pub(super) fn build_unchecked(self) -> RowGroupMetaData {
780 self.0
781 }
782}
783
784#[derive(Debug, Clone, PartialEq)]
786pub struct ColumnChunkMetaData {
787 column_descr: ColumnDescPtr,
788 encodings: EncodingMask,
789 file_path: Option<String>,
790 file_offset: i64,
791 num_values: i64,
792 compression: CompressionCodec,
793 total_compressed_size: i64,
794 total_uncompressed_size: i64,
795 data_page_offset: i64,
796 index_page_offset: Option<i64>,
797 dictionary_page_offset: Option<i64>,
798 statistics: Option<Statistics>,
799 geo_statistics: Option<Box<geo_statistics::GeospatialStatistics>>,
800 encoding_stats: Option<ParquetPageEncodingStats>,
801 bloom_filter_offset: Option<i64>,
802 bloom_filter_length: Option<i32>,
803 offset_index_offset: Option<i64>,
804 offset_index_length: Option<i32>,
805 column_index_offset: Option<i64>,
806 column_index_length: Option<i32>,
807 unencoded_byte_array_data_bytes: Option<i64>,
808 repetition_level_histogram: Option<LevelHistogram>,
809 definition_level_histogram: Option<LevelHistogram>,
810 #[cfg(feature = "encryption")]
811 column_crypto_metadata: Option<Box<ColumnCryptoMetaData>>,
812 #[cfg(feature = "encryption")]
813 encrypted_column_metadata: Option<Vec<u8>>,
814 #[cfg(feature = "encryption")]
818 plaintext_footer_mode: bool,
819}
820
821#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
830pub struct LevelHistogram {
831 inner: Vec<i64>,
832}
833
834impl LevelHistogram {
835 pub fn try_new(max_level: i16) -> Option<Self> {
841 if max_level > 0 {
842 Some(Self {
843 inner: vec![0; max_level as usize + 1],
844 })
845 } else {
846 None
847 }
848 }
849 pub fn values(&self) -> &[i64] {
851 &self.inner
852 }
853
854 pub fn into_inner(self) -> Vec<i64> {
856 self.inner
857 }
858
859 pub fn get(&self, index: usize) -> Option<i64> {
866 self.inner.get(index).copied()
867 }
868
869 pub fn add(&mut self, other: &Self) {
874 assert_eq!(self.len(), other.len());
875 for (dst, src) in self.inner.iter_mut().zip(other.inner.iter()) {
876 *dst += src;
877 }
878 }
879
880 pub fn len(&self) -> usize {
882 self.inner.len()
883 }
884
885 pub fn is_empty(&self) -> bool {
887 self.inner.is_empty()
888 }
889
890 pub fn reset(&mut self) {
892 self.inner.fill(0);
893 }
894
895 #[inline]
897 pub fn increment_by(&mut self, level: i16, count: i64) {
898 self.inner[level as usize] += count;
899 }
900
901 #[deprecated(since = "58.2.0", note = "Use `increment_by` instead")]
907 pub fn update_from_levels(&mut self, levels: &[i16]) {
908 for &level in levels {
909 self.increment_by(level, 1);
910 }
911 }
912}
913
914impl From<Vec<i64>> for LevelHistogram {
915 fn from(inner: Vec<i64>) -> Self {
916 Self { inner }
917 }
918}
919
920impl From<LevelHistogram> for Vec<i64> {
921 fn from(value: LevelHistogram) -> Self {
922 value.into_inner()
923 }
924}
925
926impl HeapSize for LevelHistogram {
927 fn heap_size(&self) -> usize {
928 self.inner.heap_size()
929 }
930}
931
932impl ColumnChunkMetaData {
934 pub fn builder(column_descr: ColumnDescPtr) -> ColumnChunkMetaDataBuilder {
936 ColumnChunkMetaDataBuilder::new(column_descr)
937 }
938
939 pub fn file_path(&self) -> Option<&str> {
944 self.file_path.as_deref()
945 }
946
947 pub fn file_offset(&self) -> i64 {
954 self.file_offset
955 }
956
957 pub fn column_type(&self) -> Type {
959 self.column_descr.physical_type()
960 }
961
962 pub fn column_path(&self) -> &ColumnPath {
964 self.column_descr.path()
965 }
966
967 pub fn column_descr(&self) -> &ColumnDescriptor {
969 self.column_descr.as_ref()
970 }
971
972 pub fn column_descr_ptr(&self) -> ColumnDescPtr {
974 self.column_descr.clone()
975 }
976
977 pub fn encodings(&self) -> impl Iterator<Item = Encoding> {
979 self.encodings.encodings()
980 }
981
982 pub fn encodings_mask(&self) -> &EncodingMask {
984 &self.encodings
985 }
986
987 pub fn num_values(&self) -> i64 {
989 self.num_values
990 }
991
992 pub fn compression(&self) -> Compression {
999 self.compression.into()
1000 }
1001
1002 pub fn compression_codec(&self) -> CompressionCodec {
1004 self.compression
1005 }
1006
1007 pub fn compressed_size(&self) -> i64 {
1009 self.total_compressed_size
1010 }
1011
1012 pub fn uncompressed_size(&self) -> i64 {
1014 self.total_uncompressed_size
1015 }
1016
1017 pub fn data_page_offset(&self) -> i64 {
1019 self.data_page_offset
1020 }
1021
1022 pub fn index_page_offset(&self) -> Option<i64> {
1024 self.index_page_offset
1025 }
1026
1027 pub fn dictionary_page_offset(&self) -> Option<i64> {
1029 self.dictionary_page_offset
1030 }
1031
1032 pub fn byte_range(&self) -> (u64, u64) {
1038 let col_start = match self.dictionary_page_offset() {
1039 Some(dictionary_page_offset) => dictionary_page_offset,
1040 None => self.data_page_offset(),
1041 };
1042 let col_len = self.compressed_size();
1043 assert!(
1044 col_start >= 0 && col_len >= 0,
1045 "column start and length should not be negative"
1046 );
1047 (col_start as u64, col_len as u64)
1048 }
1049
1050 pub fn statistics(&self) -> Option<&Statistics> {
1053 self.statistics.as_ref()
1054 }
1055
1056 pub fn geo_statistics(&self) -> Option<&geo_statistics::GeospatialStatistics> {
1059 self.geo_statistics.as_deref()
1060 }
1061
1062 pub fn page_encoding_stats(&self) -> Option<&Vec<PageEncodingStats>> {
1069 match self.encoding_stats.as_ref() {
1070 Some(ParquetPageEncodingStats::Full(stats)) => Some(stats),
1071 _ => None,
1072 }
1073 }
1074
1075 pub fn page_encoding_stats_mask(&self) -> Option<&EncodingMask> {
1105 match self.encoding_stats.as_ref() {
1106 Some(ParquetPageEncodingStats::Mask(stats)) => Some(stats),
1107 _ => None,
1108 }
1109 }
1110
1111 pub fn bloom_filter_offset(&self) -> Option<i64> {
1113 self.bloom_filter_offset
1114 }
1115
1116 pub fn bloom_filter_length(&self) -> Option<i32> {
1118 self.bloom_filter_length
1119 }
1120
1121 pub fn column_index_offset(&self) -> Option<i64> {
1123 self.column_index_offset
1124 }
1125
1126 pub fn column_index_length(&self) -> Option<i32> {
1128 self.column_index_length
1129 }
1130
1131 pub fn column_index_range(&self) -> Option<Range<u64>> {
1137 let offset = u64::try_from(self.column_index_offset?).ok()?;
1138 let length = u64::try_from(self.column_index_length?).ok()?;
1139 Some(offset..(offset + length))
1140 }
1141
1142 pub fn offset_index_offset(&self) -> Option<i64> {
1144 self.offset_index_offset
1145 }
1146
1147 pub fn offset_index_length(&self) -> Option<i32> {
1149 self.offset_index_length
1150 }
1151
1152 pub fn offset_index_range(&self) -> Option<Range<u64>> {
1158 let offset = u64::try_from(self.offset_index_offset?).ok()?;
1159 let length = u64::try_from(self.offset_index_length?).ok()?;
1160 Some(offset..(offset + length))
1161 }
1162
1163 pub fn unencoded_byte_array_data_bytes(&self) -> Option<i64> {
1168 self.unencoded_byte_array_data_bytes
1169 }
1170
1171 pub fn repetition_level_histogram(&self) -> Option<&LevelHistogram> {
1177 self.repetition_level_histogram.as_ref()
1178 }
1179
1180 pub fn definition_level_histogram(&self) -> Option<&LevelHistogram> {
1186 self.definition_level_histogram.as_ref()
1187 }
1188
1189 #[cfg(feature = "encryption")]
1191 pub fn crypto_metadata(&self) -> Option<&ColumnCryptoMetaData> {
1192 self.column_crypto_metadata.as_deref()
1193 }
1194
1195 pub fn into_builder(self) -> ColumnChunkMetaDataBuilder {
1197 ColumnChunkMetaDataBuilder::from(self)
1198 }
1199}
1200
1201pub struct ColumnChunkMetaDataBuilder(ColumnChunkMetaData);
1220
1221impl ColumnChunkMetaDataBuilder {
1222 fn new(column_descr: ColumnDescPtr) -> Self {
1226 Self(ColumnChunkMetaData {
1227 column_descr,
1228 encodings: Default::default(),
1229 file_path: None,
1230 file_offset: 0,
1231 num_values: 0,
1232 compression: CompressionCodec::UNCOMPRESSED,
1233 total_compressed_size: 0,
1234 total_uncompressed_size: 0,
1235 data_page_offset: 0,
1236 index_page_offset: None,
1237 dictionary_page_offset: None,
1238 statistics: None,
1239 geo_statistics: None,
1240 encoding_stats: None,
1241 bloom_filter_offset: None,
1242 bloom_filter_length: None,
1243 offset_index_offset: None,
1244 offset_index_length: None,
1245 column_index_offset: None,
1246 column_index_length: None,
1247 unencoded_byte_array_data_bytes: None,
1248 repetition_level_histogram: None,
1249 definition_level_histogram: None,
1250 #[cfg(feature = "encryption")]
1251 column_crypto_metadata: None,
1252 #[cfg(feature = "encryption")]
1253 encrypted_column_metadata: None,
1254 #[cfg(feature = "encryption")]
1255 plaintext_footer_mode: false,
1256 })
1257 }
1258
1259 pub fn set_encodings(mut self, encodings: Vec<Encoding>) -> Self {
1261 self.0.encodings = EncodingMask::new_from_encodings(encodings.iter());
1262 self
1263 }
1264
1265 pub fn set_encodings_mask(mut self, encodings: EncodingMask) -> Self {
1267 self.0.encodings = encodings;
1268 self
1269 }
1270
1271 pub fn set_file_path(mut self, value: String) -> Self {
1273 self.0.file_path = Some(value);
1274 self
1275 }
1276
1277 pub fn set_num_values(mut self, value: i64) -> Self {
1279 self.0.num_values = value;
1280 self
1281 }
1282
1283 pub fn set_compression(mut self, value: Compression) -> Self {
1285 self.0.compression = value.into();
1286 self
1287 }
1288
1289 pub fn set_compression_codec(mut self, value: CompressionCodec) -> Self {
1291 self.0.compression = value;
1292 self
1293 }
1294
1295 pub fn set_total_compressed_size(mut self, value: i64) -> Self {
1297 self.0.total_compressed_size = value;
1298 self
1299 }
1300
1301 pub fn set_total_uncompressed_size(mut self, value: i64) -> Self {
1303 self.0.total_uncompressed_size = value;
1304 self
1305 }
1306
1307 pub fn set_data_page_offset(mut self, value: i64) -> Self {
1309 self.0.data_page_offset = value;
1310 self
1311 }
1312
1313 pub fn set_dictionary_page_offset(mut self, value: Option<i64>) -> Self {
1315 self.0.dictionary_page_offset = value;
1316 self
1317 }
1318
1319 pub fn set_index_page_offset(mut self, value: Option<i64>) -> Self {
1321 self.0.index_page_offset = value;
1322 self
1323 }
1324
1325 pub fn set_statistics(mut self, value: Statistics) -> Self {
1327 self.0.statistics = Some(value);
1328 self
1329 }
1330
1331 pub fn set_geo_statistics(mut self, value: Box<geo_statistics::GeospatialStatistics>) -> Self {
1333 self.0.geo_statistics = Some(value);
1334 self
1335 }
1336
1337 pub fn clear_statistics(mut self) -> Self {
1339 self.0.statistics = None;
1340 self
1341 }
1342
1343 pub fn set_page_encoding_stats(mut self, value: Vec<PageEncodingStats>) -> Self {
1347 self.0.encoding_stats = Some(ParquetPageEncodingStats::Full(value));
1348 self
1349 }
1350
1351 pub fn set_page_encoding_stats_mask(mut self, value: EncodingMask) -> Self {
1355 self.0.encoding_stats = Some(ParquetPageEncodingStats::Mask(value));
1356 self
1357 }
1358
1359 pub fn clear_page_encoding_stats(mut self) -> Self {
1361 self.0.encoding_stats = None;
1362 self
1363 }
1364
1365 pub fn set_bloom_filter_offset(mut self, value: Option<i64>) -> Self {
1367 self.0.bloom_filter_offset = value;
1368 self
1369 }
1370
1371 pub fn set_bloom_filter_length(mut self, value: Option<i32>) -> Self {
1373 self.0.bloom_filter_length = value;
1374 self
1375 }
1376
1377 pub fn set_offset_index_offset(mut self, value: Option<i64>) -> Self {
1379 self.0.offset_index_offset = value;
1380 self
1381 }
1382
1383 pub fn set_offset_index_length(mut self, value: Option<i32>) -> Self {
1385 self.0.offset_index_length = value;
1386 self
1387 }
1388
1389 pub fn set_column_index_offset(mut self, value: Option<i64>) -> Self {
1391 self.0.column_index_offset = value;
1392 self
1393 }
1394
1395 pub fn set_column_index_length(mut self, value: Option<i32>) -> Self {
1397 self.0.column_index_length = value;
1398 self
1399 }
1400
1401 pub fn set_unencoded_byte_array_data_bytes(mut self, value: Option<i64>) -> Self {
1403 self.0.unencoded_byte_array_data_bytes = value;
1404 self
1405 }
1406
1407 pub fn set_repetition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1409 self.0.repetition_level_histogram = value;
1410 self
1411 }
1412
1413 pub fn set_definition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1415 self.0.definition_level_histogram = value;
1416 self
1417 }
1418
1419 #[cfg(feature = "encryption")]
1420 pub fn set_column_crypto_metadata(mut self, value: Option<ColumnCryptoMetaData>) -> Self {
1422 self.0.column_crypto_metadata = value.map(Box::new);
1423 self
1424 }
1425
1426 #[cfg(feature = "encryption")]
1427 pub fn set_encrypted_column_metadata(mut self, value: Option<Vec<u8>>) -> Self {
1429 self.0.encrypted_column_metadata = value;
1430 self
1431 }
1432
1433 pub fn build(self) -> Result<ColumnChunkMetaData> {
1435 Ok(self.0)
1436 }
1437}
1438
1439pub struct ColumnIndexBuilder {
1444 column_type: Type,
1445 null_pages: Vec<bool>,
1446 min_values: Vec<Vec<u8>>,
1447 max_values: Vec<Vec<u8>>,
1448 null_counts: Vec<i64>,
1449 nan_counts: Vec<Option<i64>>,
1450 boundary_order: BoundaryOrder,
1451 repetition_level_histograms: Option<Vec<i64>>,
1453 definition_level_histograms: Option<Vec<i64>>,
1455 valid: bool,
1463}
1464
1465impl ColumnIndexBuilder {
1466 pub fn new(column_type: Type) -> Self {
1468 ColumnIndexBuilder {
1469 column_type,
1470 null_pages: Vec::new(),
1471 min_values: Vec::new(),
1472 max_values: Vec::new(),
1473 null_counts: Vec::new(),
1474 nan_counts: Vec::new(),
1475 boundary_order: BoundaryOrder::UNORDERED,
1476 repetition_level_histograms: None,
1477 definition_level_histograms: None,
1478 valid: true,
1479 }
1480 }
1481
1482 pub fn append(
1489 &mut self,
1490 null_page: bool,
1491 min_value: Vec<u8>,
1492 max_value: Vec<u8>,
1493 null_count: i64,
1494 nan_count: Option<i64>,
1495 ) {
1496 self.null_pages.push(null_page);
1497 self.min_values.push(min_value);
1498 self.max_values.push(max_value);
1499 self.null_counts.push(null_count);
1500 self.nan_counts.push(nan_count);
1501 }
1502
1503 pub fn append_histograms(
1508 &mut self,
1509 repetition_level_histogram: &Option<LevelHistogram>,
1510 definition_level_histogram: &Option<LevelHistogram>,
1511 ) {
1512 if !self.valid {
1513 return;
1514 }
1515 if let Some(rep_lvl_hist) = repetition_level_histogram {
1516 let hist = self.repetition_level_histograms.get_or_insert(Vec::new());
1517 hist.reserve(rep_lvl_hist.len());
1518 hist.extend(rep_lvl_hist.values());
1519 }
1520 if let Some(def_lvl_hist) = definition_level_histogram {
1521 let hist = self.definition_level_histograms.get_or_insert(Vec::new());
1522 hist.reserve(def_lvl_hist.len());
1523 hist.extend(def_lvl_hist.values());
1524 }
1525 }
1526
1527 pub fn set_boundary_order(&mut self, boundary_order: BoundaryOrder) {
1529 self.boundary_order = boundary_order;
1530 }
1531
1532 pub fn to_invalid(&mut self) {
1534 self.valid = false;
1535 }
1536
1537 pub fn valid(&self) -> bool {
1539 self.valid
1540 }
1541
1542 pub fn build(self) -> Result<ColumnIndexMetaData> {
1546 Ok(match self.column_type {
1547 Type::BOOLEAN => {
1548 let index = self.build_page_index(false)?;
1549 ColumnIndexMetaData::BOOLEAN(index)
1550 }
1551 Type::INT32 => {
1552 let index = self.build_page_index(false)?;
1553 ColumnIndexMetaData::INT32(index)
1554 }
1555 Type::INT64 => {
1556 let index = self.build_page_index(false)?;
1557 ColumnIndexMetaData::INT64(index)
1558 }
1559 Type::INT96 => {
1560 let index = self.build_page_index(false)?;
1561 ColumnIndexMetaData::INT96(index)
1562 }
1563 Type::FLOAT => {
1564 let index = self.build_page_index(true)?;
1565 ColumnIndexMetaData::FLOAT(index)
1566 }
1567 Type::DOUBLE => {
1568 let index = self.build_page_index(true)?;
1569 ColumnIndexMetaData::DOUBLE(index)
1570 }
1571 Type::BYTE_ARRAY => {
1572 let index = self.build_byte_array_index(false)?;
1573 ColumnIndexMetaData::BYTE_ARRAY(index)
1574 }
1575 Type::FIXED_LEN_BYTE_ARRAY => {
1576 let index = self.build_byte_array_index(true)?;
1577 ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index)
1578 }
1579 })
1580 }
1581
1582 fn build_nan_counts(nan_counts: &[Option<i64>]) -> Option<Vec<i64>> {
1583 let has_some = nan_counts.iter().any(|x| x.is_some());
1584 let has_none = nan_counts.iter().any(|x| x.is_none());
1585
1586 if has_some && !has_none {
1587 Some(nan_counts.iter().map(|x| x.unwrap()).collect())
1588 } else if !has_some && has_none {
1589 None
1590 } else {
1591 debug_assert!(
1592 false,
1593 "Mixed Some/None in nan_counts - caller should provide consistent values"
1594 );
1595 Some(nan_counts.iter().map(|x| x.unwrap_or(0)).collect())
1596 }
1597 }
1598
1599 fn build_page_index<T>(self, may_have_nan: bool) -> Result<PrimitiveColumnIndex<T>>
1600 where
1601 T: ParquetValueType,
1602 {
1603 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1604 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1605
1606 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1611 Self::build_nan_counts(&self.nan_counts)
1612 } else {
1613 None
1614 };
1615
1616 PrimitiveColumnIndex::try_new(
1617 self.null_pages,
1618 self.boundary_order,
1619 Some(self.null_counts),
1620 nan_counts,
1621 self.repetition_level_histograms,
1622 self.definition_level_histograms,
1623 min_values,
1624 max_values,
1625 )
1626 }
1627
1628 fn build_byte_array_index(self, may_have_nan: bool) -> Result<ByteArrayColumnIndex> {
1629 let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1630 let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1631
1632 let nan_counts = if may_have_nan && !self.nan_counts.is_empty() {
1637 Self::build_nan_counts(&self.nan_counts)
1638 } else {
1639 None
1640 };
1641
1642 ByteArrayColumnIndex::try_new(
1643 self.null_pages,
1644 self.boundary_order,
1645 Some(self.null_counts),
1646 nan_counts,
1647 self.repetition_level_histograms,
1648 self.definition_level_histograms,
1649 min_values,
1650 max_values,
1651 )
1652 }
1653}
1654
1655impl From<ColumnChunkMetaData> for ColumnChunkMetaDataBuilder {
1656 fn from(value: ColumnChunkMetaData) -> Self {
1657 ColumnChunkMetaDataBuilder(value)
1658 }
1659}
1660
1661pub struct OffsetIndexBuilder {
1665 offset_array: Vec<i64>,
1666 compressed_page_size_array: Vec<i32>,
1667 first_row_index_array: Vec<i64>,
1668 unencoded_byte_array_data_bytes_array: Option<Vec<i64>>,
1669 current_first_row_index: i64,
1670}
1671
1672impl Default for OffsetIndexBuilder {
1673 fn default() -> Self {
1674 Self::new()
1675 }
1676}
1677
1678impl OffsetIndexBuilder {
1679 pub fn new() -> Self {
1681 OffsetIndexBuilder {
1682 offset_array: Vec::new(),
1683 compressed_page_size_array: Vec::new(),
1684 first_row_index_array: Vec::new(),
1685 unencoded_byte_array_data_bytes_array: None,
1686 current_first_row_index: 0,
1687 }
1688 }
1689
1690 pub fn append_row_count(&mut self, row_count: i64) {
1692 let current_page_row_index = self.current_first_row_index;
1693 self.first_row_index_array.push(current_page_row_index);
1694 self.current_first_row_index += row_count;
1695 }
1696
1697 pub fn append_offset_and_size(&mut self, offset: i64, compressed_page_size: i32) {
1699 self.offset_array.push(offset);
1700 self.compressed_page_size_array.push(compressed_page_size);
1701 }
1702
1703 pub fn append_unencoded_byte_array_data_bytes(
1705 &mut self,
1706 unencoded_byte_array_data_bytes: Option<i64>,
1707 ) {
1708 if let Some(val) = unencoded_byte_array_data_bytes {
1709 self.unencoded_byte_array_data_bytes_array
1710 .get_or_insert(Vec::new())
1711 .push(val);
1712 }
1713 }
1714
1715 pub fn build(self) -> OffsetIndexMetaData {
1717 let locations = self
1718 .offset_array
1719 .iter()
1720 .zip(self.compressed_page_size_array.iter())
1721 .zip(self.first_row_index_array.iter())
1722 .map(|((offset, size), row_index)| PageLocation {
1723 offset: *offset,
1724 compressed_page_size: *size,
1725 first_row_index: *row_index,
1726 })
1727 .collect::<Vec<_>>();
1728 OffsetIndexMetaData {
1729 page_locations: locations,
1730 unencoded_byte_array_data_bytes: self.unencoded_byte_array_data_bytes_array,
1731 }
1732 }
1733}
1734
1735#[cfg(test)]
1736mod tests {
1737 use super::*;
1738 use crate::basic::{PageType, SortOrder};
1739 use crate::file::metadata::page_index::PageIndexBuilder;
1740 use crate::file::metadata::thrift::tests::{
1741 read_column_chunk, read_column_chunk_with_options, read_row_group,
1742 };
1743
1744 #[test]
1745 #[expect(deprecated)]
1746 fn test_level_histogram_update_from_levels_compat() {
1747 let mut histogram = LevelHistogram::try_new(2).unwrap();
1748 histogram.update_from_levels(&[0, 2, 1, 2, 2]);
1749 assert_eq!(histogram.values(), &[1, 1, 3]);
1750 }
1751
1752 #[test]
1753 fn test_row_group_metadata_thrift_conversion() {
1754 let schema_descr = get_test_schema_descr();
1755
1756 let mut columns = vec![];
1757 for ptr in schema_descr.columns() {
1758 let column = ColumnChunkMetaData::builder(ptr.clone()).build().unwrap();
1759 columns.push(column);
1760 }
1761 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
1762 .set_num_rows(1000)
1763 .set_total_byte_size(2000)
1764 .set_column_metadata(columns)
1765 .set_ordinal(1)
1766 .build()
1767 .unwrap();
1768
1769 let mut buf = Vec::new();
1770 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1771 row_group_meta.write_thrift(&mut writer).unwrap();
1772
1773 let row_group_res = read_row_group(&buf, schema_descr).unwrap();
1774
1775 assert_eq!(row_group_res, row_group_meta);
1776 }
1777
1778 #[test]
1779 fn test_row_group_metadata_thrift_conversion_empty() {
1780 let schema_descr = get_test_schema_descr();
1781
1782 let row_group_meta = RowGroupMetaData::builder(schema_descr).build();
1783
1784 assert!(row_group_meta.is_err());
1785 if let Err(e) = row_group_meta {
1786 assert_eq!(
1787 format!("{e}"),
1788 "Parquet error: Column length mismatch: 2 != 0"
1789 );
1790 }
1791 }
1792
1793 #[test]
1795 fn test_row_group_metadata_thrift_corrupted() {
1796 let schema_descr_2cols = Arc::new(SchemaDescriptor::new(Arc::new(
1797 SchemaType::group_type_builder("schema")
1798 .with_fields(vec![
1799 Arc::new(
1800 SchemaType::primitive_type_builder("a", Type::INT32)
1801 .build()
1802 .unwrap(),
1803 ),
1804 Arc::new(
1805 SchemaType::primitive_type_builder("b", Type::INT32)
1806 .build()
1807 .unwrap(),
1808 ),
1809 ])
1810 .build()
1811 .unwrap(),
1812 )));
1813
1814 let schema_descr_3cols = Arc::new(SchemaDescriptor::new(Arc::new(
1815 SchemaType::group_type_builder("schema")
1816 .with_fields(vec![
1817 Arc::new(
1818 SchemaType::primitive_type_builder("a", Type::INT32)
1819 .build()
1820 .unwrap(),
1821 ),
1822 Arc::new(
1823 SchemaType::primitive_type_builder("b", Type::INT32)
1824 .build()
1825 .unwrap(),
1826 ),
1827 Arc::new(
1828 SchemaType::primitive_type_builder("c", Type::INT32)
1829 .build()
1830 .unwrap(),
1831 ),
1832 ])
1833 .build()
1834 .unwrap(),
1835 )));
1836
1837 let row_group_meta_2cols = RowGroupMetaData::builder(schema_descr_2cols.clone())
1838 .set_num_rows(1000)
1839 .set_total_byte_size(2000)
1840 .set_column_metadata(vec![
1841 ColumnChunkMetaData::builder(schema_descr_2cols.column(0))
1842 .build()
1843 .unwrap(),
1844 ColumnChunkMetaData::builder(schema_descr_2cols.column(1))
1845 .build()
1846 .unwrap(),
1847 ])
1848 .set_ordinal(1)
1849 .build()
1850 .unwrap();
1851 let mut buf = Vec::new();
1852 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1853 row_group_meta_2cols.write_thrift(&mut writer).unwrap();
1854
1855 let err = read_row_group(&buf, schema_descr_3cols)
1856 .unwrap_err()
1857 .to_string();
1858 assert_eq!(
1859 err,
1860 "Parquet error: Column count mismatch. Schema has 3 columns while Row Group has 2"
1861 );
1862 }
1863
1864 #[test]
1865 fn test_column_chunk_metadata_thrift_conversion() {
1866 let column_descr = get_test_schema_descr().column(0);
1867 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1868 .set_encodings_mask(EncodingMask::new_from_encodings(
1869 [Encoding::PLAIN, Encoding::RLE].iter(),
1870 ))
1871 .set_file_path("file_path".to_owned())
1872 .set_num_values(1000)
1873 .set_compression_codec(CompressionCodec::SNAPPY)
1874 .set_total_compressed_size(2000)
1875 .set_total_uncompressed_size(3000)
1876 .set_data_page_offset(4000)
1877 .set_dictionary_page_offset(Some(5000))
1878 .set_page_encoding_stats(vec![
1879 PageEncodingStats {
1880 page_type: PageType::DATA_PAGE,
1881 encoding: Encoding::PLAIN,
1882 count: 3,
1883 },
1884 PageEncodingStats {
1885 page_type: PageType::DATA_PAGE,
1886 encoding: Encoding::RLE,
1887 count: 5,
1888 },
1889 ])
1890 .set_bloom_filter_offset(Some(6000))
1891 .set_bloom_filter_length(Some(25))
1892 .set_offset_index_offset(Some(7000))
1893 .set_offset_index_length(Some(25))
1894 .set_column_index_offset(Some(8000))
1895 .set_column_index_length(Some(25))
1896 .set_unencoded_byte_array_data_bytes(Some(2000))
1897 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
1898 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
1899 .build()
1900 .unwrap();
1901
1902 let mut buf = Vec::new();
1903 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1904 col_metadata.write_thrift(&mut writer).unwrap();
1905 let col_chunk_res = read_column_chunk(&buf, column_descr.clone()).unwrap();
1906
1907 let expected_metadata = ColumnChunkMetaData::builder(column_descr)
1908 .set_encodings_mask(EncodingMask::new_from_encodings(
1909 [Encoding::PLAIN, Encoding::RLE].iter(),
1910 ))
1911 .set_file_path("file_path".to_owned())
1912 .set_num_values(1000)
1913 .set_compression_codec(CompressionCodec::SNAPPY)
1914 .set_total_compressed_size(2000)
1915 .set_total_uncompressed_size(3000)
1916 .set_data_page_offset(4000)
1917 .set_dictionary_page_offset(Some(5000))
1918 .set_page_encoding_stats_mask(EncodingMask::new_from_encodings(
1919 [Encoding::PLAIN, Encoding::RLE].iter(),
1920 ))
1921 .set_bloom_filter_offset(Some(6000))
1922 .set_bloom_filter_length(Some(25))
1923 .set_offset_index_offset(Some(7000))
1924 .set_offset_index_length(Some(25))
1925 .set_column_index_offset(Some(8000))
1926 .set_column_index_length(Some(25))
1927 .set_unencoded_byte_array_data_bytes(Some(2000))
1928 .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
1929 .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
1930 .build()
1931 .unwrap();
1932
1933 assert_eq!(col_chunk_res, expected_metadata);
1934 }
1935
1936 #[test]
1937 fn test_column_chunk_metadata_thrift_conversion_full_stats() {
1938 let column_descr = get_test_schema_descr().column(0);
1939 let stats = vec![
1940 PageEncodingStats {
1941 page_type: PageType::DATA_PAGE,
1942 encoding: Encoding::PLAIN,
1943 count: 3,
1944 },
1945 PageEncodingStats {
1946 page_type: PageType::DATA_PAGE,
1947 encoding: Encoding::RLE,
1948 count: 5,
1949 },
1950 ];
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_num_values(1000)
1956 .set_compression_codec(CompressionCodec::SNAPPY)
1957 .set_total_compressed_size(2000)
1958 .set_total_uncompressed_size(3000)
1959 .set_data_page_offset(4000)
1960 .set_page_encoding_stats(stats)
1961 .build()
1962 .unwrap();
1963
1964 let mut buf = Vec::new();
1965 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1966 col_metadata.write_thrift(&mut writer).unwrap();
1967
1968 let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
1969 let col_chunk_res =
1970 read_column_chunk_with_options(&buf, column_descr, Some(&options)).unwrap();
1971
1972 assert_eq!(col_chunk_res, col_metadata);
1973 }
1974
1975 #[test]
1976 fn test_column_chunk_metadata_thrift_conversion_empty() {
1977 let column_descr = get_test_schema_descr().column(0);
1978
1979 let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1980 .build()
1981 .unwrap();
1982
1983 let mut buf = Vec::new();
1984 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1985 col_metadata.write_thrift(&mut writer).unwrap();
1986 let col_chunk_res = read_column_chunk(&buf, column_descr).unwrap();
1987
1988 assert_eq!(col_chunk_res, col_metadata);
1989 }
1990
1991 #[test]
1992 fn test_compressed_size() {
1993 let schema_descr = get_test_schema_descr();
1994
1995 let mut columns = vec![];
1996 for column_descr in schema_descr.columns() {
1997 let column = ColumnChunkMetaData::builder(column_descr.clone())
1998 .set_total_compressed_size(500)
1999 .set_total_uncompressed_size(700)
2000 .build()
2001 .unwrap();
2002 columns.push(column);
2003 }
2004 let row_group_meta = RowGroupMetaData::builder(schema_descr)
2005 .set_num_rows(1000)
2006 .set_column_metadata(columns)
2007 .build()
2008 .unwrap();
2009
2010 let compressed_size_res = row_group_meta.compressed_size();
2011 let compressed_size_exp: i64 = 1000;
2012
2013 assert_eq!(compressed_size_res, compressed_size_exp);
2014 }
2015
2016 #[test]
2017 fn test_memory_size() {
2018 let schema_descr = get_test_schema_descr();
2019
2020 let columns = schema_descr
2021 .columns()
2022 .iter()
2023 .map(|column_descr| {
2024 ColumnChunkMetaData::builder(column_descr.clone())
2025 .set_statistics(Statistics::new::<i32>(None, None, None, None, false))
2026 .build()
2027 })
2028 .collect::<Result<Vec<_>>>()
2029 .unwrap();
2030 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2031 .set_num_rows(1000)
2032 .set_column_metadata(columns)
2033 .build()
2034 .unwrap();
2035 let row_group_meta = vec![row_group_meta];
2036
2037 let version = 2;
2038 let num_rows = 1000;
2039 let created_by = Some(String::from("test harness"));
2040 let key_value_metadata = Some(vec![KeyValue::new(
2041 String::from("Foo"),
2042 Some(String::from("bar")),
2043 )]);
2044 let column_orders = Some(vec![
2045 ColumnOrder::UNDEFINED,
2046 ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED),
2047 ]);
2048 let file_metadata = FileMetaData::new(
2049 version,
2050 num_rows,
2051 created_by,
2052 key_value_metadata,
2053 schema_descr.clone(),
2054 column_orders,
2055 );
2056
2057 let columns_with_stats = schema_descr
2059 .columns()
2060 .iter()
2061 .map(|column_descr| {
2062 ColumnChunkMetaData::builder(column_descr.clone())
2063 .set_statistics(Statistics::new::<i32>(
2064 Some(0),
2065 Some(100),
2066 None,
2067 None,
2068 false,
2069 ))
2070 .build()
2071 })
2072 .collect::<Result<Vec<_>>>()
2073 .unwrap();
2074
2075 let row_group_meta_with_stats = RowGroupMetaData::builder(schema_descr)
2076 .set_num_rows(1000)
2077 .set_column_metadata(columns_with_stats)
2078 .build()
2079 .unwrap();
2080 let row_group_meta_with_stats = vec![row_group_meta_with_stats];
2081
2082 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata.clone())
2083 .set_row_groups(row_group_meta_with_stats)
2084 .build();
2085
2086 #[cfg(not(feature = "encryption"))]
2088 let base_expected_size = 2766;
2089 #[cfg(feature = "encryption")]
2090 let base_expected_size = 2934;
2091
2092 assert_eq!(parquet_meta.memory_size(), base_expected_size);
2093
2094 let mut page_index = PageIndexBuilder::new(1, 1);
2095
2096 let mut column_index = ColumnIndexBuilder::new(Type::BOOLEAN);
2097 column_index.append(false, vec![1u8], vec![2u8, 3u8], 4, None);
2098 let column_index = column_index.build().unwrap();
2099 {
2100 let ColumnIndexMetaData::BOOLEAN(_) = column_index else {
2101 panic!("wrong type of column index")
2102 };
2103 }
2104
2105 page_index.put_column_index(column_index, 0, 0);
2106
2107 let mut offset_index = OffsetIndexBuilder::new();
2109 offset_index.append_row_count(1);
2110 offset_index.append_offset_and_size(2, 3);
2111 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2112 offset_index.append_row_count(1);
2113 offset_index.append_offset_and_size(2, 3);
2114 offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2115 let offset_index = offset_index.build();
2116 page_index.put_offset_index(offset_index, 0, 0);
2117
2118 let parquet_meta = ParquetMetaDataBuilder::new(file_metadata)
2119 .set_row_groups(row_group_meta)
2120 .set_page_index(Some(Arc::new(page_index.build())))
2121 .build();
2122
2123 #[cfg(not(feature = "encryption"))]
2125 let bigger_expected_size = 3280;
2126 #[cfg(feature = "encryption")]
2127 let bigger_expected_size = 3448;
2128
2129 assert!(bigger_expected_size > base_expected_size);
2131 assert_eq!(parquet_meta.memory_size(), bigger_expected_size);
2132 }
2133
2134 #[test]
2135 #[cfg(feature = "encryption")]
2136 fn test_memory_size_with_decryptor() {
2137 use crate::encryption::decrypt::FileDecryptionProperties;
2138 use crate::file::metadata::thrift::encryption::AesGcmV1;
2139
2140 let schema_descr = get_test_schema_descr();
2141
2142 let columns = schema_descr
2143 .columns()
2144 .iter()
2145 .map(|column_descr| ColumnChunkMetaData::builder(column_descr.clone()).build())
2146 .collect::<Result<Vec<_>>>()
2147 .unwrap();
2148 let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2149 .set_num_rows(1000)
2150 .set_column_metadata(columns)
2151 .build()
2152 .unwrap();
2153 let row_group_meta = vec![row_group_meta];
2154
2155 let version = 2;
2156 let num_rows = 1000;
2157 let aad_file_unique = vec![1u8; 8];
2158 let aad_prefix = vec![2u8; 8];
2159 let encryption_algorithm = EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 {
2160 aad_prefix: Some(aad_prefix.clone()),
2161 aad_file_unique: Some(aad_file_unique.clone()),
2162 supply_aad_prefix: Some(true),
2163 });
2164 let footer_key_metadata = Some(vec![3u8; 8]);
2165 let file_metadata =
2166 FileMetaData::new(version, num_rows, None, None, schema_descr.clone(), None)
2167 .with_encryption_algorithm(Some(encryption_algorithm))
2168 .with_footer_signing_key_metadata(footer_key_metadata.clone());
2169
2170 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2171 .set_row_groups(row_group_meta.clone())
2172 .build();
2173
2174 let base_expected_size = 2042;
2175 assert_eq!(parquet_meta_data.memory_size(), base_expected_size);
2176
2177 let footer_key = b"0123456789012345";
2178 let column_key = b"1234567890123450";
2179 let mut decryption_properties_builder =
2180 FileDecryptionProperties::builder(footer_key.to_vec())
2181 .with_aad_prefix(aad_prefix.clone());
2182 for column in schema_descr.columns() {
2183 decryption_properties_builder = decryption_properties_builder
2184 .with_column_key(&column.path().string(), column_key.to_vec());
2185 }
2186 let decryption_properties = decryption_properties_builder.build().unwrap();
2187 let decryptor = FileDecryptor::new(
2188 &decryption_properties,
2189 footer_key_metadata.as_deref(),
2190 aad_file_unique,
2191 aad_prefix,
2192 )
2193 .unwrap();
2194
2195 let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2196 .set_row_groups(row_group_meta.clone())
2197 .set_file_decryptor(Some(decryptor))
2198 .build();
2199
2200 let expected_size_with_decryptor = 3056;
2201 assert!(expected_size_with_decryptor > base_expected_size);
2202
2203 assert_eq!(
2204 parquet_meta_data.memory_size(),
2205 expected_size_with_decryptor
2206 );
2207 }
2208
2209 fn get_test_schema_descr() -> SchemaDescPtr {
2211 let schema = SchemaType::group_type_builder("schema")
2212 .with_fields(vec![
2213 Arc::new(
2214 SchemaType::primitive_type_builder("a", Type::INT32)
2215 .build()
2216 .unwrap(),
2217 ),
2218 Arc::new(
2219 SchemaType::primitive_type_builder("b", Type::INT32)
2220 .build()
2221 .unwrap(),
2222 ),
2223 ])
2224 .build()
2225 .unwrap();
2226
2227 Arc::new(SchemaDescriptor::new(Arc::new(schema)))
2228 }
2229
2230 #[test]
2231 fn test_page_index_builder_partial_population() {
2232 let mut builder = PageIndexBuilder::new(1, 1);
2233
2234 let mut offset_idx_builder = OffsetIndexBuilder::new();
2236 offset_idx_builder.append_row_count(50);
2237 offset_idx_builder.append_offset_and_size(1000, 500);
2238 let offset_idx = offset_idx_builder.build();
2239 builder.put_offset_index(offset_idx, 0, 0);
2240
2241 let page_index = builder.build();
2242
2243 assert!(!page_index.has_column_indexes());
2245 assert!(page_index.has_offset_indexes());
2247 }
2248
2249 #[test]
2250 fn test_page_index_builder_empty_to_none() {
2251 let builder = PageIndexBuilder::new(2, 2);
2253
2254 let page_index = builder.build();
2255
2256 assert!(!page_index.has_column_indexes());
2258 assert!(!page_index.has_offset_indexes());
2259 }
2260
2261 #[test]
2262 fn test_rebuild_page_index() {
2263 let mut builder = PageIndexBuilder::new(1, 2);
2265
2266 let mut col_idx_builder = ColumnIndexBuilder::new(Type::INT32);
2268 col_idx_builder.append(false, vec![1, 0, 0, 0], vec![100, 0, 0, 0], 10, None);
2269 let col_idx = col_idx_builder.build().unwrap();
2270 builder.put_column_index(col_idx.clone(), 0, 0);
2271
2272 let mut offset_idx_builder = OffsetIndexBuilder::new();
2274 offset_idx_builder.append_row_count(50);
2275 offset_idx_builder.append_offset_and_size(1000, 500);
2276 let offset_idx = offset_idx_builder.build();
2277 builder.put_offset_index(offset_idx.clone(), 0, 0);
2278
2279 let page_index = builder.build();
2280
2281 assert!(page_index.is_complete());
2283
2284 assert!(page_index.column_index(0, 0).is_some());
2286 assert!(page_index.offset_index(0, 0).is_some());
2287 assert!(page_index.column_index(0, 1).is_none());
2289 assert!(page_index.offset_index(0, 1).is_none());
2290
2291 let mut builder = PageIndexBuilder::new_from(page_index);
2292
2293 builder.put_column_index(col_idx, 0, 1);
2295 builder.put_offset_index(offset_idx, 0, 1);
2296 let page_index = builder.build();
2297
2298 assert!(page_index.column_index(0, 0).is_some());
2300 assert!(page_index.offset_index(0, 0).is_some());
2301 assert!(page_index.column_index(0, 1).is_some());
2302 assert!(page_index.offset_index(0, 1).is_some());
2303 }
2304}