Clarify docs in session accessors

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

Author: gatesn

Cargo.lock

@@ -24,6 +24,7 @@ vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }
 vortex-mask = { workspace = true }
 vortex-scalar = { workspace = true }
+vortex-session = { workspace = true }
 
 [lints]
 workspace = true

encodings/runend/Cargo.toml

@@ -11,6 +11,7 @@ pub mod compress;
 mod compute;
 mod iter;
 mod ops;
+mod rules;
 
 #[doc(hidden)]
 pub mod _benchmarking {
@@ -23,11 +24,17 @@ pub mod _benchmarking {
 use vortex_array::ArrayBufferVisitor;
 use vortex_array::ArrayChildVisitor;
 use vortex_array::Canonical;
+use vortex_array::session::ArraySession;
+use vortex_array::session::ArraySessionExt;
+use vortex_array::vtable::ArrayVTableExt;
 use vortex_array::vtable::EncodeVTable;
 use vortex_array::vtable::VisitorVTable;
 use vortex_error::VortexResult;
+use vortex_session::SessionExt;
+use vortex_session::VortexSession;
 
 use crate::compress::runend_encode;
+use crate::rules::RunEndScalarFnRule;
 
 impl EncodeVTable<RunEndVTable> for RunEndVTable {
     fn encode(
@@ -59,6 +66,15 @@ impl VisitorVTable<RunEndVTable> for RunEndVTable {
     }
 }
 
+/// Initialize run-end encoding in the given session.
+pub fn initialize(session: &mut VortexSession) {
+    session.arrays().register(RunEndVTable.as_vtable());
+    session
+        .get_mut::<ArraySession>()
+        .optimizer_mut()
+        .register_parent_rule(RunEndScalarFnRule);
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_array::ProstMetadata;

encodings/runend/src/lib.rs

@@ -0,0 +1,85 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::arrays::AnyScalarFn;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::ConstantVTable;
+use vortex_array::arrays::ScalarFnArray;
+use vortex_array::optimizer::rules::ArrayParentReduceRule;
+use vortex_array::optimizer::rules::Exact;
+use vortex_dtype::DType;
+use vortex_error::VortexResult;
+
+use crate::RunEndArray;
+use crate::RunEndVTable;
+
+/// A rule to push down scalar functions through run-end encoding into the values array.
+///
+/// This only works if all other children of the scalar function array are constants.
+#[derive(Debug)]
+pub(crate) struct RunEndScalarFnRule;
+
+impl ArrayParentReduceRule<Exact<RunEndVTable>, AnyScalarFn> for RunEndScalarFnRule {
+    fn child(&self) -> Exact<RunEndVTable> {
+        Exact::from(&RunEndVTable)
+    }
+
+    fn parent(&self) -> AnyScalarFn {
+        AnyScalarFn
+    }
+
+    fn reduce_parent(
+        &self,
+        run_end: &RunEndArray,
+        parent: &ScalarFnArray,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        for (idx, child) in parent.children().iter().enumerate() {
+            if idx == child_idx {
+                // Skip ourselves
+                continue;
+            }
+
+            if !child.is::<ConstantVTable>() {
+                // We can only push down if all other children are constants
+                return Ok(None);
+            }
+        }
+
+        // TODO(ngates): relax this constraint and implement run-end decoding for all vector types.
+        if !matches!(parent.dtype(), DType::Bool(_) | DType::Primitive(..)) {
+            return Ok(None);
+        }
+
+        let values_len = run_end.values().len();
+        let mut new_children = parent.children();
+        for (idx, child) in new_children.iter_mut().enumerate() {
+            if idx == child_idx {
+                // Replace ourselves with run end values
+                *child = run_end.values().clone();
+                continue;
+            }
+
+            // Replace other children with their constant scalar value with length adjusted
+            // to the length of the run end values.
+            let constant = child.as_::<ConstantVTable>();
+            *child = ConstantArray::new(constant.scalar().clone(), values_len).into_array();
+        }
+
+        let new_values =
+            ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children, values_len)?
+                .into_array();
+
+        Ok(Some(
+            RunEndArray::try_new_offset_length(
+                run_end.ends().clone(),
+                new_values,
+                run_end.offset(),
+                run_end.len(),
+            )?
+            .into_array(),
+        ))
+    }
+}

encodings/runend/src/rules.rs

@@ -9,6 +9,7 @@ use vortex_dtype::PType;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
+use vortex_mask::Mask;
 use vortex_vector::Vector;
 use vortex_vector::VectorMut;
 use vortex_vector::VectorMutOps;
@@ -20,6 +21,7 @@ use crate::arrays::PrimitiveArray;
 use crate::kernel::BindCtx;
 use crate::kernel::Kernel;
 use crate::kernel::KernelRef;
+use crate::kernel::PushDownResult;
 use crate::serde::ArrayChildren;
 use crate::validity::Validity;
 use crate::vtable;
@@ -157,4 +159,8 @@ impl Kernel for ChunkedKernel {
         }
         Ok(vector.freeze())
     }
+
+    fn push_down_filter(self: Box<Self>, _selection: &Mask) -> VortexResult<PushDownResult> {
+        Ok(PushDownResult::NotPushed(self))
+    }
 }

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

@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_compute::filter::Filter;
+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,
+}
+
+impl FilterKernel {
+    pub fn new(child: KernelRef, mask: Mask) -> Self {
+        Self { child, mask }
+    }
+}
+
+impl Kernel for FilterKernel {
+    fn execute(self: Box<Self>) -> VortexResult<Vector> {
+        Ok(Filter::filter(&self.child.execute()?, &self.mask))
+    }
+
+    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,
+            })),
+            PushDownResult::Pushed(new_k) => PushDownResult::Pushed(new_k),
+        })
+    }
+}

vortex-array/src/arrays/filter/kernel.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/mod.rs

@@ -24,9 +24,10 @@ 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::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;
@@ -82,22 +83,46 @@ impl VTable for FilterVTable {
         &self,
         dtype: &DType,
         len: usize,
-        metadata: &Self::Metadata,
+        selection_mask: &Mask,
         _buffers: &[BufferHandle],
         children: &dyn ArrayChildren,
     ) -> VortexResult<Self::Array> {
-        let child = children.get(0, dtype, len)?;
+        assert_eq!(len, selection_mask.true_count());
+        let child = children.get(0, dtype, selection_mask.len())?;
         Ok(FilterArray {
             child,
-            mask: metadata.clone(),
+            mask: selection_mask.clone(),
             stats: Default::default(),
         })
     }
 
     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))))
+
+        // NOTE(ngates): for now we keep the same behavior as develop where we push-down any
+        //  query with <20% true values.
+        let pushdown = array.mask.density() < 0.2;
+
+        if pushdown {
+            // Try to push down the filter to the child if it's cheaper.
+            child = match child.push_down_filter(&mask)? {
+                PushDownResult::Pushed(new_k) => {
+                    tracing::debug!("Filter push down kernel:\n{:?}", new_k);
+                    return Ok(new_k);
+                }
+                PushDownResult::NotPushed(child) => {
+                    tracing::warn!(
+                        "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)))
     }
 }
 

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

@@ -44,4 +44,9 @@ impl ScalarFnArray {
             stats: Default::default(),
         })
     }
+
+    /// Get the scalar function bound to this array.
+    pub fn scalar_fn(&self) -> &ScalarFn {
+        &self.scalar_fn
+    }
 }

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

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-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;
@@ -12,6 +12,7 @@ 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)]
@@ -37,14 +38,17 @@ pub(super) enum KernelInput {
 
 impl Kernel for ScalarFnKernel {
     fn execute(self: Box<Self>) -> VortexResult<Vector> {
-        let datums: Vec<_> = self
-            .inputs
-            .into_iter()
-            .map(|input| match input {
-                KernelInput::Scalar(s) => Ok(Datum::Scalar(s)),
-                KernelInput::Vector(k) => k.execute().map(Datum::Vector),
-            })
-            .try_collect()?;
+        let mut datums: Vec<Datum> = Vec::with_capacity(self.inputs.len());
+        for input in self.inputs {
+            match input {
+                KernelInput::Scalar(s) => {
+                    datums.push(Datum::Scalar(s));
+                }
+                KernelInput::Vector(kernel) => {
+                    datums.push(Datum::Vector(kernel.execute()?));
+                }
+            }
+        }
 
         let args = ExecutionArgs {
             datums,
@@ -55,4 +59,26 @@ impl Kernel for ScalarFnKernel {
 
         Ok(self.scalar_fn.execute(args)?.ensure_vector(self.row_count))
     }
+
+    fn push_down_filter(self: Box<Self>, selection: &Mask) -> VortexResult<PushDownResult> {
+        let mut new_inputs = Vec::with_capacity(self.inputs.len());
+        for input in self.inputs {
+            match input {
+                KernelInput::Scalar(s) => {
+                    new_inputs.push(KernelInput::Scalar(s.clone()));
+                }
+                KernelInput::Vector(k) => {
+                    new_inputs.push(KernelInput::Vector(k.force_push_down_filter(selection)?));
+                }
+            }
+        }
+
+        Ok(PushDownResult::Pushed(Box::new(ScalarFnKernel {
+            scalar_fn: self.scalar_fn,
+            inputs: new_inputs,
+            input_dtypes: self.input_dtypes,
+            row_count: selection.true_count(),
+            return_dtype: self.return_dtype,
+        })))
+    }
 }