From 6ba92fef72cf5587ed6dbe1ed448227eb70f5b48 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 15:59:02 -0400 Subject: [PATCH 1/5] 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 e141e8cf786..34c14c6f6f5 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -24,10 +24,7 @@ use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::constant_unsigned_usize; 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; @@ -65,7 +62,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(), @@ -78,7 +75,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>, @@ -92,56 +93,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()) } @@ -577,83 +600,6 @@ fn new_offset_value(value: usize) -> T { T::from_usize(value).vortex_expect("output offset fits selected offset type") } -// 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 b7afe5e994006803fb489572a344f3f0addbd4df Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:07:41 -0400 Subject: [PATCH 2/5] 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 34c14c6f6f5..57a3c466473 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -616,6 +616,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; @@ -704,6 +705,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 17b31d0b8b791d3a95fc5844992c4f3baf3b4681 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:15:24 -0400 Subject: [PATCH 3/5] 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 57a3c466473..6615e615365 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::Columnar; @@ -24,11 +24,13 @@ use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::constant_unsigned_usize; 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 @@ -52,10 +54,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()); @@ -66,7 +74,8 @@ impl TakeExecute for List { array, offsets.as_view(), indices.as_view(), - ctx, + new_validity, + &validity_mask, ) .map(Some) }) @@ -83,16 +92,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(); @@ -107,8 +109,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); @@ -117,13 +119,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_(); @@ -158,14 +153,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 39d8de7740fdbc7948dacde0da53bf515deb5f7c Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:48:36 -0400 Subject: [PATCH 4/5] 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 6615e615365..1a6e4bfe3aa 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -24,13 +24,11 @@ use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::constant_unsigned_usize; 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 @@ -55,11 +53,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()); From e88d96f26ede1e233642f4d4feb3413c77c86b3d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 20 Jul 2026 12:54:57 -0400 Subject: [PATCH 5/5] Rely on List offset invariants in List take Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 1a6e4bfe3aa..b973b6c2c85 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -3,6 +3,7 @@ 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; @@ -105,7 +106,7 @@ fn take_with_piecewise_elements< for (&data_idx, is_valid) in indices.iter().zip_eq(validity_mask.iter()) { if !is_valid { - new_offsets.push(new_offset_value::(current_offset)?); + new_offsets.push(new_offset_value::(current_offset)); element_starts.push(0); element_lengths.push(0); continue; @@ -117,14 +118,12 @@ fn take_with_piecewise_elements< 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}"))?; + let length = stop - start; 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)?); + new_offsets.push(new_offset_value::(current_offset)); element_starts.push(start as u64); element_lengths.push(length as u64); }