Skip to content

Commit de11d9c

Browse files
Dandandanclaude
andauthored
perf(parquet): Vectorize dict-index bounds check in RleDecoder::get_batch_with_dict (up to -7.9%) (#9746)
# Which issue does this PR close? - Close: #9747 # Rationale for this change Rewrite the code to generate more SIMD instructions / amortize loop/branching overhead. # What changes are included in this PR? One-file change in `parquet/src/encodings/rle.rs`: - u32 max-reduction bounds check over `CHUNK = 16` indices - `#[cold] #[inline(never)] fn oob` for the panic path - `get_unchecked` gather after the vectorised check | type | main | this-pr | Δ | | ---------------- | ---------------- | ---------------- | -----:| | UInt64Array | 60.6 ± 7.75 µs | 55.8 ± 0.58 µs | −7.9% | | Int64Array | 58.2 ± 0.77 µs | 55.3 ± 0.34 µs | −5.0% | | INT32 Decimal128 | 88.5 ± 2.07 µs | 84.3 ± 0.35 µs | −4.7% | | Int32Array | 49.6 ± 0.47 µs | 47.4 ± 0.50 µs | −4.4% | | UInt8Array | 53.0 ± 0.42 µs | 50.9 ± 0.60 µs | −4.0% | | Int16Array | 53.6 ± 0.49 µs | 51.7 ± 1.38 µs | −3.6% | | UInt16Array | 53.2 ± 0.34 µs | 51.7 ± 0.90 µs | −2.8% | | UInt32Array | 49.7 ± 4.23 µs | 48.4 ± 1.23 µs | −2.6% | | INT64 Decimal128 | 96.4 ± 2.10 µs | 94.9 ± 1.47 µs | −1.5% | | Int8Array | 53.2 ± 0.61 µs | 52.8 ± 5.21 µs | −0.7% | # Are these changes tested? Existing tests # Are there any user-facing changes? No — no API change; decoded output is identical, panic behaviour on out-of-bounds indices is preserved. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 9d3a4d9 commit de11d9c

1 file changed

Lines changed: 38 additions & 10 deletions

File tree

parquet/src/encodings/rle.rs

Lines changed: 38 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -484,7 +484,16 @@ impl RleDecoder {
484484
if self.rle_left > 0 {
485485
let num_values = cmp::min(max_values - values_read, self.rle_left as usize);
486486
let dict_idx = self.current_value.unwrap() as usize;
487-
let dict_value = dict[dict_idx].clone();
487+
let dict_value = dict
488+
.get(dict_idx)
489+
.ok_or_else(|| {
490+
general_err!(
491+
"dictionary index out of bounds: the len is {} but the index is {}",
492+
dict.len(),
493+
dict_idx
494+
)
495+
})?
496+
.clone();
488497

489498
buffer[values_read..values_read + num_values].fill(dict_value);
490499

@@ -514,16 +523,30 @@ impl RleDecoder {
514523
break;
515524
}
516525
{
526+
#[cold]
527+
#[inline(never)]
528+
fn oob(max_idx: u32, dict_len: usize) -> ParquetError {
529+
general_err!(
530+
"dictionary index out of bounds: the len is {} but the index is {}",
531+
dict_len,
532+
max_idx
533+
)
534+
}
535+
const CHUNK: usize = 16;
517536
let out = &mut buffer[values_read..values_read + num_values];
518537
let idx = &index_buf[..num_values];
519-
let mut out_chunks = out.chunks_exact_mut(8);
520-
let idx_chunks = idx.chunks_exact(8);
538+
let dict_len = dict.len();
539+
let mut out_chunks = out.chunks_exact_mut(CHUNK);
540+
let idx_chunks = idx.chunks_exact(CHUNK);
521541
for (out_chunk, idx_chunk) in out_chunks.by_ref().zip(idx_chunks) {
522-
let dict_len = dict.len();
523-
assert!(
524-
idx_chunk.iter().all(|&i| (i as usize) < dict_len),
525-
"dictionary index out of bounds"
526-
);
542+
// u32 max-reduction instead of `.all(|&i| ..)`: `.all`
543+
// short-circuits and blocks autovectorisation. Negative
544+
// i32 cast to u32 becomes a large value so the bounds
545+
// check still rejects it.
546+
let max_idx = idx_chunk.iter().fold(0u32, |acc, &i| acc.max(i as u32));
547+
if (max_idx as usize) >= dict_len {
548+
return Err(oob(max_idx, dict_len));
549+
}
527550
for (b, i) in out_chunk.iter_mut().zip(idx_chunk.iter()) {
528551
// SAFETY: all indices checked above to be in bounds
529552
b.clone_from(unsafe { dict.get_unchecked(*i as usize) });
@@ -532,9 +555,14 @@ impl RleDecoder {
532555
for (b, i) in out_chunks
533556
.into_remainder()
534557
.iter_mut()
535-
.zip(idx.chunks_exact(8).remainder().iter())
558+
.zip(idx.chunks_exact(CHUNK).remainder().iter())
536559
{
537-
b.clone_from(&dict[*i as usize]);
560+
let dict_idx = *i as usize;
561+
if dict_idx >= dict_len {
562+
return Err(oob(*i as u32, dict_len));
563+
}
564+
// SAFETY: bounds checked above
565+
b.clone_from(unsafe { dict.get_unchecked(dict_idx) });
538566
}
539567
}
540568
self.bit_packed_left -= num_values as u32;

0 commit comments

Comments
 (0)