parquet/arrow/array_reader/
byte_array.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
5use crate::arrow::array_reader::{ArrayReader, read_records, skip_records};
6use crate::arrow::buffer::bit_util::sign_extend_be;
7use crate::arrow::buffer::offset_buffer::OffsetBuffer;
8use crate::arrow::decoder::{DeltaByteArrayDecoder, DictIndexDecoder};
9use crate::arrow::record_reader::GenericRecordReader;
10use crate::arrow::schema::parquet_to_arrow_field;
11use crate::basic::{ConvertedType, Encoding};
12use crate::column::page::PageIterator;
13use crate::column::reader::decoder::ColumnValueDecoder;
14use crate::data_type::Int32Type;
15use crate::encodings::decoding::{Decoder, DeltaBitPackDecoder};
16use crate::errors::{ParquetError, Result};
17use crate::schema::types::ColumnDescPtr;
18use arrow_array::{
19    Array, ArrayRef, BinaryArray, Decimal128Array, Decimal256Array, OffsetSizeTrait,
20};
21use arrow_buffer::i256;
22use arrow_schema::DataType as ArrowType;
23use bytes::Bytes;
24use std::any::Any;
25use std::sync::Arc;
26
27/// Returns an [`ArrayReader`] that decodes the provided byte array column
28pub fn make_byte_array_reader(
29    pages: Box<dyn PageIterator>,
30    column_desc: ColumnDescPtr,
31    arrow_type: Option<ArrowType>,
32    batch_size: usize,
33    padding_threshold: Option<i16>,
34) -> Result<Box<dyn ArrayReader>> {
35    // Check if Arrow type is specified, else create it from Parquet type
36    let data_type = match arrow_type {
37        Some(t) => t,
38        None => parquet_to_arrow_field(column_desc.as_ref())?
39            .data_type()
40            .clone(),
41    };
42
43    match data_type {
44        ArrowType::Binary
45        | ArrowType::Utf8
46        | ArrowType::Decimal128(_, _)
47        | ArrowType::Decimal256(_, _) => {
48            let mut reader = GenericRecordReader::new(column_desc, batch_size);
49            if let Some(threshold) = padding_threshold {
50                reader.set_padding_threshold(threshold);
51            }
52            Ok(Box::new(ByteArrayReader::<i32>::new(
53                pages, data_type, reader,
54            )))
55        }
56        ArrowType::LargeUtf8 | ArrowType::LargeBinary => {
57            let mut reader = GenericRecordReader::new(column_desc, batch_size);
58            if let Some(threshold) = padding_threshold {
59                reader.set_padding_threshold(threshold);
60            }
61            Ok(Box::new(ByteArrayReader::<i64>::new(
62                pages, data_type, reader,
63            )))
64        }
65        _ => Err(general_err!(
66            "invalid data type for byte array reader - {}",
67            data_type
68        )),
69    }
70}
71
72/// An [`ArrayReader`] for variable length byte arrays
73struct ByteArrayReader<I: OffsetSizeTrait> {
74    data_type: ArrowType,
75    pages: Box<dyn PageIterator>,
76    def_levels_buffer: Option<Vec<i16>>,
77    rep_levels_buffer: Option<Vec<i16>>,
78    record_reader: GenericRecordReader<OffsetBuffer<I>, ByteArrayColumnValueDecoder<I>>,
79}
80
81impl<I: OffsetSizeTrait> ByteArrayReader<I> {
82    fn new(
83        pages: Box<dyn PageIterator>,
84        data_type: ArrowType,
85        record_reader: GenericRecordReader<OffsetBuffer<I>, ByteArrayColumnValueDecoder<I>>,
86    ) -> Self {
87        Self {
88            data_type,
89            pages,
90            def_levels_buffer: None,
91            rep_levels_buffer: None,
92            record_reader,
93        }
94    }
95}
96
97impl<I: OffsetSizeTrait> ArrayReader for ByteArrayReader<I> {
98    fn as_any(&self) -> &dyn Any {
99        self
100    }
101
102    fn get_data_type(&self) -> &ArrowType {
103        &self.data_type
104    }
105
106    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
107        read_records(&mut self.record_reader, self.pages.as_mut(), batch_size)
108    }
109
110    fn consume_batch(&mut self) -> Result<ArrayRef> {
111        let buffer = self.record_reader.consume_record_data();
112        let null_buffer = self.record_reader.consume_compact_bitmap();
113        self.def_levels_buffer = self.record_reader.consume_def_levels();
114        self.rep_levels_buffer = self.record_reader.consume_rep_levels();
115        self.record_reader.reset();
116
117        let array: ArrayRef = match self.data_type {
118            // Apply conversion to all elements regardless of null slots as the conversions
119            // are infallible. This improves performance by avoiding a branch in the inner
120            // loop (see docs for `PrimitiveArray::from_unary`).
121            ArrowType::Decimal128(p, s) => {
122                let array = buffer.into_array(null_buffer, ArrowType::Binary);
123                let binary = array.as_any().downcast_ref::<BinaryArray>().unwrap();
124                // Null slots will have 0 length, so we need to check for that in the lambda
125                // or sign_extend_be will panic.
126                let decimal = Decimal128Array::from_unary(binary, |x| match x.len() {
127                    0 => i128::default(),
128                    _ => i128::from_be_bytes(sign_extend_be(x)),
129                })
130                .with_precision_and_scale(p, s)?;
131                Arc::new(decimal)
132            }
133            ArrowType::Decimal256(p, s) => {
134                let array = buffer.into_array(null_buffer, ArrowType::Binary);
135                let binary = array.as_any().downcast_ref::<BinaryArray>().unwrap();
136                // Null slots will have 0 length, so we need to check for that in the lambda
137                // or sign_extend_be will panic.
138                let decimal = Decimal256Array::from_unary(binary, |x| match x.len() {
139                    0 => i256::default(),
140                    _ => i256::from_be_bytes(sign_extend_be(x)),
141                })
142                .with_precision_and_scale(p, s)?;
143                Arc::new(decimal)
144            }
145            _ => buffer.into_array(null_buffer, self.data_type.clone()),
146        };
147
148        Ok(array)
149    }
150
151    fn skip_records(&mut self, num_records: usize) -> Result<usize> {
152        skip_records(&mut self.record_reader, self.pages.as_mut(), num_records)
153    }
154
155    fn get_def_levels(&self) -> Option<&[i16]> {
156        self.def_levels_buffer.as_deref()
157    }
158
159    fn get_rep_levels(&self) -> Option<&[i16]> {
160        self.rep_levels_buffer.as_deref()
161    }
162
163    fn max_def_level(&self) -> i16 {
164        self.record_reader.max_def_level()
165    }
166}
167
168/// A [`ColumnValueDecoder`] for variable length byte arrays
169struct ByteArrayColumnValueDecoder<I: OffsetSizeTrait> {
170    dict: Option<OffsetBuffer<I>>,
171    decoder: Option<ByteArrayDecoder>,
172    validate_utf8: bool,
173}
174
175impl<I: OffsetSizeTrait> ColumnValueDecoder for ByteArrayColumnValueDecoder<I> {
176    type Buffer = OffsetBuffer<I>;
177
178    fn new(desc: &ColumnDescPtr) -> Self {
179        let validate_utf8 = desc.converted_type() == ConvertedType::UTF8;
180        Self {
181            dict: None,
182            decoder: None,
183            validate_utf8,
184        }
185    }
186
187    fn set_dict(
188        &mut self,
189        buf: Bytes,
190        num_values: u32,
191        encoding: Encoding,
192        _is_sorted: bool,
193    ) -> Result<()> {
194        if !matches!(
195            encoding,
196            Encoding::PLAIN | Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY
197        ) {
198            return Err(nyi_err!(
199                "Invalid/Unsupported encoding type for dictionary: {}",
200                encoding
201            ));
202        }
203
204        let mut buffer = OffsetBuffer::with_capacity(0);
205        let mut decoder = ByteArrayDecoderPlain::new(
206            buf,
207            num_values as usize,
208            Some(num_values as usize),
209            self.validate_utf8,
210        );
211        decoder.read(&mut buffer, usize::MAX)?;
212        self.dict = Some(buffer);
213        Ok(())
214    }
215
216    fn set_data(
217        &mut self,
218        encoding: Encoding,
219        data: Bytes,
220        num_levels: usize,
221        num_values: Option<usize>,
222    ) -> Result<()> {
223        self.decoder = Some(ByteArrayDecoder::new(
224            encoding,
225            data,
226            num_levels,
227            num_values,
228            self.validate_utf8,
229        )?);
230        Ok(())
231    }
232
233    fn read(&mut self, out: &mut Self::Buffer, num_values: usize) -> Result<usize> {
234        let decoder = self
235            .decoder
236            .as_mut()
237            .ok_or_else(|| general_err!("no decoder set"))?;
238
239        decoder.read(out, num_values, self.dict.as_ref())
240    }
241
242    fn skip_values(&mut self, num_values: usize) -> Result<usize> {
243        let decoder = self
244            .decoder
245            .as_mut()
246            .ok_or_else(|| general_err!("no decoder set"))?;
247
248        decoder.skip(num_values, self.dict.as_ref())
249    }
250}
251
252/// A generic decoder from uncompressed parquet value data to [`OffsetBuffer`]
253pub enum ByteArrayDecoder {
254    Plain(ByteArrayDecoderPlain),
255    Dictionary(ByteArrayDecoderDictionary),
256    DeltaLength(ByteArrayDecoderDeltaLength),
257    DeltaByteArray(ByteArrayDecoderDelta),
258}
259
260impl ByteArrayDecoder {
261    pub fn new(
262        encoding: Encoding,
263        data: Bytes,
264        num_levels: usize,
265        num_values: Option<usize>,
266        validate_utf8: bool,
267    ) -> Result<Self> {
268        let decoder = match encoding {
269            Encoding::PLAIN => ByteArrayDecoder::Plain(ByteArrayDecoderPlain::new(
270                data,
271                num_levels,
272                num_values,
273                validate_utf8,
274            )),
275            Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => ByteArrayDecoder::Dictionary(
276                ByteArrayDecoderDictionary::new(data, num_levels, num_values)?,
277            ),
278            Encoding::DELTA_LENGTH_BYTE_ARRAY => ByteArrayDecoder::DeltaLength(
279                ByteArrayDecoderDeltaLength::new(data, validate_utf8)?,
280            ),
281            Encoding::DELTA_BYTE_ARRAY => {
282                ByteArrayDecoder::DeltaByteArray(ByteArrayDecoderDelta::new(data, validate_utf8)?)
283            }
284            _ => {
285                return Err(general_err!(
286                    "unsupported encoding for byte array: {}",
287                    encoding
288                ));
289            }
290        };
291
292        Ok(decoder)
293    }
294
295    /// Read up to `len` values to `out` with the optional dictionary
296    pub fn read<I: OffsetSizeTrait>(
297        &mut self,
298        out: &mut OffsetBuffer<I>,
299        len: usize,
300        dict: Option<&OffsetBuffer<I>>,
301    ) -> Result<usize> {
302        match self {
303            ByteArrayDecoder::Plain(d) => d.read(out, len),
304            ByteArrayDecoder::Dictionary(d) => {
305                let dict =
306                    dict.ok_or_else(|| general_err!("missing dictionary page for column"))?;
307
308                d.read(out, dict, len)
309            }
310            ByteArrayDecoder::DeltaLength(d) => d.read(out, len),
311            ByteArrayDecoder::DeltaByteArray(d) => d.read(out, len),
312        }
313    }
314
315    /// Skip `len` values
316    pub fn skip<I: OffsetSizeTrait>(
317        &mut self,
318        len: usize,
319        dict: Option<&OffsetBuffer<I>>,
320    ) -> Result<usize> {
321        match self {
322            ByteArrayDecoder::Plain(d) => d.skip(len),
323            ByteArrayDecoder::Dictionary(d) => {
324                let dict =
325                    dict.ok_or_else(|| general_err!("missing dictionary page for column"))?;
326
327                d.skip(dict, len)
328            }
329            ByteArrayDecoder::DeltaLength(d) => d.skip(len),
330            ByteArrayDecoder::DeltaByteArray(d) => d.skip(len),
331        }
332    }
333}
334
335/// Decoder from [`Encoding::PLAIN`] data to [`OffsetBuffer`]
336pub struct ByteArrayDecoderPlain {
337    buf: Bytes,
338    offset: usize,
339    validate_utf8: bool,
340
341    /// This is a maximum as the null count is not always known, e.g. value data from
342    /// a v1 data page
343    max_remaining_values: usize,
344}
345
346impl ByteArrayDecoderPlain {
347    pub fn new(
348        buf: Bytes,
349        num_levels: usize,
350        num_values: Option<usize>,
351        validate_utf8: bool,
352    ) -> Self {
353        Self {
354            buf,
355            validate_utf8,
356            offset: 0,
357            max_remaining_values: num_values.unwrap_or(num_levels),
358        }
359    }
360
361    pub fn read<I: OffsetSizeTrait>(
362        &mut self,
363        output: &mut OffsetBuffer<I>,
364        len: usize,
365    ) -> Result<usize> {
366        let initial_values_length = output.values.len();
367
368        let to_read = len.min(self.max_remaining_values);
369        output.offsets.reserve(to_read);
370
371        let remaining_bytes = self.buf.len() - self.offset;
372        if remaining_bytes == 0 {
373            return Ok(0);
374        }
375
376        let estimated_bytes = remaining_bytes
377            .checked_mul(to_read)
378            .map(|x| x / self.max_remaining_values)
379            .unwrap_or_default();
380
381        output.values.reserve(estimated_bytes);
382
383        let mut read = 0;
384
385        let buf = self.buf.as_ref();
386        while self.offset < self.buf.len() && read != to_read {
387            if self.offset + 4 > buf.len() {
388                return Err(ParquetError::EOF("eof decoding byte array".into()));
389            }
390            let len_bytes: [u8; 4] = buf[self.offset..self.offset + 4].try_into().unwrap();
391            let len = u32::from_le_bytes(len_bytes);
392
393            let start_offset = self.offset + 4;
394            let end_offset = start_offset + len as usize;
395            if end_offset > buf.len() {
396                return Err(ParquetError::EOF("eof decoding byte array".into()));
397            }
398
399            output.try_push(&buf[start_offset..end_offset], self.validate_utf8)?;
400
401            self.offset = end_offset;
402            read += 1;
403        }
404        self.max_remaining_values -= read;
405
406        if self.validate_utf8 {
407            output.check_valid_utf8(initial_values_length)?;
408        }
409        Ok(read)
410    }
411
412    pub fn skip(&mut self, to_skip: usize) -> Result<usize> {
413        let to_skip = to_skip.min(self.max_remaining_values);
414        let mut skip = 0;
415        let buf = self.buf.as_ref();
416
417        while self.offset < self.buf.len() && skip != to_skip {
418            if self.offset + 4 > buf.len() {
419                return Err(ParquetError::EOF("eof decoding byte array".into()));
420            }
421            let len_bytes: [u8; 4] = buf[self.offset..self.offset + 4].try_into().unwrap();
422            let len = u32::from_le_bytes(len_bytes) as usize;
423            skip += 1;
424            self.offset = self.offset + 4 + len;
425        }
426        self.max_remaining_values -= skip;
427        Ok(skip)
428    }
429}
430
431/// Decoder from [`Encoding::DELTA_LENGTH_BYTE_ARRAY`] data to [`OffsetBuffer`]
432pub struct ByteArrayDecoderDeltaLength {
433    lengths: Vec<i32>,
434    data: Bytes,
435    length_offset: usize,
436    data_offset: usize,
437    validate_utf8: bool,
438}
439
440impl ByteArrayDecoderDeltaLength {
441    fn new(data: Bytes, validate_utf8: bool) -> Result<Self> {
442        let mut len_decoder = DeltaBitPackDecoder::<Int32Type>::new();
443        len_decoder.set_data(data.clone(), 0)?;
444        let values = len_decoder.values_left();
445
446        let mut lengths = vec![0; values];
447        len_decoder.get(&mut lengths)?;
448
449        let mut total_bytes = len_decoder.get_offset();
450
451        for l in &lengths {
452            if *l < 0 {
453                return Err(ParquetError::General(
454                    "negative delta length byte array length".to_string(),
455                ));
456            }
457            total_bytes = total_bytes
458                .checked_add(*l as usize)
459                .ok_or_else(|| general_err!("delta length byte array length overflow"))?;
460        }
461
462        if total_bytes > data.len() {
463            return Err(ParquetError::General(
464                "Insufficient delta length byte array bytes".to_string(),
465            ));
466        }
467
468        Ok(Self {
469            lengths,
470            data,
471            validate_utf8,
472            length_offset: 0,
473            data_offset: len_decoder.get_offset(),
474        })
475    }
476
477    fn read<I: OffsetSizeTrait>(
478        &mut self,
479        output: &mut OffsetBuffer<I>,
480        len: usize,
481    ) -> Result<usize> {
482        let initial_values_length = output.values.len();
483
484        let to_read = len.min(self.lengths.len() - self.length_offset);
485        let src_lengths = &self.lengths[self.length_offset..self.length_offset + to_read];
486        let total_bytes: usize = src_lengths.iter().map(|x| *x as usize).sum();
487
488        Self::checked_offset::<I>(initial_values_length, total_bytes)?;
489
490        // Reserve capacity for both offsets and values upfront
491        output.offsets.reserve(to_read);
492        output.values.reserve(total_bytes);
493
494        // Delta length data is contiguous — copy all value bytes at once
495        let data_end = self.data_offset + total_bytes;
496        output
497            .values
498            .extend_from_slice(&self.data.as_ref()[self.data_offset..data_end]);
499
500        // Compute and extend offsets in batch using extend
501        let base_offset = initial_values_length;
502        let mut running = base_offset;
503        output.offsets.extend(src_lengths.iter().map(|length| {
504            running += *length as usize;
505            // Nonnegative lengths and the checked final offset make every offset valid.
506            I::usize_as(running)
507        }));
508
509        self.data_offset = data_end;
510        self.length_offset += to_read;
511
512        if self.validate_utf8 {
513            output.check_valid_utf8(initial_values_length)?;
514        }
515        Ok(to_read)
516    }
517
518    fn checked_offset<I: OffsetSizeTrait>(offset: usize, additional: usize) -> Result<I> {
519        offset
520            .checked_add(additional)
521            .and_then(I::from_usize)
522            .ok_or_else(|| general_err!("index overflow decoding byte array"))
523    }
524
525    fn skip(&mut self, to_skip: usize) -> Result<usize> {
526        let remain_values = self.lengths.len() - self.length_offset;
527        let to_skip = remain_values.min(to_skip);
528
529        let src_lengths = &self.lengths[self.length_offset..self.length_offset + to_skip];
530        let total_bytes: usize = src_lengths.iter().map(|x| *x as usize).sum();
531
532        self.data_offset += total_bytes;
533        self.length_offset += to_skip;
534        Ok(to_skip)
535    }
536}
537
538/// Decoder from [`Encoding::DELTA_BYTE_ARRAY`] to [`OffsetBuffer`]
539pub struct ByteArrayDecoderDelta {
540    decoder: DeltaByteArrayDecoder,
541    validate_utf8: bool,
542}
543
544impl ByteArrayDecoderDelta {
545    fn new(data: Bytes, validate_utf8: bool) -> Result<Self> {
546        Ok(Self {
547            decoder: DeltaByteArrayDecoder::new(data)?,
548            validate_utf8,
549        })
550    }
551
552    fn read<I: OffsetSizeTrait>(
553        &mut self,
554        output: &mut OffsetBuffer<I>,
555        len: usize,
556    ) -> Result<usize> {
557        let initial_values_length = output.values.len();
558        output.offsets.reserve(len.min(self.decoder.remaining()));
559
560        let read = self
561            .decoder
562            .read(len, |bytes| output.try_push(bytes, self.validate_utf8))?;
563
564        if self.validate_utf8 {
565            output.check_valid_utf8(initial_values_length)?;
566        }
567        Ok(read)
568    }
569
570    fn skip(&mut self, to_skip: usize) -> Result<usize> {
571        self.decoder.skip(to_skip)
572    }
573}
574
575/// Decoder from [`Encoding::RLE_DICTIONARY`] to [`OffsetBuffer`]
576pub struct ByteArrayDecoderDictionary {
577    decoder: DictIndexDecoder,
578}
579
580impl ByteArrayDecoderDictionary {
581    fn new(data: Bytes, num_levels: usize, num_values: Option<usize>) -> Result<Self> {
582        Ok(Self {
583            decoder: DictIndexDecoder::new(data, num_levels, num_values)?,
584        })
585    }
586
587    fn read<I: OffsetSizeTrait>(
588        &mut self,
589        output: &mut OffsetBuffer<I>,
590        dict: &OffsetBuffer<I>,
591        len: usize,
592    ) -> Result<usize> {
593        // All data must be NULL
594        if dict.is_empty() {
595            return Ok(0);
596        }
597
598        // Pre-reserve offsets capacity to avoid per-chunk reallocation
599        output.offsets.reserve(len);
600
601        self.decoder.read(len, |keys| {
602            output.extend_from_dictionary(keys, dict.offsets.as_slice(), dict.values.as_slice())
603        })
604    }
605
606    fn skip<I: OffsetSizeTrait>(
607        &mut self,
608        dict: &OffsetBuffer<I>,
609        to_skip: usize,
610    ) -> Result<usize> {
611        // All data must be NULL
612        if dict.is_empty() {
613            return Ok(0);
614        }
615
616        self.decoder.skip(to_skip)
617    }
618}
619
620#[cfg(test)]
621mod tests {
622    use super::*;
623    use crate::arrow::array_reader::test_util::{byte_array_all_encodings, utf8_column};
624    use crate::arrow::record_reader::buffer::ValuesBuffer;
625    use arrow_array::{Array, StringArray};
626    use arrow_buffer::Buffer;
627
628    #[test]
629    fn test_byte_array_decoder() {
630        let (pages, encoded_dictionary) =
631            byte_array_all_encodings(vec!["hello", "world", "a", "b"]);
632
633        let column_desc = utf8_column();
634        let mut decoder = ByteArrayColumnValueDecoder::new(&column_desc);
635
636        decoder
637            .set_dict(encoded_dictionary, 4, Encoding::RLE_DICTIONARY, false)
638            .unwrap();
639
640        for (encoding, page) in pages {
641            let mut output = OffsetBuffer::<i32>::with_capacity(0);
642            decoder.set_data(encoding, page, 4, Some(4)).unwrap();
643
644            assert_eq!(decoder.read(&mut output, 1).unwrap(), 1);
645
646            assert_eq!(output.values.as_slice(), b"hello");
647            assert_eq!(output.offsets.as_slice(), &[0, 5]);
648
649            assert_eq!(decoder.read(&mut output, 1).unwrap(), 1);
650            assert_eq!(output.values.as_slice(), b"helloworld");
651            assert_eq!(output.offsets.as_slice(), &[0, 5, 10]);
652
653            assert_eq!(decoder.read(&mut output, 2).unwrap(), 2);
654            assert_eq!(output.values.as_slice(), b"helloworldab");
655            assert_eq!(output.offsets.as_slice(), &[0, 5, 10, 11, 12]);
656
657            assert_eq!(decoder.read(&mut output, 4).unwrap(), 0);
658
659            let valid = [false, false, true, true, false, true, true, false, false];
660            let valid_buffer = Buffer::from_iter(valid.iter().copied());
661
662            output
663                .pad_nulls(0, 4, valid.len(), valid_buffer.as_slice())
664                .unwrap();
665            let array = output.into_array(Some(valid_buffer), ArrowType::Utf8);
666            let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
667
668            assert_eq!(
669                strings.iter().collect::<Vec<_>>(),
670                vec![
671                    None,
672                    None,
673                    Some("hello"),
674                    Some("world"),
675                    None,
676                    Some("a"),
677                    Some("b"),
678                    None,
679                    None,
680                ]
681            );
682        }
683    }
684
685    #[test]
686    fn test_byte_array_decoder_skip() {
687        let (pages, encoded_dictionary) =
688            byte_array_all_encodings(vec!["hello", "world", "a", "b"]);
689
690        let column_desc = utf8_column();
691        let mut decoder = ByteArrayColumnValueDecoder::new(&column_desc);
692
693        decoder
694            .set_dict(encoded_dictionary, 4, Encoding::RLE_DICTIONARY, false)
695            .unwrap();
696
697        for (encoding, page) in pages {
698            let mut output = OffsetBuffer::<i32>::with_capacity(0);
699            decoder.set_data(encoding, page, 4, Some(4)).unwrap();
700
701            assert_eq!(decoder.read(&mut output, 1).unwrap(), 1);
702
703            assert_eq!(output.values.as_slice(), b"hello");
704            assert_eq!(output.offsets.as_slice(), &[0, 5]);
705
706            assert_eq!(decoder.skip_values(1).unwrap(), 1);
707            assert_eq!(decoder.skip_values(1).unwrap(), 1);
708
709            assert_eq!(decoder.read(&mut output, 1).unwrap(), 1);
710            assert_eq!(output.values.as_slice(), b"hellob");
711            assert_eq!(output.offsets.as_slice(), &[0, 5, 6]);
712
713            assert_eq!(decoder.read(&mut output, 4).unwrap(), 0);
714
715            let valid = [false, false, true, true, false, false];
716            let valid_buffer = Buffer::from_iter(valid.iter().copied());
717
718            output
719                .pad_nulls(0, 2, valid.len(), valid_buffer.as_slice())
720                .unwrap();
721            let array = output.into_array(Some(valid_buffer), ArrowType::Utf8);
722            let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
723
724            assert_eq!(
725                strings.iter().collect::<Vec<_>>(),
726                vec![None, None, Some("hello"), Some("b"), None, None,]
727            );
728        }
729    }
730
731    #[test]
732    fn test_byte_array_decoder_nulls() {
733        let (pages, encoded_dictionary) = byte_array_all_encodings(Vec::<&str>::new());
734
735        let column_desc = utf8_column();
736        let mut decoder = ByteArrayColumnValueDecoder::new(&column_desc);
737
738        decoder
739            .set_dict(encoded_dictionary, 4, Encoding::RLE_DICTIONARY, false)
740            .unwrap();
741
742        // test nulls read
743        for (encoding, page) in pages.clone() {
744            let mut output = OffsetBuffer::<i32>::with_capacity(0);
745            decoder.set_data(encoding, page, 4, None).unwrap();
746            assert_eq!(decoder.read(&mut output, 1024).unwrap(), 0);
747        }
748
749        // test nulls skip
750        for (encoding, page) in pages {
751            decoder.set_data(encoding, page, 4, None).unwrap();
752            assert_eq!(decoder.skip_values(1024).unwrap(), 0);
753        }
754    }
755
756    #[test]
757    fn test_plain_decoder_reports_values_actually_read() {
758        // The page claims to contain two values, but its buffer contains only
759        // one complete PLAIN-encoded BYTE_ARRAY value.
760        let buffer = Bytes::from_static(&[3, 0, 0, 0, b'f', b'o', b'o']);
761        let mut decoder = ByteArrayDecoderPlain::new(buffer, 2, Some(2), false);
762        let mut output = OffsetBuffer::<i32>::with_capacity(2);
763
764        assert_eq!(decoder.read(&mut output, 2).unwrap(), 1);
765        assert_eq!(output.values.as_slice(), b"foo");
766        assert_eq!(output.offsets.as_slice(), &[0, 3]);
767        assert_eq!(decoder.max_remaining_values, 1);
768    }
769
770    #[test]
771    fn test_delta_length_checked_offset() {
772        let max = i32::MAX as usize;
773        assert_eq!(
774            ByteArrayDecoderDeltaLength::checked_offset::<i32>(max - 1, 1).unwrap(),
775            i32::MAX
776        );
777        assert_eq!(
778            ByteArrayDecoderDeltaLength::checked_offset::<i32>(max, 0).unwrap(),
779            i32::MAX
780        );
781        for (offset, additional) in [(max, 1), (0, max + 1), (usize::MAX, 1)] {
782            let error =
783                ByteArrayDecoderDeltaLength::checked_offset::<i32>(offset, additional).unwrap_err();
784            assert_eq!(
785                error.to_string(),
786                "Parquet error: index overflow decoding byte array"
787            );
788        }
789        assert_eq!(
790            ByteArrayDecoderDeltaLength::checked_offset::<i64>(max, 1).unwrap(),
791            i32::MAX as i64 + 1
792        );
793        assert!(ByteArrayDecoderDeltaLength::checked_offset::<i64>(usize::MAX, 1).is_err());
794    }
795
796    #[test]
797    fn test_delta_length_read_append_and_skip() {
798        let (pages, _) = byte_array_all_encodings(vec!["ab", "", "cde", "f"]);
799        let (_, page) = pages
800            .into_iter()
801            .find(|(encoding, _)| *encoding == Encoding::DELTA_LENGTH_BYTE_ARRAY)
802            .unwrap();
803        let mut decoder = ByteArrayDecoderDeltaLength::new(page, true).unwrap();
804        let mut output = OffsetBuffer::<i32>::with_capacity(0);
805        output.try_push(b"prefix", true).unwrap();
806
807        assert_eq!(decoder.read(&mut output, 2).unwrap(), 2);
808        assert_eq!(output.values.as_slice(), b"prefixab");
809        assert_eq!(output.offsets.as_slice(), &[0, 6, 8, 8]);
810        assert_eq!(decoder.skip(1).unwrap(), 1);
811        assert_eq!(decoder.read(&mut output, 1).unwrap(), 1);
812        assert_eq!(output.values.as_slice(), b"prefixabf");
813        assert_eq!(output.offsets.as_slice(), &[0, 6, 8, 8, 9]);
814        assert_eq!(decoder.read(&mut output, 1).unwrap(), 0);
815    }
816
817    // The checked_offset unit tests cover arithmetic; this also verifies that
818    // read rejects a real overflow before changing output or decoder state.
819    // Keep this regression available for manual runs without requiring an
820    // allocation of about 2 GiB in routine CI.
821    #[test]
822    #[ignore = "allocates about 2 GiB; run --ignored --test-threads=1"]
823    fn test_delta_length_offset_overflow_preserves_state() {
824        let (pages, _) = byte_array_all_encodings(vec!["", "a"]);
825        let (_, page) = pages
826            .into_iter()
827            .find(|(encoding, _)| *encoding == Encoding::DELTA_LENGTH_BYTE_ARRAY)
828            .unwrap();
829        let mut decoder = ByteArrayDecoderDeltaLength::new(page, false).unwrap();
830        let data_offset = decoder.data_offset;
831        let mut output = OffsetBuffer::<i32> {
832            offsets: vec![0, i32::MAX],
833            values: vec![0; i32::MAX as usize],
834        };
835
836        let error = decoder.read(&mut output, 2).unwrap_err();
837        assert_eq!(
838            error.to_string(),
839            "Parquet error: index overflow decoding byte array"
840        );
841        assert_eq!(output.values.len(), i32::MAX as usize);
842        assert_eq!(output.offsets, vec![0, i32::MAX]);
843        assert_eq!(decoder.data_offset, data_offset);
844        assert_eq!(decoder.length_offset, 0);
845
846        assert_eq!(decoder.read(&mut output, 1).unwrap(), 1);
847        assert_eq!(output.offsets, vec![0, i32::MAX, i32::MAX]);
848        assert_eq!(decoder.skip(1).unwrap(), 1);
849        assert_eq!(decoder.read(&mut output, 1).unwrap(), 0);
850    }
851}
parquet/arrow/array_reader/byte_array.rs

parquet/arrow/array_reader/
byte_array.rs
