viirya commented on code in PR #12046:
URL: https://github.com/apache/arrow-rs/pull/12046#discussion_r4500915503
##########
parquet/src/arrow/buffer/offset_buffer.rs:
##########
@@ -94,26 +94,62 @@ impl<I: OffsetSizeTrait> OffsetBuffer<I> {
dict_values: &[u8],
) -> Result<()> {
self.offsets.reserve(keys.len());
-
+ let offsets_start = self.offsets.len();
+ let offsets_dst = self.offsets.as_mut_ptr();
+
+ // Values and offsets are written through raw pointers, as this is
+ // much faster than `extend_from_slice` and `push` for short values.
+ // `len` bytes of `self.values` and `offsets_start + count` offsets
+ // are initialized, and the lengths are updated before returning.
+ let mut len = self.values.len();
+ let mut dst = self.values.as_mut_ptr();
+ let mut count = 0;
+ let mut result = Ok(());
for key in keys {
let index = key.as_usize();
if index + 1 >= dict_offsets.len() {
- return Err(general_err!(
+ result = Err(general_err!(
"dictionary key beyond bounds of dictionary: 0..{}",
dict_offsets.len().saturating_sub(1)
));
+ break;
}
let start_offset = dict_offsets[index].as_usize();
let end_offset = dict_offsets[index + 1].as_usize();
-
- // Dictionary values are verified when decoding dictionary page
- self.values
- .extend_from_slice(&dict_values[start_offset..end_offset]);
- let index_offset = I::from_usize(self.values.len())
- .ok_or_else(|| general_err!("index overflow decoding byte
array"))?;
- self.offsets.push(index_offset);
+ if start_offset > end_offset || end_offset > dict_values.len() {
+ result = Err(general_err!("invalid dictionary offsets"));
+ break;
+ }
+ let value_len = end_offset - start_offset;
+ let Some(offset) = I::from_usize(len + value_len) else {
+ result = Err(general_err!("index overflow decoding byte
array"));
+ break;
+ };
+
+ // `copy_value` may write up to `COPY_CHUNK` bytes past the value
+ if self.values.capacity() - len < value_len + COPY_CHUNK {
+ // SAFETY: the first `len` bytes are initialized
+ unsafe { self.values.set_len(len) };
+ self.values.reserve(value_len + COPY_CHUNK);
Review Comment:
Good catch, reproduced: 1,024 values of 16 bytes ended up in a 32 KiB values
buffer vs 16 KiB on main. Fixed in 1330c4226: it now reserves only `value_len`,
and `copy_value` takes the remaining capacity and only copies whole chunks when
they fit in both the source and the destination, otherwise it does an exact
copy. Empty values no longer allocate.
`test_offset_buffer_extend_from_dictionary_capacity` covers both, plus a short
value landing exactly at the end of the capacity, which is the case where the
destination bound matters. Performance is about the same as before the fix
(numbers in the description).
##########
parquet/src/arrow/buffer/offset_buffer.rs:
##########
@@ -94,26 +94,62 @@ impl<I: OffsetSizeTrait> OffsetBuffer<I> {
dict_values: &[u8],
) -> Result<()> {
self.offsets.reserve(keys.len());
-
+ let offsets_start = self.offsets.len();
+ let offsets_dst = self.offsets.as_mut_ptr();
+
+ // Values and offsets are written through raw pointers, as this is
+ // much faster than `extend_from_slice` and `push` for short values.
+ // `len` bytes of `self.values` and `offsets_start + count` offsets
+ // are initialized, and the lengths are updated before returning.
+ let mut len = self.values.len();
+ let mut dst = self.values.as_mut_ptr();
+ let mut count = 0;
+ let mut result = Ok(());
for key in keys {
let index = key.as_usize();
if index + 1 >= dict_offsets.len() {
- return Err(general_err!(
+ result = Err(general_err!(
"dictionary key beyond bounds of dictionary: 0..{}",
Review Comment:
Reproduced, and it was a bit worse than you describe: on the caught panic
the values had been committed up to the last reallocation but the offsets
hadn't, so the two were out of step (40 bytes with offsets `[0]` in your
probe). Now the check is `index >= dict_offsets.len().saturating_sub(1)`, so a
negative key takes the existing error path with the valid prefix committed, and
`-1` is in the error test. The reallocation branch also commits the offsets
together with the values now, so a panic in `reserve` leaves them consistent as
well.
##########
parquet/src/arrow/buffer/offset_buffer.rs:
##########
@@ -144,6 +180,51 @@ impl<I: OffsetSizeTrait> OffsetBuffer<I> {
}
}
+/// Number of bytes [`copy_value`] copies at once
+const COPY_CHUNK: usize = 16;
+
+/// Copies `src[start..start + len]` to `dst`
+///
+/// Values of up to `4 * COPY_CHUNK` bytes are copied in whole `COPY_CHUNK`
+/// byte chunks when `src` has enough bytes past `start`, as this is faster
+/// than `memcpy` for short values. This may write up to `COPY_CHUNK` bytes
+/// past `dst + len`.
+///
+/// # Safety
+///
+/// `start + len <= src.len()`, and `dst` must be valid for writes of
+/// `len + COPY_CHUNK` bytes
Review Comment:
Added, together with the new destination bound.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
