From 5d7e3b82669d3f3758ac7faea4929358400f16ec Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 23 Oct 2026 14:26:08 +0100 Subject: [PATCH] Widen list take piece bounds to u64 before dispatch The piecewise list take compiled its gather kernels for every combination of piece start, piece length, list offset and output offset type (4 x 4 x 4 x 4). The kernels read each piece start and length only once per piece, so the take now widens them to u64 and dispatches only on the offset types. The widening does not copy when the pieces are already u64, which is the type that the list take produces. vortex-array release __text: 14,012,988 -> 12,798,204 bytes (-1.21 MB). Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5.5 --- vortex-array/src/arrays/list/compute/take.rs | 229 ++++++++----------- 1 file changed, 94 insertions(+), 135 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 591e1cf638d..a25268b4340 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools as _; +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -27,6 +29,7 @@ use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::dtype::IntegerPType; +use crate::dtype::PType; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; @@ -162,152 +165,54 @@ fn take_slices( let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); let output_len = indices_ref.len(); + // The piece starts and lengths are read once per piece, so they are widened to `u64` instead + // of compiling the gather kernels for each of their widths. + let starts = pieces_as_u64(&starts); let taken = match lengths { Columnar::Constant(lengths) => { let length = constant_unsigned_usize(&lengths); - take_slices_constant_start_dispatch( - array, - &starts, - length, - &offsets, - indices_ref, - output_len, - &data_validity, - )? + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take_slices_constant_length::( + array, + &starts, + length, + offsets.as_slice::(), + indices_ref, + output_len, + &data_validity, + ) + })? } Columnar::Canonical(lengths) => { - let lengths = lengths.into_primitive(); - take_slices_start_dispatch( - array, - &starts, - &lengths, - &offsets, - indices_ref, - output_len, - &data_validity, - )? + let lengths = pieces_as_u64(&lengths.into_primitive()); + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take_slices_typed::( + array, + &starts, + &lengths, + offsets.as_slice::(), + indices_ref, + output_len, + &data_validity, + ) + })? } }; Ok(Some(taken)) } -fn take_slices_constant_start_dispatch( - array: ArrayView<'_, List>, - starts: &PrimitiveArray, - length: usize, - offsets: &PrimitiveArray, - indices_ref: &ArrayRef, - output_len: usize, - data_validity: &Mask, -) -> VortexResult { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - take_slices_constant_offset_dispatch::( - array, - starts, - length, - offsets, - indices_ref, - output_len, - data_validity, - ) - }) -} - -fn take_slices_constant_offset_dispatch( - array: ArrayView<'_, List>, - starts: &PrimitiveArray, - length: usize, - offsets: &PrimitiveArray, - indices_ref: &ArrayRef, - output_len: usize, - data_validity: &Mask, -) -> VortexResult -where - S: UnsignedPType, -{ - match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { - take_slices_constant_length::( - array, - starts.as_slice::(), - length, - offsets.as_slice::(), - indices_ref, - output_len, - data_validity, - ) - }) -} - -fn take_slices_start_dispatch( - array: ArrayView<'_, List>, - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - offsets: &PrimitiveArray, - indices_ref: &ArrayRef, - output_len: usize, - data_validity: &Mask, -) -> VortexResult { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - take_slices_length_dispatch::( - array, - starts, - lengths, - offsets, - indices_ref, - output_len, - data_validity, - ) - }) -} - -fn take_slices_length_dispatch( - array: ArrayView<'_, List>, - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - offsets: &PrimitiveArray, - indices_ref: &ArrayRef, - output_len: usize, - data_validity: &Mask, -) -> VortexResult -where - S: UnsignedPType, -{ - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - take_slices_offset_dispatch::( - array, - starts, - lengths, - offsets, - indices_ref, - output_len, - data_validity, - ) - }) -} +/// Returns the unsigned piece values as `u64`. This does not copy when they are already `u64`. +fn pieces_as_u64(pieces: &PrimitiveArray) -> Buffer { + if pieces.ptype() == PType::U64 { + return pieces.to_buffer::(); + } -fn take_slices_offset_dispatch( - array: ArrayView<'_, List>, - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - offsets: &PrimitiveArray, - indices_ref: &ArrayRef, - output_len: usize, - data_validity: &Mask, -) -> VortexResult -where - S: UnsignedPType, - L: UnsignedPType, -{ - match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { - take_slices_typed::( - array, - starts.as_slice::(), - lengths.as_slice::(), - offsets.as_slice::(), - indices_ref, - output_len, - data_validity, - ) + match_each_unsigned_integer_ptype!(pieces.ptype(), |P| { + pieces + .as_slice::

() + .iter() + .map(|&piece| piece.as_()) + .collect() }) } @@ -850,10 +755,12 @@ mod test { use std::sync::LazyLock; use rstest::rstest; + use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; + use crate::ArrayRef; use crate::IntoArray as _; use crate::VortexSessionExecute; use crate::arrays::BoolArray; @@ -1113,6 +1020,58 @@ mod test { ); } + /// The take widens narrow piece starts and lengths to `u64`. + #[rstest] + #[case::u8_pieces( + buffer![1u8, 0].into_array(), + buffer![2u8, 1].into_array(), + buffer![2i32, 3, 4, 0, 1], + buffer![0u32, 3, 3, 5] + )] + #[case::mixed_pieces( + buffer![1u16, 0].into_array(), + buffer![2u32, 1].into_array(), + buffer![2i32, 3, 4, 0, 1], + buffer![0u32, 3, 3, 5] + )] + #[case::constant_length( + buffer![3u16, 0].into_array(), + ConstantArray::new(1u8, 2).into_array(), + buffer![5i32, 6, 0, 1], + buffer![0u32, 2, 4] + )] + fn piecewise_sequence_take_narrow_pieces( + #[case] starts: ArrayRef, + #[case] lengths: ArrayRef, + #[case] expected_elements: Buffer, + #[case] expected_offsets: Buffer, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let list = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + &mut ctx, + )? + .into_array(); + let output_len = expected_offsets.len() - 1; + let indices = PiecewiseSequenceArray::try_new( + starts, + lengths, + ConstantArray::new(1u64, 2).into_array(), + output_len, + )?; + let expected = ListArray::try_new( + expected_elements.into_array(), + expected_offsets.into_array(), + Validity::NonNullable, + &mut ctx, + )?; + + assert_arrays_eq!(list.take(indices.into_array())?, expected, &mut ctx); + Ok(()) + } + #[test] fn piecewise_sequence_take() { let mut ctx = SESSION.create_execution_ctx();