1#[cfg(feature = "arrow")]
6use crate::arrow::ProjectionMask;
7#[cfg(feature = "encryption")]
8use crate::encryption::decrypt::FileDecryptionProperties;
9use crate::errors::{ParquetError, Result};
10use crate::file::FOOTER_SIZE;
11#[cfg(feature = "arrow")]
12use crate::file::metadata::dictionary::decode_dictionary_page;
13use crate::file::metadata::page_index::{PageIndex, PageIndexProvider};
14use crate::file::metadata::parser::decode_metadata;
15use crate::file::metadata::thrift::parquet_schema_from_bytes;
16use crate::file::metadata::{
17 FooterTail, ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataPushDecoder,
18};
19use crate::file::reader::ChunkReader;
20use crate::schema::types::SchemaDescriptor;
21#[cfg(feature = "arrow")]
22use arrow_array::ArrayRef;
23use bytes::Bytes;
24use std::collections::BTreeSet;
25use std::iter::FusedIterator;
26use std::slice::Iter;
27use std::sync::Arc;
28use std::{io::Read, ops::Range};
29
30use crate::DecodeResult;
31#[cfg(all(feature = "async", feature = "arrow"))]
32use crate::arrow::async_reader::{MetadataFetch, MetadataSuffixFetch};
33
34#[derive(Default, Debug)]
67pub struct ParquetMetaDataReader {
68 metadata: Option<ParquetMetaData>,
69 column_index: PageIndexPolicy,
70 offset_index: PageIndexPolicy,
71 column_index_mask: ColumnChunkMask,
72 offset_index_mask: ColumnChunkMask,
73 prefetch_hint: Option<usize>,
74 metadata_options: Option<Arc<ParquetMetaDataOptions>>,
75 metadata_size: Option<usize>,
78 #[cfg(feature = "encryption")]
79 file_decryption_properties: Option<Arc<FileDecryptionProperties>>,
80}
81
82#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
88pub enum PageIndexPolicy {
89 #[default]
91 Skip,
92 Optional,
94 Required,
101}
102
103impl From<bool> for PageIndexPolicy {
104 fn from(value: bool) -> Self {
105 match value {
106 true => Self::Required,
107 false => Self::Skip,
108 }
109 }
110}
111
112#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
147pub struct ColumnChunkMask {
148 row_groups: Option<Arc<[u32]>>,
154 columns: Option<Arc<[u32]>>,
155}
156
157impl ColumnChunkMask {
158 pub fn all() -> Self {
160 Self::default()
161 }
162
163 pub fn none() -> Self {
165 Self {
166 row_groups: Some(Arc::from([])),
167 columns: Some(Arc::from([])),
168 }
169 }
170
171 pub fn columns(columns: impl IntoIterator<Item = usize>) -> Self {
175 Self {
176 row_groups: None,
177 columns: Self::iter_to_set(columns),
178 }
179 }
180
181 pub fn row_groups(row_groups: impl IntoIterator<Item = usize>) -> Self {
185 Self {
186 row_groups: Self::iter_to_set(row_groups),
187 columns: None,
188 }
189 }
190
191 pub fn row_groups_and_columns(
195 row_groups: impl IntoIterator<Item = usize>,
196 columns: impl IntoIterator<Item = usize>,
197 ) -> Self {
198 Self {
199 row_groups: Self::iter_to_set(row_groups),
200 columns: Self::iter_to_set(columns),
201 }
202 }
203
204 pub fn includes_row_group(&self, idx: usize, num_row_groups: usize) -> bool {
206 idx < num_row_groups && Self::includes_index(self.row_groups.as_ref(), idx)
207 }
208
209 pub fn includes_column(&self, idx: usize, num_columns: usize) -> bool {
211 idx < num_columns && Self::includes_index(self.columns.as_ref(), idx)
212 }
213
214 fn includes_index(keep: Option<&Arc<[u32]>>, idx: usize) -> bool {
215 let Ok(idx) = i32::try_from(idx) else {
217 return false;
218 };
219 let idx = idx as u32;
220 keep.is_none_or(|keep| keep.binary_search(&idx).is_ok())
221 }
222
223 pub fn is_all(&self) -> bool {
225 self.row_groups.is_none() && self.columns.is_none()
226 }
227
228 pub fn is_none(&self) -> bool {
234 self.row_groups.as_ref().is_some_and(|i| i.is_empty())
235 || self.columns.as_ref().is_some_and(|i| i.is_empty())
236 }
237
238 #[cfg(feature = "arrow")]
240 pub fn from_projection(projection: &ProjectionMask, schema: &SchemaDescriptor) -> Self {
241 Self::columns((0..schema.num_columns()).filter(|&i| projection.leaf_included(i)))
242 }
243
244 pub fn row_group_indices(&self, num_row_groups: usize) -> impl Iterator<Item = usize> + '_ {
247 Self::axis_indices(self.row_groups.as_deref(), num_row_groups)
248 }
249
250 pub fn column_indices(&self, num_columns: usize) -> impl Iterator<Item = usize> + '_ {
253 Self::axis_indices(self.columns.as_deref(), num_columns)
254 }
255
256 fn axis_indices(axis: Option<&[u32]>, len: usize) -> AxisIndices<'_> {
257 match axis {
258 None => AxisIndices::All(0..len),
259 Some(indices) => {
260 let end = indices.partition_point(|&i| (i as usize) < len);
261 AxisIndices::Selected(indices[..end].iter())
262 }
263 }
264 }
265
266 fn iter_to_set(indices: impl IntoIterator<Item = usize>) -> Option<Arc<[u32]>> {
269 Some(
270 indices
271 .into_iter()
272 .filter_map(|i| i32::try_from(i).ok()) .map(|i| i as u32) .collect::<BTreeSet<_>>()
275 .into_iter()
276 .collect::<Vec<_>>()
277 .into(),
278 )
279 }
280}
281
282enum AxisIndices<'a> {
284 All(Range<usize>),
285 Selected(Iter<'a, u32>),
286}
287
288impl Iterator for AxisIndices<'_> {
289 type Item = usize;
290
291 fn next(&mut self) -> Option<Self::Item> {
292 match self {
293 Self::All(indices) => indices.next(),
294 Self::Selected(indices) => indices.next().map(|&i| i as usize),
295 }
296 }
297
298 fn size_hint(&self) -> (usize, Option<usize>) {
299 match self {
300 Self::All(indices) => indices.size_hint(),
301 Self::Selected(indices) => indices.size_hint(),
302 }
303 }
304}
305
306impl ExactSizeIterator for AxisIndices<'_> {}
307impl FusedIterator for AxisIndices<'_> {}
308
309impl ParquetMetaDataReader {
310 pub fn new() -> Self {
312 Default::default()
313 }
314
315 pub fn new_with_metadata(metadata: ParquetMetaData) -> Self {
318 Self {
319 metadata: Some(metadata),
320 ..Default::default()
321 }
322 }
323
324 pub fn with_page_index_policy(self, policy: PageIndexPolicy) -> Self {
326 self.with_column_index_policy(policy)
327 .with_offset_index_policy(policy)
328 }
329
330 pub fn with_column_index_policy(mut self, policy: PageIndexPolicy) -> Self {
332 self.column_index = policy;
333 self
334 }
335
336 pub fn with_offset_index_policy(mut self, policy: PageIndexPolicy) -> Self {
338 self.offset_index = policy;
339 self
340 }
341
342 pub fn with_page_index_mask(self, mask: ColumnChunkMask) -> Self {
350 self.with_column_index_mask(mask.clone())
351 .with_offset_index_mask(mask)
352 }
353
354 pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self {
359 self.column_index_mask = mask;
360 self
361 }
362
363 pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self {
368 self.offset_index_mask = mask;
369 self
370 }
371
372 pub fn with_metadata_options(mut self, options: Option<ParquetMetaDataOptions>) -> Self {
374 self.metadata_options = options.map(Arc::new);
375 self
376 }
377
378 pub fn with_prefetch_hint(mut self, prefetch: Option<usize>) -> Self {
390 self.prefetch_hint = prefetch;
391 self
392 }
393
394 #[cfg(feature = "encryption")]
398 pub fn with_decryption_properties(
399 mut self,
400 properties: Option<std::sync::Arc<FileDecryptionProperties>>,
401 ) -> Self {
402 self.file_decryption_properties = properties;
403 self
404 }
405
406 pub fn has_metadata(&self) -> bool {
408 self.metadata.is_some()
409 }
410
411 pub fn finish(&mut self) -> Result<ParquetMetaData> {
413 self.metadata
414 .take()
415 .ok_or_else(|| general_err!("could not parse parquet metadata"))
416 }
417
418 pub fn parse_and_finish<R: ChunkReader>(mut self, reader: &R) -> Result<ParquetMetaData> {
437 self.try_parse(reader)?;
438 self.finish()
439 }
440
441 pub fn try_parse<R: ChunkReader>(&mut self, reader: &R) -> Result<()> {
447 self.try_parse_sized(reader, reader.len())
448 }
449
450 pub fn try_parse_sized<R: ChunkReader>(&mut self, reader: &R, file_size: u64) -> Result<()> {
523 self.metadata = match self.parse_metadata(reader) {
524 Ok(metadata) => Some(metadata),
525 Err(ParquetError::NeedMoreData(needed)) => {
526 return if file_size == reader.len() || needed as u64 > file_size {
529 Err(eof_err!(
530 "Parquet file too small. Size is {} but need {}",
531 file_size,
532 needed
533 ))
534 } else {
535 Err(ParquetError::NeedMoreData(needed))
537 };
538 }
539 Err(e) => return Err(e),
540 };
541
542 if self.column_index == PageIndexPolicy::Skip && self.offset_index == PageIndexPolicy::Skip
544 {
545 return Ok(());
546 }
547
548 self.read_page_indexes_sized(reader, file_size)
549 }
550
551 pub fn read_page_indexes<R: ChunkReader>(&mut self, reader: &R) -> Result<()> {
561 self.read_page_indexes_sized(reader, reader.len())
562 }
563
564 pub fn read_page_index_and_finish<R: ChunkReader>(
612 self,
613 reader: &R,
614 ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
615 self.read_page_index_sized_and_finish(reader, reader.len())
616 }
617
618 pub fn read_page_index_sized_and_finish<R: ChunkReader>(
626 mut self,
627 reader: &R,
628 file_size: u64,
629 ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
630 self.read_page_index_sized_impl(reader, file_size)
631 }
632
633 fn read_page_index_sized_impl<R: ChunkReader>(
634 &mut self,
635 reader: &R,
636 file_size: u64,
637 ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
638 let Some(metadata) = self.metadata.take() else {
639 return Err(general_err!(
640 "Tried to read page indexes without ParquetMetaData metadata"
641 ));
642 };
643
644 let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)?
645 .with_offset_index_policy(self.offset_index)
646 .with_column_index_policy(self.column_index)
647 .with_offset_index_mask(self.offset_index_mask.clone())
648 .with_column_index_mask(self.column_index_mask.clone())
649 .with_metadata_options(self.metadata_options.clone());
650 let mut push_decoder = self.prepare_push_decoder(push_decoder);
651
652 let range = match needs_page_index_data(&mut push_decoder)? {
653 NeedsIndexData::No(page_index) => return Ok(page_index),
654 NeedsIndexData::Yes(range) => range,
655 };
656
657 let file_range = file_size.saturating_sub(reader.len())..file_size;
658 if !(file_range.contains(&range.start) && file_range.contains(&range.end)) {
659 return if range.end > file_size {
660 Err(eof_err!(
661 "Parquet file too small. Range {range:?} is beyond file bounds {file_size}",
662 ))
663 } else {
664 Err(ParquetError::NeedMoreData(
665 (file_size - range.start).try_into()?,
666 ))
667 };
668 }
669
670 if let Some(metadata_size) = self.metadata_size {
671 let metadata_range = file_size.saturating_sub(metadata_size as u64)..file_size;
672 if range.end > metadata_range.start {
673 return Err(eof_err!(
674 "Parquet file too small. Page index range {range:?} overlaps with file metadata {metadata_range:?}",
675 ));
676 }
677 }
678
679 let bytes_needed = usize::try_from(range.end - range.start)?;
680 let bytes = reader.get_bytes(range.start - file_range.start, bytes_needed)?;
681 push_decoder.push_range(range, bytes)?;
682 parse_page_index_data(&mut push_decoder)
683 }
684
685 pub fn read_page_indexes_sized<R: ChunkReader>(
694 &mut self,
695 reader: &R,
696 file_size: u64,
697 ) -> Result<()> {
698 let (mut metadata, page_index) = self.read_page_index_sized_impl(reader, file_size)?;
699 metadata
700 .set_page_index(page_index.map(|index| Arc::new(index) as Arc<dyn PageIndexProvider>));
701 self.metadata = Some(metadata);
702 Ok(())
703 }
704
705 #[cfg(all(feature = "async", feature = "arrow"))]
712 pub async fn load_and_finish<F: MetadataFetch>(
713 mut self,
714 fetch: F,
715 file_size: u64,
716 ) -> Result<ParquetMetaData> {
717 self.try_load(fetch, file_size).await?;
718 self.finish()
719 }
720
721 #[cfg(all(feature = "async", feature = "arrow"))]
728 pub async fn load_via_suffix_and_finish<F: MetadataSuffixFetch>(
729 mut self,
730 fetch: F,
731 ) -> Result<ParquetMetaData> {
732 self.try_load_via_suffix(fetch).await?;
733 self.finish()
734 }
735 #[cfg(all(feature = "async", feature = "arrow"))]
741 pub async fn try_load<F: MetadataFetch>(&mut self, mut fetch: F, file_size: u64) -> Result<()> {
742 let (metadata, remainder) = self.load_metadata(&mut fetch, file_size).await?;
743
744 self.metadata = Some(metadata);
745
746 if self.column_index == PageIndexPolicy::Skip && self.offset_index == PageIndexPolicy::Skip
748 {
749 return Ok(());
750 }
751
752 self.load_page_index_with_remainder(fetch, remainder).await
753 }
754
755 #[cfg(all(feature = "async", feature = "arrow"))]
761 pub async fn try_load_via_suffix<F: MetadataSuffixFetch>(
762 &mut self,
763 mut fetch: F,
764 ) -> Result<()> {
765 let metadata = self.load_metadata_via_suffix(&mut fetch).await?;
766
767 self.metadata = Some(metadata);
768
769 if self.column_index == PageIndexPolicy::Skip && self.offset_index == PageIndexPolicy::Skip
771 {
772 return Ok(());
773 }
774
775 self.load_page_index_with_remainder(fetch, None).await
776 }
777
778 #[cfg(all(feature = "async", feature = "arrow"))]
784 pub async fn load_page_index<F: MetadataFetch>(&mut self, fetch: F) -> Result<()> {
785 self.load_page_index_with_remainder(fetch, None).await
786 }
787
788 #[cfg(all(feature = "async", feature = "arrow"))]
797 pub async fn load_page_index_and_finish<F: MetadataFetch>(
798 mut self,
799 fetch: F,
800 ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
801 self.load_page_index_with_remainder_impl(fetch, None).await
802 }
803
804 #[cfg(feature = "arrow")]
822 pub fn read_column_dictionary<R: ChunkReader>(
823 reader: &R,
824 metadata: &ParquetMetaData,
825 row_group_idx: usize,
826 column_idx: usize,
827 ) -> Result<Option<ArrayRef>> {
828 let Some(range) = dictionary_page_byte_range(metadata, row_group_idx, column_idx)? else {
829 return Ok(None);
830 };
831 let length = usize::try_from(range.end - range.start)?;
832 let buffer = reader.get_bytes(range.start, length)?;
833 decode_dictionary_page(buffer, metadata, row_group_idx, column_idx).map(Some)
834 }
835
836 #[cfg(all(feature = "async", feature = "arrow"))]
838 pub async fn read_column_dictionary_async<F: MetadataFetch>(
839 mut fetch: F,
840 metadata: &ParquetMetaData,
841 row_group_idx: usize,
842 column_idx: usize,
843 ) -> Result<Option<ArrayRef>> {
844 let Some(range) = dictionary_page_byte_range(metadata, row_group_idx, column_idx)? else {
845 return Ok(None);
846 };
847 let buffer = fetch.fetch(range).await?;
848 decode_dictionary_page(buffer, metadata, row_group_idx, column_idx).map(Some)
849 }
850
851 #[cfg(all(feature = "async", feature = "arrow"))]
852 async fn load_page_index_with_remainder<F: MetadataFetch>(
853 &mut self,
854 fetch: F,
855 remainder: Option<(usize, Bytes)>,
856 ) -> Result<()> {
857 let (mut metadata, page_index) = self
858 .load_page_index_with_remainder_impl(fetch, remainder)
859 .await?;
860 metadata
861 .set_page_index(page_index.map(|index| Arc::new(index) as Arc<dyn PageIndexProvider>));
862 self.metadata = Some(metadata);
863 Ok(())
864 }
865
866 #[cfg(all(feature = "async", feature = "arrow"))]
867 async fn load_page_index_with_remainder_impl<F: MetadataFetch>(
868 &mut self,
869 mut fetch: F,
870 remainder: Option<(usize, Bytes)>,
871 ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
872 let Some(metadata) = self.metadata.take() else {
873 return Err(general_err!("Footer metadata is not present"));
874 };
875
876 let file_size = u64::MAX;
879 let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)?
880 .with_offset_index_policy(self.offset_index)
881 .with_column_index_policy(self.column_index)
882 .with_offset_index_mask(self.offset_index_mask.clone())
883 .with_column_index_mask(self.column_index_mask.clone())
884 .with_metadata_options(self.metadata_options.clone());
885 let mut push_decoder = self.prepare_push_decoder(push_decoder);
886
887 let range = match needs_page_index_data(&mut push_decoder)? {
889 NeedsIndexData::No(page_index) => return Ok(page_index),
890 NeedsIndexData::Yes(range) => range,
891 };
892
893 let bytes = match &remainder {
894 Some((remainder_start, remainder)) if *remainder_start as u64 <= range.start => {
895 let remainder_start = *remainder_start as u64;
896 let offset = usize::try_from(range.start - remainder_start)?;
897 let end = usize::try_from(range.end - remainder_start)?;
898 if end > remainder.len() {
899 return Err(general_err!(
900 "Corrupted parquet file: index data range ({:?}) exceeds remainder length ({})",
901 range,
902 remainder.len()
903 ));
904 }
905 remainder.slice(offset..end)
906 }
907 _ => fetch.fetch(range.clone()).await?,
909 };
910
911 if bytes.len() as u64 != range.end - range.start {
913 return Err(general_err!(
914 "Corrupted parquet file: index data length mismatch, expected {}, got {}",
915 range.end - range.start,
916 bytes.len()
917 ));
918 }
919 push_decoder.push_range(range.clone(), bytes)?;
920 parse_page_index_data(&mut push_decoder)
921 }
922
923 fn parse_metadata<R: ChunkReader>(&mut self, chunk_reader: &R) -> Result<ParquetMetaData> {
926 let file_size = chunk_reader.len();
928 if file_size < (FOOTER_SIZE as u64) {
929 return Err(ParquetError::NeedMoreData(FOOTER_SIZE));
930 }
931
932 let mut footer = [0_u8; FOOTER_SIZE];
933 chunk_reader
934 .get_read(file_size - FOOTER_SIZE as u64)?
935 .read_exact(&mut footer)?;
936
937 let footer = FooterTail::try_new(&footer)?;
938 let metadata_len = footer.metadata_length();
939 let footer_metadata_len = FOOTER_SIZE + metadata_len;
940 self.metadata_size = Some(footer_metadata_len);
941
942 if footer_metadata_len as u64 > file_size {
943 return Err(ParquetError::NeedMoreData(footer_metadata_len));
944 }
945
946 let start = file_size - footer_metadata_len as u64;
947 let bytes = chunk_reader.get_bytes(start, metadata_len)?;
948 self.decode_footer_metadata(bytes, file_size, footer)
949 }
950
951 pub fn metadata_size(&self) -> Option<usize> {
954 self.metadata_size
955 }
956
957 #[cfg(all(feature = "async", feature = "arrow"))]
961 fn get_prefetch_size(&self) -> usize {
962 if let Some(prefetch) = self.prefetch_hint
963 && prefetch > FOOTER_SIZE
964 {
965 return prefetch;
966 }
967 FOOTER_SIZE
968 }
969
970 #[cfg(all(feature = "async", feature = "arrow"))]
971 async fn load_metadata<F: MetadataFetch>(
972 &self,
973 fetch: &mut F,
974 file_size: u64,
975 ) -> Result<(ParquetMetaData, Option<(usize, Bytes)>)> {
976 let prefetch = self.get_prefetch_size() as u64;
977
978 if file_size < FOOTER_SIZE as u64 {
979 return Err(eof_err!("file size of {} is less than footer", file_size));
980 }
981
982 let footer_start = file_size.saturating_sub(prefetch);
986
987 let suffix = fetch.fetch(footer_start..file_size).await?;
988 let suffix_len = suffix.len();
989 let fetch_len = (file_size - footer_start)
990 .try_into()
991 .expect("footer size should never be larger than u32");
992 if suffix_len < fetch_len {
993 return Err(eof_err!(
994 "metadata requires {} bytes, but could only read {}",
995 fetch_len,
996 suffix_len
997 ));
998 }
999
1000 let mut footer = [0; FOOTER_SIZE];
1001 footer.copy_from_slice(&suffix[suffix_len - FOOTER_SIZE..suffix_len]);
1002
1003 let footer = FooterTail::try_new(&footer)?;
1004 let length = footer.metadata_length();
1005
1006 if file_size < (length + FOOTER_SIZE) as u64 {
1007 return Err(eof_err!(
1008 "file size of {} is less than footer + metadata {}",
1009 file_size,
1010 length + FOOTER_SIZE
1011 ));
1012 }
1013
1014 if length > suffix_len - FOOTER_SIZE {
1016 let metadata_start = file_size - (length + FOOTER_SIZE) as u64;
1017 let meta = fetch
1018 .fetch(metadata_start..(file_size - FOOTER_SIZE as u64))
1019 .await?;
1020 Ok((self.decode_footer_metadata(meta, file_size, footer)?, None))
1021 } else {
1022 let metadata_start = (file_size - (length + FOOTER_SIZE) as u64 - footer_start)
1023 .try_into()
1024 .expect("metadata length should never be larger than u32");
1025 let slice = suffix.slice(metadata_start..suffix_len - FOOTER_SIZE);
1026 Ok((
1027 self.decode_footer_metadata(slice, file_size, footer)?,
1028 Some((footer_start as usize, suffix.slice(..metadata_start))),
1029 ))
1030 }
1031 }
1032
1033 #[cfg(all(feature = "async", feature = "arrow"))]
1034 async fn load_metadata_via_suffix<F: MetadataSuffixFetch>(
1038 &self,
1039 fetch: &mut F,
1040 ) -> Result<ParquetMetaData> {
1041 let prefetch = self.get_prefetch_size();
1042
1043 let suffix = fetch.fetch_suffix(prefetch).await?;
1044 let suffix_len = suffix.len();
1045
1046 if suffix_len < FOOTER_SIZE {
1047 return Err(eof_err!(
1048 "footer metadata requires {} bytes, but could only read {}",
1049 FOOTER_SIZE,
1050 suffix_len
1051 ));
1052 }
1053
1054 let mut footer = [0; FOOTER_SIZE];
1055 footer.copy_from_slice(&suffix[suffix_len - FOOTER_SIZE..suffix_len]);
1056
1057 let footer = FooterTail::try_new(&footer)?;
1058 let length = footer.metadata_length();
1059 let file_size = (length + FOOTER_SIZE) as u64;
1062
1063 let metadata_offset = length + FOOTER_SIZE;
1065 if length > suffix_len - FOOTER_SIZE {
1066 let meta = fetch.fetch_suffix(metadata_offset).await?;
1067
1068 if meta.len() < metadata_offset {
1069 return Err(eof_err!(
1070 "metadata requires {} bytes, but could only read {}",
1071 metadata_offset,
1072 meta.len()
1073 ));
1074 }
1075
1076 let meta = meta.slice(0..length);
1078 Ok(self.decode_footer_metadata(meta, file_size, footer)?)
1079 } else {
1080 let metadata_start = suffix_len - metadata_offset;
1081 let slice = suffix.slice(metadata_start..suffix_len - FOOTER_SIZE);
1082 Ok(self.decode_footer_metadata(slice, file_size, footer)?)
1083 }
1084 }
1085
1086 pub(crate) fn decode_footer_metadata(
1100 &self,
1101 buf: Bytes,
1102 file_size: u64,
1103 footer_tail: FooterTail,
1104 ) -> Result<ParquetMetaData> {
1105 let ending_offset = file_size.checked_sub(FOOTER_SIZE as u64).ok_or_else(|| {
1110 general_err!(
1111 "file size {file_size} is smaller than footer size {}",
1112 FOOTER_SIZE
1113 )
1114 })?;
1115
1116 let starting_offset = ending_offset.checked_sub(buf.len() as u64).ok_or_else(|| {
1117 general_err!(
1118 "file size {file_size} is smaller than buffer size {} + footer size {}",
1119 buf.len(),
1120 FOOTER_SIZE
1121 )
1122 })?;
1123
1124 let range = starting_offset..ending_offset;
1125
1126 let push_decoder =
1127 ParquetMetaDataPushDecoder::try_new_with_footer_tail(file_size, footer_tail)?
1128 .with_page_index_policy(PageIndexPolicy::Skip)
1130 .with_metadata_options(self.metadata_options.clone());
1131
1132 let mut push_decoder = self.prepare_push_decoder(push_decoder);
1133 push_decoder.push_range(range, buf)?;
1134 match push_decoder.try_decode()? {
1135 DecodeResult::Data(metadata) => Ok(metadata),
1136 DecodeResult::Finished => Err(general_err!(
1137 "could not parse parquet metadata -- previously finished"
1138 )),
1139 DecodeResult::NeedsData(ranges) => Err(general_err!(
1140 "could not parse parquet metadata, needs ranges {:?}",
1141 ranges
1142 )),
1143 }
1144 }
1145
1146 #[cfg(feature = "encryption")]
1148 fn prepare_push_decoder(
1149 &self,
1150 push_decoder: ParquetMetaDataPushDecoder,
1151 ) -> ParquetMetaDataPushDecoder {
1152 push_decoder.with_file_decryption_properties(
1153 self.file_decryption_properties
1154 .as_ref()
1155 .map(std::sync::Arc::clone),
1156 )
1157 }
1158 #[cfg(not(feature = "encryption"))]
1159 fn prepare_push_decoder(
1160 &self,
1161 push_decoder: ParquetMetaDataPushDecoder,
1162 ) -> ParquetMetaDataPushDecoder {
1163 push_decoder
1164 }
1165
1166 pub fn decode_metadata(buf: &[u8]) -> Result<ParquetMetaData> {
1174 decode_metadata(buf, None)
1175 }
1176
1177 pub fn decode_metadata_with_options(
1182 buf: &[u8],
1183 options: Option<&ParquetMetaDataOptions>,
1184 ) -> Result<ParquetMetaData> {
1185 decode_metadata(buf, options)
1186 }
1187
1188 pub fn decode_schema(buf: &[u8]) -> Result<Arc<SchemaDescriptor>> {
1191 Ok(Arc::new(parquet_schema_from_bytes(buf)?))
1192 }
1193}
1194
1195enum NeedsIndexData<T> {
1198 No(T),
1200 Yes(Range<u64>),
1202}
1203
1204fn needs_page_index_data(
1207 push_decoder: &mut ParquetMetaDataPushDecoder,
1208) -> Result<NeedsIndexData<(ParquetMetaData, Option<PageIndex>)>> {
1209 match push_decoder.try_decode_page_index()? {
1210 DecodeResult::NeedsData(ranges) => {
1211 let range = ranges
1212 .into_iter()
1213 .reduce(|a, b| a.start.min(b.start)..a.end.max(b.end))
1214 .ok_or_else(|| general_err!("Internal error: no ranges provided"))?;
1215 Ok(NeedsIndexData::Yes(range))
1216 }
1217 DecodeResult::Data(page_index) => Ok(NeedsIndexData::No(page_index)),
1218 DecodeResult::Finished => Err(general_err!("Internal error: decoder was finished")),
1219 }
1220}
1221
1222fn parse_page_index_data(
1225 push_decoder: &mut ParquetMetaDataPushDecoder,
1226) -> Result<(ParquetMetaData, Option<PageIndex>)> {
1227 match push_decoder.try_decode_page_index()? {
1228 DecodeResult::NeedsData(_) => Err(general_err!(
1229 "Internal error: decoder still needs data after reading required range"
1230 )),
1231 DecodeResult::Data(page_index) => Ok(page_index),
1232 DecodeResult::Finished => Err(general_err!("Internal error: decoder was finished")),
1233 }
1234}
1235
1236#[cfg(feature = "arrow")]
1241fn dictionary_page_byte_range(
1242 metadata: &ParquetMetaData,
1243 row_group_idx: usize,
1244 column_idx: usize,
1245) -> Result<Option<Range<u64>>> {
1246 let column_metadata = metadata.row_group(row_group_idx).column(column_idx);
1247 let column_descriptor = column_metadata.column_descr();
1248
1249 if column_descriptor.physical_type() != crate::basic::Type::BYTE_ARRAY {
1250 return Ok(None);
1251 }
1252 let Some(start) = column_metadata.dictionary_page_offset() else {
1253 return Ok(None);
1254 };
1255 let start: u64 = start
1256 .try_into()
1257 .map_err(|_| ParquetError::General("Dictionary page offset is invalid".to_string()))?;
1258 let end: u64 = column_metadata
1259 .data_page_offset()
1260 .try_into()
1261 .map_err(|_| ParquetError::General("Data page offset is invalid".to_string()))?;
1262 if end < start {
1263 return Err(ParquetError::General(
1264 "Data page offset precedes dictionary page offset".to_string(),
1265 ));
1266 }
1267
1268 Ok(Some(start..end))
1269}
1270
1271#[cfg(test)]
1272mod tests {
1273 use super::*;
1274 use crate::file::reader::Length;
1275 use crate::util::test_common::file_util::get_test_file;
1276 use std::ops::Range;
1277
1278 #[test]
1279 fn test_parse_metadata_size_smaller_than_footer() {
1280 let test_file = tempfile::tempfile().unwrap();
1281 let err = ParquetMetaDataReader::new()
1282 .parse_metadata(&test_file)
1283 .unwrap_err();
1284 assert!(matches!(err, ParquetError::NeedMoreData(FOOTER_SIZE)));
1285 }
1286
1287 #[test]
1288 fn test_parse_metadata_corrupt_footer() {
1289 let data = Bytes::from(vec![1, 2, 3, 4, 5, 6, 7, 8]);
1290 let reader_result = ParquetMetaDataReader::new().parse_metadata(&data);
1291 assert_eq!(
1292 reader_result.unwrap_err().to_string(),
1293 "Parquet error: Invalid Parquet file. Corrupt footer"
1294 );
1295 }
1296
1297 #[test]
1298 fn test_parse_metadata_invalid_start() {
1299 let test_file = Bytes::from(vec![255, 0, 0, 0, b'P', b'A', b'R', b'1']);
1300 let err = ParquetMetaDataReader::new()
1301 .parse_metadata(&test_file)
1302 .unwrap_err();
1303 assert!(matches!(err, ParquetError::NeedMoreData(263)));
1304 }
1305
1306 #[test]
1307 #[cfg_attr(miri, ignore)] fn test_try_parse() {
1309 let file = get_test_file("alltypes_tiny_pages.parquet");
1310 let len = file.len();
1311
1312 let mut reader =
1313 ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
1314
1315 let bytes_for_range = |range: Range<u64>| {
1316 file.get_bytes(range.start, (range.end - range.start).try_into().unwrap())
1317 .unwrap()
1318 };
1319
1320 let bytes = bytes_for_range(0..len);
1322 reader.try_parse(&bytes).unwrap();
1323 let metadata = reader.finish().unwrap();
1324 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1325
1326 let bytes = bytes_for_range(320000..len);
1328 reader.try_parse_sized(&bytes, len).unwrap();
1329 let metadata = reader.finish().unwrap();
1330 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1331
1332 let bytes = bytes_for_range(323583..len);
1334 reader.try_parse_sized(&bytes, len).unwrap();
1335 let metadata = reader.finish().unwrap();
1336 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1337
1338 let bytes = bytes_for_range(323584..len);
1340 match reader.try_parse_sized(&bytes, len).unwrap_err() {
1342 ParquetError::NeedMoreData(needed) => {
1344 let bytes = bytes_for_range(len - needed as u64..len);
1345 reader.try_parse_sized(&bytes, len).unwrap();
1346 let metadata = reader.finish().unwrap();
1347 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1348 }
1349 _ => panic!("unexpected error"),
1350 }
1351
1352 let mut reader =
1354 ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
1355 let mut bytes = bytes_for_range(452505..len);
1356 loop {
1357 match reader.try_parse_sized(&bytes, len) {
1358 Ok(()) => break,
1359 Err(ParquetError::NeedMoreData(needed)) => {
1360 bytes = bytes_for_range(len - needed as u64..len);
1361 if reader.has_metadata() {
1362 reader.read_page_indexes_sized(&bytes, len).unwrap();
1363 break;
1364 }
1365 }
1366 _ => panic!("unexpected error"),
1367 }
1368 }
1369 let metadata = reader.finish().unwrap();
1370 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1371
1372 let bytes = bytes_for_range(323584..len);
1374 let reader_result = reader.try_parse_sized(&bytes, len - 323584).unwrap_err();
1375 assert_eq!(
1376 reader_result.to_string(),
1377 "EOF: Parquet file too small. Range 323583..452504 is beyond file bounds 130649"
1378 );
1379
1380 let mut reader = ParquetMetaDataReader::new();
1382 let bytes = bytes_for_range(452505..len);
1383 match reader.try_parse_sized(&bytes, len).unwrap_err() {
1385 ParquetError::NeedMoreData(needed) => {
1387 let bytes = bytes_for_range(len - needed as u64..len);
1388 reader.try_parse_sized(&bytes, len).unwrap();
1389 reader.finish().unwrap();
1390 }
1391 _ => panic!("unexpected error"),
1392 }
1393
1394 let reader_result = reader.try_parse(&bytes).unwrap_err();
1396 assert_eq!(
1397 reader_result.to_string(),
1398 "EOF: Parquet file too small. Size is 1728 but need 1729"
1399 );
1400
1401 let bytes = bytes_for_range(0..1000);
1403 let reader_result = reader.try_parse_sized(&bytes, len).unwrap_err();
1404 assert_eq!(
1405 reader_result.to_string(),
1406 "Parquet error: Invalid Parquet file. Corrupt footer"
1407 );
1408
1409 let bytes = bytes_for_range(452510..len);
1411 let reader_result = reader.try_parse_sized(&bytes, len - 452505).unwrap_err();
1412 assert_eq!(
1413 reader_result.to_string(),
1414 "EOF: Parquet file too small. Size is 1728 but need 1729"
1415 );
1416 }
1417
1418 #[test]
1419 #[cfg_attr(miri, ignore)] fn test_read_page_index_and_finish() {
1421 use crate::file::metadata::page_index::PageIndexProvider;
1422
1423 let file = get_test_file("alltypes_tiny_pages.parquet");
1424 let metadata = ParquetMetaDataReader::new()
1425 .with_page_index_policy(PageIndexPolicy::Required)
1426 .parse_and_finish(&file)
1427 .unwrap();
1428 let original_page_index = Arc::clone(metadata.page_index().unwrap());
1429
1430 let reader = ParquetMetaDataReader::new_with_metadata(metadata);
1431 let (metadata, page_index) = reader.read_page_index_and_finish(&file).unwrap();
1432 assert!(page_index.is_none());
1433 assert!(Arc::ptr_eq(
1434 metadata.page_index().unwrap(),
1435 &original_page_index
1436 ));
1437
1438 let reader = ParquetMetaDataReader::new_with_metadata(metadata)
1439 .with_page_index_policy(PageIndexPolicy::Required)
1440 .with_column_index_mask(ColumnChunkMask::columns([0]))
1441 .with_offset_index_mask(ColumnChunkMask::columns([0]));
1442 let (metadata, page_index) = reader.read_page_index_and_finish(&file).unwrap();
1443 let page_index = page_index.unwrap();
1444 assert!(Arc::ptr_eq(
1445 metadata.page_index().unwrap(),
1446 &original_page_index
1447 ));
1448
1449 assert!(page_index.column_index(0, 0).is_some());
1450 assert!(page_index.offset_index(0, 0).is_some());
1451 assert!(page_index.column_index(0, 1).is_none());
1452 assert!(page_index.offset_index(0, 1).is_none());
1453 }
1454
1455 #[test]
1456 fn test_chunk_mask() {
1457 let mask = ColumnChunkMask::row_groups_and_columns([0], [1]);
1458 assert!(mask.includes_row_group(0, 1));
1459 assert!(!mask.includes_row_group(1, 1));
1460 assert!(!mask.includes_column(0, 2));
1461 assert!(mask.includes_column(1, 2));
1462
1463 let mask = ColumnChunkMask::columns([]);
1464 assert!(!mask.includes_column(0, 3));
1465 assert_eq!(mask.column_indices(3).count(), 0);
1466
1467 let all = ColumnChunkMask::all();
1468 assert!(all.is_all());
1469 assert!(all.includes_column(2, 3));
1470 assert!(!all.includes_column(3, 3));
1471
1472 assert_eq!(all.row_group_indices(3).collect::<Vec<_>>(), [0, 1, 2]);
1473 assert_eq!(all.column_indices(2).collect::<Vec<_>>(), [0, 1]);
1474
1475 let none = ColumnChunkMask::none();
1476 assert_eq!(none.row_group_indices(3).count(), 0);
1477 assert_eq!(none.column_indices(3).count(), 0);
1478
1479 let sparse = ColumnChunkMask::row_groups_and_columns([0, 2, 4], [1, 3, 5]);
1480 assert_eq!(sparse.row_group_indices(3).collect::<Vec<_>>(), [0, 2]);
1481 assert_eq!(sparse.column_indices(4).collect::<Vec<_>>(), [1, 3]);
1482 }
1483}
1484
1485#[cfg(all(feature = "async", feature = "arrow", test))]
1486mod async_tests {
1487 use super::*;
1488
1489 use arrow::{array::Int32Array, datatypes::DataType};
1490 use arrow_array::RecordBatch;
1491 use arrow_schema::{Field, Schema};
1492 use bytes::Bytes;
1493 use futures::FutureExt;
1494 use futures::future::BoxFuture;
1495 use std::fs::File;
1496 use std::future::Future;
1497 use std::io::{Read, Seek, SeekFrom};
1498 use std::ops::Range;
1499 use std::sync::Arc;
1500 use std::sync::atomic::{AtomicUsize, Ordering};
1501 use tempfile::NamedTempFile;
1502
1503 use crate::arrow::ArrowWriter;
1504 use crate::file::properties::WriterProperties;
1505 use crate::file::reader::Length;
1506 use crate::util::test_common::file_util::get_test_file;
1507
1508 struct MetadataFetchFn<F>(F);
1509
1510 impl<F, Fut> MetadataFetch for MetadataFetchFn<F>
1511 where
1512 F: FnMut(Range<u64>) -> Fut + Send,
1513 Fut: Future<Output = Result<Bytes>> + Send,
1514 {
1515 fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
1516 async move { self.0(range).await }.boxed()
1517 }
1518 }
1519
1520 struct MetadataSuffixFetchFn<F1, F2>(F1, F2);
1521
1522 impl<F1, Fut, F2> MetadataFetch for MetadataSuffixFetchFn<F1, F2>
1523 where
1524 F1: FnMut(Range<u64>) -> Fut + Send,
1525 Fut: Future<Output = Result<Bytes>> + Send,
1526 F2: Send,
1527 {
1528 fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
1529 async move { self.0(range).await }.boxed()
1530 }
1531 }
1532
1533 impl<F1, Fut, F2> MetadataSuffixFetch for MetadataSuffixFetchFn<F1, F2>
1534 where
1535 F1: FnMut(Range<u64>) -> Fut + Send,
1536 F2: FnMut(usize) -> Fut + Send,
1537 Fut: Future<Output = Result<Bytes>> + Send,
1538 {
1539 fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result<Bytes>> {
1540 async move { self.1(suffix).await }.boxed()
1541 }
1542 }
1543
1544 fn read_range(file: &mut File, range: Range<u64>) -> Result<Bytes> {
1545 file.seek(SeekFrom::Start(range.start))?;
1546 let len = range.end - range.start;
1547 let mut buf = Vec::with_capacity(len.try_into().unwrap());
1548 file.take(len).read_to_end(&mut buf)?;
1549 Ok(buf.into())
1550 }
1551
1552 fn read_suffix(file: &mut File, suffix: usize) -> Result<Bytes> {
1553 let file_len = file.len();
1554 file.seek(SeekFrom::End(0 - suffix.min(file_len as _) as i64))?;
1556 let mut buf = Vec::with_capacity(suffix);
1557 file.take(suffix as _).read_to_end(&mut buf)?;
1558 Ok(buf.into())
1559 }
1560
1561 #[tokio::test]
1562 async fn test_load_page_index_and_finish() {
1563 use crate::file::metadata::page_index::PageIndexProvider;
1564
1565 let mut file = get_test_file("alltypes_tiny_pages.parquet");
1566 let metadata = ParquetMetaDataReader::new()
1567 .parse_and_finish(&file)
1568 .unwrap();
1569 let reader = ParquetMetaDataReader::new_with_metadata(metadata)
1570 .with_page_index_policy(PageIndexPolicy::Required)
1571 .with_column_index_mask(ColumnChunkMask::columns([0]))
1572 .with_offset_index_mask(ColumnChunkMask::columns([0]));
1573 let fetch = MetadataFetchFn(|range| futures::future::ready(read_range(&mut file, range)));
1574
1575 let page_index = reader
1576 .load_page_index_and_finish(fetch)
1577 .await
1578 .unwrap()
1579 .1
1580 .unwrap();
1581
1582 assert!(page_index.column_index(0, 0).is_some());
1583 assert!(page_index.offset_index(0, 0).is_some());
1584 assert!(page_index.column_index(0, 1).is_none());
1585 assert!(page_index.offset_index(0, 1).is_none());
1586 }
1587
1588 #[tokio::test]
1589 async fn test_simple() {
1590 let mut file = get_test_file("nulls.snappy.parquet");
1591 let len = file.len();
1592
1593 let expected = ParquetMetaDataReader::new()
1594 .parse_and_finish(&file)
1595 .unwrap();
1596 let expected = expected.file_metadata().schema();
1597 let fetch_count = AtomicUsize::new(0);
1598
1599 let mut fetch = |range| {
1600 fetch_count.fetch_add(1, Ordering::SeqCst);
1601 futures::future::ready(read_range(&mut file, range))
1602 };
1603
1604 let input = MetadataFetchFn(&mut fetch);
1605 let actual = ParquetMetaDataReader::new()
1606 .load_and_finish(input, len)
1607 .await
1608 .unwrap();
1609 assert_eq!(actual.file_metadata().schema(), expected);
1610 assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1611
1612 fetch_count.store(0, Ordering::SeqCst);
1614 let input = MetadataFetchFn(&mut fetch);
1615 let actual = ParquetMetaDataReader::new()
1616 .with_prefetch_hint(Some(7))
1617 .load_and_finish(input, len)
1618 .await
1619 .unwrap();
1620 assert_eq!(actual.file_metadata().schema(), expected);
1621 assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1622
1623 fetch_count.store(0, Ordering::SeqCst);
1625 let input = MetadataFetchFn(&mut fetch);
1626 let actual = ParquetMetaDataReader::new()
1627 .with_prefetch_hint(Some(10))
1628 .load_and_finish(input, len)
1629 .await
1630 .unwrap();
1631 assert_eq!(actual.file_metadata().schema(), expected);
1632 assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1633
1634 fetch_count.store(0, Ordering::SeqCst);
1636 let input = MetadataFetchFn(&mut fetch);
1637 let actual = ParquetMetaDataReader::new()
1638 .with_prefetch_hint(Some(500))
1639 .load_and_finish(input, len)
1640 .await
1641 .unwrap();
1642 assert_eq!(actual.file_metadata().schema(), expected);
1643 assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1644
1645 fetch_count.store(0, Ordering::SeqCst);
1647 let input = MetadataFetchFn(&mut fetch);
1648 let actual = ParquetMetaDataReader::new()
1649 .with_prefetch_hint(Some(428))
1650 .load_and_finish(input, len)
1651 .await
1652 .unwrap();
1653 assert_eq!(actual.file_metadata().schema(), expected);
1654 assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1655
1656 let input = MetadataFetchFn(&mut fetch);
1657 let err = ParquetMetaDataReader::new()
1658 .load_and_finish(input, 4)
1659 .await
1660 .unwrap_err()
1661 .to_string();
1662 assert_eq!(err, "EOF: file size of 4 is less than footer");
1663
1664 let input = MetadataFetchFn(&mut fetch);
1665 let err = ParquetMetaDataReader::new()
1666 .load_and_finish(input, 20)
1667 .await
1668 .unwrap_err()
1669 .to_string();
1670 assert_eq!(err, "Parquet error: Invalid Parquet file. Corrupt footer");
1671 }
1672
1673 #[tokio::test]
1674 async fn test_suffix() {
1675 let mut file = get_test_file("nulls.snappy.parquet");
1676 let mut file2 = file.try_clone().unwrap();
1677
1678 let expected = ParquetMetaDataReader::new()
1679 .parse_and_finish(&file)
1680 .unwrap();
1681 let expected = expected.file_metadata().schema();
1682 let fetch_count = AtomicUsize::new(0);
1683 let suffix_fetch_count = AtomicUsize::new(0);
1684
1685 let mut fetch = |range| {
1686 fetch_count.fetch_add(1, Ordering::SeqCst);
1687 futures::future::ready(read_range(&mut file, range))
1688 };
1689 let mut suffix_fetch = |suffix| {
1690 suffix_fetch_count.fetch_add(1, Ordering::SeqCst);
1691 futures::future::ready(read_suffix(&mut file2, suffix))
1692 };
1693
1694 let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1695 let actual = ParquetMetaDataReader::new()
1696 .load_via_suffix_and_finish(input)
1697 .await
1698 .unwrap();
1699 assert_eq!(actual.file_metadata().schema(), expected);
1700 assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1701 assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 2);
1702
1703 fetch_count.store(0, Ordering::SeqCst);
1705 suffix_fetch_count.store(0, Ordering::SeqCst);
1706 let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1707 let actual = ParquetMetaDataReader::new()
1708 .with_prefetch_hint(Some(7))
1709 .load_via_suffix_and_finish(input)
1710 .await
1711 .unwrap();
1712 assert_eq!(actual.file_metadata().schema(), expected);
1713 assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1714 assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 2);
1715
1716 fetch_count.store(0, Ordering::SeqCst);
1718 suffix_fetch_count.store(0, Ordering::SeqCst);
1719 let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1720 let actual = ParquetMetaDataReader::new()
1721 .with_prefetch_hint(Some(10))
1722 .load_via_suffix_and_finish(input)
1723 .await
1724 .unwrap();
1725 assert_eq!(actual.file_metadata().schema(), expected);
1726 assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1727 assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 2);
1728
1729 fetch_count.store(0, Ordering::SeqCst);
1731 suffix_fetch_count.store(0, Ordering::SeqCst);
1732 let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1733 let actual = ParquetMetaDataReader::new()
1734 .with_prefetch_hint(Some(500))
1735 .load_via_suffix_and_finish(input)
1736 .await
1737 .unwrap();
1738 assert_eq!(actual.file_metadata().schema(), expected);
1739 assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1740 assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 1);
1741
1742 fetch_count.store(0, Ordering::SeqCst);
1744 suffix_fetch_count.store(0, Ordering::SeqCst);
1745 let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1746 let actual = ParquetMetaDataReader::new()
1747 .with_prefetch_hint(Some(428))
1748 .load_via_suffix_and_finish(input)
1749 .await
1750 .unwrap();
1751 assert_eq!(actual.file_metadata().schema(), expected);
1752 assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1753 assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 1);
1754 }
1755
1756 #[cfg(feature = "encryption")]
1757 #[tokio::test]
1758 async fn test_suffix_with_encryption() {
1759 let mut file = get_test_file("uniform_encryption.parquet.encrypted");
1760 let mut file2 = file.try_clone().unwrap();
1761
1762 let mut fetch = |range| futures::future::ready(read_range(&mut file, range));
1763 let mut suffix_fetch = |suffix| futures::future::ready(read_suffix(&mut file2, suffix));
1764
1765 let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1766
1767 let key_code: &[u8] = b"0123456789012345";
1768 let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec())
1769 .build()
1770 .unwrap();
1771
1772 let expected = ParquetMetaDataReader::new()
1774 .with_decryption_properties(Some(decryption_properties))
1775 .load_via_suffix_and_finish(input)
1776 .await
1777 .unwrap();
1778 assert_eq!(expected.num_row_groups(), 1);
1779 }
1780
1781 #[tokio::test]
1782 async fn test_page_index() {
1783 let mut file = get_test_file("alltypes_tiny_pages.parquet");
1784 let len = file.len();
1785 let fetch_count = AtomicUsize::new(0);
1786 let mut fetch = |range| {
1787 fetch_count.fetch_add(1, Ordering::SeqCst);
1788 futures::future::ready(read_range(&mut file, range))
1789 };
1790
1791 let f = MetadataFetchFn(&mut fetch);
1792 let mut loader =
1793 ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
1794 loader.try_load(f, len).await.unwrap();
1795 assert_eq!(fetch_count.load(Ordering::SeqCst), 3);
1796 let metadata = loader.finish().unwrap();
1797 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1798
1799 fetch_count.store(0, Ordering::SeqCst);
1801 let f = MetadataFetchFn(&mut fetch);
1802 let mut loader = ParquetMetaDataReader::new()
1803 .with_page_index_policy(PageIndexPolicy::Required)
1804 .with_prefetch_hint(Some(1729));
1805 loader.try_load(f, len).await.unwrap();
1806 assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1807 let metadata = loader.finish().unwrap();
1808 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1809
1810 fetch_count.store(0, Ordering::SeqCst);
1812 let f = MetadataFetchFn(&mut fetch);
1813 let mut loader = ParquetMetaDataReader::new()
1814 .with_page_index_policy(PageIndexPolicy::Required)
1815 .with_prefetch_hint(Some(130649));
1816 loader.try_load(f, len).await.unwrap();
1817 assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1818 let metadata = loader.finish().unwrap();
1819 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1820
1821 fetch_count.store(0, Ordering::SeqCst);
1823 let f = MetadataFetchFn(&mut fetch);
1824 let metadata = ParquetMetaDataReader::new()
1825 .with_page_index_policy(PageIndexPolicy::Required)
1826 .with_prefetch_hint(Some(130650))
1827 .load_and_finish(f, len)
1828 .await
1829 .unwrap();
1830 assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1831 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1832
1833 fetch_count.store(0, Ordering::SeqCst);
1835 let f = MetadataFetchFn(&mut fetch);
1836 let metadata = ParquetMetaDataReader::new()
1837 .with_page_index_policy(PageIndexPolicy::Required)
1838 .with_prefetch_hint(Some((len - 1000) as usize)) .load_and_finish(f, len)
1840 .await
1841 .unwrap();
1842 assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1843 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1844
1845 fetch_count.store(0, Ordering::SeqCst);
1847 let f = MetadataFetchFn(&mut fetch);
1848 let metadata = ParquetMetaDataReader::new()
1849 .with_page_index_policy(PageIndexPolicy::Required)
1850 .with_prefetch_hint(Some(len as usize)) .load_and_finish(f, len)
1852 .await
1853 .unwrap();
1854 assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1855 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1856
1857 fetch_count.store(0, Ordering::SeqCst);
1859 let f = MetadataFetchFn(&mut fetch);
1860 let metadata = ParquetMetaDataReader::new()
1861 .with_page_index_policy(PageIndexPolicy::Required)
1862 .with_prefetch_hint(Some((len + 1000) as usize)) .load_and_finish(f, len)
1864 .await
1865 .unwrap();
1866 assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1867 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1868 }
1869
1870 #[tokio::test]
1871 async fn test_page_index_via_suffix_with_prefetch() {
1872 let mut file = get_test_file("alltypes_tiny_pages.parquet");
1873 let mut suffix_file = file.try_clone().unwrap();
1874 let len = file.len();
1875 let fetch_count = AtomicUsize::new(0);
1876 let suffix_fetch_count = AtomicUsize::new(0);
1877
1878 let mut fetch = |range| {
1879 fetch_count.fetch_add(1, Ordering::SeqCst);
1880 futures::future::ready(read_range(&mut file, range))
1881 };
1882 let mut suffix_fetch = |suffix| {
1883 suffix_fetch_count.fetch_add(1, Ordering::SeqCst);
1884 futures::future::ready(read_suffix(&mut suffix_file, suffix))
1885 };
1886
1887 let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1888 let metadata = ParquetMetaDataReader::new()
1889 .with_page_index_policy(PageIndexPolicy::Required)
1890 .with_prefetch_hint(Some((len - 1000) as usize))
1891 .load_via_suffix_and_finish(input)
1892 .await
1893 .unwrap();
1894
1895 assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 1);
1896 assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1897 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1898 }
1899
1900 fn write_parquet_file(offset_index_disabled: bool) -> Result<NamedTempFile> {
1901 let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1902 let batch = RecordBatch::try_new(
1903 schema.clone(),
1904 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
1905 )?;
1906
1907 let file = NamedTempFile::new().unwrap();
1908
1909 let props = WriterProperties::builder()
1911 .set_offset_index_disabled(offset_index_disabled)
1912 .build();
1913
1914 let mut writer = ArrowWriter::try_new(file.reopen()?, schema, Some(props))?;
1915 writer.write(&batch)?;
1916 writer.close()?;
1917
1918 Ok(file)
1919 }
1920
1921 fn read_and_check(file: &File, policy: PageIndexPolicy) -> Result<ParquetMetaData> {
1922 let mut reader = ParquetMetaDataReader::new().with_page_index_policy(policy);
1923 reader.try_parse(file)?;
1924 reader.finish()
1925 }
1926
1927 #[test]
1928 fn test_page_index_policy() {
1929 let f = write_parquet_file(false).unwrap();
1931 read_and_check(f.as_file(), PageIndexPolicy::Required).unwrap();
1932 read_and_check(f.as_file(), PageIndexPolicy::Optional).unwrap();
1933 read_and_check(f.as_file(), PageIndexPolicy::Skip).unwrap();
1934
1935 let f = write_parquet_file(true).unwrap();
1937 let res = read_and_check(f.as_file(), PageIndexPolicy::Required);
1938 assert!(matches!(
1939 res,
1940 Err(ParquetError::General(e)) if e == "missing offset index"
1941 ));
1942 read_and_check(f.as_file(), PageIndexPolicy::Optional).unwrap();
1943 read_and_check(f.as_file(), PageIndexPolicy::Skip).unwrap();
1944 }
1945}