Add TakeSlices for fixed-size-list take

vortex-array/benches/take_fsl.rs

@@ -6,23 +6,36 @@
 //! 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)]
 
 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<VortexSession> = 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<T>(list_size: usize, num_lists: usize) -> FixedSizeListArray
+where
+    T: NativePType + FromPrimitive,
+{
     let total_elements = list_size * num_lists;
-    let elements: Buffer<i64> = (0..total_elements as i64).collect();
+    let elements: Buffer<T> = (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<u64> {
 }
 
 #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)]
-fn take_fsl_random<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize) {
-    let fsl = create_fsl(LIST_SIZE, NUM_LISTS);
+fn take_fsl_f16_random<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize) {
+    take_fsl_random::<f16, LIST_SIZE>(bencher, num_indices);
+}
+
+#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)]
+fn take_fsl_u8_random<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize) {
+    take_fsl_random::<u8, LIST_SIZE>(bencher, num_indices);
+}
+
+#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)]
+fn take_fsl_u32_random<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize) {
+    take_fsl_random::<u32, LIST_SIZE>(bencher, num_indices);
+}
+
+#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)]
+fn take_fsl_u64_random<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize) {
+    take_fsl_random::<u64, LIST_SIZE>(bencher, num_indices);
+}
+
+fn take_fsl_random<T, const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize)
+where
+    T: NativePType + FromPrimitive,
+{
+    let fsl = create_fsl::<T>(LIST_SIZE, NUM_LISTS);
     let indices = create_random_indices(num_indices, NUM_LISTS);
     let indices_array = indices.into_array();
 
     bencher
+        .counter(BytesCount::of_many::<T>(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<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize)
         });
 }
 
+#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)]
+fn take_fsl_f16_force_per_index<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize) {
+    let fsl = create_fsl::<f16>(LIST_SIZE, NUM_LISTS);
+    let indices = create_random_indices(num_indices, NUM_LISTS);
+
+    bencher
+        .counter(BytesCount::of_many::<f16>(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::<LIST_SIZE, E>(array, indices)
+            })
+            .into_array()
+            .execute::<RecursiveCanonical>(execution_ctx)
+            .unwrap()
+        });
+}
+
+#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)]
+fn take_fsl_f16_force_take_slices<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize) {
+    let fsl = create_fsl::<f16>(LIST_SIZE, NUM_LISTS);
+    let indices = create_random_indices(num_indices, NUM_LISTS);
+
+    bencher
+        .counter(BytesCount::of_many::<f16>(num_indices * LIST_SIZE))
+        .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx()))
+        .bench_refs(|(array, indices, execution_ctx)| {
+            take_fsl_f16_take_slices_strategy::<LIST_SIZE>(array, indices)
+                .into_array()
+                .execute::<RecursiveCanonical>(execution_ctx)
+                .unwrap()
+        });
+}
+
+#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)]
+fn take_fsl_f16_force_manual_range_copy<const LIST_SIZE: usize>(
+    bencher: Bencher,
+    num_indices: usize,
+) {
+    let fsl = create_fsl::<f16>(LIST_SIZE, NUM_LISTS);
+    let indices = create_random_indices(num_indices, NUM_LISTS);
+
+    bencher
+        .counter(BytesCount::of_many::<f16>(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::<LIST_SIZE>(array, indices, execution_ctx)
+                .into_array()
+                .execute::<RecursiveCanonical>(execution_ctx)
+                .unwrap()
+        });
+}
+
+fn take_fsl_f16_per_index_strategy<const LIST_SIZE: usize, E: IntegerPType>(
+    array: &FixedSizeListArray,
+    indices: &Buffer<u64>,
+) -> FixedSizeListArray {
+    let mut element_indices = BufferMut::<E>::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<const LIST_SIZE: usize>(
+    array: &FixedSizeListArray,
+    indices: &Buffer<u64>,
+) -> FixedSizeListArray {
+    let starts = indices
+        .as_ref()
+        .iter()
+        .map(|&idx| idx as usize * LIST_SIZE)
+        .collect::<Vec<_>>();
+    let run_count = starts.len();
+    let starts = starts
+        .into_iter()
+        .map(|start| start as u64)
+        .collect::<Vec<_>>();
+    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<const LIST_SIZE: usize>(
+    array: &FixedSizeListArray,
+    indices: &Buffer<u64>,
+    execution_ctx: &mut ExecutionCtx,
+) -> FixedSizeListArray {
+    let elements = array
+        .elements()
+        .clone()
+        .execute::<PrimitiveArray>(execution_ctx)
+        .unwrap();
+    let source = elements.as_slice::<f16>();
+    let mut values = BufferMut::<f16>::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<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize) {
+fn take_fsl_f16_nullable_random<const LIST_SIZE: usize>(bencher: Bencher, num_indices: usize) {
     let total_elements = LIST_SIZE * NUM_LISTS;
-    let elements: Buffer<i64> = (0..total_elements as i64).collect();
+    let elements: Buffer<f16> = (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<const LIST_SIZE: usize>(bencher: Bencher, num_indice
     let indices_array = indices.into_array();
 
     bencher
+        .counter(BytesCount::of_many::<f16>(num_indices * LIST_SIZE))
         .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx()))
         .bench_refs(|(array, indices, execution_ctx)| {
             array

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)]

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<Option<ArrayRef>> {
+        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::<S, L>(array, starts, lengths, output_len, ctx)
+            })
+        })
+        .map(Some)
+    }
+}
+
+fn take_slices_typed<S, L>(
+    array: ArrayView<'_, Bool>,
+    starts: &ArrayRef,
+    lengths: &ArrayRef,
+    output_len: usize,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef>
+where
+    S: IntegerPType,
+    L: IntegerPType,
+{
+    let starts = starts.clone().execute::<PrimitiveArray>(ctx)?;
+    let lengths = lengths.clone().execute::<PrimitiveArray>(ctx)?;
+    let values = gather_bits(
+        &array.to_bit_buffer(),
+        starts.as_slice::<S>(),
+        lengths.as_slice::<L>(),
+        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<S, L>(
+    source: &BitBuffer,
+    starts: &[S],
+    lengths: &[L],
+    output_len: usize,
+) -> VortexResult<BitBuffer>
+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())
+}