Skip to content

Use PiecewiseSequence for List take elements #8833

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
danking wants to merge 6 commits into
base: develop
Choose a base branch
from
+104 −130
Open

Use PiecewiseSequence for List take elements #8833

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
2f29140
Use PiecewiseSequence for List take elements
danking Jul 17, 2026
90c43ab
Cover null List take placeholders
danking Jul 17, 2026
fb435f1
Normalize List take indices before range rebuild
danking Jul 17, 2026
201a5f0
Avoid normalizing unused null list indices
danking Jul 17, 2026
8d542fe
Rely on List offset invariants in List take
danking Jul 20, 2026
e0096ca
Use valid null List ranges in take test
danking Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
234 changes: 104 additions & 130 deletions vortex-array/src/arrays/list/compute/take.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,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;
Expand Down Expand Up @@ -56,21 +53,24 @@ impl TakeExecute for List {
return Ok(Some(taken));
}

let new_validity = array.validity()?.take(indices)?;
let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
let indices = indices.reinterpret_cast(indices.ptype().to_unsigned());
let offsets = array.offsets().clone().execute::<PrimitiveArray>(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());

match_each_unsigned_integer_ptype!(offsets.ptype(), |O| {
match_each_unsigned_integer_ptype!(indices.ptype(), |I| {
match_smallest_offset_type!(total_approx, |OutputOffsetType| {
_take::<I, O, OutputOffsetType>(
take_with_piecewise_elements::<I, O, OutputOffsetType>(
array,
offsets.as_view(),
indices.as_view(),
ctx,
new_validity,
&validity_mask,
Comment thread
robert3005 marked this conversation as resolved.
)
.map(Some)
})
Expand All @@ -79,71 +79,74 @@ impl TakeExecute for List {
}
}

fn _take<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
fn take_with_piecewise_elements<
I: IntegerPType,
O: IntegerPType,
OutputOffsetType: IntegerPType,
>(
array: ArrayView<'_, List>,
offsets_array: ArrayView<'_, Primitive>,
indices_array: ArrayView<'_, Primitive>,
ctx: &mut ExecutionCtx,
new_validity: Validity,
validity_mask: &Mask,
) -> VortexResult<ArrayRef> {
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)?;

if !indices_validity.all_true() || !data_validity.all_true() {
return _take_nullable::<I, O, OutputOffsetType>(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::<OutputOffsetType>::with_capacity(
Nullability::NonNullable,
indices.len(),
);
let mut elements_to_take =
PrimitiveBuilder::with_capacity(Nullability::NonNullable, 2 * indices.len());

let mut current_offset = OutputOffsetType::zero();
new_offsets.append_zero();
let offsets_capacity = indices
.len()
.checked_add(1)
.ok_or_else(|| vortex_err!("List take offsets length overflow"))?;
let mut new_offsets = BufferMut::<OutputOffsetType>::with_capacity(offsets_capacity);

@robert3005 robert3005 Jul 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe, just maybe this name is too long? OutOffset?

let mut element_starts = BufferMut::<u64>::with_capacity(indices.len());
let mut element_lengths = BufferMut::<u64>::with_capacity(indices.len());

let mut current_offset = 0usize;
new_offsets.push(OutputOffsetType::zero());

for (&data_idx, is_valid) in indices.iter().zip_eq(validity_mask.iter()) {
if !is_valid {
new_offsets.push(new_offset_value::<OutputOffsetType>(current_offset));
element_starts.push(0);
element_lengths.push(0);
continue;
}

for &data_idx in indices {
let data_idx: usize = data_idx.as_();

let start = offsets[data_idx];
let stop = offsets[data_idx + 1];
let start: usize = start.as_();
let stop: usize = stop.as_();
let length = stop - start;

// 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::<OutputOffsetType>(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_offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array();
let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array();

let new_elements = array.elements().take(elements_to_take)?;
// 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())?;

Ok(ListArray::try_new(
new_elements,
new_offsets,
array.validity()?.take(indices_array.array())?,
)?
.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, new_validity) }.into_array())
}

fn take_slices(
Expand Down Expand Up @@ -842,83 +845,6 @@ fn new_offset_value<T: IntegerPType>(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<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
array: ArrayView<'_, List>,
offsets_array: ArrayView<'_, Primitive>,
indices_array: ArrayView<'_, Primitive>,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
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::<OutputOffsetType>::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::<O>::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;
Expand All @@ -937,6 +863,7 @@ mod test {
use crate::arrays::PiecewiseSequenceArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::listview::ListViewArrayExt;
use crate::assert_arrays_eq;
use crate::compute::conformance::take::test_take_conformance;
use crate::dtype::DType;
use crate::dtype::Nullability;
Expand Down Expand Up @@ -1025,6 +952,53 @@ 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_uses_valid_empty_output_range() {
let mut ctx = array_session().create_execution_ctx();
let list = ListArray::new(
buffer![1i32, 2, 7, 8].into_array(),
buffer![0u32, 2, 4].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(
Expand Down
Loading