viirya opened a new pull request, #12046:
URL: https://github.com/apache/arrow-rs/pull/12046
# Which issue does this PR close?
- Closes #12045
# Rationale for this change
`OffsetBuffer::extend_from_dictionary` turns dictionary keys into the values
and offsets of a `StringArray` / `BinaryArray`. It calls `extend_from_slice`
and `push` per value, so each value pays for a `memcpy` call with a data
dependent length and two capacity checks. Dictionary encoded strings are
usually short, so this overhead dominates. See the issue for details.
# What changes are included in this PR?
`extend_from_dictionary` now writes values and offsets through raw pointers
in a single pass:
- Offsets: `keys.len()` slots are reserved up front (as before), written
with `ptr::write`, and committed with one `set_len` at the end.
- Values: before each copy it checks that `self.values` has room for the
value plus 16 bytes of slack, and reserves when it doesn't. Since
`Vec::reserve` grows geometrically, this branch is rarely taken.
- A new `copy_value` helper copies values of up to 64 bytes as one to four
16-byte chunks with fixed-size `copy_nonoverlapping`. Those compile to plain
vector loads and stores. Anything else uses `memcpy`. The chunks are explicitly
unrolled because LLVM turns the equivalent loop back into a `memcpy` call.
Key and offset validation and the offset overflow check are the same as
before, all done before anything is written for that value. A dictionary offset
outside `dict_values` used to panic when slicing; it now returns an error.
Today's callers always pass a valid dictionary.
Safety argument:
- **Source reads:** a value is only copied in chunks if the source has
enough bytes from `start`: at least 16 for values of up to 16 bytes, and at
least `len.next_multiple_of(16)` for values of 17 to 64 bytes. Otherwise it
falls back to an exact-length `memcpy`. A zero-length value at the very end of
`dict_values` therefore copies 0 bytes. The `len > 16` condition on the
multi-chunk branch is what stops a zero-length value from going down that
branch.
- **Destination writes:** chunked copies write at most `len + 16` bytes, and
the capacity check guarantees `capacity >= len_so_far + value_len + 16`. When
the buffer has to grow, `set_len(len_so_far)` runs first so the reallocation
keeps every byte written so far, and the destination pointer is reloaded
afterwards.
- **`set_len`:** values and offsets are only ever written contiguously, so
the final `set_len` calls only cover initialized memory. The final `set_len`
also runs on the error paths, so an early return leaves the buffer consistent:
every value before the bad key is kept with its offset, just as before this
change.
- **Overflow:** the offset for each value is checked with `I::from_usize(len
+ value_len)` before the copy, as before. `len + value_len` can't wrap, since
both terms are at most `isize::MAX`.
# Are these changes tested?
Yes. Two new unit tests:
- `test_offset_buffer_extend_from_dictionary_lengths` covers value lengths
on both sides of each multiple of 16 up to 100 bytes, including zero-length and
short values at the end of the dictionary. It runs over several calls, so some
calls start with little spare capacity. The dictionary buffer is exactly sized,
so a read past its end is detected under AddressSanitizer.
- `test_offset_buffer_extend_from_dictionary_err` checks that a bad key or
bad dictionary offset leaves the earlier values in place and the buffer valid.
I ran the tests under AddressSanitizer (`RUSTFLAGS=-Zsanitizer=address cargo
+nightly test -p parquet --lib --target aarch64-apple-darwin offset_buffer`).
Each of these mutations makes them fail with a heap or global buffer overflow:
dropping the source length check on the single-chunk path, dropping the `len >
16` guard on the multi-chunk path, and dropping the 16 byte destination slack.
Benchmarks: Apple M4 Max, `main` and this PR run back to back. The upstream
`arrow_reader` numbers are the mean of two rounds.
`cargo bench -p parquet --features "arrow test_common experimental" --bench
arrow_reader -- 'dictionary encoded'` (dictionary values are 18-19 bytes)
| benchmark | main | this PR | change |
|---|---|---|---|
| StringArray, mandatory, no NULLs | 300.8 µs | 154.5 µs | −48.6% |
| StringArray, optional, no NULLs | 302.7 µs | 156.6 µs | −48.3% |
| StringArray, optional, half NULLs | 291.2 µs | 211.6 µs | −27.3% |
| BinaryArray, mandatory, no NULLs | 311.4 µs | 145.8 µs | −53.2% |
| BinaryArray, optional, no NULLs | 313.3 µs | 147.7 µs | −52.9% |
| BinaryArray, optional, half NULLs | 288.6 µs | 219.1 µs | −24.1% |
| BinaryArray, mandatory, no NULLs, large values (64 KiB) | 5.75 ms | 5.61
ms | −2.5% (noise) |
To check other value lengths, I temporarily changed the benchmark's
dictionary values to `format!("{x:0>N}")` (N bytes; this change isn't part of
the PR). StringArray results:
| value length | mandatory, no NULLs | optional, no NULLs | optional, half
NULLs |
|---|---|---|---|
| 8 B | 216.7 → 107.8 µs (−50%) | 218.3 → 110.1 µs (−50%) | 244.2 → 186.2 µs
(−24%) |
| 19 B | 254.9 → 149.3 µs (−41%) | 256.2 → 154.4 µs (−40%) | 263.4 → 214.3
µs (−19%) |
| 48 B | 343.0 → 205.3 µs (−40%) | 344.3 → 209.6 µs (−39%) | 296.4 → 223.4
µs (−25%) |
| 64 B | 249.3 → 233.0 µs (−7%) | 251.7 → 235.1 µs (−7%) | 257.8 → 252.1 µs
(−2%) |
| 100 B | 288.4 → 264.3 µs (−8%) | 284.6 → 268.9 µs (−6%) | 303.4 → 298.5 µs
(−2%) |
| 128 B | 288.6 → 277.5 µs (−4%) | 285.7 → 281.9 µs (−1%) | 323.5 → 313.8 µs
(−3%) |
# Are there any significant user-facing changes?
No.
This pull request and its description were written by Isaac.
--
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]
