sunchao commented on code in PR #12046:
URL: https://github.com/apache/arrow-rs/pull/12046#discussion_r4498303795
##########
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:
A malformed 32-bit RLE key can decode to `-1_i32`, which becomes
`usize::MAX` here. The `index + 1` check can panic before the shared length
updates. The panic already exists on the base, but the deferred writes change
what survives if it is caught. In a paired probe with keys `[0, 1, -1]` and two
40-byte values, the base retained 80 bytes and offsets `[0, 40, 80]`. This head
retained 40 bytes and offsets `[0]`. Checking `index >=
dict_offsets.len().saturating_sub(1)` would send this key through the existing
error path and commit the valid prefix together.
##########
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:
Can we reserve only `value_len` here? I wrote a dictionary-encoded UTF-8
column alternating two 16-byte strings and read it with
`ParquetRecordBatchReader` at the default batch size. Each 1,024-row batch held
the same 16,384 logical bytes, but `StringArray::values().capacity()` went from
16,384 bytes at the merge base to 32,768 bytes here. The unconditional slack
triggers geometric growth even though the 16-byte copy fits exactly, and
`Buffer::from_vec` retains the allocation in the returned array.
The chunked copy could check the destination space as well as the source
space, then fall back to an exact copy when the rounded write would exceed it.
With that reservation, an empty value can append its offset without allocating
or copying any bytes.
##########
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:
Please include the non-overlap precondition in `# Safety`. The destination
must be disjoint from the source bytes the helper reads, including any bytes
copied beyond the logical value. The current caller uses separate allocations
and satisfies that requirement.
--
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]
