From 8ce4ad9a6b3f5867c7fce715962184903fa79bb5 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 7 Jul 2026 12:30:03 -0400 Subject: [PATCH 01/21] 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 e19b8c543d1803c487ba886fa794c7cc9b0a45bd Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 7 Jul 2026 16:50:59 -0400 Subject: [PATCH 02/21] 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 a93ea19dc7b473a186b9b59f339849b3b98d9ba1 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Tue, 7 Jul 2026 17:06:43 -0400 Subject: [PATCH 03/21] 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 44cd3d444b248b9a351b399bf4d35929424e2f41 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 16:07:58 -0400 Subject: [PATCH 04/21] 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 6b70cba507084d5707b208d9dc3b5f438077aac2 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 16:28:32 -0400 Subject: [PATCH 05/21] 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 ce21f59a33a261113938452388cf71c09629f63d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 16:31:31 -0400 Subject: [PATCH 06/21] 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 b0ecab88d32dd5b55b6cb0c3c9cf7d108e1a10fb Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 16:33:02 -0400 Subject: [PATCH 07/21] 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 4285b937414273bdafe10b9ea3fda03a20dd26db Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 17:04:00 -0400 Subject: [PATCH 08/21] 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 b7c367ef8f6..36805c400e2 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -45,6 +45,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 868d30b7b41529b9dafd1feee14903d01abd8256 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 9 Jul 2026 17:16:49 -0400 Subject: [PATCH 09/21] 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 36805c400e2..9e663e87ff5 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 65cd73c555c69971d752a77830112a47e2b101d8 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 11:31:08 -0400 Subject: [PATCH 10/21] 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 16df2a942c1..1fb40227caa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9902,6 +9902,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 069f33a57e9..0b5e6fde99b 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -51,6 +51,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 9e663e87ff5..fba14169f7c 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 72d0c72f90b0e56d5d600ed33be57218e5d7180c Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 12:09:47 -0400 Subject: [PATCH 11/21] 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 e9e5c75efd0ed48125dbe30e7bdcc95d448cf0d5 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 13:31:16 -0400 Subject: [PATCH 12/21] 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 db8c84cfd27fe1e799e899b4a2116678b9f5da49 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 13:38:18 -0400 Subject: [PATCH 13/21] 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 0d4e192b7d2246c4f2123b00e1f254d87a72d64d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:05:33 -0400 Subject: [PATCH 14/21] 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 de1b3f06fffc8fbb698fe14e880206371459df86 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:08:46 -0400 Subject: [PATCH 15/21] 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 fba14169f7c..b7c367ef8f6 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; @@ -51,7 +45,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 45cc0afa118e300dc482b83715da748d609d2987 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:14:23 -0400 Subject: [PATCH 16/21] 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 a9601f67bfc5c2d25bb16a27bebb2eb4e54b1aed Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:23:27 -0400 Subject: [PATCH 17/21] 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 3910c8b4edd8b9722fe93060a0bdcc1536acd303 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 14:37:01 -0400 Subject: [PATCH 18/21] 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 53c40ac2301115aeb14bd869816e650272622699 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 16:36:52 -0400 Subject: [PATCH 19/21] 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 5b2446f22c89bc0225ab2e8bb7094817cc363380 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 10 Jul 2026 17:32:22 -0400 Subject: [PATCH 20/21] 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 feb6a4f1509ddf6a22516271b60ccce6edbe7bba Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Sat, 11 Jul 2026 01:50:20 +0100 Subject: [PATCH 21/21] 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];