Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
reserve pushable containers in parquet extend_from_decoder
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Nov 17, 2022
1 parent 0ba4f8e commit 7a60529
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 1 deletion.
6 changes: 6 additions & 0 deletions src/io/parquet/read/deserialize/binary/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,9 @@ impl<O: Offset> Offsets<O> {
}

impl<O: Offset> Pushable<O> for Offsets<O> {
fn reserve(&mut self, additional: usize) {
self.0.reserve(additional)
}
#[inline]
fn len(&self) -> usize {
self.0.len() - 1
Expand Down Expand Up @@ -88,6 +91,9 @@ impl<O: Offset> Binary<O> {
}

impl<'a, O: Offset> Pushable<&'a [u8]> for Binary<O> {
fn reserve(&mut self, _additional: usize) {
// a no-op as it is not really clear how to resize a binary
}
#[inline]
fn len(&self) -> usize {
self.len()
Expand Down
3 changes: 3 additions & 0 deletions src/io/parquet/read/deserialize/fixed_size_binary/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,9 @@ impl FixedSizeBinary {
}

impl<'a> Pushable<&'a [u8]> for FixedSizeBinary {
fn reserve(&mut self, additional: usize) {
self.values.reserve(additional * self.size);
}
#[inline]
fn push(&mut self, value: &[u8]) {
debug_assert_eq!(value.len(), self.size);
Expand Down
10 changes: 9 additions & 1 deletion src/io/parquet/read/deserialize/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,14 +30,17 @@ pub fn not_implemented(page: &DataPage) -> Error {

/// A private trait representing structs that can receive elements.
pub(super) trait Pushable<T>: Sized {
//fn reserve(&mut self, additional: usize);
fn reserve(&mut self, additional: usize);
fn push(&mut self, value: T);
fn len(&self) -> usize;
fn push_null(&mut self);
fn extend_constant(&mut self, additional: usize, value: T);
}

impl Pushable<bool> for MutableBitmap {
fn reserve(&mut self, additional: usize) {
MutableBitmap::reserve(self, additional)
}
#[inline]
fn len(&self) -> usize {
self.len()
Expand All @@ -60,6 +63,9 @@ impl Pushable<bool> for MutableBitmap {
}

impl<A: Copy + Default> Pushable<A> for Vec<A> {
fn reserve(&mut self, additional: usize) {
Vec::reserve(self, additional)
}
#[inline]
fn len(&self) -> usize {
self.len()
Expand Down Expand Up @@ -304,6 +310,7 @@ pub(super) fn extend_from_decoder<'a, T: Default, P: Pushable<T>, I: Iterator<It
// consume `length` items
let iter = BitmapIter::new(values, offset, length);
let iter = Zip::new(iter, &mut values_iter);
pushable.reserve(length);

for item in iter {
if let Some(item) = item {
Expand All @@ -319,6 +326,7 @@ pub(super) fn extend_from_decoder<'a, T: Default, P: Pushable<T>, I: Iterator<It
FilteredHybridEncoded::Repeated { is_set, length } => {
validity.extend_constant(length, is_set);
if is_set {
pushable.reserve(length);
(0..length).for_each(|_| pushable.push(values_iter.next().unwrap()));
} else {
pushable.extend_constant(length, T::default());
Expand Down

0 comments on commit 7a60529

Please sign in to comment.