parquet/arrow/decoder/
delta_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 bytes::Bytes;
6
7use crate::data_type::Int32Type;
8use crate::encodings::decoding::{Decoder, DeltaBitPackDecoder};
9use crate::errors::{ParquetError, Result};
10
11/// Decoder for `Encoding::DELTA_BYTE_ARRAY`
12pub struct DeltaByteArrayDecoder {
13    prefix_lengths: Vec<i32>,
14    suffix_lengths: Vec<i32>,
15    data: Bytes,
16    length_offset: usize,
17    data_offset: usize,
18    last_value: Vec<u8>,
19}
20
21impl DeltaByteArrayDecoder {
22    /// Create a new [`DeltaByteArrayDecoder`] with the provided data page
23    pub fn new(data: Bytes) -> Result<Self> {
24        let mut prefix = DeltaBitPackDecoder::<Int32Type>::new();
25        prefix.set_data(data.clone(), 0)?;
26
27        let num_prefix = prefix.values_left();
28        let mut prefix_lengths = vec![0; num_prefix];
29        assert_eq!(prefix.get(&mut prefix_lengths)?, num_prefix);
30
31        let mut suffix = DeltaBitPackDecoder::<Int32Type>::new();
32        suffix.set_data(data.slice(prefix.get_offset()..), 0)?;
33
34        let num_suffix = suffix.values_left();
35        let mut suffix_lengths = vec![0; num_suffix];
36        assert_eq!(suffix.get(&mut suffix_lengths)?, num_suffix);
37
38        if num_prefix != num_suffix {
39            return Err(general_err!(format!(
40                "inconsistent DELTA_BYTE_ARRAY lengths, prefixes: {num_prefix}, suffixes: {num_suffix}"
41            )));
42        }
43
44        assert_eq!(prefix_lengths.len(), suffix_lengths.len());
45
46        Ok(Self {
47            prefix_lengths,
48            suffix_lengths,
49            data,
50            length_offset: 0,
51            data_offset: prefix.get_offset() + suffix.get_offset(),
52            last_value: vec![],
53        })
54    }
55
56    /// Returns the number of values remaining
57    pub fn remaining(&self) -> usize {
58        self.prefix_lengths.len() - self.length_offset
59    }
60
61    /// Read up to `len` values, returning the number of values read
62    /// and calling `f` with each decoded byte slice
63    ///
64    /// Will short-circuit and return on error
65    pub fn read<F>(&mut self, len: usize, mut f: F) -> Result<usize>
66    where
67        F: FnMut(&[u8]) -> Result<()>,
68    {
69        let to_read = len.min(self.remaining());
70
71        let length_range = self.length_offset..self.length_offset + to_read;
72        let iter = self.prefix_lengths[length_range.clone()]
73            .iter()
74            .zip(&self.suffix_lengths[length_range]);
75
76        let data = self.data.as_ref();
77
78        for (prefix_length, suffix_length) in iter {
79            let prefix_length = *prefix_length as usize;
80            let suffix_length = *suffix_length as usize;
81
82            if self.data_offset + suffix_length > self.data.len() {
83                return Err(ParquetError::EOF("eof decoding byte array".into()));
84            }
85
86            self.last_value.truncate(prefix_length);
87            self.last_value
88                .extend_from_slice(&data[self.data_offset..self.data_offset + suffix_length]);
89            f(&self.last_value)?;
90
91            self.data_offset += suffix_length;
92        }
93
94        self.length_offset += to_read;
95        Ok(to_read)
96    }
97
98    /// Skip up to `to_skip` values, returning the number of values skipped
99    pub fn skip(&mut self, to_skip: usize) -> Result<usize> {
100        let to_skip = to_skip.min(self.prefix_lengths.len() - self.length_offset);
101
102        let length_range = self.length_offset..self.length_offset + to_skip;
103        let iter = self.prefix_lengths[length_range.clone()]
104            .iter()
105            .zip(&self.suffix_lengths[length_range]);
106
107        let data = self.data.as_ref();
108
109        for (prefix_length, suffix_length) in iter {
110            let prefix_length = *prefix_length as usize;
111            let suffix_length = *suffix_length as usize;
112
113            if self.data_offset + suffix_length > self.data.len() {
114                return Err(ParquetError::EOF("eof decoding byte array".into()));
115            }
116
117            self.last_value.truncate(prefix_length);
118            self.last_value
119                .extend_from_slice(&data[self.data_offset..self.data_offset + suffix_length]);
120            self.data_offset += suffix_length;
121        }
122        self.length_offset += to_skip;
123        Ok(to_skip)
124    }
125}
parquet/arrow/decoder/delta_byte_array.rs

parquet/arrow/decoder/
delta_byte_array.rs
