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 5dc491c28c5..c11bf70964b 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,28 @@ 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::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; +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() { @@ -35,16 +48,25 @@ 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, 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 +84,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 +124,166 @@ 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 starts = indices + .as_ref() + .iter() + .map(|&idx| idx as usize * LIST_SIZE) + .collect::>(); + let run_count = starts.len(); + let starts = starts + .into_iter() + .map(|start| start as u64) + .collect::>(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(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`. + 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 +295,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/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/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/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.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 193730f62c4..49dcf1bfdab 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -2,55 +2,52 @@ // 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_panic; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; +use crate::RecursiveCanonical; use crate::array::ArrayView; use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; +use crate::arrays::TakeSlicesArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; +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 runs +/// and delegates the execution strategy to the elements child via materialized `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, @@ -58,82 +55,123 @@ 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()); - // Make sure to handle degenerate case where lists have size 0 (these can take fast paths). if list_size == 0 { - debug_assert!( - 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. - 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(), + 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, ) - } 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) - } } + .into_array()) } -/// Takes from an array when both the array and indices are non-nullable. -fn take_non_nullable_fsl( +fn take_empty_non_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, indices_array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let list_size = array.list_size() as usize; - let indices: &[I] = indices_array.as_slice::(); - let new_len = indices.len(); + debug_assert_ne!(array.list_size(), 0); + debug_assert!(array.is_empty()); - // Build the element indices directly without validity tracking. - let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); + validate_valid_indices::(&indices_array, 0, ctx)?; - // 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_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 + }; - let list_start = data_idx * list_size; - let list_end = (data_idx + 1) * list_size; + // 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()) +} - // 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")) - }; - } +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_non_empty_fsl::(array, indices_array, ctx) + } else { + take_non_nullable_non_empty_fsl::(array, indices_array, ctx) } +} - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); +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(); + let indices: &[I] = indices_array.as_slice::(); + let new_len = indices.len(); + let expected_elements_len = take_elements_len(new_len, list_size)?; + let mut starts = Vec::with_capacity(new_len); + + 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(start); + } - 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); + let new_elements = take_element_runs( + array.elements(), + starts, + list_size, + expected_elements_len, + ctx, + )?; + ensure_elements_len(new_elements.len(), expected_elements_len)?; - // Both inputs are non-nullable, so the result is non-nullable. + // SAFETY: `starts` contains one checked run of `list_size` elements for each output row, + // `new_elements` has `new_len * list_size` elements, and non-nullable validity has no length. Ok(unsafe { FixedSizeListArray::new_unchecked( new_elements, @@ -145,77 +183,168 @@ 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_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(); 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), - }; + .execute_mask(array.as_ref().len(), 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 mut starts = 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) }; + // Null output rows still need placeholder child elements so the FSL elements length stays + // `rows * list_size`. This path has a non-empty source, so 0 is in bounds and validity hides it. + for (&data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { + if !is_index_valid { + starts.push(0); new_validity_builder.append(false); - } 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); + continue; + } + + 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(0); + new_validity_builder.append(false); + continue; } - } - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); + starts.push(start); + new_validity_builder.append(true); + } - 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); + let new_elements = take_element_runs( + array.elements(), + starts, + list_size, + expected_elements_len, + ctx, + )?; + 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 indices_validity_mask( + indices_array: &ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let indices_len = indices_array.as_ref().len(); + match indices_array.validity()? { + Validity::NonNullable | Validity::AllValid => Ok(Mask::new_true(indices_len)), + Validity::AllInvalid => Ok(Mask::new_false(indices_len)), + Validity::Array(a) => Ok(a.execute::(ctx)?.execute_mask(ctx)), + } +} + +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_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(|| { + vortex_err!( + "FixedSizeList take element range overflow for index {data_idx} and list size {list_size}" + ) + })?; + start.checked_add(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take element range overflow for index {data_idx} and list size {list_size}" + ) + })?; + Ok(start) +} + +fn check_index_in_bounds(data_idx: usize, array_len: usize) -> VortexResult<()> { + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); + } + Ok(()) +} + +fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef { + let mut builder = builder_with_capacity(array.elements().dtype(), len); + builder.append_defaults(len); + builder.finish() +} + +fn take_element_runs( + elements: &ArrayRef, + starts: Vec, + length: usize, + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let run_count = starts.len(); + let starts = starts + .into_iter() + .map(|start| start as u64) + .collect::>(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(length as u64, run_count).into_array(); + + // 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 array, and `output_len` was computed as `run_count * length`. + Ok( + unsafe { TakeSlicesArray::new_unchecked(elements.clone(), starts, lengths, output_len) } + .into_array() + .execute::(ctx)? + .0 + .into_array(), + ) +} + +fn index_to_usize(index: I) -> VortexResult { + index + .to_usize() + .ok_or_else(|| vortex_err!("FixedSizeList take index {index} does not fit in usize")) +} 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/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index af8f94f2d5f..ba2aa7f89bf 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -2,7 +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; @@ -10,19 +18,36 @@ 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::FixedSizeList; use crate::arrays::FixedSizeListArray; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; +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. @@ -112,6 +137,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 +182,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 +237,23 @@ 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. +#[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( FixedSizeListArray::new( @@ -181,24 +286,300 @@ 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_empty_source_rejects_valid_index_after_null_index() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + 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_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); + 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, + Validity::NonNullable, + 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. #[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 +612,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/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/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 342c9eccf19..a19bff08ec3 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -5,16 +5,18 @@ use num_traits::FromPrimitive; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; -use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; +use crate::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::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -140,19 +142,21 @@ impl ListViewArray { // for sizes as well. match_each_unsigned_integer_ptype!(sizes_ptype.to_unsigned(), |S| { match offsets_ptype.to_unsigned() { - PType::U8 => self.naive_rebuild::(ctx), - PType::U16 => self.naive_rebuild::(ctx), - PType::U32 => self.naive_rebuild::(ctx), - PType::U64 => self.naive_rebuild::(ctx), + PType::U8 => self.rebuild_with_take_or_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 [`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( + /// 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 { @@ -164,20 +168,14 @@ impl ListViewArray { .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) + self.rebuild_with_take_slices::(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; @@ -243,112 +241,83 @@ impl ListViewArray { .reinterpret_cast(size_ptype) .into_array(); - // SAFETY: same invariants as `rebuild_list_by_list` — offsets are sequential and - // non-overlapping, all (offset, size) pairs reference valid elements, and the validity - // array is preserved from the original. + // SAFETY: offsets are sequential and non-overlapping, all (offset, size) pairs reference + // valid elements, and the validity array is preserved from the original. Ok(unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?) .with_zero_copy_to_list(true) }) } - /// Rebuilds elements list-by-list: canonicalize elements upfront, then for each list `slice` - /// the relevant range and `append_to_builder` into a typed builder. - fn rebuild_list_by_list( + /// Rebuilds elements using a single contiguous-run gather over the element child. + /// + /// `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, ) -> VortexResult { - let element_dtype = self - .dtype() - .as_list_element_opt() - .vortex_expect("somehow had a canonical list that was not a list"); - let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype()); let size_ptype = self.sizes().dtype().as_ptype(); let offsets_canonical = self.offsets().clone().execute::(ctx)?; let offsets_canonical = offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned()); - let offsets_slice = offsets_canonical.as_slice::(); let sizes_canonical = self.sizes().clone().execute::(ctx)?; let sizes_canonical = sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); - let sizes_slice = sizes_canonical.as_slice::(); - - let len = offsets_slice.len(); - - let mut new_offsets = BufferMut::::with_capacity(len); - // TODO(connor)[ListView]: Do we really need to do this? - // The only reason we need to rebuild the sizes here is that the validity may indicate that - // a list is null even though it has a non-zero size. This rebuild will set the size of all - // null lists to 0. - let mut new_sizes = BufferMut::::with_capacity(len); - - // Canonicalize the elements up front as we will be slicing the elements quite a lot. - let elements_canonical = self - .elements() - .clone() - .execute::(ctx)? - .into_array(); - // Note that we do not know what the exact capacity should be of the new elements since - // there could be overlaps in the existing `ListViewArray`. - let mut new_elements_builder = - builder_with_capacity(element_dtype.as_ref(), self.elements().len()); - - // Resolve validity to a mask once instead of probing it per row (see `rebuild_with_take`). - let validity = self.validity()?.execute_mask(len, ctx)?; - - let mut n_elements = NewOffset::zero(); - for index in 0..len { - if !validity.value(index) { - // For NULL lists, place them after the previous item's data to maintain the - // no-overlap invariant for zero-copy to `ListArray` arrays. - new_offsets.push(n_elements); - new_sizes.push(S::zero()); - continue; - } - - let offset = offsets_slice[index]; - let size = sizes_slice[index]; - - let start = offset.as_(); - let stop = start + size.as_(); + let len = offsets_canonical.len(); - new_offsets.push(n_elements); - new_sizes.push(size); - elements_canonical - .slice(start..stop)? - .append_to_builder(new_elements_builder.as_mut(), ctx)?; + let validity = self.validity()?; - n_elements += num_traits::cast(size).vortex_expect("Cast failed"); - } + // 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_mask = validity.execute_mask(len, ctx)?; + let ranges = match_each_unsigned_integer_ptype!(offsets_canonical.ptype(), |O| { + rebuild_ranges::( + offsets_canonical.as_slice::(), + sizes_canonical.as_slice::(), + &validity_mask, + ) + })?; + + let RebuildRanges { + new_offsets, + new_sizes, + starts, + lengths, + } = 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(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()?) + ListViewArray::new_unchecked(elements, offsets, sizes, validity) .with_zero_copy_to_list(true) }) } @@ -414,6 +383,70 @@ impl ListViewArray { } } +struct RebuildRanges { + new_offsets: BufferMut, + new_sizes: BufferMut, + starts: Vec, + lengths: 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 lengths = 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); + lengths.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"))?; + 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); + lengths.push(length); + n_elements += num_traits::cast(size).vortex_expect("Cast failed"); + } + + Ok(RebuildRanges { + new_offsets, + new_sizes, + starts, + lengths, + }) +} + #[cfg(test)] mod tests { use vortex_buffer::BitBuffer; @@ -593,7 +626,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 +730,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/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/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..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/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 new file mode 100644 index 00000000000..4c317a724e5 --- /dev/null +++ b/vortex-array/src/arrays/struct_/compute/take_slices.rs @@ -0,0 +1,38 @@ +// 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::Struct; +use crate::arrays::StructArray; +use crate::arrays::TakeSlicesArray; +use crate::arrays::struct_::StructArrayExt; +use crate::arrays::take_slices::TakeSlicesReduce; +use crate::arrays::take_slices::check_index_arrays; + +impl TakeSlicesReduce for Struct { + fn take_slices( + array: ArrayView<'_, Self>, + starts: &ArrayRef, + lengths: &ArrayRef, + output_len: usize, + ) -> VortexResult> { + check_index_arrays(starts, lengths)?; + + let fields = array + .iter_unmasked_fields() + .map(|field| { + TakeSlicesArray::try_new(field.clone(), starts.clone(), lengths.clone(), output_len) + .map(IntoArray::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) + .map(Some) + } +} 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..b82d893fe67 --- /dev/null +++ b/vortex-array/src/arrays/take_slices/array.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::smallvec; +use vortex_error::VortexExpect; +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; + +/// The child array selected by the run sequence. +pub(super) const CHILD_SLOT: usize = 0; +/// The start index for each child run. +pub(super) const STARTS_SLOT: usize = 1; +/// The length for each child run. +pub(super) const LENGTHS_SLOT: usize = 2; +pub(super) const NUM_SLOTS: usize = 3; +pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "starts", "lengths"]; + +/// Extension methods for [`TakeSlices`] arrays. +pub trait TakeSlicesArrayExt: TypedArrayRef { + /// 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 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 length for each child run. + fn lengths(&self) -> &ArrayRef { + self.as_ref().slots()[LENGTHS_SLOT] + .as_ref() + .vortex_expect("validated take-slices lengths slot") + } +} +impl> TakeSlicesArrayExt for T {} + +impl Array { + /// Constructs a new `TakeSlicesArray` from start/length arrays and caller-provided output length. + /// + /// Construction validates only the structural array invariants. Index values are interpreted + /// when the lazy gather is executed. + pub fn try_new( + child: ArrayRef, + starts: ArrayRef, + lengths: ArrayRef, + len: usize, + ) -> VortexResult { + let dtype = child.dtype().clone(); + Array::try_from_parts( + ArrayParts::new(TakeSlices, dtype, len, EmptyArrayData).with_slots(smallvec![ + Some(child), + Some(starts), + Some(lengths) + ]), + ) + } + + /// Constructs a new `TakeSlicesArray` without validation. + /// + /// # Safety + /// + /// 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, + starts: ArrayRef, + lengths: ArrayRef, + len: usize, + ) -> Self { + let dtype = child.dtype().clone(); + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(TakeSlices, dtype, len, EmptyArrayData).with_slots(smallvec![ + Some(child), + Some(starts), + Some(lengths) + ]), + ) + } + } +} 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..d215cb0616f --- /dev/null +++ b/vortex-array/src/arrays/take_slices/kernel.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reduce and execute adaptors 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; +use crate::optimizer::rules::ArrayParentReduceRule; + +/// 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`, 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. + 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 new file mode 100644 index 00000000000..ebb4ecacdbe --- /dev/null +++ b/vortex-array/src/arrays/take_slices/mod.rs @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lazy gather of contiguous child ranges. +//! +//! `TakeSlicesArray` represents the concatenation of +//! `values[starts[i]..starts[i] + lengths[i]]` for each range row. Ranges may overlap, repeat, +//! and appear in any order. + +mod array; +mod kernel; +mod vtable; + +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; +use vortex_error::vortex_err; +pub use vtable::*; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::dtype::IntegerPType; + +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 starts and lengths must have equal length, got starts {} and lengths {}", + starts.len(), + lengths.len() + ); + Ok(()) +} + +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 {}", + indices.dtype() + ); + Ok(()) + } + other => vortex_bail!( + "TakeSlicesArray {name} must be a non-nullable unsigned integer, got {other}" + ), + } +} + +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 new file mode 100644 index 00000000000..4c39f758673 --- /dev/null +++ b/vortex-array/src/arrays/take_slices/tests.rs @@ -0,0 +1,616 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +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] +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 = 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); + 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 = 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); + Ok(()) +} + +#[test] +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 = take_slices(&array, &[(1, 3), (4, 2)])?; + 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_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 = take_slices(&array, &[(0, 1), (0, 1)])?; + let expected = PrimitiveArray::from_iter([7i32, 7]); + + assert_arrays_eq!(actual, expected, &mut ctx); + 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 = take_slices(&array, &[(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(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]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +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..6).into_array(); + let starts = PrimitiveArray::from_iter([7u64]).into_array(); + let lengths = PrimitiveArray::from_iter([0u64]).into_array(); + + let take_slices = TakeSlicesArray::try_new(array, starts, lengths, 0)?.into_array(); + + assert!(take_slices.is::()); + assert_eq!(take_slices.len(), 0); + assert!(take_slices.execute::(&mut ctx).is_err()); + Ok(()) +} + +#[test] +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(); + let lengths = ConstantArray::new(3u64, 3).into_array(); + + 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(()) +} + +#[test] +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(); + let lengths = PrimitiveArray::from_iter([1u64, 3, 5]).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); + Ok(()) +} + +#[test] +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(); + let lengths = ConstantArray::new(2u64, 3).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); + 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_reduce_pushes_into_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, + )?; + + 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(()) +} + +#[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(); + let array = PrimitiveArray::from_iter(0i32..6).into_array(); + + let actual = take_slices(&array, &[])?; + let expected = PrimitiveArray::from_iter(std::iter::empty::()); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +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 = take_slices(&array, &[])?; + let expected = PrimitiveArray::from_iter(std::iter::empty::()); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +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(); + + 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 = take_slices(&array, &[(4, 3)])?; + assert!( + out_of_bounds_non_empty + .execute::(&mut ctx) + .is_err() + ); + 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(); + 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)])?; + + 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( + vortex_buffer::buffer![2u8, 0, 1, 2, 0, 1].into_array(), + PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), + )? + .into_array(); + + let actual = take_slices(&dict, &[(1, 3), (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( + vortex_buffer::buffer![0u8, 1, 2, 0].into_array(), + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array(), + )? + .into_array(); + + 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 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) +} + +fn variant_storage(values: impl IntoIterator) -> 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 new file mode 100644 index 00000000000..714ac1d3513 --- /dev/null +++ b/vortex-array/src/arrays/take_slices/vtable.rs @@ -0,0 +1,308 @@ +// 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; +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 crate::ArrayParts; +use crate::ArrayRef; +use crate::EmptyArrayData; +use crate::IntoArray; +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::PrimitiveArray; +use crate::arrays::take_slices::TakeSlicesArrayExt; +use crate::arrays::take_slices::array::CHILD_SLOT; +use crate::arrays::take_slices::array::LENGTHS_SLOT; +use crate::arrays::take_slices::array::NUM_SLOTS; +use crate::arrays::take_slices::array::SLOT_NAMES; +use crate::arrays::take_slices::array::STARTS_SLOT; +use crate::arrays::take_slices::check_index_arrays; +use crate::arrays::take_slices::checked_range_end; +use crate::arrays::take_slices::index_value_to_usize; +use crate::arrays::take_slices::validate_index_ranges; +use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::builder_with_capacity_in; +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; + +/// A [`TakeSlices`]-encoded Vortex array. +pub type TakeSlicesArray = Array; + +/// Contiguous-range gather selection encoding. +/// +/// Like [`crate::arrays::Slice`], this is a lazy compute encoding and is not serialized as a file +/// encoding. Execute it before a file-writing path when a materialized physical representation is +/// required. +#[derive(Clone, Debug)] +pub struct TakeSlices; + +impl VTable for TakeSlices { + type TypedArrayData = EmptyArrayData; + 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" + ); + 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"); + vortex_ensure!( + child.dtype() == dtype, + "TakeSlicesArray dtype {} does not match outer dtype {}", + child.dtype(), + dtype + ); + let starts = slots[STARTS_SLOT] + .as_ref() + .vortex_expect("validated starts slot"); + let lengths = slots[LENGTHS_SLOT] + .as_ref() + .vortex_expect("validated lengths slot"); + check_index_arrays(starts, lengths)?; + 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 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())) + } +} + +impl OperationsVTable for TakeSlices { + fn scalar_at( + array: ArrayView<'_, TakeSlices>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + scalar_at_selected_range(array, index, ctx)? + .ok_or_else(|| vortex_err!("TakeSlicesArray scalar index {index} out of bounds")) + } +} + +impl ValidityVTable for TakeSlices { + fn validity(array: ArrayView<'_, TakeSlices>) -> VortexResult { + array + .child() + .validity()? + .take_slices(array.starts(), array.lengths(), array.len()) + } +} + +fn append_selected_ranges( + array: ArrayView<'_, TakeSlices>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> 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 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_array_ranges( + child: &ArrayRef, + starts: &[S], + lengths: &[L], + output_len: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + S: IntegerPType, + L: IntegerPType, +{ + 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 = start + length; + 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 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 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_array_ranges( + child: &ArrayRef, + index: usize, + starts: &[S], + lengths: &[L], + ctx: &mut ExecutionCtx, +) -> VortexResult> +where + S: IntegerPType, + L: IntegerPType, +{ + let mut logical_start = 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)?; + 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_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 = 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) + .map(Some); + } + Ok(None) +} 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)); } 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>, 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`. /// 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..05765bcc921 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -224,6 +224,29 @@ impl Validity { } } + /// Select validity values by concatenating a sequence of contiguous ranges. + #[allow(clippy::disallowed_methods)] + pub fn take_slices( + &self, + starts: &ArrayRef, + lengths: &ArrayRef, + len: usize, + ) -> VortexResult { + match self { + v @ (Self::NonNullable | Self::AllValid | Self::AllInvalid) => Ok(v.clone()), + Self::Array(is_valid) => Ok(Self::Array( + crate::arrays::TakeSlicesArray::try_new( + is_valid.clone(), + starts.clone(), + lengths.clone(), + len, + )? + .into_array() + .optimize()?, + )), + } + } + // Invert the validity pub fn not(&self) -> VortexResult { match self { 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 }