1use crate::basic::{PageType, Type};
22use crate::bloom_filter::Sbbf;
23use crate::column::page::{Page, PageMetadata, PageReader};
24use crate::compression::{Codec, create_codec};
25#[cfg(feature = "encryption")]
26use crate::encryption::decrypt::{CryptoContext, read_and_decrypt};
27use crate::errors::{ParquetError, Result};
28use crate::file::metadata::page_index::RowGroupPageIndex;
29use crate::file::metadata::thrift::PageHeader;
30use crate::file::page_index::offset_index::PageLocation;
31use crate::file::statistics;
32use crate::file::{
33 metadata::*,
34 properties::{ReaderProperties, ReaderPropertiesPtr},
35 reader::*,
36};
37#[cfg(feature = "encryption")]
38use crate::parquet_thrift::ThriftSliceInputProtocol;
39use crate::parquet_thrift::{ReadThrift, ThriftReadInputProtocol};
40use crate::record::Row;
41use crate::record::reader::RowIter;
42use crate::schema::types::{SchemaDescPtr, Type as SchemaType};
43use bytes::Bytes;
44use std::collections::VecDeque;
45use std::{fs::File, io::Read, path::Path, sync::Arc};
46
47impl TryFrom<File> for SerializedFileReader<File> {
48 type Error = ParquetError;
49
50 fn try_from(file: File) -> Result<Self> {
51 Self::new(file)
52 }
53}
54
55impl TryFrom<&Path> for SerializedFileReader<File> {
56 type Error = ParquetError;
57
58 fn try_from(path: &Path) -> Result<Self> {
59 let file = File::open(path)?;
60 Self::try_from(file)
61 }
62}
63
64impl TryFrom<String> for SerializedFileReader<File> {
65 type Error = ParquetError;
66
67 fn try_from(path: String) -> Result<Self> {
68 Self::try_from(Path::new(&path))
69 }
70}
71
72impl TryFrom<&str> for SerializedFileReader<File> {
73 type Error = ParquetError;
74
75 fn try_from(path: &str) -> Result<Self> {
76 Self::try_from(Path::new(&path))
77 }
78}
79
80impl IntoIterator for SerializedFileReader<File> {
83 type Item = Result<Row>;
84 type IntoIter = RowIter<'static>;
85
86 fn into_iter(self) -> Self::IntoIter {
87 RowIter::from_file_into(Box::new(self))
88 }
89}
90
91pub struct SerializedFileReader<R: ChunkReader> {
96 chunk_reader: Arc<R>,
97 metadata: Arc<ParquetMetaData>,
98 props: ReaderPropertiesPtr,
99}
100
101pub type ReadGroupPredicate = Box<dyn FnMut(&RowGroupMetaData, usize) -> bool>;
105
106#[derive(Default)]
110pub struct ReadOptionsBuilder {
111 predicates: Vec<ReadGroupPredicate>,
112 enable_page_index: bool,
113 props: Option<ReaderProperties>,
114 metadata_options: ParquetMetaDataOptions,
115}
116
117impl ReadOptionsBuilder {
118 pub fn new() -> Self {
120 Self::default()
121 }
122
123 pub fn with_predicate(mut self, predicate: ReadGroupPredicate) -> Self {
126 self.predicates.push(predicate);
127 self
128 }
129
130 pub fn with_range(mut self, start: i64, end: i64) -> Self {
137 assert!(start < end);
138 let predicate = move |rg: &RowGroupMetaData, _: usize| {
139 let mid = get_midpoint_offset(rg);
140 mid >= start && mid < end
141 };
142 self.predicates.push(Box::new(predicate));
143 self
144 }
145
146 pub fn with_page_index(mut self) -> Self {
151 self.enable_page_index = true;
152 self
153 }
154
155 pub fn with_reader_properties(mut self, properties: ReaderProperties) -> Self {
157 self.props = Some(properties);
158 self
159 }
160
161 pub fn with_parquet_schema(mut self, schema: SchemaDescPtr) -> Self {
164 self.metadata_options.set_schema(schema);
165 self
166 }
167
168 pub fn with_encoding_stats_as_mask(mut self, val: bool) -> Self {
177 self.metadata_options.set_encoding_stats_as_mask(val);
178 self
179 }
180
181 pub fn with_encoding_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
186 self.metadata_options.set_encoding_stats_policy(policy);
187 self
188 }
189
190 pub fn with_column_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
195 self.metadata_options.set_column_stats_policy(policy);
196 self
197 }
198
199 pub fn with_size_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
204 self.metadata_options.set_size_stats_policy(policy);
205 self
206 }
207
208 pub fn build(self) -> ReadOptions {
210 let props = self
211 .props
212 .unwrap_or_else(|| ReaderProperties::builder().build());
213 ReadOptions {
214 predicates: self.predicates,
215 enable_page_index: self.enable_page_index,
216 props,
217 metadata_options: self.metadata_options,
218 }
219 }
220}
221
222pub struct ReadOptions {
227 predicates: Vec<ReadGroupPredicate>,
228 enable_page_index: bool,
229 props: ReaderProperties,
230 metadata_options: ParquetMetaDataOptions,
231}
232
233impl<R: 'static + ChunkReader> SerializedFileReader<R> {
234 pub fn new(chunk_reader: R) -> Result<Self> {
237 let metadata = ParquetMetaDataReader::new().parse_and_finish(&chunk_reader)?;
238 let props = Arc::new(ReaderProperties::builder().build());
239 Ok(Self {
240 chunk_reader: Arc::new(chunk_reader),
241 metadata: Arc::new(metadata),
242 props,
243 })
244 }
245
246 pub fn new_with_options(chunk_reader: R, options: ReadOptions) -> Result<Self> {
249 let mut metadata = ParquetMetaDataReader::new()
250 .with_metadata_options(Some(options.metadata_options.clone()))
251 .parse_and_finish(&chunk_reader)?;
252 let mut predicates = options.predicates;
253
254 if !predicates.is_empty() {
255 let row_groups = metadata.row_groups();
256 let mut filtered = None;
257
258 for (i, rg_meta) in row_groups.iter().enumerate() {
260 let keep = predicates.iter_mut().all(|predicate| predicate(rg_meta, i));
261 if !keep {
262 filtered.get_or_insert_with(|| row_groups[..i].to_vec());
263 } else if let Some(kept) = &mut filtered {
264 kept.push(rg_meta.clone());
265 }
266 }
267
268 if let Some(row_groups) = filtered {
269 metadata = metadata.into_builder().set_row_groups(row_groups).build();
270 }
271 }
272
273 if options.enable_page_index {
275 let mut reader = ParquetMetaDataReader::new_with_metadata(metadata)
276 .with_page_index_policy(PageIndexPolicy::Required);
277 reader.read_page_indexes(&chunk_reader)?;
278 metadata = reader.finish()?;
279 }
280
281 Ok(Self {
282 chunk_reader: Arc::new(chunk_reader),
283 metadata: Arc::new(metadata),
284 props: Arc::new(options.props),
285 })
286 }
287}
288
289fn get_midpoint_offset(meta: &RowGroupMetaData) -> i64 {
291 let col = meta.column(0);
292 let mut offset = col.data_page_offset();
293 if let Some(dic_offset) = col.dictionary_page_offset()
294 && offset > dic_offset
295 {
296 offset = dic_offset
297 }
298 offset + meta.compressed_size() / 2
299}
300
301impl<R: 'static + ChunkReader> FileReader for SerializedFileReader<R> {
302 fn metadata(&self) -> &ParquetMetaData {
303 &self.metadata
304 }
305
306 fn num_row_groups(&self) -> usize {
307 self.metadata.num_row_groups()
308 }
309
310 fn get_row_group(&self, i: usize) -> Result<Box<dyn RowGroupReader + '_>> {
311 let row_group_metadata = self.metadata.row_group(i);
312 let props = Arc::clone(&self.props);
314 let f = Arc::clone(&self.chunk_reader);
315 let page_index = self.metadata.page_index_for_row_group(i);
316 Ok(Box::new(SerializedRowGroupReader::new(
317 f,
318 row_group_metadata,
319 page_index,
320 props,
321 )?))
322 }
323
324 fn get_row_iter(&self, projection: Option<SchemaType>) -> Result<RowIter<'_>> {
325 RowIter::from_file(projection, self)
326 }
327}
328
329pub struct SerializedRowGroupReader<'a, R: ChunkReader> {
331 chunk_reader: Arc<R>,
332 metadata: &'a RowGroupMetaData,
333 page_index: RowGroupPageIndex,
334 props: ReaderPropertiesPtr,
335 bloom_filters: Vec<Option<Sbbf>>,
336}
337
338impl<'a, R: ChunkReader> SerializedRowGroupReader<'a, R> {
339 pub fn new(
341 chunk_reader: Arc<R>,
342 metadata: &'a RowGroupMetaData,
343 page_index: RowGroupPageIndex,
344 props: ReaderPropertiesPtr,
345 ) -> Result<Self> {
346 let bloom_filters = if props.read_bloom_filter() {
347 metadata
348 .columns()
349 .iter()
350 .map(|col| Sbbf::read_from_column_chunk(col, &*chunk_reader))
351 .collect::<Result<Vec<_>>>()?
352 } else {
353 std::iter::repeat_n(None, metadata.columns().len()).collect()
354 };
355 Ok(Self {
356 chunk_reader,
357 metadata,
358 page_index,
359 props,
360 bloom_filters,
361 })
362 }
363}
364
365impl<R: 'static + ChunkReader> RowGroupReader for SerializedRowGroupReader<'_, R> {
366 fn metadata(&self) -> &RowGroupMetaData {
367 self.metadata
368 }
369
370 fn num_columns(&self) -> usize {
371 self.metadata.num_columns()
372 }
373
374 fn get_column_page_reader(&self, i: usize) -> Result<Box<dyn PageReader>> {
376 let col = self.metadata.column(i);
377
378 let page_locations = self.page_index.page_locations(i).cloned();
380
381 let props = Arc::clone(&self.props);
382 Ok(Box::new(SerializedPageReader::new_with_properties(
383 Arc::clone(&self.chunk_reader),
384 col,
385 usize::try_from(self.metadata.num_rows())?,
386 page_locations,
387 props,
388 )?))
389 }
390
391 fn get_column_bloom_filter(&self, i: usize) -> Option<&Sbbf> {
393 self.bloom_filters[i].as_ref()
394 }
395
396 fn get_row_iter(&self, projection: Option<SchemaType>) -> Result<RowIter<'_>> {
397 RowIter::from_row_group(projection, self)
398 }
399}
400
401pub(crate) fn decode_page(
403 page_header: PageHeader,
404 buffer: Bytes,
405 physical_type: Type,
406 decompressor: Option<&mut Box<dyn Codec>>,
407) -> Result<Page> {
408 #[cfg(feature = "crc")]
410 if let Some(expected_crc) = page_header.crc {
411 let crc = crc32fast::hash(&buffer);
412 if crc != expected_crc as u32 {
413 return Err(general_err!("Page CRC checksum mismatch"));
414 }
415 }
416
417 let (offset, can_decompress): (usize, bool) = match page_header.data_page_header_v2 {
424 Some(ref header_v2) => {
425 if header_v2.definition_levels_byte_length < 0
426 || header_v2.repetition_levels_byte_length < 0
427 || header_v2.definition_levels_byte_length + header_v2.repetition_levels_byte_length
428 > page_header.uncompressed_page_size
429 {
430 return Err(general_err!(
431 "DataPage v2 header contains implausible values \
432 for definition_levels_byte_length ({}) \
433 and repetition_levels_byte_length ({}) \
434 given DataPage header provides uncompressed_page_size ({})",
435 header_v2.definition_levels_byte_length,
436 header_v2.repetition_levels_byte_length,
437 page_header.uncompressed_page_size
438 ));
439 }
440 (
441 usize::try_from(
442 header_v2.definition_levels_byte_length
443 + header_v2.repetition_levels_byte_length,
444 )?,
445 header_v2.is_compressed.unwrap_or(true),
447 )
448 }
449 None => (0, true),
450 };
451
452 let buffer = match decompressor {
453 Some(decompressor) if can_decompress => {
454 let uncompressed_page_size = usize::try_from(page_header.uncompressed_page_size)?;
455 if offset > buffer.len() || offset > uncompressed_page_size {
456 return Err(general_err!("Invalid page header"));
457 }
458 let decompressed_size = uncompressed_page_size - offset;
459 let mut decompressed = Vec::with_capacity(uncompressed_page_size);
460 decompressed.extend_from_slice(&buffer[..offset]);
461 if decompressed_size > 0 {
464 let compressed = &buffer[offset..];
465 decompressor.decompress(compressed, &mut decompressed, Some(decompressed_size))?;
466 }
467
468 if decompressed.len() != uncompressed_page_size {
469 return Err(general_err!(
470 "Actual decompressed size doesn't match the expected one ({} vs {})",
471 decompressed.len(),
472 uncompressed_page_size
473 ));
474 }
475
476 Bytes::from(decompressed)
477 }
478 _ => buffer,
479 };
480
481 let result = match page_header.r#type {
482 PageType::DICTIONARY_PAGE => {
483 let dict_header = page_header.dictionary_page_header.as_ref().ok_or_else(|| {
484 ParquetError::General("Missing dictionary page header".to_string())
485 })?;
486 let is_sorted = dict_header.is_sorted.unwrap_or(false);
487 Page::DictionaryPage {
488 buf: buffer,
489 num_values: dict_header.num_values.try_into()?,
490 encoding: dict_header.encoding,
491 is_sorted,
492 }
493 }
494 PageType::DATA_PAGE => {
495 let header = page_header
496 .data_page_header
497 .ok_or_else(|| ParquetError::General("Missing V1 data page header".to_string()))?;
498 Page::DataPage {
499 buf: buffer,
500 num_values: header.num_values.try_into()?,
501 encoding: header.encoding,
502 def_level_encoding: header.definition_level_encoding,
503 rep_level_encoding: header.repetition_level_encoding,
504 statistics: statistics::from_thrift_page_stats(physical_type, header.statistics)?,
505 }
506 }
507 PageType::DATA_PAGE_V2 => {
508 let header = page_header
509 .data_page_header_v2
510 .ok_or_else(|| ParquetError::General("Missing V2 data page header".to_string()))?;
511 let is_compressed = header.is_compressed.unwrap_or(true);
512 Page::DataPageV2 {
513 buf: buffer,
514 num_values: header.num_values.try_into()?,
515 encoding: header.encoding,
516 num_nulls: header.num_nulls.try_into()?,
517 num_rows: header.num_rows.try_into()?,
518 def_levels_byte_len: header.definition_levels_byte_length.try_into()?,
519 rep_levels_byte_len: header.repetition_levels_byte_length.try_into()?,
520 is_compressed,
521 statistics: statistics::from_thrift_page_stats(physical_type, header.statistics)?,
522 }
523 }
524 PageType::INDEX_PAGE => {
525 return Err(general_err!(
527 "Page type {:?} is not supported",
528 page_header.r#type
529 ));
530 }
531 };
532
533 Ok(result)
534}
535
536enum DeferredDictionaryPage {
558 Header {
561 data_start: u64,
562 header: Box<PageHeader>,
563 },
564 Location(PageLocation),
567}
568
569enum SerializedPageReaderState {
570 Values {
571 offset: u64,
574
575 remaining_bytes: u64,
578
579 next_page_header: Option<Box<PageHeader>>,
581
582 page_index: usize,
584
585 require_dictionary: bool,
587 },
588 Pages {
589 page_locations: VecDeque<PageLocation>,
591 dictionary_page: Option<PageLocation>,
593 total_rows: usize,
595 page_index: usize,
597 },
598}
599
600#[derive(Default)]
601pub(crate) struct SerializedPageReaderContext {
602 pub(crate) read_stats: bool,
604 #[cfg(feature = "encryption")]
606 pub(crate) crypto_context: Option<Arc<CryptoContext>>,
607}
608
609pub struct SerializedPageReader<R: ChunkReader> {
611 reader: Arc<R>,
613
614 deferred_dictionary: Option<DeferredDictionaryPage>,
617
618 decompressor: Option<Box<dyn Codec>>,
620
621 physical_type: Type,
623
624 state: SerializedPageReaderState,
625
626 context: SerializedPageReaderContext,
627}
628
629impl<R: ChunkReader> SerializedPageReader<R> {
630 pub fn new(
632 reader: Arc<R>,
633 column_chunk_metadata: &ColumnChunkMetaData,
634 total_rows: usize,
635 page_locations: Option<Vec<PageLocation>>,
636 ) -> Result<Self> {
637 let props = Arc::new(ReaderProperties::builder().build());
638 SerializedPageReader::new_with_properties(
639 reader,
640 column_chunk_metadata,
641 total_rows,
642 page_locations,
643 props,
644 )
645 }
646
647 #[cfg(all(feature = "arrow", not(feature = "encryption")))]
649 pub(crate) fn add_crypto_context(
650 self,
651 _rg_idx: usize,
652 _column_idx: usize,
653 _parquet_meta_data: &ParquetMetaData,
654 _column_chunk_metadata: &ColumnChunkMetaData,
655 ) -> Result<SerializedPageReader<R>> {
656 Ok(self)
657 }
658
659 #[cfg(feature = "encryption")]
661 pub(crate) fn add_crypto_context(
662 mut self,
663 rg_idx: usize,
664 column_idx: usize,
665 parquet_meta_data: &ParquetMetaData,
666 column_chunk_metadata: &ColumnChunkMetaData,
667 ) -> Result<SerializedPageReader<R>> {
668 let Some(file_decryptor) = parquet_meta_data.file_decryptor() else {
669 return Ok(self);
670 };
671 let Some(crypto_metadata) = column_chunk_metadata.crypto_metadata() else {
672 return Ok(self);
673 };
674 let crypto_context =
675 CryptoContext::for_column(file_decryptor, crypto_metadata, rg_idx, column_idx)?;
676 self.context.crypto_context = Some(Arc::new(crypto_context));
677 Ok(self)
678 }
679
680 pub fn new_with_properties(
682 reader: Arc<R>,
683 meta: &ColumnChunkMetaData,
684 total_rows: usize,
685 page_locations: Option<Vec<PageLocation>>,
686 props: ReaderPropertiesPtr,
687 ) -> Result<Self> {
688 let decompressor = create_codec(meta.compression(), props.codec_options())?;
689 let (start, len) = meta.byte_range();
690
691 let state = match page_locations {
692 Some(locations) => {
693 let dictionary_page = match locations.first() {
696 Some(dict_offset) if dict_offset.offset as u64 != start => Some(PageLocation {
697 offset: start as i64,
698 compressed_page_size: (dict_offset.offset as u64 - start) as i32,
699 first_row_index: 0,
700 }),
701 _ => None,
702 };
703
704 SerializedPageReaderState::Pages {
705 page_locations: locations.into(),
706 dictionary_page,
707 total_rows,
708 page_index: 0,
709 }
710 }
711 None => SerializedPageReaderState::Values {
712 offset: start,
713 remaining_bytes: len,
714 next_page_header: None,
715 page_index: 0,
716 require_dictionary: meta.dictionary_page_offset().is_some(),
717 },
718 };
719 let mut context = SerializedPageReaderContext::default();
720 if props.read_page_stats() {
721 context.read_stats = true;
722 }
723 Ok(Self {
724 reader,
725 deferred_dictionary: None,
726 decompressor,
727 state,
728 physical_type: meta.column_type(),
729 context,
730 })
731 }
732
733 #[cfg(test)]
739 fn peek_next_page_offset(&mut self) -> Result<Option<u64>> {
740 match &mut self.state {
741 SerializedPageReaderState::Values {
742 offset,
743 remaining_bytes,
744 next_page_header,
745 page_index,
746 require_dictionary,
747 } => {
748 loop {
749 if *remaining_bytes == 0 {
750 return Ok(None);
751 }
752 return if let Some(header) = next_page_header.as_ref() {
753 if let Ok(_page_meta) = PageMetadata::try_from(&**header) {
754 Ok(Some(*offset))
755 } else {
756 *next_page_header = None;
758 continue;
759 }
760 } else {
761 let mut read = self.reader.get_read(*offset)?;
762 let (header_len, header) = Self::read_page_header_len(
763 &self.context,
764 &mut read,
765 *page_index,
766 *require_dictionary,
767 )?;
768 *offset += header_len as u64;
769 *remaining_bytes -= header_len as u64;
770 let page_meta = if let Ok(_page_meta) = PageMetadata::try_from(&header) {
771 Ok(Some(*offset))
772 } else {
773 continue;
775 };
776 *next_page_header = Some(Box::new(header));
777 page_meta
778 };
779 }
780 }
781 SerializedPageReaderState::Pages {
782 page_locations,
783 dictionary_page,
784 ..
785 } => {
786 if let Some(page) = dictionary_page {
787 Ok(Some(page.offset as u64))
788 } else if let Some(page) = page_locations.front() {
789 Ok(Some(page.offset as u64))
790 } else {
791 Ok(None)
792 }
793 }
794 }
795 }
796
797 fn read_page_header_len<T: Read>(
798 context: &SerializedPageReaderContext,
799 input: &mut T,
800 page_index: usize,
801 dictionary_page: bool,
802 ) -> Result<(usize, PageHeader)> {
803 struct TrackedRead<R> {
805 inner: R,
806 bytes_read: usize,
807 }
808
809 impl<R: Read> Read for TrackedRead<R> {
810 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
811 let v = self.inner.read(buf)?;
812 self.bytes_read += v;
813 Ok(v)
814 }
815 }
816
817 let mut tracked = TrackedRead {
818 inner: input,
819 bytes_read: 0,
820 };
821 let header = context.read_page_header(&mut tracked, page_index, dictionary_page)?;
822 Ok((tracked.bytes_read, header))
823 }
824}
825
826pub(crate) fn read_page_header_len_from_bytes(
834 context: &SerializedPageReaderContext,
835 buffer: &[u8],
836 page_index: usize,
837 dictionary_page: bool,
838) -> Result<(usize, PageHeader)> {
839 let mut input = std::io::Cursor::new(buffer);
840 let header = context.read_page_header(&mut input, page_index, dictionary_page)?;
841 let header_len = input.position() as usize;
842 Ok((header_len, header))
843}
844
845#[cfg(not(feature = "encryption"))]
846impl SerializedPageReaderContext {
847 pub(crate) fn read_page_header<T: Read>(
848 &self,
849 input: &mut T,
850 _page_index: usize,
851 _dictionary_page: bool,
852 ) -> Result<PageHeader> {
853 let mut prot = ThriftReadInputProtocol::new(input);
854 if self.read_stats {
855 Ok(PageHeader::read_thrift(&mut prot)?)
856 } else {
857 Ok(PageHeader::read_thrift_without_stats(&mut prot)?)
858 }
859 }
860
861 pub(crate) fn decrypt_page_data<T>(
862 &self,
863 buffer: T,
864 _page_index: usize,
865 _dictionary_page: bool,
866 ) -> Result<T> {
867 Ok(buffer)
868 }
869}
870
871#[cfg(feature = "encryption")]
872impl SerializedPageReaderContext {
873 pub(crate) fn read_page_header<T: Read>(
874 &self,
875 input: &mut T,
876 page_index: usize,
877 dictionary_page: bool,
878 ) -> Result<PageHeader> {
879 match self.page_crypto_context(page_index, dictionary_page) {
880 None => {
881 let mut prot = ThriftReadInputProtocol::new(input);
882 if self.read_stats {
883 Ok(PageHeader::read_thrift(&mut prot)?)
884 } else {
885 use crate::file::metadata::thrift::PageHeader;
886
887 Ok(PageHeader::read_thrift_without_stats(&mut prot)?)
888 }
889 }
890 Some(page_crypto_context) => {
891 let data_decryptor = page_crypto_context.data_decryptor();
892 let aad = page_crypto_context.create_page_header_aad()?;
893
894 let buf = read_and_decrypt(data_decryptor, input, aad.as_ref()).map_err(|_| {
895 ParquetError::General(format!(
896 "Error decrypting page header for column {}, decryption key may be wrong",
897 page_crypto_context.column_ordinal
898 ))
899 })?;
900
901 let mut prot = ThriftSliceInputProtocol::new(buf.as_slice());
902 if self.read_stats {
903 Ok(PageHeader::read_thrift(&mut prot)?)
904 } else {
905 Ok(PageHeader::read_thrift_without_stats(&mut prot)?)
906 }
907 }
908 }
909 }
910
911 pub(crate) fn decrypt_page_data<T>(
912 &self,
913 buffer: T,
914 page_index: usize,
915 dictionary_page: bool,
916 ) -> Result<T>
917 where
918 T: AsRef<[u8]>,
919 T: From<Vec<u8>>,
920 {
921 let page_crypto_context = self.page_crypto_context(page_index, dictionary_page);
922 if let Some(page_crypto_context) = page_crypto_context {
923 let decryptor = page_crypto_context.data_decryptor();
924 let aad = page_crypto_context.create_page_aad()?;
925 let decrypted = decryptor.decrypt(buffer.as_ref(), &aad)?;
926 Ok(T::from(decrypted))
927 } else {
928 Ok(buffer)
929 }
930 }
931
932 fn page_crypto_context(
933 &self,
934 page_index: usize,
935 dictionary_page: bool,
936 ) -> Option<Arc<CryptoContext>> {
937 self.crypto_context.as_ref().map(|c| {
938 Arc::new(if dictionary_page {
939 c.for_dictionary_page()
940 } else {
941 c.with_page_ordinal(page_index)
942 })
943 })
944 }
945}
946
947impl<R: ChunkReader> Iterator for SerializedPageReader<R> {
948 type Item = Result<Page>;
949
950 fn next(&mut self) -> Option<Self::Item> {
951 self.get_next_page().transpose()
952 }
953}
954
955fn verify_page_header_len(header_len: usize, remaining_bytes: u64) -> Result<()> {
956 if header_len as u64 > remaining_bytes {
957 return Err(eof_err!("Invalid page header"));
958 }
959 Ok(())
960}
961
962pub(crate) fn verify_page_size(
963 compressed_size: i32,
964 uncompressed_size: i32,
965 remaining_bytes: u64,
966) -> Result<()> {
967 if compressed_size < 0 || compressed_size as u64 > remaining_bytes || uncompressed_size < 0 {
971 return Err(eof_err!("Invalid page header"));
972 }
973 Ok(())
974}
975
976impl<R: ChunkReader> PageReader for SerializedPageReader<R> {
977 fn get_next_page(&mut self) -> Result<Option<Page>> {
978 loop {
979 let page = match &mut self.state {
980 SerializedPageReaderState::Values {
981 offset,
982 remaining_bytes: remaining,
983 next_page_header,
984 page_index,
985 require_dictionary,
986 } => {
987 if *remaining == 0 {
988 return Ok(None);
989 }
990
991 let mut read = self.reader.get_read(*offset)?;
992 let header = if let Some(header) = next_page_header.take() {
993 *header
994 } else {
995 let (header_len, header) = Self::read_page_header_len(
996 &self.context,
997 &mut read,
998 *page_index,
999 *require_dictionary,
1000 )?;
1001 verify_page_header_len(header_len, *remaining)?;
1002 *offset += header_len as u64;
1003 *remaining -= header_len as u64;
1004 header
1005 };
1006 verify_page_size(
1007 header.compressed_page_size,
1008 header.uncompressed_page_size,
1009 *remaining,
1010 )?;
1011 let data_len = header.compressed_page_size as usize;
1012 let data_start = *offset;
1013 *offset += data_len as u64;
1014 *remaining -= data_len as u64;
1015
1016 if header.r#type == PageType::INDEX_PAGE {
1017 continue;
1018 }
1019
1020 let buffer = self.reader.get_bytes(data_start, data_len)?;
1021
1022 let buffer =
1023 self.context
1024 .decrypt_page_data(buffer, *page_index, *require_dictionary)?;
1025
1026 let page = decode_page(
1027 header,
1028 buffer,
1029 self.physical_type,
1030 self.decompressor.as_mut(),
1031 )?;
1032 if page.is_data_page() {
1033 *page_index += 1;
1034 } else if page.is_dictionary_page() {
1035 *require_dictionary = false;
1036 }
1037 page
1038 }
1039 SerializedPageReaderState::Pages {
1040 page_locations,
1041 dictionary_page,
1042 page_index,
1043 ..
1044 } => {
1045 let (front, is_dictionary_page) = match dictionary_page.take() {
1046 Some(front) => (front, true),
1047 None => match page_locations.pop_front() {
1048 Some(front) => (front, false),
1049 None => return Ok(None),
1050 },
1051 };
1052
1053 let page_len = usize::try_from(front.compressed_page_size)?;
1054 let buffer = self.reader.get_bytes(front.offset as u64, page_len)?;
1055
1056 let (offset, header) = read_page_header_len_from_bytes(
1057 &self.context,
1058 buffer.as_ref(),
1059 *page_index,
1060 is_dictionary_page,
1061 )?;
1062 let bytes = buffer.slice(offset..);
1063 let bytes =
1064 self.context
1065 .decrypt_page_data(bytes, *page_index, is_dictionary_page)?;
1066
1067 if !is_dictionary_page {
1068 *page_index += 1;
1069 }
1070 decode_page(
1071 header,
1072 bytes,
1073 self.physical_type,
1074 self.decompressor.as_mut(),
1075 )?
1076 }
1077 };
1078
1079 return Ok(Some(page));
1080 }
1081 }
1082
1083 fn peek_next_page(&mut self) -> Result<Option<PageMetadata>> {
1084 match &mut self.state {
1085 SerializedPageReaderState::Values {
1086 offset,
1087 remaining_bytes,
1088 next_page_header,
1089 page_index,
1090 require_dictionary,
1091 } => {
1092 loop {
1093 if *remaining_bytes == 0 {
1094 return Ok(None);
1095 }
1096 return if let Some(header) = next_page_header.as_ref() {
1097 if let Ok(page_meta) = (&**header).try_into() {
1098 Ok(Some(page_meta))
1099 } else {
1100 *next_page_header = None;
1102 continue;
1103 }
1104 } else {
1105 let mut read = self.reader.get_read(*offset)?;
1106 let (header_len, header) = Self::read_page_header_len(
1107 &self.context,
1108 &mut read,
1109 *page_index,
1110 *require_dictionary,
1111 )?;
1112 verify_page_header_len(header_len, *remaining_bytes)?;
1113 *offset += header_len as u64;
1114 *remaining_bytes -= header_len as u64;
1115 let page_meta = if let Ok(page_meta) = (&header).try_into() {
1116 Ok(Some(page_meta))
1117 } else {
1118 continue;
1120 };
1121 *next_page_header = Some(Box::new(header));
1122 page_meta
1123 };
1124 }
1125 }
1126 SerializedPageReaderState::Pages {
1127 page_locations,
1128 dictionary_page,
1129 total_rows,
1130 page_index: _,
1131 } => {
1132 if dictionary_page.is_some() {
1133 Ok(Some(PageMetadata {
1134 num_rows: None,
1135 num_levels: None,
1136 is_dict: true,
1137 }))
1138 } else if let Some(page) = page_locations.front() {
1139 let next_rows = page_locations
1140 .get(1)
1141 .map(|x| x.first_row_index as usize)
1142 .unwrap_or(*total_rows);
1143
1144 Ok(Some(PageMetadata {
1145 num_rows: Some(next_rows - page.first_row_index as usize),
1146 num_levels: None,
1147 is_dict: false,
1148 }))
1149 } else {
1150 Ok(None)
1151 }
1152 }
1153 }
1154 }
1155
1156 fn skip_next_page(&mut self) -> Result<()> {
1157 let deferred = match &mut self.state {
1160 SerializedPageReaderState::Values {
1161 offset,
1162 remaining_bytes,
1163 next_page_header,
1164 page_index,
1165 require_dictionary,
1166 } => {
1167 let mut deferred = None;
1168 if let Some(buffered_header) = next_page_header.take() {
1169 verify_page_size(
1170 buffered_header.compressed_page_size,
1171 buffered_header.uncompressed_page_size,
1172 *remaining_bytes,
1173 )?;
1174 let body_len = buffered_header.compressed_page_size as u64;
1177 if buffered_header.r#type == PageType::DICTIONARY_PAGE {
1178 deferred = Some(DeferredDictionaryPage::Header {
1179 data_start: *offset,
1180 header: buffered_header,
1181 });
1182 }
1183 *offset += body_len;
1184 *remaining_bytes -= body_len;
1185 } else {
1186 let mut read = self.reader.get_read(*offset)?;
1187 let (header_len, header) = Self::read_page_header_len(
1188 &self.context,
1189 &mut read,
1190 *page_index,
1191 *require_dictionary,
1192 )?;
1193 verify_page_header_len(header_len, *remaining_bytes)?;
1194 verify_page_size(
1195 header.compressed_page_size,
1196 header.uncompressed_page_size,
1197 *remaining_bytes,
1198 )?;
1199 let data_page_size = header.compressed_page_size as u64;
1200 let is_dict = header.r#type == PageType::DICTIONARY_PAGE;
1201 if is_dict {
1202 deferred = Some(DeferredDictionaryPage::Header {
1203 data_start: *offset + header_len as u64,
1204 header: Box::new(header),
1205 });
1206 }
1207 *offset += header_len as u64 + data_page_size;
1208 *remaining_bytes -= header_len as u64 + data_page_size;
1209 }
1210 if deferred.is_some() {
1211 *require_dictionary = false;
1212 } else {
1213 *page_index += 1;
1214 }
1215 deferred
1216 }
1217 SerializedPageReaderState::Pages {
1218 page_locations,
1219 dictionary_page,
1220 page_index,
1221 ..
1222 } => {
1223 if dictionary_page.is_some() {
1224 dictionary_page.take().map(DeferredDictionaryPage::Location)
1226 } else {
1227 if page_locations.pop_front().is_some() {
1229 *page_index += 1;
1230 }
1231 None
1232 }
1233 }
1234 };
1235 if deferred.is_some() {
1236 self.deferred_dictionary = deferred;
1237 }
1238 Ok(())
1239 }
1240
1241 fn supports_deferred_dictionary(&self) -> bool {
1242 match &self.state {
1243 SerializedPageReaderState::Values { .. } => true,
1245 SerializedPageReaderState::Pages {
1249 dictionary_page, ..
1250 } => dictionary_page.is_some(),
1251 }
1252 }
1253
1254 fn take_deferred_dictionary(&mut self) -> Result<Option<Page>> {
1255 let Some(deferred) = self.deferred_dictionary.take() else {
1256 return Ok(None);
1257 };
1258 let page = match deferred {
1259 DeferredDictionaryPage::Header { data_start, header } => {
1260 let data_len = header.compressed_page_size as usize;
1261 let buffer = self.reader.get_bytes(data_start, data_len)?;
1262 let buffer = self.context.decrypt_page_data(buffer, 0, true)?;
1263 decode_page(
1264 *header,
1265 buffer,
1266 self.physical_type,
1267 self.decompressor.as_mut(),
1268 )?
1269 }
1270 DeferredDictionaryPage::Location(front) => {
1271 let page_len = usize::try_from(front.compressed_page_size)?;
1272 let buffer = self.reader.get_bytes(front.offset as u64, page_len)?;
1273 let (offset, header) =
1274 read_page_header_len_from_bytes(&self.context, buffer.as_ref(), 0, true)?;
1275 let bytes = buffer.slice(offset..);
1276 let bytes = self.context.decrypt_page_data(bytes, 0, true)?;
1277 decode_page(
1278 header,
1279 bytes,
1280 self.physical_type,
1281 self.decompressor.as_mut(),
1282 )?
1283 }
1284 };
1285 Ok(Some(page))
1286 }
1287
1288 fn at_record_boundary(&mut self) -> Result<bool> {
1289 match &mut self.state {
1290 SerializedPageReaderState::Values { .. } => match self.peek_next_page()? {
1291 None => Ok(true),
1292 Some(metadata) => Ok(metadata.num_rows.is_some()),
1295 },
1296 SerializedPageReaderState::Pages { .. } => Ok(true),
1297 }
1298 }
1299}
1300
1301#[cfg(test)]
1302mod tests {
1303 use std::cell::RefCell;
1304 use std::collections::HashSet;
1305 use std::rc::Rc;
1306
1307 use bytes::Buf;
1308
1309 use crate::file::page_index::column_index::{
1310 ByteArrayColumnIndex, ColumnIndexMetaData, PrimitiveColumnIndex,
1311 };
1312 use crate::file::properties::{EnabledStatistics, WriterProperties};
1313
1314 use crate::basic::{self, BoundaryOrder, ColumnOrder, Compression, Encoding, SortOrder};
1315 use crate::column::reader::ColumnReader;
1316 use crate::data_type::private::ParquetValueType;
1317 use crate::data_type::{AsBytes, FixedLenByteArrayType, Int32Type};
1318 use crate::file::metadata::thrift::DataPageHeaderV2;
1319 use crate::file::writer::SerializedFileWriter;
1320 use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
1321 use crate::record::RowAccessor;
1322 use crate::schema::parser::parse_message_type;
1323 use crate::util::test_common::file_util::{get_test_file, get_test_path};
1324
1325 use super::*;
1326
1327 #[test]
1328 fn test_decode_page_invalid_offset() {
1329 let page_header = PageHeader {
1330 r#type: PageType::DATA_PAGE_V2,
1331 uncompressed_page_size: 10,
1332 compressed_page_size: 10,
1333 data_page_header: None,
1334 index_page_header: None,
1335 dictionary_page_header: None,
1336 crc: None,
1337 data_page_header_v2: Some(DataPageHeaderV2 {
1338 num_nulls: 0,
1339 num_rows: 0,
1340 num_values: 0,
1341 encoding: Encoding::PLAIN,
1342 definition_levels_byte_length: 11,
1343 repetition_levels_byte_length: 0,
1344 is_compressed: None,
1345 statistics: None,
1346 }),
1347 };
1348
1349 let buffer = Bytes::new();
1350 let err = decode_page(page_header, buffer, Type::INT32, None).unwrap_err();
1351 assert!(
1352 err.to_string()
1353 .contains("DataPage v2 header contains implausible values")
1354 );
1355 }
1356
1357 #[test]
1358 fn test_decode_unsupported_page() {
1359 let mut page_header = PageHeader {
1360 r#type: PageType::INDEX_PAGE,
1361 uncompressed_page_size: 10,
1362 compressed_page_size: 10,
1363 data_page_header: None,
1364 index_page_header: None,
1365 dictionary_page_header: None,
1366 crc: None,
1367 data_page_header_v2: None,
1368 };
1369 let buffer = Bytes::new();
1370 let err = decode_page(page_header.clone(), buffer.clone(), Type::INT32, None).unwrap_err();
1371 assert_eq!(
1372 err.to_string(),
1373 "Parquet error: Page type INDEX_PAGE is not supported"
1374 );
1375
1376 page_header.data_page_header_v2 = Some(DataPageHeaderV2 {
1377 num_nulls: 0,
1378 num_rows: 0,
1379 num_values: 0,
1380 encoding: Encoding::PLAIN,
1381 definition_levels_byte_length: 11,
1382 repetition_levels_byte_length: 0,
1383 is_compressed: None,
1384 statistics: None,
1385 });
1386 let err = decode_page(page_header, buffer, Type::INT32, None).unwrap_err();
1387 assert!(
1388 err.to_string()
1389 .contains("DataPage v2 header contains implausible values")
1390 );
1391 }
1392
1393 #[test]
1394 fn test_cursor_and_file_has_the_same_behaviour() {
1395 let mut buf: Vec<u8> = Vec::new();
1396 get_test_file("alltypes_plain.parquet")
1397 .read_to_end(&mut buf)
1398 .unwrap();
1399 let cursor = Bytes::from(buf);
1400 let read_from_cursor = SerializedFileReader::new(cursor).unwrap();
1401
1402 let test_file = get_test_file("alltypes_plain.parquet");
1403 let read_from_file = SerializedFileReader::new(test_file).unwrap();
1404
1405 let file_iter = read_from_file.get_row_iter(None).unwrap();
1406 let cursor_iter = read_from_cursor.get_row_iter(None).unwrap();
1407
1408 for (a, b) in file_iter.zip(cursor_iter) {
1409 assert_eq!(a.unwrap(), b.unwrap())
1410 }
1411 }
1412
1413 #[test]
1414 fn test_file_reader_try_from() {
1415 let test_file = get_test_file("alltypes_plain.parquet");
1417 let test_path_buf = get_test_path("alltypes_plain.parquet");
1418 let test_path = test_path_buf.as_path();
1419 let test_path_str = test_path.to_str().unwrap();
1420
1421 let reader = SerializedFileReader::try_from(test_file);
1422 assert!(reader.is_ok());
1423
1424 let reader = SerializedFileReader::try_from(test_path);
1425 assert!(reader.is_ok());
1426
1427 let reader = SerializedFileReader::try_from(test_path_str);
1428 assert!(reader.is_ok());
1429
1430 let reader = SerializedFileReader::try_from(test_path_str.to_string());
1431 assert!(reader.is_ok());
1432
1433 let test_path = Path::new("invalid.parquet");
1435 let test_path_str = test_path.to_str().unwrap();
1436
1437 let reader = SerializedFileReader::try_from(test_path);
1438 assert!(reader.is_err());
1439
1440 let reader = SerializedFileReader::try_from(test_path_str);
1441 assert!(reader.is_err());
1442
1443 let reader = SerializedFileReader::try_from(test_path_str.to_string());
1444 assert!(reader.is_err());
1445 }
1446
1447 #[test]
1448 fn test_file_reader_into_iter() {
1449 let path = get_test_path("alltypes_plain.parquet");
1450 let reader = SerializedFileReader::try_from(path.as_path()).unwrap();
1451 let iter = reader.into_iter();
1452 let values: Vec<_> = iter.flat_map(|x| x.unwrap().get_int(0)).collect();
1453
1454 assert_eq!(values, &[4, 5, 6, 7, 2, 3, 0, 1]);
1455 }
1456
1457 #[test]
1458 fn test_file_reader_into_iter_project() {
1459 let path = get_test_path("alltypes_plain.parquet");
1460 let reader = SerializedFileReader::try_from(path.as_path()).unwrap();
1461 let schema = "message schema { OPTIONAL INT32 id; }";
1462 let proj = parse_message_type(schema).ok();
1463 let iter = reader.into_iter().project(proj).unwrap();
1464 let values: Vec<_> = iter.flat_map(|x| x.unwrap().get_int(0)).collect();
1465
1466 assert_eq!(values, &[4, 5, 6, 7, 2, 3, 0, 1]);
1467 }
1468
1469 #[test]
1470 fn test_reuse_file_chunk() {
1471 let test_file = get_test_file("alltypes_plain.parquet");
1475 let reader = SerializedFileReader::new(test_file).unwrap();
1476 let row_group = reader.get_row_group(0).unwrap();
1477
1478 let mut page_readers = Vec::new();
1479 for i in 0..row_group.num_columns() {
1480 page_readers.push(row_group.get_column_page_reader(i).unwrap());
1481 }
1482
1483 for mut page_reader in page_readers {
1486 assert!(page_reader.get_next_page().is_ok());
1487 }
1488 }
1489
1490 #[test]
1491 fn test_file_reader() {
1492 let test_file = get_test_file("alltypes_plain.parquet");
1493 let reader_result = SerializedFileReader::new(test_file);
1494 assert!(reader_result.is_ok());
1495 let reader = reader_result.unwrap();
1496
1497 let metadata = reader.metadata();
1499 assert_eq!(metadata.num_row_groups(), 1);
1500
1501 let file_metadata = metadata.file_metadata();
1503 assert!(file_metadata.created_by().is_some());
1504 assert_eq!(
1505 file_metadata.created_by().unwrap(),
1506 "impala version 1.3.0-INTERNAL (build 8a48ddb1eff84592b3fc06bc6f51ec120e1fffc9)"
1507 );
1508 assert!(file_metadata.key_value_metadata().is_none());
1509 assert_eq!(file_metadata.num_rows(), 8);
1510 assert_eq!(file_metadata.version(), 1);
1511 assert_eq!(file_metadata.column_orders(), None);
1512
1513 let row_group_metadata = metadata.row_group(0);
1515 assert_eq!(row_group_metadata.num_columns(), 11);
1516 assert_eq!(row_group_metadata.num_rows(), 8);
1517 assert_eq!(row_group_metadata.total_byte_size(), 671);
1518 for i in 0..row_group_metadata.num_columns() {
1520 assert_eq!(file_metadata.column_order(i), ColumnOrder::UNDEFINED);
1521 }
1522
1523 let row_group_reader_result = reader.get_row_group(0);
1525 assert!(row_group_reader_result.is_ok());
1526 let row_group_reader: Box<dyn RowGroupReader> = row_group_reader_result.unwrap();
1527 assert_eq!(
1528 row_group_reader.num_columns(),
1529 row_group_metadata.num_columns()
1530 );
1531 assert_eq!(
1532 row_group_reader.metadata().total_byte_size(),
1533 row_group_metadata.total_byte_size()
1534 );
1535
1536 let page_reader_0_result = row_group_reader.get_column_page_reader(0);
1539 assert!(page_reader_0_result.is_ok());
1540 let mut page_reader_0: Box<dyn PageReader> = page_reader_0_result.unwrap();
1541 let mut page_count = 0;
1542 while let Some(page) = page_reader_0.get_next_page().unwrap() {
1543 let is_expected_page = match page {
1544 Page::DictionaryPage {
1545 buf,
1546 num_values,
1547 encoding,
1548 is_sorted,
1549 } => {
1550 assert_eq!(buf.len(), 32);
1551 assert_eq!(num_values, 8);
1552 assert_eq!(encoding, Encoding::PLAIN_DICTIONARY);
1553 assert!(!is_sorted);
1554 true
1555 }
1556 Page::DataPage {
1557 buf,
1558 num_values,
1559 encoding,
1560 def_level_encoding,
1561 rep_level_encoding,
1562 statistics,
1563 } => {
1564 assert_eq!(buf.len(), 11);
1565 assert_eq!(num_values, 8);
1566 assert_eq!(encoding, Encoding::PLAIN_DICTIONARY);
1567 assert_eq!(def_level_encoding, Encoding::RLE);
1568 #[expect(deprecated)]
1569 let expected_rep_level_encoding = Encoding::BIT_PACKED;
1570 assert_eq!(rep_level_encoding, expected_rep_level_encoding);
1571 assert!(statistics.is_none());
1572 true
1573 }
1574 Page::DataPageV2 { .. } => false,
1575 };
1576 assert!(is_expected_page);
1577 page_count += 1;
1578 }
1579 assert_eq!(page_count, 2);
1580 }
1581
1582 #[test]
1583 fn test_file_reader_datapage_v2() {
1584 let test_file = get_test_file("datapage_v2.snappy.parquet");
1585 let reader_result = SerializedFileReader::new(test_file);
1586 assert!(reader_result.is_ok());
1587 let reader = reader_result.unwrap();
1588
1589 let metadata = reader.metadata();
1591 assert_eq!(metadata.num_row_groups(), 1);
1592
1593 let file_metadata = metadata.file_metadata();
1595 assert!(file_metadata.created_by().is_some());
1596 assert_eq!(
1597 file_metadata.created_by().unwrap(),
1598 "parquet-mr version 1.8.1 (build 4aba4dae7bb0d4edbcf7923ae1339f28fd3f7fcf)"
1599 );
1600 assert!(file_metadata.key_value_metadata().is_some());
1601 assert_eq!(file_metadata.key_value_metadata().unwrap().len(), 1);
1602
1603 assert_eq!(file_metadata.num_rows(), 5);
1604 assert_eq!(file_metadata.version(), 1);
1605 assert_eq!(file_metadata.column_orders(), None);
1606
1607 let row_group_metadata = metadata.row_group(0);
1608
1609 for i in 0..row_group_metadata.num_columns() {
1611 assert_eq!(file_metadata.column_order(i), ColumnOrder::UNDEFINED);
1612 }
1613
1614 let row_group_reader_result = reader.get_row_group(0);
1616 assert!(row_group_reader_result.is_ok());
1617 let row_group_reader: Box<dyn RowGroupReader> = row_group_reader_result.unwrap();
1618 assert_eq!(
1619 row_group_reader.num_columns(),
1620 row_group_metadata.num_columns()
1621 );
1622 assert_eq!(
1623 row_group_reader.metadata().total_byte_size(),
1624 row_group_metadata.total_byte_size()
1625 );
1626
1627 let page_reader_0_result = row_group_reader.get_column_page_reader(0);
1630 assert!(page_reader_0_result.is_ok());
1631 let mut page_reader_0: Box<dyn PageReader> = page_reader_0_result.unwrap();
1632 let mut page_count = 0;
1633 while let Some(page) = page_reader_0.get_next_page().unwrap() {
1634 let is_expected_page = match page {
1635 Page::DictionaryPage {
1636 buf,
1637 num_values,
1638 encoding,
1639 is_sorted,
1640 } => {
1641 assert_eq!(buf.len(), 7);
1642 assert_eq!(num_values, 1);
1643 assert_eq!(encoding, Encoding::PLAIN);
1644 assert!(!is_sorted);
1645 true
1646 }
1647 Page::DataPageV2 {
1648 buf,
1649 num_values,
1650 encoding,
1651 num_nulls,
1652 num_rows,
1653 def_levels_byte_len,
1654 rep_levels_byte_len,
1655 is_compressed,
1656 statistics,
1657 } => {
1658 assert_eq!(buf.len(), 4);
1659 assert_eq!(num_values, 5);
1660 assert_eq!(encoding, Encoding::RLE_DICTIONARY);
1661 assert_eq!(num_nulls, 1);
1662 assert_eq!(num_rows, 5);
1663 assert_eq!(def_levels_byte_len, 2);
1664 assert_eq!(rep_levels_byte_len, 0);
1665 assert!(is_compressed);
1666 assert!(statistics.is_none()); true
1668 }
1669 Page::DataPage { .. } => false,
1670 };
1671 assert!(is_expected_page);
1672 page_count += 1;
1673 }
1674 assert_eq!(page_count, 2);
1675 }
1676
1677 #[cfg_attr(miri, ignore)] #[test]
1679 fn test_file_reader_empty_compressed_datapage_v2() {
1680 let test_file = get_test_file("page_v2_empty_compressed.parquet");
1682 let reader_result = SerializedFileReader::new(test_file);
1683 assert!(reader_result.is_ok());
1684 let reader = reader_result.unwrap();
1685
1686 let metadata = reader.metadata();
1688 assert_eq!(metadata.num_row_groups(), 1);
1689
1690 let file_metadata = metadata.file_metadata();
1692 assert!(file_metadata.created_by().is_some());
1693 assert_eq!(
1694 file_metadata.created_by().unwrap(),
1695 "parquet-cpp-arrow version 14.0.2"
1696 );
1697 assert!(file_metadata.key_value_metadata().is_some());
1698 assert_eq!(file_metadata.key_value_metadata().unwrap().len(), 1);
1699
1700 assert_eq!(file_metadata.num_rows(), 10);
1701 assert_eq!(file_metadata.version(), 2);
1702 let expected_order = ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED);
1703 assert_eq!(
1704 file_metadata.column_orders(),
1705 Some(vec![expected_order].as_ref())
1706 );
1707
1708 let row_group_metadata = metadata.row_group(0);
1709
1710 for i in 0..row_group_metadata.num_columns() {
1712 assert_eq!(file_metadata.column_order(i), expected_order);
1713 }
1714
1715 let row_group_reader_result = reader.get_row_group(0);
1717 assert!(row_group_reader_result.is_ok());
1718 let row_group_reader: Box<dyn RowGroupReader> = row_group_reader_result.unwrap();
1719 assert_eq!(
1720 row_group_reader.num_columns(),
1721 row_group_metadata.num_columns()
1722 );
1723 assert_eq!(
1724 row_group_reader.metadata().total_byte_size(),
1725 row_group_metadata.total_byte_size()
1726 );
1727
1728 let page_reader_0_result = row_group_reader.get_column_page_reader(0);
1730 assert!(page_reader_0_result.is_ok());
1731 let mut page_reader_0: Box<dyn PageReader> = page_reader_0_result.unwrap();
1732 let mut page_count = 0;
1733 while let Some(page) = page_reader_0.get_next_page().unwrap() {
1734 let is_expected_page = match page {
1735 Page::DictionaryPage {
1736 buf,
1737 num_values,
1738 encoding,
1739 is_sorted,
1740 } => {
1741 assert_eq!(buf.len(), 0);
1742 assert_eq!(num_values, 0);
1743 assert_eq!(encoding, Encoding::PLAIN);
1744 assert!(!is_sorted);
1745 true
1746 }
1747 Page::DataPageV2 {
1748 buf,
1749 num_values,
1750 encoding,
1751 num_nulls,
1752 num_rows,
1753 def_levels_byte_len,
1754 rep_levels_byte_len,
1755 is_compressed,
1756 statistics,
1757 } => {
1758 assert_eq!(buf.len(), 3);
1759 assert_eq!(num_values, 10);
1760 assert_eq!(encoding, Encoding::RLE_DICTIONARY);
1761 assert_eq!(num_nulls, 10);
1762 assert_eq!(num_rows, 10);
1763 assert_eq!(def_levels_byte_len, 2);
1764 assert_eq!(rep_levels_byte_len, 0);
1765 assert!(is_compressed);
1766 assert!(statistics.is_none()); true
1768 }
1769 Page::DataPage { .. } => false,
1770 };
1771 assert!(is_expected_page);
1772 page_count += 1;
1773 }
1774 assert_eq!(page_count, 2);
1775 }
1776
1777 #[test]
1778 fn test_file_reader_empty_datapage_v2() {
1779 let test_file = get_test_file("datapage_v2_empty_datapage.snappy.parquet");
1781 let reader_result = SerializedFileReader::new(test_file);
1782 assert!(reader_result.is_ok());
1783 let reader = reader_result.unwrap();
1784
1785 let metadata = reader.metadata();
1787 assert_eq!(metadata.num_row_groups(), 1);
1788
1789 let file_metadata = metadata.file_metadata();
1791 assert!(file_metadata.created_by().is_some());
1792 assert_eq!(
1793 file_metadata.created_by().unwrap(),
1794 "parquet-mr version 1.13.1 (build db4183109d5b734ec5930d870cdae161e408ddba)"
1795 );
1796 assert!(file_metadata.key_value_metadata().is_some());
1797 assert_eq!(file_metadata.key_value_metadata().unwrap().len(), 2);
1798
1799 assert_eq!(file_metadata.num_rows(), 1);
1800 assert_eq!(file_metadata.version(), 1);
1801 let expected_order = ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED);
1802 assert_eq!(
1803 file_metadata.column_orders(),
1804 Some(vec![expected_order].as_ref())
1805 );
1806
1807 let row_group_metadata = metadata.row_group(0);
1808
1809 for i in 0..row_group_metadata.num_columns() {
1811 assert_eq!(file_metadata.column_order(i), expected_order);
1812 }
1813
1814 let row_group_reader_result = reader.get_row_group(0);
1816 assert!(row_group_reader_result.is_ok());
1817 let row_group_reader: Box<dyn RowGroupReader> = row_group_reader_result.unwrap();
1818 assert_eq!(
1819 row_group_reader.num_columns(),
1820 row_group_metadata.num_columns()
1821 );
1822 assert_eq!(
1823 row_group_reader.metadata().total_byte_size(),
1824 row_group_metadata.total_byte_size()
1825 );
1826
1827 let page_reader_0_result = row_group_reader.get_column_page_reader(0);
1829 assert!(page_reader_0_result.is_ok());
1830 let mut page_reader_0: Box<dyn PageReader> = page_reader_0_result.unwrap();
1831 let mut page_count = 0;
1832 while let Some(page) = page_reader_0.get_next_page().unwrap() {
1833 let is_expected_page = match page {
1834 Page::DataPageV2 {
1835 buf,
1836 num_values,
1837 encoding,
1838 num_nulls,
1839 num_rows,
1840 def_levels_byte_len,
1841 rep_levels_byte_len,
1842 is_compressed,
1843 statistics,
1844 } => {
1845 assert_eq!(buf.len(), 2);
1846 assert_eq!(num_values, 1);
1847 assert_eq!(encoding, Encoding::PLAIN);
1848 assert_eq!(num_nulls, 1);
1849 assert_eq!(num_rows, 1);
1850 assert_eq!(def_levels_byte_len, 2);
1851 assert_eq!(rep_levels_byte_len, 0);
1852 assert!(is_compressed);
1853 assert!(statistics.is_none());
1854 true
1855 }
1856 _ => false,
1857 };
1858 assert!(is_expected_page);
1859 page_count += 1;
1860 }
1861 assert_eq!(page_count, 1);
1862 }
1863
1864 fn get_serialized_page_reader<R: ChunkReader>(
1865 file_reader: &SerializedFileReader<R>,
1866 row_group_idx: usize,
1867 column: usize,
1868 ) -> Result<SerializedPageReader<R>> {
1869 let row_group = {
1870 let row_group_metadata = file_reader.metadata.row_group(row_group_idx);
1871 let props = Arc::clone(&file_reader.props);
1872 let f = Arc::clone(&file_reader.chunk_reader);
1873 let page_index = file_reader.metadata.page_index_for_row_group(row_group_idx);
1874 SerializedRowGroupReader::new(f, row_group_metadata, page_index, props)?
1875 };
1876
1877 let col = row_group.metadata.column(column);
1878 let page_locations = file_reader
1879 .metadata
1880 .page_index()
1881 .map(|pi| pi.page_locations(row_group_idx, column))
1882 .unwrap_or(None)
1883 .cloned();
1884
1885 let props = Arc::clone(&row_group.props);
1886 SerializedPageReader::new_with_properties(
1887 Arc::clone(&row_group.chunk_reader),
1888 col,
1889 usize::try_from(row_group.metadata.num_rows())?,
1890 page_locations,
1891 props,
1892 )
1893 }
1894
1895 #[test]
1896 fn test_peek_next_page_offset_matches_actual() -> Result<()> {
1897 let test_file = get_test_file("alltypes_plain.parquet");
1898 let reader = SerializedFileReader::new(test_file)?;
1899
1900 let mut offset_set = HashSet::new();
1901 let num_row_groups = reader.metadata.num_row_groups();
1902 for row_group in 0..num_row_groups {
1903 let num_columns = reader.metadata.row_group(row_group).num_columns();
1904 for column in 0..num_columns {
1905 let mut page_reader = get_serialized_page_reader(&reader, row_group, column)?;
1906
1907 while let Ok(Some(page_offset)) = page_reader.peek_next_page_offset() {
1908 match &page_reader.state {
1909 SerializedPageReaderState::Pages {
1910 page_locations,
1911 dictionary_page,
1912 ..
1913 } => {
1914 if let Some(page) = dictionary_page {
1915 assert_eq!(page.offset as u64, page_offset);
1916 } else if let Some(page) = page_locations.front() {
1917 assert_eq!(page.offset as u64, page_offset);
1918 } else {
1919 unreachable!()
1920 }
1921 }
1922 SerializedPageReaderState::Values {
1923 offset,
1924 next_page_header,
1925 ..
1926 } => {
1927 assert!(next_page_header.is_some());
1928 assert_eq!(*offset, page_offset);
1929 }
1930 }
1931 let page = page_reader.get_next_page()?;
1932 assert!(page.is_some());
1933 let newly_inserted = offset_set.insert(page_offset);
1934 assert!(newly_inserted);
1935 }
1936 }
1937 }
1938
1939 Ok(())
1940 }
1941
1942 #[test]
1943 fn test_page_iterator() {
1944 let file = get_test_file("alltypes_plain.parquet");
1945 let file_reader = Arc::new(SerializedFileReader::new(file).unwrap());
1946
1947 let mut page_iterator = FilePageIterator::new(0, file_reader.clone()).unwrap();
1948
1949 let page = page_iterator.next();
1951 assert!(page.is_some());
1952 assert!(page.unwrap().is_ok());
1953
1954 let page = page_iterator.next();
1956 assert!(page.is_none());
1957
1958 let row_group_indices = Box::new(0..1);
1959 let mut page_iterator =
1960 FilePageIterator::with_row_groups(0, row_group_indices, file_reader).unwrap();
1961
1962 let page = page_iterator.next();
1964 assert!(page.is_some());
1965 assert!(page.unwrap().is_ok());
1966
1967 let page = page_iterator.next();
1969 assert!(page.is_none());
1970 }
1971
1972 #[test]
1973 fn test_file_reader_key_value_metadata() {
1974 let file = get_test_file("binary.parquet");
1975 let file_reader = Arc::new(SerializedFileReader::new(file).unwrap());
1976
1977 let metadata = file_reader
1978 .metadata
1979 .file_metadata()
1980 .key_value_metadata()
1981 .unwrap();
1982
1983 assert_eq!(metadata.len(), 3);
1984
1985 assert_eq!(metadata[0].key, "parquet.proto.descriptor");
1986
1987 assert_eq!(metadata[1].key, "writer.model.name");
1988 assert_eq!(metadata[1].value, Some("protobuf".to_owned()));
1989
1990 assert_eq!(metadata[2].key, "parquet.proto.class");
1991 assert_eq!(metadata[2].value, Some("foo.baz.Foobaz$Event".to_owned()));
1992 }
1993
1994 #[test]
1995 fn test_file_reader_optional_metadata() {
1996 let file = get_test_file("data_index_bloom_encoding_stats.parquet");
1998 let options = ReadOptionsBuilder::new()
1999 .with_encoding_stats_as_mask(false)
2000 .build();
2001 let file_reader = Arc::new(SerializedFileReader::new_with_options(file, options).unwrap());
2002
2003 let row_group_metadata = file_reader.metadata.row_group(0);
2004 let col0_metadata = row_group_metadata.column(0);
2005
2006 assert_eq!(col0_metadata.bloom_filter_offset().unwrap(), 192);
2008
2009 let page_encoding_stats = &col0_metadata.page_encoding_stats().unwrap()[0];
2011
2012 assert_eq!(page_encoding_stats.page_type, basic::PageType::DATA_PAGE);
2013 assert_eq!(page_encoding_stats.encoding, Encoding::PLAIN);
2014 assert_eq!(page_encoding_stats.count, 1);
2015
2016 assert_eq!(col0_metadata.column_index_offset().unwrap(), 156);
2018 assert_eq!(col0_metadata.column_index_length().unwrap(), 25);
2019
2020 assert_eq!(col0_metadata.offset_index_offset().unwrap(), 181);
2022 assert_eq!(col0_metadata.offset_index_length().unwrap(), 11);
2023 }
2024
2025 #[test]
2026 fn test_file_reader_page_stats_mask() {
2027 let file = get_test_file("alltypes_tiny_pages.parquet");
2028 let options = ReadOptionsBuilder::new()
2029 .with_encoding_stats_as_mask(true)
2030 .build();
2031 let file_reader = Arc::new(SerializedFileReader::new_with_options(file, options).unwrap());
2032
2033 let row_group_metadata = file_reader.metadata.row_group(0);
2034
2035 let page_encoding_stats = row_group_metadata
2037 .column(0)
2038 .page_encoding_stats_mask()
2039 .unwrap();
2040 assert!(page_encoding_stats.is_only(Encoding::PLAIN));
2041 let page_encoding_stats = row_group_metadata
2042 .column(2)
2043 .page_encoding_stats_mask()
2044 .unwrap();
2045 assert!(page_encoding_stats.is_only(Encoding::PLAIN_DICTIONARY));
2046 }
2047
2048 #[test]
2049 fn test_file_reader_page_stats_skipped() {
2050 let file = get_test_file("alltypes_tiny_pages.parquet");
2051
2052 let options = ReadOptionsBuilder::new()
2054 .with_encoding_stats_policy(ParquetStatisticsPolicy::SkipAll)
2055 .with_column_stats_policy(ParquetStatisticsPolicy::SkipAll)
2056 .build();
2057 let file_reader = Arc::new(
2058 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(),
2059 );
2060
2061 let row_group_metadata = file_reader.metadata.row_group(0);
2062 for column in row_group_metadata.columns() {
2063 assert!(column.page_encoding_stats().is_none());
2064 assert!(column.page_encoding_stats_mask().is_none());
2065 assert!(column.statistics().is_none());
2066 }
2067
2068 let options = ReadOptionsBuilder::new()
2070 .with_encoding_stats_as_mask(true)
2071 .with_encoding_stats_policy(ParquetStatisticsPolicy::skip_except(&[0]))
2072 .with_column_stats_policy(ParquetStatisticsPolicy::skip_except(&[0]))
2073 .build();
2074 let file_reader = Arc::new(
2075 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(),
2076 );
2077
2078 let row_group_metadata = file_reader.metadata.row_group(0);
2079 for (idx, column) in row_group_metadata.columns().iter().enumerate() {
2080 assert!(column.page_encoding_stats().is_none());
2081 assert_eq!(column.page_encoding_stats_mask().is_some(), idx == 0);
2082 assert_eq!(column.statistics().is_some(), idx == 0);
2083 }
2084 }
2085
2086 #[test]
2087 fn test_file_reader_size_stats_skipped() {
2088 let file = get_test_file("repeated_primitive_no_list.parquet");
2089
2090 let options = ReadOptionsBuilder::new()
2092 .with_size_stats_policy(ParquetStatisticsPolicy::SkipAll)
2093 .build();
2094 let file_reader = Arc::new(
2095 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(),
2096 );
2097
2098 let row_group_metadata = file_reader.metadata.row_group(0);
2099 for column in row_group_metadata.columns() {
2100 assert!(column.repetition_level_histogram().is_none());
2101 assert!(column.definition_level_histogram().is_none());
2102 assert!(column.unencoded_byte_array_data_bytes().is_none());
2103 }
2104
2105 let options = ReadOptionsBuilder::new()
2107 .with_encoding_stats_as_mask(true)
2108 .with_size_stats_policy(ParquetStatisticsPolicy::skip_except(&[1]))
2109 .build();
2110 let file_reader = Arc::new(
2111 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(),
2112 );
2113
2114 let row_group_metadata = file_reader.metadata.row_group(0);
2115 for (idx, column) in row_group_metadata.columns().iter().enumerate() {
2116 assert_eq!(column.repetition_level_histogram().is_some(), idx == 1);
2117 assert_eq!(column.definition_level_histogram().is_some(), idx == 1);
2118 assert_eq!(column.unencoded_byte_array_data_bytes().is_some(), idx == 1);
2119 }
2120 }
2121
2122 #[test]
2123 fn test_file_reader_with_no_filter() -> Result<()> {
2124 let test_file = get_test_file("alltypes_plain.parquet");
2125 let origin_reader = SerializedFileReader::new(test_file)?;
2126 let metadata = origin_reader.metadata();
2128 assert_eq!(metadata.num_row_groups(), 1);
2129 Ok(())
2130 }
2131
2132 #[test]
2133 fn test_file_reader_filter_row_groups_with_predicate() -> Result<()> {
2134 let test_file = get_test_file("alltypes_plain.parquet");
2135 let read_options = ReadOptionsBuilder::new()
2136 .with_predicate(Box::new(|_, _| false))
2137 .build();
2138 let reader = SerializedFileReader::new_with_options(test_file, read_options)?;
2139 let metadata = reader.metadata();
2140 assert_eq!(metadata.num_row_groups(), 0);
2141 Ok(())
2142 }
2143
2144 #[test]
2145 fn test_file_reader_filter_row_groups_with_range() -> Result<()> {
2146 let test_file = get_test_file("alltypes_plain.parquet");
2147 let origin_reader = SerializedFileReader::new(test_file)?;
2148 let metadata = origin_reader.metadata();
2150 assert_eq!(metadata.num_row_groups(), 1);
2151 let mid = get_midpoint_offset(metadata.row_group(0));
2152
2153 let test_file = get_test_file("alltypes_plain.parquet");
2154 let read_options = ReadOptionsBuilder::new().with_range(0, mid + 1).build();
2155 let reader = SerializedFileReader::new_with_options(test_file, read_options)?;
2156 let metadata = reader.metadata();
2157 assert_eq!(metadata.num_row_groups(), 1);
2158
2159 let test_file = get_test_file("alltypes_plain.parquet");
2160 let read_options = ReadOptionsBuilder::new().with_range(0, mid).build();
2161 let reader = SerializedFileReader::new_with_options(test_file, read_options)?;
2162 let metadata = reader.metadata();
2163 assert_eq!(metadata.num_row_groups(), 0);
2164 Ok(())
2165 }
2166
2167 #[test]
2168 #[cfg_attr(miri, ignore)] fn test_file_reader_filter_row_groups_and_range() -> Result<()> {
2170 let test_file = get_test_file("alltypes_tiny_pages.parquet");
2171 let origin_reader = SerializedFileReader::new(test_file)?;
2172 let metadata = origin_reader.metadata();
2173 let mid = get_midpoint_offset(metadata.row_group(0));
2174
2175 let test_file = get_test_file("alltypes_tiny_pages.parquet");
2177 let read_options = ReadOptionsBuilder::new()
2178 .with_page_index()
2179 .with_predicate(Box::new(|_, _| true))
2180 .with_range(mid, mid + 1)
2181 .build();
2182 let reader = SerializedFileReader::new_with_options(test_file, read_options)?;
2183 let metadata = reader.metadata();
2184 assert_eq!(metadata.num_row_groups(), 1);
2185 assert!(metadata.page_index().is_some());
2186
2187 let test_file = get_test_file("alltypes_tiny_pages.parquet");
2189 let read_options = ReadOptionsBuilder::new()
2190 .with_page_index()
2191 .with_predicate(Box::new(|_, _| true))
2192 .with_range(0, mid)
2193 .build();
2194 let reader = SerializedFileReader::new_with_options(test_file, read_options)?;
2195 let metadata = reader.metadata();
2196 assert_eq!(metadata.num_row_groups(), 0);
2197 assert!(metadata.page_index().is_none());
2198
2199 let test_file = get_test_file("alltypes_tiny_pages.parquet");
2201 let read_options = ReadOptionsBuilder::new()
2202 .with_page_index()
2203 .with_predicate(Box::new(|_, _| false))
2204 .with_range(mid, mid + 1)
2205 .build();
2206 let reader = SerializedFileReader::new_with_options(test_file, read_options)?;
2207 let metadata = reader.metadata();
2208 assert_eq!(metadata.num_row_groups(), 0);
2209 assert!(metadata.page_index().is_none());
2210
2211 let test_file = get_test_file("alltypes_tiny_pages.parquet");
2213 let read_options = ReadOptionsBuilder::new()
2214 .with_page_index()
2215 .with_predicate(Box::new(|_, _| false))
2216 .with_range(0, mid)
2217 .build();
2218 let reader = SerializedFileReader::new_with_options(test_file, read_options)?;
2219 let metadata = reader.metadata();
2220 assert_eq!(metadata.num_row_groups(), 0);
2221 assert!(metadata.page_index().is_none());
2222 Ok(())
2223 }
2224
2225 #[test]
2226 fn test_file_reader_invalid_metadata() {
2227 let data = [
2228 255, 172, 1, 0, 50, 82, 65, 73, 1, 0, 0, 0, 169, 168, 168, 162, 87, 255, 16, 0, 0, 0,
2229 80, 65, 82, 49,
2230 ];
2231 let ret = SerializedFileReader::new(Bytes::copy_from_slice(&data));
2232 assert_eq!(
2233 ret.err().unwrap().to_string(),
2234 "Parquet error: Expected list element type of Struct but got List"
2235 );
2236 }
2237
2238 #[test]
2239 fn test_page_index_reader() {
2256 let test_file = get_test_file("data_index_bloom_encoding_stats.parquet");
2257 let builder = ReadOptionsBuilder::new();
2258 let options = builder.with_page_index().build();
2260 let reader_result = SerializedFileReader::new_with_options(test_file, options);
2261 let reader = reader_result.unwrap();
2262
2263 let metadata = reader.metadata();
2265 assert_eq!(metadata.num_row_groups(), 1);
2266
2267 let page_index = metadata.page_index().expect("page index should be present");
2268
2269 let Some(ColumnIndexMetaData::BYTE_ARRAY(index)) = page_index.column_index(0, 0) else {
2271 unreachable!()
2272 };
2273
2274 assert_eq!(index.boundary_order, BoundaryOrder::ASCENDING);
2275
2276 assert_eq!(index.num_pages(), 1);
2278
2279 let min = index.min_value(0).unwrap();
2280 let max = index.max_value(0).unwrap();
2281 assert_eq!(b"Hello", min.as_bytes());
2282 assert_eq!(b"today", max.as_bytes());
2283
2284 let offset_index = page_index
2286 .offset_index(0, 0)
2287 .expect("offset index should be present");
2288 let page_offset = offset_index
2289 .page_locations()
2290 .first()
2291 .expect("offset index too small");
2292
2293 assert_eq!(4, page_offset.offset);
2294 assert_eq!(152, page_offset.compressed_page_size);
2295 assert_eq!(0, page_offset.first_row_index);
2296 }
2297
2298 #[test]
2299 #[cfg_attr(miri, ignore)] fn test_page_index_reader_all_type() {
2301 let test_file = get_test_file("alltypes_tiny_pages_plain.parquet");
2302 let builder = ReadOptionsBuilder::new();
2303 let options = builder.with_page_index().build();
2305 let reader_result = SerializedFileReader::new_with_options(test_file, options);
2306 let reader = reader_result.unwrap();
2307
2308 let metadata = reader.metadata();
2310 assert_eq!(metadata.num_row_groups(), 1);
2311
2312 let page_index = metadata.page_index().unwrap();
2313 let row_group_offset_indexes = metadata.page_index_for_row_group(0);
2314
2315 let row_group_metadata = metadata.row_group(0);
2317
2318 let ci = page_index.column_index(0, 0).unwrap();
2320 assert!(!ci.is_sorted());
2321 assert!(matches!(
2322 ci.get_boundary_order(),
2323 Some(BoundaryOrder::UNORDERED)
2324 ));
2325 if let ColumnIndexMetaData::INT32(index) = ci {
2326 check_native_page_index(
2327 index,
2328 325,
2329 get_row_group_min_max_bytes(row_group_metadata, 0),
2330 BoundaryOrder::UNORDERED,
2331 );
2332 assert_eq!(
2333 row_group_offset_indexes.page_locations(0).unwrap().len(),
2334 325
2335 );
2336 } else {
2337 unreachable!()
2338 }
2339 let ci = page_index.column_index(0, 1).unwrap();
2341 assert!(ci.is_sorted());
2342 if let ColumnIndexMetaData::BOOLEAN(index) = ci {
2343 assert_eq!(index.num_pages(), 82);
2344 assert_eq!(
2345 row_group_offset_indexes.page_locations(1).unwrap().len(),
2346 82
2347 );
2348 } else {
2349 unreachable!()
2350 }
2351 let ci = page_index.column_index(0, 2).unwrap();
2353 assert!(ci.is_sorted());
2354 if let ColumnIndexMetaData::INT32(index) = ci {
2355 check_native_page_index(
2356 index,
2357 325,
2358 get_row_group_min_max_bytes(row_group_metadata, 2),
2359 BoundaryOrder::ASCENDING,
2360 );
2361 assert_eq!(
2362 row_group_offset_indexes.page_locations(2).unwrap().len(),
2363 325
2364 );
2365 } else {
2366 unreachable!()
2367 }
2368 let ci = page_index.column_index(0, 3).unwrap();
2370 assert!(ci.is_sorted());
2371 if let ColumnIndexMetaData::INT32(index) = ci {
2372 check_native_page_index(
2373 index,
2374 325,
2375 get_row_group_min_max_bytes(row_group_metadata, 3),
2376 BoundaryOrder::ASCENDING,
2377 );
2378 assert_eq!(
2379 row_group_offset_indexes.page_locations(3).unwrap().len(),
2380 325
2381 );
2382 } else {
2383 unreachable!()
2384 }
2385 let ci = page_index.column_index(0, 4).unwrap();
2387 assert!(ci.is_sorted());
2388 if let ColumnIndexMetaData::INT32(index) = ci {
2389 check_native_page_index(
2390 index,
2391 325,
2392 get_row_group_min_max_bytes(row_group_metadata, 4),
2393 BoundaryOrder::ASCENDING,
2394 );
2395 assert_eq!(
2396 row_group_offset_indexes.page_locations(4).unwrap().len(),
2397 325
2398 );
2399 } else {
2400 unreachable!()
2401 }
2402 let ci = page_index.column_index(0, 5).unwrap();
2404 assert!(!ci.is_sorted());
2405 if let ColumnIndexMetaData::INT64(index) = ci {
2406 check_native_page_index(
2407 index,
2408 528,
2409 get_row_group_min_max_bytes(row_group_metadata, 5),
2410 BoundaryOrder::UNORDERED,
2411 );
2412 assert_eq!(
2413 row_group_offset_indexes.page_locations(5).unwrap().len(),
2414 528
2415 );
2416 } else {
2417 unreachable!()
2418 }
2419 let ci = page_index.column_index(0, 6).unwrap();
2421 assert!(ci.is_sorted());
2422 if let ColumnIndexMetaData::FLOAT(index) = ci {
2423 check_native_page_index(
2424 index,
2425 325,
2426 get_row_group_min_max_bytes(row_group_metadata, 6),
2427 BoundaryOrder::ASCENDING,
2428 );
2429 assert_eq!(
2430 row_group_offset_indexes.page_locations(6).unwrap().len(),
2431 325
2432 );
2433 } else {
2434 unreachable!()
2435 }
2436 let ci = page_index.column_index(0, 7).unwrap();
2438 assert!(!ci.is_sorted());
2439 if let ColumnIndexMetaData::DOUBLE(index) = ci {
2440 check_native_page_index(
2441 index,
2442 528,
2443 get_row_group_min_max_bytes(row_group_metadata, 7),
2444 BoundaryOrder::UNORDERED,
2445 );
2446 assert_eq!(
2447 row_group_offset_indexes.page_locations(7).unwrap().len(),
2448 528
2449 );
2450 } else {
2451 unreachable!()
2452 }
2453 let ci = page_index.column_index(0, 8).unwrap();
2455 assert!(!ci.is_sorted());
2456 if let ColumnIndexMetaData::BYTE_ARRAY(index) = ci {
2457 check_byte_array_page_index(
2458 index,
2459 974,
2460 get_row_group_min_max_bytes(row_group_metadata, 8),
2461 BoundaryOrder::UNORDERED,
2462 );
2463 assert_eq!(
2464 row_group_offset_indexes.page_locations(8).unwrap().len(),
2465 974
2466 );
2467 } else {
2468 unreachable!()
2469 }
2470 let ci = page_index.column_index(0, 9).unwrap();
2472 assert!(ci.is_sorted());
2473 if let ColumnIndexMetaData::BYTE_ARRAY(index) = ci {
2474 check_byte_array_page_index(
2475 index,
2476 352,
2477 get_row_group_min_max_bytes(row_group_metadata, 9),
2478 BoundaryOrder::ASCENDING,
2479 );
2480 assert_eq!(
2481 row_group_offset_indexes.page_locations(9).unwrap().len(),
2482 352
2483 );
2484 } else {
2485 unreachable!()
2486 }
2487 assert!(page_index.column_index(0, 10).is_none());
2490 let ci = page_index.column_index(0, 11).unwrap();
2492 assert!(ci.is_sorted());
2493 if let ColumnIndexMetaData::INT32(index) = ci {
2494 check_native_page_index(
2495 index,
2496 325,
2497 get_row_group_min_max_bytes(row_group_metadata, 11),
2498 BoundaryOrder::ASCENDING,
2499 );
2500 assert_eq!(
2501 row_group_offset_indexes.page_locations(11).unwrap().len(),
2502 325
2503 );
2504 } else {
2505 unreachable!()
2506 }
2507 let ci = page_index.column_index(0, 12).unwrap();
2509 assert!(!ci.is_sorted());
2510 if let ColumnIndexMetaData::INT32(index) = ci {
2511 check_native_page_index(
2512 index,
2513 325,
2514 get_row_group_min_max_bytes(row_group_metadata, 12),
2515 BoundaryOrder::UNORDERED,
2516 );
2517 assert_eq!(
2518 row_group_offset_indexes.page_locations(12).unwrap().len(),
2519 325
2520 );
2521 } else {
2522 unreachable!()
2523 }
2524 }
2525
2526 fn check_native_page_index<T: ParquetValueType>(
2527 row_group_index: &PrimitiveColumnIndex<T>,
2528 page_size: usize,
2529 min_max: (&[u8], &[u8]),
2530 boundary_order: BoundaryOrder,
2531 ) {
2532 assert_eq!(row_group_index.num_pages() as usize, page_size);
2533 assert_eq!(row_group_index.boundary_order, boundary_order);
2534 assert!(row_group_index.min_values().iter().all(|x| {
2535 x >= &T::try_from_le_slice(min_max.0).unwrap()
2536 && x <= &T::try_from_le_slice(min_max.1).unwrap()
2537 }));
2538 }
2539
2540 fn check_byte_array_page_index(
2541 row_group_index: &ByteArrayColumnIndex,
2542 page_size: usize,
2543 min_max: (&[u8], &[u8]),
2544 boundary_order: BoundaryOrder,
2545 ) {
2546 assert_eq!(row_group_index.num_pages() as usize, page_size);
2547 assert_eq!(row_group_index.boundary_order, boundary_order);
2548 for i in 0..row_group_index.num_pages() as usize {
2549 let x = row_group_index.min_value(i).unwrap();
2550 assert!(x >= min_max.0 && x <= min_max.1);
2551 }
2552 }
2553
2554 fn get_row_group_min_max_bytes(r: &RowGroupMetaData, col_num: usize) -> (&[u8], &[u8]) {
2555 let statistics = r.column(col_num).statistics().unwrap();
2556 (
2557 statistics.min_bytes_opt().unwrap_or_default(),
2558 statistics.max_bytes_opt().unwrap_or_default(),
2559 )
2560 }
2561
2562 #[test]
2563 #[cfg_attr(miri, ignore)] fn test_skip_next_page_with_dictionary_page() {
2565 let test_file = get_test_file("alltypes_tiny_pages.parquet");
2566 let builder = ReadOptionsBuilder::new();
2567 let options = builder.with_page_index().build();
2569 let reader_result = SerializedFileReader::new_with_options(test_file, options);
2570 let reader = reader_result.unwrap();
2571
2572 let row_group_reader = reader.get_row_group(0).unwrap();
2573
2574 let mut column_page_reader = row_group_reader.get_column_page_reader(9).unwrap();
2576
2577 let mut vec = vec![];
2578
2579 let meta = column_page_reader.peek_next_page().unwrap().unwrap();
2581 assert!(meta.is_dict);
2582
2583 column_page_reader.skip_next_page().unwrap();
2585
2586 let page = column_page_reader.get_next_page().unwrap().unwrap();
2588 assert!(matches!(page.page_type(), basic::PageType::DATA_PAGE));
2589
2590 for _i in 0..351 {
2592 let meta = column_page_reader.peek_next_page().unwrap().unwrap();
2594 assert!(!meta.is_dict); vec.push(meta);
2596
2597 let page = column_page_reader.get_next_page().unwrap().unwrap();
2598 assert!(matches!(page.page_type(), basic::PageType::DATA_PAGE));
2599 }
2600
2601 assert!(column_page_reader.peek_next_page().unwrap().is_none());
2603 assert!(column_page_reader.get_next_page().unwrap().is_none());
2604
2605 assert_eq!(vec.len(), 351);
2607 }
2608
2609 #[test]
2610 #[cfg_attr(miri, ignore)] fn test_skip_page_with_offset_index() {
2612 let test_file = get_test_file("alltypes_tiny_pages_plain.parquet");
2613 let builder = ReadOptionsBuilder::new();
2614 let options = builder.with_page_index().build();
2616 let reader_result = SerializedFileReader::new_with_options(test_file, options);
2617 let reader = reader_result.unwrap();
2618
2619 let row_group_reader = reader.get_row_group(0).unwrap();
2620
2621 let mut column_page_reader = row_group_reader.get_column_page_reader(4).unwrap();
2623
2624 let mut vec = vec![];
2625
2626 for i in 0..325 {
2627 if i % 2 == 0 {
2628 vec.push(column_page_reader.get_next_page().unwrap().unwrap());
2629 } else {
2630 column_page_reader.skip_next_page().unwrap();
2631 }
2632 }
2633 assert!(column_page_reader.peek_next_page().unwrap().is_none());
2635 assert!(column_page_reader.get_next_page().unwrap().is_none());
2636
2637 assert_eq!(vec.len(), 163);
2638 }
2639
2640 #[test]
2641 fn test_skip_page_without_offset_index() {
2642 let test_file = get_test_file("alltypes_tiny_pages_plain.parquet");
2643
2644 let reader_result = SerializedFileReader::new(test_file);
2646 let reader = reader_result.unwrap();
2647
2648 let row_group_reader = reader.get_row_group(0).unwrap();
2649
2650 let mut column_page_reader = row_group_reader.get_column_page_reader(4).unwrap();
2652
2653 let mut vec = vec![];
2654
2655 for i in 0..325 {
2656 if i % 2 == 0 {
2657 vec.push(column_page_reader.get_next_page().unwrap().unwrap());
2658 } else {
2659 column_page_reader.peek_next_page().unwrap().unwrap();
2660 column_page_reader.skip_next_page().unwrap();
2661 }
2662 }
2663 assert!(column_page_reader.peek_next_page().unwrap().is_none());
2665 assert!(column_page_reader.get_next_page().unwrap().is_none());
2666
2667 assert_eq!(vec.len(), 163);
2668 }
2669
2670 #[test]
2671 #[cfg_attr(miri, ignore)] fn test_peek_page_with_dictionary_page() {
2673 let test_file = get_test_file("alltypes_tiny_pages.parquet");
2674 let builder = ReadOptionsBuilder::new();
2675 let options = builder.with_page_index().build();
2677 let reader_result = SerializedFileReader::new_with_options(test_file, options);
2678 let reader = reader_result.unwrap();
2679 let row_group_reader = reader.get_row_group(0).unwrap();
2680
2681 let mut column_page_reader = row_group_reader.get_column_page_reader(9).unwrap();
2683
2684 let mut vec = vec![];
2685
2686 let meta = column_page_reader.peek_next_page().unwrap().unwrap();
2687 assert!(meta.is_dict);
2688 let page = column_page_reader.get_next_page().unwrap().unwrap();
2689 assert!(matches!(page.page_type(), basic::PageType::DICTIONARY_PAGE));
2690
2691 for i in 0..352 {
2692 let meta = column_page_reader.peek_next_page().unwrap().unwrap();
2693 if i != 351 {
2696 assert!((meta.num_rows == Some(21)) || (meta.num_rows == Some(20)));
2697 } else {
2698 assert_eq!(meta.num_rows, Some(10));
2701 }
2702 assert!(!meta.is_dict);
2703 vec.push(meta);
2704 let page = column_page_reader.get_next_page().unwrap().unwrap();
2705 assert!(matches!(page.page_type(), basic::PageType::DATA_PAGE));
2706 }
2707
2708 assert!(column_page_reader.peek_next_page().unwrap().is_none());
2710 assert!(column_page_reader.get_next_page().unwrap().is_none());
2711
2712 assert_eq!(vec.len(), 352);
2713 }
2714
2715 #[test]
2716 fn test_peek_page_with_dictionary_page_without_offset_index() {
2717 let test_file = get_test_file("alltypes_tiny_pages.parquet");
2718
2719 let reader_result = SerializedFileReader::new(test_file);
2720 let reader = reader_result.unwrap();
2721 let row_group_reader = reader.get_row_group(0).unwrap();
2722
2723 let mut column_page_reader = row_group_reader.get_column_page_reader(9).unwrap();
2725
2726 let mut vec = vec![];
2727
2728 let meta = column_page_reader.peek_next_page().unwrap().unwrap();
2729 assert!(meta.is_dict);
2730 let page = column_page_reader.get_next_page().unwrap().unwrap();
2731 assert!(matches!(page.page_type(), basic::PageType::DICTIONARY_PAGE));
2732
2733 for i in 0..352 {
2734 let meta = column_page_reader.peek_next_page().unwrap().unwrap();
2735 if i != 351 {
2738 assert!((meta.num_levels == Some(21)) || (meta.num_levels == Some(20)));
2739 } else {
2740 assert_eq!(meta.num_levels, Some(10));
2743 }
2744 assert!(!meta.is_dict);
2745 vec.push(meta);
2746 let page = column_page_reader.get_next_page().unwrap().unwrap();
2747 assert!(matches!(page.page_type(), basic::PageType::DATA_PAGE));
2748 }
2749
2750 assert!(column_page_reader.peek_next_page().unwrap().is_none());
2752 assert!(column_page_reader.get_next_page().unwrap().is_none());
2753
2754 assert_eq!(vec.len(), 352);
2755 }
2756
2757 #[test]
2758 fn test_fixed_length_index() {
2759 let message_type = "
2760 message test_schema {
2761 OPTIONAL FIXED_LEN_BYTE_ARRAY (11) value (DECIMAL(25,2));
2762 }
2763 ";
2764
2765 let schema = parse_message_type(message_type).unwrap();
2766 let mut out = Vec::with_capacity(1024);
2767 let mut writer =
2768 SerializedFileWriter::new(&mut out, Arc::new(schema), Default::default()).unwrap();
2769
2770 let mut r = writer.next_row_group().unwrap();
2771 let mut c = r.next_column().unwrap().unwrap();
2772 c.typed::<FixedLenByteArrayType>()
2773 .write_batch(
2774 &[vec![0; 11].into(), vec![5; 11].into(), vec![3; 11].into()],
2775 Some(&[1, 1, 0, 1]),
2776 None,
2777 )
2778 .unwrap();
2779 c.close().unwrap();
2780 r.close().unwrap();
2781 writer.close().unwrap();
2782
2783 let b = Bytes::from(out);
2784 let options = ReadOptionsBuilder::new().with_page_index().build();
2785 let reader = SerializedFileReader::new_with_options(b, options).unwrap();
2786 let page_index = reader.metadata().page_index().unwrap();
2787
2788 match page_index.column_index(0, 0) {
2789 Some(ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(v)) => {
2790 assert_eq!(v.num_pages(), 1);
2791 assert_eq!(v.null_count(0).unwrap(), 1);
2792 assert_eq!(v.min_value(0).unwrap(), &[0; 11]);
2793 assert_eq!(v.max_value(0).unwrap(), &[5; 11]);
2794 }
2795 _ => unreachable!(),
2796 }
2797 }
2798
2799 #[test]
2800 fn test_multi_gz() {
2801 let file = get_test_file("concatenated_gzip_members.parquet");
2802 let reader = SerializedFileReader::new(file).unwrap();
2803 let row_group_reader = reader.get_row_group(0).unwrap();
2804 match row_group_reader.get_column_reader(0).unwrap() {
2805 ColumnReader::Int64ColumnReader(mut reader) => {
2806 let mut buffer = Vec::with_capacity(1024);
2807 let mut def_levels = Vec::with_capacity(1024);
2808 let (num_records, num_values, num_levels) = reader
2809 .read_records(1024, Some(&mut def_levels), None, &mut buffer)
2810 .unwrap();
2811
2812 assert_eq!(num_records, 513);
2813 assert_eq!(num_values, 513);
2814 assert_eq!(num_levels, 513);
2815
2816 let expected: Vec<i64> = (1..514).collect();
2817 assert_eq!(&buffer, &expected);
2818 }
2819 _ => unreachable!(),
2820 }
2821 }
2822
2823 #[test]
2824 #[cfg_attr(miri, ignore)] fn test_byte_stream_split_extended() {
2826 let path = format!(
2827 "{}/byte_stream_split_extended.gzip.parquet",
2828 arrow::util::test_util::parquet_test_data(),
2829 );
2830 let file = File::open(path).unwrap();
2831 let reader = Box::new(SerializedFileReader::new(file).expect("Failed to create reader"));
2832
2833 let mut iter = reader
2835 .get_row_iter(None)
2836 .expect("Failed to create row iterator");
2837
2838 let mut start = 0;
2839 let end = reader.metadata().file_metadata().num_rows();
2840
2841 let check_row = |row: Result<Row, ParquetError>| {
2842 assert!(row.is_ok());
2843 let r = row.unwrap();
2844 assert_eq!(r.get_float16(0).unwrap(), r.get_float16(1).unwrap());
2845 assert_eq!(r.get_float(2).unwrap(), r.get_float(3).unwrap());
2846 assert_eq!(r.get_double(4).unwrap(), r.get_double(5).unwrap());
2847 assert_eq!(r.get_int(6).unwrap(), r.get_int(7).unwrap());
2848 assert_eq!(r.get_long(8).unwrap(), r.get_long(9).unwrap());
2849 assert_eq!(r.get_bytes(10).unwrap(), r.get_bytes(11).unwrap());
2850 assert_eq!(r.get_decimal(12).unwrap(), r.get_decimal(13).unwrap());
2851 };
2852
2853 while start < end {
2854 match iter.next() {
2855 Some(row) => check_row(row),
2856 None => break,
2857 }
2858 start += 1;
2859 }
2860 }
2861
2862 #[test]
2863 fn test_filtered_rowgroup_metadata() {
2864 let message_type = "
2865 message test_schema {
2866 REQUIRED INT32 a;
2867 }
2868 ";
2869 let schema = Arc::new(parse_message_type(message_type).unwrap());
2870 let props = Arc::new(
2871 WriterProperties::builder()
2872 .set_statistics_enabled(EnabledStatistics::Page)
2873 .build(),
2874 );
2875 let mut file: File = tempfile::tempfile().unwrap();
2876 let mut file_writer = SerializedFileWriter::new(&mut file, schema, props).unwrap();
2877 let data = [1, 2, 3, 4, 5];
2878
2879 for idx in 0..5 {
2881 let data_i: Vec<i32> = data.iter().map(|x| x * (idx + 1)).collect();
2882 let mut row_group_writer = file_writer.next_row_group().unwrap();
2883 if let Some(mut writer) = row_group_writer.next_column().unwrap() {
2884 writer
2885 .typed::<Int32Type>()
2886 .write_batch(data_i.as_slice(), None, None)
2887 .unwrap();
2888 writer.close().unwrap();
2889 }
2890 row_group_writer.close().unwrap();
2891 file_writer.flushed_row_groups();
2892 }
2893 let file_metadata = file_writer.close().unwrap();
2894
2895 assert_eq!(file_metadata.file_metadata().num_rows(), 25);
2896 assert_eq!(file_metadata.num_row_groups(), 5);
2897
2898 for filter in [false, true] {
2901 let calls = Rc::new(RefCell::new(Vec::new()));
2902 let first_calls = Rc::clone(&calls);
2903 let second_calls = Rc::clone(&calls);
2904 let mut next_index = 0;
2905 let read_options = ReadOptionsBuilder::new()
2906 .with_page_index()
2907 .with_predicate(Box::new(move |rg_meta, i| {
2908 assert_eq!(i, next_index);
2909 assert_eq!(rg_meta.ordinal(), Some(i as i32));
2910 next_index += 1;
2911 first_calls.borrow_mut().push((i, 0));
2912 !filter || i != 1
2913 }))
2914 .with_predicate(Box::new(move |_, i| {
2915 second_calls.borrow_mut().push((i, 1));
2916 !filter || i != 3
2917 }))
2918 .build();
2919 let reader =
2920 SerializedFileReader::new_with_options(file.try_clone().unwrap(), read_options)
2921 .unwrap();
2922 let metadata = reader.metadata();
2923 let expected_ordinals = if filter {
2924 vec![0, 2, 4]
2925 } else {
2926 vec![0, 1, 2, 3, 4]
2927 };
2928 let ordinals: Vec<_> = metadata
2929 .row_groups()
2930 .iter()
2931 .map(|rg| rg.ordinal().unwrap())
2932 .collect();
2933 assert_eq!(ordinals, expected_ordinals);
2934 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
2935
2936 let expected_calls: Vec<_> = (0..5)
2937 .flat_map(|i| {
2938 if filter && i == 1 {
2939 vec![(i, 0)]
2940 } else {
2941 vec![(i, 0), (i, 1)]
2942 }
2943 })
2944 .collect();
2945 assert_eq!(*calls.borrow(), expected_calls);
2946 }
2947
2948 let read_options = ReadOptionsBuilder::new()
2950 .with_page_index()
2951 .with_predicate(Box::new(|rgmeta, _| rgmeta.ordinal().unwrap_or(0) == 2))
2952 .build();
2953 let reader =
2954 SerializedFileReader::new_with_options(file.try_clone().unwrap(), read_options)
2955 .unwrap();
2956 let metadata = reader.metadata();
2957
2958 assert_eq!(metadata.num_row_groups(), 1);
2960 assert_eq!(metadata.row_group(0).ordinal(), Some(2));
2961
2962 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
2964 let page_index = metadata.page_index().unwrap();
2965
2966 let col_stats = metadata.row_group(0).column(0).statistics().unwrap();
2967 let pg_idx = page_index.column_index(0, 0);
2968 let off_idx_i = page_index.offset_index(0, 0);
2969
2970 match pg_idx {
2972 Some(ColumnIndexMetaData::INT32(int_idx)) => {
2973 let min = col_stats.min_bytes_opt().unwrap().get_i32_le();
2974 let max = col_stats.max_bytes_opt().unwrap().get_i32_le();
2975 assert_eq!(int_idx.min_value(0), Some(min).as_ref());
2976 assert_eq!(int_idx.max_value(0), Some(max).as_ref());
2977 }
2978 _ => panic!("wrong stats type"),
2979 }
2980
2981 assert_eq!(
2983 off_idx_i.as_ref().unwrap().page_locations[0].offset,
2984 metadata.row_group(0).column(0).data_page_offset()
2985 );
2986
2987 let read_options = ReadOptionsBuilder::new()
2989 .with_page_index()
2990 .with_predicate(Box::new(|rgmeta, _| rgmeta.ordinal().unwrap_or(0) % 2 == 1))
2991 .build();
2992 let reader =
2993 SerializedFileReader::new_with_options(file.try_clone().unwrap(), read_options)
2994 .unwrap();
2995 let metadata = reader.metadata();
2996
2997 assert_eq!(metadata.num_row_groups(), 2);
2999 assert_eq!(metadata.row_group(0).ordinal(), Some(1));
3000 assert_eq!(metadata.row_group(1).ordinal(), Some(3));
3001
3002 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
3004
3005 let page_index = metadata.page_index().unwrap();
3006
3007 for rg_idx in 0..metadata.num_row_groups() {
3008 let col_stats = metadata.row_group(rg_idx).column(0).statistics().unwrap();
3009 let pg_idx = page_index.column_index(rg_idx, 0);
3010 let off_idx_i = page_index.offset_index(rg_idx, 0);
3011
3012 match pg_idx {
3014 Some(ColumnIndexMetaData::INT32(int_idx)) => {
3015 let min = col_stats.min_bytes_opt().unwrap().get_i32_le();
3016 let max = col_stats.max_bytes_opt().unwrap().get_i32_le();
3017 assert_eq!(int_idx.min_value(0), Some(min).as_ref());
3018 assert_eq!(int_idx.max_value(0), Some(max).as_ref());
3019 }
3020 _ => panic!("wrong stats type"),
3021 }
3022
3023 assert_eq!(
3025 off_idx_i.as_ref().unwrap().page_locations[0].offset,
3026 metadata.row_group(rg_idx).column(0).data_page_offset()
3027 );
3028 }
3029 }
3030
3031 #[test]
3032 fn test_reuse_schema() {
3033 let file = get_test_file("alltypes_plain.parquet");
3034 let file_reader = SerializedFileReader::new(file.try_clone().unwrap()).unwrap();
3035 let schema = file_reader.metadata().file_metadata().schema_descr_ptr();
3036 let expected = file_reader.metadata;
3037
3038 let options = ReadOptionsBuilder::new()
3039 .with_parquet_schema(schema)
3040 .build();
3041 let file_reader = SerializedFileReader::new_with_options(file, options).unwrap();
3042
3043 assert_eq!(expected.as_ref(), file_reader.metadata.as_ref());
3044 assert!(Arc::ptr_eq(
3046 &expected.file_metadata().schema_descr_ptr(),
3047 &file_reader.metadata.file_metadata().schema_descr_ptr()
3048 ));
3049 }
3050
3051 #[test]
3052 fn test_read_unknown_logical_type() {
3053 let file = get_test_file("unknown-logical-type.parquet");
3054 let reader = SerializedFileReader::new(file).expect("Error opening file");
3055
3056 let schema = reader.metadata().file_metadata().schema_descr();
3057 assert_eq!(
3058 schema.column(0).logical_type_ref(),
3059 Some(&basic::LogicalType::String)
3060 );
3061 assert_eq!(
3062 schema.column(1).logical_type_ref(),
3063 Some(&basic::LogicalType::_Unknown { field_id: 2555 })
3064 );
3065 assert_eq!(schema.column(1).physical_type(), Type::BYTE_ARRAY);
3066
3067 let mut iter = reader
3068 .get_row_iter(None)
3069 .expect("Failed to create row iterator");
3070
3071 let mut num_rows = 0;
3072 while iter.next().is_some() {
3073 num_rows += 1;
3074 }
3075 assert_eq!(num_rows, reader.metadata().file_metadata().num_rows());
3076 }
3077
3078 fn dict_column_file() -> (Bytes, std::ops::Range<usize>) {
3084 let message_type = "
3085 message test_schema {
3086 REQUIRED INT32 a;
3087 }
3088 ";
3089 let schema = Arc::new(parse_message_type(message_type).unwrap());
3090 let props = Arc::new(
3091 WriterProperties::builder()
3092 .set_compression(Compression::SNAPPY)
3093 .build(),
3094 );
3095 let mut out = Vec::new();
3096 let mut writer = SerializedFileWriter::new(&mut out, schema, props).unwrap();
3097 let mut rg = writer.next_row_group().unwrap();
3098 let mut col = rg.next_column().unwrap().unwrap();
3099 col.typed::<Int32Type>()
3100 .write_batch(&[1, 2, 3, 4, 5], None, None)
3101 .unwrap();
3102 col.close().unwrap();
3103 rg.close().unwrap();
3104 writer.close().unwrap();
3105 let bytes = Bytes::from(out);
3106
3107 let reader = SerializedFileReader::new(bytes.clone()).unwrap();
3108 let col_meta = reader.metadata().row_group(0).column(0).clone();
3109 let dict_offset = col_meta.dictionary_page_offset().unwrap() as usize;
3110 let data_offset = col_meta.data_page_offset() as usize;
3111 let mut prot = ThriftSliceInputProtocol::new(&bytes[dict_offset..data_offset]);
3113 let _header = PageHeader::read_thrift(&mut prot).unwrap();
3114 let header_len = (data_offset - dict_offset) - prot.as_slice().len();
3115 (bytes, dict_offset + header_len..data_offset)
3116 }
3117
3118 fn corrupt_dictionary_body(bytes: &Bytes, body: &std::ops::Range<usize>) -> Bytes {
3124 let mut v = bytes.to_vec();
3125 v[body.clone()].fill(0);
3126 Bytes::from(v)
3127 }
3128
3129 fn int32_reader(bytes: Bytes) -> crate::column::reader::ColumnReaderImpl<Int32Type> {
3130 let reader = SerializedFileReader::new(bytes.clone()).unwrap();
3131 let meta = reader.metadata();
3132 let descr = meta.file_metadata().schema_descr().column(0);
3133 let col_meta = meta.row_group(0).column(0).clone();
3134 let total_rows = meta.row_group(0).num_rows() as usize;
3135 let pages =
3136 SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, None).unwrap();
3137 crate::column::reader::ColumnReaderImpl::<Int32Type>::new(descr, Box::new(pages))
3138 }
3139
3140 #[test]
3145 fn skipping_a_whole_chunk_never_decodes_the_dictionary() {
3146 let (bytes, body) = dict_column_file();
3147 let corrupted = corrupt_dictionary_body(&bytes, &body);
3148 let mut reader = int32_reader(corrupted);
3149 let skipped = reader.skip_records(5).unwrap();
3150 assert_eq!(skipped, 5);
3151 }
3152
3153 #[test]
3157 fn a_partial_skip_pays_the_deferred_dictionary() {
3158 let (bytes, body) = dict_column_file();
3159 let corrupted = corrupt_dictionary_body(&bytes, &body);
3160 let mut reader = int32_reader(corrupted);
3161 let err = reader.skip_records(2).unwrap_err().to_string();
3162 assert!(
3163 err.contains("corrupt") || err.contains("size"),
3164 "expected the deferred dictionary decode to fail on \
3165 decompression, got: {err}"
3166 );
3167 }
3168
3169 #[test]
3172 fn skip_then_read_decodes_correct_values() {
3173 let (bytes, _) = dict_column_file();
3174 let mut reader = int32_reader(bytes);
3175 assert_eq!(reader.skip_records(2).unwrap(), 2);
3176 let mut values = Vec::new();
3177 let (records, _, read) = reader.read_records(3, None, None, &mut values).unwrap();
3178 assert_eq!(records, 3);
3179 assert_eq!(read, 3);
3180 assert_eq!(values, vec![3, 4, 5]);
3181 }
3182
3183 #[test]
3186 fn plain_read_is_unchanged() {
3187 let (bytes, _) = dict_column_file();
3188 let mut reader = int32_reader(bytes);
3189 let mut values = Vec::new();
3190 let (records, _, read) = reader.read_records(5, None, None, &mut values).unwrap();
3191 assert_eq!((records, read), (5, 5));
3192 assert_eq!(values, vec![1, 2, 3, 4, 5]);
3193 }
3194
3195 struct NonDeferringPageReader(SerializedPageReader<Bytes>);
3201
3202 impl Iterator for NonDeferringPageReader {
3203 type Item = Result<Page>;
3204
3205 fn next(&mut self) -> Option<Self::Item> {
3206 self.get_next_page().transpose()
3207 }
3208 }
3209
3210 impl PageReader for NonDeferringPageReader {
3211 fn get_next_page(&mut self) -> Result<Option<Page>> {
3212 self.0.get_next_page()
3213 }
3214
3215 fn peek_next_page(&mut self) -> Result<Option<PageMetadata>> {
3216 self.0.peek_next_page()
3217 }
3218
3219 fn skip_next_page(&mut self) -> Result<()> {
3220 self.0.skip_next_page()
3221 }
3222 }
3223
3224 fn non_deferring_int32_reader(
3225 bytes: Bytes,
3226 ) -> crate::column::reader::ColumnReaderImpl<Int32Type> {
3227 let reader = SerializedFileReader::new(bytes.clone()).unwrap();
3228 let meta = reader.metadata();
3229 let descr = meta.file_metadata().schema_descr().column(0);
3230 let col_meta = meta.row_group(0).column(0).clone();
3231 let total_rows = meta.row_group(0).num_rows() as usize;
3232 let pages =
3233 SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, None).unwrap();
3234 crate::column::reader::ColumnReaderImpl::<Int32Type>::new(
3235 descr,
3236 Box::new(NonDeferringPageReader(pages)),
3237 )
3238 }
3239
3240 #[test]
3248 fn a_reader_that_cannot_defer_still_gets_its_dictionary() {
3249 let (bytes, _) = dict_column_file();
3250 let mut reader = non_deferring_int32_reader(bytes);
3251
3252 assert_eq!(reader.skip_records(2).unwrap(), 2);
3253
3254 let mut values = Vec::new();
3255 let (records, _, read) = reader.read_records(3, None, None, &mut values).unwrap();
3256 assert_eq!(records, 3);
3257 assert_eq!(read, 3);
3258 assert_eq!(values, vec![3, 4, 5]);
3259 }
3260
3261 #[test]
3265 fn a_reader_that_cannot_defer_survives_a_whole_chunk_skip() {
3266 let (bytes, _) = dict_column_file();
3267 let mut reader = non_deferring_int32_reader(bytes);
3268
3269 assert_eq!(reader.skip_records(5).unwrap(), 5);
3270
3271 let mut values = Vec::new();
3272 let (records, _, _) = reader.read_records(1, None, None, &mut values).unwrap();
3273 assert_eq!(records, 0, "chunk is exhausted");
3274 assert!(values.is_empty());
3275 }
3276
3277 #[test]
3281 fn values_state_supports_deferral() {
3282 let (bytes, _) = dict_column_file();
3283 let reader = SerializedFileReader::new(bytes.clone()).unwrap();
3284 let meta = reader.metadata();
3285 let col_meta = meta.row_group(0).column(0).clone();
3286 let total_rows = meta.row_group(0).num_rows() as usize;
3287 let pages =
3288 SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, None).unwrap();
3289
3290 assert!(pages.supports_deferred_dictionary());
3291 }
3292
3293 #[test]
3297 fn offset_index_state_with_a_dictionary_offset_supports_deferral() {
3298 let (bytes, _) = dict_column_file();
3299 let reader = SerializedFileReader::new(bytes.clone()).unwrap();
3300 let meta = reader.metadata();
3301 let col_meta = meta.row_group(0).column(0).clone();
3302 let total_rows = meta.row_group(0).num_rows() as usize;
3303 let (dict_offset, data_offset) = (
3304 col_meta.dictionary_page_offset().unwrap(),
3305 col_meta.data_page_offset(),
3306 );
3307 assert_ne!(
3308 dict_offset, data_offset,
3309 "fixture must record a dictionary offset"
3310 );
3311
3312 let locations = vec![PageLocation {
3313 offset: data_offset,
3314 compressed_page_size: 0,
3315 first_row_index: 0,
3316 }];
3317 let pages =
3318 SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, Some(locations))
3319 .unwrap();
3320
3321 assert!(pages.supports_deferred_dictionary());
3322 }
3323
3324 #[test]
3330 fn offset_index_state_without_a_dictionary_offset_declines_deferral() {
3331 let (bytes, _) = dict_column_file();
3332 let reader = SerializedFileReader::new(bytes.clone()).unwrap();
3333 let meta = reader.metadata();
3334 let total_rows = meta.row_group(0).num_rows() as usize;
3335
3336 let col_meta = meta.row_group(0).column(0).clone();
3339 let stripped = col_meta
3340 .clone()
3341 .into_builder()
3342 .set_dictionary_page_offset(None)
3343 .build()
3344 .unwrap();
3345
3346 let locations = vec![PageLocation {
3347 offset: stripped.data_page_offset(),
3348 compressed_page_size: 0,
3349 first_row_index: 0,
3350 }];
3351 let pages =
3352 SerializedPageReader::new(Arc::new(bytes), &stripped, total_rows, Some(locations))
3353 .unwrap();
3354
3355 assert!(
3356 !pages.supports_deferred_dictionary(),
3357 "an offset-less dictionary cannot be recovered from the offset-index state"
3358 );
3359 }
3364
3365 #[test]
3370 fn alternating_skip_and_read_decodes_correct_values() {
3371 let (bytes, _) = dict_column_file();
3372 let mut reader = int32_reader(bytes);
3373
3374 assert_eq!(reader.skip_records(1).unwrap(), 1);
3375 let mut first = Vec::new();
3376 let (records, _, _) = reader.read_records(2, None, None, &mut first).unwrap();
3377 assert_eq!(records, 2);
3378 assert_eq!(first, vec![2, 3]);
3379
3380 assert_eq!(reader.skip_records(1).unwrap(), 1);
3381 let mut second = Vec::new();
3382 let (records, _, _) = reader.read_records(1, None, None, &mut second).unwrap();
3383 assert_eq!(records, 1);
3384 assert_eq!(second, vec![5]);
3385 }
3386
3387 struct HalfForwardingPageReader(SerializedPageReader<Bytes>);
3391
3392 impl Iterator for HalfForwardingPageReader {
3393 type Item = Result<Page>;
3394
3395 fn next(&mut self) -> Option<Self::Item> {
3396 self.get_next_page().transpose()
3397 }
3398 }
3399
3400 impl PageReader for HalfForwardingPageReader {
3401 fn get_next_page(&mut self) -> Result<Option<Page>> {
3402 self.0.get_next_page()
3403 }
3404
3405 fn peek_next_page(&mut self) -> Result<Option<PageMetadata>> {
3406 self.0.peek_next_page()
3407 }
3408
3409 fn skip_next_page(&mut self) -> Result<()> {
3410 self.0.skip_next_page()
3411 }
3412
3413 fn take_deferred_dictionary(&mut self) -> Result<Option<Page>> {
3414 self.0.take_deferred_dictionary()
3415 }
3416 }
3417
3418 #[test]
3423 fn forwarding_only_take_deferred_dictionary_is_still_correct() {
3424 let (bytes, _) = dict_column_file();
3425 let reader = SerializedFileReader::new(bytes.clone()).unwrap();
3426 let meta = reader.metadata();
3427 let descr = meta.file_metadata().schema_descr().column(0);
3428 let col_meta = meta.row_group(0).column(0).clone();
3429 let total_rows = meta.row_group(0).num_rows() as usize;
3430 let pages =
3431 SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, None).unwrap();
3432 let mut reader = crate::column::reader::ColumnReaderImpl::<Int32Type>::new(
3433 descr,
3434 Box::new(HalfForwardingPageReader(pages)),
3435 );
3436
3437 assert_eq!(reader.skip_records(2).unwrap(), 2);
3438 let mut values = Vec::new();
3439 let (records, _, _) = reader.read_records(3, None, None, &mut values).unwrap();
3440 assert_eq!(records, 3);
3441 assert_eq!(values, vec![3, 4, 5]);
3442 }
3443}