AVX2 take handles indices that are equal to the index type max value (#7359)

The in the avx2 take code correctly states that we need to fitler out
values
greater than but then it also filters out values that are equal, this
fixes the
edge case for values that are equal to the max index value

---------

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

vortex-array/src/arrays/primitive/compute/take/avx2.rs

@@ -14,6 +14,7 @@ use std::arch::x86_64::_mm_setzero_si128;
 use std::arch::x86_64::_mm_shuffle_epi32;
 use std::arch::x86_64::_mm_storeu_si128;
 use std::arch::x86_64::_mm_unpacklo_epi64;
+use std::arch::x86_64::_mm256_andnot_si256;
 use std::arch::x86_64::_mm256_cmpgt_epi32;
 use std::arch::x86_64::_mm256_cmpgt_epi64;
 use std::arch::x86_64::_mm256_cvtepu8_epi32;
@@ -118,7 +119,6 @@ where
 /// # Safety
 ///
 /// The caller must ensure the `avx2` feature is enabled.
-#[allow(dead_code, unused_variables, reason = "TODO(connor): Implement this")]
 #[target_feature(enable = "avx2")]
 #[doc(hidden)]
 unsafe fn take_avx2<V: NativePType, I: UnsignedPType>(buffer: &[V], indices: &[I]) -> Buffer<V> {
@@ -135,6 +135,10 @@ unsafe fn take_avx2<V: NativePType, I: UnsignedPType>(buffer: &[V], indices: &[I
         }};
     }
 
+    if buffer.is_empty() {
+        return Buffer::zeroed(indices.len());
+    }
+
     match (I::PTYPE, V::PTYPE) {
         // Int value types. Only 32 and 64 bit types are supported.
         (PType::U8, PType::I32) => dispatch_avx2!(u8, i32),
@@ -223,7 +227,7 @@ macro_rules! impl_gather {
                     // Create a vec of the max idx.
                     let max_idx_vec = unsafe { $splat(max_idx as _) };
                     // Create a mask for valid indices (where the max_idx > provided index).
-                    let invalid_mask = unsafe { $mask_indices(max_idx_vec, indices_vec) };
+                    let invalid_mask = unsafe { _mm256_andnot_si256($mask_indices(indices_vec, max_idx_vec), $splat(-1)) };
                     let invalid_mask = {
                         let $mask_var = invalid_mask;
                         $mask_cvt
@@ -579,4 +583,22 @@ mod avx2_tests {
         index_type => u64,
         value_types => u32, i32, u64, i64, f32, f64
     );
+
+    #[test]
+    fn test_avx2_take_last_valid_index_u8() {
+        let values: Vec<i64> = (0..(255 + 1)).collect();
+        let indices: Vec<u8> = vec![255; 20];
+
+        let result = unsafe { take_avx2(&values, &indices) };
+        assert_eq!(&vec![255; indices.len()], result.as_slice());
+    }
+
+    #[test]
+    fn test_avx2_take_last_valid_index_u16() {
+        let values: Vec<i64> = (0..(65535 + 1)).collect();
+        let indices: Vec<u16> = vec![65535; 20];
+
+        let result = unsafe { take_avx2(&values, &indices) };
+        assert_eq!(&vec![65535; indices.len()], result.as_slice());
+    }
 }