diff --git a/Cargo.lock b/Cargo.lock index e79211b5436..11a58bb49f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9783,6 +9783,7 @@ dependencies = [ "reqwest 0.13.4", "serde", "serde_json", + "sha2 0.11.0", "spatialbench", "spatialbench-arrow", "sysinfo", diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index 5dc491c28c5..cfd434b16c4 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::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::half::f16; +use vortex_array::match_smallest_offset_type; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_session::VortexSession; fn main() { @@ -39,18 +52,45 @@ const NUM_LISTS: usize = 500; const NUM_INDICES: &[usize] = &[100, 1_000]; /// Fixed size list lengths (elements per list). -const LIST_SIZES: &[usize] = &[16, 64, 256, 1024, 4096]; +const LIST_SIZES: &[usize] = &[16, 64, 128, 256, 512, 1024, 2048, 4096]; + +/// F16 list lengths for isolating the per-index, piecewise, and manual range-copy strategies. +const F16_STRATEGY_LIST_SIZES: &[usize] = &[1, 2, 4, 8, 16, 64, 128, 256, 512, 1024, 2048, 4096]; + +/// F16 strategy benchmarks keep a smaller take width so the forced slow strategies stay cheap. +const F16_STRATEGY_NUM_INDICES: &[usize] = &[10]; /// 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, +{ + create_fsl_with_validity::(list_size, num_lists, Validity::NonNullable) +} + +fn create_fsl_with_validity( + list_size: usize, + num_lists: usize, + validity: Validity, +) -> FixedSizeListArray +where + T: NativePType + FromPrimitive, +{ + let total_elements = list_size * num_lists; + 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, validity, num_lists) +} + +fn create_i64_fsl_with_validity( + list_size: usize, + num_lists: usize, + validity: Validity, +) -> FixedSizeListArray { let total_elements = list_size * num_lists; let elements: Buffer = (0..total_elements as i64).collect(); - FixedSizeListArray::new( - elements.into_array(), - list_size as u32, - Validity::NonNullable, - num_lists, - ) + FixedSizeListArray::new(elements.into_array(), list_size as u32, validity, num_lists) } /// Creates random indices for taking from the array. @@ -63,11 +103,50 @@ 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); + let fsl = create_i64_fsl_with_validity(LIST_SIZE, NUM_LISTS, Validity::NonNullable); + bench_take_fsl_random::(bencher, num_indices, fsl); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_f16_random(bencher: Bencher, num_indices: usize) { + take_fsl_random_typed::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u8_random(bencher: Bencher, num_indices: usize) { + take_fsl_random_typed::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u32_random(bencher: Bencher, num_indices: usize) { + take_fsl_random_typed::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u64_random(bencher: Bencher, num_indices: usize) { + take_fsl_random_typed::(bencher, num_indices); +} + +fn take_fsl_random_typed(bencher: Bencher, num_indices: usize) +where + T: NativePType + FromPrimitive, +{ + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + bench_take_fsl_random::(bencher, num_indices, fsl); +} + +fn bench_take_fsl_random( + bencher: Bencher, + num_indices: usize, + fsl: FixedSizeListArray, +) where + T: NativePType, +{ 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,28 +158,185 @@ fn take_fsl_random(bencher: Bencher, num_indices: usize) }); } -#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] -fn take_fsl_nullable_random(bencher: Bencher, num_indices: usize) { - let total_elements = LIST_SIZE * NUM_LISTS; - let elements: Buffer = (0..total_elements as i64).collect(); - - // Create validity with ~10% nulls - let mut rng = StdRng::seed_from_u64(123); - let validity = Validity::from_iter((0..NUM_LISTS).map(|_| rng.random_ratio(9, 10))); +#[divan::bench(args = F16_STRATEGY_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); - let fsl = FixedSizeListArray::new(elements.into_array(), LIST_SIZE as u32, validity, 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 = F16_STRATEGY_NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_piecewise_sequence( + bencher: Bencher, + num_indices: usize, +) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); let indices = create_random_indices(num_indices, NUM_LISTS); - let indices_array = indices.into_array(); bencher - .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { - array - .clone() - .take(indices.clone()) + take_fsl_f16_piecewise_sequence_strategy::(array, indices) + .into_array() + .execute::(execution_ctx) .unwrap() + }); +} + +#[divan::bench(args = F16_STRATEGY_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_piecewise_sequence_strategy( + array: &FixedSizeListArray, + indices: &Buffer, +) -> FixedSizeListArray { + let starts = indices + .as_ref() + .iter() + .map(|&idx| idx * LIST_SIZE as u64) + .collect::>(); + let run_count = starts.len(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(LIST_SIZE as u64, run_count).into_array(); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); + + // SAFETY: benchmark indices are generated in-bounds; lengths and multiplier 1 are + // non-nullable unsigned constants; output length is exactly `indices.len() * LIST_SIZE`. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + starts, + lengths, + multipliers, + indices.len() * LIST_SIZE, + ) + } + .into_array(); + let elements = array.elements().take(element_indices).unwrap(); + + // 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) { + // Create validity with ~10% nulls + let mut rng = StdRng::seed_from_u64(123); + let validity = Validity::from_iter((0..NUM_LISTS).map(|_| rng.random_ratio(9, 10))); + + let fsl = create_i64_fsl_with_validity(LIST_SIZE, NUM_LISTS, validity); + bench_take_fsl_random::(bencher, num_indices, fsl); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_f16_nullable_random(bencher: Bencher, num_indices: usize) { + take_fsl_nullable_random_typed::(bencher, num_indices); +} + +fn take_fsl_nullable_random_typed(bencher: Bencher, num_indices: usize) +where + T: NativePType + FromPrimitive, +{ + // Create validity with ~10% nulls + let mut rng = StdRng::seed_from_u64(123); + let validity = Validity::from_iter((0..NUM_LISTS).map(|_| rng.random_ratio(9, 10))); + + let fsl = create_fsl_with_validity::(LIST_SIZE, NUM_LISTS, validity); + bench_take_fsl_random::(bencher, num_indices, fsl); +} diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 8469c7e25c2..14d90a5e7b1 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -4,23 +4,32 @@ use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; use vortex_buffer::BitBufferView; use vortex_buffer::get_bit; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::Mask; use crate::ArrayRef; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::builtins::ArrayBuiltins; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; impl TakeExecute for Bool { @@ -29,6 +38,12 @@ impl TakeExecute for Bool { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let indices_nulls_zeroed = match indices.validity()?.execute_mask(indices.len(), ctx)? { Mask::AllTrue(_) => indices.clone(), Mask::AllFalse(_) => { @@ -55,6 +70,49 @@ impl TakeExecute for Bool { } } +fn take_contiguous_ranges( + array: ArrayView<'_, Bool>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let source = array.to_bit_buffer(); + let output_len = indices_ref.len(); + let buffer = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_bit_slices_constant_length( + &source, + starts.as_slice::(), + length, + output_len, + )? + }) + } + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + take_bit_slices( + &source, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )? + }) + }) + } + }; + + Ok(Some( + BoolArray::new(buffer, array.validity()?.take(indices_ref)?).into_array(), + )) +} + fn take_valid_indices>(bools: BitBufferView<'_>, indices: &[I]) -> BitBuffer { // For boolean arrays that roughly fit into a single page (at least, on Linux), it's worth // the overhead to convert to a Vec. @@ -82,6 +140,58 @@ fn take_bool_impl>(bools: BitBufferView<'_>, indices: &[I] }) } +fn take_bit_slices_constant_length( + source: &BitBuffer, + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut values = BitBufferMut::with_capacity(output_len); + for &start in starts { + let start = start.as_(); + values.append_buffer(&source.slice(start..).slice(..length)); + } + + Ok(values.freeze()) +} + +fn take_bit_slices( + source: &BitBuffer, + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, +{ + let mut values = BitBufferMut::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + values.append_buffer(&source.slice(start..).slice(..length)); + } + + vortex_ensure!( + values.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + values.len() + ); + Ok(values.freeze()) +} + #[cfg(test)] mod test { use rstest::rstest; diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index ce786f8e4ab..31523c02a8a 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -1,21 +1,32 @@ // 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 vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::ArrayRef; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Decimal; use crate::arrays::DecimalArray; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::dtype::IntegerPType; use crate::dtype::NativeDecimalType; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_decimal_value_type; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; +use crate::validity::Validity; impl TakeExecute for Decimal { fn take( @@ -23,6 +34,12 @@ impl TakeExecute for Decimal { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let indices = indices.clone().execute::(ctx)?; let validity = array.validity()?.take(&indices.clone().into_array())?; @@ -42,10 +59,180 @@ impl TakeExecute for Decimal { } } +fn take_contiguous_ranges( + array: ArrayView<'_, Decimal>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let validity = array.validity()?.take(indices_ref)?; + let output_len = indices_ref.len(); + let taken = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths); + take_slices_constant_length(array, &starts, length, validity, output_len)? + } + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + take_slices(array, &starts, &lengths, validity, output_len)? + } + }; + Ok(Some(taken)) +} + +fn take_slices_constant_length( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + length: usize, + validity: Validity, + output_len: usize, +) -> VortexResult { + match_each_decimal_value_type!(array.values_type(), |D| { + take_slices_constant_length_typed::(array, starts, length, validity, output_len) + }) +} + +fn take_slices_constant_length_typed( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + length: usize, + validity: Validity, + output_len: usize, +) -> VortexResult +where + D: NativeDecimalType, +{ + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + let values = take_slices_constant_length_to_buffer::( + starts.as_slice::(), + length, + array.buffer::().as_slice(), + output_len, + )?; + + // SAFETY: contiguous gather preserves the decimal dtype and value representation. + Ok( + unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) } + .into_array(), + ) + }) +} + +fn take_slices( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult { + match_each_decimal_value_type!(array.values_type(), |D| { + take_slices_typed::(array, starts, lengths, validity, output_len) + }) +} + +fn take_slices_typed( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult +where + D: NativeDecimalType, +{ + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_slices_start_typed::(array, starts, lengths, validity, output_len) + }) +} + +fn take_slices_start_typed( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult +where + D: NativeDecimalType, + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + let values = take_slices_to_buffer::( + starts.as_slice::(), + lengths.as_slice::(), + array.buffer::().as_slice(), + output_len, + )?; + + // SAFETY: contiguous gather preserves the decimal dtype and value representation. + Ok( + unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) } + .into_array(), + ) + }) +} + fn take_to_buffer(indices: &[I], values: &[T]) -> Buffer { indices.iter().map(|idx| values[idx.as_()]).collect() } +fn take_slices_constant_length_to_buffer( + starts: &[S], + length: usize, + values: &[T], + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, + T: NativeDecimalType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut result = BufferMut::::with_capacity(output_len); + for &start in starts { + let start = start.as_(); + result.extend_from_slice(&values[start..][..length]); + } + + Ok(result.freeze()) +} + +fn take_slices_to_buffer( + starts: &[S], + lengths: &[L], + values: &[T], + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, + L: UnsignedPType, + T: NativeDecimalType, +{ + let mut result = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + result.extend_from_slice(&values[start..][..length]); + } + + vortex_ensure!( + result.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + result.len() + ); + Ok(result.freeze()) +} + #[cfg(test)] mod tests { use rstest::rstest; 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..5f2e0a50835 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -1,221 +1,251 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::BitBufferMut; +use itertools::Itertools; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_panic; -use vortex_mask::Mask; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; -use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -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::DType; use crate::dtype::IntegerPType; use crate::executor::ExecutionCtx; -use crate::match_each_unsigned_integer_ptype; -use crate::match_smallest_offset_type; +use crate::match_each_integer_ptype; +use crate::optimizer::ArrayOptimizer; 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 `PiecewiseSequenceArray` indices. 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) - }) - }) - .map(Some) + if array.is_empty() { + return take_empty_fsl(array, indices, ctx).map(Some); + } + + take_non_empty_fsl(array, indices, ctx).map(Some) } } -/// Dispatches to the appropriate take implementation based on list size and nullability. -fn take_with_indices( +fn take_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - let list_size = array.list_size() as usize; + debug_assert!(array.is_empty()); - let indices_array = indices.clone().execute::(ctx)?; - // Reinterpret to unsigned so `as_slice::` (with unsigned `I`) matches; values are unchanged. - let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned()); - - // Make sure to handle degenerate case where lists have size 0 (these can take fast paths). - if list_size == 0 { - debug_assert!( - array.elements().is_empty(), - "degenerate list must have empty elements" + let new_len = indices.len(); + if new_len != 0 { + let indices_validity = indices.validity()?.execute_mask(new_len, ctx)?; + vortex_ensure!( + indices_validity.all_false(), + "cannot take valid indices from an empty FixedSizeList" ); + } - // 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(), + let list_size = array.list_size() as usize; + let elements_len = new_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" ) + })?; + let new_elements = default_elements(array, elements_len); + let new_validity = if new_len == 0 { + array.validity()?.take(indices)? } else { - // 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) - } + Validity::AllInvalid + }; + + // SAFETY: empty output needs no child values; otherwise the index validity mask proves every + // output row is null. Placeholder child elements have the exact length required by FSL. + unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } + .into_array() + .optimize_ctx(ctx.session()) } -/// Takes from an array when both the array and indices are non-nullable. -fn take_non_nullable_fsl( +fn take_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, - indices_array: ArrayView<'_, Primitive>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let list_size = array.list_size() as usize; - let indices: &[I] = indices_array.as_slice::(); - let new_len = indices.len(); - - // Build the element indices directly without validity tracking. - let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); + debug_assert!(!array.is_empty()); - // Build the element indices for each list. - for data_idx in indices { - let data_idx = data_idx - .to_usize() - .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx)); - - let list_start = data_idx * list_size; - let list_end = (data_idx + 1) * list_size; + let DType::Primitive(ptype, _) = indices.dtype() else { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + }; + if !ptype.is_int() { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()); + } - // 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")) - }; - } + if array.list_size() == 0 { + return take_non_empty_degenerate_fsl(array, indices, ctx); } - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); + let indices_array = indices.clone().execute::(ctx)?; + match_each_integer_ptype!(indices_array.ptype(), |I| { + take_non_empty_non_degenerate_fsl::(array, indices, indices_array.as_view(), ctx) + }) +} - 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); +fn take_non_empty_degenerate_fsl( + array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + debug_assert!(!array.is_empty()); + debug_assert_eq!(array.list_size(), 0); + vortex_ensure!( + array.elements().is_empty(), + "degenerate list must have empty elements" + ); - // Both inputs are non-nullable, so the result is non-nullable. - Ok(unsafe { + let indices_array = indices.clone().execute::(ctx)?; + match_each_integer_ptype!(indices_array.ptype(), |I| { + bounds_check_valid_indices::(&indices_array.as_view(), array.as_ref().len(), ctx) + })?; + let new_validity = array.validity()?.take(indices)?; + let new_len = indices_array.len(); + + // SAFETY: degenerate FSL inputs have no elements, valid index payloads were checked against + // the source length, and `Validity::take` produces validity for `new_len`. + unsafe { FixedSizeListArray::new_unchecked( - new_elements, + array.elements().clone(), array.list_size(), - Validity::NonNullable, + new_validity, new_len, ) } - .into_array()) + .into_array() + .optimize_ctx(ctx.session()) } -/// Takes from an array when either the array or indices are nullable. -fn take_nullable_fsl( +fn take_non_empty_non_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, indices_array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { + debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); + + let (new_elements, new_len) = + take_non_empty_non_degenerate_elements::(array, indices_array, ctx)?; + let new_validity = array.validity()?.take(indices)?; + + // SAFETY: `new_elements` has `new_len * list_size` elements, and `Validity::take` produces + // validity for `new_len`. + unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) + } + .into_array() + .optimize_ctx(ctx.session()) +} + +fn take_non_empty_non_degenerate_elements( + array: ArrayView<'_, FixedSizeList>, + indices_array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ArrayRef, usize)> { + debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); + let list_size = array.list_size() as usize; + let array_len = array.as_ref().len(); let indices: &[I] = indices_array.as_slice::(); let new_len = indices.len(); - - let 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), - }; - - // 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 new_validity_builder = BitBufferMut::with_capacity(new_len); - - // Build the element indices while tracking which lists are null. - for (data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { - let data_idx = data_idx - .to_usize() - .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx)); - - // The list is null if the index is null or the indexed element is null. - if !is_index_valid || !array_validity.value(data_idx) { - // Append placeholder zeros for null lists. These will be masked by the validity array. - // We cannot use append_nulls here as explained above. - unsafe { elements_indices.push_n_unchecked(E::zero(), list_size) }; - new_validity_builder.append(false); - } else { - // Append the actual element indices for this list. - let list_start = data_idx * list_size; - let list_end = (data_idx + 1) * list_size; - - // Expand the list into individual element indices. - for i in list_start..list_end { - // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity. - unsafe { - elements_indices.push_unchecked(E::from_usize(i).vortex_expect("i < list_end")) - }; + let elements_len = new_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" + ) + })?; + + let indices_validity = indices_array.validity()?.execute_mask(new_len, ctx)?; + let starts = indices + .iter() + .zip_eq(indices_validity.iter()) + .map(|(&data_idx, is_index_valid)| { + if !is_index_valid { + return Ok(0); } - new_validity_builder.append(true); - } - } + let data_idx: usize = data_idx.as_(); + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); + } + Ok((data_idx * list_size) as u64) + }) + .process_results(|iter| iter.collect::>())?; - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); + let new_elements = + take_element_runs(array.elements(), starts.freeze(), list_size, elements_len)?; - let elements_indices_array = PrimitiveArray::new(elements_indices, Validity::NonNullable); - let new_elements = array.elements().take(elements_indices_array.into_array())?; - debug_assert_eq!(new_elements.len(), new_len * list_size); + Ok((new_elements, new_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)); +fn bounds_check_valid_indices( + indices_array: &ArrayView<'_, Primitive>, + array_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let indices: &[I] = indices_array.as_slice::(); + let indices_validity = indices_array.validity()?.execute_mask(indices.len(), ctx)?; - Ok(unsafe { - FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) + for (&data_idx, is_index_valid) in indices.iter().zip_eq(indices_validity.iter()) { + if is_index_valid { + let data_idx = data_idx.as_(); + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); + } + } } - .into_array()) + 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: Buffer, + length: usize, + output_len: usize, +) -> VortexResult { + let run_count = starts.len(); + let starts = PrimitiveArray::new(starts, Validity::NonNullable).into_array(); + let lengths = ConstantArray::new(length as u64, run_count).into_array(); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); + + // SAFETY: callers produced one start per output row after validating list indices against the + // source FSL length. `length` and multiplier 1 are represented as non-nullable unsigned + // constant arrays, and `output_len` was computed as `run_count * length`. + let indices = + unsafe { PiecewiseSequenceArray::new_unchecked(starts, lengths, multipliers, output_len) }; + elements.take(indices.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..93262a74b37 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -147,6 +147,28 @@ fn test_take_fsl_with_null_indices_preserves_elements() { assert_arrays_eq!(expected, result, &mut ctx); } +#[test] +fn test_take_fsl_null_index_ignores_out_of_bounds_payload() { + 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![1u32, 99, 2], + Validity::from_iter([true, false, true]), + ) + .into_array(); + let result = fsl.take(indices).unwrap(); + + let expected = FixedSizeListArray::new( + buffer![3i32, 4, 0, 0, 5, 6].into_array(), + 2, + Validity::from_iter([true, false, true]), + 3, + ); + 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. #[rstest] diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index b914c5e9651..05b7b13edf4 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -9,18 +9,23 @@ //! element. use itertools::Itertools; +use num_traits::AsPrimitive; use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use crate::ArrayRef; +use crate::Columnar; use crate::array::ArrayView; +use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; +use crate::scalar::PValue; pub mod array; mod vtable; @@ -61,10 +66,48 @@ pub(crate) fn execute_index_arrays( let starts = array.starts().clone().execute::(ctx)?; let lengths = array.lengths().clone().execute::(ctx)?; let multipliers = array.multipliers().clone().execute::(ctx)?; - check_index_arrays(starts.as_ref(), lengths.as_ref(), multipliers.as_ref())?; Ok((starts, lengths, multipliers)) } +pub(crate) fn maybe_contiguous_slices( + array: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if !is_constant_one(array.multipliers()) { + return Ok(None); + } + + let starts = array.starts().clone().execute::(ctx)?; + let lengths = array.lengths().clone().execute::(ctx)?; + Ok(Some((starts, lengths))) +} + +pub(crate) fn is_constant_one(multipliers: &ArrayRef) -> bool { + let Some(scalar) = multipliers.as_constant() else { + return false; + }; + matches!( + scalar.as_primitive_opt().and_then(|scalar| scalar.pvalue()), + Some(PValue::U8(1) | PValue::U16(1) | PValue::U32(1) | PValue::U64(1)) + ) +} + +pub(crate) fn constant_unsigned_usize(array: &ConstantArray) -> usize { + let pvalue = array + .scalar() + .as_primitive_opt() + .and_then(|scalar| scalar.pvalue()) + .vortex_expect("validated PiecewiseSequence length constants are primitive"); + + match pvalue { + PValue::U8(value) => value as usize, + PValue::U16(value) => value as usize, + PValue::U32(value) => value as usize, + PValue::U64(value) => value.as_(), + _ => unreachable!("validated PiecewiseSequence length constants are unsigned"), + } +} + fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { vortex_ensure!( array.dtype().is_unsigned_int(), diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index ab201943e3d..eef5d4af12d 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -1,22 +1,57 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::ByteBufferMut; +use num_traits::AsPrimitive; use vortex_buffer::buffer; use vortex_error::VortexResult; -use vortex_session::registry::ReadContext; use crate::ArrayContext; +use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::BoolArray; use crate::arrays::ConstantArray; -use crate::arrays::PiecewiseSequence; +use crate::arrays::DecimalArray; +use crate::arrays::FixedSizeListArray; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinArray; +use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; use crate::serde::SerializeOptions; -use crate::serde::SerializedArray; +use crate::validity::Validity; + +fn piecewise_indices( + starts: impl IntoIterator, + lengths: &[u64], +) -> VortexResult { + let len = lengths + .iter() + .map(|&length| -> usize { length.as_() }) + .sum(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = PrimitiveArray::from_iter(lengths.iter().copied()).into_array(); + let multipliers = ConstantArray::new(1u64, lengths.len()).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, lengths, multipliers, len)?.into_array()) +} + +fn piecewise_indices_constant_length( + starts: impl IntoIterator, + length: u64, +) -> VortexResult { + let starts = starts.into_iter().collect::>(); + let length_usize: usize = length.as_(); + let len = starts.len() * length_usize; + let piece_count = starts.len(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(length, piece_count).into_array(); + let multipliers = ConstantArray::new(1u64, piece_count).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, lengths, multipliers, len)?.into_array()) +} #[test] fn materializes_piecewise_indices() -> VortexResult<()> { @@ -117,7 +152,7 @@ fn execution_checks_declared_length() -> VortexResult<()> { } #[test] -fn serde_roundtrip_preserves_piecewise_indices() -> VortexResult<()> { +fn serialization_is_not_supported() -> VortexResult<()> { let array = PiecewiseSequenceArray::try_new( buffer![3u32, 15, 21].into_array(), buffer![2u16, 0, 2].into_array(), @@ -125,30 +160,193 @@ fn serde_roundtrip_preserves_piecewise_indices() -> VortexResult<()> { 4, )? .into_array(); - let dtype = array.dtype().clone(); - let len = array.len(); - - let array_ctx = ArrayContext::empty(); - let serialized = array.serialize(&array_ctx, &array_session(), &SerializeOptions::default())?; - - let mut concat = ByteBufferMut::empty(); - for buffer in serialized { - concat.extend_from_slice(buffer.as_ref()); - } - - let parts = SerializedArray::try_from(concat.freeze())?; - let decoded = parts.decode( - &dtype, - len, - &ReadContext::new(array_ctx.to_ids()), - &array_session(), - )?; - - assert!(decoded.is::()); + + let err = array + .serialize( + &ArrayContext::empty(), + &array_session(), + &SerializeOptions::default(), + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("Array vortex.piecewise-sequence does not support serialization"), + "{err}" + ); + Ok(()) +} + +#[test] +fn primitive_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = PrimitiveArray::from_iter(0i32..20).into_array(); + let taken = values.take(piecewise_indices([3, 10], &[2, 3])?)?; + + assert_arrays_eq!( + taken, + PrimitiveArray::from_iter([3i32, 4, 10, 11, 12]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn primitive_take_handles_non_unit_multiplier() -> VortexResult<()> { + let values = PrimitiveArray::from_iter(0i32..20).into_array(); + let indices = PiecewiseSequenceArray::try_new( + buffer![3u64].into_array(), + buffer![3u64].into_array(), + buffer![2u64].into_array(), + 3, + )? + .into_array(); + let taken = values.take(indices)?; + + assert_arrays_eq!( + taken, + PrimitiveArray::from_iter([3i32, 5, 7]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn bool_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = BoolArray::from_iter([true, false, true, true, false, false]).into_array(); + let taken = values.take(piecewise_indices([1, 4], &[2, 2])?)?; + + assert_arrays_eq!( + taken, + BoolArray::from_iter([false, true, false, false]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn decimal_take_consumes_piecewise_indices() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(19, 2); + let values = DecimalArray::from_iter([10i128, 20, 30, 40, 50, 60], decimal_dtype).into_array(); + let taken = values.take(piecewise_indices([1, 4], &[2, 1])?)?; + + assert_arrays_eq!( + taken, + DecimalArray::from_iter([20i128, 30, 50], decimal_dtype).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn varbinview_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = VarBinViewArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let taken = values.take(piecewise_indices([1, 3], &[2, 1])?)?; + + assert_arrays_eq!( + taken, + VarBinViewArray::from_iter( + ["bb", "ccc", "dddd"].map(Some), + DType::Utf8(Nullability::NonNullable) + ) + .into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn varbin_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = VarBinArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let taken = values.take(piecewise_indices([1, 3], &[2, 1])?)?; + assert_arrays_eq!( - decoded, - PrimitiveArray::from_iter([3u64, 5, 21, 24]).into_array(), + taken, + VarBinArray::from_iter( + ["bb", "ccc", "dddd"].map(Some), + DType::Utf8(Nullability::NonNullable) + ) + .into_array(), &mut array_session().create_execution_ctx() ); Ok(()) } + +#[test] +fn contiguous_take_consumers_support_constant_piecewise_lengths() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let indices = piecewise_indices_constant_length([1, 3], 2)?; + + let primitive = PrimitiveArray::from_iter(0i32..10).into_array(); + assert_arrays_eq!( + primitive.take(indices.clone())?, + PrimitiveArray::from_iter([1i32, 2, 3, 4]).into_array(), + &mut ctx + ); + + let bools = BoolArray::from_iter([false, true, false, true, false]).into_array(); + assert_arrays_eq!( + bools.take(indices.clone())?, + BoolArray::from_iter([true, false, true, false]).into_array(), + &mut ctx + ); + + let decimal_dtype = DecimalDType::new(19, 2); + let decimals = DecimalArray::from_iter([10i128, 20, 30, 40, 50], decimal_dtype).into_array(); + assert_arrays_eq!( + decimals.take(indices.clone())?, + DecimalArray::from_iter([20i128, 30, 40, 50], decimal_dtype).into_array(), + &mut ctx + ); + + let varbinview = VarBinViewArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + assert_arrays_eq!( + varbinview.take(indices.clone())?, + VarBinViewArray::from_iter( + ["bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(), + &mut ctx + ); + + let varbin = VarBinArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + assert_arrays_eq!( + varbin.take(indices)?, + VarBinArray::from_iter( + ["bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(), + &mut ctx + ); + + Ok(()) +} + +#[test] +fn fixed_size_list_take_builds_piecewise_element_indices() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter(0i32..12).into_array(); + let array = FixedSizeListArray::new(elements, 3, Validity::NonNullable, 4).into_array(); + let taken = array.take(buffer![1u64, 3].into_array())?; + + let expected_elements = PrimitiveArray::from_iter([3i32, 4, 5, 9, 10, 11]).into_array(); + let expected = + FixedSizeListArray::new(expected_elements, 3, Validity::NonNullable, 2).into_array(); + assert_arrays_eq!(taken, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} diff --git a/vortex-array/src/arrays/piecewise_sequence/vtable.rs b/vortex-array/src/arrays/piecewise_sequence/vtable.rs index ac2eba61116..d68a00815f3 100644 --- a/vortex-array/src/arrays/piecewise_sequence/vtable.rs +++ b/vortex-array/src/arrays/piecewise_sequence/vtable.rs @@ -2,8 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; -use prost::Message; -use smallvec::smallvec; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -26,7 +24,6 @@ use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; use crate::arrays::PrimitiveArray; -use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; use crate::arrays::piecewise_sequence::array::PiecewiseSequenceSlots; use crate::arrays::piecewise_sequence::check_index_arrays; use crate::arrays::piecewise_sequence::execute_index_arrays; @@ -44,18 +41,6 @@ use crate::validity::Validity; /// A [`PiecewiseSequence`]-encoded Vortex index array. pub type PiecewiseSequenceArray = Array; -#[derive(Clone, prost::Message)] -struct PiecewiseSequenceMetadata { - #[prost(uint64, tag = "1")] - num_pieces: u64, - #[prost(enumeration = "PType", tag = "2")] - starts_ptype: i32, - #[prost(enumeration = "PType", tag = "3")] - lengths_ptype: i32, - #[prost(enumeration = "PType", tag = "4")] - multipliers_ptype: i32, -} - #[derive(Clone, Debug)] pub struct PiecewiseSequence; @@ -125,77 +110,22 @@ impl VTable for PiecewiseSequence { } fn serialize( - array: ArrayView<'_, Self>, + _array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - Ok(Some( - PiecewiseSequenceMetadata { - num_pieces: u64::try_from(array.starts().len()).map_err(|_| { - vortex_err!( - "PiecewiseSequenceArray piece count {} overflowed u64", - array.starts().len() - ) - })?, - starts_ptype: PType::try_from(array.starts().dtype())? as i32, - lengths_ptype: PType::try_from(array.lengths().dtype())? as i32, - multipliers_ptype: PType::try_from(array.multipliers().dtype())? as i32, - } - .encode_to_vec(), - )) + Ok(None) } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - let metadata = PiecewiseSequenceMetadata::decode(metadata)?; - vortex_ensure!( - dtype == &DType::from(PType::U64), - "PiecewiseSequenceArray dtype must be u64, got {dtype}" - ); - vortex_ensure!( - buffers.is_empty(), - "PiecewiseSequenceArray expects no buffers, got {}", - buffers.len() - ); - vortex_ensure!( - children.len() == PiecewiseSequenceSlots::NAMES.len(), - "PiecewiseSequenceArray expects {} children, got {}", - PiecewiseSequenceSlots::NAMES.len(), - children.len() - ); - - let num_pieces = usize::try_from(metadata.num_pieces).map_err(|_| { - vortex_err!( - "PiecewiseSequenceArray piece count {} does not fit in usize", - metadata.num_pieces - ) - })?; - let starts = children.get( - PiecewiseSequenceSlots::STARTS, - &metadata.starts_ptype().into(), - num_pieces, - )?; - let lengths = children.get( - PiecewiseSequenceSlots::LENGTHS, - &metadata.lengths_ptype().into(), - num_pieces, - )?; - let multipliers = children.get( - PiecewiseSequenceSlots::MULTIPLIERS, - &metadata.multipliers_ptype().into(), - num_pieces, - )?; - - Ok( - ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData) - .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), - ) + vortex_bail!("PiecewiseSequenceArray is not serializable") } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index dd562281608..70401fcd367 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -6,25 +6,35 @@ mod avx2; use std::sync::LazyLock; +use itertools::Itertools as _; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::Mask; use crate::ArrayRef; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::ConstantArray; +use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; +use crate::dtype::NativePType; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; use crate::match_each_native_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; use crate::validity::Validity; @@ -79,6 +89,12 @@ impl TakeExecute for Primitive { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let DType::Primitive(ptype, null) = indices.dtype() else { vortex_bail!("Invalid indices dtype: {}", indices.dtype()) }; @@ -123,6 +139,30 @@ impl TakeExecute for Primitive { } } +fn take_contiguous_ranges( + array: ArrayView<'_, Primitive>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let validity = array.validity()?.take(indices_ref)?; + let output_len = indices_ref.len(); + let taken = match lengths { + Columnar::Constant(lengths) => { + let length = constant_unsigned_usize(&lengths); + take_slices_constant_length(array, &starts, length, validity, output_len)? + } + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + take_slices(array, &starts, &lengths, validity, output_len)? + } + }; + Ok(Some(taken)) +} + // Compiler may see this as unused based on enabled features #[inline(always)] fn take_primitive_scalar(buffer: &[T], indices: &[I]) -> Buffer { @@ -144,6 +184,142 @@ fn take_primitive_scalar(buffer: &[T], indices: &[I]) result.freeze() } +fn take_slices( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult { + match_each_native_ptype!(array.ptype(), |T| { + take_slices_typed::(array, starts, lengths, validity, output_len) + }) +} + +fn take_slices_typed( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult +where + T: NativePType, +{ + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_slices_start_typed::(array, starts, lengths, validity, output_len) + }) +} + +fn take_slices_start_typed( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult +where + T: NativePType, + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + let values = take_slices_to_buffer::( + array.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + Ok(PrimitiveArray::new(values, validity).into_array()) + }) +} + +fn take_slices_to_buffer( + source: &[T], + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + T: Copy, + S: UnsignedPType, + L: UnsignedPType, +{ + let mut values = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + values.extend_from_slice(&source[start..][..length]); + } + + vortex_ensure!( + values.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + values.len() + ); + Ok(values.freeze()) +} + +fn take_slices_constant_length( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + length: usize, + validity: Validity, + output_len: usize, +) -> VortexResult { + match_each_native_ptype!(array.ptype(), |T| { + take_slices_constant_length_typed::(array, starts, length, validity, output_len) + }) +} + +fn take_slices_constant_length_typed( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + length: usize, + validity: Validity, + output_len: usize, +) -> VortexResult +where + T: NativePType, +{ + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + let values = take_slices_constant_length_to_buffer::( + array.as_slice::(), + starts.as_slice::(), + length, + output_len, + )?; + Ok(PrimitiveArray::new(values, validity).into_array()) + }) +} + +fn take_slices_constant_length_to_buffer( + source: &[T], + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult> +where + T: Copy, + S: UnsignedPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut values = BufferMut::::with_capacity(output_len); + for &start in starts { + let start = start.as_(); + values.extend_from_slice(&source[start..][..length]); + } + + Ok(values.freeze()) +} + #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] #[cfg(test)] mod test { diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 02f043b8e68..d3a23c976f1 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -58,6 +58,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 }