1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287
use parquet_format_safe::DataPageHeaderV2;
use streaming_decompression;
use crate::compression::{self, Compression};
use crate::error::{Error, Result};
use crate::page::{CompressedPage, DataPage, DataPageHeader, DictPage, Page};
use crate::FallibleStreamingIterator;
use super::page::PageIterator;
fn decompress_v1(compressed: &[u8], compression: Compression, buffer: &mut [u8]) -> Result<()> {
compression::decompress(compression, compressed, buffer)
}
fn decompress_v2(
compressed: &[u8],
page_header: &DataPageHeaderV2,
compression: Compression,
buffer: &mut [u8],
) -> Result<()> {
// When processing data page v2, depending on enabled compression for the
// page, we should account for uncompressed data ('offset') of
// repetition and definition levels.
//
// We always use 0 offset for other pages other than v2, `true` flag means
// that compression will be applied if decompressor is defined
let offset = (page_header.definition_levels_byte_length
+ page_header.repetition_levels_byte_length) as usize;
// When is_compressed flag is missing the page is considered compressed
let can_decompress = page_header.is_compressed.unwrap_or(true);
if can_decompress {
if offset > buffer.len() || offset > compressed.len() {
return Err(Error::OutOfSpec(
"V2 Page Header reported incorrect offset to compressed data".to_string(),
));
}
(&mut buffer[..offset]).copy_from_slice(&compressed[..offset]);
compression::decompress(compression, &compressed[offset..], &mut buffer[offset..])?;
} else {
if buffer.len() != compressed.len() {
return Err(Error::OutOfSpec(
"V2 Page Header reported incorrect decompressed size".to_string(),
));
}
buffer.copy_from_slice(compressed);
}
Ok(())
}
/// decompresses a [`CompressedDataPage`] into `buffer`.
/// If the page is un-compressed, `buffer` is swapped instead.
/// Returns whether the page was decompressed.
pub fn decompress_buffer(
compressed_page: &mut CompressedPage,
buffer: &mut Vec<u8>,
) -> Result<bool> {
if compressed_page.compression() != Compression::Uncompressed {
// prepare the compression buffer
let read_size = compressed_page.uncompressed_size();
if read_size > buffer.capacity() {
// dealloc and ignore region, replacing it by a new region.
// This won't reallocate - it frees and calls `alloc_zeroed`
*buffer = vec![0; read_size];
} else if read_size > buffer.len() {
// fill what we need with zeros so that we can use them in `Read`.
// This won't reallocate
buffer.resize(read_size, 0);
} else {
buffer.truncate(read_size);
}
match compressed_page {
CompressedPage::Data(compressed_page) => match compressed_page.header() {
DataPageHeader::V1(_) => {
decompress_v1(&compressed_page.buffer, compressed_page.compression, buffer)?
}
DataPageHeader::V2(header) => decompress_v2(
&compressed_page.buffer,
header,
compressed_page.compression,
buffer,
)?,
},
CompressedPage::Dict(page) => decompress_v1(&page.buffer, page.compression(), buffer)?,
}
Ok(true)
} else {
// page.buffer is already decompressed => swap it with `buffer`, making `page.buffer` the
// decompression buffer and `buffer` the decompressed buffer
std::mem::swap(compressed_page.buffer(), buffer);
Ok(false)
}
}
fn create_page(compressed_page: CompressedPage, buffer: Vec<u8>) -> Page {
match compressed_page {
CompressedPage::Data(page) => Page::Data(DataPage::new_read(
page.header,
buffer,
page.descriptor,
page.selected_rows,
)),
CompressedPage::Dict(page) => Page::Dict(DictPage {
buffer,
num_values: page.num_values,
is_sorted: page.is_sorted,
}),
}
}
/// Decompresses the page, using `buffer` for decompression.
/// If `page.buffer.len() == 0`, there was no decompression and the buffer was moved.
/// Else, decompression took place.
pub fn decompress(mut compressed_page: CompressedPage, buffer: &mut Vec<u8>) -> Result<Page> {
decompress_buffer(&mut compressed_page, buffer)?;
Ok(create_page(compressed_page, std::mem::take(buffer)))
}
fn decompress_reuse<P: PageIterator>(
mut compressed_page: CompressedPage,
iterator: &mut P,
buffer: &mut Vec<u8>,
) -> Result<(Page, bool)> {
let was_decompressed = decompress_buffer(&mut compressed_page, buffer)?;
if was_decompressed {
iterator.swap_buffer(compressed_page.buffer())
};
let new_page = create_page(compressed_page, std::mem::take(buffer));
Ok((new_page, was_decompressed))
}
/// Decompressor that allows re-using the page buffer of [`PageIterator`].
/// # Implementation
/// The implementation depends on whether a page is compressed or not.
/// > `PageReader(a)`, `CompressedPage(b)`, `Decompressor(c)`, `DecompressedPage(d)`
/// ### un-compressed pages:
/// > page iter: `a` is swapped with `b`
/// > decompress iter: `b` is swapped with `d`, `b` is swapped with `a`
/// therefore:
/// * `PageReader` has its buffer back
/// * `Decompressor`'s buffer is un-used
/// * `DecompressedPage` has the same data as `CompressedPage` had
/// ### compressed pages:
/// > page iter: `a` is swapped with `b`
/// > decompress iter:
/// > * `b` is decompressed into `c`
/// > * `b` is swapped with `a`
/// > * `c` is moved to `d`
/// > * (next iteration): `d` is moved to `c`
/// therefore, while the page is available:
/// * `PageReader` has its buffer back
/// * `Decompressor`'s buffer empty
/// * `DecompressedPage` has the decompressed buffer
/// after the page is used:
/// * `PageReader` has its buffer back
/// * `Decompressor` has its buffer back
/// * `DecompressedPage` has an empty buffer
pub struct Decompressor<P: PageIterator> {
iter: P,
buffer: Vec<u8>,
current: Option<Page>,
was_decompressed: bool,
}
impl<P: PageIterator> Decompressor<P> {
/// Creates a new [`Decompressor`].
pub fn new(iter: P, buffer: Vec<u8>) -> Self {
Self {
iter,
buffer,
current: None,
was_decompressed: false,
}
}
/// Returns two buffers: the first buffer corresponds to the page buffer,
/// the second to the decompression buffer.
pub fn into_buffers(mut self) -> (Vec<u8>, Vec<u8>) {
let mut page_buffer = vec![];
self.iter.swap_buffer(&mut page_buffer);
(page_buffer, self.buffer)
}
}
impl<P: PageIterator> FallibleStreamingIterator for Decompressor<P> {
type Item = Page;
type Error = Error;
fn advance(&mut self) -> Result<()> {
if let Some(page) = self.current.as_mut() {
if self.was_decompressed {
self.buffer = std::mem::take(page.buffer());
} else {
self.iter.swap_buffer(page.buffer());
}
}
let next = self
.iter
.next()
.map(|x| {
x.and_then(|x| {
let (page, was_decompressed) =
decompress_reuse(x, &mut self.iter, &mut self.buffer)?;
self.was_decompressed = was_decompressed;
Ok(page)
})
})
.transpose()?;
self.current = next;
Ok(())
}
fn get(&self) -> Option<&Self::Item> {
self.current.as_ref()
}
}
type _Decompressor<I> = streaming_decompression::Decompressor<
CompressedPage,
Page,
fn(CompressedPage, &mut Vec<u8>) -> Result<Page>,
Error,
I,
>;
impl streaming_decompression::Compressed for CompressedPage {
#[inline]
fn is_compressed(&self) -> bool {
self.compression() != Compression::Uncompressed
}
}
impl streaming_decompression::Decompressed for Page {
#[inline]
fn buffer_mut(&mut self) -> &mut Vec<u8> {
self.buffer()
}
}
/// A [`FallibleStreamingIterator`] that decompresses [`CompressedPage`] into [`DataPage`].
/// # Implementation
/// This decompressor uses an internal [`Vec<u8>`] to perform decompressions which
/// is re-used across pages, so that a single allocation is required.
/// If the pages are not compressed, the internal buffer is not used.
pub struct BasicDecompressor<I: Iterator<Item = Result<CompressedPage>>> {
iter: _Decompressor<I>,
}
impl<I> BasicDecompressor<I>
where
I: Iterator<Item = Result<CompressedPage>>,
{
/// Returns a new [`BasicDecompressor`].
pub fn new(iter: I, buffer: Vec<u8>) -> Self {
Self {
iter: _Decompressor::new(iter, buffer, decompress),
}
}
/// Returns its internal buffer, consuming itself.
pub fn into_inner(self) -> Vec<u8> {
self.iter.into_inner()
}
}
impl<I> FallibleStreamingIterator for BasicDecompressor<I>
where
I: Iterator<Item = Result<CompressedPage>>,
{
type Item = Page;
type Error = Error;
fn advance(&mut self) -> Result<()> {
self.iter.advance()
}
fn get(&self) -> Option<&Self::Item> {
self.iter.get()
}
}