Lazy take for operators (#5692)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/array/mod.rs

@@ -31,6 +31,7 @@ use crate::DynArrayHash;
 use crate::arrays::BoolVTable;
 use crate::arrays::ConstantVTable;
 use crate::arrays::DecimalVTable;
+use crate::arrays::DictArray;
 use crate::arrays::ExtensionVTable;
 use crate::arrays::FilterArray;
 use crate::arrays::FixedSizeListVTable;
@@ -96,8 +97,11 @@ pub trait Array:
     /// Performs a constant-time slice of the array.
     fn slice(&self, range: Range<usize>) -> ArrayRef;
 
-    /// Performs a constant-time filter of the array.
-    fn filter(&self, mask: &Mask) -> VortexResult<ArrayRef>;
+    /// Wraps the array in a [`FilterArray`] such that it is logically filtered by the given mask.
+    fn filter(&self, mask: Mask) -> VortexResult<ArrayRef>;
+
+    /// Wraps the array in a [`DictArray`] such that it is logically taken by the given indices.
+    fn take(&self, indices: ArrayRef) -> VortexResult<ArrayRef>;
 
     /// Fetch the scalar at the given index.
     ///
@@ -232,10 +236,14 @@ impl Array for Arc<dyn Array> {
         self.as_ref().slice(range)
     }
 
-    fn filter(&self, mask: &Mask) -> VortexResult<ArrayRef> {
+    fn filter(&self, mask: Mask) -> VortexResult<ArrayRef> {
         self.as_ref().filter(mask)
     }
 
+    fn take(&self, indices: ArrayRef) -> VortexResult<ArrayRef> {
+        self.as_ref().take(indices)
+    }
+
     #[inline]
     fn scalar_at(&self, index: usize) -> Scalar {
         self.as_ref().scalar_at(index)
@@ -505,10 +513,13 @@ impl<V: VTable> Array for ArrayAdapter<V> {
         sliced
     }
 
-    fn filter(&self, mask: &Mask) -> VortexResult<ArrayRef> {
+    fn filter(&self, mask: Mask) -> VortexResult<ArrayRef> {
         vortex_ensure!(self.len() == mask.len(), "Filter mask length mismatch");
-        Ok(V::filter(&self.0, mask)?
-            .unwrap_or_else(|| FilterArray::new(self.to_array(), mask.clone()).into_array()))
+        Ok(FilterArray::new(self.to_array(), mask).into_array())
+    }
+
+    fn take(&self, indices: ArrayRef) -> VortexResult<ArrayRef> {
+        Ok(DictArray::try_new(indices, self.to_array())?.into_array())
     }
 
     fn scalar_at(&self, index: usize) -> Scalar {

vortex-array/src/vtable/mod.rs

@@ -26,10 +26,8 @@ pub use visitor::*;
 use vortex_buffer::BufferHandle;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
-use vortex_mask::Mask;
 
 use crate::Array;
-use crate::ArrayRef;
 use crate::IntoArray;
 use crate::VectorExecutor;
 use crate::kernel::BindCtx;
@@ -152,21 +150,6 @@ pub trait VTable: 'static + Sized + Send + Sync + Debug {
             canonical.into_array().execute_vector(&session)
         }))
     }
-
-    /// Return an array filtered using the given mask.
-    ///
-    /// This should be implemented if the array is able to push down filtering operations to its
-    /// children. If not, return `Ok(None)` and the caller will handle filtering.
-    ///
-    /// NOTE: in the future, filter push-down will happen over the physical execution plan in
-    /// order to avoid cases where pushing down filters actually becomes more expensive. This
-    /// happens because filtering is very slow, and performing it in every leaf of a tree
-    /// can be more expensive than filtering once at the root.
-    fn filter(array: &Self::Array, mask: &Mask) -> VortexResult<Option<ArrayRef>> {
-        _ = array;
-        _ = mask;
-        Ok(None)
-    }
 }
 
 /// Placeholder type used to indicate when a particular vtable is not supported by the encoding.

vortex-layout/src/layouts/dict/reader.rs

@@ -21,6 +21,7 @@ use vortex_array::compute::min_max;
 use vortex_array::compute::take;
 use vortex_array::expr::Expression;
 use vortex_array::expr::root;
+use vortex_array::mask::MaskExecutor;
 use vortex_array::session::ArraySessionExt;
 use vortex_dtype::DType;
 use vortex_dtype::FieldMask;
@@ -183,35 +184,41 @@ impl LayoutReader for DictReader {
             MaskFuture::new_true(mask.len()),
         )?;
 
+        let session = self.session.clone();
+
         Ok(MaskFuture::new(mask.len(), async move {
             // Join on the I/O futures first, before the mask.
             let (codes, values) = try_join!(codes_eval, values_eval.map_err(VortexError::from))?;
             let mask = mask.await?;
 
-            // Short-circuit when the values are all true/false.
-            if values.all_valid()
-                && let Some(MinMaxResult { min, max }) = min_max(&values)?
-            {
-                #[expect(clippy::bool_comparison, reason = "easy to follow")]
-                if max.as_bool().value().vortex_expect("non null") == false {
-                    // All values are false
-                    return Ok(Mask::AllFalse(mask.len()));
-                }
-                #[expect(clippy::bool_comparison, reason = "easy to follow")]
-                if min.as_bool().value().vortex_expect("not null") == true {
-                    // All values are true, but we still need to respect codes validity
-                    return Ok(mask.bitand(&codes.validity_mask()));
+            let dict_mask = if *USE_VORTEX_OPERATORS {
+                values.take(codes)?.execute_mask_optimized(&session)?
+            } else {
+                // Short-circuit when the values are all true/false.
+                if values.all_valid()
+                    && let Some(MinMaxResult { min, max }) = min_max(&values)?
+                {
+                    #[expect(clippy::bool_comparison, reason = "easy to follow")]
+                    if max.as_bool().value().vortex_expect("non null") == false {
+                        // All values are false
+                        return Ok(Mask::AllFalse(mask.len()));
+                    }
+                    #[expect(clippy::bool_comparison, reason = "easy to follow")]
+                    if min.as_bool().value().vortex_expect("not null") == true {
+                        // All values are true, but we still need to respect codes validity
+                        return Ok(mask.bitand(&codes.validity_mask()));
+                    }
                 }
-            }
 
-            // Creating a mask from the dict array would canonicalize it,
-            // it should be fine for now as long as values is already canonical,
-            // so different row ranges do not canonicalize to the same array
-            // multiple times.
-            // TODO(joe): fixme casting null to false is *VERY* unsound, if the expression in the filter
-            // can inspect nulls (e.g. `is_null`).
-            // See `FlatEvaluation` for more details.
-            let dict_mask = take(&values, &codes)?.try_to_mask_fill_null_false()?;
+                // Creating a mask from the dict array would canonicalize it,
+                // it should be fine for now as long as values is already canonical,
+                // so different row ranges do not canonicalize to the same array
+                // multiple times.
+                // TODO(joe): fixme casting null to false is *VERY* unsound, if the expression in the filter
+                // can inspect nulls (e.g. `is_null`).
+                // See `FlatEvaluation` for more details.
+                take(&values, &codes)?.try_to_mask_fill_null_false()?
+            };
 
             Ok(mask.bitand(&dict_mask))
         }))

vortex-layout/src/layouts/flat/reader.rs

@@ -232,7 +232,7 @@ impl LayoutReader for FlatReader {
 
                 // Filter the array based on the row mask.
                 if !mask.all_true() {
-                    array = array.filter(&mask)?;
+                    array = array.filter(mask)?;
                 }
 
                 tracing::debug!("Project Array:\n{}", array.display_tree());