Use PiecewiseSequential indices for FSL take

Signed-off-by: Daniel King <dan@spiraldb.com>

Author: danking

vortex-array/src/arrays/bool/compute/take.rs

@@ -4,6 +4,7 @@
 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;
@@ -15,12 +16,17 @@ use crate::array::ArrayView;
 use crate::arrays::Bool;
 use crate::arrays::BoolArray;
 use crate::arrays::ConstantArray;
+use crate::arrays::PiecewiseSequential;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::bool::BoolArrayExt;
 use crate::arrays::dict::TakeExecute;
+use crate::arrays::piecewise_sequential::execute_index_arrays;
+use crate::arrays::piecewise_sequential::index_value_to_usize;
+use crate::arrays::piecewise_sequential::validate_index_ranges;
 use crate::builtins::ArrayBuiltins;
 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 +35,10 @@ impl TakeExecute for Bool {
         indices: &ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
+        if let Some(piecewise_indices) = indices.as_opt::<PiecewiseSequential>() {
+            return take_piecewise_sequential(array, piecewise_indices, indices, ctx).map(Some);
+        }
+
         let indices_nulls_zeroed = match indices.validity()?.execute_mask(indices.len(), ctx)? {
             Mask::AllTrue(_) => indices.clone(),
             Mask::AllFalse(_) => {
@@ -55,6 +65,27 @@ impl TakeExecute for Bool {
     }
 }
 
+fn take_piecewise_sequential(
+    array: ArrayView<'_, Bool>,
+    indices: ArrayView<'_, PiecewiseSequential>,
+    indices_ref: &ArrayRef,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef> {
+    let (starts, lengths) = execute_index_arrays(indices, ctx)?;
+    let buffer = match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
+        match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
+            take_piecewise_bits(
+                &array.to_bit_buffer(),
+                starts.as_slice::<S>(),
+                lengths.as_slice::<L>(),
+                indices_ref.len(),
+            )?
+        })
+    });
+
+    Ok(BoolArray::new(buffer, array.validity()?.take(indices_ref)?).into_array())
+}
+
 fn take_valid_indices<I: AsPrimitive<usize>>(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<bool>.
@@ -82,6 +113,28 @@ fn take_bool_impl<I: AsPrimitive<usize>>(bools: BitBufferView<'_>, indices: &[I]
     })
 }
 
+fn take_piecewise_bits<S, L>(
+    source: &BitBuffer,
+    starts: &[S],
+    lengths: &[L],
+    output_len: usize,
+) -> VortexResult<BitBuffer>
+where
+    S: crate::dtype::UnsignedPType,
+    L: crate::dtype::UnsignedPType,
+{
+    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);
+        let length = index_value_to_usize(length);
+        values.append_buffer(&source.slice(start..start + length));
+    }
+
+    Ok(values.freeze())
+}
+
 #[cfg(test)]
 mod test {
     use rstest::rstest;

vortex-array/src/arrays/decimal/compute/take.rs

@@ -1,28 +1,39 @@
 // 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::PiecewiseSequential;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::dict::TakeExecute;
+use crate::arrays::piecewise_sequential::execute_index_arrays;
+use crate::arrays::piecewise_sequential::index_value_to_usize;
+use crate::arrays::piecewise_sequential::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_integer_ptype;
+use crate::match_each_unsigned_integer_ptype;
 
 impl TakeExecute for Decimal {
     fn take(
         array: ArrayView<'_, Decimal>,
         indices: &ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
+        if let Some(piecewise_indices) = indices.as_opt::<PiecewiseSequential>() {
+            return take_piecewise_sequential(array, piecewise_indices, indices, ctx).map(Some);
+        }
+
         let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
         let validity = array.validity()?.take(&indices.clone().into_array())?;
 
@@ -42,10 +53,61 @@ impl TakeExecute for Decimal {
     }
 }
 
+fn take_piecewise_sequential(
+    array: ArrayView<'_, Decimal>,
+    indices: ArrayView<'_, PiecewiseSequential>,
+    indices_ref: &ArrayRef,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef> {
+    let (starts, lengths) = execute_index_arrays(indices, ctx)?;
+    match_each_decimal_value_type!(array.values_type(), |D| {
+        match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
+            match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
+                let values = take_piecewise_to_buffer::<S, L, D>(
+                    starts.as_slice::<S>(),
+                    lengths.as_slice::<L>(),
+                    array.buffer::<D>().as_slice(),
+                    indices_ref.len(),
+                )?;
+                let validity = array.validity()?.take(indices_ref)?;
+
+                // 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<I: IntegerPType, T: NativeDecimalType>(indices: &[I], values: &[T]) -> Buffer<T> {
     indices.iter().map(|idx| values[idx.as_()]).collect()
 }
 
+fn take_piecewise_to_buffer<S, L, T>(
+    starts: &[S],
+    lengths: &[L],
+    values: &[T],
+    output_len: usize,
+) -> VortexResult<Buffer<T>>
+where
+    S: crate::dtype::UnsignedPType,
+    L: crate::dtype::UnsignedPType,
+    T: NativeDecimalType,
+{
+    validate_index_ranges(values.len(), starts, lengths, output_len)?;
+
+    let mut result = BufferMut::<T>::with_capacity(output_len);
+    for (&start, &length) in starts.iter().zip_eq(lengths) {
+        let start = index_value_to_usize(start);
+        let length = index_value_to_usize(length);
+        result.extend_from_slice(&values[start..start + length]);
+    }
+
+    Ok(result.freeze())
+}
+
 #[cfg(test)]
 mod tests {
     use rstest::rstest;