Clarify docs in session accessors

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

Author: gatesn

vortex-array/src/arrays/filter/kernel.rs

@@ -2,7 +2,6 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_compute::filter::Filter;
-use vortex_dtype::DType;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 use vortex_vector::Vector;
@@ -15,13 +14,11 @@ use crate::kernel::PushDownResult;
 pub struct FilterKernel {
     child: KernelRef,
     mask: Mask,
-    // Used for estimating filter cost
-    dtype: DType,
 }
 
 impl FilterKernel {
-    pub fn new(child: KernelRef, mask: Mask, dtype: DType) -> Self {
-        Self { child, mask, dtype }
+    pub fn new(child: KernelRef, mask: Mask) -> Self {
+        Self { child, mask }
     }
 }
 
@@ -36,7 +33,6 @@ impl Kernel for FilterKernel {
             PushDownResult::NotPushed(k) => PushDownResult::NotPushed(Box::new(FilterKernel {
                 child: k,
                 mask: new_mask,
-                dtype: self.dtype.clone(),
             })),
             PushDownResult::Pushed(new_k) => PushDownResult::Pushed(new_k),
         })

vortex-array/src/arrays/filter/vtable.rs

@@ -122,11 +122,7 @@ impl VTable for FilterVTable {
         }
 
         // Otherwise, wrap up the child in a filter kernel.
-        Ok(Box::new(FilterKernel::new(
-            child,
-            mask,
-            array.dtype().clone(),
-        )))
+        Ok(Box::new(FilterKernel::new(child, mask)))
     }
 }
 

vortex-array/src/arrays/scalar_fn/kernel.rs

@@ -8,7 +8,6 @@ use vortex_vector::Datum;
 use vortex_vector::Scalar;
 use vortex_vector::Vector;
 
-use crate::arrays::FilterKernel;
 use crate::expr::ExecutionArgs;
 use crate::expr::ScalarFn;
 use crate::kernel::Kernel;
@@ -63,20 +62,14 @@ impl Kernel for ScalarFnKernel {
 
     fn push_down_filter(self: Box<Self>, selection: &Mask) -> VortexResult<PushDownResult> {
         let mut new_inputs = Vec::with_capacity(self.inputs.len());
-        for (input, dtype) in self.inputs.into_iter().zip(&self.input_dtypes) {
+        for input in self.inputs {
             match input {
                 KernelInput::Scalar(s) => {
                     new_inputs.push(KernelInput::Scalar(s.clone()));
                 }
-                KernelInput::Vector(k) => match k.push_down_filter(selection)? {
-                    PushDownResult::Pushed(new_k) => {
-                        new_inputs.push(KernelInput::Vector(new_k));
-                    }
-                    PushDownResult::NotPushed(k) => {
-                        let new_k = FilterKernel::new(k, selection.clone(), dtype.clone());
-                        new_inputs.push(KernelInput::Vector(Box::new(new_k)));
-                    }
-                },
+                KernelInput::Vector(k) => {
+                    new_inputs.push(KernelInput::Vector(k.force_push_down_filter(selection)?));
+                }
             }
         }
 

vortex-array/src/arrays/struct_/vtable/mod.rs

@@ -5,17 +5,21 @@ use std::sync::Arc;
 
 use itertools::Itertools;
 use vortex_buffer::BufferHandle;
+use vortex_compute::filter::Filter;
 use vortex_dtype::DType;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_mask::Mask;
+use vortex_vector::Vector;
 use vortex_vector::struct_::StructVector;
 
 use crate::EmptyMetadata;
 use crate::arrays::struct_::StructArray;
 use crate::kernel::BindCtx;
+use crate::kernel::Kernel;
 use crate::kernel::KernelRef;
-use crate::kernel::kernel;
+use crate::kernel::PushDownResult;
 use crate::serde::ArrayChildren;
 use crate::validity::Validity;
 use crate::vtable;
@@ -115,13 +119,42 @@ impl VTable for StructVTable {
             .try_collect()?;
         let validity_mask = array.validity_mask();
 
-        Ok(kernel(move || {
-            // SAFETY: we know that all field lengths match the struct array length, and the validity
-            let fields = fields.into_iter().map(|k| k.execute()).try_collect()?;
-            Ok(unsafe { StructVector::new_unchecked(Arc::new(fields), validity_mask) }.into())
+        Ok(Box::new(StructKernel {
+            fields,
+            validity_mask,
         }))
     }
 }
 
+#[derive(Debug)]
+struct StructKernel {
+    fields: Box<[KernelRef]>,
+    // TODO(ngates): hold a kernel that computes the mask.
+    validity_mask: Mask,
+}
+
+impl Kernel for StructKernel {
+    fn execute(self: Box<Self>) -> VortexResult<Vector> {
+        // SAFETY: we know that all field lengths match the struct array length, and the validity
+        let fields = self.fields.into_iter().map(|k| k.execute()).try_collect()?;
+        Ok(unsafe { StructVector::new_unchecked(Arc::new(fields), self.validity_mask) }.into())
+    }
+
+    fn push_down_filter(self: Box<Self>, selection: &Mask) -> VortexResult<PushDownResult> {
+        let fields = self
+            .fields
+            .into_iter()
+            .map(|k| k.force_push_down_filter(selection))
+            .try_collect()?;
+
+        let validity_mask = self.validity_mask.filter(selection);
+
+        Ok(PushDownResult::Pushed(Box::new(StructKernel {
+            fields,
+            validity_mask,
+        })))
+    }
+}
+
 #[derive(Debug)]
 pub struct StructVTable;

vortex-array/src/kernel/mod.rs

@@ -17,6 +17,8 @@ use vortex_mask::Mask;
 use vortex_session::VortexSession;
 use vortex_vector::Vector;
 
+use crate::arrays::FilterKernel;
+
 /// A boxed reference to a kernel.
 pub type KernelRef = Box<dyn Kernel>;
 
@@ -52,3 +54,13 @@ impl BindCtx {
         &self.session
     }
 }
+
+impl dyn Kernel + '_ {
+    /// Push-down a filter mask, or else wrap up the kernel to apply the filter later.
+    pub fn force_push_down_filter(self: Box<Self>, selection: &Mask) -> VortexResult<KernelRef> {
+        match self.push_down_filter(selection)? {
+            PushDownResult::Pushed(k) => Ok(k),
+            PushDownResult::NotPushed(k) => Ok(Box::new(FilterKernel::new(k, selection.clone()))),
+        }
+    }
+}