diff --git a/vortex-array/src/arrays/extension/compute/rules.rs b/vortex-array/src/arrays/extension/compute/rules.rs index 2d02d1ae7a7..0699f2f68f4 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -11,6 +11,7 @@ use crate::arrays::ConstantArray; use crate::arrays::Extension; use crate::arrays::ExtensionArray; use crate::arrays::Filter; +use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; @@ -50,6 +51,7 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(Extension)), ParentRuleSet::lift(&MaskReduceAdaptor(Extension)), ParentRuleSet::lift(&SliceReduceAdaptor(Extension)), + ParentRuleSet::lift(&TakeReduceAdaptor(Extension)), ]); /// Push filter operations into the storage array of an extension array. diff --git a/vortex-array/src/arrays/extension/compute/take.rs b/vortex-array/src/arrays/extension/compute/take.rs index 8fb0da11222..d9ac1b4f600 100644 --- a/vortex-array/src/arrays/extension/compute/take.rs +++ b/vortex-array/src/arrays/extension/compute/take.rs @@ -10,8 +10,24 @@ use crate::array::ArrayView; use crate::arrays::Extension; use crate::arrays::ExtensionArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::dict::TakeReduce; use crate::arrays::extension::ExtensionArrayExt; +impl TakeReduce for Extension { + fn take(array: ArrayView<'_, Extension>, indices: &ArrayRef) -> VortexResult> { + let taken_storage = array.storage_array().take(indices.clone())?; + Ok(Some( + ExtensionArray::new( + array + .ext_dtype() + .with_nullability(taken_storage.dtype().nullability()), + taken_storage, + ) + .into_array(), + )) + } +} + impl TakeExecute for Extension { fn take( array: ArrayView<'_, Extension>, diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 14e89d4c238..369a743d4dc 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -1,26 +1,39 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +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 crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::ConstantArray; use crate::arrays::List; use crate::arrays::ListArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::list::ListArrayExt; +use crate::arrays::piecewise_sequence::UnitMultiplierLengths; +use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::validate_index_ranges; +use crate::arrays::piecewise_sequence::validate_index_ranges_constant; 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; use crate::match_smallest_offset_type; +use crate::validity::Validity; // TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call // the `ListView::take` compute function once `ListView` is more stable. @@ -37,6 +50,12 @@ impl TakeExecute for List { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let indices = indices.clone().execute::(ctx)?; let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); let offsets = array.offsets().clone().execute::(ctx)?; @@ -127,6 +146,448 @@ fn _take( .into_array()) } +fn take_piecewise_sequence( + array: ArrayView<'_, List>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let data_validity = array + .list_validity() + .execute_mask(array.as_ref().len(), ctx)?; + if !data_validity.all_true() { + return Ok(None); + } + + let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + return Ok(None); + }; + let offsets = array.offsets().clone().execute::(ctx)?; + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let output_len = indices_ref.len(); + + let taken = match &lengths { + UnitMultiplierLengths::Constant(length) => take_piecewise_sequence_constant_dispatch( + array, + &starts, + *length, + &offsets, + indices_ref, + output_len, + )?, + UnitMultiplierLengths::Array(lengths) => take_piecewise_sequence_lengths_dispatch( + array, + &starts, + lengths, + &offsets, + indices_ref, + output_len, + )?, + }; + Ok(Some(taken)) +} + +fn take_piecewise_sequence_constant_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_piecewise_sequence_constant_start_dispatch::( + array, + starts, + length, + offsets, + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_constant_start_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take_piecewise_sequence_constant_length::( + array, + starts.as_slice::(), + length, + offsets.as_slice::(), + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_lengths_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_piecewise_sequence_lengths_start_dispatch::( + array, + starts, + lengths, + offsets, + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_lengths_start_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + take_piecewise_sequence_lengths_start_length_dispatch::( + array, + starts, + lengths, + offsets, + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_lengths_start_length_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take_piecewise_sequence_typed::( + array, + starts.as_slice::(), + lengths.as_slice::(), + offsets.as_slice::(), + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_constant_length( + array: ArrayView<'_, List>, + starts: &[S], + length: usize, + offsets: &[Offset], + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, +{ + validate_index_ranges_constant(array.len(), starts, length, output_len)?; + let total_elements = + piecewise_list_elements_len_constant(array.elements().len(), offsets, starts, length)?; + let validity = array.validity()?.take(indices_ref)?; + + match_smallest_offset_type!(total_elements, |OutputOffset| { + let gathered = gather_piecewise_list_constant_length::( + array.elements(), + offsets, + starts, + length, + output_len, + total_elements, + )?; + + // SAFETY: output offsets are rebuilt from valid monotonic source offsets; output elements + // are exactly the gathered child ranges referenced by those offsets; validity has one bit + // per output row. + Ok( + unsafe { ListArray::new_unchecked(gathered.elements, gathered.offsets, validity) } + .into_array(), + ) + }) +} + +fn take_piecewise_sequence_typed( + array: ArrayView<'_, List>, + starts: &[S], + lengths: &[L], + offsets: &[Offset], + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, +{ + validate_index_ranges(array.len(), starts, lengths, output_len)?; + let total_elements = + piecewise_list_elements_len(array.elements().len(), offsets, starts, lengths)?; + + match_smallest_offset_type!(total_elements, |OutputOffset| { + let gathered = gather_piecewise_list::( + array.elements(), + offsets, + starts, + lengths, + output_len, + total_elements, + )?; + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: output offsets are rebuilt from valid monotonic source offsets; output elements + // are exactly the gathered child ranges referenced by those offsets; validity has one bit + // per output row. + Ok( + unsafe { ListArray::new_unchecked(gathered.elements, gathered.offsets, validity) } + .into_array(), + ) + }) +} + +struct GatheredList { + elements: ArrayRef, + offsets: ArrayRef, +} + +fn piecewise_list_elements_len_constant( + elements_len: usize, + offsets: &[Offset], + starts: &[S], + length: usize, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, +{ + let mut total = 0usize; + for &start in starts { + let start: usize = start.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let element_start: usize = offsets[start].as_(); + let element_end: usize = offsets[end].as_(); + vortex_ensure!( + element_start <= element_end && element_end <= elements_len, + "List offsets range {element_start}..{element_end} exceeds elements length {elements_len}", + ); + total = total + .checked_add(element_end - element_start) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + Ok(total) +} + +fn piecewise_list_elements_len( + elements_len: usize, + offsets: &[Offset], + starts: &[S], + lengths: &[L], +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, +{ + let mut total = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let element_start: usize = offsets[start].as_(); + let element_end: usize = offsets[end].as_(); + vortex_ensure!( + element_start <= element_end && element_end <= elements_len, + "List offsets range {element_start}..{element_end} exceeds elements length {elements_len}", + ); + total = total + .checked_add(element_end - element_start) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + Ok(total) +} + +fn gather_piecewise_list_constant_length( + elements: &ArrayRef, + offsets: &[Offset], + starts: &[S], + length: usize, + output_len: usize, + total_elements: usize, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, + OutputOffset: IntegerPType, +{ + let offsets_capacity = output_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(starts.len()); + let mut element_lengths = BufferMut::::with_capacity(starts.len()); + let mut output_elements = 0usize; + + new_offsets.push(OutputOffset::zero()); + for &start in starts { + let start: usize = start.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let element_start: usize = offsets[start].as_(); + let element_end: usize = offsets[end].as_(); + for &offset in &offsets[start + 1..=end] { + let offset: usize = offset.as_(); + let relative = offset + .checked_sub(element_start) + .ok_or_else(|| vortex_err!("List offsets are not monotonic at offset {offset}"))?; + let output_offset = output_elements + .checked_add(relative) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + let element_length = element_end - element_start; + element_starts.push(element_start as u64); + element_lengths.push(element_length as u64); + output_elements = output_elements + .checked_add(element_length) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + debug_assert_eq!(output_elements, total_elements); + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array(); + // SAFETY: element ranges are derived from validated source list offsets, and total_elements is + // the sum of the gathered element range lengths. Multiplier 1 preserves contiguous ranges. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + element_starts.into_array(), + element_lengths.into_array(), + multipliers, + total_elements, + ) + }; + let elements = elements.take(element_indices.into_array())?; + + Ok(GatheredList { elements, offsets }) +} + +fn gather_piecewise_list( + elements: &ArrayRef, + offsets: &[Offset], + starts: &[S], + lengths: &[L], + output_len: usize, + total_elements: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, + OutputOffset: IntegerPType, +{ + let offsets_capacity = output_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(starts.len()); + let mut element_lengths = BufferMut::::with_capacity(lengths.len()); + let mut output_elements = 0usize; + + new_offsets.push(OutputOffset::zero()); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let element_start: usize = offsets[start].as_(); + let element_end: usize = offsets[end].as_(); + for &offset in &offsets[start + 1..=end] { + let offset: usize = offset.as_(); + let relative = offset + .checked_sub(element_start) + .ok_or_else(|| vortex_err!("List offsets are not monotonic at offset {offset}"))?; + let output_offset = output_elements + .checked_add(relative) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + let element_length = element_end - element_start; + element_starts.push(element_start as u64); + element_lengths.push(element_length as u64); + output_elements = output_elements + .checked_add(element_length) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + debug_assert_eq!(output_elements, total_elements); + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array(); + // SAFETY: element ranges are derived from validated source list offsets, and total_elements is + // the sum of the gathered element range lengths. Multiplier 1 preserves contiguous ranges. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + element_starts.into_array(), + element_lengths.into_array(), + multipliers, + total_elements, + ) + }; + let elements = elements.take(element_indices.into_array())?; + + Ok(GatheredList { elements, offsets }) +} + +fn new_offset_value(value: usize) -> VortexResult { + T::from_usize(value).ok_or_else(|| { + vortex_err!( + "List take offset value {value} does not fit in {}", + T::PTYPE + ) + }) +} + // 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)] @@ -215,8 +676,10 @@ mod test { use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; + use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DType; @@ -403,6 +866,54 @@ mod test { ); } + #[test] + fn piecewise_sequence_take() { + let mut ctx = array_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, + ) + .unwrap() + .into_array(); + let idx = PiecewiseSequenceArray::try_new( + buffer![1u64, 0].into_array(), + buffer![2u64, 1].into_array(), + ConstantArray::new(1u64, 2).into_array(), + 3, + ) + .unwrap() + .into_array(); + + let result = list + .take(idx) + .unwrap() + .execute::(&mut ctx) + .unwrap(); + + let element_dtype: Arc = Arc::new(I32.into()); + assert_eq!( + result.execute_scalar(0, &mut ctx).unwrap(), + Scalar::list( + Arc::clone(&element_dtype), + vec![2i32.into(), 3.into(), 4.into()], + Nullability::NonNullable + ) + ); + assert_eq!( + result.execute_scalar(1, &mut ctx).unwrap(), + Scalar::list(Arc::clone(&element_dtype), vec![], Nullability::NonNullable) + ); + assert_eq!( + result.execute_scalar(2, &mut ctx).unwrap(), + Scalar::list( + element_dtype, + vec![0i32.into(), 1.into()], + Nullability::NonNullable + ) + ); + } + #[test] fn test_take_empty_array() { let list = ListArray::try_new( diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 342c9eccf19..15eb9feab38 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -5,16 +5,17 @@ use num_traits::FromPrimitive; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; -use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -140,19 +141,18 @@ impl ListViewArray { // for sizes as well. match_each_unsigned_integer_ptype!(sizes_ptype.to_unsigned(), |S| { match offsets_ptype.to_unsigned() { - PType::U8 => self.naive_rebuild::(ctx), - PType::U16 => self.naive_rebuild::(ctx), - PType::U32 => self.naive_rebuild::(ctx), - PType::U64 => self.naive_rebuild::(ctx), + PType::U8 => self.rebuild_with_take_or_piecewise::(ctx), + PType::U16 => self.rebuild_with_take_or_piecewise::(ctx), + PType::U32 => self.rebuild_with_take_or_piecewise::(ctx), + PType::U64 => self.rebuild_with_take_or_piecewise::(ctx), _ => unreachable!("invalid offsets PType"), } }) } /// Picks between [`rebuild_with_take`](Self::rebuild_with_take) and - /// [`rebuild_list_by_list`](Self::rebuild_list_by_list) based on element dtype and average - /// list size. - fn naive_rebuild( + /// [`rebuild_with_piecewise`](Self::rebuild_with_piecewise) based on average list size. + fn rebuild_with_take_or_piecewise( &self, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -167,12 +167,12 @@ impl ListViewArray { if Self::should_use_take(total, self.len()) { self.rebuild_with_take::(ctx) } else { - self.rebuild_list_by_list::(ctx) + self.rebuild_with_piecewise::(ctx) } } /// Returns `true` when we are confident that `rebuild_with_take` will - /// outperform `rebuild_list_by_list`. + /// outperform `rebuild_with_piecewise`. /// /// Take is dramatically faster for small lists (often 10-100×) because it /// avoids per-list builder overhead. LBL is the safer default for larger @@ -243,86 +243,69 @@ impl ListViewArray { .reinterpret_cast(size_ptype) .into_array(); - // SAFETY: same invariants as `rebuild_list_by_list` — offsets are sequential and - // non-overlapping, all (offset, size) pairs reference valid elements, and the validity - // array is preserved from the original. + // SAFETY: offsets are sequential and non-overlapping, all (offset, size) pairs reference + // valid elements, and the validity array is preserved from the original. Ok(unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?) .with_zero_copy_to_list(true) }) } - /// Rebuilds elements list-by-list: canonicalize elements upfront, then for each list `slice` - /// the relevant range and `append_to_builder` into a typed builder. - fn rebuild_list_by_list( + /// Rebuilds elements using one contiguous-run take over the element child. + fn rebuild_with_piecewise( &self, ctx: &mut ExecutionCtx, ) -> VortexResult { - let element_dtype = self - .dtype() - .as_list_element_opt() - .vortex_expect("somehow had a canonical list that was not a list"); - let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype()); let size_ptype = self.sizes().dtype().as_ptype(); let offsets_canonical = self.offsets().clone().execute::(ctx)?; let offsets_canonical = offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned()); - let offsets_slice = offsets_canonical.as_slice::(); let sizes_canonical = self.sizes().clone().execute::(ctx)?; let sizes_canonical = sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); - let sizes_slice = sizes_canonical.as_slice::(); - - let len = offsets_slice.len(); - - let mut new_offsets = BufferMut::::with_capacity(len); - // TODO(connor)[ListView]: Do we really need to do this? - // The only reason we need to rebuild the sizes here is that the validity may indicate that - // a list is null even though it has a non-zero size. This rebuild will set the size of all - // null lists to 0. - let mut new_sizes = BufferMut::::with_capacity(len); - - // Canonicalize the elements up front as we will be slicing the elements quite a lot. - let elements_canonical = self - .elements() - .clone() - .execute::(ctx)? - .into_array(); - - // Note that we do not know what the exact capacity should be of the new elements since - // there could be overlaps in the existing `ListViewArray`. - let mut new_elements_builder = - builder_with_capacity(element_dtype.as_ref(), self.elements().len()); - - // Resolve validity to a mask once instead of probing it per row (see `rebuild_with_take`). - let validity = self.validity()?.execute_mask(len, ctx)?; - - let mut n_elements = NewOffset::zero(); - for index in 0..len { - if !validity.value(index) { - // For NULL lists, place them after the previous item's data to maintain the - // no-overlap invariant for zero-copy to `ListArray` arrays. - new_offsets.push(n_elements); - new_sizes.push(S::zero()); - continue; - } - let offset = offsets_slice[index]; - let size = sizes_slice[index]; + let len = offsets_canonical.len(); + let validity = self.validity()?; + let validity_mask = validity.execute_mask(len, ctx)?; - let start = offset.as_(); - let stop = start + size.as_(); - - new_offsets.push(n_elements); - new_sizes.push(size); - elements_canonical - .slice(start..stop)? - .append_to_builder(new_elements_builder.as_mut(), ctx)?; - - n_elements += num_traits::cast(size).vortex_expect("Cast failed"); - } + let ranges = match_each_unsigned_integer_ptype!(offsets_canonical.ptype(), |O| { + rebuild_ranges::( + offsets_canonical.as_slice::(), + sizes_canonical.as_slice::(), + &validity_mask, + ) + })?; + + let RebuildRanges { + new_offsets, + new_sizes, + starts, + lengths, + elements_len, + } = ranges; + let constant_length = lengths + .first() + .copied() + .filter(|first| lengths.iter().all(|length| *length == *first)); + let lengths = match constant_length { + Some(length) => ConstantArray::new(length, starts.len()).into_array(), + None => lengths.into_array(), + }; + + // SAFETY: range starts and lengths are derived from valid ListView metadata; elements_len + // is the sum of all generated range lengths. Multiplier 1 preserves contiguous ranges. + let multipliers = ConstantArray::new(1u64, starts.len()).into_array(); + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + starts.into_array(), + lengths, + multipliers, + elements_len, + ) + }; + let elements = self.elements().take(element_indices.into_array())?; // Built unsigned; reinterpret back to the signed-preserving result types. let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) @@ -331,24 +314,11 @@ impl ListViewArray { let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable) .reinterpret_cast(size_ptype) .into_array(); - let elements = new_elements_builder.finish(); - - debug_assert_eq!( - n_elements.as_(), - elements.len(), - "The accumulated elements somehow had the wrong length" - ); - // SAFETY: - // - All offsets are sequential and non-overlapping (`n_elements` tracks running total). - // - Each `offset[i] + size[i]` equals `offset[i+1]` for all valid indices (including null - // lists). - // - All elements referenced by (offset, size) pairs exist within the new `elements` array. - // - The validity array is preserved from the original array unchanged - // - The array satisfies the zero-copy-to-list property by having sorted offsets, no gaps, - // and no overlaps. + // SAFETY: offsets are sequential and non-overlapping, all (offset, size) pairs reference + // valid elements, and validity is preserved from the original array. Ok(unsafe { - ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?) + ListViewArray::new_unchecked(elements, offsets, sizes, validity) .with_zero_copy_to_list(true) }) } @@ -414,6 +384,69 @@ impl ListViewArray { } } +struct RebuildRanges { + new_offsets: BufferMut, + new_sizes: BufferMut, + starts: BufferMut, + lengths: BufferMut, + elements_len: usize, +} + +fn rebuild_ranges( + offsets: &[O], + sizes: &[S], + validity_mask: &Mask, +) -> VortexResult> +where + NewOffset: IntegerPType, + O: IntegerPType, + S: IntegerPType, +{ + let len = offsets.len(); + let mut new_offsets = BufferMut::::with_capacity(len); + let mut new_sizes = BufferMut::::with_capacity(len); + let mut starts = BufferMut::::with_capacity(len); + let mut lengths = BufferMut::::with_capacity(len); + let mut n_elements = NewOffset::zero(); + let mut elements_len = 0usize; + + for (index, is_valid) in validity_mask.iter().enumerate() { + if !is_valid { + new_offsets.push(n_elements); + new_sizes.push(S::zero()); + starts.push(0); + lengths.push(0); + continue; + } + + let size = sizes[index]; + let start: usize = offsets[index].as_(); + let length: usize = size.as_(); + start.checked_add(length).ok_or_else(|| { + vortex_err!( + "ListView rebuild element range overflow for start {start} and length {length}" + ) + })?; + elements_len = elements_len + .checked_add(length) + .ok_or_else(|| vortex_err!("ListView rebuild elements length overflow"))?; + + new_offsets.push(n_elements); + new_sizes.push(size); + starts.push(start as u64); + lengths.push(length as u64); + n_elements += num_traits::cast(size).vortex_expect("Cast failed"); + } + + Ok(RebuildRanges { + new_offsets, + new_sizes, + starts, + lengths, + elements_len, + }) +} + #[cfg(test)] mod tests { use vortex_buffer::BitBuffer; @@ -561,6 +594,35 @@ mod tests { Ok(()) } + #[test] + fn test_rebuild_flatten_large_lists_with_piecewise_indices() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter(0i32..300).into_array(); + let offsets = PrimitiveArray::from_iter(vec![10u32, 0]).into_array(); + let sizes = PrimitiveArray::from_iter(vec![150u32, 130]).into_array(); + let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); + + let mut ctx = SESSION.create_execution_ctx(); + let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; + + assert_eq!(flattened.elements().len(), 280); + assert_eq!(flattened.offset_at(0), 0); + assert_eq!(flattened.size_at(0), 150); + assert_eq!(flattened.offset_at(1), 150); + assert_eq!(flattened.size_at(1), 130); + + assert_arrays_eq!( + flattened.list_elements_at(0)?, + PrimitiveArray::from_iter(10i32..160), + &mut ctx + ); + assert_arrays_eq!( + flattened.list_elements_at(1)?, + PrimitiveArray::from_iter(0i32..130), + &mut ctx + ); + Ok(()) + } + #[test] fn test_rebuild_with_trailing_nulls_regression() -> VortexResult<()> { // Regression test for issue #5412 diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index 82ead343f96..7a2a9401ec5 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -16,6 +16,7 @@ use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; use crate::arrays::PiecewiseSequence; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; @@ -354,6 +355,27 @@ fn contiguous_take_consumers_support_constant_piecewise_lengths() -> VortexResul Ok(()) } +#[test] +fn list_take_consumes_constant_piecewise_lengths() -> VortexResult<()> { + let mut ctx = array_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, + )? + .into_array(); + let list_indices = piecewise_indices_constant_length([1, 0], 2)?; + let expected = ListArray::try_new( + buffer![2i32, 3, 4, 0, 1, 2, 3, 4].into_array(), + buffer![0u32, 3, 3, 5, 8].into_array(), + Validity::NonNullable, + )? + .into_array(); + assert_arrays_eq!(list.take(list_indices)?, expected, &mut ctx); + + Ok(()) +} + #[test] fn fixed_size_list_take_builds_piecewise_element_indices() -> VortexResult<()> { let elements = PrimitiveArray::from_iter(0i32..12).into_array();