Feature: add a take compute trait (#5503)

In the new operator world, `take` as a compute function just becomes a
`DictArray` (since arrays are lazy), so it is nice to have `take` as a
standalone trait (like `filter`).

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

Author: connortsui20

vortex-compute/src/filter/mod.rs

@@ -11,16 +11,16 @@ mod slice_mut;
 mod vector;
 
 /// Function for filtering based on a selection mask.
-pub trait Filter<By: ?Sized> {
+pub trait Filter<Selection: ?Sized> {
     /// The result type after performing the operation.
     type Output;
 
-    /// Filters the vector using the provided mask, returning a new value.
+    /// Filters an object using the provided mask, returning a new value.
     ///
     /// The result value will have length equal to the true count of the provided mask.
     ///
     /// # Panics
     ///
     /// If the length of the mask does not equal the length of the value being filtered.
-    fn filter(self, selection: &By) -> Self::Output;
+    fn filter(self, selection: &Selection) -> Self::Output;
 }

vortex-compute/src/lib.rs

@@ -16,3 +16,4 @@ pub mod expand;
 pub mod filter;
 pub mod logical;
 pub mod mask;
+pub mod take;

vortex-compute/src/take/mod.rs

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Take function.
+
+/// Function for taking based on indices (which can have different representations).
+pub trait Take<Indices: ?Sized> {
+    /// The result type after performing the operation.
+    type Output;
+
+    /// Creates a new object using the elements from the input indexed by `indices`.
+    ///
+    /// For example, if we have an array `[1, 2, 3, 4, 5]` and `indices` `[4, 2]`, the resulting
+    /// array would be `[5, 3]`.
+    ///
+    /// The output should have the same length as the `indices`.
+    ///
+    /// # Panics
+    ///
+    /// This should panic if an index in `indices` is out-of-bounds with respect to `self`.
+    fn take(self, indices: &Indices) -> Self::Output;
+}