make take generic when indices is a primitive vector

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

Author: connortsui20

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

@@ -1,5 +1,26 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_vector::match_each_unsigned_pvector;
+use vortex_vector::primitive::PVector;
+use vortex_vector::primitive::PrimitiveVector;
+
+use crate::take::Take;
+
+mod bool;
 mod primitive;
 mod pvector;
+
+impl<T> Take<PrimitiveVector> for &T
+where
+    for<'a> &'a T: Take<PVector<u8>, Output = T>,
+    for<'a> &'a T: Take<PVector<u16>, Output = T>,
+    for<'a> &'a T: Take<PVector<u32>, Output = T>,
+    for<'a> &'a T: Take<PVector<u64>, Output = T>,
+{
+    type Output = T;
+
+    fn take(self, indices: &PrimitiveVector) -> T {
+        match_each_unsigned_pvector!(indices, |iv| { self.take(iv) })
+    }
+}

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

@@ -9,14 +9,6 @@ use vortex_vector::primitive::PrimitiveVector;
 
 use crate::take::Take;
 
-impl Take<PrimitiveVector> for &PrimitiveVector {
-    type Output = PrimitiveVector;
-
-    fn take(self, indices: &PrimitiveVector) -> PrimitiveVector {
-        match_each_pvector!(self, |v| { v.take(indices).into() })
-    }
-}
-
 impl<I: UnsignedPType> Take<PVector<I>> for &PrimitiveVector {
     type Output = PrimitiveVector;
 

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

@@ -4,20 +4,10 @@
 use vortex_dtype::NativePType;
 use vortex_dtype::UnsignedPType;
 use vortex_vector::VectorOps;
-use vortex_vector::match_each_unsigned_pvector;
 use vortex_vector::primitive::PVector;
-use vortex_vector::primitive::PrimitiveVector;
 
 use crate::take::Take;
 
-impl<T: NativePType> Take<PrimitiveVector> for &PVector<T> {
-    type Output = PVector<T>;
-
-    fn take(self, indices: &PrimitiveVector) -> PVector<T> {
-        match_each_unsigned_pvector!(indices, |iv| { self.take(iv) })
-    }
-}
-
 impl<T: NativePType, I: UnsignedPType> Take<PVector<I>> for &PVector<T> {
     type Output = PVector<T>;