From 9c1de8d1844bf918f180fa1d03ef79c218c5d5f3 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 15:59:02 -0400 Subject: [PATCH 1/4] Use PiecewiseSequence for List take elements Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 184 +++++++------------ 1 file changed, 65 insertions(+), 119 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index ff9549d6a34..b78b5847630 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -23,10 +23,7 @@ use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::ConstantOrArray; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::builders::ArrayBuilder; -use crate::builders::PrimitiveBuilder; use crate::dtype::IntegerPType; -use crate::dtype::Nullability; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; @@ -64,7 +61,7 @@ impl TakeExecute for List { match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { match_each_unsigned_integer_ptype!(indices.ptype(), |I| { match_smallest_offset_type!(total_approx, |OutputOffsetType| { - _take::( + take_with_piecewise_elements::( array, offsets.as_view(), indices.as_view(), @@ -77,7 +74,11 @@ impl TakeExecute for List { } } -fn _take( +fn take_with_piecewise_elements< + I: IntegerPType, + O: IntegerPType, + OutputOffsetType: IntegerPType, +>( array: ArrayView<'_, List>, offsets_array: ArrayView<'_, Primitive>, indices_array: ArrayView<'_, Primitive>, @@ -91,56 +92,78 @@ fn _take( .vortex_expect("Failed to compute validity mask") .execute_mask(indices_array.as_ref().len(), ctx)?; - if !indices_validity.all_true() || !data_validity.all_true() { - return _take_nullable::(array, offsets_array, indices_array, ctx); - } - let offsets: &[O] = offsets_array.as_slice(); let indices: &[I] = indices_array.as_slice(); - let mut new_offsets = PrimitiveBuilder::::with_capacity( - Nullability::NonNullable, - indices.len(), - ); - let mut elements_to_take = - PrimitiveBuilder::with_capacity(Nullability::NonNullable, 2 * indices.len()); + let offsets_capacity = indices + .len() + .checked_add(1) + .ok_or_else(|| vortex_err!("List take offsets length overflow"))?; + let mut new_offsets = BufferMut::::with_capacity(offsets_capacity); + let mut element_starts = BufferMut::::with_capacity(indices.len()); + let mut element_lengths = BufferMut::::with_capacity(indices.len()); + + let mut current_offset = 0usize; + new_offsets.push(OutputOffsetType::zero()); - let mut current_offset = OutputOffsetType::zero(); - new_offsets.append_zero(); + for (&data_idx, index_valid) in indices.iter().zip_eq(indices_validity.iter()) { + if !index_valid { + new_offsets.push(new_offset_value::(current_offset)?); + element_starts.push(0); + element_lengths.push(0); + continue; + } - for &data_idx in indices { let data_idx: usize = data_idx.as_(); + if !data_validity.value(data_idx) { + new_offsets.push(new_offset_value::(current_offset)?); + element_starts.push(0); + element_lengths.push(0); + continue; + } + let start = offsets[data_idx]; let stop = offsets[data_idx + 1]; + let start: usize = start.as_(); + let stop: usize = stop.as_(); + let length = stop + .checked_sub(start) + .ok_or_else(|| vortex_err!("List offsets are not monotonic at offset {stop}"))?; - // Annoyingly, we can't turn (start..end) into a range, so we're doing that manually. - // - // We could convert start and end to usize, but that would impose a potentially - // harder constraint - now we don't care if they fit into usize as long as their - // difference does. - let additional: usize = (stop - start).as_(); - - // TODO(0ax1): optimize this - elements_to_take.reserve_exact(additional); - for i in 0..additional { - elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional")); - } - current_offset += - OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion"); - new_offsets.append_value(current_offset); + current_offset = current_offset + .checked_add(length) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + new_offsets.push(new_offset_value::(current_offset)?); + element_starts.push(start as u64); + element_lengths.push(length as u64); } - let elements_to_take = elements_to_take.finish(); - let new_offsets = new_offsets.finish(); - - let new_elements = array.elements().take(elements_to_take)?; + let new_offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array(); - Ok(ListArray::try_new( - new_elements, - new_offsets, - array.validity()?.take(indices_array.array())?, - )? + // SAFETY: valid source rows contribute ranges derived from list offsets; null index/source + // rows contribute zero-length placeholder ranges. `current_offset` is the sum of all generated + // element lengths, and multiplier 1 preserves contiguous ranges. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + element_starts.into_array(), + element_lengths.into_array(), + multipliers, + current_offset, + ) + }; + let new_elements = array.elements().take(element_indices.into_array())?; + + // SAFETY: offsets are rebuilt from the gathered element ranges and have one entry per output + // row plus the leading zero; validity is produced by the usual take-validity path. + Ok(unsafe { + ListArray::new_unchecked( + new_elements, + new_offsets, + array.validity()?.take(indices_array.array())?, + ) + } .into_array()) } @@ -603,83 +626,6 @@ fn new_offset_value(value: usize) -> VortexResult { }) } -// Kept out-of-line: as a single-callsite generic helper it would otherwise be inlined into every -// monomorphization of `_take`, duplicating the entire nullable path across all specializations. -#[inline(never)] -fn _take_nullable( - array: ArrayView<'_, List>, - offsets_array: ArrayView<'_, Primitive>, - indices_array: ArrayView<'_, Primitive>, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let offsets: &[O] = offsets_array.as_slice(); - let indices: &[I] = indices_array.as_slice(); - let data_validity = array - .list_validity() - .execute_mask(array.as_ref().len(), ctx)?; - let indices_validity = indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - .execute_mask(indices_array.as_ref().len(), ctx)?; - - let mut new_offsets = PrimitiveBuilder::::with_capacity( - Nullability::NonNullable, - indices.len(), - ); - - // This will be the indices we push down to the child array to call `take` with. - // - // There are 2 things to note here: - // - We do not know how many elements we need to take from our child since lists are variable - // size: thus we arbitrarily choose a capacity of `2 * # of indices`. - // - The type of the primitive builder needs to fit the largest offset of the (parent) - // `ListArray`, so we make this `PrimitiveBuilder` generic over `O` (instead of `I`). - let mut elements_to_take = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 2 * indices.len()); - - let mut current_offset = OutputOffsetType::zero(); - new_offsets.append_zero(); - - for (data_idx, index_valid) in indices.iter().zip(indices_validity.iter()) { - if !index_valid { - new_offsets.append_value(current_offset); - continue; - } - - let data_idx: usize = data_idx.as_(); - - if !data_validity.value(data_idx) { - new_offsets.append_value(current_offset); - continue; - } - - let start = offsets[data_idx]; - let stop = offsets[data_idx + 1]; - - // See the note in `_take` on the reasoning. - let additional: usize = (stop - start).as_(); - - elements_to_take.reserve_exact(additional); - for i in 0..additional { - elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional")); - } - current_offset += - OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion"); - new_offsets.append_value(current_offset); - } - - let elements_to_take = elements_to_take.finish(); - let new_offsets = new_offsets.finish(); - let new_elements = array.elements().take(elements_to_take)?; - - Ok(ListArray::try_new( - new_elements, - new_offsets, - array.validity()?.take(indices_array.array())?, - )? - .into_array()) -} - #[cfg(test)] mod test { use std::sync::Arc; From ff27447e611e1ffc82d116868d60a56262f3d778 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:07:41 -0400 Subject: [PATCH 2/4] Cover null List take placeholders Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 50 ++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index b78b5847630..6260f86a681 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -642,6 +642,7 @@ mod test { use crate::arrays::ListViewArray; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DType; use crate::dtype::Nullability; @@ -730,6 +731,55 @@ mod test { ); } + #[test] + fn null_index_ignores_out_of_bounds_payload() { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::try_new( + buffer![1i32, 2, 3, 4].into_array(), + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + + let idx = PrimitiveArray::new( + buffer![1u32, 99, 0], + Validity::from_iter([true, false, true]), + ) + .into_array(); + let result = list.take(idx).unwrap(); + + let expected = ListArray::new( + buffer![3i32, 4, 1, 2].into_array(), + buffer![0u32, 2, 2, 4].into_array(), + Validity::from_iter([true, false, true]), + ); + assert_arrays_eq!(expected, result, &mut ctx); + } + + #[test] + fn null_source_row_ignores_invalid_offset_payload() { + let mut ctx = array_session().create_execution_ctx(); + let list = unsafe { + ListArray::new_unchecked( + buffer![1i32, 2].into_array(), + buffer![0u32, 2, 999].into_array(), + Validity::from_iter([true, false]), + ) + } + .into_array(); + + let idx = buffer![0u32, 1].into_array(); + let result = list.take(idx).unwrap(); + + let expected = ListArray::new( + buffer![1i32, 2].into_array(), + buffer![0u32, 2, 2].into_array(), + Validity::from_iter([true, false]), + ); + assert_arrays_eq!(expected, result, &mut ctx); + } + #[test] fn change_validity() { let list = ListArray::try_new( From e781b619830a21ba73d698fe043024024bbef91e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:15:24 -0400 Subject: [PATCH 3/4] Normalize List take indices before range rebuild Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 46 ++++++++------------ 1 file changed, 17 insertions(+), 29 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 6260f86a681..aa583fe886e 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -3,10 +3,10 @@ use itertools::Itertools as _; use vortex_buffer::BufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; @@ -23,11 +23,13 @@ use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::ConstantOrArray; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::match_smallest_offset_type; +use crate::scalar::Scalar; use crate::validity::Validity; // TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call @@ -51,10 +53,16 @@ impl TakeExecute for List { return Ok(Some(taken)); } - let indices = indices.clone().execute::(ctx)?; + let new_validity = array.validity()?.take(indices)?; + let normalized_indices = indices + .clone() + .mask(new_validity.to_array(indices.len()))? + .fill_null(Scalar::from(0).cast(indices.dtype())?)?; + let indices = normalized_indices.execute::(ctx)?; let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); let offsets = array.offsets().clone().execute::(ctx)?; let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let validity_mask = new_validity.execute_mask(indices.len(), ctx)?; // This is an over-approximation of the total number of elements in the resulting array. let total_approx = array.elements().len().saturating_mul(indices.len()); @@ -65,7 +73,8 @@ impl TakeExecute for List { array, offsets.as_view(), indices.as_view(), - ctx, + new_validity, + &validity_mask, ) .map(Some) }) @@ -82,16 +91,9 @@ fn take_with_piecewise_elements< array: ArrayView<'_, List>, offsets_array: ArrayView<'_, Primitive>, indices_array: ArrayView<'_, Primitive>, - ctx: &mut ExecutionCtx, + new_validity: Validity, + validity_mask: &Mask, ) -> VortexResult { - let data_validity = array - .list_validity() - .execute_mask(array.as_ref().len(), ctx)?; - let indices_validity = indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - .execute_mask(indices_array.as_ref().len(), ctx)?; - let offsets: &[O] = offsets_array.as_slice(); let indices: &[I] = indices_array.as_slice(); @@ -106,8 +108,8 @@ fn take_with_piecewise_elements< let mut current_offset = 0usize; new_offsets.push(OutputOffsetType::zero()); - for (&data_idx, index_valid) in indices.iter().zip_eq(indices_validity.iter()) { - if !index_valid { + for (&data_idx, is_valid) in indices.iter().zip_eq(validity_mask.iter()) { + if !is_valid { new_offsets.push(new_offset_value::(current_offset)?); element_starts.push(0); element_lengths.push(0); @@ -116,13 +118,6 @@ fn take_with_piecewise_elements< let data_idx: usize = data_idx.as_(); - if !data_validity.value(data_idx) { - new_offsets.push(new_offset_value::(current_offset)?); - element_starts.push(0); - element_lengths.push(0); - continue; - } - let start = offsets[data_idx]; let stop = offsets[data_idx + 1]; let start: usize = start.as_(); @@ -157,14 +152,7 @@ fn take_with_piecewise_elements< // SAFETY: offsets are rebuilt from the gathered element ranges and have one entry per output // row plus the leading zero; validity is produced by the usual take-validity path. - Ok(unsafe { - ListArray::new_unchecked( - new_elements, - new_offsets, - array.validity()?.take(indices_array.array())?, - ) - } - .into_array()) + Ok(unsafe { ListArray::new_unchecked(new_elements, new_offsets, new_validity) }.into_array()) } fn take_piecewise_sequence( From 243ed198b6bfa4e54eae5ce0ab209ee89bb2d976 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:48:36 -0400 Subject: [PATCH 4/4] Avoid normalizing unused null list indices Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index aa583fe886e..56b6ce22511 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -23,13 +23,11 @@ use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::ConstantOrArray; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::match_smallest_offset_type; -use crate::scalar::Scalar; use crate::validity::Validity; // TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call @@ -54,11 +52,7 @@ impl TakeExecute for List { } let new_validity = array.validity()?.take(indices)?; - let normalized_indices = indices - .clone() - .mask(new_validity.to_array(indices.len()))? - .fill_null(Scalar::from(0).cast(indices.dtype())?)?; - let indices = normalized_indices.execute::(ctx)?; + let indices = indices.clone().execute::(ctx)?; let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); let offsets = array.offsets().clone().execute::(ctx)?; let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned());