Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

parquet: improve BOOLEAN writing logic and report error on encoding fail #443

Merged
merged 5 commits into from
Jun 16, 2021
Merged

parquet: improve BOOLEAN writing logic and report error on encoding fail #443

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
db8af9b
improve BOOLEAN writing logic and report error on encoding fail
garyanaplan Jun 10, 2021
a4ec2d7
only manipulate the bit_writer for BOOLEAN data
garyanaplan Jun 10, 2021
eb11b2b
better isolation of changes
garyanaplan Jun 10, 2021
06d9a33
add test for boolean writer
garyanaplan Jun 14, 2021
da8c665
fix capacity calculation error in bool encoding
garyanaplan Jun 15, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
6 changes: 5 additions & 1 deletion parquet/src/data_type.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -662,7 +662,11 @@ pub(crate) mod private {
bit_writer: &mut BitWriter,
) -> Result<()> {
for value in values {
bit_writer.put_value(*value as u64, 1);
if !bit_writer.put_value(*value as u64, 1) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since put_value returns false if there isn't enough space, you might be able to avoid errors with something like:

for value in values {
  if !bit_writer.put_value(*value as u64, 1) {
    bit_writer.extend(256)
    bit_writer.put_value(*value as u64, 1)
  }
}

Rather than returning an error

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yea, we can either do this or make sure up front that there's enough capacity to write. One minor concern is putting the if branch inside the for loop might hurt the performance.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I found it hard to think of a good way to test this with the fix in place.

I preferred the "don't auto expand memory at the point of failure" approach because I'm fairly conservative and didn't want to make a change that was too wide in impact without a better understanding of the code. i.e.: my fix specifically targeted the error I reported and made it possible to detect in other locations.

I think a better fix would be to (somehow) pre-size the vector or avoid having to size a vector for all the bytes that could be written, but that would be a much bigger scope to the fix.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

leaving the code as is seems fine to me

return Err(ParquetError::EOF(
"unable to put boolean value".to_string(),
));
}
}
Ok(())
}
Expand Down
6 changes: 6 additions & 0 deletions parquet/src/encodings/encoding.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,7 @@ pub struct PlainEncoder<T: DataType> {
buffer: ByteBuffer,
bit_writer: BitWriter,
desc: ColumnDescPtr,
bw_bytes_written: usize,
_phantom: PhantomData<T>,
}

Expand All @@ -124,6 +125,7 @@ impl<T: DataType> PlainEncoder<T> {
buffer: byte_buffer,
bit_writer: BitWriter::new(256),
desc,
bw_bytes_written: 0,
_phantom: PhantomData,
}
}
Expand Down Expand Up @@ -153,7 +155,11 @@ impl<T: DataType> Encoder<T> for PlainEncoder<T> {

#[inline]
fn put(&mut self, values: &[T::T]) -> Result<()> {
if self.bw_bytes_written + values.len() >= self.bit_writer.capacity() {
self.bit_writer.extend(256);
}
T::T::encode(values, &mut self.buffer, &mut self.bit_writer)?;
self.bw_bytes_written += values.len();
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm going to add a comment myself! :-)

I just realised that I only want to do this checking if the encoding is for a Boolean, otherwise it's wasted work/memory. I'll think of the best way to achieve that.

Ok(())
}
}
Expand Down
14 changes: 14 additions & 0 deletions parquet/src/util/bit_util.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,20 @@ impl BitWriter {
}
}

/// Extend buffer size
#[inline]
pub fn extend(&mut self, increment: usize) {
self.max_bytes += increment;
let extra = vec![0; increment];
self.buffer.extend(extra);
}

/// Report buffer size
#[inline]
pub fn capacity(&mut self) -> usize {
self.max_bytes
}

/// Consumes and returns the current buffer.
#[inline]
pub fn consume(mut self) -> Vec<u8> {
Expand Down