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.

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

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

Merged
merged 5 commits into from
Jun 16, 2021
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
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
9 changes: 8 additions & 1 deletion parquet/src/data_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -661,8 +661,15 @@ pub(crate) mod private {
_: &mut W,
bit_writer: &mut BitWriter,
) -> Result<()> {
if bit_writer.bytes_written() + values.len() >= bit_writer.capacity() {
Copy link
Member

Choose a reason for hiding this comment

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

Seems here values.len is the number of bits to be written? should we use values.len() / 8?

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 think this calculation is entirely in terms of bytes, so units should all be correct as is.

Copy link
Member

Choose a reason for hiding this comment

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

Hmm I'm sorry but can you elaborate why the unit of values is also byte?

bit_writer.extend(256);
}
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
14 changes: 14 additions & 0 deletions parquet/src/util/bit_util.rs
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