Filter pushdown

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

Author: gatesn

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

@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: Apache-2.0
+// 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;
+
+use crate::kernel::Kernel;
+use crate::kernel::KernelRef;
+use crate::kernel::PushDownResult;
+
+#[derive(Debug)]
+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 }
+    }
+}
+
+impl Kernel for FilterKernel {
+    fn execute(self: Box<Self>) -> VortexResult<Vector> {
+        Ok(Filter::filter(&self.child.execute()?, &self.mask))
+    }
+
+    fn cost_estimate(&self, selection: &Mask) -> f64 {
+        cost_for_dtype(&self.dtype, selection)
+    }
+
+    fn push_down_filter(self: Box<Self>, selection: &Mask) -> VortexResult<PushDownResult> {
+        let new_mask = self.mask.intersect_by_rank(selection);
+        Ok(match self.child.push_down_filter(&new_mask)? {
+            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),
+        })
+    }
+}
+
+fn cost_for_dtype(dtype: &DType, selection: &Mask) -> f64 {
+    match dtype {
+        DType::Null => 0.0,
+        DType::Extension(ext) => cost_for_dtype(ext.storage_dtype(), selection),
+        _ => f64::INFINITY,
+    }
+}

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

@@ -2,7 +2,9 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 mod array;
+mod kernel;
 mod vtable;
 
 pub use array::*;
+pub use kernel::*;
 pub use vtable::*;

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

@@ -7,26 +7,18 @@ use std::ops::Range;
 use vortex_buffer::BufferHandle;
 use vortex_compute::filter::Filter;
 use vortex_dtype::DType;
+use vortex_error::vortex_bail;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
 
-use crate::Array;
-use crate::ArrayBufferVisitor;
-use crate::ArrayChildVisitor;
-use crate::ArrayEq;
-use crate::ArrayHash;
-use crate::ArrayRef;
-use crate::Canonical;
-use crate::IntoArray;
-use crate::Precision;
-use crate::arrays::LEGACY_SESSION;
 use crate::arrays::filter::array::FilterArray;
+use crate::arrays::filter::kernel::FilterKernel;
+use crate::arrays::LEGACY_SESSION;
 use crate::kernel::BindCtx;
 use crate::kernel::KernelRef;
-use crate::kernel::kernel;
+use crate::kernel::PushDownResult;
 use crate::serde::ArrayChildren;
 use crate::stats::StatsSetRef;
 use crate::vectors::VectorIntoArray;
@@ -41,6 +33,15 @@ use crate::vtable::OperationsVTable;
 use crate::vtable::VTable;
 use crate::vtable::ValidityVTable;
 use crate::vtable::VisitorVTable;
+use crate::Array;
+use crate::ArrayBufferVisitor;
+use crate::ArrayChildVisitor;
+use crate::ArrayEq;
+use crate::ArrayHash;
+use crate::ArrayRef;
+use crate::Canonical;
+use crate::IntoArray;
+use crate::Precision;
 
 vtable!(Filter);
 
@@ -95,9 +96,40 @@ impl VTable for FilterVTable {
     }
 
     fn bind_kernel(array: &Self::Array, ctx: &mut BindCtx) -> VortexResult<KernelRef> {
-        let child = array.child.bind_kernel(ctx)?;
+        let mut child = array.child.bind_kernel(ctx)?;
         let mask = array.mask.clone();
-        Ok(kernel(move || Ok(Filter::filter(&child.execute()?, &mask))))
+
+        let full_cost = child.cost_estimate(&Mask::new_true(array.child.len()));
+        let pushdown_cost = child.cost_estimate(&mask);
+        log::debug!(
+            "Filter kernel cost estimate: full={}, pushdown={}",
+            full_cost,
+            pushdown_cost
+        );
+
+        if pushdown_cost < full_cost {
+            // Try to push down the filter to the child if it's cheaper.
+            child = match child.push_down_filter(&mask)? {
+                PushDownResult::Pushed(new_k) => {
+                    log::debug!("Filter push down kernel:\n{:?}", new_k);
+                    return Ok(new_k);
+                }
+                PushDownResult::NotPushed(child) => {
+                    log::debug!(
+                        "Filter pushdown was cheaper but not supported by child array {}",
+                        array.child.display_tree()
+                    );
+                    child
+                }
+            };
+        }
+
+        // Otherwise, wrap up the child in a filter kernel.
+        Ok(Box::new(FilterKernel::new(
+            child,
+            mask,
+            array.dtype().clone(),
+        )))
     }
 }
 

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

@@ -4,14 +4,17 @@
 use itertools::Itertools;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
+use vortex_mask::Mask;
 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;
 use crate::kernel::KernelRef;
+use crate::kernel::PushDownResult;
 
 /// A CPU kernel for executing scalar functions.
 #[derive(Debug)]
@@ -55,4 +58,46 @@ impl Kernel for ScalarFnKernel {
 
         Ok(self.scalar_fn.execute(args)?.ensure_vector(self.row_count))
     }
+
+    fn cost_estimate(&self, selection: &Mask) -> f64 {
+        let self_cost = self.scalar_fn.cost_estimate(selection);
+        let child_cost = self
+            .inputs
+            .iter()
+            .map(|input| match input {
+                KernelInput::Scalar(_) => 0.0,
+                KernelInput::Vector(k) => k.cost_estimate(selection),
+            })
+            .sum::<f64>();
+
+        self_cost + child_cost
+    }
+
+    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) {
+            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)));
+                    }
+                },
+            }
+        }
+
+        Ok(PushDownResult::Pushed(Box::new(ScalarFnKernel {
+            scalar_fn: self.scalar_fn,
+            inputs: new_inputs,
+            input_dtypes: self.input_dtypes,
+            row_count: self.row_count,
+            return_dtype: self.return_dtype,
+        })))
+    }
 }

vortex-array/src/arrays/scalar_fn/vtable/canonical.rs

@@ -4,15 +4,15 @@
 use itertools::Itertools;
 use vortex_error::VortexExpect;
 
-use crate::Array;
-use crate::Canonical;
-use crate::arrays::LEGACY_SESSION;
 use crate::arrays::scalar_fn::array::ScalarFnArray;
 use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
+use crate::arrays::LEGACY_SESSION;
 use crate::executor::VectorExecutor;
 use crate::expr::ExecutionArgs;
 use crate::vectors::VectorIntoArray;
 use crate::vtable::CanonicalVTable;
+use crate::Array;
+use crate::Canonical;
 
 impl CanonicalVTable<ScalarFnVTable> for ScalarFnVTable {
     fn canonicalize(array: &ScalarFnArray) -> Canonical {
@@ -34,12 +34,12 @@ impl CanonicalVTable<ScalarFnVTable> for ScalarFnVTable {
             return_dtype: array.dtype.clone(),
         };
 
+        let len = array.len;
         let result_vector = array
             .scalar_fn
             .execute(ctx)
             .vortex_expect("Canonicalize should be fallible")
-            .into_vector()
-            .vortex_expect("Canonicalize should return a vector");
+            .ensure_vector(len);
 
         result_vector.into_array(&array.dtype).to_canonical()
     }

vortex-array/src/expr/exprs/get_item.rs

@@ -9,17 +9,18 @@ use vortex_dtype::DType;
 use vortex_dtype::FieldName;
 use vortex_dtype::FieldPath;
 use vortex_dtype::Nullability;
+use vortex_error::vortex_err;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_error::vortex_err;
+use vortex_mask::Mask;
 use vortex_proto::expr as pb;
 use vortex_vector::Datum;
 use vortex_vector::ScalarOps;
 use vortex_vector::VectorOps;
 
-use crate::ArrayRef;
-use crate::ToCanonical;
 use crate::compute::mask;
+use crate::expr::exprs::root::root;
+use crate::expr::stats::Stat;
 use crate::expr::Arity;
 use crate::expr::ChildName;
 use crate::expr::ExecutionArgs;
@@ -30,8 +31,8 @@ use crate::expr::SimplifyCtx;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
-use crate::expr::exprs::root::root;
-use crate::expr::stats::Stat;
+use crate::ArrayRef;
+use crate::ToCanonical;
 
 pub struct GetItem;
 
@@ -189,6 +190,11 @@ impl VTable for GetItem {
         // If this type-checks its infallible.
         false
     }
+
+    fn cost_estimate(&self, _options: &Self::Options, _selection: &Mask) -> f64 {
+        // This is largely a metadata-only operation.
+        0.0
+    }
 }
 
 /// Creates an expression that accesses a field from the root array.
@@ -226,15 +232,15 @@ mod tests {
     use vortex_dtype::StructFields;
     use vortex_scalar::Scalar;
 
-    use crate::Array;
-    use crate::IntoArray;
     use crate::arrays::StructArray;
     use crate::expr::exprs::binary::checked_add;
     use crate::expr::exprs::get_item::get_item;
     use crate::expr::exprs::literal::lit;
     use crate::expr::exprs::pack::pack;
     use crate::expr::exprs::root::root;
     use crate::validity::Validity;
+    use crate::Array;
+    use crate::IntoArray;
 
     fn test_array() -> StructArray {
         StructArray::from_fields(&[

vortex-array/src/expr/exprs/merge.rs

@@ -11,17 +11,16 @@ use vortex_dtype::DType;
 use vortex_dtype::FieldNames;
 use vortex_dtype::Nullability;
 use vortex_dtype::StructFields;
+use vortex_error::vortex_bail;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
+use vortex_mask::Mask;
 use vortex_utils::aliases::hash_set::HashSet;
 use vortex_vector::Datum;
 
-use crate::Array;
-use crate::ArrayRef;
-use crate::IntoArray as _;
-use crate::ToCanonical;
 use crate::arrays::StructArray;
+use crate::expr::get_item;
+use crate::expr::pack;
 use crate::expr::Arity;
 use crate::expr::ChildName;
 use crate::expr::ExecutionArgs;
@@ -30,9 +29,11 @@ use crate::expr::Expression;
 use crate::expr::SimplifyCtx;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
-use crate::expr::get_item;
-use crate::expr::pack;
 use crate::validity::Validity;
+use crate::Array;
+use crate::ArrayRef;
+use crate::IntoArray as _;
+use crate::ToCanonical;
 
 /// Merge zero or more expressions that ALL return structs.
 ///
@@ -185,6 +186,11 @@ impl VTable for Merge {
         todo!()
     }
 
+    fn cost_estimate(&self, _options: &Self::Options, _selection: &Mask) -> f64 {
+        // This is largely a metadata-only operation.
+        0.0
+    }
+
     fn simplify(
         &self,
         options: &Self::Options,
@@ -298,21 +304,21 @@ mod tests {
     use vortex_dtype::PType::I64;
     use vortex_dtype::PType::U32;
     use vortex_dtype::PType::U64;
-    use vortex_error::VortexResult;
     use vortex_error::vortex_bail;
+    use vortex_error::VortexResult;
 
     use super::merge;
-    use crate::Array;
-    use crate::IntoArray;
-    use crate::ToCanonical;
     use crate::arrays::PrimitiveArray;
     use crate::arrays::StructArray;
-    use crate::expr::Expression;
-    use crate::expr::Pack;
     use crate::expr::exprs::get_item::get_item;
-    use crate::expr::exprs::merge::DuplicateHandling;
     use crate::expr::exprs::merge::merge_opts;
+    use crate::expr::exprs::merge::DuplicateHandling;
     use crate::expr::exprs::root::root;
+    use crate::expr::Expression;
+    use crate::expr::Pack;
+    use crate::Array;
+    use crate::IntoArray;
+    use crate::ToCanonical;
 
     fn primitive_field(array: &dyn Array, field_path: &[&str]) -> VortexResult<PrimitiveArray> {
         let mut field_path = field_path.iter();