From fda912d95f4b8353349022b88e629f6044ba551b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 16 Apr 2026 20:10:12 +0200 Subject: [PATCH 1/8] Speed up ByteView dictionary decoder with chunks_exact gather MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the `extend(keys.iter().map(...))` loop in `ByteViewArrayDecoderDictionary::read` with a `chunks_exact(8)` loop that bulk-validates each chunk's keys, then uses `get_unchecked` gather plus raw-pointer writes. Matches the pattern in `RleDecoder::get_batch_with_dict`. Drops per-element bounds check, per-element `error.is_none()` branch, and `Vec::extend`'s per-push capacity check. Invalid keys now return an error eagerly via a cold helper instead of zero-filling and deferring. Dictionary-decode microbenchmarks (parquet/benches/arrow_reader.rs): BinaryView mandatory, no NULLs 102.91 µs -> 74.29 µs -27.8% BinaryView optional, no NULLs 104.63 µs -> 76.65 µs -26.9% BinaryView optional, half NULLs 143.25 µs -> 132.46 µs -7.3% StringView mandatory, no NULLs 105.98 µs -> 73.87 µs -28.8% StringView optional, no NULLs 104.62 µs -> 76.34 µs -27.4% StringView optional, half NULLs 141.86 µs -> 131.85 µs -7.1% Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/arrow/array_reader/byte_view_array.rs | 137 +++++++++++++----- 1 file changed, 100 insertions(+), 37 deletions(-) diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs b/parquet/src/arrow/array_reader/byte_view_array.rs index 1933654118f3..6f94a725d3f9 100644 --- a/parquet/src/arrow/array_reader/byte_view_array.rs +++ b/parquet/src/arrow/array_reader/byte_view_array.rs @@ -450,6 +450,39 @@ impl ByteViewArrayDecoderPlain { } } +/// Rewrite `view`'s buffer index by `base` when the dictionary buffers were +/// appended later than position 0 in the output buffer list. Inlined views +/// (length ≤ 12) carry their data in the high 96 bits and must be copied +/// verbatim. +#[inline(always)] +fn adjust_buffer_index(view: u128, base: u32) -> u128 { + let len = view as u32; + if len <= 12 { + view + } else { + let mut bv = ByteView::from(view); + bv.buffer_index += base; + bv.into() + } +} + +/// Slow-path error constructor for a chunk whose validity check failed. Kept +/// out of the hot loop so the fast path stays small. +#[cold] +#[inline(never)] +fn invalid_dict_key(chunk: &[i32], dict_len: usize) -> ParquetError { + let bad = chunk + .iter() + .copied() + .find(|&k| (k as usize) >= dict_len) + .unwrap_or(0); + general_err!( + "invalid key={} for dictionary of length {}", + bad, + dict_len + ) +} + pub struct ByteViewArrayDecoderDictionary { decoder: DictIndexDecoder, } @@ -500,52 +533,82 @@ impl ByteViewArrayDecoderDictionary { // then the base_buffer_idx is 5 - 2 = 3 let base_buffer_idx = output.buffers.len() as u32 - dict.buffers.len() as u32; - // Pre-reserve output capacity to avoid per-chunk reallocation in extend + // Pre-reserve output capacity so the gather loop can write through a raw + // pointer without `Vec::extend`'s per-element capacity checks. output.views.reserve(len); - let mut error = None; + let dict_views: &[u128] = dict.views.as_slice(); + let dict_len = dict_views.len(); + let read = self.decoder.read(len, |keys| { + // SAFETY: `output.views.reserve(len)` was called above and the + // outer loop ensures we never write more than `len` views. + let out_ptr = unsafe { output.views.as_mut_ptr().add(output.views.len()) }; + + // Process 8-key chunks with a bulk validity check that the compiler + // can autovectorise, then use `get_unchecked` in the gather loop. + // Mirrors the pattern in `RleDecoder::get_batch_with_dict`. + const CHUNK: usize = 8; + let mut chunks = keys.chunks_exact(CHUNK); + let mut written = 0usize; + if base_buffer_idx == 0 { - // the dictionary buffers are the last buffers in output, we can directly use the views - output - .views - .extend(keys.iter().map(|k| match dict.views.get(*k as usize) { - Some(&view) => view, - None => { - if error.is_none() { - error = Some(general_err!("invalid key={} for dictionary", *k)); - } - 0 + for chunk in chunks.by_ref() { + if !chunk.iter().all(|&k| (k as usize) < dict_len) { + return Err(invalid_dict_key(chunk, dict_len)); + } + for (i, &k) in chunk.iter().enumerate() { + // SAFETY: bounds checked above. + unsafe { + let view = *dict_views.get_unchecked(k as usize); + out_ptr.add(written + i).write(view); } - })); - Ok(()) + } + written += CHUNK; + } + for &k in chunks.remainder() { + let view = *dict_views + .get(k as usize) + .ok_or_else(|| general_err!("invalid key={k} for dictionary"))?; + // SAFETY: remainder writes stay within the reserved range. + unsafe { out_ptr.add(written).write(view) }; + written += 1; + } } else { - output - .views - .extend(keys.iter().map(|k| match dict.views.get(*k as usize) { - Some(&view) => { - let len = view as u32; - if len <= 12 { - view - } else { - let mut view = ByteView::from(view); - view.buffer_index += base_buffer_idx; - view.into() - } + for chunk in chunks.by_ref() { + if !chunk.iter().all(|&k| (k as usize) < dict_len) { + return Err(invalid_dict_key(chunk, dict_len)); + } + for (i, &k) in chunk.iter().enumerate() { + // SAFETY: bounds checked above. + unsafe { + let view = *dict_views.get_unchecked(k as usize); + out_ptr + .add(written + i) + .write(adjust_buffer_index(view, base_buffer_idx)); } - None => { - if error.is_none() { - error = Some(general_err!("invalid key={} for dictionary", *k)); - } - 0 - } - })); - Ok(()) + } + written += CHUNK; + } + for &k in chunks.remainder() { + let view = *dict_views + .get(k as usize) + .ok_or_else(|| general_err!("invalid key={k} for dictionary"))?; + // SAFETY: remainder writes stay within the reserved range. + unsafe { + out_ptr + .add(written) + .write(adjust_buffer_index(view, base_buffer_idx)) + }; + written += 1; + } } + + // SAFETY: we wrote exactly `written == keys.len()` new views. + debug_assert_eq!(written, keys.len()); + unsafe { output.views.set_len(output.views.len() + written) }; + Ok(()) })?; - if let Some(e) = error { - return Err(e); - } Ok(read) } From fe1728d50ca382bb522b999ad306041ad186e3e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 16 Apr 2026 20:20:54 +0200 Subject: [PATCH 2/8] Branchless adjust_buffer_index and SIMD-friendly bounds check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two small follow-ups to the chunked-gather rewrite, both driven by inspecting the aarch64 asm: 1) Rewrite `adjust_buffer_index` without an `if/else` so LLVM emits a `csel` in the hot chunked loop. Previously the main 8-key gather went through an out-of-line block with a conditional branch per view; now each view is 5 branchless instructions (ldp/cmp/csel/ add/stp). 2) Replace `chunk.iter().all(|&k| cond)` with a max-reduction over `u32` keys. `.all()` short-circuits, which blocks vectorisation — LLVM emitted 8 sequential `ldrsw+cmp+b.ls`. The max-reduction compiles on aarch64 NEON to: ldp q1, q0, [x1] ; one load, 8 keys umax.4s v2, v1, v0 ; pairwise lane max umaxv.4s s2, v2 ; horizontal reduce cmp w13, w22 ; one compare b.hs ; one branch The NEON registers are then reused for the gather (`fmov`/`mov.s v[i]`) so keys are loaded exactly once. Casting keys via `k as u32` correctly rejects any negative i32 (corrupt data) because a negative value becomes a large u32. Microbenchmark deltas over the previous commit (criterion, aarch64): BinaryView mandatory, no NULLs 74.29 µs -> 72.96 µs -1.8% BinaryView optional, no NULLs 76.65 µs -> 75.01 µs -2.1% StringView mandatory, no NULLs 73.87 µs -> 72.27 µs -2.2% StringView optional, no NULLs 76.34 µs -> 75.41 µs -1.2% Cumulative vs. main HEAD (89b1497484): BinaryView mandatory, no NULLs 102.91 µs -> 72.96 µs -29.2% BinaryView optional, no NULLs 104.63 µs -> 75.01 µs -28.4% BinaryView optional, half NULLs 143.25 µs -> 133.06 µs -7.4% StringView mandatory, no NULLs 105.98 µs -> 72.27 µs -30.7% StringView optional, no NULLs 104.62 µs -> 75.41 µs -29.2% StringView optional, half NULLs 141.86 µs -> 132.20 µs -6.8% Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/arrow/array_reader/byte_view_array.rs | 33 ++++++++++++------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs b/parquet/src/arrow/array_reader/byte_view_array.rs index 6f94a725d3f9..a95d009f548d 100644 --- a/parquet/src/arrow/array_reader/byte_view_array.rs +++ b/parquet/src/arrow/array_reader/byte_view_array.rs @@ -453,17 +453,13 @@ impl ByteViewArrayDecoderPlain { /// Rewrite `view`'s buffer index by `base` when the dictionary buffers were /// appended later than position 0 in the output buffer list. Inlined views /// (length ≤ 12) carry their data in the high 96 bits and must be copied -/// verbatim. +/// verbatim. Written branchlessly so LLVM emits `csel`/`cmov` inside the +/// hot chunked gather loop instead of a per-view conditional branch. #[inline(always)] fn adjust_buffer_index(view: u128, base: u32) -> u128 { - let len = view as u32; - if len <= 12 { - view - } else { - let mut bv = ByteView::from(view); - bv.buffer_index += base; - bv.into() - } + // View layout: bits [0..32] = len, [64..96] = buffer_index (long-view only). + let is_long = ((view as u32) > 12) as u128; + view.wrapping_add((is_long * base as u128) << 64) } /// Slow-path error constructor for a chunk whose validity check failed. Kept @@ -551,10 +547,21 @@ impl ByteViewArrayDecoderDictionary { const CHUNK: usize = 8; let mut chunks = keys.chunks_exact(CHUNK); let mut written = 0usize; + // Cast to u32 so that any negative i32 (corrupt data) compares as a + // very large value and fails the check. + let dict_len_u32 = dict_len as u32; if base_buffer_idx == 0 { for chunk in chunks.by_ref() { - if !chunk.iter().all(|&k| (k as usize) < dict_len) { + // Branchless max-reduction over 8 keys: LLVM emits a SIMD + // umax sequence on aarch64/x86_64 instead of the short- + // circuited `.all()` form which compiles to a chain of + // per-key `cmp + b.ls`. + let mut max_key = 0u32; + for &k in chunk { + max_key = max_key.max(k as u32); + } + if max_key >= dict_len_u32 { return Err(invalid_dict_key(chunk, dict_len)); } for (i, &k) in chunk.iter().enumerate() { @@ -576,7 +583,11 @@ impl ByteViewArrayDecoderDictionary { } } else { for chunk in chunks.by_ref() { - if !chunk.iter().all(|&k| (k as usize) < dict_len) { + let mut max_key = 0u32; + for &k in chunk { + max_key = max_key.max(k as u32); + } + if max_key >= dict_len_u32 { return Err(invalid_dict_key(chunk, dict_len)); } for (i, &k) in chunk.iter().enumerate() { From 90e095bcaa5b26dc37571fc2238ec2aee2316c6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 17 Apr 2026 06:22:09 +0200 Subject: [PATCH 3/8] Remove unused ByteView import Co-Authored-By: Claude Opus 4.7 (1M context) --- parquet/src/arrow/array_reader/byte_view_array.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs b/parquet/src/arrow/array_reader/byte_view_array.rs index a95d009f548d..c143fd35fdb4 100644 --- a/parquet/src/arrow/array_reader/byte_view_array.rs +++ b/parquet/src/arrow/array_reader/byte_view_array.rs @@ -30,7 +30,6 @@ use crate::schema::types::ColumnDescPtr; use crate::util::utf8::check_valid_utf8; use arrow_array::{ArrayRef, builder::make_view}; use arrow_buffer::Buffer; -use arrow_data::ByteView; use arrow_schema::DataType as ArrowType; use bytes::Bytes; use std::any::Any; From 0fa7d13ca74edb09b9f5e8ca23cf08fb43dcd5e7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 17 Apr 2026 06:39:23 +0200 Subject: [PATCH 4/8] Use fold for max-key reduction in dict gather Co-Authored-By: Claude Opus 4.7 (1M context) --- parquet/src/arrow/array_reader/byte_view_array.rs | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs b/parquet/src/arrow/array_reader/byte_view_array.rs index c143fd35fdb4..91970f79f172 100644 --- a/parquet/src/arrow/array_reader/byte_view_array.rs +++ b/parquet/src/arrow/array_reader/byte_view_array.rs @@ -556,10 +556,7 @@ impl ByteViewArrayDecoderDictionary { // umax sequence on aarch64/x86_64 instead of the short- // circuited `.all()` form which compiles to a chain of // per-key `cmp + b.ls`. - let mut max_key = 0u32; - for &k in chunk { - max_key = max_key.max(k as u32); - } + let max_key = chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); if max_key >= dict_len_u32 { return Err(invalid_dict_key(chunk, dict_len)); } @@ -582,10 +579,7 @@ impl ByteViewArrayDecoderDictionary { } } else { for chunk in chunks.by_ref() { - let mut max_key = 0u32; - for &k in chunk { - max_key = max_key.max(k as u32); - } + let max_key = chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); if max_key >= dict_len_u32 { return Err(invalid_dict_key(chunk, dict_len)); } From 7a27b7c233fb86e0d962fcf348b6f18cd8c7ff00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 23 Apr 2026 07:40:48 +0200 Subject: [PATCH 5/8] Use CHUNK=16 and MaybeUninit + get_unchecked_mut for ByteView dict gather Raises the chunk size from 8 to 16 to match #9746's finding for the RLE dict gather, and replaces the raw-pointer writes with a spare- capacity slice of MaybeUninit so the unsafe surface is confined to one slice index and one set_len. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/arrow/array_reader/byte_view_array.rs | 106 ++++++++---------- 1 file changed, 45 insertions(+), 61 deletions(-) diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs b/parquet/src/arrow/array_reader/byte_view_array.rs index 95a84d3829c3..8a60bd1eaacb 100644 --- a/parquet/src/arrow/array_reader/byte_view_array.rs +++ b/parquet/src/arrow/array_reader/byte_view_array.rs @@ -450,20 +450,15 @@ impl ByteViewArrayDecoderPlain { } } -/// Rewrite `view`'s buffer index by `base` when the dictionary buffers were -/// appended later than position 0 in the output buffer list. Inlined views -/// (length ≤ 12) carry their data in the high 96 bits and must be copied -/// verbatim. Written branchlessly so LLVM emits `csel`/`cmov` inside the -/// hot chunked gather loop instead of a per-view conditional branch. +/// Branchlessly add `base` to the buffer-index field of a long view +/// (inline views with length ≤ 12 carry data in the high bits and are +/// left untouched). #[inline(always)] fn adjust_buffer_index(view: u128, base: u32) -> u128 { - // View layout: bits [0..32] = len, [64..96] = buffer_index (long-view only). let is_long = ((view as u32) > 12) as u128; view.wrapping_add((is_long * base as u128) << 64) } -/// Slow-path error constructor for a chunk whose validity check failed. Kept -/// out of the hot loop so the fast path stays small. #[cold] #[inline(never)] fn invalid_dict_key(chunk: &[i32], dict_len: usize) -> ParquetError { @@ -529,91 +524,80 @@ impl ByteViewArrayDecoderDictionary { // then the base_buffer_idx is 5 - 2 = 3 let base_buffer_idx = output.buffers.len() as u32 - dict.buffers.len() as u32; - // Pre-reserve output capacity so the gather loop can write through a raw - // pointer without `Vec::extend`'s per-element capacity checks. + // Pre-reserve output capacity to avoid per-chunk reallocation in extend output.views.reserve(len); let dict_views: &[u128] = dict.views.as_slice(); let dict_len = dict_views.len(); + let mut out_offset = 0usize; let read = self.decoder.read(len, |keys| { - // SAFETY: `output.views.reserve(len)` was called above and the - // outer loop ensures we never write more than `len` views. - let out_ptr = unsafe { output.views.as_mut_ptr().add(output.views.len()) }; - - // Process 8-key chunks with a bulk validity check that the compiler - // can autovectorise, then use `get_unchecked` in the gather loop. - // Mirrors the pattern in `RleDecoder::get_batch_with_dict`. - const CHUNK: usize = 8; - let mut chunks = keys.chunks_exact(CHUNK); - let mut written = 0usize; - // Cast to u32 so that any negative i32 (corrupt data) compares as a - // very large value and fails the check. + // SAFETY: `reserve(len)` above + callbacks summing to `len` means + // spare capacity is always at least `keys.len()` from `out_offset`. + let out: &mut [std::mem::MaybeUninit] = unsafe { + output + .views + .spare_capacity_mut() + .get_unchecked_mut(out_offset..out_offset + keys.len()) + }; + out_offset += keys.len(); + + const CHUNK: usize = 16; + let mut out_chunks = out.chunks_exact_mut(CHUNK); + let mut key_chunks = keys.chunks_exact(CHUNK); + // Cast to u32 so negative i32 (corrupt data) compares as a large value. let dict_len_u32 = dict_len as u32; if base_buffer_idx == 0 { - for chunk in chunks.by_ref() { - // Branchless max-reduction over 8 keys: LLVM emits a SIMD - // umax sequence on aarch64/x86_64 instead of the short- - // circuited `.all()` form which compiles to a chain of - // per-key `cmp + b.ls`. - let max_key = chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); + for (out_chunk, key_chunk) in out_chunks.by_ref().zip(key_chunks.by_ref()) { + let max_key = key_chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); if max_key >= dict_len_u32 { - return Err(invalid_dict_key(chunk, dict_len)); + return Err(invalid_dict_key(key_chunk, dict_len)); } - for (i, &k) in chunk.iter().enumerate() { + for (dst, &k) in out_chunk.iter_mut().zip(key_chunk.iter()) { // SAFETY: bounds checked above. - unsafe { - let view = *dict_views.get_unchecked(k as usize); - out_ptr.add(written + i).write(view); - } + dst.write(unsafe { *dict_views.get_unchecked(k as usize) }); } - written += CHUNK; } - for &k in chunks.remainder() { + for (dst, &k) in out_chunks + .into_remainder() + .iter_mut() + .zip(key_chunks.remainder().iter()) + { let view = *dict_views .get(k as usize) .ok_or_else(|| general_err!("invalid key={k} for dictionary"))?; - // SAFETY: remainder writes stay within the reserved range. - unsafe { out_ptr.add(written).write(view) }; - written += 1; + dst.write(view); } } else { - for chunk in chunks.by_ref() { - let max_key = chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); + for (out_chunk, key_chunk) in out_chunks.by_ref().zip(key_chunks.by_ref()) { + let max_key = key_chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); if max_key >= dict_len_u32 { - return Err(invalid_dict_key(chunk, dict_len)); + return Err(invalid_dict_key(key_chunk, dict_len)); } - for (i, &k) in chunk.iter().enumerate() { + for (dst, &k) in out_chunk.iter_mut().zip(key_chunk.iter()) { // SAFETY: bounds checked above. - unsafe { - let view = *dict_views.get_unchecked(k as usize); - out_ptr - .add(written + i) - .write(adjust_buffer_index(view, base_buffer_idx)); - } + let view = unsafe { *dict_views.get_unchecked(k as usize) }; + dst.write(adjust_buffer_index(view, base_buffer_idx)); } - written += CHUNK; } - for &k in chunks.remainder() { + for (dst, &k) in out_chunks + .into_remainder() + .iter_mut() + .zip(key_chunks.remainder().iter()) + { let view = *dict_views .get(k as usize) .ok_or_else(|| general_err!("invalid key={k} for dictionary"))?; - // SAFETY: remainder writes stay within the reserved range. - unsafe { - out_ptr - .add(written) - .write(adjust_buffer_index(view, base_buffer_idx)) - }; - written += 1; + dst.write(adjust_buffer_index(view, base_buffer_idx)); } } - // SAFETY: we wrote exactly `written == keys.len()` new views. - debug_assert_eq!(written, keys.len()); - unsafe { output.views.set_len(output.views.len() + written) }; Ok(()) })?; + // SAFETY: decoder.read wrote exactly `read` views via dst.write. + debug_assert_eq!(out_offset, read); + unsafe { output.views.set_len(output.views.len() + read) }; Ok(read) } From 2d4e134600beb23cfe6db2b9f5d6ea7175e916d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 23 Apr 2026 08:10:24 +0200 Subject: [PATCH 6/8] Fuse RLE decode with view gather via MaybeUninit spare capacity Change `RleDecoder::get_batch_with_dict` (pub(crate)) to take `&mut [MaybeUninit]` so callers can gather directly into `Vec::spare_capacity_mut()` without zero-initialising first. In `ByteViewArrayDecoderDictionary::read`, the common `base_buffer_idx == 0` case now calls a new `DictIndexDecoder::read_with_dict` that delegates to `get_batch_with_dict`, skipping the intermediate index-buffer pass. The `base_buffer_idx != 0` branch keeps the chunked-gather fallback. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/arrow/array_reader/byte_view_array.rs | 78 +++++++++---------- parquet/src/arrow/decoder/dictionary_index.rs | 72 +++++++++++++++++ parquet/src/encodings/decoding.rs | 12 ++- parquet/src/encodings/rle.rs | 57 +++++++------- 4 files changed, 152 insertions(+), 67 deletions(-) diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs b/parquet/src/arrow/array_reader/byte_view_array.rs index 8a60bd1eaacb..451d269a060f 100644 --- a/parquet/src/arrow/array_reader/byte_view_array.rs +++ b/parquet/src/arrow/array_reader/byte_view_array.rs @@ -529,6 +529,27 @@ impl ByteViewArrayDecoderDictionary { let dict_views: &[u128] = dict.views.as_slice(); let dict_len = dict_views.len(); + + if base_buffer_idx == 0 { + // Fused path: RLE decode + view gather in one pass via + // `RleDecoder::get_batch_with_dict`, writing directly into spare + // capacity (no zero-init) and skipping the intermediate index + // buffer for RLE runs. + let base = output.views.len(); + // SAFETY: `reserve(len)` above ensures the spare slice is at + // least `len` long. + let spare = unsafe { + output + .views + .spare_capacity_mut() + .get_unchecked_mut(..len) + }; + let read = self.decoder.read_with_dict(len, dict_views, spare)?; + // SAFETY: `read_with_dict` wrote exactly `read` views. + unsafe { output.views.set_len(base + read) }; + return Ok(read); + } + let mut out_offset = 0usize; let read = self.decoder.read(len, |keys| { @@ -548,50 +569,27 @@ impl ByteViewArrayDecoderDictionary { // Cast to u32 so negative i32 (corrupt data) compares as a large value. let dict_len_u32 = dict_len as u32; - if base_buffer_idx == 0 { - for (out_chunk, key_chunk) in out_chunks.by_ref().zip(key_chunks.by_ref()) { - let max_key = key_chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); - if max_key >= dict_len_u32 { - return Err(invalid_dict_key(key_chunk, dict_len)); - } - for (dst, &k) in out_chunk.iter_mut().zip(key_chunk.iter()) { - // SAFETY: bounds checked above. - dst.write(unsafe { *dict_views.get_unchecked(k as usize) }); - } - } - for (dst, &k) in out_chunks - .into_remainder() - .iter_mut() - .zip(key_chunks.remainder().iter()) - { - let view = *dict_views - .get(k as usize) - .ok_or_else(|| general_err!("invalid key={k} for dictionary"))?; - dst.write(view); + for (out_chunk, key_chunk) in out_chunks.by_ref().zip(key_chunks.by_ref()) { + let max_key = key_chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); + if max_key >= dict_len_u32 { + return Err(invalid_dict_key(key_chunk, dict_len)); } - } else { - for (out_chunk, key_chunk) in out_chunks.by_ref().zip(key_chunks.by_ref()) { - let max_key = key_chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); - if max_key >= dict_len_u32 { - return Err(invalid_dict_key(key_chunk, dict_len)); - } - for (dst, &k) in out_chunk.iter_mut().zip(key_chunk.iter()) { - // SAFETY: bounds checked above. - let view = unsafe { *dict_views.get_unchecked(k as usize) }; - dst.write(adjust_buffer_index(view, base_buffer_idx)); - } - } - for (dst, &k) in out_chunks - .into_remainder() - .iter_mut() - .zip(key_chunks.remainder().iter()) - { - let view = *dict_views - .get(k as usize) - .ok_or_else(|| general_err!("invalid key={k} for dictionary"))?; + for (dst, &k) in out_chunk.iter_mut().zip(key_chunk.iter()) { + // SAFETY: bounds checked above. + let view = unsafe { *dict_views.get_unchecked(k as usize) }; dst.write(adjust_buffer_index(view, base_buffer_idx)); } } + for (dst, &k) in out_chunks + .into_remainder() + .iter_mut() + .zip(key_chunks.remainder().iter()) + { + let view = *dict_views + .get(k as usize) + .ok_or_else(|| general_err!("invalid key={k} for dictionary"))?; + dst.write(adjust_buffer_index(view, base_buffer_idx)); + } Ok(()) })?; diff --git a/parquet/src/arrow/decoder/dictionary_index.rs b/parquet/src/arrow/decoder/dictionary_index.rs index 7a4b77f89d59..987793ae1ace 100644 --- a/parquet/src/arrow/decoder/dictionary_index.rs +++ b/parquet/src/arrow/decoder/dictionary_index.rs @@ -91,6 +91,78 @@ impl DictIndexDecoder { Ok(values_read) } + /// Read up to `len` values directly into `output`, gathering through `dict` + /// in a single pass. Avoids expanding RLE runs into the index buffer and + /// writes through a `MaybeUninit` slice so the caller does not need to + /// zero-initialise output capacity. + pub fn read_with_dict( + &mut self, + len: usize, + dict: &[T], + output: &mut [std::mem::MaybeUninit], + ) -> Result { + use crate::errors::ParquetError; + let total_to_read = len.min(self.max_remaining_values); + let mut values_read = 0; + + // Drain any leftover indices buffered from a prior `read` call before + // switching to the direct-gather path. Uses the same CHUNK=16 + + // max-reduction pattern as `RleDecoder::get_batch_with_dict`. + let leftover = self.index_buf_len - self.index_offset; + if leftover > 0 { + let n = leftover.min(total_to_read); + let keys = &self.index_buf[self.index_offset..self.index_offset + n]; + let out = &mut output[..n]; + let dict_len = dict.len(); + let dict_len_u32 = dict_len as u32; + + const CHUNK: usize = 16; + let mut out_chunks = out.chunks_exact_mut(CHUNK); + let mut key_chunks = keys.chunks_exact(CHUNK); + for (out_chunk, key_chunk) in out_chunks.by_ref().zip(key_chunks.by_ref()) { + let max_key = key_chunk.iter().fold(0u32, |acc, &k| acc.max(k as u32)); + if max_key >= dict_len_u32 { + return Err(ParquetError::General(format!( + "dictionary index out of bounds: the len is {dict_len} but the index is {max_key}" + ))); + } + for (dst, &k) in out_chunk.iter_mut().zip(key_chunk.iter()) { + // SAFETY: bounds checked above. + dst.write(unsafe { dict.get_unchecked(k as usize) }.clone()); + } + } + for (dst, &k) in out_chunks + .into_remainder() + .iter_mut() + .zip(key_chunks.remainder().iter()) + { + let idx = k as usize; + if idx >= dict_len { + return Err(ParquetError::General(format!( + "dictionary index out of bounds: the len is {dict_len} but the index is {idx}" + ))); + } + // SAFETY: bounds checked above. + dst.write(unsafe { dict.get_unchecked(idx) }.clone()); + } + + self.index_offset += n; + values_read += n; + } + + if values_read < total_to_read { + let got = self.decoder.get_batch_with_dict( + dict, + &mut output[values_read..total_to_read], + total_to_read - values_read, + )?; + values_read += got; + } + + self.max_remaining_values -= values_read; + Ok(values_read) + } + /// Skip up to `to_skip` values, returning the number of values skipped pub fn skip(&mut self, to_skip: usize) -> Result { let to_skip = to_skip.min(self.max_remaining_values); diff --git a/parquet/src/encodings/decoding.rs b/parquet/src/encodings/decoding.rs index f7f4d9be4726..0f7e8101aa85 100644 --- a/parquet/src/encodings/decoding.rs +++ b/parquet/src/encodings/decoding.rs @@ -405,7 +405,17 @@ impl Decoder for DictDecoder { let rle = self.rle_decoder.as_mut().unwrap(); let num_values = cmp::min(buffer.len(), self.num_values); - rle.get_batch_with_dict(&self.dictionary[..], buffer, num_values) + // SAFETY: reinterpreting `&mut [T]` as `&mut [MaybeUninit]` is sound + // because every initialised `T` is a valid `MaybeUninit`; `get_batch_with_dict` + // only writes through the slice, and we do not read through the original + // reference after this call. + let uninit: &mut [std::mem::MaybeUninit] = unsafe { + std::slice::from_raw_parts_mut( + buffer.as_mut_ptr().cast::>(), + buffer.len(), + ) + }; + rle.get_batch_with_dict(&self.dictionary[..], uninit, num_values) } /// Number of values left in this decoder stream diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs index 806b41a353b4..1aa8f833228b 100644 --- a/parquet/src/encodings/rle.rs +++ b/parquet/src/encodings/rle.rs @@ -477,11 +477,11 @@ impl RleDecoder { pub fn get_batch_with_dict( &mut self, dict: &[T], - buffer: &mut [T], + buffer: &mut [std::mem::MaybeUninit], max_values: usize, ) -> Result where - T: Default + Clone, + T: Clone, { debug_assert!(buffer.len() >= max_values); @@ -492,18 +492,17 @@ impl RleDecoder { if self.rle_left > 0 { let num_values = cmp::min(max_values - values_read, self.rle_left as usize); let dict_idx = self.current_value.unwrap() as usize; - let dict_value = dict - .get(dict_idx) - .ok_or_else(|| { - general_err!( - "dictionary index out of bounds: the len is {} but the index is {}", - dict.len(), - dict_idx - ) - })? - .clone(); - - buffer[values_read..values_read + num_values].fill(dict_value); + let dict_value = dict.get(dict_idx).ok_or_else(|| { + general_err!( + "dictionary index out of bounds: the len is {} but the index is {}", + dict.len(), + dict_idx + ) + })?; + + for slot in &mut buffer[values_read..values_read + num_values] { + slot.write(dict_value.clone()); + } self.rle_left -= num_values as u32; values_read += num_values; @@ -557,7 +556,7 @@ impl RleDecoder { } for (b, i) in out_chunk.iter_mut().zip(idx_chunk.iter()) { // SAFETY: all indices checked above to be in bounds - b.clone_from(unsafe { dict.get_unchecked(*i as usize) }); + b.write(unsafe { dict.get_unchecked(*i as usize) }.clone()); } } for (b, i) in out_chunks @@ -570,7 +569,7 @@ impl RleDecoder { return Err(oob(*i as u32, dict_len)); } // SAFETY: bounds checked above - b.clone_from(unsafe { dict.get_unchecked(dict_idx) }); + b.write(unsafe { dict.get_unchecked(dict_idx) }.clone()); } } self.bit_packed_left -= num_values as u32; @@ -624,6 +623,16 @@ mod tests { use crate::util::bit_util::ceil; use rand::{self, Rng, SeedableRng, distr::StandardUniform, rng}; + use std::mem::MaybeUninit; + + /// Reinterpret an initialised slice as a `MaybeUninit` slice for calls to + /// `get_batch_with_dict`. Sound because every `T` is a valid `MaybeUninit` + /// and the callee only writes. + fn as_uninit(s: &mut [T]) -> &mut [MaybeUninit] { + unsafe { + std::slice::from_raw_parts_mut(s.as_mut_ptr().cast::>(), s.len()) + } + } const MAX_WIDTH: usize = 32; @@ -772,7 +781,7 @@ mod tests { decoder.set_data(data.into()).unwrap(); let mut buffer = vec![0; 12]; let expected = vec![10, 10, 10, 20, 20, 20, 20, 30, 30, 30, 30, 30]; - let result = decoder.get_batch_with_dict::(&dict, &mut buffer, 12); + let result = decoder.get_batch_with_dict::(&dict, as_uninit(&mut buffer), 12); assert!(result.is_ok()); assert_eq!(buffer, expected); @@ -788,7 +797,7 @@ mod tests { "ddd", "eee", "fff", "ddd", "eee", "fff", "ddd", "eee", "fff", "eee", "fff", "fff", ]; let result = - decoder.get_batch_with_dict::<&str>(dict.as_slice(), buffer.as_mut_slice(), 12); + decoder.get_batch_with_dict::<&str>(dict.as_slice(), as_uninit(&mut buffer), 12); assert!(result.is_ok()); assert_eq!(buffer, expected); } @@ -806,7 +815,7 @@ mod tests { let skipped = decoder.skip(2).expect("skipping two values"); assert_eq!(skipped, 2); let remainder = decoder - .get_batch_with_dict::(&dict, &mut buffer, 10) + .get_batch_with_dict::(&dict, as_uninit(&mut buffer), 10) .expect("getting remainder"); assert_eq!(remainder, 10); assert_eq!(buffer, expected); @@ -823,11 +832,7 @@ mod tests { let skipped = decoder.skip(4).expect("skipping four values"); assert_eq!(skipped, 4); let remainder = decoder - .get_batch_with_dict::<&str>( - dict.as_slice(), - buffer.as_mut_slice(), - BIT_PACK_GROUP_SIZE, - ) + .get_batch_with_dict::<&str>(dict.as_slice(), as_uninit(&mut buffer), BIT_PACK_GROUP_SIZE) .expect("getting remainder"); assert_eq!(remainder, BIT_PACK_GROUP_SIZE); assert_eq!(buffer, expected); @@ -986,7 +991,7 @@ mod tests { let dict: Vec = (0..256).collect(); let mut output = vec![0_u16; 100]; let read = decoder - .get_batch_with_dict(&dict, &mut output, 100) + .get_batch_with_dict(&dict, as_uninit(&mut output), 100) .unwrap(); assert_eq!(read, 20); @@ -1056,7 +1061,7 @@ mod tests { decoder.set_data(buffer).unwrap(); let r = decoder - .get_batch_with_dict(&[0, 23], &mut decoded, num_values) + .get_batch_with_dict(&[0, 23], as_uninit(&mut decoded), num_values) .unwrap(); assert_eq!(r, num_values); assert_eq!(vec![23; num_values], decoded); From 7e166a560ec976d7a216296ddf923a3ee95fd17c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 23 Apr 2026 09:12:55 +0200 Subject: [PATCH 7/8] Move index_buf allocation into bit-packed branch The 1024-entry scratch is only used when decoding bit-packed runs. Moving the `get_or_insert_with` call inside the `else if self.bit_packed_left > 0` branch means RLE-only streams skip the allocation entirely, and the `Option` discriminant check is paid only where the buffer is actually read. Relies on Rust's disjoint field borrows to hold both `self.bit_reader` and `self.index_buf` mutably at once. Co-Authored-By: Claude Opus 4.7 (1M context) --- parquet/src/encodings/rle.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs index 1aa8f833228b..069c5fc85ad6 100644 --- a/parquet/src/encodings/rle.rs +++ b/parquet/src/encodings/rle.rs @@ -487,8 +487,6 @@ impl RleDecoder { let mut values_read = 0; while values_read < max_values { - let index_buf = self.index_buf.get_or_insert_with(|| Box::new([0; 1024])); - if self.rle_left > 0 { let num_values = cmp::min(max_values - values_read, self.rle_left as usize); let dict_idx = self.current_value.unwrap() as usize; @@ -511,6 +509,7 @@ impl RleDecoder { .bit_reader .as_mut() .ok_or_else(|| general_err!("bit_reader should be set"))?; + let index_buf = self.index_buf.get_or_insert_with(|| Box::new([0; 1024])); loop { let to_read = index_buf From e344717d6a3bee62efcfb621391efb12182ef8c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 23 Apr 2026 09:27:17 +0200 Subject: [PATCH 8/8] cargo fmt Co-Authored-By: Claude Opus 4.7 (1M context) --- parquet/src/arrow/array_reader/byte_view_array.rs | 13 ++----------- parquet/src/encodings/rle.rs | 10 ++++++---- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs b/parquet/src/arrow/array_reader/byte_view_array.rs index 451d269a060f..801d62edd06b 100644 --- a/parquet/src/arrow/array_reader/byte_view_array.rs +++ b/parquet/src/arrow/array_reader/byte_view_array.rs @@ -467,11 +467,7 @@ fn invalid_dict_key(chunk: &[i32], dict_len: usize) -> ParquetError { .copied() .find(|&k| (k as usize) >= dict_len) .unwrap_or(0); - general_err!( - "invalid key={} for dictionary of length {}", - bad, - dict_len - ) + general_err!("invalid key={} for dictionary of length {}", bad, dict_len) } pub struct ByteViewArrayDecoderDictionary { @@ -538,12 +534,7 @@ impl ByteViewArrayDecoderDictionary { let base = output.views.len(); // SAFETY: `reserve(len)` above ensures the spare slice is at // least `len` long. - let spare = unsafe { - output - .views - .spare_capacity_mut() - .get_unchecked_mut(..len) - }; + let spare = unsafe { output.views.spare_capacity_mut().get_unchecked_mut(..len) }; let read = self.decoder.read_with_dict(len, dict_views, spare)?; // SAFETY: `read_with_dict` wrote exactly `read` views. unsafe { output.views.set_len(base + read) }; diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs index 069c5fc85ad6..42c53b644719 100644 --- a/parquet/src/encodings/rle.rs +++ b/parquet/src/encodings/rle.rs @@ -628,9 +628,7 @@ mod tests { /// `get_batch_with_dict`. Sound because every `T` is a valid `MaybeUninit` /// and the callee only writes. fn as_uninit(s: &mut [T]) -> &mut [MaybeUninit] { - unsafe { - std::slice::from_raw_parts_mut(s.as_mut_ptr().cast::>(), s.len()) - } + unsafe { std::slice::from_raw_parts_mut(s.as_mut_ptr().cast::>(), s.len()) } } const MAX_WIDTH: usize = 32; @@ -831,7 +829,11 @@ mod tests { let skipped = decoder.skip(4).expect("skipping four values"); assert_eq!(skipped, 4); let remainder = decoder - .get_batch_with_dict::<&str>(dict.as_slice(), as_uninit(&mut buffer), BIT_PACK_GROUP_SIZE) + .get_batch_with_dict::<&str>( + dict.as_slice(), + as_uninit(&mut buffer), + BIT_PACK_GROUP_SIZE, + ) .expect("getting remainder"); assert_eq!(remainder, BIT_PACK_GROUP_SIZE); assert_eq!(buffer, expected);