Skip to main content

parquet/column/
reader.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Contains column reader API.
6
7use 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
24/// Column reader for a Parquet type.
25pub enum ColumnReader {
26    /// Column reader for boolean type
27    BoolColumnReader(ColumnReaderImpl<BoolType>),
28    /// Column reader for int32 type
29    Int32ColumnReader(ColumnReaderImpl<Int32Type>),
30    /// Column reader for int64 type
31    Int64ColumnReader(ColumnReaderImpl<Int64Type>),
32    /// Column reader for int96 type
33    Int96ColumnReader(ColumnReaderImpl<Int96Type>),
34    /// Column reader for float type
35    FloatColumnReader(ColumnReaderImpl<FloatType>),
36    /// Column reader for double type
37    DoubleColumnReader(ColumnReaderImpl<DoubleType>),
38    /// Column reader for byte array type
39    ByteArrayColumnReader(ColumnReaderImpl<ByteArrayType>),
40    /// Column reader for fixed length byte array type
41    FixedLenByteArrayColumnReader(ColumnReaderImpl<FixedLenByteArrayType>),
42}
43
44/// Gets a specific column reader corresponding to column descriptor `col_descr`. The
45/// column reader will read from pages in `col_page_reader`.
46pub 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
78/// Gets a typed column reader for the specific type `T`, by "up-casting" `col_reader` of
79/// non-generic type to a generic column reader type `ColumnReaderImpl`.
80///
81/// # Panics
82///
83/// Panics if actual enum value for `col_reader` does not match the type `T`.
84pub 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
93/// Typed value reader for a particular primitive column.
94pub type ColumnReaderImpl<T> = GenericColumnReader<
95    RepetitionLevelDecoderImpl,
96    DefinitionLevelDecoderImpl,
97    ColumnValueDecoderImpl<T>,
98>;
99
100/// Reads data for a given column chunk, using the provided decoders:
101///
102/// - R: `ColumnLevelDecoder` used to decode repetition levels
103/// - D: `ColumnLevelDecoder` used to decode definition levels
104/// - V: `ColumnValueDecoder` used to decode value data
105pub struct GenericColumnReader<R, D, V> {
106    descr: ColumnDescPtr,
107
108    page_reader: Box<dyn PageReader>,
109
110    /// The total number of values stored in the data page.
111    num_buffered_values: usize,
112
113    /// The number of values from the current data page that has been decoded into memory
114    /// so far.
115    num_decoded_values: usize,
116
117    /// True if the end of the current data page denotes the end of a record
118    has_record_delimiter: bool,
119
120    /// The decoder for the definition levels if any
121    def_level_decoder: Option<D>,
122
123    /// The decoder for the repetition levels if any
124    rep_level_decoder: Option<R>,
125
126    /// The decoder for the values
127    values_decoder: V,
128
129    /// Width and remaining payload bytes for a nullable V1 FLBA page whose
130    /// physical count is checked during normal level decoding, not a pre-scan.
131    remaining_fixed_len_bytes: Option<(usize, usize)>,
132}
133
134impl<V> GenericColumnReader<RepetitionLevelDecoderImpl, DefinitionLevelDecoderImpl, V>
135where
136    V: ColumnValueDecoder,
137{
138    /// Creates new column reader based on column descriptor and page reader.
139    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    /// Read up to `max_records` whole records, returning the number of complete
185    /// records, non-null values and levels decoded. All levels for a given record
186    /// will be read, i.e. the next repetition level, if any, will be 0
187    ///
188    /// If the max definition level is 0, `def_levels` will be ignored and the number of records,
189    /// non-null values and levels decoded will all be equal, otherwise `def_levels` will be
190    /// populated with the number of levels read, with an error returned if it is `None`.
191    ///
192    /// If the max repetition level is 0, `rep_levels` will be ignored and the number of records
193    /// and levels decoded will both be equal, otherwise `rep_levels` will be populated with
194    /// the number of levels read, with an error returned if it is `None`.
195    ///
196    /// `values` will be contiguously populated with the non-null values. Note that if the column
197    /// is not required, this may be less than either `max_records` or the number of levels read.
198    ///
199    /// Nullable V1 FLBA payload sizes may be validated incrementally. Extra values or
200    /// truncated levels can therefore be reported only when the rest of a page is consumed.
201    /// Pages that could require legacy PLAIN repair are still validated eagerly.
202    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    /// Like [`Self::read_records`], but invokes `reserve_values` after levels
219    /// are decoded and before values are read.
220    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                        // The fact that we're still looping implies there must be some levels to read.
250                        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    /// Skips over `num_records` records, where records are delimited by repetition levels of 0
308    ///
309    /// # Returns
310    ///
311    /// Returns the number of records skipped
312    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                // Defer dictionary decoding when the page reader can recover it later.
321                // Otherwise, preserve the eager behavior required by existing readers.
322                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                // If page has less rows than the remaining records to
332                // be skipped, skip entire page
333                let rows = metadata.num_rows.or_else(|| {
334                    // If no repetition levels, num_levels == num_rows
335                    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                // because self.num_buffered_values == self.num_decoded_values means
348                // we need reads a new page and set up the decoders for levels
349                if !self.read_new_page()? {
350                    return Ok(num_records - remaining_records);
351                }
352            }
353
354            // start skip values in page level
355
356            // The number of levels in the current data page
357            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                    // No repetition levels, so each level corresponds to a row
373                    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                // Exhausted buffered page - no need to advance other decoders
384                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    /// Account for actual physical values decoded or skipped from definition levels.
418    /// The payload length is a byte budget, not evidence of the expected value count.
419    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    /// Reads the next page as a dictionary page and installs it.
445    ///
446    /// Used on the skip path for readers that do not retain skipped
447    /// dictionaries, and so must consume the page through `get_next_page`
448    /// while it is still available. Returns an error if the next page is not
449    /// a dictionary page.
450    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    /// Installs a dictionary previously deferred by the page reader.
467    /// A no-op when nothing was deferred.
468    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    /// Reads a new page and set up the decoders for levels, values or dictionary.
486    /// Returns false if there's no page left.
487    fn read_new_page(&mut self) -> Result<bool> {
488        loop {
489            match self.page_reader.get_next_page()? {
490                // No more page to read
491                None => return Ok(false),
492                Some(current_page) => {
493                    match current_page {
494                        // 1. Dictionary page: configure dictionary for this page.
495                        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                        // 2. Data page v1
505                        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                        // 3. Data page v2
593                        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                            // DataPage v2 only supports RLE encoding for repetition
624                            // levels
625                            if self.descr.max_rep_level() > 0 {
626                                // Technically a DataPage v2 should not write a record
627                                // across multiple pages, however, the parquet writer
628                                // used to do this so we preserve backwards compatibility
629                                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                            // DataPage v2 only supports RLE encoding for definition
639                            // levels
640                            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    /// Check whether there is more data to read from this column,
678    /// If the current page is fully decoded, this will load the next page
679    /// (if it exists) into the buffer
680    #[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            // TODO: should we return false if read_new_page() = true and
684            // num_buffered_values = 0?
685            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
696/// Checks that a partial record can still be flushed into the caller's record budget.
697///
698/// Reaching the end of a page with a record delimiter means the decoder stopped because
699/// it ran out of levels, not records, so it cannot have used up the whole budget. A page
700/// whose repetition levels disagree with its record count breaks that, and flushing the
701/// partial record would then take the count past what was asked for.
702fn 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    /// These read/skip assertions apply to both canonical and normalized values.
791    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                // Count-dependent corruption can be detected after the first
810                // value, regardless of the original value representation.
811                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        // Finishing by skipping checks both missing and extra values, including
836        // an all-null page with an unexpected whole value in its payload.
837        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        // Ten canonical values look exactly like five length-prefixed values.
863        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; // exceed any packed padding
905        }
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    /// Shared page/level assertions, independent of the value representation.
933    /// Keeping the canonical test separate lets it survive removal of legacy support.
934    #[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            // Null parents, empty containers, null elements and repeated values.
968            (3, vec![0, 1, 2, 3, 3, 1, 3], vec![0, 0, 0, 1, 1, 0, 1]),
969            // Exercise counting across multiple bounded definition-level batches.
970            (
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                        // Exceed even the final bit-packed run's padding.
1110                        *num_values += 8;
1111                    }
1112                    // Incremental V1 validation detects missing levels when
1113                    // the page is consumed, not necessarily on the first row.
1114                    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                    // Byte-alignment errors remain eager. Nullable V1 errors
1132                    // requiring the physical count can be delayed until page end.
1133                    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                    // A valid dictionary and dictionary-encoded page must not
1172                    // conceal malformed PLAIN data after dictionary fallback.
1173                    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        // Say length is 10, but buffer is only 4
1210        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        // Say length is 4, but buffer is only 3
1218        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 to generate test cases
1232    macro_rules! test {
1233        // branch for generating i32 cases
1234        ($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        // branch for generating i64 cases
1251        ($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 cases when column descriptor has MAX_DEF_LEVEL = 0 and MAX_REP_LEVEL = 0
1363    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 cases when column descriptor has MAX_DEF_LEVEL = 0 and MAX_REP_LEVEL = 0
1464    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        // This test covers scenario when buffering new page results in setting number
1637        // of decoded values to 0, resulting on reading `batch_size` of values, but it is
1638        // larger than we can insert into slice (affects values and levels).
1639        //
1640        // Note: values are chosen to reproduce the issue.
1641        //
1642        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    // ----------------------------------------------------------------------
1668    // Helper methods to make pages and test
1669    //
1670    // # Overview
1671    //
1672    // Most of the test functionality is implemented in `ColumnReaderTester`, which
1673    // provides some general data page test methods:
1674    // - `test_read_batch_general`
1675    // - `test_read_batch`
1676    //
1677    // There are also some high level wrappers that are part of `ColumnReaderTester`:
1678    // - `plain_v1` -> call `test_read_batch_general` with data page v1 and plain encoding
1679    // - `plain_v2` -> call `test_read_batch_general` with data page v2 and plain encoding
1680    // - `dict_v1` -> call `test_read_batch_general` with data page v1 + dictionary page
1681    // - `dict_v2` -> call `test_read_batch_general` with data page v2 + dictionary page
1682    //
1683    // And even higher level wrappers that simplify testing of almost the same test cases:
1684    // - `get_test_int32_type`, provides dummy schema type
1685    // - `get_test_int64_type`, provides dummy schema type
1686    // - `test_read_batch_int32`, wrapper for `read_batch` tests, since they are basically
1687    //   the same, just different def/rep levels and batch size.
1688    //
1689    // # Page assembly
1690    //
1691    // Page construction and generation of values, definition and repetition levels
1692    // happens in `make_pages` function.
1693    // All values are randomly generated based on provided min/max, levels are calculated
1694    // based on provided max level for column descriptor (which is basically either int32
1695    // or int64 type in tests) and `levels_per_page` variable.
1696    //
1697    // We use `DataPageBuilder` and its implementation `DataPageBuilderImpl` to actually
1698    // turn values, definition and repetition levels into data pages (either v1 or v2).
1699    //
1700    // Those data pages are then stored as part of `TestPageReader` (we just pass vector
1701    // of generated pages directly), which implements `PageReader` interface.
1702    //
1703    // # Comparison
1704    //
1705    // This allows us to pass test page reader into column reader, so we can test
1706    // functionality of column reader - see `test_read_batch`, where we create column
1707    // reader -> typed column reader, buffer values in `read_batch` method and compare
1708    // output with generated data.
1709
1710    // Returns dummy Parquet `Type` for primitive field, because most of our tests use
1711    // INT32 physical type.
1712    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    // Returns dummy Parquet `Type` for INT64 physical type.
1722    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    // Tests `read_batch()` functionality for INT32.
1732    //
1733    // This is a high level wrapper on `ColumnReaderTester` that allows us to specify some
1734    // boilerplate code for setting up definition/repetition levels and column descriptor.
1735    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        // Method to generate and test data pages v1
1780        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        // Method to generate and test data pages v2
1802        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        // Method to generate and test dictionary page + data pages v1
1824        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        // Method to generate and test dictionary page + data pages v2
1846        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        // Helper function for the general case of `read_batch()` where `values`,
1868        // `def_levels` and `rep_levels` are always provided with enough space.
1869        #[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        // Helper function to test `read_batch()` method with custom buffers for values,
1887        // definition and repetition levels.
1888        #[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    /// Regression test for <https://github.com/apache/arrow-rs/issues/9370>
1968    ///
1969    /// Reproduces the production scenario: all DataPage v2 pages for a
1970    /// list column (rep_level=1) read without an offset index (i.e.
1971    /// `at_record_boundary` returns false for non-last pages).
1972    ///
1973    /// When a prior operation (here `skip_records(1)`) loads a v2 page,
1974    /// and a subsequent `skip_records` exhausts the remaining levels on
1975    /// that page, the rep level decoder is left with `has_partial=true`.
1976    /// Because `has_record_delimiter` is false, the partial is not
1977    /// flushed during level-based processing. When the next v2 page is
1978    /// then peeked with `num_rows` available, the whole-page-skip
1979    /// shortcut must flush the pending partial first. Otherwise:
1980    ///
1981    /// 1. The skip over-counts (skips N+1 records instead of N), and
1982    /// 2. The stale `has_partial` causes a subsequent `read_records` to
1983    ///    produce a "phantom" record with 0 values.
1984    #[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        // Column descriptor for a list element column (rep=1, def=1)
1992        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        // Helper: build a DataPage v2 for this list column.
2004        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        // All pages are DataPage v2 (matching the production scenario where
2035        // parquet-rs writes only v2 data pages and no offset index is loaded,
2036        // so at_record_boundary() returns false for non-last pages).
2037
2038        // Page 1 (v2): 2 records × 2 elements = [10,20], [30,40]
2039        let page1 = make_v2_page(&[0, 1, 0, 1], &[1, 1, 1, 1], &[10, 20, 30, 40], 2);
2040
2041        // Page 2 (v2): 2 records × 2 elements = [50,60], [70,80]
2042        let page2 = make_v2_page(&[0, 1, 0, 1], &[1, 1, 1, 1], &[50, 60, 70, 80], 2);
2043
2044        // Page 3 (v2): 1 record × 2 elements = [90,100]
2045        let page3 = make_v2_page(&[0, 1], &[1, 1], &[90, 100], 1);
2046
2047        // 5 records total: [10,20], [30,40], [50,60], [70,80], [90,100]
2048        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        // Step 1 — skip 1 record:
2054        //   Peek page 1: num_rows=2, remaining=1 → rows(2) > remaining(1),
2055        //   so the page is LOADED (not whole-page-skipped).
2056        //   Level-based skip consumes rep levels [0,1] for record [10,20],
2057        //   stopping at the 0 that starts record [30,40].
2058        let skipped = typed_reader.skip_records(1).unwrap();
2059        assert_eq!(skipped, 1);
2060
2061        // Step 2 — skip 2 more records ([30,40] and [50,60]):
2062        //   Mid-page in page 1 with 2 remaining levels [0,1] for [30,40].
2063        //   skip_rep_levels(2, 2): the leading 0 does NOT act as a record
2064        //   delimiter (has_partial=false, idx==0), so count_records returns
2065        //   (true, 0, 2) — all levels consumed, has_partial=true, 0 records.
2066        //
2067        //   has_record_delimiter is false → no flush at page boundary.
2068        //   Page 1 exhausted → peek page 2 (v2, num_rows=2).
2069        //
2070        //   With fix: flush_partial → remaining 2→1, page 2 NOT skipped
2071        //   (rows=2 > remaining=1). Load page 2, skip 1 record [50,60].
2072        //
2073        //   Without fix: rows(2) <= remaining(2) → page 2 whole-page-skipped,
2074        //   over-counting by 1. has_partial stays true (stale from page 1).
2075        let skipped = typed_reader.skip_records(2).unwrap();
2076        assert_eq!(skipped, 2);
2077
2078        // Step 3 — read 1 record:
2079        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        // Without the fix: (1, 0, 0) — phantom record from stale has_partial;
2088        //   the rep=0 on page 3 "completes" the phantom, yielding 0 values.
2089        // With the fix:    (1, 2, 2) — correctly reads record [70, 80].
2090        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}