add tests

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

Author: connortsui20

vortex-compute/src/take/mod.rs

@@ -3,8 +3,11 @@
 
 //! Take function.
 
+mod bit_buffer;
 mod buffer;
+mod mask;
 pub mod slice;
+mod vector;
 
 /// Function for taking based on indices (which can have different representations).
 pub trait Take<Indices: ?Sized> {

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

@@ -11,6 +11,9 @@ mod bool;
 mod primitive;
 mod pvector;
 
+#[cfg(test)]
+mod tests;
+
 impl<T> Take<PrimitiveVector> for &T
 where
     for<'a> &'a T: Take<PVector<u8>, Output = T>,

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

@@ -0,0 +1,136 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_vector::VectorMutOps;
+use vortex_vector::bool::BoolVector;
+use vortex_vector::bool::BoolVectorMut;
+use vortex_vector::primitive::PVector;
+use vortex_vector::primitive::PVectorMut;
+use vortex_vector::primitive::PrimitiveVector;
+
+use crate::take::Take;
+
+/// Tests `Take` on `PVector` with mixed validity in both data and indices.
+///
+/// This test covers:
+/// - Taking from a vector with some null values.
+/// - Using indices that have some null values.
+/// - Verifying that nulls in the data propagate correctly.
+/// - Verifying that nulls in indices result in null outputs.
+/// - Using non-sequential indices to test the general case.
+#[test]
+fn test_pvector_take_with_nullable_indices() {
+    // Data: [10, null, 30, 40, null, 60]
+    let data: PVectorMut<i32> = [Some(10), None, Some(30), Some(40), None, Some(60)]
+        .into_iter()
+        .collect();
+    let data = data.freeze();
+
+    // Indices: [0, null, 2, 5, null] (u32 indices with nulls)
+    let indices: PVectorMut<u32> = [Some(0), None, Some(2), Some(5), None]
+        .into_iter()
+        .collect();
+    let indices = indices.freeze();
+
+    let result = (&data).take(&indices);
+
+    // Expected: [10, null, 30, 60, null]
+    // - Index 0 -> data[0] = 10 (valid)
+    // - Index null -> null (null index produces null)
+    // - Index 2 -> data[2] = 30 (valid)
+    // - Index 5 -> data[5] = 60 (valid)
+    // - Index null -> null (null index produces null)
+    assert_eq!(result.get(0), Some(&10));
+    assert_eq!(result.get(1), None); // Null index.
+    assert_eq!(result.get(2), Some(&30));
+    assert_eq!(result.get(3), Some(&60));
+    assert_eq!(result.get(4), None); // Null index.
+}
+
+/// Tests `Take` on `PVector` using `PrimitiveVector` indices (type-erased).
+///
+/// This ensures the generic `Take<PrimitiveVector>` impl works correctly by dispatching to the
+/// typed implementation.
+#[test]
+fn test_pvector_take_with_primitive_vector_indices() {
+    // Data: [100, 200, null, 400, 500]
+    let data: PVectorMut<i64> = [Some(100), Some(200), None, Some(400), Some(500)]
+        .into_iter()
+        .collect();
+    let data = data.freeze();
+
+    // Indices as PrimitiveVector (u16): [4, 2, 0, 1]
+    let indices: PVectorMut<u16> = [4u16, 2, 0, 1].into_iter().collect();
+    let indices: PrimitiveVector = indices.freeze().into();
+
+    let result: PVector<i64> = (&data).take(&indices);
+
+    // Expected: [500, null, 100, 200]
+    assert_eq!(result.get(0), Some(&500));
+    assert_eq!(result.get(1), None); // data[2] is null.
+    assert_eq!(result.get(2), Some(&100));
+    assert_eq!(result.get(3), Some(&200));
+}
+
+/// Tests `Take` on `BoolVector` with mixed validity in both data and indices.
+///
+/// This test covers:
+/// - Taking from a boolean vector with some null values.
+/// - Using indices that have some null values.
+/// - Verifying that nulls in the data propagate correctly.
+/// - Verifying that nulls in indices result in null outputs.
+#[test]
+fn test_bool_vector_take_with_nullable_indices() {
+    // Data: [true, null, false, true, null, false]
+    let data: BoolVectorMut = [Some(true), None, Some(false), Some(true), None, Some(false)]
+        .into_iter()
+        .collect();
+    let data = data.freeze();
+
+    // Indices: [5, null, 0, 3, null, 2] (u32 indices with nulls)
+    let indices: PVectorMut<u32> = [Some(5), None, Some(0), Some(3), None, Some(2)]
+        .into_iter()
+        .collect();
+    let indices = indices.freeze();
+
+    let result = (&data).take(&indices);
+
+    // Expected: [false, null, true, true, null, false]
+    // - Index 5 -> data[5] = false (valid)
+    // - Index null -> null (null index produces null)
+    // - Index 0 -> data[0] = true (valid)
+    // - Index 3 -> data[3] = true (valid)
+    // - Index null -> null (null index produces null)
+    // - Index 2 -> data[2] = false (valid)
+    assert_eq!(result.get(0), Some(false));
+    assert_eq!(result.get(1), None); // Null index.
+    assert_eq!(result.get(2), Some(true));
+    assert_eq!(result.get(3), Some(true));
+    assert_eq!(result.get(4), None); // Null index.
+    assert_eq!(result.get(5), Some(false));
+}
+
+/// Tests `Take` on `BoolVector` using `PrimitiveVector` indices (type-erased).
+///
+/// This ensures the generic `Take<PrimitiveVector>` impl works correctly for `BoolVector`.
+#[test]
+fn test_bool_vector_take_with_primitive_vector_indices() {
+    // Data: [true, false, null, true, false]
+    let data: BoolVectorMut = [Some(true), Some(false), None, Some(true), Some(false)]
+        .into_iter()
+        .collect();
+    let data = data.freeze();
+
+    // Indices as PrimitiveVector (u64): [4, 2, 1, 0, 3]
+    let indices: PVectorMut<u64> = [4u64, 2, 1, 0, 3].into_iter().collect();
+    let indices: PrimitiveVector = indices.freeze().into();
+
+    let result: BoolVector = (&data).take(&indices);
+
+    // Expected: [false, null, false, true, true]
+    assert_eq!(result.get(0), Some(false)); // data[4]
+    assert_eq!(result.get(1), None); // data[2] is null.
+    assert_eq!(result.get(2), Some(false)); // data[1]
+    assert_eq!(result.get(3), Some(true)); // data[0]
+    assert_eq!(result.get(4), Some(true)); // data[3]
+}

vortex-vector/src/bool/vector.rs

@@ -73,6 +73,22 @@ impl BoolVector {
     pub fn bits(&self) -> &BitBuffer {
         &self.bits
     }
+
+    /// Gets a nullable element at the given index, panicking on out-of-bounds.
+    ///
+    /// If the element at the given index is null, returns `None`. Otherwise, returns `Some(x)`,
+    /// where `x: bool`.
+    ///
+    /// Note that this `get` method is different from the standard library [`slice::get`], which
+    /// returns `None` if the index is out of bounds. This method will panic if the index is out of
+    /// bounds, and return `None` if the element is null.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the index is out of bounds.
+    pub fn get(&self, index: usize) -> Option<bool> {
+        self.validity.value(index).then(|| self.bits.value(index))
+    }
 }
 
 impl VectorOps for BoolVector {