From 8d9a0d64eea333221361c64d7ef99d52557d9998 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 7 Jul 2026 12:30:03 -0400 Subject: [PATCH 01/37] Add TakeSlices for fixed-size-list take Introduce a lazy TakeSlices array for ordered child ranges, wire FixedSizeList take to emit element ranges, and add a primitive child execution kernel for contiguous range copies. Expand FSL take benchmarks and cover TakeSlices edge cases, nullable validity, generic slice-based execution, and nullable FSL take with nullable indices. Signed-off-by: Daniel King --- vortex-array/benches/take_fsl.rs | 199 +++++++++++- vortex-array/src/array/erased.rs | 9 + .../arrays/fixed_size_list/compute/take.rs | 301 +++++++++++------- .../src/arrays/fixed_size_list/tests/take.rs | 240 ++++++++++++-- vortex-array/src/arrays/mod.rs | 4 + .../src/arrays/primitive/compute/mod.rs | 1 + .../arrays/primitive/compute/take_slices.rs | 37 +++ .../src/arrays/primitive/vtable/kernel.rs | 7 + vortex-array/src/arrays/take_slices/array.rs | 114 +++++++ vortex-array/src/arrays/take_slices/mod.rs | 126 ++++++++ vortex-array/src/arrays/take_slices/rules.rs | 71 +++++ vortex-array/src/arrays/take_slices/tests.rs | 205 ++++++++++++ vortex-array/src/arrays/take_slices/vtable.rs | 208 ++++++++++++ vortex-array/src/session/mod.rs | 2 + vortex-array/src/validity.rs | 8 + 15 files changed, 1392 insertions(+), 140 deletions(-) create mode 100644 vortex-array/src/arrays/primitive/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/take_slices/array.rs create mode 100644 vortex-array/src/arrays/take_slices/mod.rs create mode 100644 vortex-array/src/arrays/take_slices/rules.rs create mode 100644 vortex-array/src/arrays/take_slices/tests.rs create mode 100644 vortex-array/src/arrays/take_slices/vtable.rs diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index 5dc491c28c5..8f0e63191cb 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -6,6 +6,7 @@ //! Parameterized over: //! - Number of indices to take //! - Fixed size list length (elements per list) +//! - Element byte width #![expect(clippy::cast_possible_truncation)] #![expect(clippy::unwrap_used)] @@ -13,16 +14,26 @@ use std::sync::LazyLock; use divan::Bencher; +use divan::counter::BytesCount; +use num_traits::FromPrimitive; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::half::f16; +use vortex_array::match_smallest_offset_type; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_session::VortexSession; fn main() { @@ -39,12 +50,20 @@ const NUM_LISTS: usize = 500; const NUM_INDICES: &[usize] = &[100, 1_000]; /// Fixed size list lengths (elements per list). -const LIST_SIZES: &[usize] = &[16, 64, 256, 1024, 4096]; +const LIST_SIZES: &[usize] = &[16, 64, 128, 256, 512, 1024, 2048, 4096]; + +/// F16 list lengths for isolating the per-index, take_slices, and manual range-copy strategies. +const F16_STRATEGY_LIST_SIZES: &[usize] = &[1, 2, 4, 8, 16, 64, 128, 256, 512, 1024, 2048, 4096]; /// Creates a FixedSizeListArray with the given list size and number of lists. -fn create_fsl(list_size: usize, num_lists: usize) -> FixedSizeListArray { +fn create_fsl(list_size: usize, num_lists: usize) -> FixedSizeListArray +where + T: NativePType + FromPrimitive, +{ let total_elements = list_size * num_lists; - let elements: Buffer = (0..total_elements as i64).collect(); + let elements: Buffer = (0..total_elements) + .map(|idx| T::from_u16((idx % 251) as u16).unwrap()) + .collect(); FixedSizeListArray::new( elements.into_array(), list_size as u32, @@ -62,12 +81,35 @@ fn create_random_indices(num_indices: usize, max_index: usize) -> Buffer { } #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] -fn take_fsl_random(bencher: Bencher, num_indices: usize) { - let fsl = create_fsl(LIST_SIZE, NUM_LISTS); +fn take_fsl_f16_random(bencher: Bencher, num_indices: usize) { + take_fsl_random::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u8_random(bencher: Bencher, num_indices: usize) { + take_fsl_random::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u32_random(bencher: Bencher, num_indices: usize) { + take_fsl_random::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u64_random(bencher: Bencher, num_indices: usize) { + take_fsl_random::(bencher, num_indices); +} + +fn take_fsl_random(bencher: Bencher, num_indices: usize) +where + T: NativePType + FromPrimitive, +{ + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); let indices = create_random_indices(num_indices, NUM_LISTS); let indices_array = indices.into_array(); bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { array @@ -79,10 +121,152 @@ fn take_fsl_random(bencher: Bencher, num_indices: usize) }); } +#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_per_index(bencher: Bencher, num_indices: usize) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + let indices = create_random_indices(num_indices, NUM_LISTS); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + match_smallest_offset_type!(array.elements().len(), |E| { + take_fsl_f16_per_index_strategy::(array, indices) + }) + .into_array() + .execute::(execution_ctx) + .unwrap() + }); +} + +#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_take_slices(bencher: Bencher, num_indices: usize) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + let indices = create_random_indices(num_indices, NUM_LISTS); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + take_fsl_f16_take_slices_strategy::(array, indices) + .into_array() + .execute::(execution_ctx) + .unwrap() + }); +} + +#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_manual_range_copy( + bencher: Bencher, + num_indices: usize, +) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + let indices = create_random_indices(num_indices, NUM_LISTS); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + take_fsl_f16_manual_range_copy_strategy::(array, indices, execution_ctx) + .into_array() + .execute::(execution_ctx) + .unwrap() + }); +} + +fn take_fsl_f16_per_index_strategy( + array: &FixedSizeListArray, + indices: &Buffer, +) -> FixedSizeListArray { + let mut element_indices = BufferMut::::with_capacity(indices.len() * LIST_SIZE); + for &idx in indices.as_ref() { + let start = idx as usize * LIST_SIZE; + let end = start + LIST_SIZE; + for element_idx in start..end { + // SAFETY: capacity is exactly `indices.len() * LIST_SIZE`, and this loop appends + // exactly `LIST_SIZE` element indices for each input index. + unsafe { element_indices.push_unchecked(E::from_usize(element_idx).unwrap()) }; + } + } + + let element_indices = + PrimitiveArray::new(element_indices.freeze(), Validity::NonNullable).into_array(); + let elements = array.elements().take(element_indices).unwrap(); + + // SAFETY: `elements` was built by taking exactly `LIST_SIZE` elements per input index, so its + // length is `indices.len() * LIST_SIZE`; the output is non-nullable by construction. + unsafe { + FixedSizeListArray::new_unchecked( + elements, + LIST_SIZE as u32, + Validity::NonNullable, + indices.len(), + ) + } +} + +fn take_fsl_f16_take_slices_strategy( + array: &FixedSizeListArray, + indices: &Buffer, +) -> FixedSizeListArray { + let slices = indices + .as_ref() + .iter() + .map(|&idx| { + let start = idx as usize * LIST_SIZE; + (start, start + LIST_SIZE) + }) + .collect(); + let elements = array.elements().take_slices(slices).unwrap(); + + // SAFETY: each generated slice has width `LIST_SIZE`, and there is one slice per input index, + // so `elements.len() == indices.len() * LIST_SIZE`. + unsafe { + FixedSizeListArray::new_unchecked( + elements, + LIST_SIZE as u32, + Validity::NonNullable, + indices.len(), + ) + } +} + +fn take_fsl_f16_manual_range_copy_strategy( + array: &FixedSizeListArray, + indices: &Buffer, + execution_ctx: &mut ExecutionCtx, +) -> FixedSizeListArray { + let elements = array + .elements() + .clone() + .execute::(execution_ctx) + .unwrap(); + let source = elements.as_slice::(); + let mut values = BufferMut::::with_capacity(indices.len() * LIST_SIZE); + + for &idx in indices.as_ref() { + let start = idx as usize * LIST_SIZE; + values.extend_from_slice(&source[start..start + LIST_SIZE]); + } + + // SAFETY: the buffer was filled with exactly `LIST_SIZE` copied values per input index, so it + // has the element length required by the constructed FSL. + unsafe { + FixedSizeListArray::new_unchecked( + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + LIST_SIZE as u32, + Validity::NonNullable, + indices.len(), + ) + } +} + #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] -fn take_fsl_nullable_random(bencher: Bencher, num_indices: usize) { +fn take_fsl_f16_nullable_random(bencher: Bencher, num_indices: usize) { let total_elements = LIST_SIZE * NUM_LISTS; - let elements: Buffer = (0..total_elements as i64).collect(); + let elements: Buffer = (0..total_elements) + .map(|idx| f16::from_u16((idx % 251) as u16).unwrap()) + .collect(); // Create validity with ~10% nulls let mut rng = StdRng::seed_from_u64(123); @@ -94,6 +278,7 @@ fn take_fsl_nullable_random(bencher: Bencher, num_indice let indices_array = indices.into_array(); bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { array diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index c3dc2e0eed7..8722ba3d115 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -41,6 +41,7 @@ use crate::arrays::FilterArray; use crate::arrays::Null; use crate::arrays::Primitive; use crate::arrays::SliceArray; +use crate::arrays::TakeSlicesArray; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::buffer::BufferHandle; @@ -264,6 +265,14 @@ impl ArrayRef { .optimize() } + /// Wraps the array in a [`TakeSlicesArray`] such that it is logically selected by ordered + /// child ranges. + pub fn take_slices(&self, slices: Vec<(usize, usize)>) -> VortexResult { + TakeSlicesArray::try_new(self.clone(), slices)? + .into_array() + .optimize() + } + /// Fetch the scalar at the given index. #[deprecated( note = "Use `execute_scalar` instead, which allows passing an execution context for more \ diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 193730f62c4..0a397b7b807 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -2,9 +2,11 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::BitBufferMut; -use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -20,37 +22,31 @@ use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builders::builder_with_capacity; use crate::dtype::IntegerPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; -use crate::match_smallest_offset_type; use crate::validity::Validity; /// Take implementation for [`FixedSizeListArray`]. /// -/// Unlike `ListView`, `FixedSizeListArray` must rebuild the elements array because it requires -/// that elements start at offset 0 and be perfectly packed without gaps. We expand list indices -/// into element indices and push them down to the child elements array. +/// `FixedSizeListArray` must rebuild its elements array because selected lists need to become +/// packed from offset 0. The FSL layer translates selected list rows into ordered element ranges +/// and delegates the execution strategy to the elements child via `take_slices`. impl TakeExecute for FixedSizeList { fn take( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let max_element_idx = array.elements().len(); - // Indices are non-negative; dispatch over the 4 unsigned widths (the executed array is - // reinterpreted to unsigned in `take_with_indices`). `E` is already unsigned. match_each_unsigned_integer_ptype!(indices.dtype().as_ptype().to_unsigned(), |I| { - match_smallest_offset_type!(max_element_idx, |E| { - take_with_indices::(array, indices, ctx) - }) + take_with_indices::(array, indices, ctx) }) .map(Some) } } -/// Dispatches to the appropriate take implementation based on list size and nullability. -fn take_with_indices( +fn take_with_indices( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, @@ -61,79 +57,57 @@ fn take_with_indices( // Reinterpret to unsigned so `as_slice::` (with unsigned `I`) matches; values are unchanged. let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned()); - // Make sure to handle degenerate case where lists have size 0 (these can take fast paths). if list_size == 0 { - debug_assert!( + vortex_ensure!( array.elements().is_empty(), "degenerate list must have empty elements" ); - // Since there are no elements to take, we just need to take on the validity map. + validate_valid_indices::(&indices_array.as_view(), array.as_ref().len(), ctx)?; let new_validity = array.validity()?.take(indices)?; let new_len = indices_array.len(); - Ok( - // SAFETY: list_size is 0, elements array is empty, and validity has the correct length. - unsafe { - FixedSizeListArray::new_unchecked( - array.elements().clone(), // Remember that this is an empty array. - array.list_size(), - new_validity, - new_len, - ) - } - .into_array(), - ) - } else { - // The result's nullability is the union of the input nullabilities. - if array.dtype().is_nullable() || indices_array.dtype().is_nullable() { - let indices_array = indices_array.as_view(); - take_nullable_fsl::(array, indices_array, ctx) - } else { - let indices_array = indices_array.as_view(); - take_non_nullable_fsl::(array, indices_array) + // SAFETY: degenerate FSL inputs have no elements, valid index payloads were checked + // against the source length, and `Validity::take` produces validity for `new_len`. + return Ok(unsafe { + FixedSizeListArray::new_unchecked( + array.elements().clone(), + array.list_size(), + new_validity, + new_len, + ) } + .into_array()); + } + + if array.dtype().is_nullable() || indices_array.dtype().is_nullable() { + take_nullable_fsl::(array, indices_array.as_view(), ctx) + } else { + take_non_nullable_fsl::(array, indices_array.as_view()) } } -/// Takes from an array when both the array and indices are non-nullable. -fn take_non_nullable_fsl( +fn take_non_nullable_fsl( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, ) -> VortexResult { let list_size = array.list_size() as usize; + let array_len = array.as_ref().len(); let indices: &[I] = indices_array.as_slice::(); let new_len = indices.len(); + let expected_elements_len = take_elements_len(new_len, list_size)?; + let mut slices = Vec::with_capacity(new_len); - // Build the element indices directly without validity tracking. - let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); - - // Build the element indices for each list. - for data_idx in indices { - let data_idx = data_idx - .to_usize() - .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx)); - - let list_start = data_idx * list_size; - let list_end = (data_idx + 1) * list_size; - - // Expand the list into individual element indices. - for i in list_start..list_end { - // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity. - unsafe { - elements_indices.push_unchecked(E::from_usize(i).vortex_expect("i < list_end")) - }; - } + for &data_idx in indices { + let data_idx = index_to_usize(data_idx); + slices.push(list_range(data_idx, list_size, array_len)?); } - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); + let new_elements = array.elements().take_slices(slices)?; + ensure_elements_len(new_elements.len(), expected_elements_len)?; - let elements_indices_array = PrimitiveArray::new(elements_indices, Validity::NonNullable); - let new_elements = array.elements().take(elements_indices_array.into_array())?; - debug_assert_eq!(new_elements.len(), new_len * list_size); - - // Both inputs are non-nullable, so the result is non-nullable. + // SAFETY: `slices` contains one checked range of `list_size` elements for each output row, + // `new_elements` has `new_len * list_size` elements, and non-nullable validity has no length. Ok(unsafe { FixedSizeListArray::new_unchecked( new_elements, @@ -145,77 +119,182 @@ fn take_non_nullable_fsl( .into_array()) } -/// Takes from an array when either the array or indices are nullable. -fn take_nullable_fsl( +fn take_nullable_fsl( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { let list_size = array.list_size() as usize; + let array_len = array.as_ref().len(); let indices: &[I] = indices_array.as_slice::(); let new_len = indices.len(); + let expected_elements_len = take_elements_len(new_len, list_size)?; let array_validity = array .fixed_size_list_validity() .execute_mask(array.as_ref().len(), ctx) .vortex_expect("Failed to compute validity mask"); - let indices_len = indices_array.as_ref().len(); - let indices_validity = match indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - { - Validity::NonNullable | Validity::AllValid => Mask::new_true(indices_len), - Validity::AllInvalid => Mask::new_false(indices_len), - Validity::Array(a) => a.execute::(ctx)?.execute_mask(ctx), - }; + let indices_validity = indices_validity_mask(&indices_array, ctx)?; - // We must use placeholder zeros for null lists to maintain the array length without - // propagating nullability to the element array's take operation. - let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); + let null_elements = null_list_elements(array, list_size); + let mut has_valid_elements = false; + let mut needs_default_elements = false; + let mut slices = Vec::with_capacity(new_len); let mut new_validity_builder = BitBufferMut::with_capacity(new_len); - // Build the element indices while tracking which lists are null. - for (data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { - let data_idx = data_idx - .to_usize() - .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx)); - - // The list is null if the index is null or the indexed element is null. - if !is_index_valid || !array_validity.value(data_idx) { - // Append placeholder zeros for null lists. These will be masked by the validity array. - // We cannot use append_nulls here as explained above. - unsafe { elements_indices.push_n_unchecked(E::zero(), list_size) }; - new_validity_builder.append(false); - } else { - // Append the actual element indices for this list. - let list_start = data_idx * list_size; - let list_end = (data_idx + 1) * list_size; - - // Expand the list into individual element indices. - for i in list_start..list_end { - // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity. - unsafe { - elements_indices.push_unchecked(E::from_usize(i).vortex_expect("i < list_end")) - }; - } - - new_validity_builder.append(true); + for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { + if !is_index_valid { + needs_default_elements |= + append_null_list_elements(null_elements, &mut slices, &mut new_validity_builder); + continue; } - } - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); + let data_idx = index_to_usize(data_idx); + let range = list_range(data_idx, list_size, array_len)?; + if !array_validity.value(data_idx) { + needs_default_elements |= + append_null_list_elements(null_elements, &mut slices, &mut new_validity_builder); + continue; + } - let elements_indices_array = PrimitiveArray::new(elements_indices, Validity::NonNullable); - let new_elements = array.elements().take(elements_indices_array.into_array())?; - debug_assert_eq!(new_elements.len(), new_len * list_size); + slices.push(range); + has_valid_elements = true; + new_validity_builder.append(true); + } + + let new_elements = if needs_default_elements { + if has_valid_elements { + vortex_bail!( + "Cannot build placeholder elements for nullable FixedSizeList take with mixed valid and null rows" + ); + } + default_elements(array, expected_elements_len) + } else { + array.elements().take_slices(slices)? + }; + ensure_elements_len(new_elements.len(), expected_elements_len)?; - // At least one input was nullable, so the result is nullable. let new_validity = Validity::from(new_validity_builder.freeze()); debug_assert!(new_validity.maybe_len().is_none_or(|vl| vl == new_len)); + // SAFETY: `new_elements` has `new_len * list_size` elements. `new_validity_builder` appends + // exactly one bit per output row, and `Validity::from` preserves that length when needed. Ok(unsafe { FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } .into_array()) } + +fn append_null_list_elements( + null_elements: NullListElements, + slices: &mut Vec<(usize, usize)>, + new_validity_builder: &mut BitBufferMut, +) -> bool { + match null_elements { + NullListElements::PlaceholderRange(range) => { + slices.push(range); + new_validity_builder.append(false); + false + } + NullListElements::DefaultsOnly => { + new_validity_builder.append(false); + true + } + } +} + +fn indices_validity_mask( + indices_array: &ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let indices_len = indices_array.as_ref().len(); + match indices_array + .validity() + .vortex_expect("Failed to compute validity mask") + { + Validity::NonNullable | Validity::AllValid => Ok(Mask::new_true(indices_len)), + Validity::AllInvalid => Ok(Mask::new_false(indices_len)), + Validity::Array(a) => Ok(a.execute::(ctx)?.execute_mask(ctx)), + } +} + +fn validate_valid_indices( + indices_array: &ArrayView<'_, Primitive>, + array_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let indices: &[I] = indices_array.as_slice::(); + let indices_validity = indices_validity_mask(indices_array, ctx)?; + + for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { + if is_index_valid { + check_index_in_bounds(index_to_usize(data_idx), array_len)?; + } + } + Ok(()) +} + +fn take_elements_len(new_len: usize, list_size: usize) -> VortexResult { + new_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" + ) + }) +} + +fn ensure_elements_len(actual: usize, expected: usize) -> VortexResult<()> { + vortex_ensure!( + actual == expected, + "FixedSizeList take elements length {actual} does not match expected length {expected}" + ); + Ok(()) +} + +fn list_range(data_idx: usize, list_size: usize, array_len: usize) -> VortexResult<(usize, usize)> { + check_index_in_bounds(data_idx, array_len)?; + + let start = data_idx.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take element range overflow for index {data_idx} and list size {list_size}" + ) + })?; + let end = start.checked_add(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take element range overflow for index {data_idx} and list size {list_size}" + ) + })?; + Ok((start, end)) +} + +fn check_index_in_bounds(data_idx: usize, array_len: usize) -> VortexResult<()> { + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); + } + Ok(()) +} + +#[derive(Clone, Copy)] +enum NullListElements { + PlaceholderRange((usize, usize)), + DefaultsOnly, +} + +fn null_list_elements(array: ArrayView<'_, FixedSizeList>, list_size: usize) -> NullListElements { + if array.elements().len() >= list_size { + NullListElements::PlaceholderRange((0, list_size)) + } else { + NullListElements::DefaultsOnly + } +} + +fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef { + let mut builder = builder_with_capacity(array.elements().dtype(), len); + builder.append_defaults(len); + builder.finish() +} + +fn index_to_usize(index: I) -> usize { + index + .to_usize() + .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", index)) +} diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index af8f94f2d5f..56c6a174e64 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -3,6 +3,8 @@ use rstest::rstest; use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; use super::common::create_basic_fsl; use super::common::create_empty_fsl; @@ -13,8 +15,13 @@ use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::DictArray; +use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; +use crate::arrays::TakeSlices; +use crate::arrays::dict::TakeExecute; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::FixedSizeListBuilder; @@ -112,6 +119,36 @@ fn test_take_degenerate_lists( } } +#[test] +fn test_take_degenerate_rejects_out_of_bounds_valid_index() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = PrimitiveArray::empty::(Nullability::NonNullable).into_array(); + let fsl = FixedSizeListArray::new(elements, 0, Validity::NonNullable, 5); + let indices = buffer![5u32].into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_take_degenerate_ignores_out_of_bounds_null_index_payload() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = PrimitiveArray::empty::(Nullability::NonNullable).into_array(); + let fsl = FixedSizeListArray::new(elements, 0, Validity::NonNullable, 5); + let indices = + PrimitiveArray::new(buffer![999u32, 1], Validity::from_iter([false, true])).into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx)? + .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; + + assert_eq!(result.len(), 2); + assert!(result.execute_scalar(0, &mut ctx)?.is_null()); + assert!(!result.execute_scalar(1, &mut ctx)?.is_null()); + Ok(()) +} + #[test] fn test_take_large_list_size() { let mut ctx = array_session().create_execution_ctx(); @@ -127,6 +164,41 @@ fn test_take_large_list_size() { assert_arrays_eq!(expected, result, &mut ctx); } +#[test] +fn test_take_range_path_large_list_size_non_nullable() { + let mut ctx = array_session().create_execution_ctx(); + let elements = PrimitiveArray::from_iter(0i32..768).into_array(); + let fsl = FixedSizeListArray::new(elements, 256, Validity::NonNullable, 3); + + let indices = buffer![2u16, 0].into_array(); + let result = fsl.take(indices).unwrap(); + + let expected_elems = PrimitiveArray::from_iter((512i32..768).chain(0..256)).into_array(); + let expected = FixedSizeListArray::new(expected_elems, 256, Validity::NonNullable, 2); + assert_arrays_eq!(expected, result, &mut ctx); +} + +#[test] +fn test_take_range_path_large_list_size_nullable() { + let mut ctx = array_session().create_execution_ctx(); + let elements = PrimitiveArray::from_iter(0i32..768).into_array(); + let fsl = FixedSizeListArray::new(elements, 256, Validity::from_iter([true, false, true]), 3); + + let indices = buffer![2u16, 1, 0].into_array(); + let result = fsl.take(indices).unwrap(); + + let expected_elems = + PrimitiveArray::from_iter((512i32..768).chain((0..256).map(|_| 0)).chain(0..256)) + .into_array(); + let expected = FixedSizeListArray::new( + expected_elems, + 256, + Validity::from_iter([true, false, true]), + 3, + ); + assert_arrays_eq!(expected, result, &mut ctx); +} + #[test] fn test_take_fsl_with_null_indices_preserves_elements() { let mut ctx = array_session().create_execution_ctx(); @@ -147,8 +219,7 @@ fn test_take_fsl_with_null_indices_preserves_elements() { assert_arrays_eq!(expected, result, &mut ctx); } -// Element index overflow: with u8 indices and list_size=16, data_idx=16 produces element index -// 16*16=256 which overflows u8. The take kernel must widen the element index type. +// List offsets must not truncate when small index types select large lists. #[rstest] #[case::non_nullable( FixedSizeListArray::new( @@ -181,24 +252,164 @@ fn test_element_index_overflow( assert_arrays_eq!(result, expected, &mut ctx); } +#[test] +fn test_take_nullable_indices_ignores_out_of_bounds_null_value() { + let mut ctx = array_session().create_execution_ctx(); + let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); + let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); + + let indices = PrimitiveArray::new( + buffer![1u64, 999, 0], + Validity::from_iter([true, false, true]), + ); + let result = fsl.take(indices.into_array()).unwrap(); + + let expected = FixedSizeListArray::new( + buffer![3i32, 4, 0, 0, 1, 2].into_array(), + 2, + Validity::from_iter([true, false, true]), + 3, + ); + assert_arrays_eq!(expected, result, &mut ctx); +} + +#[test] +fn test_take_rejects_overflowing_valid_index() { + let mut ctx = array_session().create_execution_ctx(); + let elements = buffer![1i32, 2, 3, 4].into_array(); + let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 2); + let overflowing_index = (usize::MAX / 2 + 1) as u64; + let indices = buffer![overflowing_index].into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx); + + assert!(result.is_err()); +} + +#[test] +fn test_take_nullable_fsl_with_nullable_indices() { + let mut ctx = array_session().create_execution_ctx(); + let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); + let fsl = FixedSizeListArray::new( + elements.into_array(), + 2, + Validity::from_iter([true, false, true]), + 3, + ); + + let indices = PrimitiveArray::new( + buffer![2u64, 999, 1, 0], + Validity::from_iter([true, false, true, true]), + ); + let result = fsl.take(indices.into_array()).unwrap(); + + let expected = FixedSizeListArray::new( + buffer![5i32, 6, 0, 0, 0, 0, 1, 2].into_array(), + 2, + Validity::from_iter([true, false, false, true]), + 4, + ); + assert_arrays_eq!(expected, result, &mut ctx); +} + +#[test] +fn test_take_empty_source_with_all_null_indices() { + let fsl = create_empty_fsl(); + let indices = PrimitiveArray::new(buffer![999u64, 123], Validity::AllInvalid); + + let result = fsl.take(indices.into_array()).unwrap(); + + assert_eq!(result.len(), 2); + for idx in 0..result.len() { + assert!( + result + .execute_scalar(idx, &mut array_session().create_execution_ctx()) + .unwrap() + .is_null() + ); + } +} + +#[test] +fn test_take_execute_empty_source_all_null_indices_builds_default_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let fsl = create_empty_fsl(); + let indices = PrimitiveArray::new(buffer![999u64, 123], Validity::AllInvalid).into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx)? + .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; + let result_fsl = result.as_::(); + + assert_eq!( + result_fsl.elements().len(), + result.len() * result_fsl.list_size() as usize + ); + for idx in 0..result.len() { + assert!(result.execute_scalar(idx, &mut ctx)?.is_null()); + } + Ok(()) +} + +#[test] +fn test_take_uses_take_slices_for_encoded_elements_child() { + let mut ctx = array_session().create_execution_ctx(); + let dict_elements = DictArray::try_new( + buffer![0u8, 1, 2, 3, 1, 2].into_array(), + PrimitiveArray::from_iter([10i32, 20, 30, 40]).into_array(), + ) + .unwrap() + .into_array(); + let fsl = FixedSizeListArray::new(dict_elements, 2, Validity::NonNullable, 3); + + let result = fsl.take(buffer![2u8, 0].into_array()).unwrap(); + + let executed = result + .clone() + .execute_until::(&mut ctx) + .unwrap(); + assert!( + executed + .as_::() + .elements() + .is::() + ); + let expected = FixedSizeListArray::new( + PrimitiveArray::from_iter([20i32, 30, 10, 20]).into_array(), + 2, + Validity::NonNullable, + 2, + ); + assert_arrays_eq!(expected, result, &mut ctx); +} + // Parameterized test for nullable array scenarios that are specific to FSL's implementation. #[rstest] #[case::nullable_mixed_elements( vec![Some(vec![1i32, 2]), None, Some(vec![5, 6])], vec![Some(2u32), Some(1), Some(0)], - vec![false, true, false] + vec![Some(vec![5i32, 6]), None, Some(vec![1, 2])] )] #[case::nullable_with_null_indices( vec![Some(vec![1i32, 2]), None, Some(vec![5, 6])], vec![Some(0u32), None, Some(1), Some(2)], - vec![false, true, true, false] + vec![Some(vec![1i32, 2]), None, None, Some(vec![5, 6])] )] fn test_take_nullable_arrays_fsl_specific( #[case] array_values: Vec>>, #[case] indices: Vec>, - #[case] expected_nulls: Vec, + #[case] expected_values: Vec>>, ) { - // Build the nullable FSL array. + let mut ctx = array_session().create_execution_ctx(); + let fsl = nullable_i32_fsl(array_values); + + let indices_array = PrimitiveArray::from_option_iter(indices); + let result = fsl.take(indices_array.into_array()).unwrap(); + let expected = nullable_i32_fsl(expected_values); + + assert_arrays_eq!(expected, result, &mut ctx); +} + +fn nullable_i32_fsl(array_values: Vec>>) -> ArrayRef { let list_size = if let Some(Some(first)) = array_values.first() { u32::try_from(first.len()).unwrap() } else { @@ -231,20 +442,5 @@ fn test_take_nullable_arrays_fsl_specific( } } - let fsl = builder.finish(); - - // Create indices (with possible nulls). - let indices_array = PrimitiveArray::from_option_iter(indices.clone()); - let result = fsl.take(indices_array.into_array()).unwrap(); - - assert_eq!(result.len(), indices.len()); - for (i, expected_null) in expected_nulls.iter().enumerate() { - assert_eq!( - result - .execute_scalar(i, &mut array_session().create_execution_ctx()) - .unwrap() - .is_null(), - *expected_null - ); - } + builder.finish() } diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 15b6b5aa6be..111a0b76722 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -104,6 +104,10 @@ pub mod slice; pub use slice::Slice; pub use slice::SliceArray; +pub mod take_slices; +pub use take_slices::TakeSlices; +pub use take_slices::TakeSlicesArray; + pub mod struct_; pub use struct_::Struct; pub use struct_::StructArray; diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index 1ca4b17d3d5..37aa3f98396 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -8,6 +8,7 @@ mod mask; pub(crate) mod rules; mod slice; mod take; +mod take_slices; mod zip; #[cfg(test)] diff --git a/vortex-array/src/arrays/primitive/compute/take_slices.rs b/vortex-array/src/arrays/primitive/compute/take_slices.rs new file mode 100644 index 00000000000..ccef3ef5bde --- /dev/null +++ b/vortex-array/src/arrays/primitive/compute/take_slices.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::take_slices::TakeSlicesExecute; +use crate::executor::ExecutionCtx; +use crate::match_each_native_ptype; + +impl TakeSlicesExecute for Primitive { + fn take_slices( + array: ArrayView<'_, Self>, + slices: &[(usize, usize)], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let validity = array.validity()?.take_slices(slices)?; + match_each_native_ptype!(array.ptype(), |T| { + let source = array.as_slice::(); + let len = slices.iter().map(|(start, end)| end - start).sum(); + let mut values = BufferMut::::with_capacity(len); + + for &(start, end) in slices { + values.extend_from_slice(&source[start..end]); + } + + Ok(Some( + PrimitiveArray::new(values.freeze(), validity).into_array(), + )) + }) + } +} diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 6382ea73794..663c8924865 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -6,7 +6,9 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::Primitive; +use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; @@ -32,5 +34,10 @@ pub(crate) fn initialize(session: &VortexSession) { FillNullExecuteAdaptor(Primitive), ); kernels.register_execute_parent_kernel(Dict.id(), Primitive, TakeExecuteAdaptor(Primitive)); + kernels.register_execute_parent_kernel( + TakeSlices.id(), + Primitive, + TakeSlicesExecuteAdaptor(Primitive), + ); kernels.register_execute_parent_kernel(Zip.id(), Primitive, ZipExecuteAdaptor(Primitive)); } diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs new file mode 100644 index 00000000000..b108d52b2f9 --- /dev/null +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::ops::Range; +use std::sync::Arc; + +use smallvec::smallvec; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::array::Array; +use crate::array::ArrayParts; +use crate::array::TypedArrayRef; +use crate::arrays::TakeSlices; + +/// The child array being selected by ordered slices. +pub(super) const CHILD_SLOT: usize = 0; +pub(super) const NUM_SLOTS: usize = 1; +pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child"]; + +/// Metadata for a [`TakeSlices`](crate::arrays::TakeSlices) array. +#[derive(Clone, Debug)] +pub struct TakeSlicesData { + pub(super) slices: Arc<[(usize, usize)]>, + len: usize, +} + +impl Display for TakeSlicesData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "nslices: {}, len: {}", self.slices.len(), self.len()) + } +} + +/// Extension methods for [`TakeSlices`](crate::arrays::TakeSlices) arrays. +pub trait TakeSlicesArrayExt: TypedArrayRef { + /// The child array selected by this ordered range list. + fn child(&self) -> &ArrayRef { + self.as_ref().slots()[CHILD_SLOT] + .as_ref() + .vortex_expect("validated take-slices child slot") + } + + /// The ordered, non-empty child ranges represented by this array. + fn slices(&self) -> &[(usize, usize)] { + &self.slices + } +} +impl> TakeSlicesArrayExt for T {} + +impl TakeSlicesData { + fn try_new(child_len: usize, slices: Vec<(usize, usize)>) -> VortexResult { + let mut len = 0usize; + for &(start, end) in &slices { + vortex_ensure!( + start < end, + "TakeSlicesArray range must be non-empty: {start}..{end}" + ); + vortex_ensure!( + end <= child_len, + "TakeSlicesArray range {start}..{end} exceeds child array length {child_len}" + ); + len = len + .checked_add(end - start) + .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; + } + + Ok(Self { + slices: Arc::from(slices.into_boxed_slice()), + len, + }) + } + + /// Returns the length of this array. + pub fn len(&self) -> usize { + self.len + } + + /// Returns `true` if this array is empty. + pub fn is_empty(&self) -> bool { + self.slices.is_empty() + } + + /// The ordered ranges used to select child values. + pub fn slices(&self) -> &[(usize, usize)] { + &self.slices + } + + /// Returns the ordered ranges as `Range` values. + pub fn slice_ranges(&self) -> impl Iterator> + '_ { + self.slices.iter().map(|&(start, end)| start..end) + } +} + +impl Array { + /// Constructs a new validated `TakeSlicesArray`. + pub fn try_new(child: ArrayRef, slices: Vec<(usize, usize)>) -> VortexResult { + let dtype = child.dtype().clone(); + let data = TakeSlicesData::try_new(child.len(), slices)?; + let len = data.len(); + + // SAFETY: `TakeSlicesData::try_new` validates range bounds and computes `len`; the outer + // dtype is copied from the child, and the required child slot is populated. + Ok(unsafe { + Array::from_parts_unchecked( + ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![Some(child)]), + ) + }) + } +} diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs new file mode 100644 index 00000000000..9957adec546 --- /dev/null +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reduce and execute adaptors for ordered slice-take operations. +//! +//! `TakeSlicesArray` represents concatenating an ordered list of non-empty child ranges. The +//! ranges preserve caller order and may overlap. Encodings that know how to serve those ranges +//! efficiently implement [`TakeSlicesReduce`] or [`TakeSlicesExecute`]. + +mod array; +mod rules; +mod vtable; + +pub use array::TakeSlicesArrayExt; +pub use array::TakeSlicesData; +use vortex_error::VortexResult; +pub use vtable::*; + +use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::kernel::ExecuteParentKernel; +use crate::matcher::Matcher; +use crate::optimizer::rules::ArrayParentReduceRule; + +/// Metadata-only implementation hook for taking ordered child ranges. +pub trait TakeSlicesReduce: VTable { + /// Take ordered slices from an array without reading buffers. + /// + /// Implementations should return `None` if serving the ranges requires buffer access. + /// + /// # Preconditions + /// + /// `slices` is guaranteed to contain only non-empty ranges in bounds for `array`. + fn take_slices( + array: ArrayView<'_, Self>, + slices: &[(usize, usize)], + ) -> VortexResult>; +} + +/// Execution implementation hook for taking ordered child ranges. +pub trait TakeSlicesExecute: VTable { + /// Take ordered slices from an array, potentially reading buffers. + /// + /// # Preconditions + /// + /// `slices` is guaranteed to contain only non-empty ranges in bounds for `array`. + fn take_slices( + array: ArrayView<'_, Self>, + slices: &[(usize, usize)], + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} + +fn trivial_take_slices( + child: &ArrayRef, + slices: &[(usize, usize)], +) -> VortexResult> { + if slices.is_empty() { + return Ok(Some(Canonical::empty(child.dtype()).into_array())); + } + + if let [(start, end)] = slices { + if *start == 0 && *end == child.len() { + return Ok(Some(child.clone())); + } + return child.slice(*start..*end).map(Some); + } + + Ok(None) +} + +/// Adaptor that wraps a [`TakeSlicesReduce`] impl as an [`ArrayParentReduceRule`]. +#[derive(Default, Debug)] +pub struct TakeSlicesReduceAdaptor(pub V); + +impl ArrayParentReduceRule for TakeSlicesReduceAdaptor +where + V: TakeSlicesReduce, +{ + type Parent = TakeSlices; + + fn reduce_parent( + &self, + array: ArrayView<'_, V>, + parent: ::Match<'_>, + child_idx: usize, + ) -> VortexResult> { + assert_eq!(child_idx, 0); + if let Some(result) = trivial_take_slices(array.array(), parent.slices())? { + return Ok(Some(result)); + } + ::take_slices(array, parent.slices()) + } +} + +/// Adaptor that wraps a [`TakeSlicesExecute`] impl as an [`ExecuteParentKernel`]. +#[derive(Default, Debug)] +pub struct TakeSlicesExecuteAdaptor(pub V); + +impl ExecuteParentKernel for TakeSlicesExecuteAdaptor +where + V: TakeSlicesExecute, +{ + type Parent = TakeSlices; + + fn execute_parent( + &self, + array: ArrayView<'_, V>, + parent: ::Match<'_>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + assert_eq!(child_idx, 0); + if let Some(result) = trivial_take_slices(array.array(), parent.slices())? { + return Ok(Some(result)); + } + ::take_slices(array, parent.slices(), ctx) + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/arrays/take_slices/rules.rs b/vortex-array/src/arrays/take_slices/rules.rs new file mode 100644 index 00000000000..5a4d350f08f --- /dev/null +++ b/vortex-array/src/arrays/take_slices/rules.rs @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::trivial_take_slices; +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::TakeSlices; +use crate::arrays::TakeSlicesArray; +use crate::arrays::take_slices::TakeSlicesArrayExt; +use crate::arrays::take_slices::TakeSlicesReduce; +use crate::arrays::take_slices::TakeSlicesReduceAdaptor; +use crate::optimizer::rules::ArrayReduceRule; +use crate::optimizer::rules::ParentRuleSet; +use crate::optimizer::rules::ReduceRuleSet; + +pub(super) const PARENT_RULES: ParentRuleSet = + ParentRuleSet::new(&[ParentRuleSet::lift(&TakeSlicesReduceAdaptor(TakeSlices))]); + +pub(super) const RULES: ReduceRuleSet = ReduceRuleSet::new(&[&TrivialTakeSlicesRule]); + +#[derive(Debug)] +struct TrivialTakeSlicesRule; + +impl ArrayReduceRule for TrivialTakeSlicesRule { + fn reduce(&self, array: ArrayView<'_, TakeSlices>) -> VortexResult> { + trivial_take_slices(array.child(), array.slices()) + } +} + +impl TakeSlicesReduce for TakeSlices { + fn take_slices( + array: ArrayView<'_, Self>, + slices: &[(usize, usize)], + ) -> VortexResult> { + let combined = project_slices(array.slices(), slices); + Ok(Some( + TakeSlicesArray::try_new(array.child().clone(), combined)?.into_array(), + )) + } +} + +fn project_slices(inner: &[(usize, usize)], outer: &[(usize, usize)]) -> Vec<(usize, usize)> { + let mut projected = Vec::new(); + + for &(outer_start, outer_end) in outer { + let mut logical_start = 0usize; + for &(inner_start, inner_end) in inner { + let inner_len = inner_end - inner_start; + let logical_end = logical_start + inner_len; + + if outer_start < logical_end && outer_end > logical_start { + let overlap_start = outer_start.max(logical_start); + let overlap_end = outer_end.min(logical_end); + projected.push(( + inner_start + (overlap_start - logical_start), + inner_start + (overlap_end - logical_start), + )); + } + + if logical_end >= outer_end { + break; + } + logical_start = logical_end; + } + } + + projected +} diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs new file mode 100644 index 00000000000..47d425953b2 --- /dev/null +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -0,0 +1,205 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::DictArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::TakeSlices; +use crate::arrays::TakeSlicesArray; +use crate::arrays::take_slices::TakeSlicesArrayExt; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; +use crate::assert_arrays_eq; +use crate::dtype::Nullability; +use crate::kernel::ExecuteParentKernel; + +#[test] +fn take_slices_preserves_order_duplicates_and_overlap() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0i32..8).into_array(); + + let actual = array.take_slices(vec![(4, 6), (1, 3), (4, 6), (2, 5)])?; + let expected = PrimitiveArray::from_iter([4i32, 5, 1, 2, 4, 5, 2, 3, 4]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn take_slices_preserves_nullable_child_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = + PrimitiveArray::from_option_iter([Some(0i32), None, Some(2), Some(3), None, Some(5)]) + .into_array(); + + let actual = array.take_slices(vec![(1, 4), (4, 6)])?; + let expected = PrimitiveArray::from_option_iter([None, Some(2), Some(3), None, Some(5)]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn take_slices_lazy_scalar_and_validity_follow_ranges() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = + PrimitiveArray::from_option_iter([Some(0i32), None, Some(2), Some(3), None, Some(5)]) + .into_array(); + + let actual = array.take_slices(vec![(1, 4), (4, 6)])?; + let validity = actual.validity()?.execute_mask(actual.len(), &mut ctx)?; + + assert!(!validity.value(0)); + assert!(validity.value(1)); + assert_eq!( + actual.execute_scalar(2, &mut ctx)?, + array.execute_scalar(3, &mut ctx)? + ); + assert!(!validity.value(3)); + assert_eq!( + actual.execute_scalar(4, &mut ctx)?, + array.execute_scalar(5, &mut ctx)? + ); + Ok(()) +} + +#[test] +fn take_slices_one_slice_reduces_to_child_slice() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0i32..6).into_array(); + + let actual = array.take_slices(vec![(1, 5)])?; + + assert!(actual.is::()); + assert!(!actual.is::()); + assert_arrays_eq!(actual, PrimitiveArray::from_iter([1i32, 2, 3, 4]), &mut ctx); + Ok(()) +} + +#[test] +fn take_slices_size_one_child_can_repeat_the_only_range() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter([7i32]).into_array(); + + let actual = array.take_slices(vec![(0, 1), (0, 1)])?; + let expected = PrimitiveArray::from_iter([7i32, 7]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn primitive_take_slices_execute_parent_copies_ranges_directly() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0i32..8).into_array(); + let parent = TakeSlicesArray::try_new(array.clone(), vec![(3, 6), (0, 2)])?.into_array(); + + let actual = TakeSlicesExecuteAdaptor(Primitive) + .execute_parent( + array.as_::(), + parent.as_::(), + 0, + &mut ctx, + )? + .ok_or_else(|| vortex_err!("Primitive TakeSlicesExecute declined multi-slice take"))?; + + assert!(actual.is::()); + assert_arrays_eq!( + actual, + PrimitiveArray::from_iter([3i32, 4, 5, 0, 1]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn take_slices_empty_ranges_return_empty_canonical() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0i32..6).into_array(); + + let actual = array.take_slices(vec![])?; + let expected = PrimitiveArray::from_iter(std::iter::empty::()); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn take_slices_empty_child_accepts_empty_ranges() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::empty::(Nullability::NonNullable).into_array(); + + let actual = array.take_slices(vec![])?; + let expected = PrimitiveArray::from_iter(std::iter::empty::()); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn take_slices_rejects_invalid_ranges() { + let array = PrimitiveArray::from_iter(0i32..6).into_array(); + + assert!(array.take_slices(vec![(2, 2)]).is_err()); + assert!(array.take_slices(vec![(4, 3)]).is_err()); + assert!(array.take_slices(vec![(4, 7)]).is_err()); +} + +#[test] +fn take_slices_of_take_slices_projects_to_original_child() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0i32..10).into_array(); + + let inner = array.take_slices(vec![(2, 5), (7, 10)])?; + let actual = inner.take_slices(vec![(1, 4), (4, 6)])?; + + let actual_take_slices = actual.as_::(); + assert!(actual_take_slices.child().is::()); + assert_eq!(actual_take_slices.slices(), &[(3, 5), (7, 8), (8, 10)]); + assert_arrays_eq!( + actual, + PrimitiveArray::from_iter([3i32, 4, 7, 8, 9]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn take_slices_generic_execution_handles_child_without_take_slices_kernel() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dict = DictArray::try_new( + buffer![2u8, 0, 1, 2, 0, 1].into_array(), + PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), + )? + .into_array(); + + let actual = dict.take_slices(vec![(1, 4), (0, 2)])?; + let expected = PrimitiveArray::from_iter([10i32, 20, 30, 30, 10]); + + assert!(actual.is::()); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn take_slices_generic_execution_preserves_nullable_encoded_child() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dict = DictArray::try_new( + buffer![0u8, 1, 2, 0].into_array(), + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array(), + )? + .into_array(); + + let actual = dict.take_slices(vec![(1, 3), (0, 2)])?; + let expected = PrimitiveArray::from_option_iter([None, Some(30), Some(10), None]); + + assert!(actual.is::()); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs new file mode 100644 index 00000000000..2be7c978224 --- /dev/null +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -0,0 +1,208 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::hash::Hash; +use std::hash::Hasher; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayEq; +use crate::ArrayHash; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::EqMode; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; +use crate::arrays::take_slices::TakeSlicesArrayExt; +use crate::arrays::take_slices::array::CHILD_SLOT; +use crate::arrays::take_slices::array::NUM_SLOTS; +use crate::arrays::take_slices::array::SLOT_NAMES; +use crate::arrays::take_slices::array::TakeSlicesData; +use crate::arrays::take_slices::rules::PARENT_RULES; +use crate::arrays::take_slices::rules::RULES; +use crate::buffer::BufferHandle; +use crate::builders::builder_with_capacity_in; +use crate::dtype::DType; +use crate::executor::ExecutionCtx; +use crate::executor::ExecutionResult; +use crate::scalar::Scalar; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +/// A [`TakeSlices`]-encoded Vortex array. +pub type TakeSlicesArray = Array; + +/// Ordered child-range selection encoding. +#[derive(Clone, Debug)] +pub struct TakeSlices; + +impl ArrayHash for TakeSlicesData { + fn array_hash(&self, state: &mut H, _accuracy: EqMode) { + self.slices.hash(state); + } +} + +impl ArrayEq for TakeSlicesData { + fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { + self.slices == other.slices + } +} + +impl VTable for TakeSlices { + type TypedArrayData = TakeSlicesData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.take_slices"); + *ID + } + + fn validate( + &self, + data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + slots.len() == NUM_SLOTS, + "TakeSlicesArray expected {NUM_SLOTS} slots, found {}", + slots.len() + ); + vortex_ensure!( + slots[CHILD_SLOT].is_some(), + "TakeSlicesArray child slot must be present" + ); + let child = slots[CHILD_SLOT] + .as_ref() + .vortex_expect("validated child slot"); + vortex_ensure!( + child.dtype() == dtype, + "TakeSlicesArray dtype {} does not match outer dtype {}", + child.dtype(), + dtype + ); + vortex_ensure!( + data.len() == len, + "TakeSlicesArray length {} does not match outer length {}", + data.len(), + len + ); + for &(start, end) in data.slices() { + vortex_ensure!( + start < end, + "TakeSlicesArray range must be non-empty: {start}..{end}" + ); + vortex_ensure!( + end <= child.len(), + "TakeSlicesArray range {start}..{end} exceeds child length {}", + child.len() + ); + } + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle { + vortex_panic!("TakeSlicesArray has no buffers") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + SLOT_NAMES[idx].to_string() + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + vortex_bail!("TakeSlices array is not serializable") + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("TakeSlices array is not serializable") + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); + for range in array.slice_ranges() { + let slice = array.child().slice(range)?; + slice.append_to_builder(builder.as_mut(), ctx)?; + } + Ok(ExecutionResult::done(builder.finish())) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } + + fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { + RULES.evaluate(array) + } +} + +impl OperationsVTable for TakeSlices { + fn scalar_at( + array: ArrayView<'_, TakeSlices>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let mut logical_start = 0usize; + for &(start, end) in array.slices() { + let len = end - start; + let logical_end = logical_start + len; + if index < logical_end { + return array + .child() + .execute_scalar(start + (index - logical_start), ctx); + } + logical_start = logical_end; + } + + vortex_panic!("TakeSlicesArray scalar index {index} out of bounds") + } +} + +impl ValidityVTable for TakeSlices { + fn validity(array: ArrayView<'_, TakeSlices>) -> VortexResult { + array.child().validity()?.take_slices(array.slices()) + } +} diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 2cb8414cbfd..77320ef2460 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -27,6 +27,7 @@ use crate::arrays::Masked; use crate::arrays::Null; use crate::arrays::Primitive; use crate::arrays::Struct; +use crate::arrays::TakeSlices; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::Variant; @@ -81,6 +82,7 @@ impl Default for ArraySession { this.register(Dict); this.register(List); this.register(Masked); + this.register(TakeSlices); this.register(VarBin); this diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 3bb47eda0ae..3a999d0f0ef 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -224,6 +224,14 @@ impl Validity { } } + /// Select validity values by concatenating ordered, non-empty ranges. + pub fn take_slices(&self, slices: &[(usize, usize)]) -> VortexResult { + match self { + v @ (Self::NonNullable | Self::AllValid | Self::AllInvalid) => Ok(v.clone()), + Self::Array(is_valid) => Ok(Self::Array(is_valid.take_slices(slices.to_vec())?)), + } + } + // Invert the validity pub fn not(&self) -> VortexResult { match self { From 470acfa5494c75fc4e63f7c71def779631fddcc6 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 7 Jul 2026 16:50:59 -0400 Subject: [PATCH 02/37] Address TakeSlices review feedback Refactor fixed-size-list take around structural cases, tighten TakeSlices validation, and extend edge-case coverage. Also document the local-reasoning guideline that would have avoided threading fallback state through the main path. Signed-off-by: Daniel King --- AGENTS.md | 16 ++ .../arrays/fixed_size_list/compute/take.rs | 178 ++++++++-------- .../src/arrays/fixed_size_list/tests/take.rs | 198 ++++++++++++++++-- .../arrays/primitive/compute/take_slices.rs | 6 +- vortex-array/src/arrays/take_slices/array.rs | 4 +- vortex-array/src/arrays/take_slices/tests.rs | 25 +++ vortex-array/src/arrays/take_slices/vtable.rs | 23 +- 7 files changed, 332 insertions(+), 118 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 30ddeeb9515..4c85be00da5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -151,6 +151,22 @@ Notes: self-explanatory code. - Keep public APIs small and consistent with neighboring crates. +## Local Reasoning and Edge Cases + +- Before adding helper state, boolean flags, or small enums to handle edge cases inside a general + loop, first ask whether the input space naturally splits into a few structural cases with + different invariants. If it does, dispatch once near the top of the function to clearly named + helpers for each case. +- Prefer helpers whose names state the invariant they rely on and whose bodies only handle that + case. A reader should be able to prove each helper correct without carrying unrelated cases in + their head. +- Avoid making a broad path responsible for exceptional setup, fallback construction, and normal + processing at the same time. If a fallback is only valid under a narrow precondition, isolate it + behind that precondition instead of threading flags through the main algorithm. +- When a branch exists only to preserve a representation invariant, put that invariant at the + branch boundary and keep the subsequent code free of defensive state that is impossible under + the branch's preconditions. + ## Performance: avoid hidden-cost accessors in hot loops The most common performance trap in this codebase is calling a *per-element accessor that diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 0a397b7b807..92a851646df 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -2,12 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::BitBufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; @@ -58,36 +56,90 @@ fn take_with_indices( let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned()); if list_size == 0 { - vortex_ensure!( - array.elements().is_empty(), - "degenerate list must have empty elements" - ); - - validate_valid_indices::(&indices_array.as_view(), array.as_ref().len(), ctx)?; - let new_validity = array.validity()?.take(indices)?; - let new_len = indices_array.len(); - - // SAFETY: degenerate FSL inputs have no elements, valid index payloads were checked - // against the source length, and `Validity::take` produces validity for `new_len`. - return Ok(unsafe { - FixedSizeListArray::new_unchecked( - array.elements().clone(), - array.list_size(), - new_validity, - new_len, - ) - } - .into_array()); + return take_degenerate_fsl::(array, indices, indices_array.as_view(), ctx); + } + + if array.is_empty() { + return take_empty_non_degenerate_fsl::(array, indices, indices_array.as_view(), ctx); + } + + take_non_empty_non_degenerate_fsl::(array, indices_array.as_view(), ctx) +} + +fn take_degenerate_fsl( + array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, + indices_array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + vortex_ensure!( + array.elements().is_empty(), + "degenerate list must have empty elements" + ); + + validate_valid_indices::(&indices_array, array.as_ref().len(), ctx)?; + let new_validity = array.validity()?.take(indices)?; + let new_len = indices_array.len(); + + // SAFETY: degenerate FSL inputs have no elements, valid index payloads were checked against + // the source length, and `Validity::take` produces validity for `new_len`. + Ok(unsafe { + FixedSizeListArray::new_unchecked( + array.elements().clone(), + array.list_size(), + new_validity, + new_len, + ) + } + .into_array()) +} + +fn take_empty_non_degenerate_fsl( + array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, + indices_array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + debug_assert_ne!(array.list_size(), 0); + debug_assert!(array.is_empty()); + + validate_valid_indices::(&indices_array, 0, ctx)?; + + let list_size = array.list_size() as usize; + let new_len = indices_array.len(); + let expected_elements_len = take_elements_len(new_len, list_size)?; + let new_elements = default_elements(array, expected_elements_len); + ensure_elements_len(new_elements.len(), expected_elements_len)?; + let new_validity = if new_len == 0 { + array.validity()?.take(indices)? + } else { + Validity::AllInvalid + }; + + // SAFETY: an empty non-degenerate source has no usable element slices. Valid indices have + // already been rejected, so non-empty output is all null and backed by default child values. + Ok(unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } + .into_array()) +} + +fn take_non_empty_non_degenerate_fsl( + array: ArrayView<'_, FixedSizeList>, + indices_array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + debug_assert_ne!(array.list_size(), 0); + debug_assert!(!array.is_empty()); if array.dtype().is_nullable() || indices_array.dtype().is_nullable() { - take_nullable_fsl::(array, indices_array.as_view(), ctx) + take_nullable_non_empty_fsl::(array, indices_array, ctx) } else { - take_non_nullable_fsl::(array, indices_array.as_view()) + take_non_nullable_non_empty_fsl::(array, indices_array) } } -fn take_non_nullable_fsl( +fn take_non_nullable_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, ) -> VortexResult { @@ -99,7 +151,7 @@ fn take_non_nullable_fsl( let mut slices = Vec::with_capacity(new_len); for &data_idx in indices { - let data_idx = index_to_usize(data_idx); + let data_idx = index_to_usize(data_idx)?; slices.push(list_range(data_idx, list_size, array_len)?); } @@ -119,7 +171,7 @@ fn take_non_nullable_fsl( .into_array()) } -fn take_nullable_fsl( +fn take_nullable_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, @@ -132,46 +184,33 @@ fn take_nullable_fsl( let array_validity = array .fixed_size_list_validity() - .execute_mask(array.as_ref().len(), ctx) - .vortex_expect("Failed to compute validity mask"); + .execute_mask(array.as_ref().len(), ctx)?; let indices_validity = indices_validity_mask(&indices_array, ctx)?; - let null_elements = null_list_elements(array, list_size); - let mut has_valid_elements = false; - let mut needs_default_elements = false; + let null_elements = (0, list_size); let mut slices = Vec::with_capacity(new_len); let mut new_validity_builder = BitBufferMut::with_capacity(new_len); for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { if !is_index_valid { - needs_default_elements |= - append_null_list_elements(null_elements, &mut slices, &mut new_validity_builder); + slices.push(null_elements); + new_validity_builder.append(false); continue; } - let data_idx = index_to_usize(data_idx); + let data_idx = index_to_usize(data_idx)?; let range = list_range(data_idx, list_size, array_len)?; if !array_validity.value(data_idx) { - needs_default_elements |= - append_null_list_elements(null_elements, &mut slices, &mut new_validity_builder); + slices.push(null_elements); + new_validity_builder.append(false); continue; } slices.push(range); - has_valid_elements = true; new_validity_builder.append(true); } - let new_elements = if needs_default_elements { - if has_valid_elements { - vortex_bail!( - "Cannot build placeholder elements for nullable FixedSizeList take with mixed valid and null rows" - ); - } - default_elements(array, expected_elements_len) - } else { - array.elements().take_slices(slices)? - }; + let new_elements = array.elements().take_slices(slices)?; ensure_elements_len(new_elements.len(), expected_elements_len)?; let new_validity = Validity::from(new_validity_builder.freeze()); @@ -185,33 +224,12 @@ fn take_nullable_fsl( .into_array()) } -fn append_null_list_elements( - null_elements: NullListElements, - slices: &mut Vec<(usize, usize)>, - new_validity_builder: &mut BitBufferMut, -) -> bool { - match null_elements { - NullListElements::PlaceholderRange(range) => { - slices.push(range); - new_validity_builder.append(false); - false - } - NullListElements::DefaultsOnly => { - new_validity_builder.append(false); - true - } - } -} - fn indices_validity_mask( indices_array: &ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { let indices_len = indices_array.as_ref().len(); - match indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - { + match indices_array.validity()? { Validity::NonNullable | Validity::AllValid => Ok(Mask::new_true(indices_len)), Validity::AllInvalid => Ok(Mask::new_false(indices_len)), Validity::Array(a) => Ok(a.execute::(ctx)?.execute_mask(ctx)), @@ -228,7 +246,7 @@ fn validate_valid_indices( for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { if is_index_valid { - check_index_in_bounds(index_to_usize(data_idx), array_len)?; + check_index_in_bounds(index_to_usize(data_idx)?, array_len)?; } } Ok(()) @@ -273,28 +291,14 @@ fn check_index_in_bounds(data_idx: usize, array_len: usize) -> VortexResult<()> Ok(()) } -#[derive(Clone, Copy)] -enum NullListElements { - PlaceholderRange((usize, usize)), - DefaultsOnly, -} - -fn null_list_elements(array: ArrayView<'_, FixedSizeList>, list_size: usize) -> NullListElements { - if array.elements().len() >= list_size { - NullListElements::PlaceholderRange((0, list_size)) - } else { - NullListElements::DefaultsOnly - } -} - fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef { let mut builder = builder_with_capacity(array.elements().dtype(), len); builder.append_defaults(len); builder.finish() } -fn index_to_usize(index: I) -> usize { +fn index_to_usize(index: I) -> VortexResult { index .to_usize() - .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", index)) + .ok_or_else(|| vortex_err!("FixedSizeList take index {index} does not fit in usize")) } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index 56c6a174e64..11bfcf218fa 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -2,9 +2,15 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use rstest::rstest; +use smallvec::smallvec; use vortex_buffer::buffer; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use super::common::create_basic_fsl; use super::common::create_empty_fsl; @@ -12,10 +18,18 @@ use super::common::create_large_fsl; use super::common::create_nullable_fsl; use super::common::create_single_element_fsl; use crate::ArrayRef; +use crate::ArrayView; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::EmptyArrayData; +use crate::array::OperationsVTable; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; use crate::array_session; -use crate::arrays::DictArray; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; @@ -23,13 +37,17 @@ use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::assert_arrays_eq; +use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::FixedSizeListBuilder; use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::executor::ExecutionCtx; +use crate::executor::ExecutionResult; use crate::scalar::Scalar; +use crate::serde::ArrayChildren; use crate::validity::Validity; // Conformance tests for common take scenarios. @@ -351,28 +369,29 @@ fn test_take_execute_empty_source_all_null_indices_builds_default_elements() -> } #[test] -fn test_take_uses_take_slices_for_encoded_elements_child() { +fn test_take_empty_source_rejects_valid_index_after_null_index() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let dict_elements = DictArray::try_new( - buffer![0u8, 1, 2, 3, 1, 2].into_array(), - PrimitiveArray::from_iter([10i32, 20, 30, 40]).into_array(), - ) - .unwrap() - .into_array(); - let fsl = FixedSizeListArray::new(dict_elements, 2, Validity::NonNullable, 3); - - let result = fsl.take(buffer![2u8, 0].into_array()).unwrap(); - - let executed = result - .clone() - .execute_until::(&mut ctx) - .unwrap(); - assert!( - executed - .as_::() - .elements() - .is::() - ); + let fsl = create_empty_fsl(); + let indices = + PrimitiveArray::new(buffer![999u64, 0], Validity::from_iter([false, true])).into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_take_uses_take_slices_for_encoded_elements_child() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let encoded_elements = NoTakeSlicesArray::wrap(buffer![10i32, 20, 30, 40, 20, 30].into_array()); + let fsl = FixedSizeListArray::new(encoded_elements, 2, Validity::NonNullable, 3); + let indices = buffer![2u8, 0].into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx)? + .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; + + assert!(result.as_::().elements().is::()); let expected = FixedSizeListArray::new( PrimitiveArray::from_iter([20i32, 30, 10, 20]).into_array(), 2, @@ -380,6 +399,141 @@ fn test_take_uses_take_slices_for_encoded_elements_child() { 2, ); assert_arrays_eq!(expected, result, &mut ctx); + Ok(()) +} + +#[derive(Clone, Debug)] +struct NoTakeSlicesArray; + +impl NoTakeSlicesArray { + fn wrap(child: ArrayRef) -> ArrayRef { + let dtype = child.dtype().clone(); + let len = child.len(); + + // SAFETY: `NoTakeSlicesArray` has one child with matching dtype and length, and no + // top-level metadata beyond `EmptyArrayData`. + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(NoTakeSlicesArray, dtype, len, EmptyArrayData) + .with_slots(smallvec![Some(child)]), + ) + } + .into_array() + } +} + +impl VTable for NoTakeSlicesArray { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.test.no_take_slices"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + slots.len() == 1, + "NoTakeSlicesArray expected one child slot" + ); + let child = slots[0] + .as_ref() + .ok_or_else(|| vortex_err!("NoTakeSlicesArray child slot must be present"))?; + vortex_ensure!( + child.dtype() == dtype, + "NoTakeSlicesArray child dtype {} does not match outer dtype {}", + child.dtype(), + dtype + ); + vortex_ensure!( + child.len() == len, + "NoTakeSlicesArray child length {} does not match outer length {}", + child.len(), + len + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle { + vortex_panic!("NoTakeSlicesArray has no buffers") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + match idx { + 0 => "child".to_string(), + _ => vortex_panic!("NoTakeSlicesArray slot index {idx} out of bounds"), + } + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + vortex_bail!("NoTakeSlicesArray is not serializable") + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("NoTakeSlicesArray is not serializable") + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array.slots()[0].clone().ok_or_else( + || vortex_err!("NoTakeSlicesArray child slot must be present"), + )?)) + } +} + +impl OperationsVTable for NoTakeSlicesArray { + fn scalar_at( + array: ArrayView<'_, NoTakeSlicesArray>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + array.as_ref().slots()[0] + .as_ref() + .ok_or_else(|| vortex_err!("NoTakeSlicesArray child slot must be present"))? + .execute_scalar(index, ctx) + } +} + +impl ValidityVTable for NoTakeSlicesArray { + fn validity(array: ArrayView<'_, NoTakeSlicesArray>) -> VortexResult { + array.as_ref().slots()[0] + .as_ref() + .ok_or_else(|| vortex_err!("NoTakeSlicesArray child slot must be present"))? + .validity() + } } // Parameterized test for nullable array scenarios that are specific to FSL's implementation. diff --git a/vortex-array/src/arrays/primitive/compute/take_slices.rs b/vortex-array/src/arrays/primitive/compute/take_slices.rs index ccef3ef5bde..d9d333fb3a2 100644 --- a/vortex-array/src/arrays/primitive/compute/take_slices.rs +++ b/vortex-array/src/arrays/primitive/compute/take_slices.rs @@ -3,6 +3,7 @@ use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::IntoArray; @@ -22,7 +23,10 @@ impl TakeSlicesExecute for Primitive { let validity = array.validity()?.take_slices(slices)?; match_each_native_ptype!(array.ptype(), |T| { let source = array.as_slice::(); - let len = slices.iter().map(|(start, end)| end - start).sum(); + let len = slices.iter().try_fold(0usize, |len, &(start, end)| { + len.checked_add(end - start) + .ok_or_else(|| vortex_err!("TakeSlices output length overflow")) + })?; let mut values = BufferMut::::with_capacity(len); for &(start, end) in slices { diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs index b108d52b2f9..ceafe06b10b 100644 --- a/vortex-array/src/arrays/take_slices/array.rs +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -23,7 +23,7 @@ pub(super) const CHILD_SLOT: usize = 0; pub(super) const NUM_SLOTS: usize = 1; pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child"]; -/// Metadata for a [`TakeSlices`](crate::arrays::TakeSlices) array. +/// Metadata for a [`TakeSlices`] array. #[derive(Clone, Debug)] pub struct TakeSlicesData { pub(super) slices: Arc<[(usize, usize)]>, @@ -36,7 +36,7 @@ impl Display for TakeSlicesData { } } -/// Extension methods for [`TakeSlices`](crate::arrays::TakeSlices) arrays. +/// Extension methods for [`TakeSlices`] arrays. pub trait TakeSlicesArrayExt: TypedArrayRef { /// The child array selected by this ordered range list. fn child(&self) -> &ArrayRef { diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs index 47d425953b2..5b2c820239f 100644 --- a/vortex-array/src/arrays/take_slices/tests.rs +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -18,6 +18,7 @@ use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::assert_arrays_eq; use crate::dtype::Nullability; use crate::kernel::ExecuteParentKernel; +use crate::validity::Validity; #[test] fn take_slices_preserves_order_duplicates_and_overlap() -> VortexResult<()> { @@ -94,6 +95,30 @@ fn take_slices_size_one_child_can_repeat_the_only_range() -> VortexResult<()> { Ok(()) } +#[test] +fn take_slices_size_one_nullable_child_can_repeat_null() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_option_iter([None::]).into_array(); + + let actual = array.take_slices(vec![(0, 1), (0, 1)])?; + let expected = PrimitiveArray::from_option_iter([None::, None]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn take_slices_preserves_all_invalid_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllInvalid).into_array(); + + let actual = array.take_slices(vec![(2, 3), (0, 2)])?; + let expected = PrimitiveArray::from_option_iter([None::, None, None]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + #[test] fn primitive_take_slices_execute_parent_copies_ranges_directly() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index 2be7c978224..7ff196c675a 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -8,6 +8,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -94,12 +95,7 @@ impl VTable for TakeSlices { child.dtype(), dtype ); - vortex_ensure!( - data.len() == len, - "TakeSlicesArray length {} does not match outer length {}", - data.len(), - len - ); + let mut computed_len = 0usize; for &(start, end) in data.slices() { vortex_ensure!( start < end, @@ -110,7 +106,22 @@ impl VTable for TakeSlices { "TakeSlicesArray range {start}..{end} exceeds child length {}", child.len() ); + computed_len = computed_len + .checked_add(end - start) + .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; } + vortex_ensure!( + data.len() == computed_len, + "TakeSlicesArray metadata length {} does not match computed range length {}", + data.len(), + computed_len + ); + vortex_ensure!( + computed_len == len, + "TakeSlicesArray computed length {} does not match outer length {}", + computed_len, + len + ); Ok(()) } From acd573a71fb4a5beaeb1c0387789af3fdeea8153 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 7 Jul 2026 17:06:43 -0400 Subject: [PATCH 03/37] Clarify TakeSlices range sequence docs Replace ambiguous ordered-range wording with caller-provided sequence terminology so the API docs do not imply sorted or disjoint ranges. Signed-off-by: Daniel King --- vortex-array/src/array/erased.rs | 4 ++-- vortex-array/src/arrays/take_slices/array.rs | 10 +++++----- vortex-array/src/arrays/take_slices/mod.rs | 17 +++++++++-------- vortex-array/src/arrays/take_slices/vtable.rs | 2 +- vortex-array/src/validity.rs | 2 +- 5 files changed, 18 insertions(+), 17 deletions(-) diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 8722ba3d115..c2245e4601a 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -265,8 +265,8 @@ impl ArrayRef { .optimize() } - /// Wraps the array in a [`TakeSlicesArray`] such that it is logically selected by ordered - /// child ranges. + /// Wraps the array in a [`TakeSlicesArray`] such that it is logically selected by a + /// caller-provided sequence of child ranges. pub fn take_slices(&self, slices: Vec<(usize, usize)>) -> VortexResult { TakeSlicesArray::try_new(self.clone(), slices)? .into_array() diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs index ceafe06b10b..00768ca797d 100644 --- a/vortex-array/src/arrays/take_slices/array.rs +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -18,7 +18,7 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::arrays::TakeSlices; -/// The child array being selected by ordered slices. +/// The child array selected by the range sequence. pub(super) const CHILD_SLOT: usize = 0; pub(super) const NUM_SLOTS: usize = 1; pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child"]; @@ -38,14 +38,14 @@ impl Display for TakeSlicesData { /// Extension methods for [`TakeSlices`] arrays. pub trait TakeSlicesArrayExt: TypedArrayRef { - /// The child array selected by this ordered range list. + /// The child array selected by this range sequence. fn child(&self) -> &ArrayRef { self.as_ref().slots()[CHILD_SLOT] .as_ref() .vortex_expect("validated take-slices child slot") } - /// The ordered, non-empty child ranges represented by this array. + /// The caller-provided, non-empty child ranges represented by this array. fn slices(&self) -> &[(usize, usize)] { &self.slices } @@ -85,12 +85,12 @@ impl TakeSlicesData { self.slices.is_empty() } - /// The ordered ranges used to select child values. + /// The ranges used to select child values. pub fn slices(&self) -> &[(usize, usize)] { &self.slices } - /// Returns the ordered ranges as `Range` values. + /// Returns the ranges as `Range` values. pub fn slice_ranges(&self) -> impl Iterator> + '_ { self.slices.iter().map(|&(start, end)| start..end) } diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index 9957adec546..538be4ae6f1 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -1,11 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Reduce and execute adaptors for ordered slice-take operations. +//! Reduce and execute adaptors for slice-sequence take operations. //! -//! `TakeSlicesArray` represents concatenating an ordered list of non-empty child ranges. The -//! ranges preserve caller order and may overlap. Encodings that know how to serve those ranges -//! efficiently implement [`TakeSlicesReduce`] or [`TakeSlicesExecute`]. +//! `TakeSlicesArray` represents concatenating a caller-provided sequence of non-empty child +//! ranges. The ranges are not required to be sorted or disjoint, and the output preserves the +//! sequence order. Encodings that know how to serve those ranges efficiently implement +//! [`TakeSlicesReduce`] or [`TakeSlicesExecute`]. mod array; mod rules; @@ -26,9 +27,9 @@ use crate::kernel::ExecuteParentKernel; use crate::matcher::Matcher; use crate::optimizer::rules::ArrayParentReduceRule; -/// Metadata-only implementation hook for taking ordered child ranges. +/// Metadata-only implementation hook for taking a sequence of child ranges. pub trait TakeSlicesReduce: VTable { - /// Take ordered slices from an array without reading buffers. + /// Take a sequence of slices from an array without reading buffers. /// /// Implementations should return `None` if serving the ranges requires buffer access. /// @@ -41,9 +42,9 @@ pub trait TakeSlicesReduce: VTable { ) -> VortexResult>; } -/// Execution implementation hook for taking ordered child ranges. +/// Execution implementation hook for taking a sequence of child ranges. pub trait TakeSlicesExecute: VTable { - /// Take ordered slices from an array, potentially reading buffers. + /// Take a sequence of slices from an array, potentially reading buffers. /// /// # Preconditions /// diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index 7ff196c675a..3e756ed74ef 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -44,7 +44,7 @@ use crate::validity::Validity; /// A [`TakeSlices`]-encoded Vortex array. pub type TakeSlicesArray = Array; -/// Ordered child-range selection encoding. +/// Child-range sequence selection encoding. #[derive(Clone, Debug)] pub struct TakeSlices; diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 3a999d0f0ef..214fb321093 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -224,7 +224,7 @@ impl Validity { } } - /// Select validity values by concatenating ordered, non-empty ranges. + /// Select validity values by concatenating a sequence of non-empty ranges. pub fn take_slices(&self, slices: &[(usize, usize)]) -> VortexResult { match self { v @ (Self::NonNullable | Self::AllValid | Self::AllInvalid) => Ok(v.clone()), From 0724ec35a55bf21202006654879fbf6a7f62eafa Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 16:07:58 -0400 Subject: [PATCH 04/37] Use run selectors for TakeSlices Switch TakeSlices to starts/lengths selector arrays and reuse the contiguous-run gather for ListView rebuilds. Signed-off-by: Daniel King --- vortex-array/benches/take_fsl.rs | 17 +- vortex-array/src/array/erased.rs | 9 +- .../arrays/fixed_size_list/compute/take.rs | 48 +++-- vortex-array/src/arrays/listview/rebuild.rs | 191 ++---------------- .../arrays/primitive/compute/take_slices.rs | 11 +- vortex-array/src/arrays/take_slices/array.rs | 85 ++++---- vortex-array/src/arrays/take_slices/mod.rs | 188 +++++++++++++++-- vortex-array/src/arrays/take_slices/rules.rs | 42 +++- vortex-array/src/arrays/take_slices/tests.rs | 88 +++++--- vortex-array/src/arrays/take_slices/vtable.rs | 56 ++--- vortex-array/src/validity.rs | 8 +- 11 files changed, 417 insertions(+), 326 deletions(-) diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index 8f0e63191cb..7767c1bbfdf 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -209,17 +209,24 @@ fn take_fsl_f16_take_slices_strategy( array: &FixedSizeListArray, indices: &Buffer, ) -> FixedSizeListArray { - let slices = indices + let starts = indices .as_ref() .iter() .map(|&idx| { let start = idx as usize * LIST_SIZE; - (start, start + LIST_SIZE) + start as u64 }) - .collect(); - let elements = array.elements().take_slices(slices).unwrap(); + .collect::>(); + let lengths = std::iter::repeat_n(LIST_SIZE as u64, indices.len()).collect::>(); + let elements = array + .elements() + .take_slices( + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(lengths).into_array(), + ) + .unwrap(); - // SAFETY: each generated slice has width `LIST_SIZE`, and there is one slice per input index, + // SAFETY: each generated run has width `LIST_SIZE`, and there is one run per input index, // so `elements.len() == indices.len() * LIST_SIZE`. unsafe { FixedSizeListArray::new_unchecked( diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index c2245e4601a..8be60976fe3 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -266,9 +266,12 @@ impl ArrayRef { } /// Wraps the array in a [`TakeSlicesArray`] such that it is logically selected by a - /// caller-provided sequence of child ranges. - pub fn take_slices(&self, slices: Vec<(usize, usize)>) -> VortexResult { - TakeSlicesArray::try_new(self.clone(), slices)? + /// caller-provided sequence of child runs. + /// + /// The output is the concatenation of `self[starts[i]..starts[i] + lengths[i]]` for each + /// selector row. + pub fn take_slices(&self, starts: ArrayRef, lengths: ArrayRef) -> VortexResult { + TakeSlicesArray::try_new(self.clone(), starts, lengths)? .into_array() .optimize() } diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 92a851646df..10d6fa30bc5 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -29,7 +29,7 @@ use crate::validity::Validity; /// Take implementation for [`FixedSizeListArray`]. /// /// `FixedSizeListArray` must rebuild its elements array because selected lists need to become -/// packed from offset 0. The FSL layer translates selected list rows into ordered element ranges +/// packed from offset 0. The FSL layer translates selected list rows into ordered element runs /// and delegates the execution strategy to the elements child via `take_slices`. impl TakeExecute for FixedSizeList { fn take( @@ -148,17 +148,18 @@ fn take_non_nullable_non_empty_fsl( let indices: &[I] = indices_array.as_slice::(); let new_len = indices.len(); let expected_elements_len = take_elements_len(new_len, list_size)?; - let mut slices = Vec::with_capacity(new_len); + let mut starts = Vec::with_capacity(new_len); for &data_idx in indices { let data_idx = index_to_usize(data_idx)?; - slices.push(list_range(data_idx, list_size, array_len)?); + let start = list_start(data_idx, list_size, array_len)?; + starts.push(usize_to_u64(start, "FixedSizeList take element start")?); } - let new_elements = array.elements().take_slices(slices)?; + let new_elements = take_element_runs(array.elements(), starts, list_size)?; ensure_elements_len(new_elements.len(), expected_elements_len)?; - // SAFETY: `slices` contains one checked range of `list_size` elements for each output row, + // SAFETY: `starts` contains one checked run of `list_size` elements for each output row, // `new_elements` has `new_len * list_size` elements, and non-nullable validity has no length. Ok(unsafe { FixedSizeListArray::new_unchecked( @@ -187,30 +188,30 @@ fn take_nullable_non_empty_fsl( .execute_mask(array.as_ref().len(), ctx)?; let indices_validity = indices_validity_mask(&indices_array, ctx)?; - let null_elements = (0, list_size); - let mut slices = Vec::with_capacity(new_len); + let null_element_start = 0u64; + let mut starts = Vec::with_capacity(new_len); let mut new_validity_builder = BitBufferMut::with_capacity(new_len); for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { if !is_index_valid { - slices.push(null_elements); + starts.push(null_element_start); new_validity_builder.append(false); continue; } let data_idx = index_to_usize(data_idx)?; - let range = list_range(data_idx, list_size, array_len)?; + let start = list_start(data_idx, list_size, array_len)?; if !array_validity.value(data_idx) { - slices.push(null_elements); + starts.push(null_element_start); new_validity_builder.append(false); continue; } - slices.push(range); + starts.push(usize_to_u64(start, "FixedSizeList take element start")?); new_validity_builder.append(true); } - let new_elements = array.elements().take_slices(slices)?; + let new_elements = take_element_runs(array.elements(), starts, list_size)?; ensure_elements_len(new_elements.len(), expected_elements_len)?; let new_validity = Validity::from(new_validity_builder.freeze()); @@ -268,7 +269,7 @@ fn ensure_elements_len(actual: usize, expected: usize) -> VortexResult<()> { Ok(()) } -fn list_range(data_idx: usize, list_size: usize, array_len: usize) -> VortexResult<(usize, usize)> { +fn list_start(data_idx: usize, list_size: usize, array_len: usize) -> VortexResult { check_index_in_bounds(data_idx, array_len)?; let start = data_idx.checked_mul(list_size).ok_or_else(|| { @@ -276,12 +277,12 @@ fn list_range(data_idx: usize, list_size: usize, array_len: usize) -> VortexResu "FixedSizeList take element range overflow for index {data_idx} and list size {list_size}" ) })?; - let end = start.checked_add(list_size).ok_or_else(|| { + start.checked_add(list_size).ok_or_else(|| { vortex_err!( "FixedSizeList take element range overflow for index {data_idx} and list size {list_size}" ) })?; - Ok((start, end)) + Ok(start) } fn check_index_in_bounds(data_idx: usize, array_len: usize) -> VortexResult<()> { @@ -297,8 +298,25 @@ fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef builder.finish() } +fn take_element_runs( + elements: &ArrayRef, + starts: Vec, + length: usize, +) -> VortexResult { + let run_count = starts.len(); + let length = usize_to_u64(length, "FixedSizeList take element run length")?; + elements.take_slices( + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(std::iter::repeat_n(length, run_count)).into_array(), + ) +} + fn index_to_usize(index: I) -> VortexResult { index .to_usize() .ok_or_else(|| vortex_err!("FixedSizeList take index {index} does not fit in usize")) } + +fn usize_to_u64(value: usize, name: &str) -> VortexResult { + u64::try_from(value).map_err(|_| vortex_err!("{name} {value} does not fit in u64")) +} diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 342c9eccf19..5b64a71aaf5 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -6,7 +6,6 @@ use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; @@ -14,7 +13,6 @@ use crate::arrays::ListViewArray; 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,55 +138,15 @@ 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 | PType::U16 | PType::U32 => self.rebuild_with_take_slices::(ctx), + PType::U64 => self.rebuild_with_take_slices::(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( - &self, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let sizes_canonical = self.sizes().clone().execute::(ctx)?; - let sizes_canonical = - sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); - let total: u64 = sizes_canonical - .as_slice::() - .iter() - .map(|s| (*s).as_() as u64) - .sum(); - if Self::should_use_take(total, self.len()) { - self.rebuild_with_take::(ctx) - } else { - self.rebuild_list_by_list::(ctx) - } - } - - /// Returns `true` when we are confident that `rebuild_with_take` will - /// outperform `rebuild_list_by_list`. - /// - /// Take is dramatically faster for small lists (often 10-100×) because it - /// avoids per-list builder overhead. LBL is the safer default for larger - /// lists since its sequential memcpy scales well. We only choose take when - /// the average list size is small enough that take clearly dominates. - fn should_use_take(total_output_elements: u64, num_lists: usize) -> bool { - if num_lists == 0 { - return true; - } - let avg = total_output_elements / num_lists as u64; - avg < 128 - } - - /// Rebuilds elements using a single bulk `take`: collect all element indices into a flat - /// `BufferMut`, perform a single `take`. - fn rebuild_with_take( + /// Rebuilds elements using a single contiguous-run gather over the element child. + fn rebuild_with_take_slices( &self, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -198,17 +156,15 @@ impl ListViewArray { 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 len = offsets_canonical.len(); let mut new_offsets = BufferMut::::with_capacity(len); let mut new_sizes = BufferMut::::with_capacity(len); - let mut take_indices = BufferMut::::with_capacity(self.elements().len()); // Resolve validity to a mask once instead of probing it per row: `execute_is_valid` // executes a scalar on every call for array-backed validity, which is O(len) work repeated @@ -216,137 +172,35 @@ impl ListViewArray { let validity = self.validity()?.execute_mask(len, ctx)?; let mut n_elements = NewOffset::zero(); - for index in 0..len { - if !validity.value(index) { + for (index, is_valid) in validity.iter().enumerate() { + if !is_valid { new_offsets.push(n_elements); new_sizes.push(S::zero()); continue; } - let offset = offsets_slice[index]; let size = sizes_slice[index]; - let start = offset.as_(); - let stop = start + size.as_(); new_offsets.push(n_elements); new_sizes.push(size); - take_indices.extend(start as u64..stop as u64); n_elements += num_traits::cast(size).vortex_expect("Cast failed"); } - let elements = self.elements().take(take_indices.into_array())?; - // Built unsigned; reinterpret back to the signed-preserving result types. - let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) - .reinterpret_cast(new_offset_ptype) - .into_array(); - let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable) - .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. - 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( - &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 + let new_sizes = new_sizes.freeze(); + let lengths = PrimitiveArray::new(new_sizes.clone(), Validity::NonNullable); + let elements = 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 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"); - } - + .take_slices(offsets_canonical.into_array(), lengths.into_array())?; // Built unsigned; reinterpret back to the signed-preserving result types. let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(new_offset_ptype) .into_array(); - let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable) + let sizes = PrimitiveArray::new(new_sizes, 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 the validity array is preserved from the original. Ok(unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?) .with_zero_copy_to_list(true) @@ -593,7 +447,7 @@ mod tests { assert_eq!(rebuilt.size_at(2), 0); // NULL has size 0 assert_eq!(rebuilt.size_at(3), 0); // NULL has size 0 - // Now rebuild with MakeExact (which calls naive_rebuild then trim_elements) + // Now rebuild with MakeExact; this should keep using the zero-copy-to-list fast path. // This should not panic (issue #5412) let exact = rebuilt.rebuild(ListViewRebuildMode::MakeExact, &mut ctx)?; @@ -697,21 +551,6 @@ mod tests { Ok(()) } - // ── should_use_take heuristic tests ──────────────────────────────────── - - #[test] - fn heuristic_zero_lists_uses_take() { - assert!(ListViewArray::should_use_take(0, 0)); - } - - #[test] - fn heuristic_small_lists_use_take() { - // avg = 127 → take - assert!(ListViewArray::should_use_take(127_000, 1_000)); - // avg = 128 → LBL - assert!(!ListViewArray::should_use_take(128_000, 1_000)); - } - /// Regression test for . /// Both offsets and sizes are u8, and offset + size exceeds u8::MAX. #[test] diff --git a/vortex-array/src/arrays/primitive/compute/take_slices.rs b/vortex-array/src/arrays/primitive/compute/take_slices.rs index d9d333fb3a2..f484f64d68b 100644 --- a/vortex-array/src/arrays/primitive/compute/take_slices.rs +++ b/vortex-array/src/arrays/primitive/compute/take_slices.rs @@ -11,16 +11,19 @@ use crate::array::ArrayView; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::take_slices::TakeSlicesExecute; +use crate::arrays::take_slices::selector_slices; use crate::executor::ExecutionCtx; use crate::match_each_native_ptype; impl TakeSlicesExecute for Primitive { fn take_slices( array: ArrayView<'_, Self>, - slices: &[(usize, usize)], - _ctx: &mut ExecutionCtx, + starts: &ArrayRef, + lengths: &ArrayRef, + ctx: &mut ExecutionCtx, ) -> VortexResult> { - let validity = array.validity()?.take_slices(slices)?; + let slices = selector_slices(array.len(), starts, lengths, ctx)?; + let validity = array.validity()?.take_slices(starts, lengths)?; match_each_native_ptype!(array.ptype(), |T| { let source = array.as_slice::(); let len = slices.iter().try_fold(0usize, |len, &(start, end)| { @@ -29,7 +32,7 @@ impl TakeSlicesExecute for Primitive { })?; let mut values = BufferMut::::with_capacity(len); - for &(start, end) in slices { + for (start, end) in slices { values.extend_from_slice(&source[start..end]); } diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs index 00768ca797d..b10c698b1c4 100644 --- a/vortex-array/src/arrays/take_slices/array.rs +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -3,76 +3,67 @@ use std::fmt::Display; use std::fmt::Formatter; -use std::ops::Range; -use std::sync::Arc; use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use crate::ArrayRef; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::arrays::TakeSlices; +use crate::arrays::take_slices::selector_output_len; -/// The child array selected by the range sequence. +/// The child array selected by the run sequence. pub(super) const CHILD_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child"]; +/// The selector naming the start of each child run. +pub(super) const STARTS_SLOT: usize = 1; +/// The selector naming the length of each child run. +pub(super) const LENGTHS_SLOT: usize = 2; +pub(super) const NUM_SLOTS: usize = 3; +pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "starts", "lengths"]; /// Metadata for a [`TakeSlices`] array. #[derive(Clone, Debug)] pub struct TakeSlicesData { - pub(super) slices: Arc<[(usize, usize)]>, len: usize, } impl Display for TakeSlicesData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "nslices: {}, len: {}", self.slices.len(), self.len()) + write!(f, "len: {}", self.len()) } } /// Extension methods for [`TakeSlices`] arrays. pub trait TakeSlicesArrayExt: TypedArrayRef { - /// The child array selected by this range sequence. + /// The child array selected by this run sequence. fn child(&self) -> &ArrayRef { self.as_ref().slots()[CHILD_SLOT] .as_ref() .vortex_expect("validated take-slices child slot") } - /// The caller-provided, non-empty child ranges represented by this array. - fn slices(&self) -> &[(usize, usize)] { - &self.slices + /// The selector naming each child run's start offset. + fn starts(&self) -> &ArrayRef { + self.as_ref().slots()[STARTS_SLOT] + .as_ref() + .vortex_expect("validated take-slices starts slot") + } + + /// The selector naming each child run's length. + fn lengths(&self) -> &ArrayRef { + self.as_ref().slots()[LENGTHS_SLOT] + .as_ref() + .vortex_expect("validated take-slices lengths slot") } } impl> TakeSlicesArrayExt for T {} impl TakeSlicesData { - fn try_new(child_len: usize, slices: Vec<(usize, usize)>) -> VortexResult { - let mut len = 0usize; - for &(start, end) in &slices { - vortex_ensure!( - start < end, - "TakeSlicesArray range must be non-empty: {start}..{end}" - ); - vortex_ensure!( - end <= child_len, - "TakeSlicesArray range {start}..{end} exceeds child array length {child_len}" - ); - len = len - .checked_add(end - start) - .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; - } - - Ok(Self { - slices: Arc::from(slices.into_boxed_slice()), - len, - }) + fn new(len: usize) -> Self { + Self { len } } /// Returns the length of this array. @@ -82,32 +73,26 @@ impl TakeSlicesData { /// Returns `true` if this array is empty. pub fn is_empty(&self) -> bool { - self.slices.is_empty() - } - - /// The ranges used to select child values. - pub fn slices(&self) -> &[(usize, usize)] { - &self.slices - } - - /// Returns the ranges as `Range` values. - pub fn slice_ranges(&self) -> impl Iterator> + '_ { - self.slices.iter().map(|&(start, end)| start..end) + self.len == 0 } } impl Array { /// Constructs a new validated `TakeSlicesArray`. - pub fn try_new(child: ArrayRef, slices: Vec<(usize, usize)>) -> VortexResult { + pub fn try_new(child: ArrayRef, starts: ArrayRef, lengths: ArrayRef) -> VortexResult { let dtype = child.dtype().clone(); - let data = TakeSlicesData::try_new(child.len(), slices)?; - let len = data.len(); + let len = selector_output_len(child.len(), &starts, &lengths)?; + let data = TakeSlicesData::new(len); - // SAFETY: `TakeSlicesData::try_new` validates range bounds and computes `len`; the outer - // dtype is copied from the child, and the required child slot is populated. + // SAFETY: `selector_output_len` validates selector dtypes, run bounds, and computes `len`; + // the outer dtype is copied from the child, and all required slots are populated. Ok(unsafe { Array::from_parts_unchecked( - ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![Some(child)]), + ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![ + Some(child), + Some(starts), + Some(lengths) + ]), ) }) } diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index 538be4ae6f1..b017fc7fe5b 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -1,11 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Reduce and execute adaptors for slice-sequence take operations. +//! Reduce and execute adaptors for contiguous-run take operations. //! -//! `TakeSlicesArray` represents concatenating a caller-provided sequence of non-empty child -//! ranges. The ranges are not required to be sorted or disjoint, and the output preserves the -//! sequence order. Encodings that know how to serve those ranges efficiently implement +//! `TakeSlicesArray` represents a gather of contiguous child runs: the output is the +//! concatenation of `values[starts[i]..starts[i] + lengths[i]]` for each selector row. Runs may +//! overlap, repeat, and appear in any order. Encodings that know how to serve those runs +//! efficiently implement //! [`TakeSlicesReduce`] or [`TakeSlicesExecute`]. mod array; @@ -15,48 +16,81 @@ mod vtable; pub use array::TakeSlicesArrayExt; pub use array::TakeSlicesData; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; pub use vtable::*; use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::IntegerPType; use crate::kernel::ExecuteParentKernel; +use crate::legacy_session; +use crate::match_each_unsigned_integer_ptype; use crate::matcher::Matcher; use crate::optimizer::rules::ArrayParentReduceRule; -/// Metadata-only implementation hook for taking a sequence of child ranges. +/// Metadata-only implementation hook for taking a sequence of child runs. pub trait TakeSlicesReduce: VTable { - /// Take a sequence of slices from an array without reading buffers. + /// Take a sequence of contiguous runs from an array without reading value buffers. /// - /// Implementations should return `None` if serving the ranges requires buffer access. + /// Implementations should return `None` if serving the runs requires value buffer access. /// /// # Preconditions /// - /// `slices` is guaranteed to contain only non-empty ranges in bounds for `array`. + /// `starts` and `lengths` are guaranteed to be equal-length, non-nullable unsigned integer + /// arrays. Each run is in bounds for `array`. fn take_slices( array: ArrayView<'_, Self>, - slices: &[(usize, usize)], + starts: &ArrayRef, + lengths: &ArrayRef, ) -> VortexResult>; } -/// Execution implementation hook for taking a sequence of child ranges. +/// Execution implementation hook for taking a sequence of child runs. pub trait TakeSlicesExecute: VTable { - /// Take a sequence of slices from an array, potentially reading buffers. + /// Take a sequence of contiguous runs from an array, potentially reading value buffers. /// /// # Preconditions /// - /// `slices` is guaranteed to contain only non-empty ranges in bounds for `array`. + /// `starts` and `lengths` are guaranteed to be equal-length, non-nullable unsigned integer + /// arrays. Each run is in bounds for `array`. fn take_slices( array: ArrayView<'_, Self>, - slices: &[(usize, usize)], + starts: &ArrayRef, + lengths: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult>; } fn trivial_take_slices( + child: &ArrayRef, + starts: &ArrayRef, + lengths: &ArrayRef, +) -> VortexResult> { + let mut ctx = legacy_execution_ctx(); + trivial_take_slices_with_ctx(child, starts, lengths, &mut ctx) +} + +fn trivial_take_slices_with_ctx( + child: &ArrayRef, + starts: &ArrayRef, + lengths: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let slices = selector_slices(child.len(), starts, lengths, ctx)?; + trivial_take_slices_from_ranges(child, &slices) +} + +fn trivial_take_slices_from_ranges( child: &ArrayRef, slices: &[(usize, usize)], ) -> VortexResult> { @@ -74,6 +108,115 @@ fn trivial_take_slices( Ok(None) } +pub(super) fn selector_output_len( + child_len: usize, + starts: &ArrayRef, + lengths: &ArrayRef, +) -> VortexResult { + let mut ctx = legacy_execution_ctx(); + let slices = selector_slices(child_len, starts, lengths, &mut ctx)?; + slices.iter().try_fold(0usize, |len, &(start, end)| { + len.checked_add(end - start) + .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow")) + }) +} + +#[allow(clippy::disallowed_methods)] +fn legacy_execution_ctx() -> ExecutionCtx { + legacy_session().create_execution_ctx() +} + +pub(super) fn selector_slices( + child_len: usize, + starts: &ArrayRef, + lengths: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + check_selector_arrays(starts, lengths)?; + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + primitive_selector_slices(child_len, starts.as_view(), lengths.as_view()) +} + +pub(super) fn check_selector_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult<()> { + check_selector_dtype("starts", starts)?; + check_selector_dtype("lengths", lengths)?; + vortex_ensure!( + starts.len() == lengths.len(), + "TakeSlicesArray selectors must have equal length, got starts {} and lengths {}", + starts.len(), + lengths.len() + ); + Ok(()) +} + +fn check_selector_dtype(name: &str, selector: &ArrayRef) -> VortexResult<()> { + match selector.dtype() { + DType::Primitive(ptype, nullability) if ptype.is_unsigned_int() => { + vortex_ensure!( + !nullability.is_nullable(), + "TakeSlicesArray {name} must be non-nullable, got {}", + selector.dtype() + ); + Ok(()) + } + other => vortex_bail!( + "TakeSlicesArray {name} must be a non-nullable unsigned integer, got {other}" + ), + } +} + +fn primitive_selector_slices( + child_len: usize, + starts: ArrayView<'_, Primitive>, + lengths: ArrayView<'_, Primitive>, +) -> VortexResult> { + vortex_ensure!( + starts.len() == lengths.len(), + "TakeSlicesArray selectors must have equal length, got starts {} and lengths {}", + starts.len(), + lengths.len() + ); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + selector_slices_typed::( + child_len, + starts.as_slice::(), + lengths.as_slice::(), + ) + }) + }) +} + +fn selector_slices_typed( + child_len: usize, + starts: &[S], + lengths: &[L], +) -> VortexResult> { + let mut slices = Vec::with_capacity(starts.len()); + for (&start, &length) in starts.iter().zip(lengths) { + let start = selector_to_usize("start", start)?; + let length = selector_to_usize("length", length)?; + let end = start.checked_add(length).ok_or_else(|| { + vortex_err!("TakeSlicesArray run overflow for start {start} and length {length}") + })?; + vortex_ensure!( + end <= child_len, + "TakeSlicesArray run {start}..{end} exceeds child array length {child_len}" + ); + if length != 0 { + slices.push((start, end)); + } + } + Ok(slices) +} + +fn selector_to_usize(name: &str, value: T) -> VortexResult { + value + .to_usize() + .ok_or_else(|| vortex_err!("TakeSlicesArray {name} selector {value} does not fit in usize")) +} + /// Adaptor that wraps a [`TakeSlicesReduce`] impl as an [`ArrayParentReduceRule`]. #[derive(Default, Debug)] pub struct TakeSlicesReduceAdaptor(pub V); @@ -90,11 +233,14 @@ where parent: ::Match<'_>, child_idx: usize, ) -> VortexResult> { - assert_eq!(child_idx, 0); - if let Some(result) = trivial_take_slices(array.array(), parent.slices())? { + if child_idx != 0 { + return Ok(None); + } + if let Some(result) = trivial_take_slices(array.array(), parent.starts(), parent.lengths())? + { return Ok(Some(result)); } - ::take_slices(array, parent.slices()) + ::take_slices(array, parent.starts(), parent.lengths()) } } @@ -115,11 +261,15 @@ where child_idx: usize, ctx: &mut ExecutionCtx, ) -> VortexResult> { - assert_eq!(child_idx, 0); - if let Some(result) = trivial_take_slices(array.array(), parent.slices())? { + if child_idx != 0 { + return Ok(None); + } + if let Some(result) = + trivial_take_slices_with_ctx(array.array(), parent.starts(), parent.lengths(), ctx)? + { return Ok(Some(result)); } - ::take_slices(array, parent.slices(), ctx) + ::take_slices(array, parent.starts(), parent.lengths(), ctx) } } diff --git a/vortex-array/src/arrays/take_slices/rules.rs b/vortex-array/src/arrays/take_slices/rules.rs index 5a4d350f08f..d578fafdc28 100644 --- a/vortex-array/src/arrays/take_slices/rules.rs +++ b/vortex-array/src/arrays/take_slices/rules.rs @@ -2,16 +2,20 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_error::vortex_err; +use super::legacy_execution_ctx; use super::trivial_take_slices; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::PrimitiveArray; use crate::arrays::TakeSlices; use crate::arrays::TakeSlicesArray; use crate::arrays::take_slices::TakeSlicesArrayExt; use crate::arrays::take_slices::TakeSlicesReduce; use crate::arrays::take_slices::TakeSlicesReduceAdaptor; +use crate::arrays::take_slices::selector_slices; use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; @@ -26,18 +30,28 @@ struct TrivialTakeSlicesRule; impl ArrayReduceRule for TrivialTakeSlicesRule { fn reduce(&self, array: ArrayView<'_, TakeSlices>) -> VortexResult> { - trivial_take_slices(array.child(), array.slices()) + trivial_take_slices(array.child(), array.starts(), array.lengths()) } } impl TakeSlicesReduce for TakeSlices { fn take_slices( array: ArrayView<'_, Self>, - slices: &[(usize, usize)], + starts: &ArrayRef, + lengths: &ArrayRef, ) -> VortexResult> { - let combined = project_slices(array.slices(), slices); + let mut ctx = legacy_execution_ctx(); + let inner = selector_slices( + array.child().len(), + array.starts(), + array.lengths(), + &mut ctx, + )?; + let outer = selector_slices(array.len(), starts, lengths, &mut ctx)?; + let combined = project_slices(&inner, &outer); + let (starts, lengths) = selectors_from_slices(&combined)?; Ok(Some( - TakeSlicesArray::try_new(array.child().clone(), combined)?.into_array(), + TakeSlicesArray::try_new(array.child().clone(), starts, lengths)?.into_array(), )) } } @@ -69,3 +83,23 @@ fn project_slices(inner: &[(usize, usize)], outer: &[(usize, usize)]) -> Vec<(us projected } + +fn selectors_from_slices(slices: &[(usize, usize)]) -> VortexResult<(ArrayRef, ArrayRef)> { + let starts = slices + .iter() + .map(|&(start, _)| selector_value(start, "start")) + .collect::>>()?; + let lengths = slices + .iter() + .map(|&(start, end)| selector_value(end - start, "length")) + .collect::>>()?; + Ok(( + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(lengths).into_array(), + )) +} + +fn selector_value(value: usize, name: &str) -> VortexResult { + u64::try_from(value) + .map_err(|_| vortex_err!("TakeSlicesArray projected {name} {value} does not fit in u64")) +} diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs index 5b2c820239f..ec76643967b 100644 --- a/vortex-array/src/arrays/take_slices/tests.rs +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -5,6 +5,7 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; +use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; @@ -25,7 +26,7 @@ fn take_slices_preserves_order_duplicates_and_overlap() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..8).into_array(); - let actual = array.take_slices(vec![(4, 6), (1, 3), (4, 6), (2, 5)])?; + let actual = take_slices(&array, &[(4, 2), (1, 2), (4, 2), (2, 3)])?; let expected = PrimitiveArray::from_iter([4i32, 5, 1, 2, 4, 5, 2, 3, 4]); assert_arrays_eq!(actual, expected, &mut ctx); @@ -39,7 +40,7 @@ fn take_slices_preserves_nullable_child_validity() -> VortexResult<()> { PrimitiveArray::from_option_iter([Some(0i32), None, Some(2), Some(3), None, Some(5)]) .into_array(); - let actual = array.take_slices(vec![(1, 4), (4, 6)])?; + let actual = take_slices(&array, &[(1, 3), (4, 2)])?; let expected = PrimitiveArray::from_option_iter([None, Some(2), Some(3), None, Some(5)]); assert_arrays_eq!(actual, expected, &mut ctx); @@ -47,13 +48,13 @@ fn take_slices_preserves_nullable_child_validity() -> VortexResult<()> { } #[test] -fn take_slices_lazy_scalar_and_validity_follow_ranges() -> VortexResult<()> { +fn take_slices_lazy_scalar_and_validity_follow_runs() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(0i32), None, Some(2), Some(3), None, Some(5)]) .into_array(); - let actual = array.take_slices(vec![(1, 4), (4, 6)])?; + let actual = take_slices(&array, &[(1, 3), (4, 2)])?; let validity = actual.validity()?.execute_mask(actual.len(), &mut ctx)?; assert!(!validity.value(0)); @@ -75,7 +76,7 @@ fn take_slices_one_slice_reduces_to_child_slice() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); - let actual = array.take_slices(vec![(1, 5)])?; + let actual = take_slices(&array, &[(1, 4)])?; assert!(actual.is::()); assert!(!actual.is::()); @@ -88,7 +89,7 @@ fn take_slices_size_one_child_can_repeat_the_only_range() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter([7i32]).into_array(); - let actual = array.take_slices(vec![(0, 1), (0, 1)])?; + let actual = take_slices(&array, &[(0, 1), (0, 1)])?; let expected = PrimitiveArray::from_iter([7i32, 7]); assert_arrays_eq!(actual, expected, &mut ctx); @@ -100,7 +101,7 @@ fn take_slices_size_one_nullable_child_can_repeat_null() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([None::]).into_array(); - let actual = array.take_slices(vec![(0, 1), (0, 1)])?; + let actual = take_slices(&array, &[(0, 1), (0, 1)])?; let expected = PrimitiveArray::from_option_iter([None::, None]); assert_arrays_eq!(actual, expected, &mut ctx); @@ -112,7 +113,7 @@ fn take_slices_preserves_all_invalid_validity() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllInvalid).into_array(); - let actual = array.take_slices(vec![(2, 3), (0, 2)])?; + let actual = take_slices(&array, &[(2, 1), (0, 2)])?; let expected = PrimitiveArray::from_option_iter([None::, None, None]); assert_arrays_eq!(actual, expected, &mut ctx); @@ -120,10 +121,11 @@ fn take_slices_preserves_all_invalid_validity() -> VortexResult<()> { } #[test] -fn primitive_take_slices_execute_parent_copies_ranges_directly() -> VortexResult<()> { +fn primitive_take_slices_execute_parent_copies_runs_directly() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..8).into_array(); - let parent = TakeSlicesArray::try_new(array.clone(), vec![(3, 6), (0, 2)])?.into_array(); + let (starts, lengths) = selector_arrays(&[(3, 3), (0, 2)])?; + let parent = TakeSlicesArray::try_new(array.clone(), starts, lengths)?.into_array(); let actual = TakeSlicesExecuteAdaptor(Primitive) .execute_parent( @@ -144,11 +146,11 @@ fn primitive_take_slices_execute_parent_copies_ranges_directly() -> VortexResult } #[test] -fn take_slices_empty_ranges_return_empty_canonical() -> VortexResult<()> { +fn take_slices_empty_runs_return_empty_canonical() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); - let actual = array.take_slices(vec![])?; + let actual = take_slices(&array, &[])?; let expected = PrimitiveArray::from_iter(std::iter::empty::()); assert_arrays_eq!(actual, expected, &mut ctx); @@ -156,11 +158,11 @@ fn take_slices_empty_ranges_return_empty_canonical() -> VortexResult<()> { } #[test] -fn take_slices_empty_child_accepts_empty_ranges() -> VortexResult<()> { +fn take_slices_empty_child_accepts_empty_runs() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::empty::(Nullability::NonNullable).into_array(); - let actual = array.take_slices(vec![])?; + let actual = take_slices(&array, &[])?; let expected = PrimitiveArray::from_iter(std::iter::empty::()); assert_arrays_eq!(actual, expected, &mut ctx); @@ -168,12 +170,15 @@ fn take_slices_empty_child_accepts_empty_ranges() -> VortexResult<()> { } #[test] -fn take_slices_rejects_invalid_ranges() { +fn take_slices_rejects_invalid_runs() -> VortexResult<()> { let array = PrimitiveArray::from_iter(0i32..6).into_array(); - assert!(array.take_slices(vec![(2, 2)]).is_err()); - assert!(array.take_slices(vec![(4, 3)]).is_err()); - assert!(array.take_slices(vec![(4, 7)]).is_err()); + assert!(take_slices(&array, &[(7, 0)]).is_err()); + assert!(take_slices(&array, &[(4, 3)]).is_err()); + let starts = PrimitiveArray::from_iter([0u64, 1]).into_array(); + let lengths = PrimitiveArray::from_iter([1u64]).into_array(); + assert!(array.take_slices(starts, lengths).is_err()); + Ok(()) } #[test] @@ -181,12 +186,19 @@ fn take_slices_of_take_slices_projects_to_original_child() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..10).into_array(); - let inner = array.take_slices(vec![(2, 5), (7, 10)])?; - let actual = inner.take_slices(vec![(1, 4), (4, 6)])?; + let inner = take_slices(&array, &[(2, 3), (7, 3)])?; + let actual = take_slices(&inner, &[(1, 3), (4, 2)])?; let actual_take_slices = actual.as_::(); assert!(actual_take_slices.child().is::()); - assert_eq!(actual_take_slices.slices(), &[(3, 5), (7, 8), (8, 10)]); + assert_eq!( + selector_values(actual_take_slices.starts(), &mut ctx)?, + vec![3, 7, 8] + ); + assert_eq!( + selector_values(actual_take_slices.lengths(), &mut ctx)?, + vec![2, 1, 2] + ); assert_arrays_eq!( actual, PrimitiveArray::from_iter([3i32, 4, 7, 8, 9]), @@ -204,7 +216,7 @@ fn take_slices_generic_execution_handles_child_without_take_slices_kernel() -> V )? .into_array(); - let actual = dict.take_slices(vec![(1, 4), (0, 2)])?; + let actual = take_slices(&dict, &[(1, 3), (0, 2)])?; let expected = PrimitiveArray::from_iter([10i32, 20, 30, 30, 10]); assert!(actual.is::()); @@ -221,10 +233,40 @@ fn take_slices_generic_execution_preserves_nullable_encoded_child() -> VortexRes )? .into_array(); - let actual = dict.take_slices(vec![(1, 3), (0, 2)])?; + let actual = take_slices(&dict, &[(1, 2), (0, 2)])?; let expected = PrimitiveArray::from_option_iter([None, Some(30), Some(10), None]); assert!(actual.is::()); assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } + +fn take_slices(array: &ArrayRef, runs: &[(usize, usize)]) -> VortexResult { + let (starts, lengths) = selector_arrays(runs)?; + array.take_slices(starts, lengths) +} + +fn selector_arrays(runs: &[(usize, usize)]) -> VortexResult<(ArrayRef, ArrayRef)> { + let starts = runs + .iter() + .map(|&(start, _)| selector_value(start, "start")) + .collect::>>()?; + let lengths = runs + .iter() + .map(|&(_, length)| selector_value(length, "length")) + .collect::>>()?; + Ok(( + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(lengths).into_array(), + )) +} + +fn selector_value(value: usize, name: &str) -> VortexResult { + u64::try_from(value) + .map_err(|_| vortex_err!("test {name} selector {value} does not fit in u64")) +} + +fn selector_values(selector: &ArrayRef, ctx: &mut crate::ExecutionCtx) -> VortexResult> { + let selector = selector.clone().execute::(ctx)?; + Ok(selector.as_slice::().to_vec()) +} diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index 3e756ed74ef..bae2619b0e5 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -8,7 +8,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -27,11 +26,15 @@ use crate::array::ValidityVTable; use crate::array::with_empty_buffers; use crate::arrays::take_slices::TakeSlicesArrayExt; use crate::arrays::take_slices::array::CHILD_SLOT; +use crate::arrays::take_slices::array::LENGTHS_SLOT; use crate::arrays::take_slices::array::NUM_SLOTS; use crate::arrays::take_slices::array::SLOT_NAMES; +use crate::arrays::take_slices::array::STARTS_SLOT; use crate::arrays::take_slices::array::TakeSlicesData; use crate::arrays::take_slices::rules::PARENT_RULES; use crate::arrays::take_slices::rules::RULES; +use crate::arrays::take_slices::selector_output_len; +use crate::arrays::take_slices::selector_slices; use crate::buffer::BufferHandle; use crate::builders::builder_with_capacity_in; use crate::dtype::DType; @@ -44,19 +47,19 @@ use crate::validity::Validity; /// A [`TakeSlices`]-encoded Vortex array. pub type TakeSlicesArray = Array; -/// Child-range sequence selection encoding. +/// Contiguous-run gather selection encoding. #[derive(Clone, Debug)] pub struct TakeSlices; impl ArrayHash for TakeSlicesData { fn array_hash(&self, state: &mut H, _accuracy: EqMode) { - self.slices.hash(state); + self.len().hash(state); } } impl ArrayEq for TakeSlicesData { fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { - self.slices == other.slices + self.len() == other.len() } } @@ -86,6 +89,14 @@ impl VTable for TakeSlices { slots[CHILD_SLOT].is_some(), "TakeSlicesArray child slot must be present" ); + vortex_ensure!( + slots[STARTS_SLOT].is_some(), + "TakeSlicesArray starts slot must be present" + ); + vortex_ensure!( + slots[LENGTHS_SLOT].is_some(), + "TakeSlicesArray lengths slot must be present" + ); let child = slots[CHILD_SLOT] .as_ref() .vortex_expect("validated child slot"); @@ -95,24 +106,16 @@ impl VTable for TakeSlices { child.dtype(), dtype ); - let mut computed_len = 0usize; - for &(start, end) in data.slices() { - vortex_ensure!( - start < end, - "TakeSlicesArray range must be non-empty: {start}..{end}" - ); - vortex_ensure!( - end <= child.len(), - "TakeSlicesArray range {start}..{end} exceeds child length {}", - child.len() - ); - computed_len = computed_len - .checked_add(end - start) - .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; - } + let starts = slots[STARTS_SLOT] + .as_ref() + .vortex_expect("validated starts slot"); + let lengths = slots[LENGTHS_SLOT] + .as_ref() + .vortex_expect("validated lengths slot"); + let computed_len = selector_output_len(child.len(), starts, lengths)?; vortex_ensure!( data.len() == computed_len, - "TakeSlicesArray metadata length {} does not match computed range length {}", + "TakeSlicesArray metadata length {} does not match computed run length {}", data.len(), computed_len ); @@ -170,8 +173,9 @@ impl VTable for TakeSlices { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); - for range in array.slice_ranges() { - let slice = array.child().slice(range)?; + let slices = selector_slices(array.child().len(), array.starts(), array.lengths(), ctx)?; + for (start, end) in slices { + let slice = array.child().slice(start..end)?; slice.append_to_builder(builder.as_mut(), ctx)?; } Ok(ExecutionResult::done(builder.finish())) @@ -196,8 +200,9 @@ impl OperationsVTable for TakeSlices { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { + let slices = selector_slices(array.child().len(), array.starts(), array.lengths(), ctx)?; let mut logical_start = 0usize; - for &(start, end) in array.slices() { + for (start, end) in slices { let len = end - start; let logical_end = logical_start + len; if index < logical_end { @@ -214,6 +219,9 @@ impl OperationsVTable for TakeSlices { impl ValidityVTable for TakeSlices { fn validity(array: ArrayView<'_, TakeSlices>) -> VortexResult { - array.child().validity()?.take_slices(array.slices()) + array + .child() + .validity()? + .take_slices(array.starts(), array.lengths()) } } diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 214fb321093..da35f5d8d83 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -224,11 +224,13 @@ impl Validity { } } - /// Select validity values by concatenating a sequence of non-empty ranges. - pub fn take_slices(&self, slices: &[(usize, usize)]) -> VortexResult { + /// Select validity values by concatenating a sequence of contiguous runs. + pub fn take_slices(&self, starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult { match self { v @ (Self::NonNullable | Self::AllValid | Self::AllInvalid) => Ok(v.clone()), - Self::Array(is_valid) => Ok(Self::Array(is_valid.take_slices(slices.to_vec())?)), + Self::Array(is_valid) => Ok(Self::Array( + is_valid.take_slices(starts.clone(), lengths.clone())?, + )), } } From 855c40bc1174530a9b903c93f8de1f3ef2462a20 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 16:28:32 -0400 Subject: [PATCH 05/37] Add TakeSlices kernels for nested children Push TakeSlices through struct fields and execute VarBinView run gathers by copying view ranges while reusing data buffers. Signed-off-by: Daniel King --- .../src/arrays/struct_/compute/rules.rs | 2 + .../src/arrays/struct_/compute/take.rs | 28 +++++++++ vortex-array/src/arrays/take_slices/tests.rs | 60 +++++++++++++++++++ .../src/arrays/varbinview/compute/take.rs | 40 +++++++++++++ .../src/arrays/varbinview/vtable/kernel.rs | 7 +++ 5 files changed, 137 insertions(+) diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 01981803239..6e700cad329 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -16,6 +16,7 @@ use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::arrays::slice::SliceReduceAdaptor; use crate::arrays::struct_::StructArrayExt; use crate::arrays::struct_::compute::cast::struct_cast_fields; +use crate::arrays::take_slices::TakeSlicesReduceAdaptor; use crate::builtins::ArrayBuiltins; use crate::matcher::Matcher; use crate::optimizer::rules::ArrayParentReduceRule; @@ -31,6 +32,7 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&MaskReduceAdaptor(Struct)), ParentRuleSet::lift(&SliceReduceAdaptor(Struct)), ParentRuleSet::lift(&TakeReduceAdaptor(Struct)), + ParentRuleSet::lift(&TakeSlicesReduceAdaptor(Struct)), ]); pub(crate) fn struct_cast_reduce_parent( diff --git a/vortex-array/src/arrays/struct_/compute/take.rs b/vortex-array/src/arrays/struct_/compute/take.rs index a63c6c85cf6..7eb70703416 100644 --- a/vortex-array/src/arrays/struct_/compute/take.rs +++ b/vortex-array/src/arrays/struct_/compute/take.rs @@ -10,6 +10,8 @@ use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::dict::TakeReduce; use crate::arrays::struct_::StructArrayExt; +use crate::arrays::take_slices::TakeSlicesReduce; +use crate::arrays::take_slices::selector_output_len; use crate::builtins::ArrayBuiltins; use crate::scalar::Scalar; use crate::validity::Validity; @@ -50,3 +52,29 @@ impl TakeReduce for Struct { .map(Some) } } + +impl TakeSlicesReduce for Struct { + fn take_slices( + array: ArrayView<'_, Struct>, + starts: &ArrayRef, + lengths: &ArrayRef, + ) -> VortexResult> { + let fields = array + .iter_unmasked_fields() + .map(|field| field.take_slices(starts.clone(), lengths.clone())) + .collect::>>()?; + let len = fields.first().map_or_else( + || selector_output_len(array.len(), starts, lengths), + |f| Ok(f.len()), + )?; + + StructArray::try_new_with_dtype( + fields, + array.struct_fields().clone(), + len, + array.validity()?.take_slices(starts, lengths)?, + ) + .map(StructArray::into_array) + .map(Some) + } +} diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs index ec76643967b..12e5376640b 100644 --- a/vortex-array/src/arrays/take_slices/tests.rs +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -12,11 +14,15 @@ use crate::array_session; use crate::arrays::DictArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; +use crate::arrays::StructArray; use crate::arrays::TakeSlices; use crate::arrays::TakeSlicesArray; +use crate::arrays::VarBinViewArray; use crate::arrays::take_slices::TakeSlicesArrayExt; use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::assert_arrays_eq; +use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::kernel::ExecuteParentKernel; use crate::validity::Validity; @@ -145,6 +151,60 @@ fn primitive_take_slices_execute_parent_copies_runs_directly() -> VortexResult<( Ok(()) } +#[test] +fn struct_take_slices_pushes_runs_to_fields() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = StructArray::try_new( + FieldNames::from(["id", "name"]), + vec![ + PrimitiveArray::from_iter(0i32..6).into_array(), + VarBinViewArray::from_iter_str(["a", "b", "c", "d", "e", "f"]).into_array(), + ], + 6, + Validity::NonNullable, + )? + .into_array(); + + let actual = take_slices(&array, &[(3, 2), (0, 2)])?; + let expected = StructArray::try_new( + FieldNames::from(["id", "name"]), + vec![ + PrimitiveArray::from_iter([3i32, 4, 0, 1]).into_array(), + VarBinViewArray::from_iter_str(["d", "e", "a", "b"]).into_array(), + ], + 4, + Validity::NonNullable, + )?; + + assert!(actual.is::()); + assert!(!actual.is::()); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn varbinview_take_slices_execute_reuses_data_buffers() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str([ + "long-string-value-000000", + "long-string-value-000001", + "long-string-value-000002", + "long-string-value-000003", + "long-string-value-000004", + "long-string-value-000005", + ]); + + let actual = take_slices(&array.clone().into_array(), &[(3, 2), (0, 2)])? + .execute::(&mut ctx)?; + + assert!(Arc::ptr_eq(actual.data_buffers(), array.data_buffers())); + assert_eq!(actual.bytes_at(0).as_slice(), b"long-string-value-000003"); + assert_eq!(actual.bytes_at(1).as_slice(), b"long-string-value-000004"); + assert_eq!(actual.bytes_at(2).as_slice(), b"long-string-value-000000"); + assert_eq!(actual.bytes_at(3).as_slice(), b"long-string-value-000001"); + Ok(()) +} + #[test] fn take_slices_empty_runs_return_empty_canonical() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index ce10abda3d0..243e127ef53 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -6,7 +6,9 @@ use std::sync::Arc; use num_traits::AsPrimitive; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -17,6 +19,8 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::take_slices::TakeSlicesExecute; +use crate::arrays::take_slices::selector_slices; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; use crate::executor::ExecutionCtx; @@ -57,6 +61,42 @@ impl TakeExecute for VarBinView { } } +impl TakeSlicesExecute for VarBinView { + fn take_slices( + array: ArrayView<'_, VarBinView>, + starts: &ArrayRef, + lengths: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let slices = selector_slices(array.len(), starts, lengths, ctx)?; + let validity = array.validity()?.take_slices(starts, lengths)?; + let len = slices.iter().try_fold(0usize, |len, &(start, end)| { + len.checked_add(end - start) + .ok_or_else(|| vortex_err!("TakeSlices output length overflow")) + })?; + + let source_views = array.views(); + let mut views = BufferMut::::with_capacity(len); + for (start, end) in slices { + views.extend_from_slice(&source_views[start..end]); + } + + // SAFETY: views are copied from a valid VarBinView array and still reference the same + // backing data buffers, so all view pointers remain valid. + unsafe { + Ok(Some( + VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(views.freeze().into_byte_buffer()), + Arc::clone(array.data_buffers()), + array.dtype().clone(), + validity, + ) + .into_array(), + )) + } + } +} + fn take_views>( views_ref: &[BinaryView], indices: &[I], diff --git a/vortex-array/src/arrays/varbinview/vtable/kernel.rs b/vortex-array/src/arrays/varbinview/vtable/kernel.rs index e09b381d590..c924f196104 100644 --- a/vortex-array/src/arrays/varbinview/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbinview/vtable/kernel.rs @@ -5,8 +5,10 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Dict; +use crate::arrays::TakeSlices; use crate::arrays::VarBinView; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; @@ -18,5 +20,10 @@ pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); kernels.register_execute_parent_kernel(Cast.id(), VarBinView, CastExecuteAdaptor(VarBinView)); kernels.register_execute_parent_kernel(Dict.id(), VarBinView, TakeExecuteAdaptor(VarBinView)); + kernels.register_execute_parent_kernel( + TakeSlices.id(), + VarBinView, + TakeSlicesExecuteAdaptor(VarBinView), + ); kernels.register_execute_parent_kernel(Zip.id(), VarBinView, ZipExecuteAdaptor(VarBinView)); } From 7b2de210352fbe91dcd5aca535caf511933a5569 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 16:31:31 -0400 Subject: [PATCH 06/37] Address TakeSlices review comments Clarify signed fixed-size-list index reinterpretation, remove redundant selector length validation, and hoist ListView validity during rebuild. Signed-off-by: Daniel King --- vortex-array/src/arrays/fixed_size_list/compute/take.rs | 3 ++- vortex-array/src/arrays/listview/rebuild.rs | 7 ++++--- vortex-array/src/arrays/take_slices/mod.rs | 6 ------ 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 10d6fa30bc5..8db5c4b4047 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -52,7 +52,8 @@ fn take_with_indices( let list_size = array.list_size() as usize; let indices_array = indices.clone().execute::(ctx)?; - // Reinterpret to unsigned so `as_slice::` (with unsigned `I`) matches; values are unchanged. + // Bit-identical reinterpret so `as_slice::` (with unsigned `I`) matches. Negative signed + // inputs become large unsigned values and are rejected by the bounds check below. let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned()); if list_size == 0 { diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 5b64a71aaf5..809ef2183f1 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -165,14 +165,15 @@ impl ListViewArray { let mut new_offsets = BufferMut::::with_capacity(len); let mut new_sizes = BufferMut::::with_capacity(len); + let validity = self.validity()?; // Resolve validity to a mask once instead of probing it per row: `execute_is_valid` // executes a scalar on every call for array-backed validity, which is O(len) work repeated // `len` times. - let validity = self.validity()?.execute_mask(len, ctx)?; + let validity_mask = validity.execute_mask(len, ctx)?; let mut n_elements = NewOffset::zero(); - for (index, is_valid) in validity.iter().enumerate() { + for (index, is_valid) in validity_mask.iter().enumerate() { if !is_valid { new_offsets.push(n_elements); new_sizes.push(S::zero()); @@ -202,7 +203,7 @@ impl ListViewArray { // 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()?) + ListViewArray::new_unchecked(elements, offsets, sizes, validity) .with_zero_copy_to_list(true) }) } diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index b017fc7fe5b..b4f9b300d59 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -171,12 +171,6 @@ fn primitive_selector_slices( starts: ArrayView<'_, Primitive>, lengths: ArrayView<'_, Primitive>, ) -> VortexResult> { - vortex_ensure!( - starts.len() == lengths.len(), - "TakeSlicesArray selectors must have equal length, got starts {} and lengths {}", - starts.len(), - lengths.len() - ); match_each_unsigned_integer_ptype!(starts.ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { selector_slices_typed::( From 01c47963a7370989345f7e1fa49c1d918aef67b0 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 16:33:02 -0400 Subject: [PATCH 07/37] Test nullable FSL take indices Add explicit coverage that nullable take indices over a non-nullable fixed-size-list produce nullable output rows. Signed-off-by: Daniel King --- .../src/arrays/fixed_size_list/tests/take.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index 11bfcf218fa..d013775cbca 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -237,6 +237,22 @@ fn test_take_fsl_with_null_indices_preserves_elements() { assert_arrays_eq!(expected, result, &mut ctx); } +#[test] +fn test_take_non_nullable_fsl_nullable_indices_makes_nullable_output() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); + let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); + + let indices = PrimitiveArray::from_option_iter([Some(2u32), None, Some(0)]); + let result = fsl.take(indices.into_array())?; + + assert_eq!(result.dtype().nullability(), Nullability::Nullable); + assert!(!result.execute_scalar(0, &mut ctx)?.is_null()); + assert!(result.execute_scalar(1, &mut ctx)?.is_null()); + assert!(!result.execute_scalar(2, &mut ctx)?.is_null()); + Ok(()) +} + // List offsets must not truncate when small index types select large lists. #[rstest] #[case::non_nullable( From 10701853595c6a966815a8c7138139efacf6dee8 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 17:04:00 -0400 Subject: [PATCH 08/37] Regenerate stale benchmark Vortex files Validate existing converted Vortex files against the source Parquet footer row count before reusing them. Regenerate files that are unreadable or have stale row counts so benchmark data caches cannot preserve empty Vortex artifacts. Signed-off-by: Daniel King --- vortex-bench/src/conversions.rs | 102 ++++++++++++++++++++++++++++---- 1 file changed, 91 insertions(+), 11 deletions(-) diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 92064c25785..a5bef782fb6 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -43,6 +43,7 @@ use vortex::dtype::extension::ExtDType; use vortex::dtype::extension::ExtDTypeRef; use vortex::error::VortexResult; use vortex::error::vortex_err; +use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexWriteOptions; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; @@ -65,6 +66,7 @@ use crate::CompactionStrategy; use crate::Format; use crate::SESSION; use crate::utils::file::idempotent_async; +use crate::utils::file::temp_download_filepath; /// Memory budget per concurrent conversion stream in GB. This is somewhat arbitary. const MEMORY_PER_STREAM_GB: u64 = 4; @@ -167,6 +169,84 @@ pub async fn convert_parquet_file_to_vortex( Ok(()) } +async fn parquet_file_row_count(parquet_path: &Path) -> anyhow::Result { + let file = File::open(parquet_path).await?; + let builder = ParquetRecordBatchStreamBuilder::new(file).await?; + let row_count = builder.metadata().file_metadata().num_rows(); + u64::try_from(row_count) + .map_err(|_| anyhow::anyhow!("Parquet file has a negative row count: {parquet_path:?}")) +} + +async fn vortex_file_row_count(vortex_path: &Path) -> anyhow::Result { + Ok(SESSION + .open_options() + .open_path(vortex_path) + .await? + .row_count()) +} + +async fn existing_vortex_file_is_current( + parquet_path: &Path, + vortex_path: &Path, +) -> anyhow::Result { + if !vortex_path.exists() { + return Ok(false); + } + + let expected_row_count = parquet_file_row_count(parquet_path).await?; + match vortex_file_row_count(vortex_path).await { + Ok(actual_row_count) if actual_row_count == expected_row_count => Ok(true), + Ok(actual_row_count) => { + info!( + parquet_path = %parquet_path.display(), + vortex_path = %vortex_path.display(), + expected_row_count, + actual_row_count, + "Regenerating Vortex file with stale row count" + ); + Ok(false) + } + Err(error) => { + info!( + parquet_path = %parquet_path.display(), + vortex_path = %vortex_path.display(), + %error, + "Regenerating unreadable Vortex file" + ); + Ok(false) + } + } +} + +async fn convert_parquet_file_to_current_vortex( + parquet_path: &Path, + output_path: &Path, + compaction: CompactionStrategy, +) -> anyhow::Result { + if existing_vortex_file_is_current(parquet_path, output_path).await? { + return Ok(output_path.to_path_buf()); + } + + if let Some(parent) = output_path.parent() { + create_dir_all(parent).await?; + } + + let temp_path = temp_download_filepath(); + info!( + parquet_path = %parquet_path.display(), + vortex_path = %output_path.display(), + ?compaction, + "Processing Parquet file as Vortex" + ); + convert_parquet_file_to_vortex(parquet_path, &temp_path, compaction).await?; + if output_path.exists() { + tokio::fs::remove_file(output_path).await?; + } + tokio::fs::rename(&temp_path, output_path).await?; + + Ok(output_path.to_path_buf()) +} + /// Whether `path` points at SpatialBench data. fn is_spatialbench(path: &Path) -> bool { path.components() @@ -221,7 +301,8 @@ fn no_dict_layout() -> Arc { /// `{input_path}/vortex-file-compressed/` (for Default compaction) or /// `{input_path}/vortex-compact/` (for Compact compaction). /// -/// The conversion is idempotent: existing Vortex files will not be regenerated. +/// Existing Vortex files are reused when their footer row count matches the source Parquet file. +/// Unreadable or stale Vortex files are regenerated. pub async fn convert_parquet_directory_to_vortex( input_path: &Path, compaction: CompactionStrategy, @@ -247,7 +328,7 @@ pub async fn convert_parquet_directory_to_vortex( .filter(|entry| entry.path().extension().is_some_and(|e| e == "parquet")); let concurrency = calculate_concurrency(); - futures::stream::iter(iter) + let conversion_results = futures::stream::iter(iter) .map(|dir_entry| { let filename = { let mut temp = dir_entry.path(); @@ -259,16 +340,12 @@ pub async fn convert_parquet_directory_to_vortex( tokio::spawn( async move { - idempotent_async(output_path.as_path(), move |vtx_file| async move { - info!( - "Processing file '{filename}' with {:?} strategy", - compaction - ); - convert_parquet_file_to_vortex(&parquet_file_path, &vtx_file, compaction) - .await - }) + convert_parquet_file_to_current_vortex( + &parquet_file_path, + &output_path, + compaction, + ) .await - .expect("Failed to write Vortex file") } .in_current_span(), ) @@ -276,6 +353,9 @@ pub async fn convert_parquet_directory_to_vortex( .buffer_unordered(concurrency) .try_collect::>() .await?; + for result in conversion_results { + result?; + } Ok(()) } From a279667da40cda8f07a208d4693edab769cc2c3e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 17:16:49 -0400 Subject: [PATCH 09/37] Tighten TakeSlices local reasoning Add context-aware TakeSlices construction and validity paths so callers with an ExecutionCtx do not fall back to legacy session execution. Centralize run selector validation/output length in RunSelectors and use it from execution kernels.\n\nClarify that TakeSlices remains a lazy non-serializable compute encoding like Slice, remove the unrelated AGENTS.md guidance change, tighten Struct length derivation, document the FSL null-placeholder run invariant, and rename the benchmark cache helper to describe its row-count-only check. Signed-off-by: Daniel King --- AGENTS.md | 16 --- vortex-array/src/array/erased.rs | 15 ++ .../arrays/fixed_size_list/compute/take.rs | 22 ++- vortex-array/src/arrays/listview/rebuild.rs | 11 +- .../arrays/primitive/compute/take_slices.rs | 16 +-- .../src/arrays/struct_/compute/take.rs | 5 +- vortex-array/src/arrays/take_slices/array.rs | 26 +++- vortex-array/src/arrays/take_slices/mod.rs | 135 +++++++++++------- vortex-array/src/arrays/take_slices/rules.rs | 11 +- vortex-array/src/arrays/take_slices/vtable.rs | 15 +- .../src/arrays/varbinview/compute/take.rs | 16 +-- vortex-array/src/validity.rs | 23 ++- vortex-bench/src/conversions.rs | 4 +- 13 files changed, 194 insertions(+), 121 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 4c85be00da5..30ddeeb9515 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -151,22 +151,6 @@ Notes: self-explanatory code. - Keep public APIs small and consistent with neighboring crates. -## Local Reasoning and Edge Cases - -- Before adding helper state, boolean flags, or small enums to handle edge cases inside a general - loop, first ask whether the input space naturally splits into a few structural cases with - different invariants. If it does, dispatch once near the top of the function to clearly named - helpers for each case. -- Prefer helpers whose names state the invariant they rely on and whose bodies only handle that - case. A reader should be able to prove each helper correct without carrying unrelated cases in - their head. -- Avoid making a broad path responsible for exceptional setup, fallback construction, and normal - processing at the same time. If a fallback is only valid under a narrow precondition, isolate it - behind that precondition instead of threading flags through the main algorithm. -- When a branch exists only to preserve a representation invariant, put that invariant at the - branch boundary and keep the subsequent code free of defensive state that is impossible under - the branch's preconditions. - ## Performance: avoid hidden-cost accessors in hot loops The most common performance trap in this codebase is calling a *per-element accessor that diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 8be60976fe3..3718b5b97fe 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -270,12 +270,27 @@ impl ArrayRef { /// /// The output is the concatenation of `self[starts[i]..starts[i] + lengths[i]]` for each /// selector row. + /// + /// Prefer [`Self::take_slices_with_ctx`] when the caller already has an execution context. pub fn take_slices(&self, starts: ArrayRef, lengths: ArrayRef) -> VortexResult { TakeSlicesArray::try_new(self.clone(), starts, lengths)? .into_array() .optimize() } + /// Wraps the array in a [`TakeSlicesArray`] using the caller's execution context to validate + /// selector arrays and compute the output length. + pub fn take_slices_with_ctx( + &self, + starts: ArrayRef, + lengths: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + TakeSlicesArray::try_new_with_ctx(self.clone(), starts, lengths, ctx)? + .into_array() + .optimize() + } + /// Fetch the scalar at the given index. #[deprecated( note = "Use `execute_scalar` instead, which allows passing an execution context for more \ diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 8db5c4b4047..a1fe846979b 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -136,13 +136,14 @@ fn take_non_empty_non_degenerate_fsl( if array.dtype().is_nullable() || indices_array.dtype().is_nullable() { take_nullable_non_empty_fsl::(array, indices_array, ctx) } else { - take_non_nullable_non_empty_fsl::(array, indices_array) + take_non_nullable_non_empty_fsl::(array, indices_array, ctx) } } fn take_non_nullable_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, ) -> VortexResult { let list_size = array.list_size() as usize; let array_len = array.as_ref().len(); @@ -157,7 +158,7 @@ fn take_non_nullable_non_empty_fsl( starts.push(usize_to_u64(start, "FixedSizeList take element start")?); } - let new_elements = take_element_runs(array.elements(), starts, list_size)?; + let new_elements = take_element_runs(array.elements(), starts, list_size, ctx)?; ensure_elements_len(new_elements.len(), expected_elements_len)?; // SAFETY: `starts` contains one checked run of `list_size` elements for each output row, @@ -189,13 +190,12 @@ fn take_nullable_non_empty_fsl( .execute_mask(array.as_ref().len(), ctx)?; let indices_validity = indices_validity_mask(&indices_array, ctx)?; - let null_element_start = 0u64; let mut starts = Vec::with_capacity(new_len); let mut new_validity_builder = BitBufferMut::with_capacity(new_len); for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { if !is_index_valid { - starts.push(null_element_start); + starts.push(null_element_run_start()); new_validity_builder.append(false); continue; } @@ -203,7 +203,7 @@ fn take_nullable_non_empty_fsl( let data_idx = index_to_usize(data_idx)?; let start = list_start(data_idx, list_size, array_len)?; if !array_validity.value(data_idx) { - starts.push(null_element_start); + starts.push(null_element_run_start()); new_validity_builder.append(false); continue; } @@ -212,7 +212,7 @@ fn take_nullable_non_empty_fsl( new_validity_builder.append(true); } - let new_elements = take_element_runs(array.elements(), starts, list_size)?; + let new_elements = take_element_runs(array.elements(), starts, list_size, ctx)?; ensure_elements_len(new_elements.len(), expected_elements_len)?; let new_validity = Validity::from(new_validity_builder.freeze()); @@ -303,15 +303,23 @@ fn take_element_runs( elements: &ArrayRef, starts: Vec, length: usize, + ctx: &mut ExecutionCtx, ) -> VortexResult { let run_count = starts.len(); let length = usize_to_u64(length, "FixedSizeList take element run length")?; - elements.take_slices( + elements.take_slices_with_ctx( PrimitiveArray::from_iter(starts).into_array(), PrimitiveArray::from_iter(std::iter::repeat_n(length, run_count)).into_array(), + ctx, ) } +fn null_element_run_start() -> u64 { + // Null output rows still need placeholder child elements so the FSL elements length stays + // `rows * list_size`. Any in-bounds run is fine because row validity hides these elements. + 0 +} + fn index_to_usize(index: I) -> VortexResult { index .to_usize() diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 809ef2183f1..2054d2e5357 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -146,6 +146,9 @@ impl ListViewArray { } /// Rebuilds elements using a single contiguous-run gather over the element child. + /// + /// `take_slices` is the rebuild primitive for ListView packing. Children with specialized + /// kernels can preserve their representation; other children fall back to generic execution. fn rebuild_with_take_slices( &self, ctx: &mut ExecutionCtx, @@ -189,9 +192,11 @@ impl ListViewArray { let new_sizes = new_sizes.freeze(); let lengths = PrimitiveArray::new(new_sizes.clone(), Validity::NonNullable); - let elements = self - .elements() - .take_slices(offsets_canonical.into_array(), lengths.into_array())?; + let elements = self.elements().take_slices_with_ctx( + offsets_canonical.into_array(), + lengths.into_array(), + ctx, + )?; // Built unsigned; reinterpret back to the signed-preserving result types. let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(new_offset_ptype) diff --git a/vortex-array/src/arrays/primitive/compute/take_slices.rs b/vortex-array/src/arrays/primitive/compute/take_slices.rs index f484f64d68b..9e94594ac25 100644 --- a/vortex-array/src/arrays/primitive/compute/take_slices.rs +++ b/vortex-array/src/arrays/primitive/compute/take_slices.rs @@ -3,15 +3,14 @@ use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_err; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; +use crate::arrays::take_slices::RunSelectors; use crate::arrays::take_slices::TakeSlicesExecute; -use crate::arrays::take_slices::selector_slices; use crate::executor::ExecutionCtx; use crate::match_each_native_ptype; @@ -22,17 +21,16 @@ impl TakeSlicesExecute for Primitive { lengths: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let slices = selector_slices(array.len(), starts, lengths, ctx)?; - let validity = array.validity()?.take_slices(starts, lengths)?; + let selectors = RunSelectors::new(array.len(), starts, lengths, ctx)?; + let validity = array + .validity()? + .take_slices_with_ctx(starts, lengths, ctx)?; match_each_native_ptype!(array.ptype(), |T| { let source = array.as_slice::(); - let len = slices.iter().try_fold(0usize, |len, &(start, end)| { - len.checked_add(end - start) - .ok_or_else(|| vortex_err!("TakeSlices output length overflow")) - })?; + let len = selectors.output_len(); let mut values = BufferMut::::with_capacity(len); - for (start, end) in slices { + for &(start, end) in selectors.slices() { values.extend_from_slice(&source[start..end]); } diff --git a/vortex-array/src/arrays/struct_/compute/take.rs b/vortex-array/src/arrays/struct_/compute/take.rs index 7eb70703416..95fad14491d 100644 --- a/vortex-array/src/arrays/struct_/compute/take.rs +++ b/vortex-array/src/arrays/struct_/compute/take.rs @@ -59,14 +59,11 @@ impl TakeSlicesReduce for Struct { starts: &ArrayRef, lengths: &ArrayRef, ) -> VortexResult> { + let len = selector_output_len(array.len(), starts, lengths)?; let fields = array .iter_unmasked_fields() .map(|field| field.take_slices(starts.clone(), lengths.clone())) .collect::>>()?; - let len = fields.first().map_or_else( - || selector_output_len(array.len(), starts, lengths), - |f| Ok(f.len()), - )?; StructArray::try_new_with_dtype( fields, diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs index b10c698b1c4..2fcaec87737 100644 --- a/vortex-array/src/arrays/take_slices/array.rs +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -9,11 +9,13 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::arrays::TakeSlices; use crate::arrays::take_slices::selector_output_len; +use crate::arrays::take_slices::selector_output_len_with_ctx; /// The child array selected by the run sequence. pub(super) const CHILD_SLOT: usize = 0; @@ -79,14 +81,32 @@ impl TakeSlicesData { impl Array { /// Constructs a new validated `TakeSlicesArray`. + /// + /// Prefer [`Self::try_new_with_ctx`] when the caller already has an execution context. pub fn try_new(child: ArrayRef, starts: ArrayRef, lengths: ArrayRef) -> VortexResult { - let dtype = child.dtype().clone(); let len = selector_output_len(child.len(), &starts, &lengths)?; + Ok(Self::new_validated(child, starts, lengths, len)) + } + + /// Constructs a new validated `TakeSlicesArray` using the caller's execution context to + /// validate selector arrays and compute the output length. + pub fn try_new_with_ctx( + child: ArrayRef, + starts: ArrayRef, + lengths: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let len = selector_output_len_with_ctx(child.len(), &starts, &lengths, ctx)?; + Ok(Self::new_validated(child, starts, lengths, len)) + } + + fn new_validated(child: ArrayRef, starts: ArrayRef, lengths: ArrayRef, len: usize) -> Self { + let dtype = child.dtype().clone(); let data = TakeSlicesData::new(len); // SAFETY: `selector_output_len` validates selector dtypes, run bounds, and computes `len`; // the outer dtype is copied from the child, and all required slots are populated. - Ok(unsafe { + unsafe { Array::from_parts_unchecked( ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![ Some(child), @@ -94,6 +114,6 @@ impl Array { Some(lengths) ]), ) - }) + } } } diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index b4f9b300d59..44818990a99 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -86,8 +86,8 @@ fn trivial_take_slices_with_ctx( lengths: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let slices = selector_slices(child.len(), starts, lengths, ctx)?; - trivial_take_slices_from_ranges(child, &slices) + let selectors = RunSelectors::new(child.len(), starts, lengths, ctx)?; + trivial_take_slices_from_ranges(child, selectors.slices()) } fn trivial_take_slices_from_ranges( @@ -114,11 +114,16 @@ pub(super) fn selector_output_len( lengths: &ArrayRef, ) -> VortexResult { let mut ctx = legacy_execution_ctx(); - let slices = selector_slices(child_len, starts, lengths, &mut ctx)?; - slices.iter().try_fold(0usize, |len, &(start, end)| { - len.checked_add(end - start) - .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow")) - }) + selector_output_len_with_ctx(child_len, starts, lengths, &mut ctx) +} + +pub(super) fn selector_output_len_with_ctx( + child_len: usize, + starts: &ArrayRef, + lengths: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + Ok(RunSelectors::new(child_len, starts, lengths, ctx)?.output_len()) } #[allow(clippy::disallowed_methods)] @@ -126,16 +131,75 @@ fn legacy_execution_ctx() -> ExecutionCtx { legacy_session().create_execution_ctx() } -pub(super) fn selector_slices( +#[derive(Debug)] +pub(super) struct RunSelectors { + slices: Vec<(usize, usize)>, + output_len: usize, +} + +impl RunSelectors { + pub(super) fn new( + child_len: usize, + starts: &ArrayRef, + lengths: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + check_selector_arrays(starts, lengths)?; + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + primitive_run_selectors(child_len, starts.as_view(), lengths.as_view()) + } + + pub(super) fn slices(&self) -> &[(usize, usize)] { + &self.slices + } + + pub(super) fn output_len(&self) -> usize { + self.output_len + } +} + +fn run_selectors_from_slices( child_len: usize, - starts: &ArrayRef, - lengths: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - check_selector_arrays(starts, lengths)?; - let starts = starts.clone().execute::(ctx)?; - let lengths = lengths.clone().execute::(ctx)?; - primitive_selector_slices(child_len, starts.as_view(), lengths.as_view()) + starts: &[S], + lengths: &[L], +) -> VortexResult { + let mut slices = Vec::with_capacity(starts.len()); + let mut output_len = 0usize; + for (&start, &length) in starts.iter().zip(lengths) { + let start = selector_to_usize("start", start)?; + let length = selector_to_usize("length", length)?; + let end = start.checked_add(length).ok_or_else(|| { + vortex_err!("TakeSlicesArray run overflow for start {start} and length {length}") + })?; + vortex_ensure!( + end <= child_len, + "TakeSlicesArray run {start}..{end} exceeds child array length {child_len}" + ); + if length != 0 { + output_len = output_len + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; + slices.push((start, end)); + } + } + Ok(RunSelectors { slices, output_len }) +} + +fn primitive_run_selectors( + child_len: usize, + starts: ArrayView<'_, Primitive>, + lengths: ArrayView<'_, Primitive>, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + run_selectors_from_slices::( + child_len, + starts.as_slice::(), + lengths.as_slice::(), + ) + }) + }) } pub(super) fn check_selector_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult<()> { @@ -166,45 +230,6 @@ fn check_selector_dtype(name: &str, selector: &ArrayRef) -> VortexResult<()> { } } -fn primitive_selector_slices( - child_len: usize, - starts: ArrayView<'_, Primitive>, - lengths: ArrayView<'_, Primitive>, -) -> VortexResult> { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - selector_slices_typed::( - child_len, - starts.as_slice::(), - lengths.as_slice::(), - ) - }) - }) -} - -fn selector_slices_typed( - child_len: usize, - starts: &[S], - lengths: &[L], -) -> VortexResult> { - let mut slices = Vec::with_capacity(starts.len()); - for (&start, &length) in starts.iter().zip(lengths) { - let start = selector_to_usize("start", start)?; - let length = selector_to_usize("length", length)?; - let end = start.checked_add(length).ok_or_else(|| { - vortex_err!("TakeSlicesArray run overflow for start {start} and length {length}") - })?; - vortex_ensure!( - end <= child_len, - "TakeSlicesArray run {start}..{end} exceeds child array length {child_len}" - ); - if length != 0 { - slices.push((start, end)); - } - } - Ok(slices) -} - fn selector_to_usize(name: &str, value: T) -> VortexResult { value .to_usize() diff --git a/vortex-array/src/arrays/take_slices/rules.rs b/vortex-array/src/arrays/take_slices/rules.rs index d578fafdc28..9f77c4e0ca1 100644 --- a/vortex-array/src/arrays/take_slices/rules.rs +++ b/vortex-array/src/arrays/take_slices/rules.rs @@ -12,10 +12,10 @@ use crate::array::ArrayView; use crate::arrays::PrimitiveArray; use crate::arrays::TakeSlices; use crate::arrays::TakeSlicesArray; +use crate::arrays::take_slices::RunSelectors; use crate::arrays::take_slices::TakeSlicesArrayExt; use crate::arrays::take_slices::TakeSlicesReduce; use crate::arrays::take_slices::TakeSlicesReduceAdaptor; -use crate::arrays::take_slices::selector_slices; use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; @@ -41,17 +41,18 @@ impl TakeSlicesReduce for TakeSlices { lengths: &ArrayRef, ) -> VortexResult> { let mut ctx = legacy_execution_ctx(); - let inner = selector_slices( + let inner = RunSelectors::new( array.child().len(), array.starts(), array.lengths(), &mut ctx, )?; - let outer = selector_slices(array.len(), starts, lengths, &mut ctx)?; - let combined = project_slices(&inner, &outer); + let outer = RunSelectors::new(array.len(), starts, lengths, &mut ctx)?; + let combined = project_slices(inner.slices(), outer.slices()); let (starts, lengths) = selectors_from_slices(&combined)?; Ok(Some( - TakeSlicesArray::try_new(array.child().clone(), starts, lengths)?.into_array(), + TakeSlicesArray::try_new_with_ctx(array.child().clone(), starts, lengths, &mut ctx)? + .into_array(), )) } } diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index bae2619b0e5..c5cf5331c39 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -24,6 +24,7 @@ use crate::array::OperationsVTable; use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; +use crate::arrays::take_slices::RunSelectors; use crate::arrays::take_slices::TakeSlicesArrayExt; use crate::arrays::take_slices::array::CHILD_SLOT; use crate::arrays::take_slices::array::LENGTHS_SLOT; @@ -34,7 +35,6 @@ use crate::arrays::take_slices::array::TakeSlicesData; use crate::arrays::take_slices::rules::PARENT_RULES; use crate::arrays::take_slices::rules::RULES; use crate::arrays::take_slices::selector_output_len; -use crate::arrays::take_slices::selector_slices; use crate::buffer::BufferHandle; use crate::builders::builder_with_capacity_in; use crate::dtype::DType; @@ -48,6 +48,9 @@ use crate::validity::Validity; pub type TakeSlicesArray = Array; /// Contiguous-run gather selection encoding. +/// +/// Like [`crate::arrays::Slice`], this is a lazy compute encoding and is not serialized as a file +/// encoding. #[derive(Clone, Debug)] pub struct TakeSlices; @@ -173,8 +176,9 @@ impl VTable for TakeSlices { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); - let slices = selector_slices(array.child().len(), array.starts(), array.lengths(), ctx)?; - for (start, end) in slices { + let selectors = + RunSelectors::new(array.child().len(), array.starts(), array.lengths(), ctx)?; + for &(start, end) in selectors.slices() { let slice = array.child().slice(start..end)?; slice.append_to_builder(builder.as_mut(), ctx)?; } @@ -200,9 +204,10 @@ impl OperationsVTable for TakeSlices { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let slices = selector_slices(array.child().len(), array.starts(), array.lengths(), ctx)?; + let selectors = + RunSelectors::new(array.child().len(), array.starts(), array.lengths(), ctx)?; let mut logical_start = 0usize; - for (start, end) in slices { + for &(start, end) in selectors.slices() { let len = end - start; let logical_end = logical_start + len; if index < logical_end { diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 243e127ef53..a21e4d1883e 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -8,7 +8,6 @@ use num_traits::AsPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -19,8 +18,8 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::take_slices::RunSelectors; use crate::arrays::take_slices::TakeSlicesExecute; -use crate::arrays::take_slices::selector_slices; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; use crate::executor::ExecutionCtx; @@ -68,16 +67,15 @@ impl TakeSlicesExecute for VarBinView { lengths: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let slices = selector_slices(array.len(), starts, lengths, ctx)?; - let validity = array.validity()?.take_slices(starts, lengths)?; - let len = slices.iter().try_fold(0usize, |len, &(start, end)| { - len.checked_add(end - start) - .ok_or_else(|| vortex_err!("TakeSlices output length overflow")) - })?; + let selectors = RunSelectors::new(array.len(), starts, lengths, ctx)?; + let validity = array + .validity()? + .take_slices_with_ctx(starts, lengths, ctx)?; + let len = selectors.output_len(); let source_views = array.views(); let mut views = BufferMut::::with_capacity(len); - for (start, end) in slices { + for &(start, end) in selectors.slices() { views.extend_from_slice(&source_views[start..end]); } diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index da35f5d8d83..2e02e59a9ab 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -225,12 +225,29 @@ impl Validity { } /// Select validity values by concatenating a sequence of contiguous runs. + /// + /// Prefer [`Self::take_slices_with_ctx`] when the caller already has an execution context. + #[allow(clippy::disallowed_methods)] pub fn take_slices(&self, starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult { + let mut ctx = legacy_session().create_execution_ctx(); + self.take_slices_with_ctx(starts, lengths, &mut ctx) + } + + /// Select validity values by concatenating a sequence of contiguous runs, using the caller's + /// execution context for array-backed validity. + pub fn take_slices_with_ctx( + &self, + starts: &ArrayRef, + lengths: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { match self { v @ (Self::NonNullable | Self::AllValid | Self::AllInvalid) => Ok(v.clone()), - Self::Array(is_valid) => Ok(Self::Array( - is_valid.take_slices(starts.clone(), lengths.clone())?, - )), + Self::Array(is_valid) => Ok(Self::Array(is_valid.take_slices_with_ctx( + starts.clone(), + lengths.clone(), + ctx, + )?)), } } diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index a5bef782fb6..57fc09f4335 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -185,7 +185,7 @@ async fn vortex_file_row_count(vortex_path: &Path) -> anyhow::Result { .row_count()) } -async fn existing_vortex_file_is_current( +async fn existing_vortex_file_has_matching_row_count( parquet_path: &Path, vortex_path: &Path, ) -> anyhow::Result { @@ -223,7 +223,7 @@ async fn convert_parquet_file_to_current_vortex( output_path: &Path, compaction: CompactionStrategy, ) -> anyhow::Result { - if existing_vortex_file_is_current(parquet_path, output_path).await? { + if existing_vortex_file_has_matching_row_count(parquet_path, output_path).await? { return Ok(output_path.to_path_buf()); } From fd480d226f5d4eebf1e5ab2c576452ef79548a11 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 11:31:08 -0400 Subject: [PATCH 10/37] Simplify TakeSlices selectors Signed-off-by: Daniel King --- Cargo.lock | 1 + vortex-array/benches/take_fsl.rs | 18 +- vortex-array/src/array/erased.rs | 66 +++-- .../arrays/fixed_size_list/compute/take.rs | 38 +-- .../src/arrays/fixed_size_list/tests/take.rs | 6 +- vortex-array/src/arrays/listview/rebuild.rs | 112 ++++++-- .../src/arrays/primitive/compute/mod.rs | 1 - .../arrays/primitive/compute/take_slices.rs | 42 --- .../src/arrays/primitive/vtable/kernel.rs | 7 - .../src/arrays/struct_/compute/rules.rs | 2 - .../src/arrays/struct_/compute/take.rs | 25 -- vortex-array/src/arrays/take_slices/array.rs | 61 ++-- vortex-array/src/arrays/take_slices/mod.rs | 270 ++---------------- vortex-array/src/arrays/take_slices/rules.rs | 106 ------- vortex-array/src/arrays/take_slices/tests.rs | 196 ++++++------- vortex-array/src/arrays/take_slices/vtable.rs | 266 +++++++++++++---- .../src/arrays/varbinview/compute/take.rs | 38 --- .../src/arrays/varbinview/vtable/kernel.rs | 7 - vortex-array/src/validity.rs | 33 ++- vortex-bench/Cargo.toml | 1 + vortex-bench/src/conversions.rs | 167 +++++++++-- 21 files changed, 670 insertions(+), 793 deletions(-) delete mode 100644 vortex-array/src/arrays/primitive/compute/take_slices.rs delete mode 100644 vortex-array/src/arrays/take_slices/rules.rs diff --git a/Cargo.lock b/Cargo.lock index e23ac0d40f8..384403520e1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9783,6 +9783,7 @@ dependencies = [ "reqwest 0.13.4", "serde", "serde_json", + "sha2 0.11.0", "spatialbench", "spatialbench-arrow", "sysinfo", diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index 7767c1bbfdf..fc2f3b373dd 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -212,19 +212,13 @@ fn take_fsl_f16_take_slices_strategy( let starts = indices .as_ref() .iter() - .map(|&idx| { - let start = idx as usize * LIST_SIZE; - start as u64 - }) + .map(|&idx| idx as usize * LIST_SIZE) .collect::>(); - let lengths = std::iter::repeat_n(LIST_SIZE as u64, indices.len()).collect::>(); - let elements = array - .elements() - .take_slices( - PrimitiveArray::from_iter(starts).into_array(), - PrimitiveArray::from_iter(lengths).into_array(), - ) - .unwrap(); + let ends = starts + .iter() + .map(|&start| start + LIST_SIZE) + .collect::>(); + let elements = array.elements().take_slices(starts, ends).unwrap(); // SAFETY: each generated run has width `LIST_SIZE`, and there is one run per input index, // so `elements.len() == indices.len() * LIST_SIZE`. diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 3718b5b97fe..7f7d8ce6cad 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -40,6 +40,7 @@ use crate::arrays::DictArray; use crate::arrays::FilterArray; use crate::arrays::Null; use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; use crate::arrays::TakeSlicesArray; use crate::arrays::VarBin; @@ -265,30 +266,44 @@ impl ArrayRef { .optimize() } - /// Wraps the array in a [`TakeSlicesArray`] such that it is logically selected by a - /// caller-provided sequence of child runs. + /// Wraps the array in a [`TakeSlicesArray`] selected by caller-provided child ranges. /// - /// The output is the concatenation of `self[starts[i]..starts[i] + lengths[i]]` for each - /// selector row. - /// - /// Prefer [`Self::take_slices_with_ctx`] when the caller already has an execution context. - pub fn take_slices(&self, starts: ArrayRef, lengths: ArrayRef) -> VortexResult { - TakeSlicesArray::try_new(self.clone(), starts, lengths)? - .into_array() - .optimize() - } - - /// Wraps the array in a [`TakeSlicesArray`] using the caller's execution context to validate - /// selector arrays and compute the output length. - pub fn take_slices_with_ctx( - &self, - starts: ArrayRef, - lengths: ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - TakeSlicesArray::try_new_with_ctx(self.clone(), starts, lengths, ctx)? - .into_array() - .optimize() + /// The output is the concatenation of `self[starts[i]..ends[i]]` for each selector row. This + /// computes the output length from `starts` and `ends`, but child bounds are checked only when + /// the lazy gather executes. + pub fn take_slices(&self, starts: Vec, ends: Vec) -> VortexResult { + vortex_ensure!( + starts.len() == ends.len(), + "TakeSlicesArray selectors must have equal length, got starts {} and ends {}", + starts.len(), + ends.len() + ); + let mut len = 0usize; + for (&start, &end) in starts.iter().zip(&ends) { + vortex_ensure!( + start <= end, + "TakeSlicesArray start {start} must be <= end {end}" + ); + len = len + .checked_add(end - start) + .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; + } + let starts = starts + .into_iter() + .map(selector_value) + .collect::>>()?; + let ends = ends + .into_iter() + .map(selector_value) + .collect::>>()?; + TakeSlicesArray::try_new( + self.clone(), + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(ends).into_array(), + len, + )? + .into_array() + .optimize() } /// Fetch the scalar at the given index. @@ -806,6 +821,11 @@ impl ArrayRef { } } +fn selector_value(value: usize) -> VortexResult { + u64::try_from(value) + .map_err(|_| vortex_err!("TakeSlicesArray selector {value} does not fit in u64")) +} + impl IntoArray for ArrayRef { #[inline(always)] fn into_array(self) -> ArrayRef { diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index a1fe846979b..352252a32ce 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -10,6 +10,7 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; +use crate::RecursiveCanonical; use crate::array::ArrayView; use crate::arrays::BoolArray; use crate::arrays::FixedSizeList; @@ -30,7 +31,7 @@ use crate::validity::Validity; /// /// `FixedSizeListArray` must rebuild its elements array because selected lists need to become /// packed from offset 0. The FSL layer translates selected list rows into ordered element runs -/// and delegates the execution strategy to the elements child via `take_slices`. +/// and delegates the execution strategy to the elements child via materialized `take_slices`. impl TakeExecute for FixedSizeList { fn take( array: ArrayView<'_, FixedSizeList>, @@ -155,7 +156,7 @@ fn take_non_nullable_non_empty_fsl( for &data_idx in indices { let data_idx = index_to_usize(data_idx)?; let start = list_start(data_idx, list_size, array_len)?; - starts.push(usize_to_u64(start, "FixedSizeList take element start")?); + starts.push(start); } let new_elements = take_element_runs(array.elements(), starts, list_size, ctx)?; @@ -208,7 +209,7 @@ fn take_nullable_non_empty_fsl( continue; } - starts.push(usize_to_u64(start, "FixedSizeList take element start")?); + starts.push(start); new_validity_builder.append(true); } @@ -301,20 +302,29 @@ fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef fn take_element_runs( elements: &ArrayRef, - starts: Vec, + starts: Vec, length: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let run_count = starts.len(); - let length = usize_to_u64(length, "FixedSizeList take element run length")?; - elements.take_slices_with_ctx( - PrimitiveArray::from_iter(starts).into_array(), - PrimitiveArray::from_iter(std::iter::repeat_n(length, run_count)).into_array(), - ctx, - ) + let ends = starts + .iter() + .map(|&start| { + start.checked_add(length).ok_or_else(|| { + vortex_err!( + "FixedSizeList take element range overflow for start {start} and length {length}" + ) + }) + }) + .collect::>>()?; + + Ok(elements + .take_slices(starts, ends)? + .execute::(ctx)? + .0 + .into_array()) } -fn null_element_run_start() -> u64 { +fn null_element_run_start() -> usize { // Null output rows still need placeholder child elements so the FSL elements length stays // `rows * list_size`. Any in-bounds run is fine because row validity hides these elements. 0 @@ -325,7 +335,3 @@ fn index_to_usize(index: I) -> VortexResult { .to_usize() .ok_or_else(|| vortex_err!("FixedSizeList take index {index} does not fit in usize")) } - -fn usize_to_u64(value: usize, name: &str) -> VortexResult { - u64::try_from(value).map_err(|_| vortex_err!("{name} {value} does not fit in u64")) -} diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index d013775cbca..ba2aa7f89bf 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -32,8 +32,8 @@ use crate::array::with_empty_buffers; use crate::array_session; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::assert_arrays_eq; @@ -398,7 +398,7 @@ fn test_take_empty_source_rejects_valid_index_after_null_index() -> VortexResult } #[test] -fn test_take_uses_take_slices_for_encoded_elements_child() -> VortexResult<()> { +fn test_take_materializes_encoded_elements_child() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let encoded_elements = NoTakeSlicesArray::wrap(buffer![10i32, 20, 30, 40, 20, 30].into_array()); let fsl = FixedSizeListArray::new(encoded_elements, 2, Validity::NonNullable, 3); @@ -407,7 +407,7 @@ fn test_take_uses_take_slices_for_encoded_elements_child() -> VortexResult<()> { let result = ::take(fsl.as_view(), &indices, &mut ctx)? .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; - assert!(result.as_::().elements().is::()); + assert!(result.as_::().elements().is::()); let expected = FixedSizeListArray::new( PrimitiveArray::from_iter([20i32, 30, 10, 20]).into_array(), 2, diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 2054d2e5357..69ca1792b6d 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -5,9 +5,12 @@ 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::ExecutionCtx; use crate::IntoArray; +use crate::RecursiveCanonical; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -147,8 +150,9 @@ impl ListViewArray { /// Rebuilds elements using a single contiguous-run gather over the element child. /// - /// `take_slices` is the rebuild primitive for ListView packing. Children with specialized - /// kernels can preserve their representation; other children fall back to generic execution. + /// `take_slices` is the rebuild primitive for ListView packing. This path materializes the + /// element gather so the rebuilt ListView has a physical packed child rather than a lazy + /// `TakeSlices` wrapper. fn rebuild_with_take_slices( &self, ctx: &mut ExecutionCtx, @@ -162,12 +166,9 @@ impl ListViewArray { 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_canonical.len(); - let mut new_offsets = BufferMut::::with_capacity(len); - let mut new_sizes = BufferMut::::with_capacity(len); let validity = self.validity()?; // Resolve validity to a mask once instead of probing it per row: `execute_is_valid` @@ -175,30 +176,23 @@ impl ListViewArray { // `len` times. let validity_mask = validity.execute_mask(len, ctx)?; - let mut n_elements = NewOffset::zero(); - for (index, is_valid) in validity_mask.iter().enumerate() { - if !is_valid { - new_offsets.push(n_elements); - new_sizes.push(S::zero()); - continue; - } - - let size = sizes_slice[index]; - - new_offsets.push(n_elements); - new_sizes.push(size); - 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 new_sizes = new_sizes.freeze(); - let lengths = PrimitiveArray::new(new_sizes.clone(), Validity::NonNullable); - let elements = self.elements().take_slices_with_ctx( - offsets_canonical.into_array(), - lengths.into_array(), - ctx, - )?; + let new_sizes = ranges.new_sizes.freeze(); + let elements = self + .elements() + .take_slices(ranges.starts, ranges.ends)? + .execute::(ctx)? + .0 + .into_array(); // Built unsigned; reinterpret back to the signed-preserving result types. - let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + let offsets = PrimitiveArray::new(ranges.new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(new_offset_ptype) .into_array(); let sizes = PrimitiveArray::new(new_sizes, Validity::NonNullable) @@ -274,6 +268,70 @@ impl ListViewArray { } } +struct RebuildRanges { + new_offsets: BufferMut, + new_sizes: BufferMut, + starts: Vec, + ends: Vec, +} + +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 = Vec::with_capacity(len); + let mut ends = Vec::with_capacity(len); + let mut n_elements = NewOffset::zero(); + + 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); + ends.push(0); + continue; + } + + let size = sizes[index]; + let start = offsets[index].to_usize().ok_or_else(|| { + vortex_err!( + "ListView rebuild offset {} does not fit in usize", + offsets[index] + ) + })?; + let length = size + .to_usize() + .ok_or_else(|| vortex_err!("ListView rebuild size {size} does not fit in usize"))?; + let end = start.checked_add(length).ok_or_else(|| { + vortex_err!( + "ListView rebuild element range overflow for start {start} and length {length}" + ) + })?; + + new_offsets.push(n_elements); + new_sizes.push(size); + starts.push(start); + ends.push(end); + n_elements += num_traits::cast(size).vortex_expect("Cast failed"); + } + + Ok(RebuildRanges { + new_offsets, + new_sizes, + starts, + ends, + }) +} + #[cfg(test)] mod tests { use vortex_buffer::BitBuffer; diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index 37aa3f98396..1ca4b17d3d5 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -8,7 +8,6 @@ mod mask; pub(crate) mod rules; mod slice; mod take; -mod take_slices; mod zip; #[cfg(test)] diff --git a/vortex-array/src/arrays/primitive/compute/take_slices.rs b/vortex-array/src/arrays/primitive/compute/take_slices.rs deleted file mode 100644 index 9e94594ac25..00000000000 --- a/vortex-array/src/arrays/primitive/compute/take_slices.rs +++ /dev/null @@ -1,42 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_buffer::BufferMut; -use vortex_error::VortexResult; - -use crate::ArrayRef; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Primitive; -use crate::arrays::PrimitiveArray; -use crate::arrays::take_slices::RunSelectors; -use crate::arrays::take_slices::TakeSlicesExecute; -use crate::executor::ExecutionCtx; -use crate::match_each_native_ptype; - -impl TakeSlicesExecute for Primitive { - fn take_slices( - array: ArrayView<'_, Self>, - starts: &ArrayRef, - lengths: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let selectors = RunSelectors::new(array.len(), starts, lengths, ctx)?; - let validity = array - .validity()? - .take_slices_with_ctx(starts, lengths, ctx)?; - match_each_native_ptype!(array.ptype(), |T| { - let source = array.as_slice::(); - let len = selectors.output_len(); - let mut values = BufferMut::::with_capacity(len); - - for &(start, end) in selectors.slices() { - values.extend_from_slice(&source[start..end]); - } - - Ok(Some( - PrimitiveArray::new(values.freeze(), validity).into_array(), - )) - }) - } -} diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 663c8924865..6382ea73794 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -6,9 +6,7 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::Primitive; -use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; @@ -34,10 +32,5 @@ pub(crate) fn initialize(session: &VortexSession) { FillNullExecuteAdaptor(Primitive), ); kernels.register_execute_parent_kernel(Dict.id(), Primitive, TakeExecuteAdaptor(Primitive)); - kernels.register_execute_parent_kernel( - TakeSlices.id(), - Primitive, - TakeSlicesExecuteAdaptor(Primitive), - ); kernels.register_execute_parent_kernel(Zip.id(), Primitive, ZipExecuteAdaptor(Primitive)); } diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 6e700cad329..01981803239 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -16,7 +16,6 @@ use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::arrays::slice::SliceReduceAdaptor; use crate::arrays::struct_::StructArrayExt; use crate::arrays::struct_::compute::cast::struct_cast_fields; -use crate::arrays::take_slices::TakeSlicesReduceAdaptor; use crate::builtins::ArrayBuiltins; use crate::matcher::Matcher; use crate::optimizer::rules::ArrayParentReduceRule; @@ -32,7 +31,6 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&MaskReduceAdaptor(Struct)), ParentRuleSet::lift(&SliceReduceAdaptor(Struct)), ParentRuleSet::lift(&TakeReduceAdaptor(Struct)), - ParentRuleSet::lift(&TakeSlicesReduceAdaptor(Struct)), ]); pub(crate) fn struct_cast_reduce_parent( diff --git a/vortex-array/src/arrays/struct_/compute/take.rs b/vortex-array/src/arrays/struct_/compute/take.rs index 95fad14491d..a63c6c85cf6 100644 --- a/vortex-array/src/arrays/struct_/compute/take.rs +++ b/vortex-array/src/arrays/struct_/compute/take.rs @@ -10,8 +10,6 @@ use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::dict::TakeReduce; use crate::arrays::struct_::StructArrayExt; -use crate::arrays::take_slices::TakeSlicesReduce; -use crate::arrays::take_slices::selector_output_len; use crate::builtins::ArrayBuiltins; use crate::scalar::Scalar; use crate::validity::Validity; @@ -52,26 +50,3 @@ impl TakeReduce for Struct { .map(Some) } } - -impl TakeSlicesReduce for Struct { - fn take_slices( - array: ArrayView<'_, Struct>, - starts: &ArrayRef, - lengths: &ArrayRef, - ) -> VortexResult> { - let len = selector_output_len(array.len(), starts, lengths)?; - let fields = array - .iter_unmasked_fields() - .map(|field| field.take_slices(starts.clone(), lengths.clone())) - .collect::>>()?; - - StructArray::try_new_with_dtype( - fields, - array.struct_fields().clone(), - len, - array.validity()?.take_slices(starts, lengths)?, - ) - .map(StructArray::into_array) - .map(Some) - } -} diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs index 2fcaec87737..3ed43dea777 100644 --- a/vortex-array/src/arrays/take_slices/array.rs +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -9,22 +9,19 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; -use crate::ExecutionCtx; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::arrays::TakeSlices; -use crate::arrays::take_slices::selector_output_len; -use crate::arrays::take_slices::selector_output_len_with_ctx; /// The child array selected by the run sequence. pub(super) const CHILD_SLOT: usize = 0; /// The selector naming the start of each child run. pub(super) const STARTS_SLOT: usize = 1; -/// The selector naming the length of each child run. -pub(super) const LENGTHS_SLOT: usize = 2; +/// The selector naming the end of each child run. +pub(super) const ENDS_SLOT: usize = 2; pub(super) const NUM_SLOTS: usize = 3; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "starts", "lengths"]; +pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "starts", "ends"]; /// Metadata for a [`TakeSlices`] array. #[derive(Clone, Debug)] @@ -54,11 +51,11 @@ pub trait TakeSlicesArrayExt: TypedArrayRef { .vortex_expect("validated take-slices starts slot") } - /// The selector naming each child run's length. - fn lengths(&self) -> &ArrayRef { - self.as_ref().slots()[LENGTHS_SLOT] + /// The selector naming each child run's end offset. + fn ends(&self) -> &ArrayRef { + self.as_ref().slots()[ENDS_SLOT] .as_ref() - .vortex_expect("validated take-slices lengths slot") + .vortex_expect("validated take-slices ends slot") } } impl> TakeSlicesArrayExt for T {} @@ -75,45 +72,29 @@ impl TakeSlicesData { /// Returns `true` if this array is empty. pub fn is_empty(&self) -> bool { - self.len == 0 + self.len() == 0 } } impl Array { - /// Constructs a new validated `TakeSlicesArray`. + /// Constructs a new `TakeSlicesArray` from selector arrays and caller-provided output length. /// - /// Prefer [`Self::try_new_with_ctx`] when the caller already has an execution context. - pub fn try_new(child: ArrayRef, starts: ArrayRef, lengths: ArrayRef) -> VortexResult { - let len = selector_output_len(child.len(), &starts, &lengths)?; - Ok(Self::new_validated(child, starts, lengths, len)) - } - - /// Constructs a new validated `TakeSlicesArray` using the caller's execution context to - /// validate selector arrays and compute the output length. - pub fn try_new_with_ctx( + /// Construction validates only the structural array invariants. Selector values are interpreted + /// when the lazy gather is executed. + pub fn try_new( child: ArrayRef, starts: ArrayRef, - lengths: ArrayRef, - ctx: &mut ExecutionCtx, + ends: ArrayRef, + len: usize, ) -> VortexResult { - let len = selector_output_len_with_ctx(child.len(), &starts, &lengths, ctx)?; - Ok(Self::new_validated(child, starts, lengths, len)) - } - - fn new_validated(child: ArrayRef, starts: ArrayRef, lengths: ArrayRef, len: usize) -> Self { let dtype = child.dtype().clone(); let data = TakeSlicesData::new(len); - - // SAFETY: `selector_output_len` validates selector dtypes, run bounds, and computes `len`; - // the outer dtype is copied from the child, and all required slots are populated. - unsafe { - Array::from_parts_unchecked( - ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![ - Some(child), - Some(starts), - Some(lengths) - ]), - ) - } + Array::try_from_parts( + ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![ + Some(child), + Some(starts), + Some(ends) + ]), + ) } } diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index 44818990a99..f5e1540dd9d 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -1,16 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Reduce and execute adaptors for contiguous-run take operations. +//! Lazy gather of contiguous child ranges. //! -//! `TakeSlicesArray` represents a gather of contiguous child runs: the output is the -//! concatenation of `values[starts[i]..starts[i] + lengths[i]]` for each selector row. Runs may -//! overlap, repeat, and appear in any order. Encodings that know how to serve those runs -//! efficiently implement -//! [`TakeSlicesReduce`] or [`TakeSlicesExecute`]. +//! `TakeSlicesArray` represents the concatenation of `values[starts[i]..ends[i]]` for each +//! selector row. Ranges may overlap, repeat, and appear in any order. mod array; -mod rules; mod vtable; pub use array::TakeSlicesArrayExt; @@ -22,194 +18,17 @@ use vortex_error::vortex_err; pub use vtable::*; use crate::ArrayRef; -use crate::Canonical; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::VortexSessionExecute; -use crate::array::ArrayView; -use crate::array::VTable; -use crate::arrays::Primitive; -use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::IntegerPType; -use crate::kernel::ExecuteParentKernel; -use crate::legacy_session; -use crate::match_each_unsigned_integer_ptype; -use crate::matcher::Matcher; -use crate::optimizer::rules::ArrayParentReduceRule; -/// Metadata-only implementation hook for taking a sequence of child runs. -pub trait TakeSlicesReduce: VTable { - /// Take a sequence of contiguous runs from an array without reading value buffers. - /// - /// Implementations should return `None` if serving the runs requires value buffer access. - /// - /// # Preconditions - /// - /// `starts` and `lengths` are guaranteed to be equal-length, non-nullable unsigned integer - /// arrays. Each run is in bounds for `array`. - fn take_slices( - array: ArrayView<'_, Self>, - starts: &ArrayRef, - lengths: &ArrayRef, - ) -> VortexResult>; -} - -/// Execution implementation hook for taking a sequence of child runs. -pub trait TakeSlicesExecute: VTable { - /// Take a sequence of contiguous runs from an array, potentially reading value buffers. - /// - /// # Preconditions - /// - /// `starts` and `lengths` are guaranteed to be equal-length, non-nullable unsigned integer - /// arrays. Each run is in bounds for `array`. - fn take_slices( - array: ArrayView<'_, Self>, - starts: &ArrayRef, - lengths: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult>; -} - -fn trivial_take_slices( - child: &ArrayRef, - starts: &ArrayRef, - lengths: &ArrayRef, -) -> VortexResult> { - let mut ctx = legacy_execution_ctx(); - trivial_take_slices_with_ctx(child, starts, lengths, &mut ctx) -} - -fn trivial_take_slices_with_ctx( - child: &ArrayRef, - starts: &ArrayRef, - lengths: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let selectors = RunSelectors::new(child.len(), starts, lengths, ctx)?; - trivial_take_slices_from_ranges(child, selectors.slices()) -} - -fn trivial_take_slices_from_ranges( - child: &ArrayRef, - slices: &[(usize, usize)], -) -> VortexResult> { - if slices.is_empty() { - return Ok(Some(Canonical::empty(child.dtype()).into_array())); - } - - if let [(start, end)] = slices { - if *start == 0 && *end == child.len() { - return Ok(Some(child.clone())); - } - return child.slice(*start..*end).map(Some); - } - - Ok(None) -} - -pub(super) fn selector_output_len( - child_len: usize, - starts: &ArrayRef, - lengths: &ArrayRef, -) -> VortexResult { - let mut ctx = legacy_execution_ctx(); - selector_output_len_with_ctx(child_len, starts, lengths, &mut ctx) -} - -pub(super) fn selector_output_len_with_ctx( - child_len: usize, - starts: &ArrayRef, - lengths: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult { - Ok(RunSelectors::new(child_len, starts, lengths, ctx)?.output_len()) -} - -#[allow(clippy::disallowed_methods)] -fn legacy_execution_ctx() -> ExecutionCtx { - legacy_session().create_execution_ctx() -} - -#[derive(Debug)] -pub(super) struct RunSelectors { - slices: Vec<(usize, usize)>, - output_len: usize, -} - -impl RunSelectors { - pub(super) fn new( - child_len: usize, - starts: &ArrayRef, - lengths: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - check_selector_arrays(starts, lengths)?; - let starts = starts.clone().execute::(ctx)?; - let lengths = lengths.clone().execute::(ctx)?; - primitive_run_selectors(child_len, starts.as_view(), lengths.as_view()) - } - - pub(super) fn slices(&self) -> &[(usize, usize)] { - &self.slices - } - - pub(super) fn output_len(&self) -> usize { - self.output_len - } -} - -fn run_selectors_from_slices( - child_len: usize, - starts: &[S], - lengths: &[L], -) -> VortexResult { - let mut slices = Vec::with_capacity(starts.len()); - let mut output_len = 0usize; - for (&start, &length) in starts.iter().zip(lengths) { - let start = selector_to_usize("start", start)?; - let length = selector_to_usize("length", length)?; - let end = start.checked_add(length).ok_or_else(|| { - vortex_err!("TakeSlicesArray run overflow for start {start} and length {length}") - })?; - vortex_ensure!( - end <= child_len, - "TakeSlicesArray run {start}..{end} exceeds child array length {child_len}" - ); - if length != 0 { - output_len = output_len - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; - slices.push((start, end)); - } - } - Ok(RunSelectors { slices, output_len }) -} - -fn primitive_run_selectors( - child_len: usize, - starts: ArrayView<'_, Primitive>, - lengths: ArrayView<'_, Primitive>, -) -> VortexResult { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - run_selectors_from_slices::( - child_len, - starts.as_slice::(), - lengths.as_slice::(), - ) - }) - }) -} - -pub(super) fn check_selector_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult<()> { +pub(super) fn check_selector_arrays(starts: &ArrayRef, ends: &ArrayRef) -> VortexResult<()> { check_selector_dtype("starts", starts)?; - check_selector_dtype("lengths", lengths)?; + check_selector_dtype("ends", ends)?; vortex_ensure!( - starts.len() == lengths.len(), - "TakeSlicesArray selectors must have equal length, got starts {} and lengths {}", + starts.len() == ends.len(), + "TakeSlicesArray selectors must have equal length, got starts {} and ends {}", starts.len(), - lengths.len() + ends.len() ); Ok(()) } @@ -230,67 +49,26 @@ fn check_selector_dtype(name: &str, selector: &ArrayRef) -> VortexResult<()> { } } -fn selector_to_usize(name: &str, value: T) -> VortexResult { +pub(super) fn selector_constant( + name: &str, + selector: &ArrayRef, +) -> VortexResult> { + selector + .as_constant() + .map(|scalar| { + scalar + .as_primitive() + .try_typed_value::()? + .ok_or_else(|| vortex_err!("TakeSlicesArray {name} constant selector is null")) + }) + .transpose() +} + +pub(super) fn selector_to_usize(name: &str, value: T) -> VortexResult { value .to_usize() .ok_or_else(|| vortex_err!("TakeSlicesArray {name} selector {value} does not fit in usize")) } -/// Adaptor that wraps a [`TakeSlicesReduce`] impl as an [`ArrayParentReduceRule`]. -#[derive(Default, Debug)] -pub struct TakeSlicesReduceAdaptor(pub V); - -impl ArrayParentReduceRule for TakeSlicesReduceAdaptor -where - V: TakeSlicesReduce, -{ - type Parent = TakeSlices; - - fn reduce_parent( - &self, - array: ArrayView<'_, V>, - parent: ::Match<'_>, - child_idx: usize, - ) -> VortexResult> { - if child_idx != 0 { - return Ok(None); - } - if let Some(result) = trivial_take_slices(array.array(), parent.starts(), parent.lengths())? - { - return Ok(Some(result)); - } - ::take_slices(array, parent.starts(), parent.lengths()) - } -} - -/// Adaptor that wraps a [`TakeSlicesExecute`] impl as an [`ExecuteParentKernel`]. -#[derive(Default, Debug)] -pub struct TakeSlicesExecuteAdaptor(pub V); - -impl ExecuteParentKernel for TakeSlicesExecuteAdaptor -where - V: TakeSlicesExecute, -{ - type Parent = TakeSlices; - - fn execute_parent( - &self, - array: ArrayView<'_, V>, - parent: ::Match<'_>, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - if child_idx != 0 { - return Ok(None); - } - if let Some(result) = - trivial_take_slices_with_ctx(array.array(), parent.starts(), parent.lengths(), ctx)? - { - return Ok(Some(result)); - } - ::take_slices(array, parent.starts(), parent.lengths(), ctx) - } -} - #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/take_slices/rules.rs b/vortex-array/src/arrays/take_slices/rules.rs deleted file mode 100644 index 9f77c4e0ca1..00000000000 --- a/vortex-array/src/arrays/take_slices/rules.rs +++ /dev/null @@ -1,106 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_error::VortexResult; -use vortex_error::vortex_err; - -use super::legacy_execution_ctx; -use super::trivial_take_slices; -use crate::ArrayRef; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::PrimitiveArray; -use crate::arrays::TakeSlices; -use crate::arrays::TakeSlicesArray; -use crate::arrays::take_slices::RunSelectors; -use crate::arrays::take_slices::TakeSlicesArrayExt; -use crate::arrays::take_slices::TakeSlicesReduce; -use crate::arrays::take_slices::TakeSlicesReduceAdaptor; -use crate::optimizer::rules::ArrayReduceRule; -use crate::optimizer::rules::ParentRuleSet; -use crate::optimizer::rules::ReduceRuleSet; - -pub(super) const PARENT_RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&TakeSlicesReduceAdaptor(TakeSlices))]); - -pub(super) const RULES: ReduceRuleSet = ReduceRuleSet::new(&[&TrivialTakeSlicesRule]); - -#[derive(Debug)] -struct TrivialTakeSlicesRule; - -impl ArrayReduceRule for TrivialTakeSlicesRule { - fn reduce(&self, array: ArrayView<'_, TakeSlices>) -> VortexResult> { - trivial_take_slices(array.child(), array.starts(), array.lengths()) - } -} - -impl TakeSlicesReduce for TakeSlices { - fn take_slices( - array: ArrayView<'_, Self>, - starts: &ArrayRef, - lengths: &ArrayRef, - ) -> VortexResult> { - let mut ctx = legacy_execution_ctx(); - let inner = RunSelectors::new( - array.child().len(), - array.starts(), - array.lengths(), - &mut ctx, - )?; - let outer = RunSelectors::new(array.len(), starts, lengths, &mut ctx)?; - let combined = project_slices(inner.slices(), outer.slices()); - let (starts, lengths) = selectors_from_slices(&combined)?; - Ok(Some( - TakeSlicesArray::try_new_with_ctx(array.child().clone(), starts, lengths, &mut ctx)? - .into_array(), - )) - } -} - -fn project_slices(inner: &[(usize, usize)], outer: &[(usize, usize)]) -> Vec<(usize, usize)> { - let mut projected = Vec::new(); - - for &(outer_start, outer_end) in outer { - let mut logical_start = 0usize; - for &(inner_start, inner_end) in inner { - let inner_len = inner_end - inner_start; - let logical_end = logical_start + inner_len; - - if outer_start < logical_end && outer_end > logical_start { - let overlap_start = outer_start.max(logical_start); - let overlap_end = outer_end.min(logical_end); - projected.push(( - inner_start + (overlap_start - logical_start), - inner_start + (overlap_end - logical_start), - )); - } - - if logical_end >= outer_end { - break; - } - logical_start = logical_end; - } - } - - projected -} - -fn selectors_from_slices(slices: &[(usize, usize)]) -> VortexResult<(ArrayRef, ArrayRef)> { - let starts = slices - .iter() - .map(|&(start, _)| selector_value(start, "start")) - .collect::>>()?; - let lengths = slices - .iter() - .map(|&(start, end)| selector_value(end - start, "length")) - .collect::>>()?; - Ok(( - PrimitiveArray::from_iter(starts).into_array(), - PrimitiveArray::from_iter(lengths).into_array(), - )) -} - -fn selector_value(value: usize, name: &str) -> VortexResult { - u64::try_from(value) - .map_err(|_| vortex_err!("TakeSlicesArray projected {name} {value} does not fit in u64")) -} diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs index 12e5376640b..00db4cee51a 100644 --- a/vortex-array/src/arrays/take_slices/tests.rs +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -1,9 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - -use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -11,20 +8,16 @@ use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::ConstantArray; use crate::arrays::DictArray; -use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::TakeSlices; use crate::arrays::TakeSlicesArray; use crate::arrays::VarBinViewArray; -use crate::arrays::take_slices::TakeSlicesArrayExt; -use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::assert_arrays_eq; use crate::dtype::FieldNames; use crate::dtype::Nullability; -use crate::kernel::ExecuteParentKernel; use crate::validity::Validity; #[test] @@ -77,19 +70,6 @@ fn take_slices_lazy_scalar_and_validity_follow_runs() -> VortexResult<()> { Ok(()) } -#[test] -fn take_slices_one_slice_reduces_to_child_slice() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let array = PrimitiveArray::from_iter(0i32..6).into_array(); - - let actual = take_slices(&array, &[(1, 4)])?; - - assert!(actual.is::()); - assert!(!actual.is::()); - assert_arrays_eq!(actual, PrimitiveArray::from_iter([1i32, 2, 3, 4]), &mut ctx); - Ok(()) -} - #[test] fn take_slices_size_one_child_can_repeat_the_only_range() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -117,7 +97,8 @@ fn take_slices_size_one_nullable_child_can_repeat_null() -> VortexResult<()> { #[test] fn take_slices_preserves_all_invalid_validity() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllInvalid).into_array(); + let array = + PrimitiveArray::new(vortex_buffer::buffer![1i32, 2, 3], Validity::AllInvalid).into_array(); let actual = take_slices(&array, &[(2, 1), (0, 2)])?; let expected = PrimitiveArray::from_option_iter([None::, None, None]); @@ -127,32 +108,64 @@ fn take_slices_preserves_all_invalid_validity() -> VortexResult<()> { } #[test] -fn primitive_take_slices_execute_parent_copies_runs_directly() -> VortexResult<()> { +fn take_slices_construction_defers_out_of_bounds_starts_to_execution() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let array = PrimitiveArray::from_iter(0i32..8).into_array(); - let (starts, lengths) = selector_arrays(&[(3, 3), (0, 2)])?; - let parent = TakeSlicesArray::try_new(array.clone(), starts, lengths)?.into_array(); - - let actual = TakeSlicesExecuteAdaptor(Primitive) - .execute_parent( - array.as_::(), - parent.as_::(), - 0, - &mut ctx, - )? - .ok_or_else(|| vortex_err!("Primitive TakeSlicesExecute declined multi-slice take"))?; - - assert!(actual.is::()); - assert_arrays_eq!( - actual, - PrimitiveArray::from_iter([3i32, 4, 5, 0, 1]), - &mut ctx - ); + let array = PrimitiveArray::from_iter(0i32..6).into_array(); + let starts = PrimitiveArray::from_iter([7u64]).into_array(); + let ends = PrimitiveArray::from_iter([7u64]).into_array(); + + let take_slices = TakeSlicesArray::try_new(array, starts, ends, 0)?.into_array(); + + assert!(take_slices.is::()); + assert_eq!(take_slices.len(), 0); + assert!(take_slices.execute::(&mut ctx).is_err()); Ok(()) } #[test] -fn struct_take_slices_pushes_runs_to_fields() -> VortexResult<()> { +fn take_slices_accepts_constant_end_selector() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0i32..6).into_array(); + let starts = PrimitiveArray::from_iter([0u64, 1, 2]).into_array(); + let ends = ConstantArray::new(3u64, 3).into_array(); + + let actual = TakeSlicesArray::try_new(array, starts, ends, 6)?.into_array(); + let expected = PrimitiveArray::from_iter([0i32, 1, 2, 1, 2, 2]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn take_slices_accepts_constant_start_selector() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0i32..6).into_array(); + let starts = ConstantArray::new(1u64, 3).into_array(); + let ends = PrimitiveArray::from_iter([2u64, 4, 6]).into_array(); + + let actual = TakeSlicesArray::try_new(array, starts, ends, 9)?.into_array(); + let expected = PrimitiveArray::from_iter([1i32, 1, 2, 3, 1, 2, 3, 4, 5]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn take_slices_accepts_constant_start_and_end_selectors() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0i32..6).into_array(); + let starts = ConstantArray::new(0u64, 3).into_array(); + let ends = ConstantArray::new(2u64, 3).into_array(); + + let actual = TakeSlicesArray::try_new(array, starts, ends, 6)?.into_array(); + let expected = PrimitiveArray::from_iter([0i32, 1, 0, 1, 0, 1]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn struct_take_slices_executes_generically() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = StructArray::try_new( FieldNames::from(["id", "name"]), @@ -176,35 +189,10 @@ fn struct_take_slices_pushes_runs_to_fields() -> VortexResult<()> { Validity::NonNullable, )?; - assert!(actual.is::()); - assert!(!actual.is::()); assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } -#[test] -fn varbinview_take_slices_execute_reuses_data_buffers() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let array = VarBinViewArray::from_iter_str([ - "long-string-value-000000", - "long-string-value-000001", - "long-string-value-000002", - "long-string-value-000003", - "long-string-value-000004", - "long-string-value-000005", - ]); - - let actual = take_slices(&array.clone().into_array(), &[(3, 2), (0, 2)])? - .execute::(&mut ctx)?; - - assert!(Arc::ptr_eq(actual.data_buffers(), array.data_buffers())); - assert_eq!(actual.bytes_at(0).as_slice(), b"long-string-value-000003"); - assert_eq!(actual.bytes_at(1).as_slice(), b"long-string-value-000004"); - assert_eq!(actual.bytes_at(2).as_slice(), b"long-string-value-000000"); - assert_eq!(actual.bytes_at(3).as_slice(), b"long-string-value-000001"); - Ok(()) -} - #[test] fn take_slices_empty_runs_return_empty_canonical() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -230,35 +218,36 @@ fn take_slices_empty_child_accepts_empty_runs() -> VortexResult<()> { } #[test] -fn take_slices_rejects_invalid_runs() -> VortexResult<()> { +fn take_slices_rejects_invalid_ranges_at_the_right_layer() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); - assert!(take_slices(&array, &[(7, 0)]).is_err()); - assert!(take_slices(&array, &[(4, 3)]).is_err()); - let starts = PrimitiveArray::from_iter([0u64, 1]).into_array(); - let lengths = PrimitiveArray::from_iter([1u64]).into_array(); - assert!(array.take_slices(starts, lengths).is_err()); + assert!(array.take_slices(vec![2], vec![1]).is_err()); + + let out_of_bounds_empty = array.take_slices(vec![7], vec![7])?; + assert!( + out_of_bounds_empty + .execute::(&mut ctx) + .is_err() + ); + + let out_of_bounds_non_empty = array.take_slices(vec![4], vec![7])?; + assert!( + out_of_bounds_non_empty + .execute::(&mut ctx) + .is_err() + ); Ok(()) } #[test] -fn take_slices_of_take_slices_projects_to_original_child() -> VortexResult<()> { +fn take_slices_of_take_slices_executes_correctly() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..10).into_array(); let inner = take_slices(&array, &[(2, 3), (7, 3)])?; let actual = take_slices(&inner, &[(1, 3), (4, 2)])?; - let actual_take_slices = actual.as_::(); - assert!(actual_take_slices.child().is::()); - assert_eq!( - selector_values(actual_take_slices.starts(), &mut ctx)?, - vec![3, 7, 8] - ); - assert_eq!( - selector_values(actual_take_slices.lengths(), &mut ctx)?, - vec![2, 1, 2] - ); assert_arrays_eq!( actual, PrimitiveArray::from_iter([3i32, 4, 7, 8, 9]), @@ -271,7 +260,7 @@ fn take_slices_of_take_slices_projects_to_original_child() -> VortexResult<()> { fn take_slices_generic_execution_handles_child_without_take_slices_kernel() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::try_new( - buffer![2u8, 0, 1, 2, 0, 1].into_array(), + vortex_buffer::buffer![2u8, 0, 1, 2, 0, 1].into_array(), PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), )? .into_array(); @@ -288,7 +277,7 @@ fn take_slices_generic_execution_handles_child_without_take_slices_kernel() -> V fn take_slices_generic_execution_preserves_nullable_encoded_child() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::try_new( - buffer![0u8, 1, 2, 0].into_array(), + vortex_buffer::buffer![0u8, 1, 2, 0].into_array(), PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array(), )? .into_array(); @@ -302,31 +291,14 @@ fn take_slices_generic_execution_preserves_nullable_encoded_child() -> VortexRes } fn take_slices(array: &ArrayRef, runs: &[(usize, usize)]) -> VortexResult { - let (starts, lengths) = selector_arrays(runs)?; - array.take_slices(starts, lengths) -} - -fn selector_arrays(runs: &[(usize, usize)]) -> VortexResult<(ArrayRef, ArrayRef)> { - let starts = runs - .iter() - .map(|&(start, _)| selector_value(start, "start")) - .collect::>>()?; - let lengths = runs + let starts = runs.iter().map(|&(start, _)| start).collect::>(); + let ends = runs .iter() - .map(|&(_, length)| selector_value(length, "length")) + .map(|&(start, length)| { + start.checked_add(length).ok_or_else(|| { + vortex_err!("test TakeSlices range overflow for start {start} and length {length}") + }) + }) .collect::>>()?; - Ok(( - PrimitiveArray::from_iter(starts).into_array(), - PrimitiveArray::from_iter(lengths).into_array(), - )) -} - -fn selector_value(value: usize, name: &str) -> VortexResult { - u64::try_from(value) - .map_err(|_| vortex_err!("test {name} selector {value} does not fit in u64")) -} - -fn selector_values(selector: &ArrayRef, ctx: &mut crate::ExecutionCtx) -> VortexResult> { - let selector = selector.clone().execute::(ctx)?; - Ok(selector.as_slice::().to_vec()) + array.take_slices(starts, ends) } diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index c5cf5331c39..f956f45b01f 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -8,6 +8,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -24,22 +25,25 @@ use crate::array::OperationsVTable; use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; -use crate::arrays::take_slices::RunSelectors; +use crate::arrays::PrimitiveArray; use crate::arrays::take_slices::TakeSlicesArrayExt; use crate::arrays::take_slices::array::CHILD_SLOT; -use crate::arrays::take_slices::array::LENGTHS_SLOT; +use crate::arrays::take_slices::array::ENDS_SLOT; use crate::arrays::take_slices::array::NUM_SLOTS; use crate::arrays::take_slices::array::SLOT_NAMES; use crate::arrays::take_slices::array::STARTS_SLOT; use crate::arrays::take_slices::array::TakeSlicesData; -use crate::arrays::take_slices::rules::PARENT_RULES; -use crate::arrays::take_slices::rules::RULES; -use crate::arrays::take_slices::selector_output_len; +use crate::arrays::take_slices::check_selector_arrays; +use crate::arrays::take_slices::selector_constant; +use crate::arrays::take_slices::selector_to_usize; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity_in; use crate::dtype::DType; +use crate::dtype::IntegerPType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; +use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -47,10 +51,11 @@ use crate::validity::Validity; /// A [`TakeSlices`]-encoded Vortex array. pub type TakeSlicesArray = Array; -/// Contiguous-run gather selection encoding. +/// Contiguous-range gather selection encoding. /// /// Like [`crate::arrays::Slice`], this is a lazy compute encoding and is not serialized as a file -/// encoding. +/// encoding. Execute it before a file-writing path when a materialized physical representation is +/// required. #[derive(Clone, Debug)] pub struct TakeSlices; @@ -97,8 +102,8 @@ impl VTable for TakeSlices { "TakeSlicesArray starts slot must be present" ); vortex_ensure!( - slots[LENGTHS_SLOT].is_some(), - "TakeSlicesArray lengths slot must be present" + slots[ENDS_SLOT].is_some(), + "TakeSlicesArray ends slot must be present" ); let child = slots[CHILD_SLOT] .as_ref() @@ -112,20 +117,14 @@ impl VTable for TakeSlices { let starts = slots[STARTS_SLOT] .as_ref() .vortex_expect("validated starts slot"); - let lengths = slots[LENGTHS_SLOT] + let ends = slots[ENDS_SLOT] .as_ref() - .vortex_expect("validated lengths slot"); - let computed_len = selector_output_len(child.len(), starts, lengths)?; + .vortex_expect("validated ends slot"); + check_selector_arrays(starts, ends)?; vortex_ensure!( - data.len() == computed_len, - "TakeSlicesArray metadata length {} does not match computed run length {}", + data.len() == len, + "TakeSlicesArray metadata length {} does not match outer length {}", data.len(), - computed_len - ); - vortex_ensure!( - computed_len == len, - "TakeSlicesArray computed length {} does not match outer length {}", - computed_len, len ); Ok(()) @@ -176,26 +175,9 @@ impl VTable for TakeSlices { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); - let selectors = - RunSelectors::new(array.child().len(), array.starts(), array.lengths(), ctx)?; - for &(start, end) in selectors.slices() { - let slice = array.child().slice(start..end)?; - slice.append_to_builder(builder.as_mut(), ctx)?; - } + append_selected_ranges(array.as_view(), builder.as_mut(), ctx)?; Ok(ExecutionResult::done(builder.finish())) } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) - } - - fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { - RULES.evaluate(array) - } } impl OperationsVTable for TakeSlices { @@ -204,21 +186,8 @@ impl OperationsVTable for TakeSlices { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let selectors = - RunSelectors::new(array.child().len(), array.starts(), array.lengths(), ctx)?; - let mut logical_start = 0usize; - for &(start, end) in selectors.slices() { - let len = end - start; - let logical_end = logical_start + len; - if index < logical_end { - return array - .child() - .execute_scalar(start + (index - logical_start), ctx); - } - logical_start = logical_end; - } - - vortex_panic!("TakeSlicesArray scalar index {index} out of bounds") + scalar_at_selected_range(array, index, ctx)? + .ok_or_else(|| vortex_err!("TakeSlicesArray scalar index {index} out of bounds")) } } @@ -227,6 +196,195 @@ impl ValidityVTable for TakeSlices { array .child() .validity()? - .take_slices(array.starts(), array.lengths()) + .take_slices(array.starts(), array.ends(), array.len()) + } +} + +fn append_selected_ranges( + array: ArrayView<'_, TakeSlices>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let starts = array.starts(); + let ends = array.ends(); + check_selector_arrays(starts, ends)?; + + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + match_each_unsigned_integer_ptype!(ends.dtype().as_ptype(), |E| { + let start = selector_constant::("start", starts)?; + let end = selector_constant::("end", ends)?; + match (start, end) { + (Some(start), Some(end)) => append_ranges::( + array.child(), + starts.len(), + |_| start, + |_| end, + builder, + ctx, + ), + (Some(start), None) => { + let ends = ends.clone().execute::(ctx)?; + let ends = ends.as_slice::(); + append_ranges::( + array.child(), + starts.len(), + |_| start, + |index| ends[index], + builder, + ctx, + ) + } + (None, Some(end)) => { + let starts = starts.clone().execute::(ctx)?; + let starts = starts.as_slice::(); + append_ranges::( + array.child(), + starts.len(), + |index| starts[index], + |_| end, + builder, + ctx, + ) + } + (None, None) => { + let starts = starts.clone().execute::(ctx)?; + let ends = ends.clone().execute::(ctx)?; + let starts = starts.as_slice::(); + let ends = ends.as_slice::(); + append_ranges::( + array.child(), + starts.len(), + |index| starts[index], + |index| ends[index], + builder, + ctx, + ) + } + } + }) + }) +} + +fn append_ranges( + child: &ArrayRef, + len: usize, + mut start_at: StartAt, + mut end_at: EndAt, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + S: IntegerPType, + E: IntegerPType, + StartAt: FnMut(usize) -> S, + EndAt: FnMut(usize) -> E, +{ + for index in 0..len { + let start = selector_to_usize("start", start_at(index))?; + let end = selector_to_usize("end", end_at(index))?; + child.slice(start..end)?.append_to_builder(builder, ctx)?; + } + Ok(()) +} + +fn scalar_at_selected_range( + array: ArrayView<'_, TakeSlices>, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let starts = array.starts(); + let ends = array.ends(); + check_selector_arrays(starts, ends)?; + + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + match_each_unsigned_integer_ptype!(ends.dtype().as_ptype(), |E| { + let start = selector_constant::("start", starts)?; + let end = selector_constant::("end", ends)?; + match (start, end) { + (Some(start), Some(end)) => scalar_from_ranges::( + array.child(), + starts.len(), + index, + |_| start, + |_| end, + ctx, + ), + (Some(start), None) => { + let ends = ends.clone().execute::(ctx)?; + let ends = ends.as_slice::(); + scalar_from_ranges::( + array.child(), + starts.len(), + index, + |_| start, + |range_index| ends[range_index], + ctx, + ) + } + (None, Some(end)) => { + let starts = starts.clone().execute::(ctx)?; + let starts = starts.as_slice::(); + scalar_from_ranges::( + array.child(), + starts.len(), + index, + |range_index| starts[range_index], + |_| end, + ctx, + ) + } + (None, None) => { + let starts = starts.clone().execute::(ctx)?; + let ends = ends.clone().execute::(ctx)?; + let starts = starts.as_slice::(); + let ends = ends.as_slice::(); + scalar_from_ranges::( + array.child(), + starts.len(), + index, + |range_index| starts[range_index], + |range_index| ends[range_index], + ctx, + ) + } + } + }) + }) +} + +fn scalar_from_ranges( + child: &ArrayRef, + len: usize, + index: usize, + mut start_at: StartAt, + mut end_at: EndAt, + ctx: &mut ExecutionCtx, +) -> VortexResult> +where + S: IntegerPType, + E: IntegerPType, + StartAt: FnMut(usize) -> S, + EndAt: FnMut(usize) -> E, +{ + let mut logical_start = 0usize; + for range_index in 0..len { + let start = selector_to_usize("start", start_at(range_index))?; + let end = selector_to_usize("end", end_at(range_index))?; + vortex_ensure!( + start <= end && end <= child.len(), + "TakeSlicesArray range {start}..{end} exceeds child array length {}", + child.len() + ); + let run_len = end - start; + let logical_end = logical_start + .checked_add(run_len) + .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; + if index < logical_end { + return child + .execute_scalar(start + (index - logical_start), ctx) + .map(Some); + } + logical_start = logical_end; } + Ok(None) } diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index a21e4d1883e..ce10abda3d0 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -6,7 +6,6 @@ use std::sync::Arc; use num_traits::AsPrimitive; use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -18,8 +17,6 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::take_slices::RunSelectors; -use crate::arrays::take_slices::TakeSlicesExecute; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; use crate::executor::ExecutionCtx; @@ -60,41 +57,6 @@ impl TakeExecute for VarBinView { } } -impl TakeSlicesExecute for VarBinView { - fn take_slices( - array: ArrayView<'_, VarBinView>, - starts: &ArrayRef, - lengths: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let selectors = RunSelectors::new(array.len(), starts, lengths, ctx)?; - let validity = array - .validity()? - .take_slices_with_ctx(starts, lengths, ctx)?; - let len = selectors.output_len(); - - let source_views = array.views(); - let mut views = BufferMut::::with_capacity(len); - for &(start, end) in selectors.slices() { - views.extend_from_slice(&source_views[start..end]); - } - - // SAFETY: views are copied from a valid VarBinView array and still reference the same - // backing data buffers, so all view pointers remain valid. - unsafe { - Ok(Some( - VarBinViewArray::new_handle_unchecked( - BufferHandle::new_host(views.freeze().into_byte_buffer()), - Arc::clone(array.data_buffers()), - array.dtype().clone(), - validity, - ) - .into_array(), - )) - } - } -} - fn take_views>( views_ref: &[BinaryView], indices: &[I], diff --git a/vortex-array/src/arrays/varbinview/vtable/kernel.rs b/vortex-array/src/arrays/varbinview/vtable/kernel.rs index c924f196104..e09b381d590 100644 --- a/vortex-array/src/arrays/varbinview/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbinview/vtable/kernel.rs @@ -5,10 +5,8 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Dict; -use crate::arrays::TakeSlices; use crate::arrays::VarBinView; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; @@ -20,10 +18,5 @@ pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); kernels.register_execute_parent_kernel(Cast.id(), VarBinView, CastExecuteAdaptor(VarBinView)); kernels.register_execute_parent_kernel(Dict.id(), VarBinView, TakeExecuteAdaptor(VarBinView)); - kernels.register_execute_parent_kernel( - TakeSlices.id(), - VarBinView, - TakeSlicesExecuteAdaptor(VarBinView), - ); kernels.register_execute_parent_kernel(Zip.id(), VarBinView, ZipExecuteAdaptor(VarBinView)); } diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 2e02e59a9ab..ebbe91ad2b6 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -224,30 +224,39 @@ impl Validity { } } - /// Select validity values by concatenating a sequence of contiguous runs. - /// - /// Prefer [`Self::take_slices_with_ctx`] when the caller already has an execution context. + /// Select validity values by concatenating a sequence of contiguous ranges. #[allow(clippy::disallowed_methods)] - pub fn take_slices(&self, starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult { + pub fn take_slices( + &self, + starts: &ArrayRef, + ends: &ArrayRef, + len: usize, + ) -> VortexResult { let mut ctx = legacy_session().create_execution_ctx(); - self.take_slices_with_ctx(starts, lengths, &mut ctx) + self.take_slices_with_ctx(starts, ends, len, &mut ctx) } - /// Select validity values by concatenating a sequence of contiguous runs, using the caller's + /// Select validity values by concatenating a sequence of contiguous ranges, using the caller's /// execution context for array-backed validity. pub fn take_slices_with_ctx( &self, starts: &ArrayRef, - lengths: &ArrayRef, + ends: &ArrayRef, + len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { match self { v @ (Self::NonNullable | Self::AllValid | Self::AllInvalid) => Ok(v.clone()), - Self::Array(is_valid) => Ok(Self::Array(is_valid.take_slices_with_ctx( - starts.clone(), - lengths.clone(), - ctx, - )?)), + Self::Array(is_valid) => Ok(Self::Array( + crate::arrays::TakeSlicesArray::try_new( + is_valid.clone(), + starts.clone(), + ends.clone(), + len, + )? + .into_array() + .optimize_ctx(ctx.session())?, + )), } } diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index f228e1571d3..dcf445eb99e 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -54,6 +54,7 @@ regex = { workspace = true } reqwest = { workspace = true, features = ["stream"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +sha2 = { workspace = true } spatialbench = { workspace = true } spatialbench-arrow = { workspace = true } spatialbench-parquet = { workspace = true } diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 57fc09f4335..886b0f48249 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -5,6 +5,7 @@ use std::fs; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::time::UNIX_EPOCH; use futures::StreamExt; use futures::TryStreamExt; @@ -13,10 +14,15 @@ use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::async_reader::ParquetRecordBatchStream; use parquet::file::metadata::KeyValue; use parquet::file::metadata::ParquetMetaData; +use serde::Deserialize; +use serde::Serialize; +use sha2::Digest; +use sha2::Sha256; use sysinfo::System; use tokio::fs::File; use tokio::fs::OpenOptions; use tokio::fs::create_dir_all; +use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; use tracing::Instrument; use tracing::info; @@ -76,6 +82,7 @@ const MIN_CONCURRENCY: u64 = 1; /// Maximum number of concurrent conversion streams. This is somewhat arbitary. const MAX_CONCURRENCY: u64 = 16; +const CONVERSION_CACHE_VERSION: u32 = 1; /// Returns the available system memory in bytes. fn available_memory_bytes() -> u64 { @@ -169,40 +176,87 @@ pub async fn convert_parquet_file_to_vortex( Ok(()) } -async fn parquet_file_row_count(parquet_path: &Path) -> anyhow::Result { +struct ParquetFileSummary { + row_count: u64, + dtype: DType, + source: SourceFileFingerprint, +} + +struct VortexFileSummary { + row_count: u64, + dtype: DType, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct SourceFileFingerprint { + len: u64, + modified_secs: u64, + modified_nanos: u32, + sha256: String, +} + +#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)] +struct ConversionManifest { + version: u32, + compaction: String, + parquet: SourceFileFingerprint, + row_count: u64, + dtype: String, +} + +async fn parquet_file_summary(parquet_path: &Path) -> anyhow::Result { let file = File::open(parquet_path).await?; let builder = ParquetRecordBatchStreamBuilder::new(file).await?; + let geo_columns = geoparquet_columns(builder.metadata()); + let dtype = tag_geo_dtype(DType::from_arrow(builder.schema().as_ref()), &geo_columns)?; let row_count = builder.metadata().file_metadata().num_rows(); - u64::try_from(row_count) - .map_err(|_| anyhow::anyhow!("Parquet file has a negative row count: {parquet_path:?}")) + let row_count = u64::try_from(row_count) + .map_err(|_| anyhow::anyhow!("Parquet file has a negative row count: {parquet_path:?}"))?; + Ok(ParquetFileSummary { + row_count, + dtype, + source: source_file_fingerprint(parquet_path).await?, + }) } -async fn vortex_file_row_count(vortex_path: &Path) -> anyhow::Result { - Ok(SESSION - .open_options() - .open_path(vortex_path) - .await? - .row_count()) +async fn vortex_file_summary(vortex_path: &Path) -> anyhow::Result { + let file = SESSION.open_options().open_path(vortex_path).await?; + Ok(VortexFileSummary { + row_count: file.row_count(), + dtype: file.dtype().clone(), + }) } -async fn existing_vortex_file_has_matching_row_count( +async fn existing_vortex_file_matches_parquet( parquet_path: &Path, vortex_path: &Path, + compaction: CompactionStrategy, ) -> anyhow::Result { if !vortex_path.exists() { return Ok(false); } - let expected_row_count = parquet_file_row_count(parquet_path).await?; - match vortex_file_row_count(vortex_path).await { - Ok(actual_row_count) if actual_row_count == expected_row_count => Ok(true), - Ok(actual_row_count) => { + let expected = parquet_file_summary(parquet_path).await?; + let expected_manifest = ConversionManifest::new(&expected, compaction); + let manifest = read_conversion_manifest(vortex_path).await; + match vortex_file_summary(vortex_path).await { + Ok(actual) + if actual.row_count == expected.row_count + && actual.dtype == expected.dtype + && manifest.as_ref().is_ok_and(|m| m == &expected_manifest) => + { + Ok(true) + } + Ok(actual) => { info!( parquet_path = %parquet_path.display(), vortex_path = %vortex_path.display(), - expected_row_count, - actual_row_count, - "Regenerating Vortex file with stale row count" + expected_row_count = expected.row_count, + actual_row_count = actual.row_count, + expected_dtype = %expected.dtype, + actual_dtype = %actual.dtype, + manifest_error = manifest.as_ref().err().map(|e| e.to_string()), + "Regenerating stale Vortex file" ); Ok(false) } @@ -218,12 +272,84 @@ async fn existing_vortex_file_has_matching_row_count( } } +impl ConversionManifest { + fn new(summary: &ParquetFileSummary, compaction: CompactionStrategy) -> Self { + Self { + version: CONVERSION_CACHE_VERSION, + compaction: format!("{compaction:?}"), + parquet: summary.source.clone(), + row_count: summary.row_count, + dtype: summary.dtype.to_string(), + } + } +} + +async fn source_file_fingerprint(path: &Path) -> anyhow::Result { + let metadata = tokio::fs::metadata(path).await?; + let modified = metadata.modified()?.duration_since(UNIX_EPOCH)?; + let mut file = File::open(path).await?; + let mut hasher = Sha256::new(); + let mut buffer = vec![0u8; 1024 * 1024]; + loop { + let read = file.read(&mut buffer).await?; + if read == 0 { + break; + } + hasher.update(&buffer[..read]); + } + + Ok(SourceFileFingerprint { + len: metadata.len(), + modified_secs: modified.as_secs(), + modified_nanos: modified.subsec_nanos(), + sha256: hex_encode(hasher.finalize()), + }) +} + +fn hex_encode(bytes: impl AsRef<[u8]>) -> String { + const HEX: &[u8; 16] = b"0123456789abcdef"; + let bytes = bytes.as_ref(); + let mut encoded = String::with_capacity(bytes.len() * 2); + for &byte in bytes { + encoded.push(HEX[(byte >> 4) as usize] as char); + encoded.push(HEX[(byte & 0x0f) as usize] as char); + } + encoded +} + +fn conversion_manifest_path(vortex_path: &Path) -> PathBuf { + let mut manifest_path = vortex_path.to_path_buf(); + let file_name = vortex_path + .file_name() + .map(|name| name.to_string_lossy()) + .unwrap_or_else(|| "vortex".into()); + manifest_path.set_file_name(format!("{file_name}.manifest.json")); + manifest_path +} + +async fn read_conversion_manifest(vortex_path: &Path) -> anyhow::Result { + let bytes = tokio::fs::read(conversion_manifest_path(vortex_path)).await?; + Ok(serde_json::from_slice(&bytes)?) +} + +async fn write_conversion_manifest( + parquet_path: &Path, + vortex_path: &Path, + compaction: CompactionStrategy, +) -> anyhow::Result<()> { + let summary = parquet_file_summary(parquet_path).await?; + let manifest = ConversionManifest::new(&summary, compaction); + let bytes = serde_json::to_vec_pretty(&manifest)?; + tokio::fs::write(conversion_manifest_path(vortex_path), bytes).await?; + Ok(()) +} + async fn convert_parquet_file_to_current_vortex( parquet_path: &Path, output_path: &Path, compaction: CompactionStrategy, ) -> anyhow::Result { - if existing_vortex_file_has_matching_row_count(parquet_path, output_path).await? { + if existing_vortex_file_matches_parquet(parquet_path, output_path, compaction).await? { return Ok(output_path.to_path_buf()); } @@ -243,6 +369,7 @@ async fn convert_parquet_file_to_current_vortex( tokio::fs::remove_file(output_path).await?; } tokio::fs::rename(&temp_path, output_path).await?; + write_conversion_manifest(parquet_path, output_path, compaction).await?; Ok(output_path.to_path_buf()) } @@ -301,8 +428,8 @@ fn no_dict_layout() -> Arc { /// `{input_path}/vortex-file-compressed/` (for Default compaction) or /// `{input_path}/vortex-compact/` (for Compact compaction). /// -/// Existing Vortex files are reused when their footer row count matches the source Parquet file. -/// Unreadable or stale Vortex files are regenerated. +/// Existing Vortex files are reused when their footer metadata and conversion manifest match the +/// source Parquet file. Unreadable or stale Vortex files are regenerated. pub async fn convert_parquet_directory_to_vortex( input_path: &Path, compaction: CompactionStrategy, From efdd335f5220bee55a29d65d8eb313a85cf7dfc7 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 12:09:47 -0400 Subject: [PATCH 11/37] Use lengths for TakeSlices selectors Signed-off-by: Daniel King --- vortex-array/benches/take_fsl.rs | 23 ++- vortex-array/src/array/erased.rs | 26 ++-- .../arrays/fixed_size_list/compute/take.rs | 56 +++++--- vortex-array/src/arrays/listview/rebuild.rs | 14 +- vortex-array/src/arrays/take_slices/array.rs | 43 ++++-- vortex-array/src/arrays/take_slices/mod.rs | 15 +- vortex-array/src/arrays/take_slices/tests.rs | 53 +++---- vortex-array/src/arrays/take_slices/vtable.rs | 134 ++++++++++-------- vortex-array/src/validity.rs | 8 +- 9 files changed, 225 insertions(+), 147 deletions(-) diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index fc2f3b373dd..de236d26716 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -24,8 +24,10 @@ use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::TakeSlicesArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::dtype::IntegerPType; use vortex_array::dtype::NativePType; @@ -214,11 +216,24 @@ fn take_fsl_f16_take_slices_strategy( .iter() .map(|&idx| idx as usize * LIST_SIZE) .collect::>(); - let ends = starts - .iter() - .map(|&start| start + LIST_SIZE) + let run_count = starts.len(); + let starts = starts + .into_iter() + .map(|start| start as u64) .collect::>(); - let elements = array.elements().take_slices(starts, ends).unwrap(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(LIST_SIZE as u64, run_count).into_array(); + // SAFETY: benchmark indices are generated in-bounds, lengths is a non-nullable unsigned + // constant selector, and output length is exactly `indices.len() * LIST_SIZE`. + let elements = unsafe { + TakeSlicesArray::new_unchecked( + array.elements().clone(), + starts, + lengths, + indices.len() * LIST_SIZE, + ) + } + .into_array(); // SAFETY: each generated run has width `LIST_SIZE`, and there is one run per input index, // so `elements.len() == indices.len() * LIST_SIZE`. diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 7f7d8ce6cad..51b4120f76e 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -268,38 +268,34 @@ impl ArrayRef { /// Wraps the array in a [`TakeSlicesArray`] selected by caller-provided child ranges. /// - /// The output is the concatenation of `self[starts[i]..ends[i]]` for each selector row. This - /// computes the output length from `starts` and `ends`, but child bounds are checked only when - /// the lazy gather executes. - pub fn take_slices(&self, starts: Vec, ends: Vec) -> VortexResult { + /// The output is the concatenation of `self[starts[i]..starts[i] + lengths[i]]` for each + /// selector row. This computes the output length from `lengths`, but child bounds are checked + /// only when the lazy gather executes. + pub fn take_slices(&self, starts: Vec, lengths: Vec) -> VortexResult { vortex_ensure!( - starts.len() == ends.len(), - "TakeSlicesArray selectors must have equal length, got starts {} and ends {}", + starts.len() == lengths.len(), + "TakeSlicesArray selectors must have equal length, got starts {} and lengths {}", starts.len(), - ends.len() + lengths.len() ); let mut len = 0usize; - for (&start, &end) in starts.iter().zip(&ends) { - vortex_ensure!( - start <= end, - "TakeSlicesArray start {start} must be <= end {end}" - ); + for &length in &lengths { len = len - .checked_add(end - start) + .checked_add(length) .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; } let starts = starts .into_iter() .map(selector_value) .collect::>>()?; - let ends = ends + let lengths = lengths .into_iter() .map(selector_value) .collect::>>()?; TakeSlicesArray::try_new( self.clone(), PrimitiveArray::from_iter(starts).into_array(), - PrimitiveArray::from_iter(ends).into_array(), + PrimitiveArray::from_iter(lengths).into_array(), len, )? .into_array() diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 352252a32ce..aae647de54b 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -13,10 +13,12 @@ use crate::IntoArray; use crate::RecursiveCanonical; use crate::array::ArrayView; use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; +use crate::arrays::TakeSlicesArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; @@ -159,7 +161,13 @@ fn take_non_nullable_non_empty_fsl( starts.push(start); } - let new_elements = take_element_runs(array.elements(), starts, list_size, ctx)?; + let new_elements = take_element_runs( + array.elements(), + starts, + list_size, + expected_elements_len, + ctx, + )?; ensure_elements_len(new_elements.len(), expected_elements_len)?; // SAFETY: `starts` contains one checked run of `list_size` elements for each output row, @@ -213,7 +221,13 @@ fn take_nullable_non_empty_fsl( new_validity_builder.append(true); } - let new_elements = take_element_runs(array.elements(), starts, list_size, ctx)?; + let new_elements = take_element_runs( + array.elements(), + starts, + list_size, + expected_elements_len, + ctx, + )?; ensure_elements_len(new_elements.len(), expected_elements_len)?; let new_validity = Validity::from(new_validity_builder.freeze()); @@ -304,24 +318,27 @@ fn take_element_runs( elements: &ArrayRef, starts: Vec, length: usize, + output_len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let ends = starts - .iter() - .map(|&start| { - start.checked_add(length).ok_or_else(|| { - vortex_err!( - "FixedSizeList take element range overflow for start {start} and length {length}" - ) - }) - }) + let run_count = starts.len(); + let starts = starts + .into_iter() + .map(selector_value) .collect::>>()?; - - Ok(elements - .take_slices(starts, ends)? - .execute::(ctx)? - .0 - .into_array()) + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(selector_value(length)?, run_count).into_array(); + + // SAFETY: callers produced one start per output row after validating list indices against the + // source FSL length. `length` is the fixed list size, represented as a non-nullable unsigned + // constant selector, and `output_len` was computed as `run_count * length`. + Ok( + unsafe { TakeSlicesArray::new_unchecked(elements.clone(), starts, lengths, output_len) } + .into_array() + .execute::(ctx)? + .0 + .into_array(), + ) } fn null_element_run_start() -> usize { @@ -335,3 +352,8 @@ fn index_to_usize(index: I) -> VortexResult { .to_usize() .ok_or_else(|| vortex_err!("FixedSizeList take index {index} does not fit in usize")) } + +fn selector_value(value: usize) -> VortexResult { + u64::try_from(value) + .map_err(|_| vortex_err!("FixedSizeList take selector {value} does not fit in u64")) +} diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 69ca1792b6d..a7d36a80dad 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -187,7 +187,7 @@ impl ListViewArray { let new_sizes = ranges.new_sizes.freeze(); let elements = self .elements() - .take_slices(ranges.starts, ranges.ends)? + .take_slices(ranges.starts, ranges.lengths)? .execute::(ctx)? .0 .into_array(); @@ -272,7 +272,7 @@ struct RebuildRanges { new_offsets: BufferMut, new_sizes: BufferMut, starts: Vec, - ends: Vec, + lengths: Vec, } fn rebuild_ranges( @@ -289,7 +289,7 @@ where let mut new_offsets = BufferMut::::with_capacity(len); let mut new_sizes = BufferMut::::with_capacity(len); let mut starts = Vec::with_capacity(len); - let mut ends = Vec::with_capacity(len); + let mut lengths = Vec::with_capacity(len); let mut n_elements = NewOffset::zero(); for (index, is_valid) in validity_mask.iter().enumerate() { @@ -297,7 +297,7 @@ where new_offsets.push(n_elements); new_sizes.push(S::zero()); starts.push(0); - ends.push(0); + lengths.push(0); continue; } @@ -311,7 +311,7 @@ where let length = size .to_usize() .ok_or_else(|| vortex_err!("ListView rebuild size {size} does not fit in usize"))?; - let end = start.checked_add(length).ok_or_else(|| { + start.checked_add(length).ok_or_else(|| { vortex_err!( "ListView rebuild element range overflow for start {start} and length {length}" ) @@ -320,7 +320,7 @@ where new_offsets.push(n_elements); new_sizes.push(size); starts.push(start); - ends.push(end); + lengths.push(length); n_elements += num_traits::cast(size).vortex_expect("Cast failed"); } @@ -328,7 +328,7 @@ where new_offsets, new_sizes, starts, - ends, + lengths, }) } diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs index 3ed43dea777..f912e9d1792 100644 --- a/vortex-array/src/arrays/take_slices/array.rs +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -18,10 +18,10 @@ use crate::arrays::TakeSlices; pub(super) const CHILD_SLOT: usize = 0; /// The selector naming the start of each child run. pub(super) const STARTS_SLOT: usize = 1; -/// The selector naming the end of each child run. -pub(super) const ENDS_SLOT: usize = 2; +/// The selector naming the length of each child run. +pub(super) const LENGTHS_SLOT: usize = 2; pub(super) const NUM_SLOTS: usize = 3; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "starts", "ends"]; +pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "starts", "lengths"]; /// Metadata for a [`TakeSlices`] array. #[derive(Clone, Debug)] @@ -51,11 +51,11 @@ pub trait TakeSlicesArrayExt: TypedArrayRef { .vortex_expect("validated take-slices starts slot") } - /// The selector naming each child run's end offset. - fn ends(&self) -> &ArrayRef { - self.as_ref().slots()[ENDS_SLOT] + /// The selector naming each child run's length. + fn lengths(&self) -> &ArrayRef { + self.as_ref().slots()[LENGTHS_SLOT] .as_ref() - .vortex_expect("validated take-slices ends slot") + .vortex_expect("validated take-slices lengths slot") } } impl> TakeSlicesArrayExt for T {} @@ -84,7 +84,7 @@ impl Array { pub fn try_new( child: ArrayRef, starts: ArrayRef, - ends: ArrayRef, + lengths: ArrayRef, len: usize, ) -> VortexResult { let dtype = child.dtype().clone(); @@ -93,8 +93,33 @@ impl Array { ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![ Some(child), Some(starts), - Some(ends) + Some(lengths) ]), ) } + + /// Constructs a new `TakeSlicesArray` without validation. + /// + /// # Safety + /// + /// The caller must ensure the child dtype is the output dtype, selector arrays are + /// non-nullable unsigned integers of equal length, and `len` is the sum of selected lengths. + pub unsafe fn new_unchecked( + child: ArrayRef, + starts: ArrayRef, + lengths: ArrayRef, + len: usize, + ) -> Self { + let dtype = child.dtype().clone(); + let data = TakeSlicesData::new(len); + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![ + Some(child), + Some(starts), + Some(lengths) + ]), + ) + } + } } diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index f5e1540dd9d..64326dc792b 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -3,8 +3,9 @@ //! Lazy gather of contiguous child ranges. //! -//! `TakeSlicesArray` represents the concatenation of `values[starts[i]..ends[i]]` for each -//! selector row. Ranges may overlap, repeat, and appear in any order. +//! `TakeSlicesArray` represents the concatenation of +//! `values[starts[i]..starts[i] + lengths[i]]` for each selector row. Ranges may overlap, repeat, +//! and appear in any order. mod array; mod vtable; @@ -21,14 +22,14 @@ use crate::ArrayRef; use crate::dtype::DType; use crate::dtype::IntegerPType; -pub(super) fn check_selector_arrays(starts: &ArrayRef, ends: &ArrayRef) -> VortexResult<()> { +pub(super) fn check_selector_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult<()> { check_selector_dtype("starts", starts)?; - check_selector_dtype("ends", ends)?; + check_selector_dtype("lengths", lengths)?; vortex_ensure!( - starts.len() == ends.len(), - "TakeSlicesArray selectors must have equal length, got starts {} and ends {}", + starts.len() == lengths.len(), + "TakeSlicesArray selectors must have equal length, got starts {} and lengths {}", starts.len(), - ends.len() + lengths.len() ); Ok(()) } diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs index 00db4cee51a..be618796598 100644 --- a/vortex-array/src/arrays/take_slices/tests.rs +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use vortex_error::vortex_err; use crate::ArrayRef; use crate::IntoArray; @@ -112,9 +111,9 @@ fn take_slices_construction_defers_out_of_bounds_starts_to_execution() -> Vortex let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); let starts = PrimitiveArray::from_iter([7u64]).into_array(); - let ends = PrimitiveArray::from_iter([7u64]).into_array(); + let lengths = PrimitiveArray::from_iter([0u64]).into_array(); - let take_slices = TakeSlicesArray::try_new(array, starts, ends, 0)?.into_array(); + let take_slices = TakeSlicesArray::try_new(array, starts, lengths, 0)?.into_array(); assert!(take_slices.is::()); assert_eq!(take_slices.len(), 0); @@ -123,14 +122,14 @@ fn take_slices_construction_defers_out_of_bounds_starts_to_execution() -> Vortex } #[test] -fn take_slices_accepts_constant_end_selector() -> VortexResult<()> { +fn take_slices_accepts_constant_length_selector() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); let starts = PrimitiveArray::from_iter([0u64, 1, 2]).into_array(); - let ends = ConstantArray::new(3u64, 3).into_array(); + let lengths = ConstantArray::new(3u64, 3).into_array(); - let actual = TakeSlicesArray::try_new(array, starts, ends, 6)?.into_array(); - let expected = PrimitiveArray::from_iter([0i32, 1, 2, 1, 2, 2]); + let actual = TakeSlicesArray::try_new(array, starts, lengths, 9)?.into_array(); + let expected = PrimitiveArray::from_iter([0i32, 1, 2, 1, 2, 3, 2, 3, 4]); assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) @@ -141,9 +140,9 @@ fn take_slices_accepts_constant_start_selector() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); let starts = ConstantArray::new(1u64, 3).into_array(); - let ends = PrimitiveArray::from_iter([2u64, 4, 6]).into_array(); + let lengths = PrimitiveArray::from_iter([1u64, 3, 5]).into_array(); - let actual = TakeSlicesArray::try_new(array, starts, ends, 9)?.into_array(); + let actual = TakeSlicesArray::try_new(array, starts, lengths, 9)?.into_array(); let expected = PrimitiveArray::from_iter([1i32, 1, 2, 3, 1, 2, 3, 4, 5]); assert_arrays_eq!(actual, expected, &mut ctx); @@ -151,13 +150,13 @@ fn take_slices_accepts_constant_start_selector() -> VortexResult<()> { } #[test] -fn take_slices_accepts_constant_start_and_end_selectors() -> VortexResult<()> { +fn take_slices_accepts_constant_start_and_length_selectors() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); let starts = ConstantArray::new(0u64, 3).into_array(); - let ends = ConstantArray::new(2u64, 3).into_array(); + let lengths = ConstantArray::new(2u64, 3).into_array(); - let actual = TakeSlicesArray::try_new(array, starts, ends, 6)?.into_array(); + let actual = TakeSlicesArray::try_new(array, starts, lengths, 6)?.into_array(); let expected = PrimitiveArray::from_iter([0i32, 1, 0, 1, 0, 1]); assert_arrays_eq!(actual, expected, &mut ctx); @@ -222,16 +221,16 @@ fn take_slices_rejects_invalid_ranges_at_the_right_layer() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); - assert!(array.take_slices(vec![2], vec![1]).is_err()); + assert!(array.take_slices(vec![0, 0], vec![usize::MAX, 1]).is_err()); - let out_of_bounds_empty = array.take_slices(vec![7], vec![7])?; + let out_of_bounds_empty = array.take_slices(vec![7], vec![0])?; assert!( out_of_bounds_empty .execute::(&mut ctx) .is_err() ); - let out_of_bounds_non_empty = array.take_slices(vec![4], vec![7])?; + let out_of_bounds_non_empty = array.take_slices(vec![4], vec![3])?; assert!( out_of_bounds_non_empty .execute::(&mut ctx) @@ -240,6 +239,19 @@ fn take_slices_rejects_invalid_ranges_at_the_right_layer() -> VortexResult<()> { Ok(()) } +#[test] +fn take_slices_execution_rejects_declared_len_mismatch() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0i32..6).into_array(); + let starts = PrimitiveArray::from_iter([0u64]).into_array(); + let lengths = PrimitiveArray::from_iter([1u64]).into_array(); + + let take_slices = TakeSlicesArray::try_new(array, starts, lengths, 0)?.into_array(); + + assert!(take_slices.execute::(&mut ctx).is_err()); + Ok(()) +} + #[test] fn take_slices_of_take_slices_executes_correctly() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -292,13 +304,6 @@ fn take_slices_generic_execution_preserves_nullable_encoded_child() -> VortexRes fn take_slices(array: &ArrayRef, runs: &[(usize, usize)]) -> VortexResult { let starts = runs.iter().map(|&(start, _)| start).collect::>(); - let ends = runs - .iter() - .map(|&(start, length)| { - start.checked_add(length).ok_or_else(|| { - vortex_err!("test TakeSlices range overflow for start {start} and length {length}") - }) - }) - .collect::>>()?; - array.take_slices(starts, ends) + let lengths = runs.iter().map(|&(_, length)| length).collect::>(); + array.take_slices(starts, lengths) } diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index f956f45b01f..ccb6ac59ebc 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -28,7 +28,7 @@ use crate::array::with_empty_buffers; use crate::arrays::PrimitiveArray; use crate::arrays::take_slices::TakeSlicesArrayExt; use crate::arrays::take_slices::array::CHILD_SLOT; -use crate::arrays::take_slices::array::ENDS_SLOT; +use crate::arrays::take_slices::array::LENGTHS_SLOT; use crate::arrays::take_slices::array::NUM_SLOTS; use crate::arrays::take_slices::array::SLOT_NAMES; use crate::arrays::take_slices::array::STARTS_SLOT; @@ -102,8 +102,8 @@ impl VTable for TakeSlices { "TakeSlicesArray starts slot must be present" ); vortex_ensure!( - slots[ENDS_SLOT].is_some(), - "TakeSlicesArray ends slot must be present" + slots[LENGTHS_SLOT].is_some(), + "TakeSlicesArray lengths slot must be present" ); let child = slots[CHILD_SLOT] .as_ref() @@ -117,10 +117,10 @@ impl VTable for TakeSlices { let starts = slots[STARTS_SLOT] .as_ref() .vortex_expect("validated starts slot"); - let ends = slots[ENDS_SLOT] + let lengths = slots[LENGTHS_SLOT] .as_ref() - .vortex_expect("validated ends slot"); - check_selector_arrays(starts, ends)?; + .vortex_expect("validated lengths slot"); + check_selector_arrays(starts, lengths)?; vortex_ensure!( data.len() == len, "TakeSlicesArray metadata length {} does not match outer length {}", @@ -175,7 +175,12 @@ impl VTable for TakeSlices { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); - append_selected_ranges(array.as_view(), builder.as_mut(), ctx)?; + let produced_len = append_selected_ranges(array.as_view(), builder.as_mut(), ctx)?; + vortex_ensure!( + produced_len == array.len(), + "TakeSlicesArray produced length {produced_len} does not match declared length {}", + array.len() + ); Ok(ExecutionResult::done(builder.finish())) } } @@ -196,7 +201,7 @@ impl ValidityVTable for TakeSlices { array .child() .validity()? - .take_slices(array.starts(), array.ends(), array.len()) + .take_slices(array.starts(), array.lengths(), array.len()) } } @@ -204,58 +209,58 @@ fn append_selected_ranges( array: ArrayView<'_, TakeSlices>, builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, -) -> VortexResult<()> { +) -> VortexResult { let starts = array.starts(); - let ends = array.ends(); - check_selector_arrays(starts, ends)?; + let lengths = array.lengths(); + check_selector_arrays(starts, lengths)?; match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { - match_each_unsigned_integer_ptype!(ends.dtype().as_ptype(), |E| { + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { let start = selector_constant::("start", starts)?; - let end = selector_constant::("end", ends)?; - match (start, end) { - (Some(start), Some(end)) => append_ranges::( + let length = selector_constant::("length", lengths)?; + match (start, length) { + (Some(start), Some(length)) => append_ranges::( array.child(), starts.len(), |_| start, - |_| end, + |_| length, builder, ctx, ), (Some(start), None) => { - let ends = ends.clone().execute::(ctx)?; - let ends = ends.as_slice::(); - append_ranges::( + let lengths = lengths.clone().execute::(ctx)?; + let lengths = lengths.as_slice::(); + append_ranges::( array.child(), starts.len(), |_| start, - |index| ends[index], + |index| lengths[index], builder, ctx, ) } - (None, Some(end)) => { + (None, Some(length)) => { let starts = starts.clone().execute::(ctx)?; let starts = starts.as_slice::(); - append_ranges::( + append_ranges::( array.child(), starts.len(), |index| starts[index], - |_| end, + |_| length, builder, ctx, ) } (None, None) => { let starts = starts.clone().execute::(ctx)?; - let ends = ends.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; let starts = starts.as_slice::(); - let ends = ends.as_slice::(); - append_ranges::( + let lengths = lengths.as_slice::(); + append_ranges::( array.child(), starts.len(), |index| starts[index], - |index| ends[index], + |index| lengths[index], builder, ctx, ) @@ -265,26 +270,33 @@ fn append_selected_ranges( }) } -fn append_ranges( +fn append_ranges( child: &ArrayRef, len: usize, mut start_at: StartAt, - mut end_at: EndAt, + mut length_at: LengthAt, builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, -) -> VortexResult<()> +) -> VortexResult where S: IntegerPType, - E: IntegerPType, + L: IntegerPType, StartAt: FnMut(usize) -> S, - EndAt: FnMut(usize) -> E, + LengthAt: FnMut(usize) -> L, { + let mut produced_len = 0usize; for index in 0..len { let start = selector_to_usize("start", start_at(index))?; - let end = selector_to_usize("end", end_at(index))?; + let length = selector_to_usize("length", length_at(index))?; + let end = start.checked_add(length).ok_or_else(|| { + vortex_err!("TakeSlicesArray range overflow for start {start} and length {length}") + })?; child.slice(start..end)?.append_to_builder(builder, ctx)?; + produced_len = produced_len + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; } - Ok(()) + Ok(produced_len) } fn scalar_at_selected_range( @@ -293,57 +305,57 @@ fn scalar_at_selected_range( ctx: &mut ExecutionCtx, ) -> VortexResult> { let starts = array.starts(); - let ends = array.ends(); - check_selector_arrays(starts, ends)?; + let lengths = array.lengths(); + check_selector_arrays(starts, lengths)?; match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { - match_each_unsigned_integer_ptype!(ends.dtype().as_ptype(), |E| { + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { let start = selector_constant::("start", starts)?; - let end = selector_constant::("end", ends)?; - match (start, end) { - (Some(start), Some(end)) => scalar_from_ranges::( + let length = selector_constant::("length", lengths)?; + match (start, length) { + (Some(start), Some(length)) => scalar_from_ranges::( array.child(), starts.len(), index, |_| start, - |_| end, + |_| length, ctx, ), (Some(start), None) => { - let ends = ends.clone().execute::(ctx)?; - let ends = ends.as_slice::(); - scalar_from_ranges::( + let lengths = lengths.clone().execute::(ctx)?; + let lengths = lengths.as_slice::(); + scalar_from_ranges::( array.child(), starts.len(), index, |_| start, - |range_index| ends[range_index], + |range_index| lengths[range_index], ctx, ) } - (None, Some(end)) => { + (None, Some(length)) => { let starts = starts.clone().execute::(ctx)?; let starts = starts.as_slice::(); - scalar_from_ranges::( + scalar_from_ranges::( array.child(), starts.len(), index, |range_index| starts[range_index], - |_| end, + |_| length, ctx, ) } (None, None) => { let starts = starts.clone().execute::(ctx)?; - let ends = ends.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; let starts = starts.as_slice::(); - let ends = ends.as_slice::(); - scalar_from_ranges::( + let lengths = lengths.as_slice::(); + scalar_from_ranges::( array.child(), starts.len(), index, |range_index| starts[range_index], - |range_index| ends[range_index], + |range_index| lengths[range_index], ctx, ) } @@ -352,32 +364,34 @@ fn scalar_at_selected_range( }) } -fn scalar_from_ranges( +fn scalar_from_ranges( child: &ArrayRef, len: usize, index: usize, mut start_at: StartAt, - mut end_at: EndAt, + mut length_at: LengthAt, ctx: &mut ExecutionCtx, ) -> VortexResult> where S: IntegerPType, - E: IntegerPType, + L: IntegerPType, StartAt: FnMut(usize) -> S, - EndAt: FnMut(usize) -> E, + LengthAt: FnMut(usize) -> L, { let mut logical_start = 0usize; for range_index in 0..len { let start = selector_to_usize("start", start_at(range_index))?; - let end = selector_to_usize("end", end_at(range_index))?; + let length = selector_to_usize("length", length_at(range_index))?; + let end = start.checked_add(length).ok_or_else(|| { + vortex_err!("TakeSlicesArray range overflow for start {start} and length {length}") + })?; vortex_ensure!( - start <= end && end <= child.len(), + end <= child.len(), "TakeSlicesArray range {start}..{end} exceeds child array length {}", child.len() ); - let run_len = end - start; let logical_end = logical_start - .checked_add(run_len) + .checked_add(length) .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; if index < logical_end { return child diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index ebbe91ad2b6..fccf60a2486 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -229,11 +229,11 @@ impl Validity { pub fn take_slices( &self, starts: &ArrayRef, - ends: &ArrayRef, + lengths: &ArrayRef, len: usize, ) -> VortexResult { let mut ctx = legacy_session().create_execution_ctx(); - self.take_slices_with_ctx(starts, ends, len, &mut ctx) + self.take_slices_with_ctx(starts, lengths, len, &mut ctx) } /// Select validity values by concatenating a sequence of contiguous ranges, using the caller's @@ -241,7 +241,7 @@ impl Validity { pub fn take_slices_with_ctx( &self, starts: &ArrayRef, - ends: &ArrayRef, + lengths: &ArrayRef, len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -251,7 +251,7 @@ impl Validity { crate::arrays::TakeSlicesArray::try_new( is_valid.clone(), starts.clone(), - ends.clone(), + lengths.clone(), len, )? .into_array() From bebfb5e38bc8ed0264bd1d20dd89a0b7acac60d2 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 13:31:16 -0400 Subject: [PATCH 12/37] Remove redundant TakeSlices metadata Signed-off-by: Daniel King --- vortex-array/src/arrays/take_slices/array.rs | 38 ++----------------- vortex-array/src/arrays/take_slices/mod.rs | 1 - vortex-array/src/arrays/take_slices/vtable.rs | 32 ++-------------- 3 files changed, 7 insertions(+), 64 deletions(-) diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs index f912e9d1792..711cc5dc045 100644 --- a/vortex-array/src/arrays/take_slices/array.rs +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -1,9 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fmt::Display; -use std::fmt::Formatter; - use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -11,6 +8,7 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::array::Array; use crate::array::ArrayParts; +use crate::array::EmptyArrayData; use crate::array::TypedArrayRef; use crate::arrays::TakeSlices; @@ -23,18 +21,6 @@ pub(super) const LENGTHS_SLOT: usize = 2; pub(super) const NUM_SLOTS: usize = 3; pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "starts", "lengths"]; -/// Metadata for a [`TakeSlices`] array. -#[derive(Clone, Debug)] -pub struct TakeSlicesData { - len: usize, -} - -impl Display for TakeSlicesData { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "len: {}", self.len()) - } -} - /// Extension methods for [`TakeSlices`] arrays. pub trait TakeSlicesArrayExt: TypedArrayRef { /// The child array selected by this run sequence. @@ -60,22 +46,6 @@ pub trait TakeSlicesArrayExt: TypedArrayRef { } impl> TakeSlicesArrayExt for T {} -impl TakeSlicesData { - fn new(len: usize) -> Self { - Self { len } - } - - /// Returns the length of this array. - pub fn len(&self) -> usize { - self.len - } - - /// Returns `true` if this array is empty. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } -} - impl Array { /// Constructs a new `TakeSlicesArray` from selector arrays and caller-provided output length. /// @@ -88,9 +58,8 @@ impl Array { len: usize, ) -> VortexResult { let dtype = child.dtype().clone(); - let data = TakeSlicesData::new(len); Array::try_from_parts( - ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![ + ArrayParts::new(TakeSlices, dtype, len, EmptyArrayData).with_slots(smallvec![ Some(child), Some(starts), Some(lengths) @@ -111,10 +80,9 @@ impl Array { len: usize, ) -> Self { let dtype = child.dtype().clone(); - let data = TakeSlicesData::new(len); unsafe { Array::from_parts_unchecked( - ArrayParts::new(TakeSlices, dtype, len, data).with_slots(smallvec![ + ArrayParts::new(TakeSlices, dtype, len, EmptyArrayData).with_slots(smallvec![ Some(child), Some(starts), Some(lengths) diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index 64326dc792b..cea93dcddcf 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -11,7 +11,6 @@ mod array; mod vtable; pub use array::TakeSlicesArrayExt; -pub use array::TakeSlicesData; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index ccb6ac59ebc..08a8eacedb9 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -1,9 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::hash::Hash; -use std::hash::Hasher; - use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -13,11 +10,9 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayEq; -use crate::ArrayHash; use crate::ArrayParts; use crate::ArrayRef; -use crate::EqMode; +use crate::EmptyArrayData; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; @@ -32,7 +27,6 @@ use crate::arrays::take_slices::array::LENGTHS_SLOT; use crate::arrays::take_slices::array::NUM_SLOTS; use crate::arrays::take_slices::array::SLOT_NAMES; use crate::arrays::take_slices::array::STARTS_SLOT; -use crate::arrays::take_slices::array::TakeSlicesData; use crate::arrays::take_slices::check_selector_arrays; use crate::arrays::take_slices::selector_constant; use crate::arrays::take_slices::selector_to_usize; @@ -59,20 +53,8 @@ pub type TakeSlicesArray = Array; #[derive(Clone, Debug)] pub struct TakeSlices; -impl ArrayHash for TakeSlicesData { - fn array_hash(&self, state: &mut H, _accuracy: EqMode) { - self.len().hash(state); - } -} - -impl ArrayEq for TakeSlicesData { - fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { - self.len() == other.len() - } -} - impl VTable for TakeSlices { - type TypedArrayData = TakeSlicesData; + type TypedArrayData = EmptyArrayData; type OperationsVTable = Self; type ValidityVTable = Self; @@ -83,9 +65,9 @@ impl VTable for TakeSlices { fn validate( &self, - data: &Self::TypedArrayData, + _data: &Self::TypedArrayData, dtype: &DType, - len: usize, + _len: usize, slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( @@ -121,12 +103,6 @@ impl VTable for TakeSlices { .as_ref() .vortex_expect("validated lengths slot"); check_selector_arrays(starts, lengths)?; - vortex_ensure!( - data.len() == len, - "TakeSlicesArray metadata length {} does not match outer length {}", - data.len(), - len - ); Ok(()) } From 8c4d42bfbc9c5e1ac54263796f1363ce7900a39b Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 13:38:18 -0400 Subject: [PATCH 13/37] Clean up TakeSlices selector conversions Signed-off-by: Daniel King --- vortex-array/src/array/erased.rs | 13 ++++------- .../arrays/fixed_size_list/compute/take.rs | 23 ++++++------------- vortex-array/src/scalar/convert/primitive.rs | 4 +++- 3 files changed, 14 insertions(+), 26 deletions(-) diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 51b4120f76e..d8dc52ed06c 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -286,12 +286,12 @@ impl ArrayRef { } let starts = starts .into_iter() - .map(selector_value) - .collect::>>()?; + .map(|start| start as u64) + .collect::>(); let lengths = lengths .into_iter() - .map(selector_value) - .collect::>>()?; + .map(|length| length as u64) + .collect::>(); TakeSlicesArray::try_new( self.clone(), PrimitiveArray::from_iter(starts).into_array(), @@ -817,11 +817,6 @@ impl ArrayRef { } } -fn selector_value(value: usize) -> VortexResult { - u64::try_from(value) - .map_err(|_| vortex_err!("TakeSlicesArray selector {value} does not fit in u64")) -} - impl IntoArray for ArrayRef { #[inline(always)] fn into_array(self) -> ArrayRef { diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index aae647de54b..34721449df4 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -202,9 +202,11 @@ fn take_nullable_non_empty_fsl( let mut starts = Vec::with_capacity(new_len); let mut new_validity_builder = BitBufferMut::with_capacity(new_len); + // Null output rows still need placeholder child elements so the FSL elements length stays + // `rows * list_size`. This path has a non-empty source, so 0 is in bounds and validity hides it. for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { if !is_index_valid { - starts.push(null_element_run_start()); + starts.push(0); new_validity_builder.append(false); continue; } @@ -212,7 +214,7 @@ fn take_nullable_non_empty_fsl( let data_idx = index_to_usize(data_idx)?; let start = list_start(data_idx, list_size, array_len)?; if !array_validity.value(data_idx) { - starts.push(null_element_run_start()); + starts.push(0); new_validity_builder.append(false); continue; } @@ -324,10 +326,10 @@ fn take_element_runs( let run_count = starts.len(); let starts = starts .into_iter() - .map(selector_value) - .collect::>>()?; + .map(|start| start as u64) + .collect::>(); let starts = PrimitiveArray::from_iter(starts).into_array(); - let lengths = ConstantArray::new(selector_value(length)?, run_count).into_array(); + let lengths = ConstantArray::new(length as u64, run_count).into_array(); // SAFETY: callers produced one start per output row after validating list indices against the // source FSL length. `length` is the fixed list size, represented as a non-nullable unsigned @@ -341,19 +343,8 @@ fn take_element_runs( ) } -fn null_element_run_start() -> usize { - // Null output rows still need placeholder child elements so the FSL elements length stays - // `rows * list_size`. Any in-bounds run is fine because row validity hides these elements. - 0 -} - fn index_to_usize(index: I) -> VortexResult { index .to_usize() .ok_or_else(|| vortex_err!("FixedSizeList take index {index} does not fit in usize")) } - -fn selector_value(value: usize) -> VortexResult { - u64::try_from(value) - .map_err(|_| vortex_err!("FixedSizeList take selector {value} does not fit in u64")) -} diff --git a/vortex-array/src/scalar/convert/primitive.rs b/vortex-array/src/scalar/convert/primitive.rs index 8b27b5fbd3b..5e0229660b6 100644 --- a/vortex-array/src/scalar/convert/primitive.rs +++ b/vortex-array/src/scalar/convert/primitive.rs @@ -147,7 +147,9 @@ primitive_scalar!(f64); // `Scalar` <---> `usize` conversions. //////////////////////////////////////////////////////////////////////////////////////////// -// NB: We cast `usize` to `u64` (which should always succeed) for better ergonomics. +// NB: Vortex represents `usize` scalar values as `u64` for better ergonomics. Rust's current +// supported target pointer widths fit in `u64`; fail to compile if that assumption changes. +static_assertions::const_assert!(usize::BITS <= u64::BITS); /// Fallible conversion from a [`ScalarValue`] into an `T`. /// From 94f73a291049fad1139531c1bbcb8b33b99f7405 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:05:33 -0400 Subject: [PATCH 14/37] Simplify TakeSlices range handling Signed-off-by: Daniel King --- vortex-array/src/array/erased.rs | 4 +- .../arrays/fixed_size_list/compute/take.rs | 2 +- vortex-array/src/arrays/take_slices/array.rs | 14 +- vortex-array/src/arrays/take_slices/mod.rs | 28 +- vortex-array/src/arrays/take_slices/tests.rs | 6 +- vortex-array/src/arrays/take_slices/vtable.rs | 280 ++++++++++++------ vortex-array/src/validity.rs | 15 +- 7 files changed, 221 insertions(+), 128 deletions(-) diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index d8dc52ed06c..0e1817b77a2 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -269,12 +269,12 @@ impl ArrayRef { /// Wraps the array in a [`TakeSlicesArray`] selected by caller-provided child ranges. /// /// The output is the concatenation of `self[starts[i]..starts[i] + lengths[i]]` for each - /// selector row. This computes the output length from `lengths`, but child bounds are checked + /// range row. This computes the output length from `lengths`, but child bounds are checked /// only when the lazy gather executes. pub fn take_slices(&self, starts: Vec, lengths: Vec) -> VortexResult { vortex_ensure!( starts.len() == lengths.len(), - "TakeSlicesArray selectors must have equal length, got starts {} and lengths {}", + "TakeSlicesArray starts and lengths must have equal length, got starts {} and lengths {}", starts.len(), lengths.len() ); diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 34721449df4..49dcf1bfdab 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -333,7 +333,7 @@ fn take_element_runs( // SAFETY: callers produced one start per output row after validating list indices against the // source FSL length. `length` is the fixed list size, represented as a non-nullable unsigned - // constant selector, and `output_len` was computed as `run_count * length`. + // constant array, and `output_len` was computed as `run_count * length`. Ok( unsafe { TakeSlicesArray::new_unchecked(elements.clone(), starts, lengths, output_len) } .into_array() diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs index 711cc5dc045..b82d893fe67 100644 --- a/vortex-array/src/arrays/take_slices/array.rs +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -14,9 +14,9 @@ use crate::arrays::TakeSlices; /// The child array selected by the run sequence. pub(super) const CHILD_SLOT: usize = 0; -/// The selector naming the start of each child run. +/// The start index for each child run. pub(super) const STARTS_SLOT: usize = 1; -/// The selector naming the length of each child run. +/// The length for each child run. pub(super) const LENGTHS_SLOT: usize = 2; pub(super) const NUM_SLOTS: usize = 3; pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "starts", "lengths"]; @@ -30,14 +30,14 @@ pub trait TakeSlicesArrayExt: TypedArrayRef { .vortex_expect("validated take-slices child slot") } - /// The selector naming each child run's start offset. + /// The start index for each child run. fn starts(&self) -> &ArrayRef { self.as_ref().slots()[STARTS_SLOT] .as_ref() .vortex_expect("validated take-slices starts slot") } - /// The selector naming each child run's length. + /// The length for each child run. fn lengths(&self) -> &ArrayRef { self.as_ref().slots()[LENGTHS_SLOT] .as_ref() @@ -47,9 +47,9 @@ pub trait TakeSlicesArrayExt: TypedArrayRef { impl> TakeSlicesArrayExt for T {} impl Array { - /// Constructs a new `TakeSlicesArray` from selector arrays and caller-provided output length. + /// Constructs a new `TakeSlicesArray` from start/length arrays and caller-provided output length. /// - /// Construction validates only the structural array invariants. Selector values are interpreted + /// Construction validates only the structural array invariants. Index values are interpreted /// when the lazy gather is executed. pub fn try_new( child: ArrayRef, @@ -71,7 +71,7 @@ impl Array { /// /// # Safety /// - /// The caller must ensure the child dtype is the output dtype, selector arrays are + /// The caller must ensure the child dtype is the output dtype, start/length arrays are /// non-nullable unsigned integers of equal length, and `len` is the sum of selected lengths. pub unsafe fn new_unchecked( child: ArrayRef, diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index cea93dcddcf..7bdc34beca4 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -4,7 +4,7 @@ //! Lazy gather of contiguous child ranges. //! //! `TakeSlicesArray` represents the concatenation of -//! `values[starts[i]..starts[i] + lengths[i]]` for each selector row. Ranges may overlap, repeat, +//! `values[starts[i]..starts[i] + lengths[i]]` for each range row. Ranges may overlap, repeat, //! and appear in any order. mod array; @@ -21,25 +21,25 @@ use crate::ArrayRef; use crate::dtype::DType; use crate::dtype::IntegerPType; -pub(super) fn check_selector_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult<()> { - check_selector_dtype("starts", starts)?; - check_selector_dtype("lengths", lengths)?; +pub(super) fn check_index_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult<()> { + check_index_dtype("starts", starts)?; + check_index_dtype("lengths", lengths)?; vortex_ensure!( starts.len() == lengths.len(), - "TakeSlicesArray selectors must have equal length, got starts {} and lengths {}", + "TakeSlicesArray starts and lengths must have equal length, got starts {} and lengths {}", starts.len(), lengths.len() ); Ok(()) } -fn check_selector_dtype(name: &str, selector: &ArrayRef) -> VortexResult<()> { - match selector.dtype() { +fn check_index_dtype(name: &str, indices: &ArrayRef) -> VortexResult<()> { + match indices.dtype() { DType::Primitive(ptype, nullability) if ptype.is_unsigned_int() => { vortex_ensure!( !nullability.is_nullable(), "TakeSlicesArray {name} must be non-nullable, got {}", - selector.dtype() + indices.dtype() ); Ok(()) } @@ -49,25 +49,25 @@ fn check_selector_dtype(name: &str, selector: &ArrayRef) -> VortexResult<()> { } } -pub(super) fn selector_constant( +pub(super) fn constant_index_value( name: &str, - selector: &ArrayRef, + indices: &ArrayRef, ) -> VortexResult> { - selector + indices .as_constant() .map(|scalar| { scalar .as_primitive() .try_typed_value::()? - .ok_or_else(|| vortex_err!("TakeSlicesArray {name} constant selector is null")) + .ok_or_else(|| vortex_err!("TakeSlicesArray {name} constant value is null")) }) .transpose() } -pub(super) fn selector_to_usize(name: &str, value: T) -> VortexResult { +pub(super) fn index_value_to_usize(name: &str, value: T) -> VortexResult { value .to_usize() - .ok_or_else(|| vortex_err!("TakeSlicesArray {name} selector {value} does not fit in usize")) + .ok_or_else(|| vortex_err!("TakeSlicesArray {name} value {value} does not fit in usize")) } #[cfg(test)] diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs index be618796598..bdabec11d44 100644 --- a/vortex-array/src/arrays/take_slices/tests.rs +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -122,7 +122,7 @@ fn take_slices_construction_defers_out_of_bounds_starts_to_execution() -> Vortex } #[test] -fn take_slices_accepts_constant_length_selector() -> VortexResult<()> { +fn take_slices_accepts_constant_length() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); let starts = PrimitiveArray::from_iter([0u64, 1, 2]).into_array(); @@ -136,7 +136,7 @@ fn take_slices_accepts_constant_length_selector() -> VortexResult<()> { } #[test] -fn take_slices_accepts_constant_start_selector() -> VortexResult<()> { +fn take_slices_accepts_constant_start() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); let starts = ConstantArray::new(1u64, 3).into_array(); @@ -150,7 +150,7 @@ fn take_slices_accepts_constant_start_selector() -> VortexResult<()> { } #[test] -fn take_slices_accepts_constant_start_and_length_selectors() -> VortexResult<()> { +fn take_slices_accepts_constant_start_and_length() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); let starts = ConstantArray::new(0u64, 3).into_array(); diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index 08a8eacedb9..1a651434ed1 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -27,9 +27,9 @@ use crate::arrays::take_slices::array::LENGTHS_SLOT; use crate::arrays::take_slices::array::NUM_SLOTS; use crate::arrays::take_slices::array::SLOT_NAMES; use crate::arrays::take_slices::array::STARTS_SLOT; -use crate::arrays::take_slices::check_selector_arrays; -use crate::arrays::take_slices::selector_constant; -use crate::arrays::take_slices::selector_to_usize; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::constant_index_value; +use crate::arrays::take_slices::index_value_to_usize; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity_in; @@ -102,7 +102,7 @@ impl VTable for TakeSlices { let lengths = slots[LENGTHS_SLOT] .as_ref() .vortex_expect("validated lengths slot"); - check_selector_arrays(starts, lengths)?; + check_index_arrays(starts, lengths)?; Ok(()) } @@ -188,29 +188,28 @@ fn append_selected_ranges( ) -> VortexResult { let starts = array.starts(); let lengths = array.lengths(); - check_selector_arrays(starts, lengths)?; + check_index_arrays(starts, lengths)?; match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { - let start = selector_constant::("start", starts)?; - let length = selector_constant::("length", lengths)?; + let start = constant_index_value::("start", starts)?; + let length = constant_index_value::("length", lengths)?; match (start, length) { - (Some(start), Some(length)) => append_ranges::( + (Some(start), Some(length)) => append_constant_start_and_length_ranges( array.child(), starts.len(), - |_| start, - |_| length, + index_value_to_usize("start", start)?, + index_value_to_usize("length", length)?, builder, ctx, ), (Some(start), None) => { let lengths = lengths.clone().execute::(ctx)?; let lengths = lengths.as_slice::(); - append_ranges::( + append_constant_start_ranges( array.child(), - starts.len(), - |_| start, - |index| lengths[index], + index_value_to_usize("start", start)?, + lengths, builder, ctx, ) @@ -218,11 +217,10 @@ fn append_selected_ranges( (None, Some(length)) => { let starts = starts.clone().execute::(ctx)?; let starts = starts.as_slice::(); - append_ranges::( + append_constant_length_ranges( array.child(), - starts.len(), - |index| starts[index], - |_| length, + starts, + index_value_to_usize("length", length)?, builder, ctx, ) @@ -232,49 +230,93 @@ fn append_selected_ranges( let lengths = lengths.clone().execute::(ctx)?; let starts = starts.as_slice::(); let lengths = lengths.as_slice::(); - append_ranges::( - array.child(), - starts.len(), - |index| starts[index], - |index| lengths[index], - builder, - ctx, - ) + append_array_ranges(array.child(), starts, lengths, builder, ctx) } } }) }) } -fn append_ranges( +fn append_constant_start_and_length_ranges( child: &ArrayRef, len: usize, - mut start_at: StartAt, - mut length_at: LengthAt, + start: usize, + length: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut produced_len = 0usize; + for _ in 0..len { + append_range(child, start, length, builder, ctx)?; + produced_len = checked_output_len(produced_len, length, "produced")?; + } + Ok(produced_len) +} + +fn append_constant_start_ranges( + child: &ArrayRef, + start: usize, + lengths: &[L], + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut produced_len = 0usize; + for &length in lengths { + let length = index_value_to_usize("length", length)?; + append_range(child, start, length, builder, ctx)?; + produced_len = checked_output_len(produced_len, length, "produced")?; + } + Ok(produced_len) +} + +fn append_constant_length_ranges( + child: &ArrayRef, + starts: &[S], + length: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut produced_len = 0usize; + for &start in starts { + let start = index_value_to_usize("start", start)?; + append_range(child, start, length, builder, ctx)?; + produced_len = checked_output_len(produced_len, length, "produced")?; + } + Ok(produced_len) +} + +fn append_array_ranges( + child: &ArrayRef, + starts: &[S], + lengths: &[L], builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult where S: IntegerPType, L: IntegerPType, - StartAt: FnMut(usize) -> S, - LengthAt: FnMut(usize) -> L, { let mut produced_len = 0usize; - for index in 0..len { - let start = selector_to_usize("start", start_at(index))?; - let length = selector_to_usize("length", length_at(index))?; - let end = start.checked_add(length).ok_or_else(|| { - vortex_err!("TakeSlicesArray range overflow for start {start} and length {length}") - })?; - child.slice(start..end)?.append_to_builder(builder, ctx)?; - produced_len = produced_len - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; + for (&start, &length) in starts.iter().zip(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + append_range(child, start, length, builder, ctx)?; + produced_len = checked_output_len(produced_len, length, "produced")?; } Ok(produced_len) } +fn append_range( + child: &ArrayRef, + start: usize, + length: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let end = checked_range_end(start, length)?; + child.slice(start..end)?.append_to_builder(builder, ctx) +} + fn scalar_at_selected_range( array: ArrayView<'_, TakeSlices>, index: usize, @@ -282,42 +324,40 @@ fn scalar_at_selected_range( ) -> VortexResult> { let starts = array.starts(); let lengths = array.lengths(); - check_selector_arrays(starts, lengths)?; + check_index_arrays(starts, lengths)?; match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { - let start = selector_constant::("start", starts)?; - let length = selector_constant::("length", lengths)?; + let start = constant_index_value::("start", starts)?; + let length = constant_index_value::("length", lengths)?; match (start, length) { - (Some(start), Some(length)) => scalar_from_ranges::( + (Some(start), Some(length)) => scalar_from_constant_start_and_length_ranges( array.child(), starts.len(), index, - |_| start, - |_| length, + index_value_to_usize("start", start)?, + index_value_to_usize("length", length)?, ctx, ), (Some(start), None) => { let lengths = lengths.clone().execute::(ctx)?; let lengths = lengths.as_slice::(); - scalar_from_ranges::( + scalar_from_constant_start_ranges( array.child(), - starts.len(), index, - |_| start, - |range_index| lengths[range_index], + index_value_to_usize("start", start)?, + lengths, ctx, ) } (None, Some(length)) => { let starts = starts.clone().execute::(ctx)?; let starts = starts.as_slice::(); - scalar_from_ranges::( + scalar_from_constant_length_ranges( array.child(), - starts.len(), index, - |range_index| starts[range_index], - |_| length, + starts, + index_value_to_usize("length", length)?, ctx, ) } @@ -326,55 +366,121 @@ fn scalar_at_selected_range( let lengths = lengths.clone().execute::(ctx)?; let starts = starts.as_slice::(); let lengths = lengths.as_slice::(); - scalar_from_ranges::( - array.child(), - starts.len(), - index, - |range_index| starts[range_index], - |range_index| lengths[range_index], - ctx, - ) + scalar_from_array_ranges(array.child(), index, starts, lengths, ctx) } } }) }) } -fn scalar_from_ranges( +fn scalar_from_constant_start_and_length_ranges( child: &ArrayRef, len: usize, index: usize, - mut start_at: StartAt, - mut length_at: LengthAt, + start: usize, + length: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let mut logical_start = 0usize; + for _ in 0..len { + if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { + return Ok(Some(scalar)); + } + logical_start = checked_output_len(logical_start, length, "logical")?; + } + Ok(None) +} + +fn scalar_from_constant_start_ranges( + child: &ArrayRef, + index: usize, + start: usize, + lengths: &[L], + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let mut logical_start = 0usize; + for &length in lengths { + let length = index_value_to_usize("length", length)?; + if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { + return Ok(Some(scalar)); + } + logical_start = checked_output_len(logical_start, length, "logical")?; + } + Ok(None) +} + +fn scalar_from_constant_length_ranges( + child: &ArrayRef, + index: usize, + starts: &[S], + length: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let mut logical_start = 0usize; + for &start in starts { + let start = index_value_to_usize("start", start)?; + if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { + return Ok(Some(scalar)); + } + logical_start = checked_output_len(logical_start, length, "logical")?; + } + Ok(None) +} + +fn scalar_from_array_ranges( + child: &ArrayRef, + index: usize, + starts: &[S], + lengths: &[L], ctx: &mut ExecutionCtx, ) -> VortexResult> where S: IntegerPType, L: IntegerPType, - StartAt: FnMut(usize) -> S, - LengthAt: FnMut(usize) -> L, { let mut logical_start = 0usize; - for range_index in 0..len { - let start = selector_to_usize("start", start_at(range_index))?; - let length = selector_to_usize("length", length_at(range_index))?; - let end = start.checked_add(length).ok_or_else(|| { - vortex_err!("TakeSlicesArray range overflow for start {start} and length {length}") - })?; - vortex_ensure!( - end <= child.len(), - "TakeSlicesArray range {start}..{end} exceeds child array length {}", - child.len() - ); - let logical_end = logical_start - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; - if index < logical_end { - return child - .execute_scalar(start + (index - logical_start), ctx) - .map(Some); + for (&start, &length) in starts.iter().zip(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { + return Ok(Some(scalar)); } - logical_start = logical_end; + logical_start = checked_output_len(logical_start, length, "logical")?; } Ok(None) } + +fn scalar_from_range( + child: &ArrayRef, + logical_start: usize, + index: usize, + start: usize, + length: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let end = checked_range_end(start, length)?; + vortex_ensure!( + end <= child.len(), + "TakeSlicesArray range {start}..{end} exceeds child array length {}", + child.len() + ); + let logical_end = checked_output_len(logical_start, length, "logical")?; + if index < logical_end { + return child + .execute_scalar(start + (index - logical_start), ctx) + .map(Some); + } + Ok(None) +} + +fn checked_range_end(start: usize, length: usize) -> VortexResult { + start.checked_add(length).ok_or_else(|| { + vortex_err!("TakeSlicesArray range overflow for start {start} and length {length}") + }) +} + +fn checked_output_len(current: usize, length: usize, name: &str) -> VortexResult { + current + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray {name} length overflow")) +} diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index fccf60a2486..05765bcc921 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -231,19 +231,6 @@ impl Validity { starts: &ArrayRef, lengths: &ArrayRef, len: usize, - ) -> VortexResult { - let mut ctx = legacy_session().create_execution_ctx(); - self.take_slices_with_ctx(starts, lengths, len, &mut ctx) - } - - /// Select validity values by concatenating a sequence of contiguous ranges, using the caller's - /// execution context for array-backed validity. - pub fn take_slices_with_ctx( - &self, - starts: &ArrayRef, - lengths: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult { match self { v @ (Self::NonNullable | Self::AllValid | Self::AllInvalid) => Ok(v.clone()), @@ -255,7 +242,7 @@ impl Validity { len, )? .into_array() - .optimize_ctx(ctx.session())?, + .optimize()?, )), } } From 0ed2b46a828ead052ece47b35effbc3c9c43091d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:08:46 -0400 Subject: [PATCH 15/37] Remove unrelated benchmark conversion changes Signed-off-by: Daniel King --- vortex-bench/src/conversions.rs | 229 ++------------------------------ 1 file changed, 11 insertions(+), 218 deletions(-) diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 886b0f48249..92064c25785 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -5,7 +5,6 @@ use std::fs; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; -use std::time::UNIX_EPOCH; use futures::StreamExt; use futures::TryStreamExt; @@ -14,15 +13,10 @@ use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::async_reader::ParquetRecordBatchStream; use parquet::file::metadata::KeyValue; use parquet::file::metadata::ParquetMetaData; -use serde::Deserialize; -use serde::Serialize; -use sha2::Digest; -use sha2::Sha256; use sysinfo::System; use tokio::fs::File; use tokio::fs::OpenOptions; use tokio::fs::create_dir_all; -use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; use tracing::Instrument; use tracing::info; @@ -49,7 +43,6 @@ use vortex::dtype::extension::ExtDType; use vortex::dtype::extension::ExtDTypeRef; use vortex::error::VortexResult; use vortex::error::vortex_err; -use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexWriteOptions; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; @@ -72,7 +65,6 @@ use crate::CompactionStrategy; use crate::Format; use crate::SESSION; use crate::utils::file::idempotent_async; -use crate::utils::file::temp_download_filepath; /// Memory budget per concurrent conversion stream in GB. This is somewhat arbitary. const MEMORY_PER_STREAM_GB: u64 = 4; @@ -82,7 +74,6 @@ const MIN_CONCURRENCY: u64 = 1; /// Maximum number of concurrent conversion streams. This is somewhat arbitary. const MAX_CONCURRENCY: u64 = 16; -const CONVERSION_CACHE_VERSION: u32 = 1; /// Returns the available system memory in bytes. fn available_memory_bytes() -> u64 { @@ -176,204 +167,6 @@ pub async fn convert_parquet_file_to_vortex( Ok(()) } -struct ParquetFileSummary { - row_count: u64, - dtype: DType, - source: SourceFileFingerprint, -} - -struct VortexFileSummary { - row_count: u64, - dtype: DType, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -struct SourceFileFingerprint { - len: u64, - modified_secs: u64, - modified_nanos: u32, - sha256: String, -} - -#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)] -struct ConversionManifest { - version: u32, - compaction: String, - parquet: SourceFileFingerprint, - row_count: u64, - dtype: String, -} - -async fn parquet_file_summary(parquet_path: &Path) -> anyhow::Result { - let file = File::open(parquet_path).await?; - let builder = ParquetRecordBatchStreamBuilder::new(file).await?; - let geo_columns = geoparquet_columns(builder.metadata()); - let dtype = tag_geo_dtype(DType::from_arrow(builder.schema().as_ref()), &geo_columns)?; - let row_count = builder.metadata().file_metadata().num_rows(); - let row_count = u64::try_from(row_count) - .map_err(|_| anyhow::anyhow!("Parquet file has a negative row count: {parquet_path:?}"))?; - Ok(ParquetFileSummary { - row_count, - dtype, - source: source_file_fingerprint(parquet_path).await?, - }) -} - -async fn vortex_file_summary(vortex_path: &Path) -> anyhow::Result { - let file = SESSION.open_options().open_path(vortex_path).await?; - Ok(VortexFileSummary { - row_count: file.row_count(), - dtype: file.dtype().clone(), - }) -} - -async fn existing_vortex_file_matches_parquet( - parquet_path: &Path, - vortex_path: &Path, - compaction: CompactionStrategy, -) -> anyhow::Result { - if !vortex_path.exists() { - return Ok(false); - } - - let expected = parquet_file_summary(parquet_path).await?; - let expected_manifest = ConversionManifest::new(&expected, compaction); - let manifest = read_conversion_manifest(vortex_path).await; - match vortex_file_summary(vortex_path).await { - Ok(actual) - if actual.row_count == expected.row_count - && actual.dtype == expected.dtype - && manifest.as_ref().is_ok_and(|m| m == &expected_manifest) => - { - Ok(true) - } - Ok(actual) => { - info!( - parquet_path = %parquet_path.display(), - vortex_path = %vortex_path.display(), - expected_row_count = expected.row_count, - actual_row_count = actual.row_count, - expected_dtype = %expected.dtype, - actual_dtype = %actual.dtype, - manifest_error = manifest.as_ref().err().map(|e| e.to_string()), - "Regenerating stale Vortex file" - ); - Ok(false) - } - Err(error) => { - info!( - parquet_path = %parquet_path.display(), - vortex_path = %vortex_path.display(), - %error, - "Regenerating unreadable Vortex file" - ); - Ok(false) - } - } -} - -impl ConversionManifest { - fn new(summary: &ParquetFileSummary, compaction: CompactionStrategy) -> Self { - Self { - version: CONVERSION_CACHE_VERSION, - compaction: format!("{compaction:?}"), - parquet: summary.source.clone(), - row_count: summary.row_count, - dtype: summary.dtype.to_string(), - } - } -} - -async fn source_file_fingerprint(path: &Path) -> anyhow::Result { - let metadata = tokio::fs::metadata(path).await?; - let modified = metadata.modified()?.duration_since(UNIX_EPOCH)?; - let mut file = File::open(path).await?; - let mut hasher = Sha256::new(); - let mut buffer = vec![0u8; 1024 * 1024]; - loop { - let read = file.read(&mut buffer).await?; - if read == 0 { - break; - } - hasher.update(&buffer[..read]); - } - - Ok(SourceFileFingerprint { - len: metadata.len(), - modified_secs: modified.as_secs(), - modified_nanos: modified.subsec_nanos(), - sha256: hex_encode(hasher.finalize()), - }) -} - -fn hex_encode(bytes: impl AsRef<[u8]>) -> String { - const HEX: &[u8; 16] = b"0123456789abcdef"; - let bytes = bytes.as_ref(); - let mut encoded = String::with_capacity(bytes.len() * 2); - for &byte in bytes { - encoded.push(HEX[(byte >> 4) as usize] as char); - encoded.push(HEX[(byte & 0x0f) as usize] as char); - } - encoded -} - -fn conversion_manifest_path(vortex_path: &Path) -> PathBuf { - let mut manifest_path = vortex_path.to_path_buf(); - let file_name = vortex_path - .file_name() - .map(|name| name.to_string_lossy()) - .unwrap_or_else(|| "vortex".into()); - manifest_path.set_file_name(format!("{file_name}.manifest.json")); - manifest_path -} - -async fn read_conversion_manifest(vortex_path: &Path) -> anyhow::Result { - let bytes = tokio::fs::read(conversion_manifest_path(vortex_path)).await?; - Ok(serde_json::from_slice(&bytes)?) -} - -async fn write_conversion_manifest( - parquet_path: &Path, - vortex_path: &Path, - compaction: CompactionStrategy, -) -> anyhow::Result<()> { - let summary = parquet_file_summary(parquet_path).await?; - let manifest = ConversionManifest::new(&summary, compaction); - let bytes = serde_json::to_vec_pretty(&manifest)?; - tokio::fs::write(conversion_manifest_path(vortex_path), bytes).await?; - Ok(()) -} - -async fn convert_parquet_file_to_current_vortex( - parquet_path: &Path, - output_path: &Path, - compaction: CompactionStrategy, -) -> anyhow::Result { - if existing_vortex_file_matches_parquet(parquet_path, output_path, compaction).await? { - return Ok(output_path.to_path_buf()); - } - - if let Some(parent) = output_path.parent() { - create_dir_all(parent).await?; - } - - let temp_path = temp_download_filepath(); - info!( - parquet_path = %parquet_path.display(), - vortex_path = %output_path.display(), - ?compaction, - "Processing Parquet file as Vortex" - ); - convert_parquet_file_to_vortex(parquet_path, &temp_path, compaction).await?; - if output_path.exists() { - tokio::fs::remove_file(output_path).await?; - } - tokio::fs::rename(&temp_path, output_path).await?; - write_conversion_manifest(parquet_path, output_path, compaction).await?; - - Ok(output_path.to_path_buf()) -} - /// Whether `path` points at SpatialBench data. fn is_spatialbench(path: &Path) -> bool { path.components() @@ -428,8 +221,7 @@ fn no_dict_layout() -> Arc { /// `{input_path}/vortex-file-compressed/` (for Default compaction) or /// `{input_path}/vortex-compact/` (for Compact compaction). /// -/// Existing Vortex files are reused when their footer metadata and conversion manifest match the -/// source Parquet file. Unreadable or stale Vortex files are regenerated. +/// The conversion is idempotent: existing Vortex files will not be regenerated. pub async fn convert_parquet_directory_to_vortex( input_path: &Path, compaction: CompactionStrategy, @@ -455,7 +247,7 @@ pub async fn convert_parquet_directory_to_vortex( .filter(|entry| entry.path().extension().is_some_and(|e| e == "parquet")); let concurrency = calculate_concurrency(); - let conversion_results = futures::stream::iter(iter) + futures::stream::iter(iter) .map(|dir_entry| { let filename = { let mut temp = dir_entry.path(); @@ -467,12 +259,16 @@ pub async fn convert_parquet_directory_to_vortex( tokio::spawn( async move { - convert_parquet_file_to_current_vortex( - &parquet_file_path, - &output_path, - compaction, - ) + idempotent_async(output_path.as_path(), move |vtx_file| async move { + info!( + "Processing file '{filename}' with {:?} strategy", + compaction + ); + convert_parquet_file_to_vortex(&parquet_file_path, &vtx_file, compaction) + .await + }) .await + .expect("Failed to write Vortex file") } .in_current_span(), ) @@ -480,9 +276,6 @@ pub async fn convert_parquet_directory_to_vortex( .buffer_unordered(concurrency) .try_collect::>() .await?; - for result in conversion_results { - result?; - } Ok(()) } From a455b450aca02ec11e28f52aaf6e1c6245c3b9fa Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:14:23 -0400 Subject: [PATCH 16/37] Remove erased TakeSlices helper Signed-off-by: Daniel King --- vortex-array/src/array/erased.rs | 38 -------------------- vortex-array/src/arrays/listview/rebuild.rs | 31 +++++++++++----- vortex-array/src/arrays/take_slices/tests.rs | 29 +++++++++++---- 3 files changed, 45 insertions(+), 53 deletions(-) diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 0e1817b77a2..c3dc2e0eed7 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -40,9 +40,7 @@ use crate::arrays::DictArray; use crate::arrays::FilterArray; use crate::arrays::Null; use crate::arrays::Primitive; -use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; -use crate::arrays::TakeSlicesArray; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::buffer::BufferHandle; @@ -266,42 +264,6 @@ impl ArrayRef { .optimize() } - /// Wraps the array in a [`TakeSlicesArray`] selected by caller-provided child ranges. - /// - /// The output is the concatenation of `self[starts[i]..starts[i] + lengths[i]]` for each - /// range row. This computes the output length from `lengths`, but child bounds are checked - /// only when the lazy gather executes. - pub fn take_slices(&self, starts: Vec, lengths: Vec) -> VortexResult { - vortex_ensure!( - starts.len() == lengths.len(), - "TakeSlicesArray starts and lengths must have equal length, got starts {} and lengths {}", - starts.len(), - lengths.len() - ); - let mut len = 0usize; - for &length in &lengths { - len = len - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow"))?; - } - let starts = starts - .into_iter() - .map(|start| start as u64) - .collect::>(); - let lengths = lengths - .into_iter() - .map(|length| length as u64) - .collect::>(); - TakeSlicesArray::try_new( - self.clone(), - PrimitiveArray::from_iter(starts).into_array(), - PrimitiveArray::from_iter(lengths).into_array(), - len, - )? - .into_array() - .optimize() - } - /// Fetch the scalar at the given index. #[deprecated( note = "Use `execute_scalar` instead, which allows passing an execution context for more \ diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index a7d36a80dad..ce265128149 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -14,6 +14,7 @@ use crate::RecursiveCanonical; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::TakeSlicesArray; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builtins::ArrayBuiltins; @@ -184,15 +185,29 @@ impl ListViewArray { ) })?; - let new_sizes = ranges.new_sizes.freeze(); - let elements = self - .elements() - .take_slices(ranges.starts, ranges.lengths)? - .execute::(ctx)? - .0 - .into_array(); + let RebuildRanges { + new_offsets, + new_sizes, + starts, + lengths, + } = ranges; + let elements_len = lengths.iter().try_fold(0usize, |acc, &length| { + acc.checked_add(length) + .ok_or_else(|| vortex_err!("ListView rebuild elements length overflow")) + })?; + let starts = + PrimitiveArray::from_iter(starts.into_iter().map(|start| start as u64)).into_array(); + let lengths = + PrimitiveArray::from_iter(lengths.into_iter().map(|length| length as u64)).into_array(); + let new_sizes = new_sizes.freeze(); + let elements = + TakeSlicesArray::try_new(self.elements().clone(), starts, lengths, elements_len)? + .into_array() + .execute::(ctx)? + .0 + .into_array(); // Built unsigned; reinterpret back to the signed-preserving result types. - let offsets = PrimitiveArray::new(ranges.new_offsets.freeze(), Validity::NonNullable) + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(new_offset_ptype) .into_array(); let sizes = PrimitiveArray::new(new_sizes, Validity::NonNullable) diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs index bdabec11d44..064985e95d8 100644 --- a/vortex-array/src/arrays/take_slices/tests.rs +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::IntoArray; @@ -221,16 +222,14 @@ fn take_slices_rejects_invalid_ranges_at_the_right_layer() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(0i32..6).into_array(); - assert!(array.take_slices(vec![0, 0], vec![usize::MAX, 1]).is_err()); - - let out_of_bounds_empty = array.take_slices(vec![7], vec![0])?; + let out_of_bounds_empty = take_slices(&array, &[(7, 0)])?; assert!( out_of_bounds_empty .execute::(&mut ctx) .is_err() ); - let out_of_bounds_non_empty = array.take_slices(vec![4], vec![3])?; + let out_of_bounds_non_empty = take_slices(&array, &[(4, 3)])?; assert!( out_of_bounds_non_empty .execute::(&mut ctx) @@ -303,7 +302,23 @@ fn take_slices_generic_execution_preserves_nullable_encoded_child() -> VortexRes } fn take_slices(array: &ArrayRef, runs: &[(usize, usize)]) -> VortexResult { - let starts = runs.iter().map(|&(start, _)| start).collect::>(); - let lengths = runs.iter().map(|&(_, length)| length).collect::>(); - array.take_slices(starts, lengths) + let len = runs.iter().try_fold(0usize, |acc, &(_, length)| { + acc.checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray length overflow")) + })?; + let starts = runs + .iter() + .map(|&(start, _)| start as u64) + .collect::>(); + let lengths = runs + .iter() + .map(|&(_, length)| length as u64) + .collect::>(); + TakeSlicesArray::try_new( + array.clone(), + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(lengths).into_array(), + len, + ) + .map(IntoArray::into_array) } From c48ac7d0ac828a1b23fc7db1ac1067342d7f0e8e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:23:27 -0400 Subject: [PATCH 17/37] Inline TakeSlices append length handling Signed-off-by: Daniel King --- vortex-array/src/arrays/take_slices/vtable.rs | 70 ++++++++++--------- 1 file changed, 38 insertions(+), 32 deletions(-) diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index 1a651434ed1..45bcedcc464 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools as _; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -247,8 +248,11 @@ fn append_constant_start_and_length_ranges( ) -> VortexResult { let mut produced_len = 0usize; for _ in 0..len { - append_range(child, start, length, builder, ctx)?; - produced_len = checked_output_len(produced_len, length, "produced")?; + let end = checked_range_end(start, length)?; + child.slice(start..end)?.append_to_builder(builder, ctx)?; + produced_len = produced_len + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; } Ok(produced_len) } @@ -263,8 +267,11 @@ fn append_constant_start_ranges( let mut produced_len = 0usize; for &length in lengths { let length = index_value_to_usize("length", length)?; - append_range(child, start, length, builder, ctx)?; - produced_len = checked_output_len(produced_len, length, "produced")?; + let end = checked_range_end(start, length)?; + child.slice(start..end)?.append_to_builder(builder, ctx)?; + produced_len = produced_len + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; } Ok(produced_len) } @@ -279,8 +286,11 @@ fn append_constant_length_ranges( let mut produced_len = 0usize; for &start in starts { let start = index_value_to_usize("start", start)?; - append_range(child, start, length, builder, ctx)?; - produced_len = checked_output_len(produced_len, length, "produced")?; + let end = checked_range_end(start, length)?; + child.slice(start..end)?.append_to_builder(builder, ctx)?; + produced_len = produced_len + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; } Ok(produced_len) } @@ -297,26 +307,18 @@ where L: IntegerPType, { let mut produced_len = 0usize; - for (&start, &length) in starts.iter().zip(lengths) { + for (&start, &length) in starts.iter().zip_eq(lengths) { let start = index_value_to_usize("start", start)?; let length = index_value_to_usize("length", length)?; - append_range(child, start, length, builder, ctx)?; - produced_len = checked_output_len(produced_len, length, "produced")?; + let end = checked_range_end(start, length)?; + child.slice(start..end)?.append_to_builder(builder, ctx)?; + produced_len = produced_len + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; } Ok(produced_len) } -fn append_range( - child: &ArrayRef, - start: usize, - length: usize, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, -) -> VortexResult<()> { - let end = checked_range_end(start, length)?; - child.slice(start..end)?.append_to_builder(builder, ctx) -} - fn scalar_at_selected_range( array: ArrayView<'_, TakeSlices>, index: usize, @@ -386,7 +388,9 @@ fn scalar_from_constant_start_and_length_ranges( if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { return Ok(Some(scalar)); } - logical_start = checked_output_len(logical_start, length, "logical")?; + logical_start = logical_start + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray logical length overflow"))?; } Ok(None) } @@ -404,7 +408,9 @@ fn scalar_from_constant_start_ranges( if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { return Ok(Some(scalar)); } - logical_start = checked_output_len(logical_start, length, "logical")?; + logical_start = logical_start + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray logical length overflow"))?; } Ok(None) } @@ -422,7 +428,9 @@ fn scalar_from_constant_length_ranges( if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { return Ok(Some(scalar)); } - logical_start = checked_output_len(logical_start, length, "logical")?; + logical_start = logical_start + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray logical length overflow"))?; } Ok(None) } @@ -439,13 +447,15 @@ where L: IntegerPType, { let mut logical_start = 0usize; - for (&start, &length) in starts.iter().zip(lengths) { + for (&start, &length) in starts.iter().zip_eq(lengths) { let start = index_value_to_usize("start", start)?; let length = index_value_to_usize("length", length)?; if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { return Ok(Some(scalar)); } - logical_start = checked_output_len(logical_start, length, "logical")?; + logical_start = logical_start + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray logical length overflow"))?; } Ok(None) } @@ -464,7 +474,9 @@ fn scalar_from_range( "TakeSlicesArray range {start}..{end} exceeds child array length {}", child.len() ); - let logical_end = checked_output_len(logical_start, length, "logical")?; + let logical_end = logical_start + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray logical length overflow"))?; if index < logical_end { return child .execute_scalar(start + (index - logical_start), ctx) @@ -478,9 +490,3 @@ fn checked_range_end(start: usize, length: usize) -> VortexResult { vortex_err!("TakeSlicesArray range overflow for start {start} and length {length}") }) } - -fn checked_output_len(current: usize, length: usize, name: &str) -> VortexResult { - current - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray {name} length overflow")) -} From 8443d999fe96738a4ee7b219d00785d9ba555eb0 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:37:01 -0400 Subject: [PATCH 18/37] Restore small ListView rebuild fast path Signed-off-by: Daniel King --- vortex-array/src/arrays/listview/rebuild.rs | 104 +++++++++++++++++++- 1 file changed, 102 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index ce265128149..a19bff08ec3 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -142,13 +142,113 @@ impl ListViewArray { // for sizes as well. match_each_unsigned_integer_ptype!(sizes_ptype.to_unsigned(), |S| { match offsets_ptype.to_unsigned() { - PType::U8 | PType::U16 | PType::U32 => self.rebuild_with_take_slices::(ctx), - PType::U64 => self.rebuild_with_take_slices::(ctx), + PType::U8 => self.rebuild_with_take_or_slices::(ctx), + PType::U16 => self.rebuild_with_take_or_slices::(ctx), + PType::U32 => self.rebuild_with_take_or_slices::(ctx), + PType::U64 => self.rebuild_with_take_or_slices::(ctx), _ => unreachable!("invalid offsets PType"), } }) } + /// Picks between flat-index `take` and contiguous-run `TakeSlices` based on average list size. + /// + /// Small lists are faster with one bulk `take` because it avoids per-run builder overhead. + /// Larger contiguous runs can benefit from `TakeSlices` because it avoids materializing every + /// element index. + fn rebuild_with_take_or_slices( + &self, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let sizes_canonical = self.sizes().clone().execute::(ctx)?; + let sizes_canonical = + sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); + let total: u64 = sizes_canonical + .as_slice::() + .iter() + .map(|s| (*s).as_() as u64) + .sum(); + + if Self::should_use_take(total, self.len()) { + self.rebuild_with_take::(ctx) + } else { + self.rebuild_with_take_slices::(ctx) + } + } + + fn should_use_take(total_output_elements: u64, num_lists: usize) -> bool { + if num_lists == 0 { + return true; + } + let avg = total_output_elements / num_lists as u64; + avg < 128 + } + + /// Rebuilds elements using a single bulk `take`: collect all element indices into a flat + /// `BufferMut`, perform a single `take`. + fn rebuild_with_take( + &self, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + 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); + let mut new_sizes = BufferMut::::with_capacity(len); + let mut take_indices = BufferMut::::with_capacity(self.elements().len()); + + // Resolve validity to a mask once instead of probing it per row: `execute_is_valid` + // executes a scalar on every call for array-backed validity, which is O(len) work repeated + // `len` times. + let validity = self.validity()?.execute_mask(len, ctx)?; + + let mut n_elements = NewOffset::zero(); + for index in 0..len { + if !validity.value(index) { + new_offsets.push(n_elements); + new_sizes.push(S::zero()); + continue; + } + + let offset = offsets_slice[index]; + let size = sizes_slice[index]; + let start = offset.as_(); + let stop = start + size.as_(); + + new_offsets.push(n_elements); + new_sizes.push(size); + take_indices.extend(start as u64..stop as u64); + n_elements += num_traits::cast(size).vortex_expect("Cast failed"); + } + + let elements = self.elements().take(take_indices.into_array())?; + // Built unsigned; reinterpret back to the signed-preserving result types. + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(new_offset_ptype) + .into_array(); + let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable) + .reinterpret_cast(size_ptype) + .into_array(); + + // 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 using a single contiguous-run gather over the element child. /// /// `take_slices` is the rebuild primitive for ListView packing. This path materializes the From 55a19af46784310563ec9b73e630c05894cdf72a Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 16:36:52 -0400 Subject: [PATCH 19/37] Add fast TakeSlices execution kernels Signed-off-by: Daniel King --- vortex-array/src/arrays/bool/compute/mod.rs | 1 + .../src/arrays/bool/compute/take_slices.rs | 93 ++++++ vortex-array/src/arrays/bool/vtable/kernel.rs | 3 + .../src/arrays/decimal/compute/mod.rs | 1 + .../src/arrays/decimal/compute/take_slices.rs | 126 +++++++++ .../src/arrays/decimal/vtable/kernel.rs | 7 + .../src/arrays/fixed_size_list/compute/mod.rs | 1 + .../fixed_size_list/compute/take_slices.rs | 138 +++++++++ .../arrays/fixed_size_list/vtable/kernel.rs | 7 + .../src/arrays/list/compute/kernels.rs | 3 + vortex-array/src/arrays/list/compute/mod.rs | 1 + .../src/arrays/list/compute/take_slices.rs | 265 ++++++++++++++++++ .../src/arrays/primitive/compute/mod.rs | 1 + .../arrays/primitive/compute/take_slices.rs | 131 +++++++++ .../src/arrays/primitive/vtable/kernel.rs | 7 + .../src/arrays/struct_/compute/mod.rs | 1 + .../src/arrays/struct_/compute/take_slices.rs | 84 ++++++ .../src/arrays/struct_/vtable/kernel.rs | 8 + vortex-array/src/arrays/take_slices/kernel.rs | 62 ++++ vortex-array/src/arrays/take_slices/mod.rs | 55 +++- vortex-array/src/arrays/take_slices/tests.rs | 175 ++++++++++++ vortex-array/src/arrays/take_slices/vtable.rs | 250 ++--------------- vortex-array/src/arrays/varbin/compute/mod.rs | 1 + .../src/arrays/varbin/compute/take.rs | 2 +- .../src/arrays/varbin/compute/take_slices.rs | 223 +++++++++++++++ .../src/arrays/varbin/vtable/kernel.rs | 7 + .../src/arrays/varbinview/compute/mod.rs | 1 + .../arrays/varbinview/compute/take_slices.rs | 106 +++++++ .../src/arrays/varbinview/vtable/kernel.rs | 7 + 29 files changed, 1531 insertions(+), 236 deletions(-) create mode 100644 vortex-array/src/arrays/bool/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/decimal/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/fixed_size_list/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/list/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/primitive/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/struct_/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/take_slices/kernel.rs create mode 100644 vortex-array/src/arrays/varbin/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/varbinview/compute/take_slices.rs diff --git a/vortex-array/src/arrays/bool/compute/mod.rs b/vortex-array/src/arrays/bool/compute/mod.rs index 6037cac8c36..369e8e6b55b 100644 --- a/vortex-array/src/arrays/bool/compute/mod.rs +++ b/vortex-array/src/arrays/bool/compute/mod.rs @@ -9,6 +9,7 @@ mod mask; pub mod rules; mod slice; mod take; +mod take_slices; mod zip; #[cfg(test)] diff --git a/vortex-array/src/arrays/bool/compute/take_slices.rs b/vortex-array/src/arrays/bool/compute/take_slices.rs new file mode 100644 index 00000000000..8a98a70ea3b --- /dev/null +++ b/vortex-array/src/arrays/bool/compute/take_slices.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools as _; +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Bool; +use crate::arrays::BoolArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::bool::BoolArrayExt; +use crate::arrays::take_slices::TakeSlicesKernel; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::index_value_to_usize; +use crate::arrays::take_slices::validate_index_ranges; +use crate::dtype::IntegerPType; +use crate::executor::ExecutionCtx; +use crate::match_each_unsigned_integer_ptype; + +impl TakeSlicesKernel for Bool { + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + check_index_arrays(starts, lengths)?; + + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { + take_slices_typed::(array, starts, lengths, output_len, ctx) + }) + }) + .map(Some) + } +} + +fn take_slices_typed( + array: ArrayView<'_, Bool>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, +{ + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + let values = gather_bits( + &array.to_bit_buffer(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + + let starts = starts.into_array(); + let lengths = lengths.into_array(); + let validity = array + .validity()? + .take_slices(&starts, &lengths, output_len)?; + + Ok(BoolArray::new(values, validity).into_array()) +} + +fn gather_bits( + source: &BitBuffer, + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, +{ + validate_index_ranges(source.len(), starts, lengths, output_len)?; + + let mut values = BitBufferMut::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let end = start + length; + values.append_buffer(&source.slice(start..end)); + } + + Ok(values.freeze()) +} diff --git a/vortex-array/src/arrays/bool/vtable/kernel.rs b/vortex-array/src/arrays/bool/vtable/kernel.rs index c19d00d5118..8767f7ff53f 100644 --- a/vortex-array/src/arrays/bool/vtable/kernel.rs +++ b/vortex-array/src/arrays/bool/vtable/kernel.rs @@ -6,7 +6,9 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Bool; use crate::arrays::Dict; +use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::binary::Binary; @@ -24,5 +26,6 @@ pub(crate) fn initialize(session: &VortexSession) { kernels.register_execute_parent_kernel(Cast.id(), Bool, CastExecuteAdaptor(Bool)); kernels.register_execute_parent_kernel(FillNull.id(), Bool, FillNullExecuteAdaptor(Bool)); kernels.register_execute_parent_kernel(Dict.id(), Bool, TakeExecuteAdaptor(Bool)); + kernels.register_execute_parent_kernel(TakeSlices.id(), Bool, TakeSlicesExecuteAdaptor(Bool)); kernels.register_execute_parent_kernel(Zip.id(), Bool, ZipExecuteAdaptor(Bool)); } diff --git a/vortex-array/src/arrays/decimal/compute/mod.rs b/vortex-array/src/arrays/decimal/compute/mod.rs index b2f21fa8398..5c2891dc03f 100644 --- a/vortex-array/src/arrays/decimal/compute/mod.rs +++ b/vortex-array/src/arrays/decimal/compute/mod.rs @@ -7,6 +7,7 @@ mod fill_null; mod mask; pub mod rules; mod take; +mod take_slices; #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/decimal/compute/take_slices.rs b/vortex-array/src/arrays/decimal/compute/take_slices.rs new file mode 100644 index 00000000000..2e47f109e70 --- /dev/null +++ b/vortex-array/src/arrays/decimal/compute/take_slices.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools as _; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::take_slices::TakeSlicesKernel; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::index_value_to_usize; +use crate::arrays::take_slices::validate_index_ranges; +use crate::dtype::IntegerPType; +use crate::dtype::NativeDecimalType; +use crate::executor::ExecutionCtx; +use crate::match_each_decimal_value_type; +use crate::match_each_unsigned_integer_ptype; + +impl TakeSlicesKernel for Decimal { + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + check_index_arrays(starts, lengths)?; + + match_each_decimal_value_type!(array.values_type(), |D| { + take_slices_for_value_type::(array, starts, lengths, output_len, ctx) + }) + } +} + +fn take_slices_for_value_type( + array: ArrayView<'_, Decimal>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> +where + T: NativeDecimalType, +{ + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + take_slices_for_start_type::(array, starts, lengths, output_len, ctx) + }) +} + +fn take_slices_for_start_type( + array: ArrayView<'_, Decimal>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> +where + T: NativeDecimalType, + S: IntegerPType, +{ + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { + take_slices_typed::(array, starts, lengths, output_len, ctx).map(Some) + }) +} + +fn take_slices_typed( + array: ArrayView<'_, Decimal>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: NativeDecimalType, + S: IntegerPType, + L: IntegerPType, +{ + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + let values = gather_values( + array.buffer::().as_slice(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + + let starts = starts.into_array(); + let lengths = lengths.into_array(); + let validity = array + .validity()? + .take_slices(&starts, &lengths, output_len)?; + + // SAFETY: contiguous gather preserves the decimal dtype and value representation. + let array = unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) }; + Ok(array.into_array()) +} + +fn gather_values( + source: &[T], + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + T: NativeDecimalType, + S: IntegerPType, + L: IntegerPType, +{ + validate_index_ranges(source.len(), starts, lengths, output_len)?; + + let mut values = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let end = start + length; + values.extend_from_slice(&source[start..end]); + } + + Ok(values.freeze()) +} diff --git a/vortex-array/src/arrays/decimal/vtable/kernel.rs b/vortex-array/src/arrays/decimal/vtable/kernel.rs index 19bb30d8c95..3d8cfc8498f 100644 --- a/vortex-array/src/arrays/decimal/vtable/kernel.rs +++ b/vortex-array/src/arrays/decimal/vtable/kernel.rs @@ -6,7 +6,9 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Decimal; use crate::arrays::Dict; +use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; @@ -22,4 +24,9 @@ pub(crate) fn initialize(session: &VortexSession) { kernels.register_execute_parent_kernel(Cast.id(), Decimal, CastExecuteAdaptor(Decimal)); kernels.register_execute_parent_kernel(FillNull.id(), Decimal, FillNullExecuteAdaptor(Decimal)); kernels.register_execute_parent_kernel(Dict.id(), Decimal, TakeExecuteAdaptor(Decimal)); + kernels.register_execute_parent_kernel( + TakeSlices.id(), + Decimal, + TakeSlicesExecuteAdaptor(Decimal), + ); } diff --git a/vortex-array/src/arrays/fixed_size_list/compute/mod.rs b/vortex-array/src/arrays/fixed_size_list/compute/mod.rs index 9a43503c4b5..880630caece 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/mod.rs @@ -6,3 +6,4 @@ mod mask; pub(crate) mod rules; mod slice; mod take; +mod take_slices; diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take_slices.rs b/vortex-array/src/arrays/fixed_size_list/compute/take_slices.rs new file mode 100644 index 00000000000..c1e7347c087 --- /dev/null +++ b/vortex-array/src/arrays/fixed_size_list/compute/take_slices.rs @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools as _; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::FixedSizeList; +use crate::arrays::FixedSizeListArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::TakeSlicesArray; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::take_slices::TakeSlicesKernel; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::index_value_to_usize; +use crate::arrays::take_slices::validate_index_ranges; +use crate::dtype::IntegerPType; +use crate::executor::ExecutionCtx; +use crate::match_each_unsigned_integer_ptype; + +impl TakeSlicesKernel for FixedSizeList { + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + check_index_arrays(starts, lengths)?; + + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { + take_slices_typed::(array, starts, lengths, output_len, ctx) + }) + }) + .map(Some) + } +} + +fn take_slices_typed( + array: ArrayView<'_, FixedSizeList>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, +{ + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + validate_index_ranges( + array.len(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + + let list_size = array.list_size() as usize; + let elements_len = output_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList TakeSlices output length overflow: {output_len} lists of size {list_size}" + ) + })?; + let elements = if list_size == 0 { + array.elements().clone() + } else { + gather_elements( + array.elements(), + starts.as_slice::(), + lengths.as_slice::(), + list_size, + elements_len, + )? + }; + + let starts = starts.into_array(); + let lengths = lengths.into_array(); + let validity = array + .validity()? + .take_slices(&starts, &lengths, output_len)?; + + // SAFETY: row ranges were validated, element ranges are the corresponding checked + // fixed-width expansions, and validity has one entry per output row. + Ok(unsafe { + FixedSizeListArray::new_unchecked(elements, array.list_size(), validity, output_len) + } + .into_array()) +} + +fn gather_elements( + elements: &ArrayRef, + starts: &[S], + lengths: &[L], + list_size: usize, + elements_len: usize, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, +{ + let mut element_starts = BufferMut::::with_capacity(starts.len()); + let mut element_lengths = BufferMut::::with_capacity(lengths.len()); + + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let element_start = start.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList TakeSlices element start overflow for start {start} and list size {list_size}" + ) + })?; + let element_length = length.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList TakeSlices element length overflow for length {length} and list size {list_size}" + ) + })?; + element_starts.push(element_start as u64); + element_lengths.push(element_length as u64); + } + + // SAFETY: row ranges have already been validated against the FSL length; multiplying by + // `list_size` maps them to valid element ranges, and `elements_len` is `output_len * list_size`. + Ok(unsafe { + TakeSlicesArray::new_unchecked( + elements.clone(), + element_starts.into_array(), + element_lengths.into_array(), + elements_len, + ) + } + .into_array()) +} diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs index 47dc3ec0651..0181b375da0 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs @@ -6,7 +6,9 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::FixedSizeList; +use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; @@ -24,4 +26,9 @@ pub(crate) fn initialize(session: &VortexSession) { FixedSizeList, TakeExecuteAdaptor(FixedSizeList), ); + kernels.register_execute_parent_kernel( + TakeSlices.id(), + FixedSizeList, + TakeSlicesExecuteAdaptor(FixedSizeList), + ); } diff --git a/vortex-array/src/arrays/list/compute/kernels.rs b/vortex-array/src/arrays/list/compute/kernels.rs index 6af858dbf49..dd0e8724938 100644 --- a/vortex-array/src/arrays/list/compute/kernels.rs +++ b/vortex-array/src/arrays/list/compute/kernels.rs @@ -7,8 +7,10 @@ use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::Filter; use crate::arrays::List; +use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; @@ -19,4 +21,5 @@ pub(crate) fn initialize(session: &VortexSession) { kernels.register_execute_parent_kernel(Cast.id(), List, CastExecuteAdaptor(List)); kernels.register_execute_parent_kernel(Filter.id(), List, FilterExecuteAdaptor(List)); kernels.register_execute_parent_kernel(Dict.id(), List, TakeExecuteAdaptor(List)); + kernels.register_execute_parent_kernel(TakeSlices.id(), List, TakeSlicesExecuteAdaptor(List)); } diff --git a/vortex-array/src/arrays/list/compute/mod.rs b/vortex-array/src/arrays/list/compute/mod.rs index ce6cd960f74..00e7a0ae203 100644 --- a/vortex-array/src/arrays/list/compute/mod.rs +++ b/vortex-array/src/arrays/list/compute/mod.rs @@ -8,6 +8,7 @@ mod mask; pub(crate) mod rules; mod slice; mod take; +mod take_slices; pub(crate) fn initialize(session: &vortex_session::VortexSession) { kernels::initialize(session); diff --git a/vortex-array/src/arrays/list/compute/take_slices.rs b/vortex-array/src/arrays/list/compute/take_slices.rs new file mode 100644 index 00000000000..88fc1c894e2 --- /dev/null +++ b/vortex-array/src/arrays/list/compute/take_slices.rs @@ -0,0 +1,265 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools as _; +use vortex_buffer::BufferMut; +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::List; +use crate::arrays::ListArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::TakeSlicesArray; +use crate::arrays::list::ListArrayExt; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::arrays::take_slices::TakeSlicesKernel; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::index_value_to_usize; +use crate::arrays::take_slices::validate_index_ranges; +use crate::dtype::IntegerPType; +use crate::dtype::PType; +use crate::executor::ExecutionCtx; +use crate::match_each_unsigned_integer_ptype; +use crate::match_smallest_offset_type; +use crate::validity::Validity; + +impl TakeSlicesKernel for List { + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + check_index_arrays(starts, lengths)?; + + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { + take_slices_typed::(array, starts, lengths, output_len, ctx) + }) + }) + .map(Some) + } +} + +fn take_slices_typed( + array: ArrayView<'_, List>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, +{ + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + validate_index_ranges( + array.len(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + + let offsets = array.offsets().clone().execute::(ctx)?; + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let total_elements = match offsets.ptype() { + PType::U8 => total_elements::( + array.elements().len(), + offsets.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + ), + PType::U16 => total_elements::( + array.elements().len(), + offsets.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + ), + PType::U32 => total_elements::( + array.elements().len(), + offsets.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + ), + PType::U64 => total_elements::( + array.elements().len(), + offsets.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + }?; + + let gathered = match_smallest_offset_type!(total_elements, |OutputOffset| { + match offsets.ptype() { + PType::U8 => gather_list::( + array.elements(), + offsets.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + total_elements, + ), + PType::U16 => gather_list::( + array.elements(), + offsets.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + total_elements, + ), + PType::U32 => gather_list::( + array.elements(), + offsets.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + total_elements, + ), + PType::U64 => gather_list::( + array.elements(), + offsets.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + total_elements, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } + })?; + + let starts = starts.into_array(); + let lengths = lengths.into_array(); + let validity = array + .validity()? + .take_slices(&starts, &lengths, output_len)?; + + // SAFETY: output offsets are rebuilt from valid, monotonic source offsets; output elements are + // exactly the gathered element ranges referenced by those offsets; validity has output_len rows. + Ok( + unsafe { ListArray::new_unchecked(gathered.elements, gathered.offsets, validity) } + .into_array(), + ) +} + +struct GatheredList { + elements: ArrayRef, + offsets: ArrayRef, +} + +fn total_elements( + elements_len: usize, + offsets: &[Offset], + starts: &[S], + lengths: &[L], +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, + Offset: IntegerPType, +{ + let mut total = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let end = start + length; + if length == 0 { + continue; + } + + let element_start = index_value_to_usize("offset", offsets[start])?; + let element_end = index_value_to_usize("offset", offsets[end])?; + 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!("TakeSlicesArray List output elements length overflow"))?; + } + Ok(total) +} + +fn gather_list( + elements: &ArrayRef, + offsets: &[Offset], + starts: &[S], + lengths: &[L], + output_len: usize, + total_elements: usize, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, + Offset: IntegerPType, + OutputOffset: IntegerPType, +{ + let offsets_capacity = output_len + .checked_add(1) + .ok_or_else(|| vortex_err!("TakeSlicesArray List 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 = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let end = start + length; + if length == 0 { + continue; + } + + let element_start = index_value_to_usize("offset", offsets[start])?; + let element_end = index_value_to_usize("offset", offsets[end])?; + for &offset in &offsets[start + 1..=end] { + let offset = index_value_to_usize("offset", offset)?; + 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!("TakeSlicesArray List 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!("TakeSlicesArray List output elements length overflow"))?; + } + debug_assert_eq!(output_elements, total_elements); + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + // SAFETY: element ranges are derived from validated list offsets, and total_elements is the sum + // of all gathered element ranges. + let elements = unsafe { + TakeSlicesArray::new_unchecked( + elements.clone(), + element_starts.into_array(), + element_lengths.into_array(), + total_elements, + ) + } + .into_array(); + + Ok(GatheredList { elements, offsets }) +} + +fn new_offset_value(value: usize) -> VortexResult { + T::from(value).ok_or_else(|| { + vortex_err!( + "TakeSlicesArray List offset value {value} does not fit in {}", + T::PTYPE + ) + }) +} diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index 1ca4b17d3d5..37aa3f98396 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -8,6 +8,7 @@ mod mask; pub(crate) mod rules; mod slice; mod take; +mod take_slices; mod zip; #[cfg(test)] diff --git a/vortex-array/src/arrays/primitive/compute/take_slices.rs b/vortex-array/src/arrays/primitive/compute/take_slices.rs new file mode 100644 index 00000000000..185285337d7 --- /dev/null +++ b/vortex-array/src/arrays/primitive/compute/take_slices.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools as _; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::take_slices::TakeSlicesKernel; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::index_value_to_usize; +use crate::arrays::take_slices::validate_index_ranges; +use crate::dtype::IntegerPType; +use crate::dtype::NativePType; +use crate::executor::ExecutionCtx; +use crate::match_each_native_ptype; +use crate::match_each_unsigned_integer_ptype; + +impl TakeSlicesKernel for Primitive { + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + execute_primitive_selected_ranges(array, starts, lengths, output_len, ctx).map(Some) + } +} + +fn execute_primitive_selected_ranges( + array: ArrayView<'_, Primitive>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match_each_native_ptype!(array.ptype(), |T| { + execute_primitive_selected_ranges_for_type::(array, starts, lengths, output_len, ctx) + }) +} + +fn execute_primitive_selected_ranges_for_type( + array: ArrayView<'_, Primitive>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + check_index_arrays(starts, lengths)?; + + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + execute_primitive_selected_ranges_for_start_type::( + array, starts, lengths, output_len, ctx, + ) + }) +} + +fn execute_primitive_selected_ranges_for_start_type( + array: ArrayView<'_, Primitive>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: NativePType, + S: IntegerPType, +{ + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { + execute_primitive_selected_ranges_typed::(array, starts, lengths, output_len, ctx) + }) +} + +fn execute_primitive_selected_ranges_typed( + array: ArrayView<'_, Primitive>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: NativePType, + S: IntegerPType, + L: IntegerPType, +{ + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + let values = primitive_array_ranges::( + array, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + let starts = starts.into_array(); + let lengths = lengths.into_array(); + let validity = array + .validity()? + .take_slices(&starts, &lengths, output_len)?; + Ok(PrimitiveArray::new(values, validity).into_array()) +} + +fn primitive_array_ranges( + array: ArrayView<'_, Primitive>, + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + T: NativePType, + S: IntegerPType, + L: IntegerPType, +{ + let source = array.as_slice::(); + validate_index_ranges(source.len(), starts, lengths, output_len)?; + + let mut values = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let end = start + length; + values.extend_from_slice(&source[start..end]); + } + + Ok(values.freeze()) +} diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 6382ea73794..663c8924865 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -6,7 +6,9 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::Primitive; +use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; @@ -32,5 +34,10 @@ pub(crate) fn initialize(session: &VortexSession) { FillNullExecuteAdaptor(Primitive), ); kernels.register_execute_parent_kernel(Dict.id(), Primitive, TakeExecuteAdaptor(Primitive)); + kernels.register_execute_parent_kernel( + TakeSlices.id(), + Primitive, + TakeSlicesExecuteAdaptor(Primitive), + ); kernels.register_execute_parent_kernel(Zip.id(), Primitive, ZipExecuteAdaptor(Primitive)); } diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index 4b9a4e396fd..8965b3ce9c3 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -6,6 +6,7 @@ mod mask; pub(crate) mod rules; mod slice; mod take; +mod take_slices; mod zip; #[cfg(test)] diff --git a/vortex-array/src/arrays/struct_/compute/take_slices.rs b/vortex-array/src/arrays/struct_/compute/take_slices.rs new file mode 100644 index 00000000000..b3140f07931 --- /dev/null +++ b/vortex-array/src/arrays/struct_/compute/take_slices.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; +use crate::arrays::StructArray; +use crate::arrays::TakeSlicesArray; +use crate::arrays::struct_::StructArrayExt; +use crate::arrays::take_slices::TakeSlicesKernel; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::validate_index_ranges; +use crate::dtype::IntegerPType; +use crate::executor::ExecutionCtx; +use crate::match_each_unsigned_integer_ptype; + +impl TakeSlicesKernel for Struct { + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + check_index_arrays(starts, lengths)?; + + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { + take_slices_typed::(array, starts, lengths, output_len, ctx) + }) + }) + .map(Some) + } +} + +fn take_slices_typed( + array: ArrayView<'_, Struct>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, +{ + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + validate_index_ranges( + array.len(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + + let starts = starts.into_array(); + let lengths = lengths.into_array(); + let fields = array + .iter_unmasked_fields() + .map(|field| { + // SAFETY: the index arrays and declared output length were validated above, and the + // child dtype becomes the outer dtype of this per-field TakeSlices array. + unsafe { + TakeSlicesArray::new_unchecked( + field.clone(), + starts.clone(), + lengths.clone(), + output_len, + ) + .into_array() + } + }) + .collect::>(); + let validity = array + .validity()? + .take_slices(&starts, &lengths, output_len)?; + + StructArray::try_new_with_dtype(fields, array.struct_fields().clone(), output_len, validity) + .map(StructArray::into_array) +} diff --git a/vortex-array/src/arrays/struct_/vtable/kernel.rs b/vortex-array/src/arrays/struct_/vtable/kernel.rs index 1e58da2d0d0..fee4d3484d0 100644 --- a/vortex-array/src/arrays/struct_/vtable/kernel.rs +++ b/vortex-array/src/arrays/struct_/vtable/kernel.rs @@ -3,7 +3,10 @@ use vortex_session::VortexSession; +use crate::ArrayVTable; use crate::arrays::Struct; +use crate::arrays::TakeSlices; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::zip::Zip; @@ -11,5 +14,10 @@ use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + TakeSlices.id(), + Struct, + TakeSlicesExecuteAdaptor(Struct), + ); kernels.register_execute_parent_kernel(Zip.id(), Struct, ZipExecuteAdaptor(Struct)); } diff --git a/vortex-array/src/arrays/take_slices/kernel.rs b/vortex-array/src/arrays/take_slices/kernel.rs new file mode 100644 index 00000000000..b18f56bc165 --- /dev/null +++ b/vortex-array/src/arrays/take_slices/kernel.rs @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execute adaptor for `TakeSlices` parent operations. + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::arrays::TakeSlices; +use crate::arrays::take_slices::TakeSlicesArrayExt; +use crate::arrays::take_slices::array::CHILD_SLOT; +use crate::kernel::ExecuteParentKernel; +use crate::matcher::Matcher; + +/// Execution kernel for child encodings that can materialize a `TakeSlices` parent directly. +pub trait TakeSlicesKernel: VTable { + /// Gather contiguous runs from `array`. + /// + /// `starts` and `lengths` are non-nullable unsigned integer arrays of equal length. `output_len` + /// is the declared length of the `TakeSlices` parent and must match the sum of selected lengths. + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} + +/// Adapter that exposes a [`TakeSlicesKernel`] implementation as an [`ExecuteParentKernel`]. +#[derive(Default, Debug)] +pub struct TakeSlicesExecuteAdaptor(pub V); + +impl ExecuteParentKernel for TakeSlicesExecuteAdaptor +where + V: TakeSlicesKernel, +{ + type Parent = TakeSlices; + + fn execute_parent( + &self, + array: ArrayView<'_, V>, + parent: ::Match<'_>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if child_idx != CHILD_SLOT { + return Ok(None); + } + + ::take_slices( + array, + parent.starts(), + parent.lengths(), + parent.len(), + ctx, + ) + } +} diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index 7bdc34beca4..20f260c544d 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -8,9 +8,13 @@ //! and appear in any order. mod array; +mod kernel; mod vtable; pub use array::TakeSlicesArrayExt; +use itertools::Itertools as _; +pub use kernel::TakeSlicesExecuteAdaptor; +pub use kernel::TakeSlicesKernel; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -49,26 +53,47 @@ fn check_index_dtype(name: &str, indices: &ArrayRef) -> VortexResult<()> { } } -pub(super) fn constant_index_value( - name: &str, - indices: &ArrayRef, -) -> VortexResult> { - indices - .as_constant() - .map(|scalar| { - scalar - .as_primitive() - .try_typed_value::()? - .ok_or_else(|| vortex_err!("TakeSlicesArray {name} constant value is null")) - }) - .transpose() -} - pub(super) fn index_value_to_usize(name: &str, value: T) -> VortexResult { value .to_usize() .ok_or_else(|| vortex_err!("TakeSlicesArray {name} value {value} does not fit in usize")) } +pub(super) fn checked_range_end(start: usize, length: usize) -> VortexResult { + start.checked_add(length).ok_or_else(|| { + vortex_err!("TakeSlicesArray range overflow for start {start} and length {length}") + }) +} + +pub(super) fn validate_index_ranges( + child_len: usize, + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult<()> +where + S: IntegerPType, + L: IntegerPType, +{ + let mut produced_len = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let end = checked_range_end(start, length)?; + vortex_ensure!( + end <= child_len, + "TakeSlicesArray range {start}..{end} exceeds child array length {child_len}", + ); + produced_len = produced_len + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; + } + vortex_ensure!( + produced_len == output_len, + "TakeSlicesArray produced length {produced_len} does not match declared length {output_len}", + ); + Ok(()) +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs index 064985e95d8..4902c212c70 100644 --- a/vortex-array/src/arrays/take_slices/tests.rs +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -8,16 +8,26 @@ use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::BoolArray; use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; use crate::arrays::DictArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::TakeSlices; use crate::arrays::TakeSlicesArray; +use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; use crate::dtype::FieldNames; use crate::dtype::Nullability; +use crate::kernel::ExecuteParentKernel; use crate::validity::Validity; #[test] @@ -164,6 +174,40 @@ fn take_slices_accepts_constant_start_and_length() -> VortexResult<()> { Ok(()) } +#[test] +fn primitive_take_slices_execute_parent_handles_only_child_slot() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let child = PrimitiveArray::from_iter(0i32..6).into_array(); + let starts = PrimitiveArray::from_iter([3u64, 0]).into_array(); + let lengths = PrimitiveArray::from_iter([2u64, 2]).into_array(); + let parent = TakeSlicesArray::try_new(child.clone(), starts, lengths, 4)?; + + let skipped = TakeSlicesExecuteAdaptor(Primitive).execute_parent( + child + .as_typed::() + .ok_or_else(|| vortex_err!("expected primitive child"))?, + parent.as_view(), + 1, + &mut ctx, + )?; + assert!(skipped.is_none()); + + let actual = TakeSlicesExecuteAdaptor(Primitive) + .execute_parent( + child + .as_typed::() + .ok_or_else(|| vortex_err!("expected primitive child"))?, + parent.as_view(), + 0, + &mut ctx, + )? + .ok_or_else(|| vortex_err!("primitive TakeSlices execute parent returned no result"))?; + let expected = PrimitiveArray::from_iter([3i32, 4, 0, 1]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + #[test] fn struct_take_slices_executes_generically() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -193,6 +237,137 @@ fn struct_take_slices_executes_generically() -> VortexResult<()> { Ok(()) } +#[test] +fn bool_take_slices_kernel_preserves_values_and_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = + BoolArray::from_iter([Some(true), None, Some(false), Some(true), None, Some(false)]) + .into_array(); + + let actual = take_slices(&array, &[(2, 3), (0, 2)])?; + let expected = BoolArray::from_iter([Some(false), Some(true), None, Some(true), None]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn decimal_take_slices_kernel_preserves_values_and_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(19, 2); + let array = DecimalArray::new( + vortex_buffer::buffer![10i128, 20, 30, 40, 50, 60], + dtype, + Validity::from_iter([true, false, true, true, false, true]), + ) + .into_array(); + + let actual = take_slices(&array, &[(2, 3), (0, 2)])?; + let expected = DecimalArray::new( + vortex_buffer::buffer![30i128, 40, 50, 10, 20], + dtype, + Validity::from_iter([true, true, false, true, false]), + ); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn varbinview_take_slices_kernel_reuses_views_and_preserves_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter( + [ + Some("short"), + None, + Some("a string that is long enough to be outlined"), + Some("middle"), + Some("another outlined string value"), + None, + ], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + + let actual = take_slices(&array, &[(2, 3), (0, 2)])?; + let expected = VarBinViewArray::from_iter( + [ + Some("a string that is long enough to be outlined"), + Some("middle"), + Some("another outlined string value"), + Some("short"), + None, + ], + DType::Utf8(Nullability::Nullable), + ); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn varbin_take_slices_kernel_rebuilds_offsets_and_preserves_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinArray::from_iter( + [Some("aa"), None, Some("ccc"), Some(""), Some("dddd"), None], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + + let actual = take_slices(&array, &[(2, 3), (0, 2)])?; + let expected = VarBinArray::from_iter( + [Some("ccc"), Some(""), Some("dddd"), Some("aa"), None], + DType::Utf8(Nullability::Nullable), + ); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn fixed_size_list_take_slices_kernel_maps_row_runs_to_element_runs() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = FixedSizeListArray::new( + PrimitiveArray::from_iter(0i32..8).into_array(), + 2, + Validity::from_iter([true, false, true, true]), + 4, + ) + .into_array(); + + let actual = take_slices(&array, &[(1, 2), (0, 1)])?; + let expected = FixedSizeListArray::new( + PrimitiveArray::from_iter([2i32, 3, 4, 5, 0, 1]).into_array(), + 2, + Validity::from_iter([false, true, true]), + 3, + ); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn list_take_slices_kernel_rebuilds_offsets_and_pushes_into_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = ListArray::new( + PrimitiveArray::from_iter(0i32..8).into_array(), + vortex_buffer::buffer![0u32, 2, 5, 5, 8].into_array(), + Validity::from_iter([true, false, true, true]), + ) + .into_array(); + + let actual = take_slices(&array, &[(1, 3), (0, 1)])?; + let expected = ListArray::new( + PrimitiveArray::from_iter([2i32, 3, 4, 5, 6, 7, 0, 1]).into_array(), + vortex_buffer::buffer![0u32, 3, 3, 6, 8].into_array(), + Validity::from_iter([false, true, true, true]), + ); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + #[test] fn take_slices_empty_runs_return_empty_canonical() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index 45bcedcc464..cc529a7f501 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -29,8 +29,9 @@ use crate::arrays::take_slices::array::NUM_SLOTS; use crate::arrays::take_slices::array::SLOT_NAMES; use crate::arrays::take_slices::array::STARTS_SLOT; use crate::arrays::take_slices::check_index_arrays; -use crate::arrays::take_slices::constant_index_value; +use crate::arrays::take_slices::checked_range_end; use crate::arrays::take_slices::index_value_to_usize; +use crate::arrays::take_slices::validate_index_ranges; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity_in; @@ -152,12 +153,7 @@ impl VTable for TakeSlices { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); - let produced_len = append_selected_ranges(array.as_view(), builder.as_mut(), ctx)?; - vortex_ensure!( - produced_len == array.len(), - "TakeSlicesArray produced length {produced_len} does not match declared length {}", - array.len() - ); + append_selected_ranges(array.as_view(), builder.as_mut(), ctx)?; Ok(ExecutionResult::done(builder.finish())) } } @@ -186,137 +182,49 @@ fn append_selected_ranges( array: ArrayView<'_, TakeSlices>, builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, -) -> VortexResult { +) -> VortexResult<()> { let starts = array.starts(); let lengths = array.lengths(); check_index_arrays(starts, lengths)?; match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { - let start = constant_index_value::("start", starts)?; - let length = constant_index_value::("length", lengths)?; - match (start, length) { - (Some(start), Some(length)) => append_constant_start_and_length_ranges( - array.child(), - starts.len(), - index_value_to_usize("start", start)?, - index_value_to_usize("length", length)?, - builder, - ctx, - ), - (Some(start), None) => { - let lengths = lengths.clone().execute::(ctx)?; - let lengths = lengths.as_slice::(); - append_constant_start_ranges( - array.child(), - index_value_to_usize("start", start)?, - lengths, - builder, - ctx, - ) - } - (None, Some(length)) => { - let starts = starts.clone().execute::(ctx)?; - let starts = starts.as_slice::(); - append_constant_length_ranges( - array.child(), - starts, - index_value_to_usize("length", length)?, - builder, - ctx, - ) - } - (None, None) => { - let starts = starts.clone().execute::(ctx)?; - let lengths = lengths.clone().execute::(ctx)?; - let starts = starts.as_slice::(); - let lengths = lengths.as_slice::(); - append_array_ranges(array.child(), starts, lengths, builder, ctx) - } - } + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + append_array_ranges( + array.child(), + starts.as_slice::(), + lengths.as_slice::(), + array.len(), + builder, + ctx, + ) }) }) } -fn append_constant_start_and_length_ranges( - child: &ArrayRef, - len: usize, - start: usize, - length: usize, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let mut produced_len = 0usize; - for _ in 0..len { - let end = checked_range_end(start, length)?; - child.slice(start..end)?.append_to_builder(builder, ctx)?; - produced_len = produced_len - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; - } - Ok(produced_len) -} - -fn append_constant_start_ranges( - child: &ArrayRef, - start: usize, - lengths: &[L], - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let mut produced_len = 0usize; - for &length in lengths { - let length = index_value_to_usize("length", length)?; - let end = checked_range_end(start, length)?; - child.slice(start..end)?.append_to_builder(builder, ctx)?; - produced_len = produced_len - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; - } - Ok(produced_len) -} - -fn append_constant_length_ranges( - child: &ArrayRef, - starts: &[S], - length: usize, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let mut produced_len = 0usize; - for &start in starts { - let start = index_value_to_usize("start", start)?; - let end = checked_range_end(start, length)?; - child.slice(start..end)?.append_to_builder(builder, ctx)?; - produced_len = produced_len - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; - } - Ok(produced_len) -} - fn append_array_ranges( child: &ArrayRef, starts: &[S], lengths: &[L], + output_len: usize, builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, -) -> VortexResult +) -> VortexResult<()> where S: IntegerPType, L: IntegerPType, { - let mut produced_len = 0usize; + validate_index_ranges(child.len(), starts, lengths, output_len)?; + for (&start, &length) in starts.iter().zip_eq(lengths) { let start = index_value_to_usize("start", start)?; let length = index_value_to_usize("length", length)?; - let end = checked_range_end(start, length)?; + let end = start + length; child.slice(start..end)?.append_to_builder(builder, ctx)?; - produced_len = produced_len - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; } - Ok(produced_len) + + Ok(()) } fn scalar_at_selected_range( @@ -330,111 +238,19 @@ fn scalar_at_selected_range( match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { - let start = constant_index_value::("start", starts)?; - let length = constant_index_value::("length", lengths)?; - match (start, length) { - (Some(start), Some(length)) => scalar_from_constant_start_and_length_ranges( - array.child(), - starts.len(), - index, - index_value_to_usize("start", start)?, - index_value_to_usize("length", length)?, - ctx, - ), - (Some(start), None) => { - let lengths = lengths.clone().execute::(ctx)?; - let lengths = lengths.as_slice::(); - scalar_from_constant_start_ranges( - array.child(), - index, - index_value_to_usize("start", start)?, - lengths, - ctx, - ) - } - (None, Some(length)) => { - let starts = starts.clone().execute::(ctx)?; - let starts = starts.as_slice::(); - scalar_from_constant_length_ranges( - array.child(), - index, - starts, - index_value_to_usize("length", length)?, - ctx, - ) - } - (None, None) => { - let starts = starts.clone().execute::(ctx)?; - let lengths = lengths.clone().execute::(ctx)?; - let starts = starts.as_slice::(); - let lengths = lengths.as_slice::(); - scalar_from_array_ranges(array.child(), index, starts, lengths, ctx) - } - } + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + scalar_from_array_ranges( + array.child(), + index, + starts.as_slice::(), + lengths.as_slice::(), + ctx, + ) }) }) } -fn scalar_from_constant_start_and_length_ranges( - child: &ArrayRef, - len: usize, - index: usize, - start: usize, - length: usize, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let mut logical_start = 0usize; - for _ in 0..len { - if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { - return Ok(Some(scalar)); - } - logical_start = logical_start - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray logical length overflow"))?; - } - Ok(None) -} - -fn scalar_from_constant_start_ranges( - child: &ArrayRef, - index: usize, - start: usize, - lengths: &[L], - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let mut logical_start = 0usize; - for &length in lengths { - let length = index_value_to_usize("length", length)?; - if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { - return Ok(Some(scalar)); - } - logical_start = logical_start - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray logical length overflow"))?; - } - Ok(None) -} - -fn scalar_from_constant_length_ranges( - child: &ArrayRef, - index: usize, - starts: &[S], - length: usize, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let mut logical_start = 0usize; - for &start in starts { - let start = index_value_to_usize("start", start)?; - if let Some(scalar) = scalar_from_range(child, logical_start, index, start, length, ctx)? { - return Ok(Some(scalar)); - } - logical_start = logical_start - .checked_add(length) - .ok_or_else(|| vortex_err!("TakeSlicesArray logical length overflow"))?; - } - Ok(None) -} - fn scalar_from_array_ranges( child: &ArrayRef, index: usize, @@ -484,9 +300,3 @@ fn scalar_from_range( } Ok(None) } - -fn checked_range_end(start: usize, length: usize) -> VortexResult { - start.checked_add(length).ok_or_else(|| { - vortex_err!("TakeSlicesArray range overflow for start {start} and length {length}") - }) -} diff --git a/vortex-array/src/arrays/varbin/compute/mod.rs b/vortex-array/src/arrays/varbin/compute/mod.rs index 480c07a3031..c828e235fa6 100644 --- a/vortex-array/src/arrays/varbin/compute/mod.rs +++ b/vortex-array/src/arrays/varbin/compute/mod.rs @@ -9,6 +9,7 @@ mod compare; mod filter; mod mask; mod take; +mod take_slices; #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 6a6454d0603..d36fe5e8557 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -27,7 +27,7 @@ use crate::validity::Validity; /// The widened offset type used for a taken `VarBinArray`: offsets are widened to at least 32 bits /// (to avoid overflow) while preserving signedness, so a signed result stays Arrow-compatible. -fn taken_offset_ptype(offsets_ptype: PType) -> PType { +pub(super) fn taken_offset_ptype(offsets_ptype: PType) -> PType { match offsets_ptype { PType::U8 | PType::U16 | PType::U32 => PType::U32, PType::U64 => PType::U64, diff --git a/vortex-array/src/arrays/varbin/compute/take_slices.rs b/vortex-array/src/arrays/varbin/compute/take_slices.rs new file mode 100644 index 00000000000..8530faaf060 --- /dev/null +++ b/vortex-array/src/arrays/varbin/compute/take_slices.rs @@ -0,0 +1,223 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools as _; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBufferMut; +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::PrimitiveArray; +use crate::arrays::VarBin; +use crate::arrays::VarBinArray; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::arrays::take_slices::TakeSlicesKernel; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::index_value_to_usize; +use crate::arrays::take_slices::validate_index_ranges; +use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::compute::take::taken_offset_ptype; +use crate::dtype::DType; +use crate::dtype::IntegerPType; +use crate::dtype::PType; +use crate::executor::ExecutionCtx; +use crate::match_each_unsigned_integer_ptype; +use crate::validity::Validity; + +impl TakeSlicesKernel for VarBin { + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + check_index_arrays(starts, lengths)?; + + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { + take_slices_typed::(array, starts, lengths, output_len, ctx) + }) + }) + .map(Some) + } +} + +fn take_slices_typed( + array: ArrayView<'_, VarBin>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, +{ + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + validate_index_ranges( + array.len(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + + let offsets = array.offsets().clone().execute::(ctx)?; + let out_offset_ptype = taken_offset_ptype(offsets.ptype()); + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + + let result = match offsets.ptype() { + PType::U8 => gather_varbin::( + array.dtype().clone(), + offsets.as_slice::(), + array.bytes().as_slice(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U16 => gather_varbin::( + array.dtype().clone(), + offsets.as_slice::(), + array.bytes().as_slice(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U32 => gather_varbin::( + array.dtype().clone(), + offsets.as_slice::(), + array.bytes().as_slice(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U64 => gather_varbin::( + array.dtype().clone(), + offsets.as_slice::(), + array.bytes().as_slice(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + }?; + + let starts = starts.into_array(); + let lengths = lengths.into_array(); + let validity = array + .validity()? + .take_slices(&starts, &lengths, output_len)?; + + // SAFETY: output offsets are built from valid input offsets, start at zero, are monotonically + // non-decreasing, and the copied data buffer has exactly the referenced byte length. + unsafe { + Ok( + VarBinArray::new_unchecked( + result.offsets, + result.data.freeze(), + result.dtype, + validity, + ) + .into_array(), + ) + } +} + +struct GatheredVarBin { + dtype: DType, + offsets: ArrayRef, + data: ByteBufferMut, +} + +fn gather_varbin( + dtype: DType, + offsets: &[Offset], + data: &[u8], + starts: &[S], + lengths: &[L], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, + Offset: IntegerPType, + NewOffset: IntegerPType, +{ + let mut new_offsets = BufferMut::::with_capacity(output_len + 1); + new_offsets.push(NewOffset::zero()); + let mut output_bytes = 0usize; + + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let end = start + length; + if length == 0 { + continue; + } + + let byte_start = index_value_to_usize("offset", offsets[start])?; + let byte_end = index_value_to_usize("offset", offsets[end])?; + vortex_ensure!( + byte_start <= byte_end && byte_end <= data.len(), + "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", + data.len() + ); + + for &offset in &offsets[start + 1..=end] { + let offset = index_value_to_usize("offset", offset)?; + let relative = offset.checked_sub(byte_start).ok_or_else(|| { + vortex_err!("VarBin offsets are not monotonic at offset {offset}") + })?; + let output_offset = output_bytes + .checked_add(relative) + .ok_or_else(|| vortex_err!("TakeSlicesArray VarBin output byte length overflow"))?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + output_bytes = output_bytes + .checked_add(byte_end - byte_start) + .ok_or_else(|| vortex_err!("TakeSlicesArray VarBin output byte length overflow"))?; + } + + let mut new_data = ByteBufferMut::with_capacity(output_bytes); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let end = start + length; + if length == 0 { + continue; + } + + let byte_start = index_value_to_usize("offset", offsets[start])?; + let byte_end = index_value_to_usize("offset", offsets[end])?; + new_data.extend_from_slice(&data[byte_start..byte_end]); + } + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(out_offset_ptype) + .into_array(); + Ok(GatheredVarBin { + dtype, + offsets, + data: new_data, + }) +} + +fn new_offset_value(value: usize) -> VortexResult { + T::from(value).ok_or_else(|| { + vortex_err!( + "TakeSlicesArray VarBin offset value {value} does not fit in {}", + T::PTYPE + ) + }) +} diff --git a/vortex-array/src/arrays/varbin/vtable/kernel.rs b/vortex-array/src/arrays/varbin/vtable/kernel.rs index 9e80abd1037..c3340142145 100644 --- a/vortex-array/src/arrays/varbin/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbin/vtable/kernel.rs @@ -6,9 +6,11 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::Filter; +use crate::arrays::TakeSlices; use crate::arrays::VarBin; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::binary::Binary; @@ -22,4 +24,9 @@ pub(crate) fn initialize(session: &VortexSession) { kernels.register_execute_parent_kernel(Binary.id(), VarBin, CompareExecuteAdaptor(VarBin)); kernels.register_execute_parent_kernel(Filter.id(), VarBin, FilterExecuteAdaptor(VarBin)); kernels.register_execute_parent_kernel(Dict.id(), VarBin, TakeExecuteAdaptor(VarBin)); + kernels.register_execute_parent_kernel( + TakeSlices.id(), + VarBin, + TakeSlicesExecuteAdaptor(VarBin), + ); } diff --git a/vortex-array/src/arrays/varbinview/compute/mod.rs b/vortex-array/src/arrays/varbinview/compute/mod.rs index e25b6eafb53..87dc3e7a24b 100644 --- a/vortex-array/src/arrays/varbinview/compute/mod.rs +++ b/vortex-array/src/arrays/varbinview/compute/mod.rs @@ -6,6 +6,7 @@ mod mask; pub(crate) mod rules; mod slice; mod take; +mod take_slices; mod zip; #[cfg(test)] diff --git a/vortex-array/src/arrays/varbinview/compute/take_slices.rs b/vortex-array/src/arrays/varbinview/compute/take_slices.rs new file mode 100644 index 00000000000..62cb1ebf3d9 --- /dev/null +++ b/vortex-array/src/arrays/varbinview/compute/take_slices.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use itertools::Itertools as _; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinView; +use crate::arrays::VarBinViewArray; +use crate::arrays::take_slices::TakeSlicesKernel; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::index_value_to_usize; +use crate::arrays::take_slices::validate_index_ranges; +use crate::arrays::varbinview::BinaryView; +use crate::buffer::BufferHandle; +use crate::dtype::IntegerPType; +use crate::executor::ExecutionCtx; +use crate::match_each_unsigned_integer_ptype; + +impl TakeSlicesKernel for VarBinView { + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + check_index_arrays(starts, lengths)?; + + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { + take_slices_typed::(array, starts, lengths, output_len, ctx) + }) + }) + .map(Some) + } +} + +fn take_slices_typed( + array: ArrayView<'_, VarBinView>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + S: IntegerPType, + L: IntegerPType, +{ + let starts = starts.clone().execute::(ctx)?; + let lengths = lengths.clone().execute::(ctx)?; + let views = gather_views( + array.views(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + + let starts = starts.into_array(); + let lengths = lengths.into_array(); + let validity = array + .validity()? + .take_slices(&starts, &lengths, output_len)?; + + // SAFETY: ranges were validated against the source views, and copied views still reference the + // same backing data buffers. + unsafe { + Ok(VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(views.into_byte_buffer()), + Arc::clone(array.data_buffers()), + array.dtype().clone(), + validity, + ) + .into_array()) + } +} + +fn gather_views( + source: &[BinaryView], + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + S: IntegerPType, + L: IntegerPType, +{ + validate_index_ranges(source.len(), starts, lengths, output_len)?; + + let mut views = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_usize("start", start)?; + let length = index_value_to_usize("length", length)?; + let end = start + length; + views.extend_from_slice(&source[start..end]); + } + + Ok(views.freeze()) +} diff --git a/vortex-array/src/arrays/varbinview/vtable/kernel.rs b/vortex-array/src/arrays/varbinview/vtable/kernel.rs index e09b381d590..c924f196104 100644 --- a/vortex-array/src/arrays/varbinview/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbinview/vtable/kernel.rs @@ -5,8 +5,10 @@ use vortex_session::VortexSession; use crate::ArrayVTable; use crate::arrays::Dict; +use crate::arrays::TakeSlices; use crate::arrays::VarBinView; use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; @@ -18,5 +20,10 @@ pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); kernels.register_execute_parent_kernel(Cast.id(), VarBinView, CastExecuteAdaptor(VarBinView)); kernels.register_execute_parent_kernel(Dict.id(), VarBinView, TakeExecuteAdaptor(VarBinView)); + kernels.register_execute_parent_kernel( + TakeSlices.id(), + VarBinView, + TakeSlicesExecuteAdaptor(VarBinView), + ); kernels.register_execute_parent_kernel(Zip.id(), VarBinView, ZipExecuteAdaptor(VarBinView)); } From 5e75825dc0a435c2a03f74180a70b2fbdfc32a94 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 17:32:22 -0400 Subject: [PATCH 20/37] Use reduce rules for TakeSlices wrapper rewrites Signed-off-by: Daniel King --- .../src/arrays/extension/compute/mod.rs | 1 + .../src/arrays/extension/compute/rules.rs | 2 + .../arrays/extension/compute/take_slices.rs | 34 +++++ vortex-array/src/arrays/masked/array.rs | 23 ++++ vortex-array/src/arrays/masked/compute/mod.rs | 1 + .../src/arrays/masked/compute/rules.rs | 2 + .../src/arrays/masked/compute/take_slices.rs | 37 ++++++ .../src/arrays/struct_/compute/rules.rs | 2 + .../src/arrays/struct_/compute/take_slices.rs | 72 ++--------- .../src/arrays/struct_/vtable/kernel.rs | 8 -- vortex-array/src/arrays/take_slices/kernel.rs | 50 +++++++- vortex-array/src/arrays/take_slices/mod.rs | 2 + vortex-array/src/arrays/take_slices/tests.rs | 119 +++++++++++++++++- vortex-array/src/arrays/take_slices/vtable.rs | 6 + .../src/arrays/variant/compute/mod.rs | 1 + .../src/arrays/variant/compute/rules.rs | 2 + .../src/arrays/variant/compute/take_slices.rs | 46 +++++++ vortex-array/src/kernel.rs | 8 +- 18 files changed, 344 insertions(+), 72 deletions(-) create mode 100644 vortex-array/src/arrays/extension/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/masked/compute/take_slices.rs create mode 100644 vortex-array/src/arrays/variant/compute/take_slices.rs diff --git a/vortex-array/src/arrays/extension/compute/mod.rs b/vortex-array/src/arrays/extension/compute/mod.rs index 256f31ac49b..9bed9b0f45e 100644 --- a/vortex-array/src/arrays/extension/compute/mod.rs +++ b/vortex-array/src/arrays/extension/compute/mod.rs @@ -8,6 +8,7 @@ mod mask; pub(crate) mod rules; mod slice; mod take; +mod take_slices; #[cfg(test)] mod test { diff --git a/vortex-array/src/arrays/extension/compute/rules.rs b/vortex-array/src/arrays/extension/compute/rules.rs index 2d02d1ae7a7..1ecb37821f7 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -14,6 +14,7 @@ use crate::arrays::Filter; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; +use crate::arrays::take_slices::TakeSlicesReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; @@ -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(&TakeSlicesReduceAdaptor(Extension)), ]); /// Push filter operations into the storage array of an extension array. diff --git a/vortex-array/src/arrays/extension/compute/take_slices.rs b/vortex-array/src/arrays/extension/compute/take_slices.rs new file mode 100644 index 00000000000..f5a04263c63 --- /dev/null +++ b/vortex-array/src/arrays/extension/compute/take_slices.rs @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Extension; +use crate::arrays::ExtensionArray; +use crate::arrays::TakeSlicesArray; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::take_slices::TakeSlicesReduce; + +impl TakeSlicesReduce for Extension { + fn take_slices( + array: ArrayView<'_, Extension>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ) -> VortexResult> { + let storage = TakeSlicesArray::try_new( + array.storage_array().clone(), + starts.clone(), + lengths.clone(), + output_len, + )? + .into_array(); + + Ok(Some( + ExtensionArray::new(array.ext_dtype().clone(), storage).into_array(), + )) + } +} diff --git a/vortex-array/src/arrays/masked/array.rs b/vortex-array/src/arrays/masked/array.rs index 077446b3f02..6dd1d219b6e 100644 --- a/vortex-array/src/arrays/masked/array.rs +++ b/vortex-array/src/arrays/masked/array.rs @@ -92,4 +92,27 @@ impl Array { ) }) } + + /// Constructs a new `MaskedArray` when the caller has already proven that `child` contains no + /// logical nulls. + /// + /// # Safety + /// + /// The caller must ensure that every child value is valid and that `validity` has `child.len()` + /// rows when it has an array-backed representation. + pub(crate) unsafe fn new_unchecked_child_all_valid( + child: ArrayRef, + validity: Validity, + ) -> VortexResult { + let dtype = child.dtype().as_nullable(); + let len = child.len(); + let validity_slot = validity_to_child(&validity, len); + let data = MaskedData::try_new(len, true, validity)?; + Ok(unsafe { + Array::from_parts_unchecked( + ArrayParts::new(Masked, dtype, len, data) + .with_slots(smallvec![Some(child), validity_slot]), + ) + }) + } } diff --git a/vortex-array/src/arrays/masked/compute/mod.rs b/vortex-array/src/arrays/masked/compute/mod.rs index 404ced01ea3..82d47fa3638 100644 --- a/vortex-array/src/arrays/masked/compute/mod.rs +++ b/vortex-array/src/arrays/masked/compute/mod.rs @@ -6,3 +6,4 @@ mod mask; pub(crate) mod rules; mod slice; mod take; +mod take_slices; diff --git a/vortex-array/src/arrays/masked/compute/rules.rs b/vortex-array/src/arrays/masked/compute/rules.rs index 3accb455c3f..097bf68a494 100644 --- a/vortex-array/src/arrays/masked/compute/rules.rs +++ b/vortex-array/src/arrays/masked/compute/rules.rs @@ -5,6 +5,7 @@ use crate::arrays::Masked; use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; +use crate::arrays::take_slices::TakeSlicesReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; @@ -13,4 +14,5 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&MaskReduceAdaptor(Masked)), ParentRuleSet::lift(&SliceReduceAdaptor(Masked)), ParentRuleSet::lift(&TakeReduceAdaptor(Masked)), + ParentRuleSet::lift(&TakeSlicesReduceAdaptor(Masked)), ]); diff --git a/vortex-array/src/arrays/masked/compute/take_slices.rs b/vortex-array/src/arrays/masked/compute/take_slices.rs new file mode 100644 index 00000000000..f69cfaa27bf --- /dev/null +++ b/vortex-array/src/arrays/masked/compute/take_slices.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Masked; +use crate::arrays::MaskedArray; +use crate::arrays::TakeSlicesArray; +use crate::arrays::masked::MaskedArraySlotsExt; +use crate::arrays::take_slices::TakeSlicesReduce; + +impl TakeSlicesReduce for Masked { + fn take_slices( + array: ArrayView<'_, Masked>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ) -> VortexResult> { + let child = TakeSlicesArray::try_new( + array.child().clone(), + starts.clone(), + lengths.clone(), + output_len, + )? + .into_array(); + let validity = array.validity()?.take_slices(starts, lengths, output_len)?; + + // SAFETY: `MaskedArray` guarantees its child has no logical nulls. Taking slices from that + // child preserves all-valid child values; nulls remain represented solely by `validity`. + unsafe { MaskedArray::new_unchecked_child_all_valid(child, validity) } + .map(IntoArray::into_array) + .map(Some) + } +} diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 01981803239..6e700cad329 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -16,6 +16,7 @@ use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::arrays::slice::SliceReduceAdaptor; use crate::arrays::struct_::StructArrayExt; use crate::arrays::struct_::compute::cast::struct_cast_fields; +use crate::arrays::take_slices::TakeSlicesReduceAdaptor; use crate::builtins::ArrayBuiltins; use crate::matcher::Matcher; use crate::optimizer::rules::ArrayParentReduceRule; @@ -31,6 +32,7 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&MaskReduceAdaptor(Struct)), ParentRuleSet::lift(&SliceReduceAdaptor(Struct)), ParentRuleSet::lift(&TakeReduceAdaptor(Struct)), + ParentRuleSet::lift(&TakeSlicesReduceAdaptor(Struct)), ]); pub(crate) fn struct_cast_reduce_parent( diff --git a/vortex-array/src/arrays/struct_/compute/take_slices.rs b/vortex-array/src/arrays/struct_/compute/take_slices.rs index b3140f07931..4c317a724e5 100644 --- a/vortex-array/src/arrays/struct_/compute/take_slices.rs +++ b/vortex-array/src/arrays/struct_/compute/take_slices.rs @@ -6,79 +6,33 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; -use crate::arrays::PrimitiveArray; use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::TakeSlicesArray; use crate::arrays::struct_::StructArrayExt; -use crate::arrays::take_slices::TakeSlicesKernel; +use crate::arrays::take_slices::TakeSlicesReduce; use crate::arrays::take_slices::check_index_arrays; -use crate::arrays::take_slices::validate_index_ranges; -use crate::dtype::IntegerPType; -use crate::executor::ExecutionCtx; -use crate::match_each_unsigned_integer_ptype; -impl TakeSlicesKernel for Struct { +impl TakeSlicesReduce for Struct { fn take_slices( array: ArrayView<'_, Self>, starts: &ArrayRef, lengths: &ArrayRef, output_len: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult> { check_index_arrays(starts, lengths)?; - match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { - take_slices_typed::(array, starts, lengths, output_len, ctx) + let fields = array + .iter_unmasked_fields() + .map(|field| { + TakeSlicesArray::try_new(field.clone(), starts.clone(), lengths.clone(), output_len) + .map(IntoArray::into_array) }) - }) - .map(Some) - } -} - -fn take_slices_typed( - array: ArrayView<'_, Struct>, - starts: &ArrayRef, - lengths: &ArrayRef, - output_len: usize, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - S: IntegerPType, - L: IntegerPType, -{ - let starts = starts.clone().execute::(ctx)?; - let lengths = lengths.clone().execute::(ctx)?; - validate_index_ranges( - array.len(), - starts.as_slice::(), - lengths.as_slice::(), - output_len, - )?; + .collect::>>()?; + let validity = array.validity()?.take_slices(starts, lengths, output_len)?; - let starts = starts.into_array(); - let lengths = lengths.into_array(); - let fields = array - .iter_unmasked_fields() - .map(|field| { - // SAFETY: the index arrays and declared output length were validated above, and the - // child dtype becomes the outer dtype of this per-field TakeSlices array. - unsafe { - TakeSlicesArray::new_unchecked( - field.clone(), - starts.clone(), - lengths.clone(), - output_len, - ) - .into_array() - } - }) - .collect::>(); - let validity = array - .validity()? - .take_slices(&starts, &lengths, output_len)?; - - StructArray::try_new_with_dtype(fields, array.struct_fields().clone(), output_len, validity) - .map(StructArray::into_array) + StructArray::try_new_with_dtype(fields, array.struct_fields().clone(), output_len, validity) + .map(StructArray::into_array) + .map(Some) + } } diff --git a/vortex-array/src/arrays/struct_/vtable/kernel.rs b/vortex-array/src/arrays/struct_/vtable/kernel.rs index fee4d3484d0..1e58da2d0d0 100644 --- a/vortex-array/src/arrays/struct_/vtable/kernel.rs +++ b/vortex-array/src/arrays/struct_/vtable/kernel.rs @@ -3,10 +3,7 @@ use vortex_session::VortexSession; -use crate::ArrayVTable; use crate::arrays::Struct; -use crate::arrays::TakeSlices; -use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::zip::Zip; @@ -14,10 +11,5 @@ use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); - kernels.register_execute_parent_kernel( - TakeSlices.id(), - Struct, - TakeSlicesExecuteAdaptor(Struct), - ); kernels.register_execute_parent_kernel(Zip.id(), Struct, ZipExecuteAdaptor(Struct)); } diff --git a/vortex-array/src/arrays/take_slices/kernel.rs b/vortex-array/src/arrays/take_slices/kernel.rs index b18f56bc165..d215cb0616f 100644 --- a/vortex-array/src/arrays/take_slices/kernel.rs +++ b/vortex-array/src/arrays/take_slices/kernel.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Execute adaptor for `TakeSlices` parent operations. +//! Reduce and execute adaptors for `TakeSlices` parent operations. use vortex_error::VortexResult; @@ -14,10 +14,54 @@ use crate::arrays::take_slices::TakeSlicesArrayExt; use crate::arrays::take_slices::array::CHILD_SLOT; use crate::kernel::ExecuteParentKernel; use crate::matcher::Matcher; +use crate::optimizer::rules::ArrayParentReduceRule; -/// Execution kernel for child encodings that can materialize a `TakeSlices` parent directly. +/// Metadata-only rewrite for child encodings that can push a `TakeSlices` parent through +/// themselves without reading buffers or executing child arrays. +pub trait TakeSlicesReduce: VTable { + /// Rewrite a contiguous-run gather from `array` to an equivalent array. + /// + /// Implementations must not inspect the values of `starts` or `lengths`; range-value errors + /// remain deferred to execution of the rewritten child `TakeSlicesArray`s. + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ) -> VortexResult>; +} + +/// Adapter that wraps a [`TakeSlicesReduce`] impl as an [`ArrayParentReduceRule`]. +#[derive(Default, Debug)] +pub struct TakeSlicesReduceAdaptor(pub V); + +impl ArrayParentReduceRule for TakeSlicesReduceAdaptor +where + V: TakeSlicesReduce, +{ + type Parent = TakeSlices; + + fn reduce_parent( + &self, + array: ArrayView<'_, V>, + parent: ::Match<'_>, + child_idx: usize, + ) -> VortexResult> { + if child_idx != CHILD_SLOT { + return Ok(None); + } + + ::take_slices(array, parent.starts(), parent.lengths(), parent.len()) + } +} + +/// Execution kernel for child encodings that can handle a `TakeSlices` parent directly. +/// +/// Implementations may either materialize the gathered values or return another equivalent +/// encoding that makes progress, such as a canonical nested array whose child is another +/// `TakeSlicesArray`. The executor will continue executing the returned array as needed. pub trait TakeSlicesKernel: VTable { - /// Gather contiguous runs from `array`. + /// Gather contiguous runs from `array`, or rewrite that gather to an equivalent array. /// /// `starts` and `lengths` are non-nullable unsigned integer arrays of equal length. `output_len` /// is the declared length of the `TakeSlices` parent and must match the sum of selected lengths. diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index 20f260c544d..ebb4ecacdbe 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -15,6 +15,8 @@ pub use array::TakeSlicesArrayExt; use itertools::Itertools as _; pub use kernel::TakeSlicesExecuteAdaptor; pub use kernel::TakeSlicesKernel; +pub use kernel::TakeSlicesReduce; +pub use kernel::TakeSlicesReduceAdaptor; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; diff --git a/vortex-array/src/arrays/take_slices/tests.rs b/vortex-array/src/arrays/take_slices/tests.rs index 4902c212c70..4c39f758673 100644 --- a/vortex-array/src/arrays/take_slices/tests.rs +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -9,25 +9,40 @@ use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; use crate::arrays::DictArray; +use crate::arrays::Extension; +use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; +use crate::arrays::Masked; +use crate::arrays::MaskedArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::TakeSlices; use crate::arrays::TakeSlicesArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; +use crate::arrays::Variant; +use crate::arrays::VariantArray; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::masked::MaskedArraySlotsExt; +use crate::arrays::struct_::StructArrayExt; use crate::arrays::take_slices::TakeSlicesExecuteAdaptor; +use crate::arrays::variant::VariantArrayExt; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; use crate::dtype::Nullability; +use crate::extension::datetime::Date; +use crate::extension::datetime::TimeUnit; use crate::kernel::ExecuteParentKernel; +use crate::scalar::Scalar; use crate::validity::Validity; #[test] @@ -209,7 +224,7 @@ fn primitive_take_slices_execute_parent_handles_only_child_slot() -> VortexResul } #[test] -fn struct_take_slices_executes_generically() -> VortexResult<()> { +fn struct_take_slices_reduce_pushes_into_fields() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = StructArray::try_new( FieldNames::from(["id", "name"]), @@ -233,6 +248,93 @@ fn struct_take_slices_executes_generically() -> VortexResult<()> { Validity::NonNullable, )?; + let reduced = actual.clone().execute::(&mut ctx)?; + let reduced_struct = reduced + .as_opt::() + .ok_or_else(|| vortex_err!("expected TakeSlices reduce to return Struct"))?; + assert!( + reduced_struct + .iter_unmasked_fields() + .all(|field| field.is::()) + ); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn extension_take_slices_reduce_pushes_into_storage() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let ext_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); + let array = ExtensionArray::new( + ext_dtype.clone(), + PrimitiveArray::from_iter(0i32..6).into_array(), + ) + .into_array(); + + let actual = take_slices(&array, &[(3, 2), (0, 2)])?; + let expected = ExtensionArray::new( + ext_dtype, + PrimitiveArray::from_iter([3i32, 4, 0, 1]).into_array(), + ); + + let reduced = actual.clone().execute::(&mut ctx)?; + let reduced_extension = reduced + .as_opt::() + .ok_or_else(|| vortex_err!("expected TakeSlices reduce to return Extension"))?; + assert!(reduced_extension.storage_array().is::()); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn masked_take_slices_reduce_preserves_mask() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = MaskedArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + Validity::from_iter([true, false, true, true, false, true]), + )? + .into_array(); + + let actual = take_slices(&array, &[(3, 2), (0, 2)])?; + let expected = MaskedArray::try_new( + PrimitiveArray::from_iter([3i32, 4, 0, 1]).into_array(), + Validity::from_iter([true, false, true, false]), + )?; + + let reduced = actual.clone().execute::(&mut ctx)?; + let reduced_masked = reduced + .as_opt::() + .ok_or_else(|| vortex_err!("expected TakeSlices reduce to return Masked"))?; + assert!(reduced_masked.child().is::()); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn variant_take_slices_reduce_pushes_into_storage_and_shredded() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = VariantArray::try_new( + variant_storage(0..6)?, + Some(PrimitiveArray::from_iter(10i32..16).into_array()), + )? + .into_array(); + + let actual = take_slices(&array, &[(3, 2), (0, 2)])?; + let expected = VariantArray::try_new( + variant_storage([3, 4, 0, 1])?, + Some(PrimitiveArray::from_iter([13i32, 14, 10, 11]).into_array()), + )?; + + let reduced = actual.clone().execute::(&mut ctx)?; + let reduced_variant = reduced + .as_opt::() + .ok_or_else(|| vortex_err!("expected TakeSlices reduce to return Variant"))?; + assert!(reduced_variant.core_storage().is::()); + assert!( + reduced_variant + .shredded() + .is_some_and(|shredded| shredded.is::()) + ); assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } @@ -497,3 +599,18 @@ fn take_slices(array: &ArrayRef, runs: &[(usize, usize)]) -> VortexResult) -> VortexResult { + let chunks = values + .into_iter() + .map(|value| { + ConstantArray::new( + Scalar::variant(Scalar::primitive(value, Nullability::NonNullable)), + 1, + ) + .into_array() + }) + .collect(); + + Ok(ChunkedArray::try_new(chunks, DType::Variant(Nullability::NonNullable))?.into_array()) +} diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index cc529a7f501..714ac1d3513 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -14,6 +14,7 @@ use vortex_session::registry::CachedId; use crate::ArrayParts; use crate::ArrayRef; use crate::EmptyArrayData; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; @@ -152,6 +153,11 @@ impl VTable for TakeSlices { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let parent = array.clone().into_array(); + if let Some(reduced) = array.child().reduce_parent(&parent, CHILD_SLOT)? { + return Ok(ExecutionResult::done(reduced)); + } + let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); append_selected_ranges(array.as_view(), builder.as_mut(), ctx)?; Ok(ExecutionResult::done(builder.finish())) diff --git a/vortex-array/src/arrays/variant/compute/mod.rs b/vortex-array/src/arrays/variant/compute/mod.rs index 5da8b60b65d..9baaa4b8a6c 100644 --- a/vortex-array/src/arrays/variant/compute/mod.rs +++ b/vortex-array/src/arrays/variant/compute/mod.rs @@ -5,3 +5,4 @@ mod filter; pub(crate) mod rules; mod slice; mod take; +mod take_slices; diff --git a/vortex-array/src/arrays/variant/compute/rules.rs b/vortex-array/src/arrays/variant/compute/rules.rs index 0a4f26b6bf8..ecbfb6244db 100644 --- a/vortex-array/src/arrays/variant/compute/rules.rs +++ b/vortex-array/src/arrays/variant/compute/rules.rs @@ -5,10 +5,12 @@ use crate::arrays::Variant; use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; +use crate::arrays::take_slices::TakeSlicesReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; pub(crate) const RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&SliceReduceAdaptor(Variant)), ParentRuleSet::lift(&FilterReduceAdaptor(Variant)), ParentRuleSet::lift(&TakeReduceAdaptor(Variant)), + ParentRuleSet::lift(&TakeSlicesReduceAdaptor(Variant)), ]); diff --git a/vortex-array/src/arrays/variant/compute/take_slices.rs b/vortex-array/src/arrays/variant/compute/take_slices.rs new file mode 100644 index 00000000000..ea5f2eca9b8 --- /dev/null +++ b/vortex-array/src/arrays/variant/compute/take_slices.rs @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::TakeSlicesArray; +use crate::arrays::Variant; +use crate::arrays::VariantArray; +use crate::arrays::take_slices::TakeSlicesReduce; +use crate::arrays::variant::VariantArrayExt; + +impl TakeSlicesReduce for Variant { + fn take_slices( + array: ArrayView<'_, Variant>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ) -> VortexResult> { + let core_storage = TakeSlicesArray::try_new( + array.core_storage().clone(), + starts.clone(), + lengths.clone(), + output_len, + )? + .into_array(); + let shredded = array + .shredded() + .map(|shredded| { + TakeSlicesArray::try_new( + shredded.clone(), + starts.clone(), + lengths.clone(), + output_len, + ) + .map(IntoArray::into_array) + }) + .transpose()?; + + Ok(Some( + VariantArray::try_new(core_storage, shredded)?.into_array(), + )) + } +} diff --git a/vortex-array/src/kernel.rs b/vortex-array/src/kernel.rs index dff9769f139..04d0f356a19 100644 --- a/vortex-array/src/kernel.rs +++ b/vortex-array/src/kernel.rs @@ -31,13 +31,19 @@ use crate::matcher::Matcher; /// /// Unlike reduce rules, parent kernels may read buffers and perform real computation. /// +/// The returned array must be semantically equivalent to the parent array and must have the same +/// logical length and dtype. It does not have to be recursively canonical or fully materialized: +/// the executor optimizes the returned array and continues execution toward the caller's requested +/// target. If a caller needs all nested children fully executed, it should request +/// `RecursiveCanonical`. +/// /// Return `Ok(None)` to decline handling (the scheduler will try the next kernel or fall /// through to the encoding's own `execute`). pub trait ExecuteParentKernel: Debug + Send + Sync + 'static { /// The parent array type this kernel handles. type Parent: Matcher; - /// Attempt to execute the parent array fused with the child array. + /// Attempt to rewrite or execute the parent array fused with the child array. fn execute_parent( &self, array: ArrayView<'_, V>, From 93abdd1903295aa46845fd9809cc38aa3a8f9fa5 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Sat, 11 Jul 2026 01:50:20 +0100 Subject: [PATCH 21/37] reduce Signed-off-by: Robert Kruszewski --- vortex-array/benches/take_fsl.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index de236d26716..c11bf70964b 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -48,8 +48,9 @@ static SESSION: LazyLock = LazyLock::new(array_session); /// Number of lists in the source array. const NUM_LISTS: usize = 500; -/// Number of indices to take. -const NUM_INDICES: &[usize] = &[100, 1_000]; +/// Number of indices to take. This keeps even the widest, longest cases below one millisecond in +/// CodSpeed's instruction-count simulation. +const NUM_INDICES: &[usize] = &[10]; /// Fixed size list lengths (elements per list). const LIST_SIZES: &[usize] = &[16, 64, 128, 256, 512, 1024, 2048, 4096]; From 52504fb32924cf1bde3886ddd34ba34a9d6926a4 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 16:52:00 -0400 Subject: [PATCH 22/37] Use slot macro for TakeSlices slots Signed-off-by: Daniel King --- vortex-array/src/arrays/take_slices/array.rs | 35 ++++++++----------- vortex-array/src/arrays/take_slices/kernel.rs | 6 ++-- vortex-array/src/arrays/take_slices/vtable.rs | 30 ++++++++-------- 3 files changed, 33 insertions(+), 38 deletions(-) diff --git a/vortex-array/src/arrays/take_slices/array.rs b/vortex-array/src/arrays/take_slices/array.rs index b82d893fe67..23620626462 100644 --- a/vortex-array/src/arrays/take_slices/array.rs +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use smallvec::smallvec; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; @@ -10,41 +9,37 @@ use crate::array::Array; use crate::array::ArrayParts; use crate::array::EmptyArrayData; use crate::array::TypedArrayRef; +use crate::array_slots; use crate::arrays::TakeSlices; -/// The child array selected by the run sequence. -pub(super) const CHILD_SLOT: usize = 0; -/// The start index for each child run. -pub(super) const STARTS_SLOT: usize = 1; -/// The length for each child run. -pub(super) const LENGTHS_SLOT: usize = 2; -pub(super) const NUM_SLOTS: usize = 3; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "starts", "lengths"]; +#[array_slots(TakeSlices)] +pub struct TakeSlicesSlots { + /// The child array selected by the run sequence. + pub child: ArrayRef, + /// The start index for each child run. + pub starts: ArrayRef, + /// The length for each child run. + pub lengths: ArrayRef, +} /// Extension methods for [`TakeSlices`] arrays. -pub trait TakeSlicesArrayExt: TypedArrayRef { +pub trait TakeSlicesArrayExt: TypedArrayRef + TakeSlicesArraySlotsExt { /// The child array selected by this run sequence. fn child(&self) -> &ArrayRef { - self.as_ref().slots()[CHILD_SLOT] - .as_ref() - .vortex_expect("validated take-slices child slot") + TakeSlicesArraySlotsExt::child(self) } /// The start index for each child run. fn starts(&self) -> &ArrayRef { - self.as_ref().slots()[STARTS_SLOT] - .as_ref() - .vortex_expect("validated take-slices starts slot") + TakeSlicesArraySlotsExt::starts(self) } /// The length for each child run. fn lengths(&self) -> &ArrayRef { - self.as_ref().slots()[LENGTHS_SLOT] - .as_ref() - .vortex_expect("validated take-slices lengths slot") + TakeSlicesArraySlotsExt::lengths(self) } } -impl> TakeSlicesArrayExt for T {} +impl + TakeSlicesArraySlotsExt> TakeSlicesArrayExt for T {} impl Array { /// Constructs a new `TakeSlicesArray` from start/length arrays and caller-provided output length. diff --git a/vortex-array/src/arrays/take_slices/kernel.rs b/vortex-array/src/arrays/take_slices/kernel.rs index d215cb0616f..be042c2ebea 100644 --- a/vortex-array/src/arrays/take_slices/kernel.rs +++ b/vortex-array/src/arrays/take_slices/kernel.rs @@ -11,7 +11,7 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::TakeSlices; use crate::arrays::take_slices::TakeSlicesArrayExt; -use crate::arrays::take_slices::array::CHILD_SLOT; +use crate::arrays::take_slices::array::TakeSlicesSlots; use crate::kernel::ExecuteParentKernel; use crate::matcher::Matcher; use crate::optimizer::rules::ArrayParentReduceRule; @@ -47,7 +47,7 @@ where parent: ::Match<'_>, child_idx: usize, ) -> VortexResult> { - if child_idx != CHILD_SLOT { + if child_idx != TakeSlicesSlots::CHILD { return Ok(None); } @@ -91,7 +91,7 @@ where child_idx: usize, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if child_idx != CHILD_SLOT { + if child_idx != TakeSlicesSlots::CHILD { return Ok(None); } diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index 714ac1d3513..0446051f74f 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -24,11 +24,7 @@ use crate::array::ValidityVTable; use crate::array::with_empty_buffers; use crate::arrays::PrimitiveArray; use crate::arrays::take_slices::TakeSlicesArrayExt; -use crate::arrays::take_slices::array::CHILD_SLOT; -use crate::arrays::take_slices::array::LENGTHS_SLOT; -use crate::arrays::take_slices::array::NUM_SLOTS; -use crate::arrays::take_slices::array::SLOT_NAMES; -use crate::arrays::take_slices::array::STARTS_SLOT; +use crate::arrays::take_slices::array::TakeSlicesSlots; use crate::arrays::take_slices::check_index_arrays; use crate::arrays::take_slices::checked_range_end; use crate::arrays::take_slices::index_value_to_usize; @@ -74,23 +70,24 @@ impl VTable for TakeSlices { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "TakeSlicesArray expected {NUM_SLOTS} slots, found {}", + slots.len() == TakeSlicesSlots::NAMES.len(), + "TakeSlicesArray expected {} slots, found {}", + TakeSlicesSlots::NAMES.len(), slots.len() ); vortex_ensure!( - slots[CHILD_SLOT].is_some(), + slots[TakeSlicesSlots::CHILD].is_some(), "TakeSlicesArray child slot must be present" ); vortex_ensure!( - slots[STARTS_SLOT].is_some(), + slots[TakeSlicesSlots::STARTS].is_some(), "TakeSlicesArray starts slot must be present" ); vortex_ensure!( - slots[LENGTHS_SLOT].is_some(), + slots[TakeSlicesSlots::LENGTHS].is_some(), "TakeSlicesArray lengths slot must be present" ); - let child = slots[CHILD_SLOT] + let child = slots[TakeSlicesSlots::CHILD] .as_ref() .vortex_expect("validated child slot"); vortex_ensure!( @@ -99,10 +96,10 @@ impl VTable for TakeSlices { child.dtype(), dtype ); - let starts = slots[STARTS_SLOT] + let starts = slots[TakeSlicesSlots::STARTS] .as_ref() .vortex_expect("validated starts slot"); - let lengths = slots[LENGTHS_SLOT] + let lengths = slots[TakeSlicesSlots::LENGTHS] .as_ref() .vortex_expect("validated lengths slot"); check_index_arrays(starts, lengths)?; @@ -130,7 +127,7 @@ impl VTable for TakeSlices { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + TakeSlicesSlots::NAMES[idx].to_string() } fn serialize( @@ -154,7 +151,10 @@ impl VTable for TakeSlices { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let parent = array.clone().into_array(); - if let Some(reduced) = array.child().reduce_parent(&parent, CHILD_SLOT)? { + if let Some(reduced) = array + .child() + .reduce_parent(&parent, TakeSlicesSlots::CHILD)? + { return Ok(ExecutionResult::done(reduced)); } From c6fa44e2c7082099a450a94261ef88524c912cd5 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 16:54:54 -0400 Subject: [PATCH 23/37] Let executor apply reduce_parent before TakeSlices fallback Signed-off-by: Daniel King --- vortex-array/src/arrays/take_slices/vtable.rs | 9 --- vortex-array/src/executor.rs | 58 +++++++++++++++++-- 2 files changed, 54 insertions(+), 13 deletions(-) diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index 0446051f74f..b1e110fdcd6 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -14,7 +14,6 @@ use vortex_session::registry::CachedId; use crate::ArrayParts; use crate::ArrayRef; use crate::EmptyArrayData; -use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; @@ -150,14 +149,6 @@ impl VTable for TakeSlices { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let parent = array.clone().into_array(); - if let Some(reduced) = array - .child() - .reduce_parent(&parent, TakeSlicesSlots::CHILD)? - { - return Ok(ExecutionResult::done(reduced)); - } - let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); append_selected_ranges(array.as_view(), builder.as_mut(), ctx)?; Ok(ExecutionResult::done(builder.finish())) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index a49f98b4d55..5d64d89272f 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -144,11 +144,13 @@ impl ArrayRef { /// Step 2a: if stack.top exists: /// parent = stack.top.parent_array /// child = current_array + /// child.reduce_parent(parent, stack.top.slot_idx), then /// kernels[(parent.encoding_id(), child.encoding_id())] /// .try_execute_parent(child, parent, stack.top.slot_idx) /// /// Step 2b: for child in current_array.children(): /// parent = current_array + /// child.reduce_parent(parent, child.slot_idx), then /// kernels[(parent.encoding_id(), child.encoding_id())] /// .try_execute_parent(child, parent, child.slot_idx) /// @@ -193,16 +195,34 @@ impl ArrayRef { } } - // Step 2a: execute_parent against the suspended parent from ExecuteSlot. + // Step 2a: reduce_parent/execute_parent against the suspended parent from ExecuteSlot. // - // When executing a child for ExecuteSlot, try execute_parent against - // the suspended parent on the stack. This lets kernels like RunEnd's + // When executing a child for ExecuteSlot, try parent rewrites against the suspended + // parent on the stack. This lets rules like Struct get_item or kernels like RunEnd's // FilterKernel fire before the child is forced to canonical. // // Skip when a builder is active: the current array has been partially // consumed by AppendChild (some slots are already in the builder), so // a parent rewrite would see inconsistent state and the builder data // would be lost when we restore frame.parent_builder. + if current_builder.is_none() + && let Some(frame) = stack.last() + && let Some(result) = + current_array.reduce_parent(&frame.parent_array, frame.slot_idx)? + { + ctx.log(format_args!( + "reduce_parent (stack) rewrote {} -> {}", + current_array, result + )); + result + .statistics() + .inherit_from(frame.parent_array.statistics()); + let frame = stack.pop().vortex_expect("just peeked"); + current_array = result.optimize_ctx(ctx.session())?; + current_builder = frame.parent_builder; + continue; + } + if current_builder.is_none() && let Some(frame) = stack.last() && let Some(result) = { @@ -225,7 +245,18 @@ impl ArrayRef { continue; } - // Step 2b: execute_parent against current_array's own children. + // Step 2b: reduce_parent/execute_parent against current_array's own children. + if current_builder.is_none() + && let Some(rewritten) = try_reduce_parent(¤t_array, ctx)? + { + ctx.log(format_args!( + "reduce_parent rewrote {} -> {}", + current_array, rewritten + )); + current_array = rewritten.optimize_ctx(ctx.session())?; + continue; + } + if current_builder.is_none() && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { @@ -590,6 +621,25 @@ fn execute_parent_for_child( Ok(None) } +/// Try reduce_parent on each occupied slot of the array. +fn try_reduce_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + for (slot_idx, slot) in array.slots().iter().enumerate() { + let Some(child) = slot else { continue }; + if let Some(reduced_parent) = child.reduce_parent(array, slot_idx)? { + ctx.log(format_args!( + "reduce_parent: slot[{}]({}) rewrote {} -> {}", + slot_idx, + child.encoding_id(), + array, + reduced_parent + )); + reduced_parent.statistics().inherit_from(array.statistics()); + return Ok(Some(reduced_parent)); + } + } + Ok(None) +} + /// Try execute_parent on each occupied slot of the array. fn try_execute_parent( array: &ArrayRef, From 7a76d3633c0bd5647f0032af93b51bdfcf1d5ce5 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 16:56:02 -0400 Subject: [PATCH 24/37] Avoid duplicate range validation in TakeSlices fallback Signed-off-by: Daniel King --- vortex-array/src/arrays/take_slices/vtable.rs | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index b1e110fdcd6..482acd08b4b 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -27,7 +27,6 @@ use crate::arrays::take_slices::array::TakeSlicesSlots; use crate::arrays::take_slices::check_index_arrays; use crate::arrays::take_slices::checked_range_end; use crate::arrays::take_slices::index_value_to_usize; -use crate::arrays::take_slices::validate_index_ranges; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity_in; @@ -212,14 +211,20 @@ where S: IntegerPType, L: IntegerPType, { - validate_index_ranges(child.len(), starts, lengths, output_len)?; - + let mut produced_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = index_value_to_usize("start", start)?; let length = index_value_to_usize("length", length)?; - let end = start + length; + let end = checked_range_end(start, length)?; + produced_len = produced_len + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; child.slice(start..end)?.append_to_builder(builder, ctx)?; } + vortex_ensure!( + produced_len == output_len, + "TakeSlicesArray produced length {produced_len} does not match declared length {output_len}", + ); Ok(()) } From 666a6aaebfedcb8fc481e009cccf10681ff45b32 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 16:57:52 -0400 Subject: [PATCH 25/37] Handle constant TakeSlices indexes in fallback Signed-off-by: Daniel King --- vortex-array/src/arrays/take_slices/vtable.rs | 134 ++++++++++++++++++ 1 file changed, 134 insertions(+) diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index 482acd08b4b..cdd6278917d 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -21,6 +21,7 @@ use crate::array::OperationsVTable; use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; +use crate::arrays::Constant; use crate::arrays::PrimitiveArray; use crate::arrays::take_slices::TakeSlicesArrayExt; use crate::arrays::take_slices::array::TakeSlicesSlots; @@ -183,6 +184,49 @@ fn append_selected_ranges( let lengths = array.lengths(); check_index_arrays(starts, lengths)?; + let constant_start = constant_index_value(starts)?; + let constant_length = constant_index_value(lengths)?; + match (constant_start, constant_length) { + (Some(start), Some(length)) => { + return append_constant_ranges( + array.child(), + start, + length, + starts.len(), + array.len(), + builder, + ctx, + ); + } + (Some(start), None) => { + return match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { + let lengths = lengths.clone().execute::(ctx)?; + append_constant_start_ranges( + array.child(), + start, + lengths.as_slice::(), + array.len(), + builder, + ctx, + ) + }); + } + (None, Some(length)) => { + return match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { + let starts = starts.clone().execute::(ctx)?; + append_constant_length_ranges( + array.child(), + starts.as_slice::(), + length, + array.len(), + builder, + ctx, + ) + }); + } + (None, None) => {} + } + match_each_unsigned_integer_ptype!(starts.dtype().as_ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.dtype().as_ptype(), |L| { let starts = starts.clone().execute::(ctx)?; @@ -199,6 +243,96 @@ fn append_selected_ranges( }) } +fn constant_index_value(array: &ArrayRef) -> VortexResult> { + array + .as_opt::() + .map(|constant| usize::try_from(constant.scalar())) + .transpose() +} + +fn append_constant_ranges( + child: &ArrayRef, + start: usize, + length: usize, + run_count: usize, + output_len: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let produced_len = run_count + .checked_mul(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; + vortex_ensure!( + produced_len == output_len, + "TakeSlicesArray produced length {produced_len} does not match declared length {output_len}", + ); + + let end = checked_range_end(start, length)?; + let slice = child.slice(start..end)?; + for _ in 0..run_count { + slice.append_to_builder(builder, ctx)?; + } + + Ok(()) +} + +fn append_constant_start_ranges( + child: &ArrayRef, + start: usize, + lengths: &[L], + output_len: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + L: IntegerPType, +{ + let mut produced_len = 0usize; + for &length in lengths { + let length = index_value_to_usize("length", length)?; + let end = checked_range_end(start, length)?; + produced_len = produced_len + .checked_add(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; + child.slice(start..end)?.append_to_builder(builder, ctx)?; + } + vortex_ensure!( + produced_len == output_len, + "TakeSlicesArray produced length {produced_len} does not match declared length {output_len}", + ); + + Ok(()) +} + +fn append_constant_length_ranges( + child: &ArrayRef, + starts: &[S], + length: usize, + output_len: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + S: IntegerPType, +{ + let produced_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("TakeSlicesArray produced length overflow"))?; + vortex_ensure!( + produced_len == output_len, + "TakeSlicesArray produced length {produced_len} does not match declared length {output_len}", + ); + + for &start in starts { + let start = index_value_to_usize("start", start)?; + let end = checked_range_end(start, length)?; + child.slice(start..end)?.append_to_builder(builder, ctx)?; + } + + Ok(()) +} + fn append_array_ranges( child: &ArrayRef, starts: &[S], From e54b90388f3d1a80b07a8a8e2d96ff8a82ac632d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 16:58:44 -0400 Subject: [PATCH 26/37] Use start-relative slices in VarBinView gather Signed-off-by: Daniel King --- vortex-array/src/arrays/varbinview/compute/take_slices.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/varbinview/compute/take_slices.rs b/vortex-array/src/arrays/varbinview/compute/take_slices.rs index 62cb1ebf3d9..8b7fd5800f5 100644 --- a/vortex-array/src/arrays/varbinview/compute/take_slices.rs +++ b/vortex-array/src/arrays/varbinview/compute/take_slices.rs @@ -98,8 +98,7 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start = index_value_to_usize("start", start)?; let length = index_value_to_usize("length", length)?; - let end = start + length; - views.extend_from_slice(&source[start..end]); + views.extend_from_slice(&source[start..][..length]); } Ok(views.freeze()) From 6161036a06a41c5c972fa6c012e39e820b0a21fa Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 17:00:05 -0400 Subject: [PATCH 27/37] Use unchecked Decimal TakeSlices copy after validation Signed-off-by: Daniel King --- .../src/arrays/decimal/compute/take_slices.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/arrays/decimal/compute/take_slices.rs b/vortex-array/src/arrays/decimal/compute/take_slices.rs index 2e47f109e70..30ffae647ba 100644 --- a/vortex-array/src/arrays/decimal/compute/take_slices.rs +++ b/vortex-array/src/arrays/decimal/compute/take_slices.rs @@ -119,8 +119,25 @@ where let start = index_value_to_usize("start", start)?; let length = index_value_to_usize("length", length)?; let end = start + length; - values.extend_from_slice(&source[start..end]); + // SAFETY: `validate_index_ranges` checked every source range and the total output length, + // so `source[start..end]` is in bounds and `values` has enough spare capacity. + unsafe { extend_from_slice_unchecked(&mut values, source.get_unchecked(start..end)) }; } Ok(values.freeze()) } + +unsafe fn extend_from_slice_unchecked(values: &mut BufferMut, source: &[T]) +where + T: NativeDecimalType, +{ + let old_len = values.len(); + let new_len = old_len + source.len(); + + // SAFETY: the caller guarantees `values` has enough spare capacity for `source.len()` values. + unsafe { + let dst = values.spare_capacity_mut().as_mut_ptr().cast::(); + std::ptr::copy_nonoverlapping(source.as_ptr(), dst, source.len()); + values.set_len(new_len); + } +} From 12cd0af9213b8de00620e99a5e2377f4791cb438 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 17:01:29 -0400 Subject: [PATCH 28/37] Use execute_mask for FSL take index validity Signed-off-by: Daniel King --- .../src/arrays/fixed_size_list/compute/take.rs | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 49dcf1bfdab..196b5271ab7 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -12,14 +12,12 @@ use crate::ArrayRef; use crate::IntoArray; use crate::RecursiveCanonical; use crate::array::ArrayView; -use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::TakeSlicesArray; -use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; @@ -247,12 +245,9 @@ fn indices_validity_mask( indices_array: &ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let indices_len = indices_array.as_ref().len(); - match indices_array.validity()? { - Validity::NonNullable | Validity::AllValid => Ok(Mask::new_true(indices_len)), - Validity::AllInvalid => Ok(Mask::new_false(indices_len)), - Validity::Array(a) => Ok(a.execute::(ctx)?.execute_mask(ctx)), - } + indices_array + .validity()? + .execute_mask(indices_array.as_ref().len(), ctx) } fn validate_valid_indices( From ed4f61ca1c745406aa61e4ad92eed8bd4beb95f3 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 17:02:35 -0400 Subject: [PATCH 29/37] Build FSL TakeSlices starts as u64 Signed-off-by: Daniel King --- .../src/arrays/fixed_size_list/compute/take.rs | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 196b5271ab7..82b9f5ae4dc 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -155,7 +155,7 @@ fn take_non_nullable_non_empty_fsl( for &data_idx in indices { let data_idx = index_to_usize(data_idx)?; - let start = list_start(data_idx, list_size, array_len)?; + let start = list_start_u64(data_idx, list_size, array_len)?; starts.push(start); } @@ -210,7 +210,7 @@ fn take_nullable_non_empty_fsl( } let data_idx = index_to_usize(data_idx)?; - let start = list_start(data_idx, list_size, array_len)?; + let start = list_start_u64(data_idx, list_size, array_len)?; if !array_validity.value(data_idx) { starts.push(0); new_validity_builder.append(false); @@ -282,7 +282,7 @@ fn ensure_elements_len(actual: usize, expected: usize) -> VortexResult<()> { Ok(()) } -fn list_start(data_idx: usize, list_size: usize, array_len: usize) -> VortexResult { +fn list_start_u64(data_idx: usize, list_size: usize, array_len: usize) -> VortexResult { check_index_in_bounds(data_idx, array_len)?; let start = data_idx.checked_mul(list_size).ok_or_else(|| { @@ -294,8 +294,8 @@ fn list_start(data_idx: usize, list_size: usize, array_len: usize) -> VortexResu vortex_err!( "FixedSizeList take element range overflow for index {data_idx} and list size {list_size}" ) - })?; - Ok(start) + })?; + Ok(start as u64) } fn check_index_in_bounds(data_idx: usize, array_len: usize) -> VortexResult<()> { @@ -313,16 +313,12 @@ fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef fn take_element_runs( elements: &ArrayRef, - starts: Vec, + starts: Vec, length: usize, output_len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { let run_count = starts.len(); - let starts = starts - .into_iter() - .map(|start| start as u64) - .collect::>(); let starts = PrimitiveArray::from_iter(starts).into_array(); let lengths = ConstantArray::new(length as u64, run_count).into_array(); From 43ec9d51cbf8bf5e9bb13df6bf103e9adee23a5d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 17:05:10 -0400 Subject: [PATCH 30/37] Use Buffer for FSL TakeSlices starts Signed-off-by: Daniel King --- .../src/arrays/fixed_size_list/compute/take.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 82b9f5ae4dc..d2e56bdeb57 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::BitBufferMut; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -151,7 +153,7 @@ fn take_non_nullable_non_empty_fsl( let indices: &[I] = indices_array.as_slice::(); let new_len = indices.len(); let expected_elements_len = take_elements_len(new_len, list_size)?; - let mut starts = Vec::with_capacity(new_len); + let mut starts = BufferMut::::with_capacity(new_len); for &data_idx in indices { let data_idx = index_to_usize(data_idx)?; @@ -161,7 +163,7 @@ fn take_non_nullable_non_empty_fsl( let new_elements = take_element_runs( array.elements(), - starts, + starts.freeze(), list_size, expected_elements_len, ctx, @@ -197,7 +199,7 @@ fn take_nullable_non_empty_fsl( .execute_mask(array.as_ref().len(), ctx)?; let indices_validity = indices_validity_mask(&indices_array, ctx)?; - let mut starts = Vec::with_capacity(new_len); + let mut starts = BufferMut::::with_capacity(new_len); let mut new_validity_builder = BitBufferMut::with_capacity(new_len); // Null output rows still need placeholder child elements so the FSL elements length stays @@ -223,7 +225,7 @@ fn take_nullable_non_empty_fsl( let new_elements = take_element_runs( array.elements(), - starts, + starts.freeze(), list_size, expected_elements_len, ctx, @@ -313,13 +315,13 @@ fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef fn take_element_runs( elements: &ArrayRef, - starts: Vec, + starts: Buffer, length: usize, output_len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { let run_count = starts.len(); - let starts = PrimitiveArray::from_iter(starts).into_array(); + let starts = PrimitiveArray::new(starts, Validity::NonNullable).into_array(); let lengths = ConstantArray::new(length as u64, run_count).into_array(); // SAFETY: callers produced one start per output row after validating list indices against the From 28cc68059517d2c2ceefeecd4c26a9d1f18e4663 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 17:07:07 -0400 Subject: [PATCH 31/37] Keep FSL TakeSlices elements lazy Signed-off-by: Daniel King --- .../src/arrays/fixed_size_list/compute/take.rs | 12 ++---------- .../src/arrays/fixed_size_list/tests/take.rs | 6 +++--- 2 files changed, 5 insertions(+), 13 deletions(-) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index d2e56bdeb57..6760f4ad58d 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -12,7 +12,6 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; -use crate::RecursiveCanonical; use crate::array::ArrayView; use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; @@ -33,7 +32,7 @@ use crate::validity::Validity; /// /// `FixedSizeListArray` must rebuild its elements array because selected lists need to become /// packed from offset 0. The FSL layer translates selected list rows into ordered element runs -/// and delegates the execution strategy to the elements child via materialized `take_slices`. +/// and delegates the execution strategy to the elements child via `take_slices`. impl TakeExecute for FixedSizeList { fn take( array: ArrayView<'_, FixedSizeList>, @@ -139,14 +138,13 @@ fn take_non_empty_non_degenerate_fsl( if array.dtype().is_nullable() || indices_array.dtype().is_nullable() { take_nullable_non_empty_fsl::(array, indices_array, ctx) } else { - take_non_nullable_non_empty_fsl::(array, indices_array, ctx) + take_non_nullable_non_empty_fsl::(array, indices_array) } } fn take_non_nullable_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, - ctx: &mut ExecutionCtx, ) -> VortexResult { let list_size = array.list_size() as usize; let array_len = array.as_ref().len(); @@ -166,7 +164,6 @@ fn take_non_nullable_non_empty_fsl( starts.freeze(), list_size, expected_elements_len, - ctx, )?; ensure_elements_len(new_elements.len(), expected_elements_len)?; @@ -228,7 +225,6 @@ fn take_nullable_non_empty_fsl( starts.freeze(), list_size, expected_elements_len, - ctx, )?; ensure_elements_len(new_elements.len(), expected_elements_len)?; @@ -318,7 +314,6 @@ fn take_element_runs( starts: Buffer, length: usize, output_len: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult { let run_count = starts.len(); let starts = PrimitiveArray::new(starts, Validity::NonNullable).into_array(); @@ -329,9 +324,6 @@ fn take_element_runs( // constant array, and `output_len` was computed as `run_count * length`. Ok( unsafe { TakeSlicesArray::new_unchecked(elements.clone(), starts, lengths, output_len) } - .into_array() - .execute::(ctx)? - .0 .into_array(), ) } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index ba2aa7f89bf..ea3620d8914 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -32,8 +32,8 @@ use crate::array::with_empty_buffers; use crate::array_session; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; -use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; +use crate::arrays::TakeSlices; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::assert_arrays_eq; @@ -398,7 +398,7 @@ fn test_take_empty_source_rejects_valid_index_after_null_index() -> VortexResult } #[test] -fn test_take_materializes_encoded_elements_child() -> VortexResult<()> { +fn test_take_preserves_lazy_elements_child() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let encoded_elements = NoTakeSlicesArray::wrap(buffer![10i32, 20, 30, 40, 20, 30].into_array()); let fsl = FixedSizeListArray::new(encoded_elements, 2, Validity::NonNullable, 3); @@ -407,7 +407,7 @@ fn test_take_materializes_encoded_elements_child() -> VortexResult<()> { let result = ::take(fsl.as_view(), &indices, &mut ctx)? .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; - assert!(result.as_::().elements().is::()); + assert!(result.as_::().elements().is::()); let expected = FixedSizeListArray::new( PrimitiveArray::from_iter([20i32, 30, 10, 20]).into_array(), 2, From feb016b79880fdbeca74d58cf547f1eda5f16cb9 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 17:09:04 -0400 Subject: [PATCH 32/37] Handle empty FSL takes generically Signed-off-by: Daniel King --- .../arrays/fixed_size_list/compute/take.rs | 52 +++++++++----- .../src/arrays/fixed_size_list/tests/take.rs | 67 ++++++++++++++----- 2 files changed, 87 insertions(+), 32 deletions(-) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 6760f4ad58d..4b8e9020e1f 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -39,18 +39,25 @@ impl TakeExecute for FixedSizeList { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - match_each_unsigned_integer_ptype!(indices.dtype().as_ptype().to_unsigned(), |I| { - take_with_indices::(array, indices, ctx) + if array.is_empty() { + return take_empty_fsl(array, indices, ctx).map(Some); + } + + let indices_ptype = indices.dtype().as_ptype().to_unsigned(); + match_each_unsigned_integer_ptype!(indices_ptype, |I| { + take_non_empty_with_indices::(array, indices, ctx) }) .map(Some) } } -fn take_with_indices( +fn take_non_empty_with_indices( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { + debug_assert!(!array.is_empty()); + let list_size = array.list_size() as usize; let indices_array = indices.clone().execute::(ctx)?; @@ -59,22 +66,20 @@ fn take_with_indices( let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned()); if list_size == 0 { - return take_degenerate_fsl::(array, indices, indices_array.as_view(), ctx); - } - - if array.is_empty() { - return take_empty_non_degenerate_fsl::(array, indices, indices_array.as_view(), ctx); + return take_non_empty_degenerate_fsl::(array, indices, indices_array.as_view(), ctx); } take_non_empty_non_degenerate_fsl::(array, indices_array.as_view(), ctx) } -fn take_degenerate_fsl( +fn take_non_empty_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { + debug_assert!(!array.is_empty()); + debug_assert_eq!(array.list_size(), 0); vortex_ensure!( array.elements().is_empty(), "degenerate list must have empty elements" @@ -97,19 +102,23 @@ fn take_degenerate_fsl( .into_array()) } -fn take_empty_non_degenerate_fsl( +fn take_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, - indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { - debug_assert_ne!(array.list_size(), 0); debug_assert!(array.is_empty()); - validate_valid_indices::(&indices_array, 0, ctx)?; + let new_len = indices.len(); + if new_len != 0 { + let indices_validity = indices.validity()?.execute_mask(new_len, ctx)?; + vortex_ensure!( + indices_validity.all_false(), + "cannot take valid indices from an empty FixedSizeList" + ); + } let list_size = array.list_size() as usize; - let new_len = indices_array.len(); let expected_elements_len = take_elements_len(new_len, list_size)?; let new_elements = default_elements(array, expected_elements_len); ensure_elements_len(new_elements.len(), expected_elements_len)?; @@ -119,8 +128,8 @@ fn take_empty_non_degenerate_fsl( Validity::AllInvalid }; - // SAFETY: an empty non-degenerate source has no usable element slices. Valid indices have - // already been rejected, so non-empty output is all null and backed by default child values. + // SAFETY: empty output needs no child values; otherwise the index validity mask proves every + // output row is null. Placeholder child elements have the exact length required by FSL. Ok(unsafe { FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } @@ -132,8 +141,8 @@ fn take_non_empty_non_degenerate_fsl( indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { - debug_assert_ne!(array.list_size(), 0); debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); if array.dtype().is_nullable() || indices_array.dtype().is_nullable() { take_nullable_non_empty_fsl::(array, indices_array, ctx) @@ -146,6 +155,9 @@ fn take_non_nullable_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, ) -> VortexResult { + debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); + let list_size = array.list_size() as usize; let array_len = array.as_ref().len(); let indices: &[I] = indices_array.as_slice::(); @@ -185,6 +197,9 @@ fn take_nullable_non_empty_fsl( indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { + debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); + let list_size = array.list_size() as usize; let array_len = array.as_ref().len(); let indices: &[I] = indices_array.as_slice::(); @@ -281,6 +296,9 @@ fn ensure_elements_len(actual: usize, expected: usize) -> VortexResult<()> { } fn list_start_u64(data_idx: usize, list_size: usize, array_len: usize) -> VortexResult { + debug_assert_ne!(array_len, 0); + debug_assert_ne!(list_size, 0); + check_index_in_bounds(data_idx, array_len)?; let start = data_idx.checked_mul(list_size).ok_or_else(|| { diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index ea3620d8914..b4c7e57cf14 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -347,43 +347,80 @@ fn test_take_nullable_fsl_with_nullable_indices() { } #[test] -fn test_take_empty_source_with_all_null_indices() { +fn test_take_empty_source_with_all_null_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let fsl = create_empty_fsl(); - let indices = PrimitiveArray::new(buffer![999u64, 123], Validity::AllInvalid); + let indices = PrimitiveArray::new(buffer![999u64, 123], Validity::AllInvalid).into_array(); - let result = fsl.take(indices.into_array()).unwrap(); + let result = ::take(fsl.as_view(), &indices, &mut ctx)? + .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; + let result_fsl = result.as_::(); assert_eq!(result.len(), 2); + assert_eq!(result.dtype().nullability(), Nullability::Nullable); + assert_eq!( + result_fsl.elements().len(), + result.len() * result_fsl.list_size() as usize + ); for idx in 0..result.len() { - assert!( - result - .execute_scalar(idx, &mut array_session().create_execution_ctx()) - .unwrap() - .is_null() - ); + assert!(result.execute_scalar(idx, &mut ctx)?.is_null()); } + Ok(()) } #[test] -fn test_take_execute_empty_source_all_null_indices_builds_default_elements() -> VortexResult<()> { +fn test_take_empty_degenerate_source_with_all_null_indices() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let fsl = create_empty_fsl(); + let elements = PrimitiveArray::empty::(Nullability::NonNullable).into_array(); + let fsl = FixedSizeListArray::new(elements, 0, Validity::NonNullable, 0); let indices = PrimitiveArray::new(buffer![999u64, 123], Validity::AllInvalid).into_array(); let result = ::take(fsl.as_view(), &indices, &mut ctx)? .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; let result_fsl = result.as_::(); - assert_eq!( - result_fsl.elements().len(), - result.len() * result_fsl.list_size() as usize - ); + assert_eq!(result.len(), 2); + assert_eq!(result_fsl.list_size(), 0); + assert_eq!(result.dtype().nullability(), Nullability::Nullable); + assert_eq!(result_fsl.elements().len(), 0); for idx in 0..result.len() { assert!(result.execute_scalar(idx, &mut ctx)?.is_null()); } Ok(()) } +#[test] +fn test_take_execute_empty_source_accepts_empty_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let fsl = create_empty_fsl(); + let indices = PrimitiveArray::from_iter::<[u64; 0]>([]).into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx)? + .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; + let result_fsl = result.as_::(); + + assert_eq!(result.len(), 0); + assert_eq!(result.dtype().nullability(), Nullability::NonNullable); + assert_eq!(result_fsl.elements().len(), 0); + Ok(()) +} + +#[test] +fn test_take_empty_source_accepts_empty_nullable_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let fsl = create_empty_fsl(); + let indices = PrimitiveArray::empty::(Nullability::Nullable).into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx)? + .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; + let result_fsl = result.as_::(); + + assert_eq!(result.len(), 0); + assert_eq!(result.dtype().nullability(), Nullability::Nullable); + assert_eq!(result_fsl.elements().len(), 0); + Ok(()) +} + #[test] fn test_take_empty_source_rejects_valid_index_after_null_index() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); From 7890cd3610cdc8bd3b16ca5868bbc23bb77f4b68 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 14 Jul 2026 17:09:56 -0400 Subject: [PATCH 33/37] Document TakeSlices selector rationale Signed-off-by: Daniel King --- vortex-array/src/arrays/take_slices/mod.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vortex-array/src/arrays/take_slices/mod.rs b/vortex-array/src/arrays/take_slices/mod.rs index ebb4ecacdbe..91438b6cc00 100644 --- a/vortex-array/src/arrays/take_slices/mod.rs +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -6,6 +6,11 @@ //! `TakeSlicesArray` represents the concatenation of //! `values[starts[i]..starts[i] + lengths[i]]` for each range row. Ranges may overlap, repeat, //! and appear in any order. +//! +//! This is equivalent to a `Take` with concatenated sequence codes, but keeps the selectors +//! proportional to the number of ranges instead of the output length. A run-end encoded code array +//! cannot represent a general range compactly because each output position within a range selects +//! a different child index. mod array; mod kernel; From ed74607ed49e4f5c8c7505a74c29347f5e50c0ee Mon Sep 17 00:00:00 2001 From: Daniel King Date: Wed, 15 Jul 2026 14:44:27 -0400 Subject: [PATCH 34/37] Reject negative signed FSL take indices Signed-off-by: Daniel King --- .../arrays/fixed_size_list/compute/take.rs | 97 +++++++++++++------ .../src/arrays/fixed_size_list/tests/take.rs | 30 ++++++ 2 files changed, 99 insertions(+), 28 deletions(-) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 4b8e9020e1f..72b12e8a2a1 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -23,9 +23,12 @@ use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::builder_with_capacity; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; +use crate::scalar::Scalar; use crate::validity::Validity; /// Take implementation for [`FixedSizeListArray`]. @@ -43,40 +46,86 @@ impl TakeExecute for FixedSizeList { return take_empty_fsl(array, indices, ctx).map(Some); } - let indices_ptype = indices.dtype().as_ptype().to_unsigned(); - match_each_unsigned_integer_ptype!(indices_ptype, |I| { - take_non_empty_with_indices::(array, indices, ctx) - }) - .map(Some) + take_non_empty_fsl(array, indices, ctx).map(Some) } } +fn take_non_empty_fsl( + array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + debug_assert!(!array.is_empty()); + + let DType::Primitive(ptype, nullability) = indices.dtype() else { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + }; + if !ptype.is_int() { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()); + } + + let indices_validity = indices.validity()?; + let indices_validity_mask = indices_validity.execute_mask(indices.len(), ctx)?; + // Null index lanes are semantically ignored. Zero them before checked signed-to-unsigned casts + // so a negative physical payload in a null lane does not fail the take. + let indices_nulls_zeroed = if indices_validity_mask.all_true() { + indices.clone() + } else { + indices + .clone() + .fill_null(Scalar::from(0).cast(indices.dtype())?)? + }; + + let unsigned_indices = if ptype.is_unsigned_int() { + indices_nulls_zeroed.execute::(ctx)? + } else { + indices_nulls_zeroed + .cast(DType::Primitive(ptype.to_unsigned(), *nullability))? + .execute::(ctx)? + }; + + match_each_unsigned_integer_ptype!(unsigned_indices.ptype(), |I| { + take_non_empty_with_indices::( + array, + indices, + unsigned_indices.as_view(), + &indices_validity_mask, + nullability.is_nullable(), + ctx, + ) + }) +} + fn take_non_empty_with_indices( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, + indices_array: ArrayView<'_, Primitive>, + indices_validity: &Mask, + indices_are_nullable: bool, ctx: &mut ExecutionCtx, ) -> VortexResult { debug_assert!(!array.is_empty()); let list_size = array.list_size() as usize; - let indices_array = indices.clone().execute::(ctx)?; - // Bit-identical reinterpret so `as_slice::` (with unsigned `I`) matches. Negative signed - // inputs become large unsigned values and are rejected by the bounds check below. - let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned()); - if list_size == 0 { - return take_non_empty_degenerate_fsl::(array, indices, indices_array.as_view(), ctx); + return take_non_empty_degenerate_fsl::(array, indices, indices_array, indices_validity); } - take_non_empty_non_degenerate_fsl::(array, indices_array.as_view(), ctx) + take_non_empty_non_degenerate_fsl::( + array, + indices_array, + indices_validity, + indices_are_nullable, + ctx, + ) } fn take_non_empty_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, indices_array: ArrayView<'_, Primitive>, - ctx: &mut ExecutionCtx, + indices_validity: &Mask, ) -> VortexResult { debug_assert!(!array.is_empty()); debug_assert_eq!(array.list_size(), 0); @@ -85,7 +134,7 @@ fn take_non_empty_degenerate_fsl( "degenerate list must have empty elements" ); - validate_valid_indices::(&indices_array, array.as_ref().len(), ctx)?; + validate_valid_indices::(&indices_array, indices_validity, array.as_ref().len())?; let new_validity = array.validity()?.take(indices)?; let new_len = indices_array.len(); @@ -139,13 +188,15 @@ fn take_empty_fsl( fn take_non_empty_non_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, + indices_validity: &Mask, + indices_are_nullable: bool, ctx: &mut ExecutionCtx, ) -> VortexResult { debug_assert!(!array.is_empty()); debug_assert_ne!(array.list_size(), 0); - if array.dtype().is_nullable() || indices_array.dtype().is_nullable() { - take_nullable_non_empty_fsl::(array, indices_array, ctx) + if array.dtype().is_nullable() || indices_are_nullable { + take_nullable_non_empty_fsl::(array, indices_array, indices_validity, ctx) } else { take_non_nullable_non_empty_fsl::(array, indices_array) } @@ -195,6 +246,7 @@ fn take_non_nullable_non_empty_fsl( fn take_nullable_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, + indices_validity: &Mask, ctx: &mut ExecutionCtx, ) -> VortexResult { debug_assert!(!array.is_empty()); @@ -209,7 +261,6 @@ fn take_nullable_non_empty_fsl( let array_validity = array .fixed_size_list_validity() .execute_mask(array.as_ref().len(), ctx)?; - let indices_validity = indices_validity_mask(&indices_array, ctx)?; let mut starts = BufferMut::::with_capacity(new_len); let mut new_validity_builder = BitBufferMut::with_capacity(new_len); @@ -254,22 +305,12 @@ fn take_nullable_non_empty_fsl( .into_array()) } -fn indices_validity_mask( - indices_array: &ArrayView<'_, Primitive>, - ctx: &mut ExecutionCtx, -) -> VortexResult { - indices_array - .validity()? - .execute_mask(indices_array.as_ref().len(), ctx) -} - fn validate_valid_indices( indices_array: &ArrayView<'_, Primitive>, + indices_validity: &Mask, array_len: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let indices: &[I] = indices_array.as_slice::(); - let indices_validity = indices_validity_mask(indices_array, ctx)?; for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { if is_index_valid { diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index b4c7e57cf14..047504a4281 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -307,6 +307,36 @@ fn test_take_nullable_indices_ignores_out_of_bounds_null_value() { assert_arrays_eq!(expected, result, &mut ctx); } +#[test] +fn test_take_rejects_negative_signed_index() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = buffer![1i32, 2, 3, 4].into_array(); + let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 2); + let indices = buffer![-1i32].into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_take_nullable_indices_ignores_negative_null_value() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = buffer![1i32, 2, 3, 4].into_array(); + let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 2); + let indices = + PrimitiveArray::new(buffer![-1i32, 1], Validity::from_iter([false, true])).into_array(); + + let result = ::take(fsl.as_view(), &indices, &mut ctx)? + .ok_or_else(|| vortex_err!("FixedSizeList TakeExecute returned no result"))?; + + assert_eq!(result.len(), 2); + assert!(result.execute_scalar(0, &mut ctx)?.is_null()); + assert!(!result.execute_scalar(1, &mut ctx)?.is_null()); + Ok(()) +} + #[test] fn test_take_rejects_overflowing_valid_index() { let mut ctx = array_session().create_execution_ctx(); From c8fe6b0925a8ca2fce8c0e4f32157f57db97d3fb Mon Sep 17 00:00:00 2001 From: Daniel King Date: Wed, 15 Jul 2026 16:20:03 -0400 Subject: [PATCH 35/37] Simplify fixed-size-list take index handling Signed-off-by: Daniel King --- .../arrays/fixed_size_list/compute/take.rs | 333 ++++++------------ 1 file changed, 99 insertions(+), 234 deletions(-) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 72b12e8a2a1..c1fd7b35720 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -1,14 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::BitBufferMut; +use itertools::Itertools; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; @@ -23,12 +22,11 @@ use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::builder_with_capacity; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; +use crate::dtype::Nullability; use crate::executor::ExecutionCtx; -use crate::match_each_unsigned_integer_ptype; -use crate::scalar::Scalar; +use crate::match_each_integer_ptype; use crate::validity::Validity; /// Take implementation for [`FixedSizeListArray`]. @@ -50,82 +48,77 @@ impl TakeExecute for FixedSizeList { } } -fn take_non_empty_fsl( +fn take_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - debug_assert!(!array.is_empty()); + debug_assert!(array.is_empty()); - let DType::Primitive(ptype, nullability) = indices.dtype() else { - vortex_bail!("Invalid indices dtype: {}", indices.dtype()) - }; - if !ptype.is_int() { - vortex_bail!("Invalid indices dtype: {}", indices.dtype()); + let new_len = indices.len(); + if new_len != 0 { + let indices_validity = indices.validity()?.execute_mask(new_len, ctx)?; + vortex_ensure!( + indices_validity.all_false(), + "cannot take valid indices from an empty FixedSizeList" + ); } - let indices_validity = indices.validity()?; - let indices_validity_mask = indices_validity.execute_mask(indices.len(), ctx)?; - // Null index lanes are semantically ignored. Zero them before checked signed-to-unsigned casts - // so a negative physical payload in a null lane does not fail the take. - let indices_nulls_zeroed = if indices_validity_mask.all_true() { - indices.clone() - } else { - indices - .clone() - .fill_null(Scalar::from(0).cast(indices.dtype())?)? - }; - - let unsigned_indices = if ptype.is_unsigned_int() { - indices_nulls_zeroed.execute::(ctx)? + let list_size = array.list_size() as usize; + let elements_len = new_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" + ) + })?; + let new_elements = default_elements(array, elements_len); + let new_validity = if new_len == 0 { + array.validity()?.take(indices)? } else { - indices_nulls_zeroed - .cast(DType::Primitive(ptype.to_unsigned(), *nullability))? - .execute::(ctx)? + Validity::AllInvalid }; - match_each_unsigned_integer_ptype!(unsigned_indices.ptype(), |I| { - take_non_empty_with_indices::( - array, - indices, - unsigned_indices.as_view(), - &indices_validity_mask, - nullability.is_nullable(), - ctx, - ) - }) + // SAFETY: empty output needs no child values; otherwise the index validity mask proves every + // output row is null. Placeholder child elements have the exact length required by FSL. + Ok(unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) + } + .into_array()) } -fn take_non_empty_with_indices( +fn take_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, - indices_array: ArrayView<'_, Primitive>, - indices_validity: &Mask, - indices_are_nullable: bool, ctx: &mut ExecutionCtx, ) -> VortexResult { debug_assert!(!array.is_empty()); - let list_size = array.list_size() as usize; + let DType::Primitive(ptype, nullability) = indices.dtype() else { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + }; + if !ptype.is_int() { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()); + } - if list_size == 0 { - return take_non_empty_degenerate_fsl::(array, indices, indices_array, indices_validity); + if array.list_size() == 0 { + return take_non_empty_degenerate_fsl(array, indices, ctx); } - take_non_empty_non_degenerate_fsl::( - array, - indices_array, - indices_validity, - indices_are_nullable, - ctx, - ) + let indices_array = indices.clone().execute::(ctx)?; + match_each_integer_ptype!(indices_array.ptype(), |I| { + take_non_empty_non_degenerate_fsl::( + array, + indices, + indices_array.as_view(), + *nullability, + ctx, + ) + }) } -fn take_non_empty_degenerate_fsl( +fn take_non_empty_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, - indices_array: ArrayView<'_, Primitive>, - indices_validity: &Mask, + ctx: &mut ExecutionCtx, ) -> VortexResult { debug_assert!(!array.is_empty()); debug_assert_eq!(array.list_size(), 0); @@ -134,7 +127,10 @@ fn take_non_empty_degenerate_fsl( "degenerate list must have empty elements" ); - validate_valid_indices::(&indices_array, indices_validity, array.as_ref().len())?; + let indices_array = indices.clone().execute::(ctx)?; + match_each_integer_ptype!(indices_array.ptype(), |I| { + bounds_check_valid_indices::(&indices_array.as_view(), array.as_ref().len(), ctx) + })?; let new_validity = array.validity()?.take(indices)?; let new_len = indices_array.len(); @@ -151,61 +147,37 @@ fn take_non_empty_degenerate_fsl( .into_array()) } -fn take_empty_fsl( +fn take_non_empty_non_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, + indices_array: ArrayView<'_, Primitive>, + indices_nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - debug_assert!(array.is_empty()); - - let new_len = indices.len(); - if new_len != 0 { - let indices_validity = indices.validity()?.execute_mask(new_len, ctx)?; - vortex_ensure!( - indices_validity.all_false(), - "cannot take valid indices from an empty FixedSizeList" - ); - } + debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); - let list_size = array.list_size() as usize; - let expected_elements_len = take_elements_len(new_len, list_size)?; - let new_elements = default_elements(array, expected_elements_len); - ensure_elements_len(new_elements.len(), expected_elements_len)?; - let new_validity = if new_len == 0 { + let (new_elements, new_len) = + take_non_empty_non_degenerate_elements::(array, indices_array, ctx)?; + let new_validity = if array.dtype().is_nullable() || indices_nullability.is_nullable() { array.validity()?.take(indices)? } else { - Validity::AllInvalid + Validity::NonNullable }; - // SAFETY: empty output needs no child values; otherwise the index validity mask proves every - // output row is null. Placeholder child elements have the exact length required by FSL. + // SAFETY: `new_elements` has `new_len * list_size` elements. `new_validity` is either + // non-nullable or was produced by `Validity::take` for `new_len`. Ok(unsafe { FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } .into_array()) } -fn take_non_empty_non_degenerate_fsl( +fn take_non_empty_non_degenerate_elements( array: ArrayView<'_, FixedSizeList>, indices_array: ArrayView<'_, Primitive>, - indices_validity: &Mask, - indices_are_nullable: bool, ctx: &mut ExecutionCtx, -) -> VortexResult { - debug_assert!(!array.is_empty()); - debug_assert_ne!(array.list_size(), 0); - - if array.dtype().is_nullable() || indices_are_nullable { - take_nullable_non_empty_fsl::(array, indices_array, indices_validity, ctx) - } else { - take_non_nullable_non_empty_fsl::(array, indices_array) - } -} - -fn take_non_nullable_non_empty_fsl( - array: ArrayView<'_, FixedSizeList>, - indices_array: ArrayView<'_, Primitive>, -) -> VortexResult { +) -> VortexResult<(ArrayRef, usize)> { debug_assert!(!array.is_empty()); debug_assert_ne!(array.list_size(), 0); @@ -213,155 +185,54 @@ fn take_non_nullable_non_empty_fsl( let array_len = array.as_ref().len(); let indices: &[I] = indices_array.as_slice::(); let new_len = indices.len(); - let expected_elements_len = take_elements_len(new_len, list_size)?; - let mut starts = BufferMut::::with_capacity(new_len); - - for &data_idx in indices { - let data_idx = index_to_usize(data_idx)?; - let start = list_start_u64(data_idx, list_size, array_len)?; - starts.push(start); - } - - let new_elements = take_element_runs( - array.elements(), - starts.freeze(), - list_size, - expected_elements_len, - )?; - ensure_elements_len(new_elements.len(), expected_elements_len)?; - - // SAFETY: `starts` contains one checked run of `list_size` elements for each output row, - // `new_elements` has `new_len * list_size` elements, and non-nullable validity has no length. - Ok(unsafe { - FixedSizeListArray::new_unchecked( - new_elements, - array.list_size(), - Validity::NonNullable, - new_len, + let elements_len = new_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" ) - } - .into_array()) -} - -fn take_nullable_non_empty_fsl( - array: ArrayView<'_, FixedSizeList>, - indices_array: ArrayView<'_, Primitive>, - indices_validity: &Mask, - ctx: &mut ExecutionCtx, -) -> VortexResult { - debug_assert!(!array.is_empty()); - debug_assert_ne!(array.list_size(), 0); - - let list_size = array.list_size() as usize; - let array_len = array.as_ref().len(); - let indices: &[I] = indices_array.as_slice::(); - let new_len = indices.len(); - let expected_elements_len = take_elements_len(new_len, list_size)?; - - let array_validity = array - .fixed_size_list_validity() - .execute_mask(array.as_ref().len(), ctx)?; - - let mut starts = BufferMut::::with_capacity(new_len); - let mut new_validity_builder = BitBufferMut::with_capacity(new_len); - - // Null output rows still need placeholder child elements so the FSL elements length stays - // `rows * list_size`. This path has a non-empty source, so 0 is in bounds and validity hides it. - for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { - if !is_index_valid { - starts.push(0); - new_validity_builder.append(false); - continue; - } - - let data_idx = index_to_usize(data_idx)?; - let start = list_start_u64(data_idx, list_size, array_len)?; - if !array_validity.value(data_idx) { - starts.push(0); - new_validity_builder.append(false); - continue; - } - - starts.push(start); - new_validity_builder.append(true); - } - - let new_elements = take_element_runs( - array.elements(), - starts.freeze(), - list_size, - expected_elements_len, - )?; - ensure_elements_len(new_elements.len(), expected_elements_len)?; - - let new_validity = Validity::from(new_validity_builder.freeze()); - debug_assert!(new_validity.maybe_len().is_none_or(|vl| vl == new_len)); + })?; - // SAFETY: `new_elements` has `new_len * list_size` elements. `new_validity_builder` appends - // exactly one bit per output row, and `Validity::from` preserves that length when needed. - Ok(unsafe { - FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) - } - .into_array()) + let indices_validity = indices_array.validity()?.execute_mask(new_len, ctx)?; + let starts = indices + .iter() + .zip_eq(indices_validity.iter()) + .map(|(&data_idx, is_index_valid)| { + if !is_index_valid { + return Ok(0); + } + + let data_idx: usize = data_idx.as_(); + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); + } + Ok((data_idx * list_size) as u64) + }) + .process_results(|iter| iter.collect::>())?; + + let new_elements = + take_element_runs(array.elements(), starts.freeze(), list_size, elements_len)?; + + Ok((new_elements, new_len)) } -fn validate_valid_indices( +fn bounds_check_valid_indices( indices_array: &ArrayView<'_, Primitive>, - indices_validity: &Mask, array_len: usize, + ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let indices: &[I] = indices_array.as_slice::(); + let indices_validity = indices_array.validity()?.execute_mask(indices.len(), ctx)?; - for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { + for (&data_idx, is_index_valid) in indices.iter().zip_eq(indices_validity.iter()) { if is_index_valid { - check_index_in_bounds(index_to_usize(data_idx)?, array_len)?; + let data_idx = data_idx.as_(); + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); + } } } Ok(()) } -fn take_elements_len(new_len: usize, list_size: usize) -> VortexResult { - new_len.checked_mul(list_size).ok_or_else(|| { - vortex_err!( - "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" - ) - }) -} - -fn ensure_elements_len(actual: usize, expected: usize) -> VortexResult<()> { - vortex_ensure!( - actual == expected, - "FixedSizeList take elements length {actual} does not match expected length {expected}" - ); - Ok(()) -} - -fn list_start_u64(data_idx: usize, list_size: usize, array_len: usize) -> VortexResult { - debug_assert_ne!(array_len, 0); - debug_assert_ne!(list_size, 0); - - check_index_in_bounds(data_idx, array_len)?; - - let start = data_idx.checked_mul(list_size).ok_or_else(|| { - vortex_err!( - "FixedSizeList take element range overflow for index {data_idx} and list size {list_size}" - ) - })?; - start.checked_add(list_size).ok_or_else(|| { - vortex_err!( - "FixedSizeList take element range overflow for index {data_idx} and list size {list_size}" - ) - })?; - Ok(start as u64) -} - -fn check_index_in_bounds(data_idx: usize, array_len: usize) -> VortexResult<()> { - if data_idx >= array_len { - vortex_bail!(OutOfBounds: data_idx, 0, array_len); - } - Ok(()) -} - fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef { let mut builder = builder_with_capacity(array.elements().dtype(), len); builder.append_defaults(len); @@ -386,9 +257,3 @@ fn take_element_runs( .into_array(), ) } - -fn index_to_usize(index: I) -> VortexResult { - index - .to_usize() - .ok_or_else(|| vortex_err!("FixedSizeList take index {index} does not fit in usize")) -} From 5431973b37f4c4580e0fff17a22e48d7b909c942 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 09:41:14 -0400 Subject: [PATCH 36/37] Localize TakeSlices reduce-parent fallback Signed-off-by: Daniel King --- vortex-array/src/arrays/take_slices/vtable.rs | 7 ++++ vortex-array/src/executor.rs | 35 ++----------------- 2 files changed, 9 insertions(+), 33 deletions(-) diff --git a/vortex-array/src/arrays/take_slices/vtable.rs b/vortex-array/src/arrays/take_slices/vtable.rs index cdd6278917d..be7ca4e302f 100644 --- a/vortex-array/src/arrays/take_slices/vtable.rs +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -149,6 +149,13 @@ impl VTable for TakeSlices { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(reduced) = array + .child() + .reduce_parent(array.as_ref(), TakeSlicesSlots::CHILD)? + { + return Ok(ExecutionResult::done(reduced)); + } + let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); append_selected_ranges(array.as_view(), builder.as_mut(), ctx)?; Ok(ExecutionResult::done(builder.finish())) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 5d64d89272f..1813e79a9df 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -139,7 +139,7 @@ impl ArrayRef { /// /// Step 2: current_builder active? /// - yes -> skip Step 2a / 2b - /// - no -> try parent kernels + /// - no -> try parent rewrites and kernels /// /// Step 2a: if stack.top exists: /// parent = stack.top.parent_array @@ -150,7 +150,6 @@ impl ArrayRef { /// /// Step 2b: for child in current_array.children(): /// parent = current_array - /// child.reduce_parent(parent, child.slot_idx), then /// kernels[(parent.encoding_id(), child.encoding_id())] /// .try_execute_parent(child, parent, child.slot_idx) /// @@ -245,18 +244,7 @@ impl ArrayRef { continue; } - // Step 2b: reduce_parent/execute_parent against current_array's own children. - if current_builder.is_none() - && let Some(rewritten) = try_reduce_parent(¤t_array, ctx)? - { - ctx.log(format_args!( - "reduce_parent rewrote {} -> {}", - current_array, rewritten - )); - current_array = rewritten.optimize_ctx(ctx.session())?; - continue; - } - + // Step 2b: execute_parent against current_array's own children. if current_builder.is_none() && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { @@ -621,25 +609,6 @@ fn execute_parent_for_child( Ok(None) } -/// Try reduce_parent on each occupied slot of the array. -fn try_reduce_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { - for (slot_idx, slot) in array.slots().iter().enumerate() { - let Some(child) = slot else { continue }; - if let Some(reduced_parent) = child.reduce_parent(array, slot_idx)? { - ctx.log(format_args!( - "reduce_parent: slot[{}]({}) rewrote {} -> {}", - slot_idx, - child.encoding_id(), - array, - reduced_parent - )); - reduced_parent.statistics().inherit_from(array.statistics()); - return Ok(Some(reduced_parent)); - } - } - Ok(None) -} - /// Try execute_parent on each occupied slot of the array. fn try_execute_parent( array: &ArrayRef, From 37634a3c6427459bde10d7971daf8fc80d964abc Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 09:41:44 -0400 Subject: [PATCH 37/37] Build ListView TakeSlices ranges as u64 buffers Signed-off-by: Daniel King --- vortex-array/src/arrays/listview/rebuild.rs | 29 +++++++++++---------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index a19bff08ec3..882920d82aa 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -290,15 +290,10 @@ impl ListViewArray { new_sizes, starts, lengths, + elements_len, } = ranges; - let elements_len = lengths.iter().try_fold(0usize, |acc, &length| { - acc.checked_add(length) - .ok_or_else(|| vortex_err!("ListView rebuild elements length overflow")) - })?; - let starts = - PrimitiveArray::from_iter(starts.into_iter().map(|start| start as u64)).into_array(); - let lengths = - PrimitiveArray::from_iter(lengths.into_iter().map(|length| length as u64)).into_array(); + let starts = PrimitiveArray::new(starts.freeze(), Validity::NonNullable).into_array(); + let lengths = PrimitiveArray::new(lengths.freeze(), Validity::NonNullable).into_array(); let new_sizes = new_sizes.freeze(); let elements = TakeSlicesArray::try_new(self.elements().clone(), starts, lengths, elements_len)? @@ -386,8 +381,9 @@ impl ListViewArray { struct RebuildRanges { new_offsets: BufferMut, new_sizes: BufferMut, - starts: Vec, - lengths: Vec, + starts: BufferMut, + lengths: BufferMut, + elements_len: usize, } fn rebuild_ranges( @@ -403,9 +399,10 @@ where let len = offsets.len(); let mut new_offsets = BufferMut::::with_capacity(len); let mut new_sizes = BufferMut::::with_capacity(len); - let mut starts = Vec::with_capacity(len); - let mut lengths = Vec::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 { @@ -431,11 +428,14 @@ where "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); - lengths.push(length); + starts.push(start as u64); + lengths.push(length as u64); n_elements += num_traits::cast(size).vortex_expect("Cast failed"); } @@ -444,6 +444,7 @@ where new_sizes, starts, lengths, + elements_len, }) }