1use bytes::Bytes;
8
9use super::page::{Page, PageReader};
10use crate::basic::*;
11use crate::column::reader::decoder::{
12 ColumnValueDecoder, ColumnValueDecoderImpl, DefinitionLevelDecoder, DefinitionLevelDecoderImpl,
13 RepetitionLevelDecoder, RepetitionLevelDecoderImpl, normalize_fixed_len_byte_array_data,
14 prepare_v1_fixed_len_byte_array,
15};
16use crate::data_type::*;
17use crate::errors::{ParquetError, Result};
18use crate::schema::types::ColumnDescPtr;
19use crate::util::bit_util::{ceil, num_required_bits, read_num_bytes};
20
21pub(crate) mod decoder;
22mod legacy_fixed_len_byte_array;
23
24pub enum ColumnReader {
26 BoolColumnReader(ColumnReaderImpl<BoolType>),
28 Int32ColumnReader(ColumnReaderImpl<Int32Type>),
30 Int64ColumnReader(ColumnReaderImpl<Int64Type>),
32 Int96ColumnReader(ColumnReaderImpl<Int96Type>),
34 FloatColumnReader(ColumnReaderImpl<FloatType>),
36 DoubleColumnReader(ColumnReaderImpl<DoubleType>),
38 ByteArrayColumnReader(ColumnReaderImpl<ByteArrayType>),
40 FixedLenByteArrayColumnReader(ColumnReaderImpl<FixedLenByteArrayType>),
42}
43
44pub fn get_column_reader(
47 col_descr: ColumnDescPtr,
48 col_page_reader: Box<dyn PageReader>,
49) -> ColumnReader {
50 match col_descr.physical_type() {
51 Type::BOOLEAN => {
52 ColumnReader::BoolColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader))
53 }
54 Type::INT32 => {
55 ColumnReader::Int32ColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader))
56 }
57 Type::INT64 => {
58 ColumnReader::Int64ColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader))
59 }
60 Type::INT96 => {
61 ColumnReader::Int96ColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader))
62 }
63 Type::FLOAT => {
64 ColumnReader::FloatColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader))
65 }
66 Type::DOUBLE => {
67 ColumnReader::DoubleColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader))
68 }
69 Type::BYTE_ARRAY => {
70 ColumnReader::ByteArrayColumnReader(ColumnReaderImpl::new(col_descr, col_page_reader))
71 }
72 Type::FIXED_LEN_BYTE_ARRAY => ColumnReader::FixedLenByteArrayColumnReader(
73 ColumnReaderImpl::new(col_descr, col_page_reader),
74 ),
75 }
76}
77
78pub fn get_typed_column_reader<T: DataType>(col_reader: ColumnReader) -> ColumnReaderImpl<T> {
85 T::get_column_reader(col_reader).unwrap_or_else(|| {
86 panic!(
87 "Failed to convert column reader into a typed column reader for `{}` type",
88 T::get_physical_type()
89 )
90 })
91}
92
93pub type ColumnReaderImpl<T> = GenericColumnReader<
95 RepetitionLevelDecoderImpl,
96 DefinitionLevelDecoderImpl,
97 ColumnValueDecoderImpl<T>,
98>;
99
100pub struct GenericColumnReader<R, D, V> {
106 descr: ColumnDescPtr,
107
108 page_reader: Box<dyn PageReader>,
109
110 num_buffered_values: usize,
112
113 num_decoded_values: usize,
116
117 has_record_delimiter: bool,
119
120 def_level_decoder: Option<D>,
122
123 rep_level_decoder: Option<R>,
125
126 values_decoder: V,
128
129 remaining_fixed_len_bytes: Option<(usize, usize)>,
132}
133
134impl<V> GenericColumnReader<RepetitionLevelDecoderImpl, DefinitionLevelDecoderImpl, V>
135where
136 V: ColumnValueDecoder,
137{
138 pub fn new(descr: ColumnDescPtr, page_reader: Box<dyn PageReader>) -> Self {
140 let values_decoder = V::new(&descr);
141
142 let def_level_decoder = (descr.max_def_level() != 0)
143 .then(|| DefinitionLevelDecoderImpl::new(descr.max_def_level()));
144
145 let rep_level_decoder = (descr.max_rep_level() != 0)
146 .then(|| RepetitionLevelDecoderImpl::new(descr.max_rep_level()));
147
148 Self::new_with_decoders(
149 descr,
150 page_reader,
151 values_decoder,
152 def_level_decoder,
153 rep_level_decoder,
154 )
155 }
156}
157
158impl<R, D, V> GenericColumnReader<R, D, V>
159where
160 R: RepetitionLevelDecoder,
161 D: DefinitionLevelDecoder,
162 V: ColumnValueDecoder,
163{
164 pub(crate) fn new_with_decoders(
165 descr: ColumnDescPtr,
166 page_reader: Box<dyn PageReader>,
167 values_decoder: V,
168 def_level_decoder: Option<D>,
169 rep_level_decoder: Option<R>,
170 ) -> Self {
171 Self {
172 descr,
173 def_level_decoder,
174 rep_level_decoder,
175 page_reader,
176 num_buffered_values: 0,
177 num_decoded_values: 0,
178 values_decoder,
179 has_record_delimiter: false,
180 remaining_fixed_len_bytes: None,
181 }
182 }
183
184 pub fn read_records(
203 &mut self,
204 max_records: usize,
205 def_levels: Option<&mut D::Buffer>,
206 rep_levels: Option<&mut R::Buffer>,
207 values: &mut V::Buffer,
208 ) -> Result<(usize, usize, usize)> {
209 self.read_records_with_reservation(
210 max_records,
211 def_levels,
212 rep_levels,
213 values,
214 |_, _, _, _| Ok(()),
215 )
216 }
217
218 pub(crate) fn read_records_with_reservation<F>(
221 &mut self,
222 max_records: usize,
223 mut def_levels: Option<&mut D::Buffer>,
224 mut rep_levels: Option<&mut R::Buffer>,
225 values: &mut V::Buffer,
226 mut reserve_values: F,
227 ) -> Result<(usize, usize, usize)>
228 where
229 F: FnMut(&mut V::Buffer, usize, usize, Option<&D::Buffer>) -> Result<()>,
230 {
231 let mut total_records_read = 0;
232 let mut total_levels_read = 0;
233 let mut total_values_read = 0;
234
235 while total_records_read < max_records && self.has_next()? {
236 let remaining_records = max_records - total_records_read;
237 let remaining_levels = self.num_buffered_values - self.num_decoded_values;
238
239 let (records_read, levels_to_read) = match self.rep_level_decoder.as_mut() {
240 Some(reader) => {
241 let out = rep_levels
242 .as_mut()
243 .ok_or_else(|| general_err!("must specify repetition levels"))?;
244
245 let (mut records_read, levels_read) =
246 reader.read_rep_levels(out, remaining_records, remaining_levels)?;
247
248 if records_read == 0 && levels_read == 0 {
249 return Err(general_err!(
251 "Insufficient repetition levels read from column"
252 ));
253 }
254 if levels_read == remaining_levels && self.has_record_delimiter {
255 check_partial_record_fits(records_read, remaining_records)?;
256 records_read += reader.flush_partial() as usize;
257 }
258 (records_read, levels_read)
259 }
260 None => {
261 let min = remaining_records.min(remaining_levels);
262 (min, min)
263 }
264 };
265
266 let values_to_read = match self.def_level_decoder.as_mut() {
267 Some(reader) => {
268 let out = def_levels
269 .as_mut()
270 .ok_or_else(|| general_err!("must specify definition levels"))?;
271
272 let (values_read, levels_read) = reader.read_def_levels(out, levels_to_read)?;
273
274 if levels_read != levels_to_read {
275 return Err(general_err!(
276 "insufficient definition levels read from column - expected {levels_to_read}, got {levels_read}"
277 ));
278 }
279
280 values_read
281 }
282 None => levels_to_read,
283 };
284
285 self.check_fixed_len_bytes(values_to_read, levels_to_read == remaining_levels)?;
286
287 let def_levels = def_levels.as_deref();
288 reserve_values(values, values_to_read, levels_to_read, def_levels)?;
289
290 let values_read = self.values_decoder.read(values, values_to_read)?;
291
292 if values_read != values_to_read {
293 return Err(general_err!(
294 "insufficient values read from column - expected: {values_to_read}, got: {values_read}",
295 ));
296 }
297
298 self.num_decoded_values += levels_to_read;
299 total_records_read += records_read;
300 total_levels_read += levels_to_read;
301 total_values_read += values_read;
302 }
303
304 Ok((total_records_read, total_values_read, total_levels_read))
305 }
306
307 pub fn skip_records(&mut self, num_records: usize) -> Result<usize> {
313 let mut remaining_records = num_records;
314 while remaining_records != 0 {
315 if self.num_buffered_values == self.num_decoded_values {
316 let Some(metadata) = self.page_reader.peek_next_page()? else {
317 return Ok(num_records - remaining_records);
318 };
319
320 if metadata.is_dict {
323 if self.page_reader.supports_deferred_dictionary() {
324 self.page_reader.skip_next_page()?;
325 } else {
326 self.read_dictionary_page()?;
327 }
328 continue;
329 }
330
331 let rows = metadata.num_rows.or_else(|| {
334 self.rep_level_decoder
336 .is_none()
337 .then_some(metadata.num_levels)?
338 });
339
340 if let Some(rows) = rows
341 && rows <= remaining_records
342 {
343 self.page_reader.skip_next_page()?;
344 remaining_records -= rows;
345 continue;
346 }
347 if !self.read_new_page()? {
350 return Ok(num_records - remaining_records);
351 }
352 }
353
354 let remaining_levels = self.num_buffered_values - self.num_decoded_values;
358
359 let (records_read, rep_levels_read) = match self.rep_level_decoder.as_mut() {
360 Some(decoder) => {
361 let (mut records_read, levels_read) =
362 decoder.skip_rep_levels(remaining_records, remaining_levels)?;
363
364 if levels_read == remaining_levels && self.has_record_delimiter {
365 check_partial_record_fits(records_read, remaining_records)?;
366 records_read += decoder.flush_partial() as usize;
367 }
368
369 (records_read, levels_read)
370 }
371 None => {
372 let levels = remaining_levels.min(remaining_records);
374 (levels, levels)
375 }
376 };
377
378 self.num_decoded_values += rep_levels_read;
379 remaining_records -= records_read;
380
381 let page_finished = self.num_buffered_values == self.num_decoded_values;
382 if page_finished && self.remaining_fixed_len_bytes.is_none() {
383 continue;
385 }
386
387 let (values_read, def_levels_read) = match self.def_level_decoder.as_mut() {
388 Some(decoder) => decoder.skip_def_levels(rep_levels_read)?,
389 None => (rep_levels_read, rep_levels_read),
390 };
391
392 if rep_levels_read != def_levels_read {
393 return Err(general_err!(
394 "levels mismatch, read {} repetition levels and {} definition levels",
395 rep_levels_read,
396 def_levels_read
397 ));
398 }
399
400 self.check_fixed_len_bytes(values_read, page_finished)?;
401 if page_finished {
402 continue;
403 }
404
405 let values = self.values_decoder.skip_values(values_read)?;
406 if values != values_read {
407 return Err(general_err!(
408 "skipped {} values, expected {}",
409 values,
410 values_read
411 ));
412 }
413 }
414 Ok(num_records - remaining_records)
415 }
416
417 fn check_fixed_len_bytes(&mut self, values: usize, page_finished: bool) -> Result<()> {
420 if let Some((width, remaining)) = self.remaining_fixed_len_bytes.as_mut() {
421 let bytes = values
422 .checked_mul(*width)
423 .ok_or_else(|| general_err!("FIXED_LEN_BYTE_ARRAY payload size overflow"))?;
424 let left = remaining.checked_sub(bytes).ok_or_else(|| {
425 general_err!(
426 "Invalid FIXED_LEN_BYTE_ARRAY PLAIN data page payload length: \
427 need {bytes} value bytes, only {remaining} remain"
428 )
429 })?;
430 if page_finished && left != 0 {
431 return Err(general_err!(
432 "Invalid FIXED_LEN_BYTE_ARRAY PLAIN data page payload length: \
433 {left} trailing value bytes"
434 ));
435 }
436 *remaining = left;
437 if page_finished {
438 self.remaining_fixed_len_bytes = None;
439 }
440 }
441 Ok(())
442 }
443
444 fn read_dictionary_page(&mut self) -> Result<()> {
451 match self.page_reader.get_next_page()? {
452 Some(Page::DictionaryPage {
453 buf,
454 num_values,
455 encoding,
456 is_sorted,
457 }) => self
458 .values_decoder
459 .set_dict(buf, num_values, encoding, is_sorted),
460 _ => Err(ParquetError::General(
461 "Invalid page. Expecting dictionary page".to_string(),
462 )),
463 }
464 }
465
466 fn install_deferred_dictionary(&mut self) -> Result<()> {
469 match self.page_reader.take_deferred_dictionary()? {
470 Some(Page::DictionaryPage {
471 buf,
472 num_values,
473 encoding,
474 is_sorted,
475 }) => self
476 .values_decoder
477 .set_dict(buf, num_values, encoding, is_sorted),
478 Some(_) => Err(ParquetError::General(
479 "Invalid page. Expecting dictionary page".to_string(),
480 )),
481 None => Ok(()),
482 }
483 }
484
485 fn read_new_page(&mut self) -> Result<bool> {
488 loop {
489 match self.page_reader.get_next_page()? {
490 None => return Ok(false),
492 Some(current_page) => {
493 match current_page {
494 Page::DictionaryPage {
496 buf,
497 num_values,
498 encoding,
499 is_sorted,
500 } => {
501 self.values_decoder
502 .set_dict(buf, num_values, encoding, is_sorted)?;
503 }
504 Page::DataPage {
506 buf,
507 num_values,
508 encoding,
509 def_level_encoding,
510 rep_level_encoding,
511 statistics: _,
512 } => {
513 if matches!(
514 encoding,
515 Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY
516 ) {
517 self.install_deferred_dictionary()?;
518 }
519 self.num_buffered_values = num_values as _;
520 self.num_decoded_values = 0;
521
522 let max_rep_level = self.descr.max_rep_level();
523 let max_def_level = self.descr.max_def_level();
524 self.remaining_fixed_len_bytes = None;
525
526 let mut offset = 0;
527
528 if max_rep_level > 0 {
529 let (bytes_read, level_data) = parse_v1_level(
530 max_rep_level,
531 num_values,
532 rep_level_encoding,
533 buf.slice(offset..),
534 )?;
535 offset += bytes_read;
536
537 self.has_record_delimiter =
538 self.page_reader.at_record_boundary()?;
539
540 self.rep_level_decoder
541 .as_mut()
542 .unwrap()
543 .set_data(rep_level_encoding, level_data)?;
544 }
545
546 let definition_levels = if max_def_level > 0 {
547 let (bytes_read, level_data) = parse_v1_level(
548 max_def_level,
549 num_values,
550 def_level_encoding,
551 buf.slice(offset..),
552 )?;
553 offset += bytes_read;
554 Some(level_data)
555 } else {
556 None
557 };
558
559 let data = buf.slice(offset..);
560 let (encoding, data) =
561 if self.descr.physical_type() == Type::FIXED_LEN_BYTE_ARRAY {
562 let width = self.descr.type_length() as usize;
563 let page = prepare_v1_fixed_len_byte_array(
564 data,
565 num_values as usize,
566 width,
567 encoding,
568 definition_levels.as_ref().map(|levels| {
569 (max_def_level, def_level_encoding, levels)
570 }),
571 )?;
572 self.remaining_fixed_len_bytes =
573 page.validate_count.then_some((width, page.data.len()));
574 (page.encoding, page.data)
575 } else {
576 (encoding, data)
577 };
578 if let Some(level_data) = definition_levels {
579 self.def_level_decoder
580 .as_mut()
581 .unwrap()
582 .set_data(def_level_encoding, level_data)?;
583 }
584 self.values_decoder.set_data(
585 encoding,
586 data,
587 num_values as usize,
588 None,
589 )?;
590 return Ok(true);
591 }
592 Page::DataPageV2 {
594 buf,
595 num_values,
596 encoding,
597 num_nulls,
598 num_rows: _,
599 def_levels_byte_len,
600 rep_levels_byte_len,
601 is_compressed: _,
602 statistics: _,
603 } => {
604 if matches!(
605 encoding,
606 Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY
607 ) {
608 self.install_deferred_dictionary()?;
609 }
610 if num_nulls > num_values {
611 return Err(general_err!(
612 "more nulls than values in page, contained {} values and {} nulls",
613 num_values,
614 num_nulls
615 ));
616 }
617
618 self.num_buffered_values = num_values as _;
619 self.num_decoded_values = 0;
620
621 self.remaining_fixed_len_bytes = None;
622
623 if self.descr.max_rep_level() > 0 {
626 self.has_record_delimiter =
630 self.page_reader.at_record_boundary()?;
631
632 self.rep_level_decoder.as_mut().unwrap().set_data(
633 Encoding::RLE,
634 buf.slice(..rep_levels_byte_len as usize),
635 )?;
636 }
637
638 if self.descr.max_def_level() > 0 {
641 self.def_level_decoder.as_mut().unwrap().set_data(
642 Encoding::RLE,
643 buf.slice(
644 rep_levels_byte_len as usize
645 ..(rep_levels_byte_len + def_levels_byte_len) as usize,
646 ),
647 )?;
648 }
649
650 let data =
651 buf.slice((rep_levels_byte_len + def_levels_byte_len) as usize..);
652 let (encoding, data) =
653 if self.descr.physical_type() == Type::FIXED_LEN_BYTE_ARRAY {
654 normalize_fixed_len_byte_array_data(
655 data,
656 (num_values - num_nulls) as usize,
657 self.descr.type_length() as usize,
658 encoding,
659 )?
660 } else {
661 (encoding, data)
662 };
663 self.values_decoder.set_data(
664 encoding,
665 data,
666 num_values as usize,
667 Some((num_values - num_nulls) as usize),
668 )?;
669 return Ok(true);
670 }
671 }
672 }
673 }
674 }
675 }
676
677 #[inline]
681 pub(crate) fn has_next(&mut self) -> Result<bool> {
682 if self.num_buffered_values == 0 || self.num_buffered_values == self.num_decoded_values {
683 if !self.read_new_page()? {
686 Ok(false)
687 } else {
688 Ok(self.num_buffered_values != 0)
689 }
690 } else {
691 Ok(true)
692 }
693 }
694}
695
696fn check_partial_record_fits(records_read: usize, remaining_records: usize) -> Result<()> {
703 if remaining_records <= records_read {
704 return Err(general_err!(
705 "page ended after {records_read} record(s), which is already all of the \
706 {remaining_records} record(s) asked for, so there is no partial record to flush"
707 ));
708 }
709 Ok(())
710}
711
712fn parse_v1_level(
713 max_level: i16,
714 num_buffered_values: u32,
715 encoding: Encoding,
716 buf: Bytes,
717) -> Result<(usize, Bytes)> {
718 match encoding {
719 Encoding::RLE => {
720 let i32_size = std::mem::size_of::<i32>();
721 if i32_size <= buf.len() {
722 let data_size = read_num_bytes::<i32>(i32_size, buf.as_ref()) as usize;
723 let end = i32_size
724 .checked_add(data_size)
725 .ok_or(general_err!("invalid level length"))?;
726 if end <= buf.len() {
727 return Ok((end, buf.slice(i32_size..end)));
728 }
729 }
730 Err(general_err!("not enough data to read levels"))
731 }
732 #[expect(deprecated)]
733 Encoding::BIT_PACKED => {
734 let bit_width = num_required_bits(max_level as u64);
735 let num_bytes = ceil(num_buffered_values as usize * bit_width as usize, 8);
736 Ok((num_bytes, buf.slice(..num_bytes)))
737 }
738 _ => Err(general_err!("invalid level encoding: {}", encoding)),
739 }
740}
741
742#[cfg(test)]
743mod tests {
744 use super::*;
745
746 use rand::distr::uniform::SampleUniform;
747 use std::{collections::VecDeque, sync::Arc};
748
749 use crate::basic::Type as PhysicalType;
750 use crate::schema::types::{ColumnDescriptor, ColumnPath, Type as SchemaType};
751 use crate::util::test_common::page_util::InMemoryPageReader;
752 use crate::util::test_common::rand_gen::make_pages;
753
754 pub(super) fn nullable_fixed_page(def: &[i16], encoding: Encoding, payload: &[u8]) -> Page {
755 use crate::encodings::levels::LevelEncoder;
756 let mut levels = LevelEncoder::v1_streaming(1);
757 levels.put_with_observer(def, |_, _| {});
758 let mut buf = levels.consume();
759 buf.extend_from_slice(payload);
760 Page::DataPage {
761 buf: buf.into(),
762 num_values: def.len() as u32,
763 encoding,
764 def_level_encoding: Encoding::RLE,
765 rep_level_encoding: Encoding::RLE,
766 statistics: None,
767 }
768 }
769
770 pub(super) fn nullable_fixed_reader(
771 pages: Vec<Page>,
772 ) -> ColumnReaderImpl<FixedLenByteArrayType> {
773 let descr = Arc::new(ColumnDescriptor::new(
774 Arc::new(
775 SchemaType::primitive_type_builder("fixed", PhysicalType::FIXED_LEN_BYTE_ARRAY)
776 .with_length(4)
777 .build()
778 .unwrap(),
779 ),
780 1,
781 0,
782 ColumnPath::from("fixed"),
783 ));
784 ColumnReaderImpl::<FixedLenByteArrayType>::new(
785 descr,
786 Box::new(InMemoryPageReader::new(pages)),
787 )
788 }
789
790 pub(super) fn check_incremental_fixed_len_byte_array_pages(
792 encoding: Encoding,
793 payload: impl Fn(usize) -> Vec<u8>,
794 ) {
795 for count in [1, 2, 3] {
796 for finish_by_skipping in [false, true] {
797 let mut reader = nullable_fixed_reader(vec![nullable_fixed_page(
798 &[1, 0, 1, 0],
799 encoding,
800 &payload(count),
801 )]);
802 let mut values = vec![];
803 assert_eq!(
804 reader
805 .read_records(1, Some(&mut vec![]), None, &mut values)
806 .unwrap(),
807 (1, 1, 1)
808 );
809 assert_eq!(values[0].data(), b"abcd");
812 assert_eq!(reader.remaining_fixed_len_bytes, Some((4, (count - 1) * 4)));
813 assert_eq!(reader.skip_records(1).unwrap(), 1);
814 let result = if finish_by_skipping {
815 reader.skip_records(2).map(|_| ())
816 } else {
817 reader
818 .read_records(2, Some(&mut vec![]), None, &mut values)
819 .map(|_| ())
820 };
821 if count == 2 {
822 result.unwrap();
823 assert!(reader.remaining_fixed_len_bytes.is_none());
824 } else {
825 assert!(result.unwrap_err().to_string().contains("payload length"));
826 }
827 }
828 }
829 }
830
831 #[test]
832 fn fixed_v1_incremental_count_validation() {
833 check_incremental_fixed_len_byte_array_pages(Encoding::PLAIN, |n| b"abcd".repeat(n));
834
835 for last in [0, 1] {
838 for payload in [b"".as_slice(), b"abcd".as_slice()] {
839 let mut reader = nullable_fixed_reader(vec![nullable_fixed_page(
840 &[0, 0, 0, last],
841 Encoding::PLAIN,
842 payload,
843 )]);
844 let mut values = vec![];
845 assert_eq!(
846 reader
847 .read_records(1, Some(&mut vec![]), None, &mut values)
848 .unwrap(),
849 (1, 0, 1)
850 );
851 assert_eq!(
852 reader.skip_records(3).is_ok(),
853 payload.len() == last as usize * 4
854 );
855 assert!(values.is_empty());
856 }
857 }
858 }
859
860 #[test]
861 fn fixed_v1_canonical_prefix_shaped_values() {
862 let raw = 4_u32.to_le_bytes().repeat(10);
864 let mut reader =
865 nullable_fixed_reader(vec![nullable_fixed_page(&[1; 10], Encoding::PLAIN, &raw)]);
866 let mut values = vec![];
867 assert_eq!(
868 reader
869 .read_records(1, Some(&mut vec![]), None, &mut values)
870 .unwrap(),
871 (1, 1, 1)
872 );
873 reader
874 .read_records(usize::MAX, Some(&mut vec![]), None, &mut values)
875 .unwrap();
876 assert_eq!(values.len(), 10);
877 assert!(values.iter().all(|v| v.data() == 4_u32.to_le_bytes()));
878 }
879
880 #[test]
881 fn fixed_v1_incremental_page_transitions_and_short_levels() {
882 let first = nullable_fixed_page(&[1, 0, 1, 0], Encoding::PLAIN, b"abcdabcd");
883 let second = nullable_fixed_page(&[1, 1], Encoding::PLAIN, b"efghefgh");
884 let mut reader = nullable_fixed_reader(vec![first.clone(), second]);
885 let mut values = vec![];
886 reader
887 .read_records(1, Some(&mut vec![]), None, &mut values)
888 .unwrap();
889 assert_eq!(reader.skip_records(3).unwrap(), 3);
890 assert!(reader.remaining_fixed_len_bytes.is_none());
891 assert_eq!(
892 reader
893 .read_records(usize::MAX, Some(&mut vec![]), None, &mut values)
894 .unwrap(),
895 (2, 2, 2)
896 );
897 assert_eq!(
898 values.iter().map(|v| v.data()).collect::<Vec<_>>(),
899 [b"abcd", b"efgh", b"efgh"]
900 );
901
902 let mut short_levels = first;
903 if let Page::DataPage { num_values, .. } = &mut short_levels {
904 *num_values += 8; }
906 let mut reader = nullable_fixed_reader(vec![short_levels]);
907 reader
908 .read_records(1, Some(&mut vec![]), None, &mut vec![])
909 .unwrap();
910 let err = reader
911 .read_records(usize::MAX, Some(&mut vec![]), None, &mut vec![])
912 .unwrap_err();
913 assert!(err.to_string().contains("definition levels"));
914 }
915
916 #[test]
917 fn fixed_plain_payload_length_validation() {
918 check_fixed_len_byte_array_pages(
919 |raw| vec![(Encoding::PLAIN, raw.to_vec())],
920 |raw| {
921 let mut extra = raw.to_vec();
922 extra.push(0);
923 let mut malformed = vec![extra];
924 if !raw.is_empty() {
925 malformed.push(raw[..raw.len() - 1].to_vec());
926 }
927 malformed
928 },
929 );
930 }
931
932 #[expect(deprecated, reason = "Cover legacy BIT_PACKED definition levels")]
935 pub(super) fn check_fixed_len_byte_array_pages(
936 payloads: impl Fn(&[u8]) -> Vec<(Encoding, Vec<u8>)>,
937 malformed_plain: impl Fn(&[u8]) -> Vec<Vec<u8>>,
938 ) {
939 use crate::encodings::levels::LevelEncoder;
940 use crate::encodings::rle::RleEncoder;
941 use crate::util::bit_util::BitWriter;
942
943 fn levels(data: &[i16], max: i16, v2: bool, encoding: Encoding) -> Vec<u8> {
944 if max == 0 {
945 return vec![];
946 }
947 if encoding == Encoding::BIT_PACKED {
948 let mut encoder = BitWriter::new(data.len());
949 for &value in data {
950 encoder.put_value(value as u64, num_required_bits(max as u64) as usize);
951 }
952 return encoder.consume();
953 }
954 let mut encoder = if v2 {
955 LevelEncoder::v2_streaming(max)
956 } else {
957 LevelEncoder::v1_streaming(max)
958 };
959 encoder.put_with_observer(data, |_, _| {});
960 encoder.consume()
961 }
962
963 for (max_def, def, rep) in [
964 (0, vec![], vec![]),
965 (1, vec![1, 0, 1, 0, 1], vec![]),
966 (1, vec![0; 5], vec![]),
967 (3, vec![0, 1, 2, 3, 3, 1, 3], vec![0, 0, 0, 1, 1, 0, 1]),
969 (
971 1,
972 (0..2051).map(|i| i16::from(i % 3 == 0)).collect(),
973 vec![],
974 ),
975 ] {
976 let max_rep = i16::from(!rep.is_empty());
977 let num_levels = if max_def == 0 { 3 } else { def.len() };
978 let num_values = if max_def == 0 {
979 num_levels
980 } else {
981 def.iter().filter(|&&v| v == max_def).count()
982 };
983 let num_rows = if max_rep == 0 {
984 num_levels
985 } else {
986 rep.iter().filter(|&&v| v == 0).count()
987 };
988 let desc = Arc::new(ColumnDescriptor::new(
989 Arc::new(
990 SchemaType::primitive_type_builder("fixed", PhysicalType::FIXED_LEN_BYTE_ARRAY)
991 .with_length(4)
992 .build()
993 .unwrap(),
994 ),
995 max_def,
996 max_rep,
997 ColumnPath::from("fixed"),
998 ));
999 let raw: Vec<u8> = b"\x04\0\0\0\xff\xff\xff\xffabcd"
1000 .as_chunks::<4>()
1001 .0
1002 .iter()
1003 .cycle()
1004 .take(num_values)
1005 .flatten()
1006 .copied()
1007 .collect();
1008 let malformed = malformed_plain(&raw);
1009
1010 for (v2, level_encoding) in [
1011 (false, Encoding::RLE),
1012 (false, Encoding::BIT_PACKED),
1013 (true, Encoding::RLE),
1014 ] {
1015 let rep_data = levels(&rep, max_rep, v2, level_encoding);
1016 let def_data = levels(&def, max_def, v2, level_encoding);
1017 let make_page = |encoding, payload: &[u8]| {
1018 let mut buf = rep_data.clone();
1019 buf.extend_from_slice(&def_data);
1020 buf.extend_from_slice(payload);
1021 if v2 {
1022 Page::DataPageV2 {
1023 buf: Bytes::from(buf),
1024 num_values: num_levels as u32,
1025 encoding,
1026 num_nulls: (num_levels - num_values) as u32,
1027 num_rows: num_rows as u32,
1028 def_levels_byte_len: def_data.len() as u32,
1029 rep_levels_byte_len: rep_data.len() as u32,
1030 is_compressed: false,
1031 statistics: None,
1032 }
1033 } else {
1034 Page::DataPage {
1035 buf: Bytes::from(buf),
1036 num_values: num_levels as u32,
1037 encoding,
1038 def_level_encoding: level_encoding,
1039 rep_level_encoding: level_encoding,
1040 statistics: None,
1041 }
1042 }
1043 };
1044 let make_reader = |pages| {
1045 ColumnReaderImpl::<FixedLenByteArrayType>::new(
1046 desc.clone(),
1047 Box::new(InMemoryPageReader::new(pages)),
1048 )
1049 };
1050 for (encoding, payload) in payloads(&raw) {
1051 let payload = payload.as_slice();
1052 let mut reader = make_reader(vec![make_page(encoding, payload)]);
1053 let (mut values, mut actual_def, mut actual_rep) = (vec![], vec![], vec![]);
1054 assert_eq!(
1055 reader
1056 .read_records(
1057 usize::MAX,
1058 Some(&mut actual_def),
1059 Some(&mut actual_rep),
1060 &mut values
1061 )
1062 .unwrap(),
1063 (num_rows, num_values, num_levels)
1064 );
1065 assert_eq!(actual_def, def);
1066 assert_eq!(actual_rep, rep);
1067 assert_eq!(
1068 values.iter().map(|v| v.data()).collect::<Vec<_>>(),
1069 raw.as_chunks::<4>().0.iter().collect::<Vec<_>>()
1070 );
1071
1072 let first_record_levels = if max_rep == 0 {
1073 1
1074 } else {
1075 rep.iter().skip(1).position(|&v| v == 0).unwrap() + 1
1076 };
1077 let first_record_values = if max_def == 0 {
1078 1
1079 } else {
1080 def[..first_record_levels]
1081 .iter()
1082 .filter(|&&v| v == max_def)
1083 .count()
1084 };
1085 let mut reader = make_reader(vec![make_page(encoding, payload)]);
1086 assert_eq!(reader.skip_records(1).unwrap(), 1);
1087 let mut values = vec![];
1088 reader
1089 .read_records(
1090 usize::MAX,
1091 Some(&mut vec![]),
1092 Some(&mut vec![]),
1093 &mut values,
1094 )
1095 .unwrap();
1096 assert_eq!(
1097 values.iter().map(|v| v.data()).collect::<Vec<_>>(),
1098 raw[first_record_values * 4..]
1099 .as_chunks::<4>()
1100 .0
1101 .iter()
1102 .collect::<Vec<_>>()
1103 );
1104 }
1105
1106 if !v2 && max_def != 0 && level_encoding == Encoding::RLE {
1107 let mut page = make_page(Encoding::PLAIN, &raw);
1108 if let Page::DataPage { num_values, .. } = &mut page {
1109 *num_values += 8;
1111 }
1112 let err = make_reader(vec![page])
1115 .read_records(
1116 usize::MAX,
1117 Some(&mut vec![]),
1118 Some(&mut vec![]),
1119 &mut vec![],
1120 )
1121 .unwrap_err()
1122 .to_string();
1123 assert!(
1124 err.contains("definition levels")
1125 || (max_rep != 0 && err.contains("repetition levels")),
1126 "{err}"
1127 );
1128 }
1129
1130 for payload in &malformed {
1131 for skip in [false, true] {
1134 let mut reader = make_reader(vec![make_page(Encoding::PLAIN, payload)]);
1135 let mut values = vec![];
1136 let mut result = if skip {
1137 reader.skip_records(1).map(|_| ())
1138 } else {
1139 reader
1140 .read_records(1, Some(&mut vec![]), Some(&mut vec![]), &mut values)
1141 .map(|_| ())
1142 };
1143 let provisional_values = values.len();
1144 if result.is_err() {
1145 assert!(values.is_empty());
1146 } else {
1147 assert!(!v2 && max_def != 0);
1148 result = if skip {
1149 reader.skip_records(usize::MAX).map(|_| ())
1150 } else {
1151 reader
1152 .read_records(
1153 usize::MAX,
1154 Some(&mut vec![]),
1155 Some(&mut vec![]),
1156 &mut values,
1157 )
1158 .map(|_| ())
1159 };
1160 }
1161 let err = result.unwrap_err().to_string();
1162 assert!(
1163 err.contains(
1164 "Invalid FIXED_LEN_BYTE_ARRAY PLAIN data page payload length"
1165 ),
1166 "v2={v2} {level_encoding} max_def={max_def} skip={skip}: {err}"
1167 );
1168 assert_eq!(values.len(), provisional_values);
1169 }
1170
1171 let mut indices = RleEncoder::new(1, 64);
1174 for _ in 0..num_values {
1175 indices.put(0);
1176 }
1177 let mut encoded = vec![1];
1178 encoded.extend(indices.consume());
1179 let mut reader = make_reader(vec![
1180 Page::DictionaryPage {
1181 buf: Bytes::from_static(b"abcd"),
1182 num_values: 1,
1183 encoding: Encoding::PLAIN,
1184 is_sorted: false,
1185 },
1186 make_page(Encoding::RLE_DICTIONARY, &encoded),
1187 make_page(Encoding::PLAIN, payload),
1188 ]);
1189 let mut values = vec![];
1190 let err = reader
1191 .read_records(
1192 usize::MAX,
1193 Some(&mut vec![]),
1194 Some(&mut vec![]),
1195 &mut values,
1196 )
1197 .unwrap_err()
1198 .to_string();
1199 assert!(err.contains("PLAIN data page payload length"), "{err}");
1200 assert_eq!(values.len(), num_values);
1201 assert!(values.iter().all(|v| v.data() == b"abcd"));
1202 }
1203 }
1204 }
1205 }
1206
1207 #[test]
1208 fn test_parse_v1_level_invalid_length() {
1209 let buf = Bytes::from(vec![10, 0, 0, 0]);
1211 let err = parse_v1_level(1, 100, Encoding::RLE, buf).unwrap_err();
1212 assert_eq!(
1213 err.to_string(),
1214 "Parquet error: not enough data to read levels"
1215 );
1216
1217 let buf = Bytes::from(vec![4, 0, 0]);
1219 let err = parse_v1_level(1, 100, Encoding::RLE, buf).unwrap_err();
1220 assert_eq!(
1221 err.to_string(),
1222 "Parquet error: not enough data to read levels"
1223 );
1224 }
1225
1226 const NUM_LEVELS: usize = 128;
1227 const NUM_PAGES: usize = 2;
1228 const MAX_DEF_LEVEL: i16 = 5;
1229 const MAX_REP_LEVEL: i16 = 5;
1230
1231 macro_rules! test {
1233 ($test_func:ident, i32, $func:ident, $def_level:expr, $rep_level:expr,
1235 $num_pages:expr, $num_levels:expr, $batch_size:expr, $min:expr, $max:expr) => {
1236 test_internal!(
1237 $test_func,
1238 Int32Type,
1239 get_test_int32_type,
1240 $func,
1241 $def_level,
1242 $rep_level,
1243 $num_pages,
1244 $num_levels,
1245 $batch_size,
1246 $min,
1247 $max
1248 );
1249 };
1250 ($test_func:ident, i64, $func:ident, $def_level:expr, $rep_level:expr,
1252 $num_pages:expr, $num_levels:expr, $batch_size:expr, $min:expr, $max:expr) => {
1253 test_internal!(
1254 $test_func,
1255 Int64Type,
1256 get_test_int64_type,
1257 $func,
1258 $def_level,
1259 $rep_level,
1260 $num_pages,
1261 $num_levels,
1262 $batch_size,
1263 $min,
1264 $max
1265 );
1266 };
1267 }
1268
1269 macro_rules! test_internal {
1270 ($test_func:ident, $ty:ident, $pty:ident, $func:ident, $def_level:expr,
1271 $rep_level:expr, $num_pages:expr, $num_levels:expr, $batch_size:expr,
1272 $min:expr, $max:expr) => {
1273 #[test]
1274 fn $test_func() {
1275 let desc = Arc::new(ColumnDescriptor::new(
1276 Arc::new($pty()),
1277 $def_level,
1278 $rep_level,
1279 ColumnPath::new(Vec::new()),
1280 ));
1281 let mut tester = ColumnReaderTester::<$ty>::new();
1282 tester.$func(desc, $num_pages, $num_levels, $batch_size, $min, $max);
1283 }
1284 };
1285 }
1286
1287 test!(
1288 test_read_plain_v1_int32,
1289 i32,
1290 plain_v1,
1291 MAX_DEF_LEVEL,
1292 MAX_REP_LEVEL,
1293 NUM_PAGES,
1294 NUM_LEVELS,
1295 16,
1296 i32::MIN,
1297 i32::MAX
1298 );
1299 test!(
1300 test_read_plain_v2_int32,
1301 i32,
1302 plain_v2,
1303 MAX_DEF_LEVEL,
1304 MAX_REP_LEVEL,
1305 NUM_PAGES,
1306 NUM_LEVELS,
1307 16,
1308 i32::MIN,
1309 i32::MAX
1310 );
1311
1312 test!(
1313 test_read_plain_v1_int32_uneven,
1314 i32,
1315 plain_v1,
1316 MAX_DEF_LEVEL,
1317 MAX_REP_LEVEL,
1318 NUM_PAGES,
1319 NUM_LEVELS,
1320 17,
1321 i32::MIN,
1322 i32::MAX
1323 );
1324 test!(
1325 test_read_plain_v2_int32_uneven,
1326 i32,
1327 plain_v2,
1328 MAX_DEF_LEVEL,
1329 MAX_REP_LEVEL,
1330 NUM_PAGES,
1331 NUM_LEVELS,
1332 17,
1333 i32::MIN,
1334 i32::MAX
1335 );
1336
1337 test!(
1338 test_read_plain_v1_int32_multi_page,
1339 i32,
1340 plain_v1,
1341 MAX_DEF_LEVEL,
1342 MAX_REP_LEVEL,
1343 NUM_PAGES,
1344 NUM_LEVELS,
1345 512,
1346 i32::MIN,
1347 i32::MAX
1348 );
1349 test!(
1350 test_read_plain_v2_int32_multi_page,
1351 i32,
1352 plain_v2,
1353 MAX_DEF_LEVEL,
1354 MAX_REP_LEVEL,
1355 NUM_PAGES,
1356 NUM_LEVELS,
1357 512,
1358 i32::MIN,
1359 i32::MAX
1360 );
1361
1362 test!(
1364 test_read_plain_v1_int32_required_non_repeated,
1365 i32,
1366 plain_v1,
1367 0,
1368 0,
1369 NUM_PAGES,
1370 NUM_LEVELS,
1371 16,
1372 i32::MIN,
1373 i32::MAX
1374 );
1375 test!(
1376 test_read_plain_v2_int32_required_non_repeated,
1377 i32,
1378 plain_v2,
1379 0,
1380 0,
1381 NUM_PAGES,
1382 NUM_LEVELS,
1383 16,
1384 i32::MIN,
1385 i32::MAX
1386 );
1387
1388 test!(
1389 test_read_plain_v1_int64,
1390 i64,
1391 plain_v1,
1392 1,
1393 1,
1394 NUM_PAGES,
1395 NUM_LEVELS,
1396 16,
1397 i64::MIN,
1398 i64::MAX
1399 );
1400 test!(
1401 test_read_plain_v2_int64,
1402 i64,
1403 plain_v2,
1404 1,
1405 1,
1406 NUM_PAGES,
1407 NUM_LEVELS,
1408 16,
1409 i64::MIN,
1410 i64::MAX
1411 );
1412
1413 test!(
1414 test_read_plain_v1_int64_uneven,
1415 i64,
1416 plain_v1,
1417 1,
1418 1,
1419 NUM_PAGES,
1420 NUM_LEVELS,
1421 17,
1422 i64::MIN,
1423 i64::MAX
1424 );
1425 test!(
1426 test_read_plain_v2_int64_uneven,
1427 i64,
1428 plain_v2,
1429 1,
1430 1,
1431 NUM_PAGES,
1432 NUM_LEVELS,
1433 17,
1434 i64::MIN,
1435 i64::MAX
1436 );
1437
1438 test!(
1439 test_read_plain_v1_int64_multi_page,
1440 i64,
1441 plain_v1,
1442 1,
1443 1,
1444 NUM_PAGES,
1445 NUM_LEVELS,
1446 512,
1447 i64::MIN,
1448 i64::MAX
1449 );
1450 test!(
1451 test_read_plain_v2_int64_multi_page,
1452 i64,
1453 plain_v2,
1454 1,
1455 1,
1456 NUM_PAGES,
1457 NUM_LEVELS,
1458 512,
1459 i64::MIN,
1460 i64::MAX
1461 );
1462
1463 test!(
1465 test_read_plain_v1_int64_required_non_repeated,
1466 i64,
1467 plain_v1,
1468 0,
1469 0,
1470 NUM_PAGES,
1471 NUM_LEVELS,
1472 16,
1473 i64::MIN,
1474 i64::MAX
1475 );
1476 test!(
1477 test_read_plain_v2_int64_required_non_repeated,
1478 i64,
1479 plain_v2,
1480 0,
1481 0,
1482 NUM_PAGES,
1483 NUM_LEVELS,
1484 16,
1485 i64::MIN,
1486 i64::MAX
1487 );
1488
1489 test!(
1490 test_read_dict_v1_int32_small,
1491 i32,
1492 dict_v1,
1493 MAX_DEF_LEVEL,
1494 MAX_REP_LEVEL,
1495 2,
1496 2,
1497 16,
1498 0,
1499 3
1500 );
1501 test!(
1502 test_read_dict_v2_int32_small,
1503 i32,
1504 dict_v2,
1505 MAX_DEF_LEVEL,
1506 MAX_REP_LEVEL,
1507 2,
1508 2,
1509 16,
1510 0,
1511 3
1512 );
1513
1514 test!(
1515 test_read_dict_v1_int32,
1516 i32,
1517 dict_v1,
1518 MAX_DEF_LEVEL,
1519 MAX_REP_LEVEL,
1520 NUM_PAGES,
1521 NUM_LEVELS,
1522 16,
1523 0,
1524 3
1525 );
1526 test!(
1527 test_read_dict_v2_int32,
1528 i32,
1529 dict_v2,
1530 MAX_DEF_LEVEL,
1531 MAX_REP_LEVEL,
1532 NUM_PAGES,
1533 NUM_LEVELS,
1534 16,
1535 0,
1536 3
1537 );
1538
1539 test!(
1540 test_read_dict_v1_int32_uneven,
1541 i32,
1542 dict_v1,
1543 MAX_DEF_LEVEL,
1544 MAX_REP_LEVEL,
1545 NUM_PAGES,
1546 NUM_LEVELS,
1547 17,
1548 0,
1549 3
1550 );
1551 test!(
1552 test_read_dict_v2_int32_uneven,
1553 i32,
1554 dict_v2,
1555 MAX_DEF_LEVEL,
1556 MAX_REP_LEVEL,
1557 NUM_PAGES,
1558 NUM_LEVELS,
1559 17,
1560 0,
1561 3
1562 );
1563
1564 test!(
1565 test_read_dict_v1_int32_multi_page,
1566 i32,
1567 dict_v1,
1568 MAX_DEF_LEVEL,
1569 MAX_REP_LEVEL,
1570 NUM_PAGES,
1571 NUM_LEVELS,
1572 512,
1573 0,
1574 3
1575 );
1576 test!(
1577 test_read_dict_v2_int32_multi_page,
1578 i32,
1579 dict_v2,
1580 MAX_DEF_LEVEL,
1581 MAX_REP_LEVEL,
1582 NUM_PAGES,
1583 NUM_LEVELS,
1584 512,
1585 0,
1586 3
1587 );
1588
1589 test!(
1590 test_read_dict_v1_int64,
1591 i64,
1592 dict_v1,
1593 MAX_DEF_LEVEL,
1594 MAX_REP_LEVEL,
1595 NUM_PAGES,
1596 NUM_LEVELS,
1597 16,
1598 0,
1599 3
1600 );
1601 test!(
1602 test_read_dict_v2_int64,
1603 i64,
1604 dict_v2,
1605 MAX_DEF_LEVEL,
1606 MAX_REP_LEVEL,
1607 NUM_PAGES,
1608 NUM_LEVELS,
1609 16,
1610 0,
1611 3
1612 );
1613
1614 #[test]
1615 fn test_read_batch_values_only() {
1616 test_read_batch_int32(16, 0, 0);
1617 }
1618
1619 #[test]
1620 fn test_read_batch_values_def_levels() {
1621 test_read_batch_int32(16, MAX_DEF_LEVEL, 0);
1622 }
1623
1624 #[test]
1625 fn test_read_batch_values_rep_levels() {
1626 test_read_batch_int32(16, 0, MAX_REP_LEVEL);
1627 }
1628
1629 #[test]
1630 fn test_read_batch_values_def_rep_levels() {
1631 test_read_batch_int32(128, MAX_DEF_LEVEL, MAX_REP_LEVEL);
1632 }
1633
1634 #[test]
1635 fn test_read_batch_adjust_after_buffering_page() {
1636 let primitive_type = get_test_int32_type();
1643 let desc = Arc::new(ColumnDescriptor::new(
1644 Arc::new(primitive_type),
1645 1,
1646 1,
1647 ColumnPath::new(Vec::new()),
1648 ));
1649
1650 let num_pages = 2;
1651 let num_levels = 4;
1652 let batch_size = 5;
1653
1654 let mut tester = ColumnReaderTester::<Int32Type>::new();
1655 tester.test_read_batch(
1656 desc,
1657 Encoding::RLE_DICTIONARY,
1658 num_pages,
1659 num_levels,
1660 batch_size,
1661 i32::MIN,
1662 i32::MAX,
1663 false,
1664 );
1665 }
1666
1667 fn get_test_int32_type() -> SchemaType {
1713 SchemaType::primitive_type_builder("a", PhysicalType::INT32)
1714 .with_repetition(Repetition::REQUIRED)
1715 .with_converted_type(ConvertedType::INT_32)
1716 .with_length(-1)
1717 .build()
1718 .expect("build() should be OK")
1719 }
1720
1721 fn get_test_int64_type() -> SchemaType {
1723 SchemaType::primitive_type_builder("a", PhysicalType::INT64)
1724 .with_repetition(Repetition::REQUIRED)
1725 .with_converted_type(ConvertedType::INT_64)
1726 .with_length(-1)
1727 .build()
1728 .expect("build() should be OK")
1729 }
1730
1731 fn test_read_batch_int32(batch_size: usize, max_def_level: i16, max_rep_level: i16) {
1736 let primitive_type = get_test_int32_type();
1737
1738 let desc = Arc::new(ColumnDescriptor::new(
1739 Arc::new(primitive_type),
1740 max_def_level,
1741 max_rep_level,
1742 ColumnPath::new(Vec::new()),
1743 ));
1744
1745 let mut tester = ColumnReaderTester::<Int32Type>::new();
1746 tester.test_read_batch(
1747 desc,
1748 Encoding::RLE_DICTIONARY,
1749 NUM_PAGES,
1750 NUM_LEVELS,
1751 batch_size,
1752 i32::MIN,
1753 i32::MAX,
1754 false,
1755 );
1756 }
1757
1758 struct ColumnReaderTester<T: DataType>
1759 where
1760 T::T: PartialOrd + SampleUniform + Copy,
1761 {
1762 rep_levels: Vec<i16>,
1763 def_levels: Vec<i16>,
1764 values: Vec<T::T>,
1765 }
1766
1767 impl<T: DataType> ColumnReaderTester<T>
1768 where
1769 T::T: PartialOrd + SampleUniform + Copy,
1770 {
1771 pub fn new() -> Self {
1772 Self {
1773 rep_levels: Vec::new(),
1774 def_levels: Vec::new(),
1775 values: Vec::new(),
1776 }
1777 }
1778
1779 fn plain_v1(
1781 &mut self,
1782 desc: ColumnDescPtr,
1783 num_pages: usize,
1784 num_levels: usize,
1785 batch_size: usize,
1786 min: T::T,
1787 max: T::T,
1788 ) {
1789 self.test_read_batch_general(
1790 desc,
1791 Encoding::PLAIN,
1792 num_pages,
1793 num_levels,
1794 batch_size,
1795 min,
1796 max,
1797 false,
1798 );
1799 }
1800
1801 fn plain_v2(
1803 &mut self,
1804 desc: ColumnDescPtr,
1805 num_pages: usize,
1806 num_levels: usize,
1807 batch_size: usize,
1808 min: T::T,
1809 max: T::T,
1810 ) {
1811 self.test_read_batch_general(
1812 desc,
1813 Encoding::PLAIN,
1814 num_pages,
1815 num_levels,
1816 batch_size,
1817 min,
1818 max,
1819 true,
1820 );
1821 }
1822
1823 fn dict_v1(
1825 &mut self,
1826 desc: ColumnDescPtr,
1827 num_pages: usize,
1828 num_levels: usize,
1829 batch_size: usize,
1830 min: T::T,
1831 max: T::T,
1832 ) {
1833 self.test_read_batch_general(
1834 desc,
1835 Encoding::RLE_DICTIONARY,
1836 num_pages,
1837 num_levels,
1838 batch_size,
1839 min,
1840 max,
1841 false,
1842 );
1843 }
1844
1845 fn dict_v2(
1847 &mut self,
1848 desc: ColumnDescPtr,
1849 num_pages: usize,
1850 num_levels: usize,
1851 batch_size: usize,
1852 min: T::T,
1853 max: T::T,
1854 ) {
1855 self.test_read_batch_general(
1856 desc,
1857 Encoding::RLE_DICTIONARY,
1858 num_pages,
1859 num_levels,
1860 batch_size,
1861 min,
1862 max,
1863 true,
1864 );
1865 }
1866
1867 #[expect(clippy::too_many_arguments)]
1870 fn test_read_batch_general(
1871 &mut self,
1872 desc: ColumnDescPtr,
1873 encoding: Encoding,
1874 num_pages: usize,
1875 num_levels: usize,
1876 batch_size: usize,
1877 min: T::T,
1878 max: T::T,
1879 use_v2: bool,
1880 ) {
1881 self.test_read_batch(
1882 desc, encoding, num_pages, num_levels, batch_size, min, max, use_v2,
1883 );
1884 }
1885
1886 #[expect(clippy::too_many_arguments)]
1889 fn test_read_batch(
1890 &mut self,
1891 desc: ColumnDescPtr,
1892 encoding: Encoding,
1893 num_pages: usize,
1894 num_levels: usize,
1895 batch_size: usize,
1896 min: T::T,
1897 max: T::T,
1898 use_v2: bool,
1899 ) {
1900 let mut pages = VecDeque::new();
1901 make_pages::<T>(
1902 desc.clone(),
1903 encoding,
1904 num_pages,
1905 num_levels,
1906 min,
1907 max,
1908 &mut self.def_levels,
1909 &mut self.rep_levels,
1910 &mut self.values,
1911 &mut pages,
1912 use_v2,
1913 );
1914 let max_def_level = desc.max_def_level();
1915 let max_rep_level = desc.max_rep_level();
1916 let page_reader = InMemoryPageReader::new(pages);
1917 let column_reader = get_column_reader(desc, Box::new(page_reader));
1918 let mut typed_column_reader = get_typed_column_reader::<T>(column_reader);
1919
1920 let mut values = Vec::new();
1921 let mut def_levels = Vec::new();
1922 let mut rep_levels = Vec::new();
1923
1924 let mut curr_values_read = 0;
1925 let mut curr_levels_read = 0;
1926 loop {
1927 let (_, values_read, levels_read) = typed_column_reader
1928 .read_records(
1929 batch_size,
1930 Some(&mut def_levels),
1931 Some(&mut rep_levels),
1932 &mut values,
1933 )
1934 .expect("read_batch() should be OK");
1935
1936 curr_values_read += values_read;
1937 curr_levels_read += levels_read;
1938
1939 if values_read == 0 && levels_read == 0 {
1940 break;
1941 }
1942 }
1943
1944 assert_eq!(values, self.values, "values content doesn't match");
1945
1946 if max_def_level > 0 {
1947 assert_eq!(
1948 def_levels, self.def_levels,
1949 "definition levels content doesn't match"
1950 );
1951 }
1952
1953 if max_rep_level > 0 {
1954 assert_eq!(
1955 rep_levels, self.rep_levels,
1956 "repetition levels content doesn't match"
1957 );
1958 }
1959
1960 assert!(
1961 curr_levels_read >= curr_values_read,
1962 "expected levels read to be greater than values read"
1963 );
1964 }
1965 }
1966
1967 #[test]
1985 fn test_skip_records_v2_page_skip_accounts_for_partial() {
1986 use crate::encodings::levels::LevelEncoder;
1987
1988 let max_rep_level: i16 = 1;
1989 let max_def_level: i16 = 1;
1990
1991 let primitive_type = SchemaType::primitive_type_builder("element", PhysicalType::INT32)
1993 .with_repetition(Repetition::REQUIRED)
1994 .build()
1995 .unwrap();
1996 let desc = Arc::new(ColumnDescriptor::new(
1997 Arc::new(primitive_type),
1998 max_def_level,
1999 max_rep_level,
2000 ColumnPath::new(vec!["list".to_string(), "element".to_string()]),
2001 ));
2002
2003 let make_v2_page =
2005 |rep_levels: &[i16], def_levels: &[i16], values: &[i32], num_rows: u32| -> Page {
2006 let mut rep_enc = LevelEncoder::v2_streaming(max_rep_level);
2007 rep_enc.put_with_observer(rep_levels, |_, _| {});
2008 let rep_bytes = rep_enc.consume();
2009
2010 let mut def_enc = LevelEncoder::v2_streaming(max_def_level);
2011 def_enc.put_with_observer(def_levels, |_, _| {});
2012 let def_bytes = def_enc.consume();
2013
2014 let val_bytes: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
2015
2016 let mut buf = Vec::new();
2017 buf.extend_from_slice(&rep_bytes);
2018 buf.extend_from_slice(&def_bytes);
2019 buf.extend_from_slice(&val_bytes);
2020
2021 Page::DataPageV2 {
2022 buf: Bytes::from(buf),
2023 num_values: rep_levels.len() as u32,
2024 encoding: Encoding::PLAIN,
2025 num_nulls: 0,
2026 num_rows,
2027 def_levels_byte_len: def_bytes.len() as u32,
2028 rep_levels_byte_len: rep_bytes.len() as u32,
2029 is_compressed: false,
2030 statistics: None,
2031 }
2032 };
2033
2034 let page1 = make_v2_page(&[0, 1, 0, 1], &[1, 1, 1, 1], &[10, 20, 30, 40], 2);
2040
2041 let page2 = make_v2_page(&[0, 1, 0, 1], &[1, 1, 1, 1], &[50, 60, 70, 80], 2);
2043
2044 let page3 = make_v2_page(&[0, 1], &[1, 1], &[90, 100], 1);
2046
2047 let pages = VecDeque::from(vec![page1, page2, page3]);
2049 let page_reader = InMemoryPageReader::new(pages);
2050 let column_reader = get_column_reader(desc, Box::new(page_reader));
2051 let mut typed_reader = get_typed_column_reader::<Int32Type>(column_reader);
2052
2053 let skipped = typed_reader.skip_records(1).unwrap();
2059 assert_eq!(skipped, 1);
2060
2061 let skipped = typed_reader.skip_records(2).unwrap();
2076 assert_eq!(skipped, 2);
2077
2078 let mut values = Vec::new();
2080 let mut def_levels = Vec::new();
2081 let mut rep_levels = Vec::new();
2082
2083 let (records, values_read, levels_read) = typed_reader
2084 .read_records(1, Some(&mut def_levels), Some(&mut rep_levels), &mut values)
2085 .unwrap();
2086
2087 assert_eq!(records, 1, "should read exactly 1 record");
2091 assert_eq!(levels_read, 2, "should read 2 levels for the record");
2092 assert_eq!(values_read, 2, "should read 2 non-null values");
2093 assert_eq!(values, vec![70, 80], "should contain 4th record's values");
2094 assert_eq!(rep_levels, vec![0, 1], "rep levels for a 2-element list");
2095 assert_eq!(def_levels, vec![1, 1], "def levels (all non-null)");
2096 }
2097}