Feature: Implement take on more vectors + stubs (#5653)

Tracking Issue: #5652

#5540 has to be merged first.

Implements `take` on the remaining vectors, except for `DecimalVector`,
`BinaryViewVector`, and `ListViewVector`

For `DecimalVector` we cannot trivially reuse the existing `take`
compute on slices and buffers because it has a bound of `NativePType`.
We can relax this bound to `Copy` instead, but that will require some
more work. Note that the current `DecimalArray` `take` implementation is
fully scalar right now so we can definitely improve that in the new
world.

For the view vectors I believe we want to think through how we "rebuild"
those in the array plan (we probably want to place rebuild nodes at good
places after just doing a `take` on the existing views).

So this PR implements:
- `FixedSizeListVector`
- `NullVector`
- `StructVector`

Edit: Now that #5666 has
merged, this also implements `DecimalVector`.

---------

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

Cargo.lock

@@ -134,7 +134,7 @@ jiff = "0.2.0"
 kanal = "0.1.1"
 lending-iterator = "0.1.7"
 libfuzzer-sys = "0.4"
-log = { version = "0.4" }
+log = { version = "0.4.21" }
 loom = { version = "0.7", features = ["checkpoint"] }
 memmap2 = "0.9.5"
 mimalloc = "0.1.42"

Cargo.toml

@@ -159,5 +159,9 @@ harness = false
 name = "take_primitive"
 harness = false
 
+[[bench]]
+name = "take_struct"
+harness = false
+
 [package.metadata.cargo-machete]
 ignored = ["getrandom_v03"]

vortex-array/Cargo.toml

@@ -0,0 +1,111 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![allow(clippy::unwrap_used)]
+
+use divan::Bencher;
+use rand::Rng;
+use rand::SeedableRng;
+use rand::distr::Uniform;
+use rand::rngs::StdRng;
+use vortex_array::IntoArray;
+use vortex_array::arrays::StructArray;
+use vortex_array::compute::take;
+use vortex_array::validity::Validity;
+use vortex_buffer::Buffer;
+use vortex_dtype::FieldNames;
+
+fn main() {
+    divan::main();
+}
+
+const ARRAY_SIZE: usize = 100_000;
+const TAKE_SIZE: usize = 1000;
+
+#[divan::bench]
+fn take_struct_simple(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(0);
+    let range = Uniform::new(0i64, 100_000_000).unwrap();
+
+    // Create single field for the struct
+    let field = (0..ARRAY_SIZE)
+        .map(|_| rng.sample(range))
+        .collect::<Buffer<i64>>()
+        .into_array();
+
+    let struct_array = StructArray::try_new(
+        FieldNames::from(["value"]),
+        vec![field],
+        ARRAY_SIZE,
+        Validity::NonNullable,
+    )
+    .unwrap();
+
+    let indices: Buffer<u64> = (0..TAKE_SIZE)
+        .map(|_| rng.random_range(0..ARRAY_SIZE) as u64)
+        .collect();
+    let indices_array = indices.into_array();
+
+    bencher
+        .with_inputs(|| (&struct_array, &indices_array))
+        .bench_refs(|(array, indices)| take(array.as_ref(), indices.as_ref()).unwrap());
+}
+
+#[divan::bench(args = [8])]
+fn take_struct_wide(bencher: Bencher, width: usize) {
+    let mut rng = StdRng::seed_from_u64(0);
+    let range = Uniform::new(0i64, 100_000_000).unwrap();
+
+    let fields: Vec<_> = (0..width)
+        .map(|_| {
+            (0..ARRAY_SIZE)
+                .map(|_| rng.sample(range))
+                .collect::<Buffer<i64>>()
+                .into_array()
+        })
+        .collect();
+
+    let field_names = FieldNames::from([
+        "field1", "field2", "field3", "field4", "field5", "field6", "field7", "field8",
+    ]);
+
+    let struct_array =
+        StructArray::try_new(field_names, fields, ARRAY_SIZE, Validity::NonNullable).unwrap();
+
+    let indices: Buffer<u64> = (0..TAKE_SIZE)
+        .map(|_| rng.random_range(0..ARRAY_SIZE) as u64)
+        .collect();
+    let indices_array = indices.into_array();
+
+    bencher
+        .with_inputs(|| (&struct_array, &indices_array))
+        .bench_refs(|(array, indices)| take(array.as_ref(), indices.as_ref()).unwrap());
+}
+
+#[divan::bench]
+fn take_struct_sequential_indices(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(0);
+    let range = Uniform::new(0i64, 100_000_000).unwrap();
+
+    // Create single field for the struct
+    let field = (0..ARRAY_SIZE)
+        .map(|_| rng.sample(range))
+        .collect::<Buffer<i64>>()
+        .into_array();
+
+    let struct_array = StructArray::try_new(
+        FieldNames::from(["value"]),
+        vec![field],
+        ARRAY_SIZE,
+        Validity::NonNullable,
+    )
+    .unwrap();
+
+    // Sequential indices for better cache performance
+    let indices: Buffer<u64> = (0..TAKE_SIZE as u64).collect();
+    let indices_array = indices.into_array();
+
+    bencher
+        .with_inputs(|| (&struct_array, &indices_array))
+        .bench_refs(|(array, indices)| take(array.as_ref(), indices.as_ref()).unwrap());
+}

vortex-array/benches/take_struct.rs

@@ -30,6 +30,8 @@ arrow-array = { workspace = true, optional = true }
 arrow-buffer = { workspace = true, optional = true }
 arrow-schema = { workspace = true, optional = true }
 half = { workspace = true }
+itertools = { workspace = true }
+log = { workspace = true }
 multiversion = { workspace = true }
 num-traits = { workspace = true }
 paste = { workspace = true }

vortex-compute/Cargo.toml

@@ -58,3 +58,29 @@ fn take_bool<I: UnsignedPType>(bools: &BitBuffer, indices: &[I]) -> BitBuffer {
         get_bit(buffer, offset + bool_idx)
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::take::Take;
+
+    #[test]
+    fn test_take_bit_buffer_take_small_and_large() {
+        use vortex_buffer::BitBuffer;
+
+        // Small buffer (uses take_byte_bool path).
+        let small: BitBuffer = [true, false, true, true, false, true, false, false]
+            .into_iter()
+            .collect();
+        let result = (&small).take(&[7u32, 0, 2, 5, 1][..]);
+
+        let values: Vec<bool> = (0..result.len()).map(|i| result.value(i)).collect();
+        assert_eq!(values, vec![false, true, true, true, false]);
+
+        // Large buffer (uses take_bool path, len > 4096).
+        let large: BitBuffer = (0..5000).map(|i| i % 3 == 0).collect();
+        let result = (&large).take(&[4999u32, 0, 1, 2, 3, 4998][..]);
+
+        let values: Vec<bool> = (0..result.len()).map(|i| result.value(i)).collect();
+        assert_eq!(values, vec![false, true, false, false, true, true]);
+    }
+}

vortex-compute/src/take/bit_buffer.rs

@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+// use std::ops::Deref;
+
+// use num_traits::AsPrimitive;
+// use vortex_buffer::Buffer;
+use vortex_dtype::UnsignedPType;
+use vortex_vector::VectorOps;
+// use vortex_vector::binaryview::BinaryView;
+use vortex_vector::binaryview::BinaryViewType;
+use vortex_vector::binaryview::BinaryViewVector;
+use vortex_vector::primitive::PVector;
+
+use crate::take::Take;
+
+impl<T: BinaryViewType, I: UnsignedPType> Take<PVector<I>> for &BinaryViewVector<T> {
+    type Output = BinaryViewVector<T>;
+
+    fn take(self, indices: &PVector<I>) -> BinaryViewVector<T> {
+        if indices.validity().all_true() {
+            self.take(indices.elements().as_slice())
+        } else {
+            take_nullable(self, indices)
+        }
+    }
+}
+
+impl<T: BinaryViewType, I: UnsignedPType> Take<[I]> for &BinaryViewVector<T> {
+    type Output = BinaryViewVector<T>;
+
+    fn take(self, _indices: &[I]) -> BinaryViewVector<T> {
+        todo!("TODO(connor): Implement `take` for `BinaryViewVector` and figure out rebuilding");
+
+        /*
+
+        let taken_views = take_views(self.views(), indices);
+        let taken_validity = self.validity().take(indices);
+
+        debug_assert_eq!(taken_views.len(), taken_validity.len());
+
+        // SAFETY: We called take on views and validity with the same indices, so the new components
+        // must have the same length. The views still point into the same buffers which we clone via
+        // Arc, so all view references remain valid.
+        unsafe {
+            BinaryViewVector::new_unchecked(taken_views, self.buffers().clone(), taken_validity)
+        }
+
+        */
+    }
+}
+
+fn take_nullable<T: BinaryViewType, I: UnsignedPType>(
+    _bvector: &BinaryViewVector<T>,
+    _indices: &PVector<I>,
+) -> BinaryViewVector<T> {
+    todo!("TODO(connor): Implement `take` for `BinaryViewVector` and figure out rebuilding");
+
+    /*
+
+    // We ignore nullability when taking the views since we can let the `Mask` implementation
+    // determine which elements are null.
+    let taken_views = take_views(bvector.views(), indices.elements().as_slice());
+    let taken_validity = bvector.validity().take(indices);
+
+    debug_assert_eq!(taken_views.len(), taken_validity.len());
+
+    // SAFETY: We used the same indices to take from both components, so they should still have the
+    // same length. The views still point into the same buffers which we clone via Arc, so all view
+    // references remain valid.
+    unsafe {
+        BinaryViewVector::new_unchecked(taken_views, bvector.buffers().clone(), taken_validity)
+    }
+
+    */
+}
+
+/*
+
+/// Takes views at the given indices.
+fn take_views<I: AsPrimitive<usize>>(
+    views: &Buffer<BinaryView>,
+    indices: &[I],
+) -> Buffer<BinaryView> {
+    let views_ref = views.deref();
+    Buffer::<BinaryView>::from_trusted_len_iter(indices.iter().map(|i| views_ref[(*i).as_()]))
+}
+
+*/

vortex-compute/src/take/vector/binaryview.rs

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::UnsignedPType;
+use vortex_vector::VectorOps;
+use vortex_vector::decimal::DecimalVector;
+use vortex_vector::match_each_dvector;
+use vortex_vector::primitive::PVector;
+
+use crate::take::Take;
+
+impl<I: UnsignedPType> Take<PVector<I>> for &DecimalVector {
+    type Output = DecimalVector;
+
+    fn take(self, indices: &PVector<I>) -> DecimalVector {
+        // If all the indices are valid, we can delegate to the slice indices implementation.
+        if indices.validity().all_true() {
+            return self.take(indices.elements().as_slice());
+        }
+
+        match_each_dvector!(self, |v| { v.take(indices).into() })
+    }
+}
+
+impl<I: UnsignedPType> Take<[I]> for &DecimalVector {
+    type Output = DecimalVector;
+
+    fn take(self, indices: &[I]) -> DecimalVector {
+        match_each_dvector!(self, |v| { v.take(indices).into() })
+    }
+}

vortex-compute/src/take/vector/decimal.rs

@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::NativeDecimalType;
+use vortex_dtype::UnsignedPType;
+use vortex_vector::VectorOps;
+use vortex_vector::decimal::DVector;
+use vortex_vector::primitive::PVector;
+
+use crate::take::Take;
+
+impl<D: NativeDecimalType, I: UnsignedPType> Take<PVector<I>> for &DVector<D> {
+    type Output = DVector<D>;
+
+    fn take(self, indices: &PVector<I>) -> DVector<D> {
+        if indices.validity().all_true() {
+            self.take(indices.elements().as_slice())
+        } else {
+            take_nullable(self, indices)
+        }
+    }
+}
+
+impl<D: NativeDecimalType, I: UnsignedPType> Take<[I]> for &DVector<D> {
+    type Output = DVector<D>;
+
+    fn take(self, indices: &[I]) -> DVector<D> {
+        let taken_elements = self.elements().take(indices);
+        let taken_validity = self.validity().take(indices);
+
+        debug_assert_eq!(taken_elements.len(), taken_validity.len());
+
+        // SAFETY: We called take on both components of the vector with the same indices, so the new
+        // components must have the same length. The elements are unchanged, so they must still be
+        // within the precision/scale bounds.
+        unsafe { DVector::new_unchecked(self.precision_scale(), taken_elements, taken_validity) }
+    }
+}
+
+fn take_nullable<D: NativeDecimalType, I: UnsignedPType>(
+    dvector: &DVector<D>,
+    indices: &PVector<I>,
+) -> DVector<D> {
+    // We ignore nullability when taking the elements since we can let the `Mask` implementation
+    // determine which elements are null.
+    let taken_elements = dvector.elements().take(indices.elements().as_slice());
+    let taken_validity = dvector.validity().take(indices);
+
+    debug_assert_eq!(taken_elements.len(), taken_validity.len());
+
+    // SAFETY: We used the same indices to take from both components, so they should still have the
+    // same length. The elements are unchanged, so they must still be within the precision/scale
+    // bounds.
+    unsafe { DVector::new_unchecked(dvector.precision_scale(), taken_elements, taken_validity) }
+}