add take struct benchmark back

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

Author: connortsui20

vortex-array/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/benches/take_struct.rs

@@ -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-compute/src/take/bit_buffer.rs

@@ -64,7 +64,7 @@ mod tests {
     use crate::take::Take;
 
     #[test]
-    fn test_bit_buffer_take_small_and_large() {
+    fn test_take_bit_buffer_take_small_and_large() {
         use vortex_buffer::BitBuffer;
 
         // Small buffer (uses take_byte_bool path).

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

@@ -24,11 +24,7 @@ impl<D: NativeDecimalType, I: UnsignedPType> Take<PVector<I>> for &DVector<D> {
 impl<D: NativeDecimalType, I: UnsignedPType> Take<[I]> for &DVector<D> {
     type Output = DVector<D>;
 
-    fn take(self, _indices: &[I]) -> DVector<D> {
-        todo!("TODO(connor): Implement `take` for `DVector` and figure out trait bounds");
-
-        /*
-
+    fn take(self, indices: &[I]) -> DVector<D> {
         let taken_elements = self.elements().take(indices);
         let taken_validity = self.validity().take(indices);
 
@@ -38,19 +34,13 @@ impl<D: NativeDecimalType, I: UnsignedPType> Take<[I]> for &DVector<D> {
         // 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: &DVector<D>,
+    indices: &PVector<I>,
 ) -> DVector<D> {
-    todo!("TODO(connor): Implement `take` for `DVector` and figure out trait bounds");
-
-    /*
-
     // 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());
@@ -62,6 +52,4 @@ fn take_nullable<D: NativeDecimalType, I: UnsignedPType>(
     // 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) }
-
-    */
 }