From c0e64cb49433d469018e5899e31cbc2c37f2db42 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 15:52:04 +0000 Subject: [PATCH 1/2] Reuse unpack helpers in BitPacked range take Follow-up to review comments on #10520: - Decode the whole-chunk middle of long runs with `BitUnpackedChunks` directly into the output instead of a hand-rolled chunk walk, and share the chunk-grouped index path with `take_primitive`. - Use `MaybeUninit` scratch instead of a zero-filled `[T; 1024]`. - Execute PiecewiseSequence lengths to `Columnar` so constant lengths are handled without materializing a lengths buffer. - Only fall back to the primitive path when every run is short, rather than when any run is short. Short runs among long ones are batched through the chunk-grouped index path. - Make `chunked_indices` use an uninitialized stack buffer rather than allocating, splitting a chunk across calls if it repeats more than 1024 indices. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RngpSDFfn7xtmA7UugiP3j --- .../fastlanes/src/bitpacking/compute/mod.rs | 58 ++-- .../fastlanes/src/bitpacking/compute/take.rs | 306 ++++++++++++------ 2 files changed, 249 insertions(+), 115 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs index fba1d8fd403..018aa0f9587 100644 --- a/encodings/fastlanes/src/bitpacking/compute/mod.rs +++ b/encodings/fastlanes/src/bitpacking/compute/mod.rs @@ -1,7 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; use std::mem::size_of; +use std::slice::from_raw_parts; use fastlanes::BitPacking; use vortex_array::dtype::NativePType; @@ -47,33 +49,35 @@ fn unpack_indices_into( } // TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. +/// Groups consecutive indices by 1024-element chunk without allocating. +/// +/// Repeated indices can select more than 1024 positions from one chunk, in which case +/// `chunk_fn` is called more than once for that chunk. fn chunked_indices( - mut indices: impl Iterator, + indices: impl Iterator, offset: usize, mut chunk_fn: F, ) { - let mut indices_within_chunk: Vec = Vec::with_capacity(1024); + let mut indices_within_chunk = [const { MaybeUninit::::uninit() }; 1024]; + let mut len = 0; + let mut current_chunk_idx = 0; + // SAFETY: Callers only read the first `len` entries, which have been written. + let initialized = + |buf: &[MaybeUninit], len| unsafe { from_raw_parts(buf.as_ptr().cast(), len) }; - let Some(first_idx) = indices.next() else { - return; - }; - - let mut current_chunk_idx = (first_idx + offset) / 1024; - indices_within_chunk.push((first_idx + offset) % 1024); for idx in indices { - let new_chunk_idx = (idx + offset) / 1024; - - if new_chunk_idx != current_chunk_idx { - chunk_fn(current_chunk_idx, &indices_within_chunk); - indices_within_chunk.clear(); + let chunk_idx = (idx + offset) / 1024; + if len != 0 && (chunk_idx != current_chunk_idx || len == indices_within_chunk.len()) { + chunk_fn(current_chunk_idx, initialized(&indices_within_chunk, len)); + len = 0; } - - current_chunk_idx = new_chunk_idx; - indices_within_chunk.push((idx + offset) % 1024); + current_chunk_idx = chunk_idx; + indices_within_chunk[len].write((idx + offset) % 1024); + len += 1; } - if !indices_within_chunk.is_empty() { - chunk_fn(current_chunk_idx, &indices_within_chunk); + if len != 0 { + chunk_fn(current_chunk_idx, initialized(&indices_within_chunk, len)); } } @@ -103,13 +107,23 @@ mod tests { #[test] fn chunk_indices_repeated() { - let mut called = false; + let mut calls = Vec::new(); chunked_indices([0; 1025].into_iter(), 0, |chunk_idx, idxs| { assert_eq!(chunk_idx, 0); - assert_eq!(idxs, [0; 1025]); - called = true; + assert!(idxs.iter().all(|&idx| idx == 0)); + calls.push(idxs.len()); + }); + // More than 1024 indices in one chunk are split across calls. + assert_eq!(calls, [1024, 1]); + } + + #[test] + fn chunk_indices_with_offset() { + let mut calls = Vec::new(); + chunked_indices([0, 5, 1000, 3000].into_iter(), 100, |chunk_idx, idxs| { + calls.push((chunk_idx, idxs.to_vec())); }); - assert!(called); + assert_eq!(calls, [(0, vec![100, 105]), (1, vec![76]), (3, vec![28])]); } #[rstest] diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 0321ede2462..1dddd75a545 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -8,6 +8,7 @@ use fastlanes::BitPacking; use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::ArrayView; +use vortex_array::Columnar; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; @@ -38,6 +39,9 @@ use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitWidthsView; use crate::FL_CHUNK_SIZE; +use crate::unpack_iter::BitPacked as BitPackedIter; +use crate::unpack_iter::BitPackingStrategy; +use crate::unpack_iter::BitUnpackedChunks; const FULL_ARRAY_DECODE_RATIO: usize = 8; @@ -109,21 +113,6 @@ fn contiguous_ranges( let Some(indices) = indices.as_opt::() else { return Ok(None); }; - let sparse_threshold = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |T| { - unpack_chunk_threshold::() as u64 - }); - // The primitive path batches short runs that share a packed chunk and selectively decodes - // isolated short runs. Reserve the range path for runs long enough to unpack a whole chunk. - if let Some(length) = indices - .lengths() - .as_opt::() - .and_then(|length| length.scalar().as_primitive().as_::()) - && length != 0 - && length <= sparse_threshold - { - return Ok(None); - } - if indices .multipliers() .as_opt::() @@ -132,29 +121,62 @@ fn contiguous_ranges( { return Ok(None); } + let sparse_threshold = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |T| { + unpack_chunk_threshold::() as u64 + }); - let lengths = indices - .lengths() - .cast(PType::U64.into())? - .execute::(ctx)?; - if lengths - .as_slice::() - .iter() - .any(|&length| length != 0 && length <= sparse_threshold) - { - return Ok(None); - } let starts = indices .starts() .cast(PType::U64.into())? .execute::(ctx)?; - let mut ranges = Vec::with_capacity(starts.len()); - let mut output_len = 0usize; - for (&start, &length) in starts - .as_slice::() - .iter() - .zip_eq(lengths.as_slice::()) - { + let starts = starts.as_slice::(); + // The primitive path batches short runs that share a packed chunk, so only take the range + // path when there is at least one run long enough to unpack a whole chunk. + match indices.lengths().clone().execute::(ctx)? { + Columnar::Constant(lengths) => { + let length = lengths + .scalar() + .as_primitive() + .as_::() + .vortex_expect("PiecewiseSequence lengths are non-nullable unsigned integers"); + if length <= sparse_threshold { + return Ok(None); + } + collect_ranges( + starts.iter().map(|&start| (start, length)), + array_len, + indices.len(), + ) + .map(Some) + } + Columnar::Canonical(lengths) => { + let lengths = lengths + .into_array() + .cast(PType::U64.into())? + .execute::(ctx)?; + let lengths = lengths.as_slice::(); + if lengths.iter().all(|&length| length <= sparse_threshold) { + return Ok(None); + } + collect_ranges( + starts.iter().copied().zip_eq(lengths.iter().copied()), + array_len, + indices.len(), + ) + .map(Some) + } + } +} + +/// Converts `(start, length)` pieces into validated, non-empty ranges. +fn collect_ranges( + pieces: impl ExactSizeIterator, + array_len: usize, + output_len: usize, +) -> VortexResult>> { + let mut ranges = Vec::with_capacity(pieces.len()); + let mut total_len = 0usize; + for (start, length) in pieces { if length == 0 { continue; } @@ -167,16 +189,16 @@ fn contiguous_ranges( end <= array_len, "PiecewiseSequenceArray range exceeds array length" ); - output_len = output_len + total_len = total_len .checked_add(length) .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; ranges.push(start..end); } - vortex_ensure_eq!(output_len, indices.len()); - Ok(Some(ranges)) + vortex_ensure_eq!(total_len, output_len); + Ok(ranges) } -fn take_ranges( +fn take_ranges>( array: ArrayView<'_, BitPacked>, bit_width: u8, ranges: &[Range], @@ -188,76 +210,121 @@ fn take_ranges( let chunk_len = 128 * bit_width / size_of::(); let offset = usize::from(array.offset()); let mut output = BufferMut::::with_capacity(output_len); - let mut unpacked = [T::zero(); FL_CHUNK_SIZE]; - let mut unpacked_chunk = None; + let mut scratch = [const { MaybeUninit::uninit() }; FL_CHUNK_SIZE]; // u64 has the largest sparse decode threshold. - let mut sparse_indices = [0usize; unpack_chunk_threshold::()]; - - // Vector gathers select runs of components. Keep those runs intact so neither the index - // buffer nor the chunk-grouping work scales with the number of selected components. - for range in ranges { - let mut start = range.start + offset; - let end = range.end + offset; - while start < end { - let chunk_idx = start / FL_CHUNK_SIZE; - let within_chunk = start % FL_CHUNK_SIZE; - let length = (end - start).min(FL_CHUNK_SIZE - within_chunk); - let packed_chunk = &packed[chunk_idx * chunk_len..][..chunk_len]; - if unpacked_chunk != Some(chunk_idx) && length <= unpack_chunk_threshold::() { - let indices = &mut sparse_indices[..length]; - for (index, position) in indices.iter_mut().zip(within_chunk..) { - *index = position; - } - unpack_indices_into(&mut output, bit_width, packed_chunk, indices); - } else { - if unpacked_chunk != Some(chunk_idx) { - // SAFETY: BitPacked validates its bit width and buffer size. The input holds - // one packed chunk and the initialized output has exactly 1024 elements. - unsafe { BitPacking::unchecked_unpack(bit_width, packed_chunk, &mut unpacked) }; - unpacked_chunk = Some(chunk_idx); - } - output.extend_from_slice(&unpacked[within_chunk..within_chunk + length]); - } - start += length; + let mut fragment_indices = [0usize; unpack_chunk_threshold::()]; + + // Vector gathers select runs of components. Long runs are decoded straight into the output so + // neither the index buffer nor the chunk-grouping work scales with the number of components. + // Short runs, and the few values a long run selects from a boundary chunk, are cheaper to + // extract individually. Consecutive short runs are expanded together so that those sharing a + // chunk are batched. + let threshold = unpack_chunk_threshold::(); + let mut short_runs_start = 0; + for (i, range) in ranges.iter().enumerate() { + if range.len() <= threshold { + continue; } + take_indices_into( + &mut output, + packed, + bit_width, + offset, + ranges[short_runs_start..i].iter().flat_map(Range::clone), + &mut scratch, + ); + short_runs_start = i + 1; + + let head_len = (FL_CHUNK_SIZE - (range.start + offset) % FL_CHUNK_SIZE) % FL_CHUNK_SIZE; + let tail_len = (range.end + offset) % FL_CHUNK_SIZE; + let middle_start = range.start + if head_len <= threshold { head_len } else { 0 }; + let middle_end = + (range.end - if tail_len <= threshold { tail_len } else { 0 }).max(middle_start); + + unpack_fragment_into( + &mut output, + packed, + bit_width, + range.start + offset..middle_start + offset, + &mut fragment_indices, + ); + let middle_len = middle_end - middle_start; + if middle_len != 0 { + let start = middle_start + offset; + let first_chunk = start / FL_CHUNK_SIZE; + let last_chunk = (middle_end + offset).div_ceil(FL_CHUNK_SIZE); + let mut chunks = BitUnpackedChunks::try_new_with_strategy( + BitPackingStrategy, + &packed[first_chunk * chunk_len..last_chunk * chunk_len], + bit_width, + start % FL_CHUNK_SIZE, + middle_len, + &mut scratch, + ) + .vortex_expect("validated range lies within the packed buffer"); + let output_start = output.len(); + chunks.decode_into(&mut output.spare_capacity_mut()[..middle_len]); + // SAFETY: `decode_into` initialized the `middle_len` values following `output_start`. + unsafe { output.set_len(output_start + middle_len) }; + } + unpack_fragment_into( + &mut output, + packed, + bit_width, + middle_end + offset..range.end + offset, + &mut fragment_indices, + ); } + take_indices_into( + &mut output, + packed, + bit_width, + offset, + ranges[short_runs_start..].iter().flat_map(Range::clone), + &mut scratch, + ); PrimitiveArray::new(output, validity) } -fn take_primitive( - array: ArrayView<'_, BitPacked>, - bit_width: u8, - indices: &PrimitiveArray, - taken_validity: Validity, - ctx: &mut ExecutionCtx, -) -> VortexResult { - if indices.is_empty() { - return Ok(PrimitiveArray::new(Buffer::::empty(), taken_validity)); +/// Appends the values at `positions`, which lie within one chunk, without decoding the whole chunk. +fn unpack_fragment_into( + output: &mut BufferMut, + packed: &[T], + bit_width: usize, + positions: Range, + indices: &mut [usize], +) { + if positions.is_empty() { + return; } - - let offset = array.offset() as usize; - let bit_width = bit_width as usize; - - let packed = array.packed_slice::(); - - // Group indices by 1024-element chunk, *without* allocating on the heap - let indices_iter = indices.as_slice::().iter().map(|i| { - i.to_usize() - .vortex_expect("index must be expressible as usize") - }); - - let mut output = BufferMut::::with_capacity(indices.len()); - let mut unpacked = [const { MaybeUninit::uninit() }; 1024]; let chunk_len = 128 * bit_width / size_of::(); + let packed = &packed[positions.start / FL_CHUNK_SIZE * chunk_len..][..chunk_len]; + let indices = &mut indices[..positions.len()]; + for (index, position) in indices.iter_mut().zip(positions) { + *index = position % FL_CHUNK_SIZE; + } + unpack_indices_into(output, bit_width, packed, indices); +} - chunked_indices(indices_iter, offset, |chunk_idx, indices_within_chunk| { +/// Appends the values at `indices` to `output`, decoding each visited chunk either fully or only +/// at the selected positions. +fn take_indices_into( + output: &mut BufferMut, + packed: &[T], + bit_width: usize, + offset: usize, + indices: impl Iterator, + unpacked: &mut [MaybeUninit; FL_CHUNK_SIZE], +) { + let chunk_len = 128 * bit_width / size_of::(); + chunked_indices(indices, offset, |chunk_idx, indices_within_chunk| { let packed = &packed[chunk_idx * chunk_len..][..chunk_len]; if indices_within_chunk.len() > unpack_chunk_threshold::() { // SAFETY: The validated bit width fits `T`. The source and destination contain one // complete FastLanes block. The call initializes every destination value. unsafe { - let dst: &mut [MaybeUninit] = &mut unpacked; + let dst: &mut [MaybeUninit] = unpacked; let dst: &mut [T] = std::mem::transmute(dst); BitPacking::unchecked_unpack(bit_width, packed, dst); } @@ -268,9 +335,35 @@ fn take_primitive( .map(|&index| unsafe { unpacked.get_unchecked(index).assume_init() }), ); } else { - unpack_indices_into(&mut output, bit_width, packed, indices_within_chunk); + unpack_indices_into(output, bit_width, packed, indices_within_chunk); } }); +} + +fn take_primitive( + array: ArrayView<'_, BitPacked>, + bit_width: u8, + indices: &PrimitiveArray, + taken_validity: Validity, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if indices.is_empty() { + return Ok(PrimitiveArray::new(Buffer::::empty(), taken_validity)); + } + + let mut output = BufferMut::::with_capacity(indices.len()); + let mut unpacked = [const { MaybeUninit::uninit() }; FL_CHUNK_SIZE]; + take_indices_into( + &mut output, + array.packed_slice::(), + usize::from(bit_width), + usize::from(array.offset()), + indices.as_slice::().iter().map(|i| { + i.to_usize() + .vortex_expect("index must be expressible as usize") + }), + &mut unpacked, + ); let unpatched_taken = if array.dtype().as_ptype().is_signed_int() { let primitive = PrimitiveArray::new(output, taken_validity); @@ -635,6 +728,33 @@ mod tests { Ok(()) } + #[rstest] + fn test_take_piecewise_mixed_run_lengths( + #[values(0, 137)] slice_offset: usize, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter((0u32..16384).map(|i| i % 127)).into_array(); + let packed = BitPackedData::encode(&values, 7, &mut ctx)?; + let packed = BitPacked::slice(packed.as_view(), slice_offset..16384, &mut ctx)?.unwrap(); + let values = values.slice(slice_offset..16384)?; + // Short runs sharing a chunk sit between long runs that cross chunk boundaries. The run at + // 1900 leaves a short tail in the chunk that the following short run also selects from. + let indices = PiecewiseSequenceArray::try_new( + buffer![10u32, 20, 1000, 30, 40, 4000, 50, 4090, 1900, 2070].into_array(), + buffer![3u32, 5, 300, 2, 0, 2048, 1, 6, 160, 3].into_array(), + ConstantArray::new(1u8, 10).into_array(), + 2528, + )? + .into_array(); + let expanded = indices + .clone() + .execute::(&mut ctx)? + .into_array(); + let actual = BitPacked::take(packed.as_::(), &indices, &mut ctx)?.unwrap(); + assert_arrays_eq!(actual, values.take(expanded)?, &mut ctx); + Ok(()) + } + #[rstest] #[case(PType::U8, 4, 16)] #[case(PType::U16, 9, 32)] From dc49835ac277c1fa900a847f56db7575bfc443e2 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 16:15:24 +0000 Subject: [PATCH 2/2] Remove stale chunked_indices duplication TODO The referenced bitpacking/kernels/mod.rs no longer exists and take.rs is the only user of chunked_indices. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RngpSDFfn7xtmA7UugiP3j --- encodings/fastlanes/src/bitpacking/compute/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs index 018aa0f9587..42ea6ea1736 100644 --- a/encodings/fastlanes/src/bitpacking/compute/mod.rs +++ b/encodings/fastlanes/src/bitpacking/compute/mod.rs @@ -48,7 +48,6 @@ fn unpack_indices_into( } } -// TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. /// Groups consecutive indices by 1024-element chunk without allocating. /// /// Repeated indices can select more than 1024 positions from one chunk, in which case