Execute parent (#5736)

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

Author: gatesn

Cargo.lock

@@ -16,6 +16,7 @@ use vortex_array::ExecutionCtx;
 use vortex_array::Precision;
 use vortex_array::ProstMetadata;
 use vortex_array::SerializeMetadata;
+use vortex_array::VectorExecutor;
 use vortex_array::patches::Patches;
 use vortex_array::patches::PatchesMetadata;
 use vortex_array::serde::ArrayChildren;

encodings/alp/src/alp/array.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 mod compute;
+mod rules;
 
 use std::hash::Hash;
 use std::ops::Range;
@@ -48,6 +49,8 @@ use vortex_error::vortex_ensure;
 use vortex_scalar::DecimalValue;
 use vortex_scalar::Scalar;
 
+use crate::decimal_byte_parts::rules::PARENT_RULES;
+
 vtable!(DecimalByteParts);
 
 #[derive(Clone, prost::Message)]
@@ -127,6 +130,14 @@ impl VTable for DecimalBytePartsVTable {
         array.msp = children.into_iter().next().vortex_expect("checked");
         Ok(())
     }
+
+    fn reduce_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        PARENT_RULES.evaluate(array, parent, child_idx)
+    }
 }
 
 /// This array encodes decimals as between 1-4 columns of primitive typed children.

encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::arrays::FilterArray;
+use vortex_array::arrays::FilterVTable;
+use vortex_array::matchers::Exact;
+use vortex_array::optimizer::rules::ArrayParentReduceRule;
+use vortex_array::optimizer::rules::ParentRuleSet;
+use vortex_error::VortexResult;
+
+use crate::DecimalBytePartsArray;
+use crate::DecimalBytePartsVTable;
+
+pub(super) const PARENT_RULES: ParentRuleSet<DecimalBytePartsVTable> =
+    ParentRuleSet::new(&[ParentRuleSet::lift(&DecimalBytePartsFilterPushDownRule)]);
+
+#[derive(Debug)]
+struct DecimalBytePartsFilterPushDownRule;
+
+impl ArrayParentReduceRule<DecimalBytePartsVTable> for DecimalBytePartsFilterPushDownRule {
+    type Parent = Exact<FilterVTable>;
+
+    fn parent(&self) -> Self::Parent {
+        Exact::from(&FilterVTable)
+    }
+
+    fn reduce_parent(
+        &self,
+        child: &DecimalBytePartsArray,
+        parent: &FilterArray,
+        _child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        // TODO(ngates): we should benchmark whether to push-down filters with "lower parts".
+        //  For now, we only push down if there are no lower parts.
+        if !child._lower_parts.is_empty() {
+            return Ok(None);
+        }
+
+        let new_msp =
+            FilterArray::new(child.msp.clone(), parent.filter_mask().clone()).into_array();
+        let new_child =
+            DecimalBytePartsArray::try_new(new_msp, *child.decimal_dtype())?.into_array();
+        Ok(Some(new_child))
+    }
+}

encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs

@@ -26,6 +26,7 @@ prost = { workspace = true }
 rand = { workspace = true, optional = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
+vortex-compute = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }
 vortex-mask = { workspace = true }

encodings/fastlanes/Cargo.toml

@@ -9,6 +9,7 @@ use prost::Message;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::ProstMetadata;
+use vortex_array::VectorExecutor;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::vtable;
 use vortex_array::vtable::ArrayId;

encodings/fastlanes/src/delta/vtable/mod.rs

@@ -24,11 +24,13 @@ use vortex_scalar::Scalar;
 use vortex_scalar::ScalarValue;
 
 use crate::FoRArray;
+use crate::r#for::vtable::rules::PARENT_RULES;
 
 mod array;
 mod canonical;
 mod encode;
 mod operations;
+mod rules;
 mod validity;
 mod visitor;
 
@@ -104,6 +106,14 @@ impl VTable for FoRVTable {
 
         FoRArray::try_new(encoded, reference)
     }
+
+    fn reduce_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        PARENT_RULES.evaluate(array, parent, child_idx)
+    }
 }
 
 #[derive(Debug)]

encodings/fastlanes/src/for/vtable/mod.rs

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::arrays::FilterArray;
+use vortex_array::arrays::FilterVTable;
+use vortex_array::matchers::Exact;
+use vortex_array::optimizer::rules::ArrayParentReduceRule;
+use vortex_array::optimizer::rules::ParentRuleSet;
+use vortex_error::VortexResult;
+
+use crate::FoRArray;
+use crate::FoRVTable;
+
+pub(super) const PARENT_RULES: ParentRuleSet<FoRVTable> =
+    ParentRuleSet::new(&[ParentRuleSet::lift(&FoRFilterPushDownRule)]);
+
+#[derive(Debug)]
+struct FoRFilterPushDownRule;
+
+impl ArrayParentReduceRule<FoRVTable> for FoRFilterPushDownRule {
+    type Parent = Exact<FilterVTable>;
+
+    fn parent(&self) -> Self::Parent {
+        Exact::from(&FilterVTable)
+    }
+
+    fn reduce_parent(
+        &self,
+        child: &FoRArray,
+        parent: &FilterArray,
+        _child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let new_array = unsafe {
+            FoRArray::new_unchecked(
+                FilterArray::new(child.encoded().clone(), parent.filter_mask().clone())
+                    .into_array(),
+                child.reference.clone(),
+            )
+        };
+        Ok(Some(new_array.into_array()))
+    }
+}

encodings/fastlanes/src/for/vtable/rules.rs

@@ -4,7 +4,6 @@
 use std::hash::Hash;
 use std::ops::Range;
 
-use num_traits::One;
 use num_traits::cast::FromPrimitive;
 use vortex_array::ArrayBufferVisitor;
 use vortex_array::ArrayChildVisitor;
@@ -15,6 +14,7 @@ use vortex_array::ExecutionCtx;
 use vortex_array::Precision;
 use vortex_array::ProstMetadata;
 use vortex_array::SerializeMetadata;
+use vortex_array::arrays::FilterVTable;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::stats::ArrayStats;
@@ -45,11 +45,14 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
+use vortex_mask::AllOr;
 use vortex_mask::Mask;
 use vortex_scalar::PValue;
 use vortex_scalar::Scalar;
 use vortex_scalar::ScalarValue;
 use vortex_vector::Vector;
+use vortex_vector::VectorMut;
+use vortex_vector::VectorMutOps;
 use vortex_vector::primitive::PVector;
 
 vtable!(Sequence);
@@ -285,20 +288,48 @@ impl VTable for SequenceVTable {
             let base = array.base().cast::<P>();
             let multiplier = array.multiplier().cast::<P>();
 
-            let values = if multiplier == <P>::one() {
-                BufferMut::from_iter(
-                    (0..array.len()).map(|i| base + <P>::from_usize(i).vortex_expect("must fit")),
-                )
-            } else {
-                BufferMut::from_iter(
-                    (0..array.len())
-                        .map(|i| base + <P>::from_usize(i).vortex_expect("must fit") * multiplier),
-                )
-            };
-
-            PVector::<P>::new(values.freeze(), Mask::new_true(array.len())).into()
+            execute_iter(base, multiplier, 0..array.len(), array.len()).into()
         }))
     }
+
+    fn execute_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        _child_idx: usize,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Vector>> {
+        // Special-case filtered execution.
+        let Some(filter) = parent.as_opt::<FilterVTable>() else {
+            return Ok(None);
+        };
+
+        match filter.filter_mask().indices() {
+            AllOr::All => Ok(None),
+            AllOr::None => Ok(Some(VectorMut::with_capacity(array.dtype(), 0).freeze())),
+            AllOr::Some(indices) => Ok(Some(match_each_native_ptype!(array.ptype(), |P| {
+                let base = array.base().cast::<P>();
+                let multiplier = array.multiplier().cast::<P>();
+                execute_iter(base, multiplier, indices.iter().copied(), indices.len()).into()
+            }))),
+        }
+    }
+}
+
+fn execute_iter<P: NativePType, I: Iterator<Item = usize>>(
+    base: P,
+    multiplier: P,
+    iter: I,
+    len: usize,
+) -> PVector<P> {
+    let values = if multiplier == <P>::one() {
+        BufferMut::from_iter(iter.map(|i| base + <P>::from_usize(i).vortex_expect("must fit")))
+    } else {
+        BufferMut::from_iter(
+            iter.map(|i| base + <P>::from_usize(i).vortex_expect("must fit") * multiplier),
+        )
+    };
+
+    PVector::<P>::new(values.freeze(), Mask::new_true(len))
 }
 
 impl BaseArrayVTable<SequenceVTable> for SequenceVTable {

encodings/sequence/src/array.rs

@@ -23,9 +23,6 @@ use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
-use vortex_vector::Vector;
-use vortex_vector::VectorOps;
-use vortex_vector::vector_matches_dtype;
 
 use crate::ArrayEq;
 use crate::ArrayHash;
@@ -52,7 +49,6 @@ use crate::compute::InvocationArgs;
 use crate::compute::IsConstantOpts;
 use crate::compute::Output;
 use crate::compute::is_constant_opts;
-use crate::executor::ExecutionCtx;
 use crate::expr::stats::Precision;
 use crate::expr::stats::Stat;
 use crate::expr::stats::StatsProviderExt;
@@ -200,11 +196,6 @@ pub trait Array:
     /// call.
     fn invoke(&self, compute_fn: &ComputeFn, args: &InvocationArgs)
     -> VortexResult<Option<Output>>;
-
-    /// Recursively execute an array using batch CPU execution.
-    ///
-    /// To invoke the top-level execution, see [`crate::executor::VectorExecutor`].
-    fn execute(&self, ctx: &mut ExecutionCtx) -> VortexResult<Vector>;
 }
 
 impl Array for Arc<dyn Array> {
@@ -318,10 +309,6 @@ impl Array for Arc<dyn Array> {
     ) -> VortexResult<Option<Output>> {
         self.as_ref().invoke(compute_fn, args)
     }
-
-    fn execute(&self, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
-        self.as_ref().execute(ctx)
-    }
 }
 
 /// A reference counted pointer to a dynamic [`Array`] trait object.
@@ -679,25 +666,6 @@ impl<V: VTable> Array for ArrayAdapter<V> {
     ) -> VortexResult<Option<Output>> {
         <V::ComputeVTable as ComputeVTable<V>>::invoke(&self.0, compute_fn, args)
     }
-
-    fn execute(&self, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
-        let result = V::execute(&self.0, ctx)?;
-
-        if cfg!(debug_assertions) {
-            vortex_ensure!(
-                result.len() == self.len(),
-                "Result length mismatch for {}",
-                self.encoding_id()
-            );
-            vortex_ensure!(
-                vector_matches_dtype(&result, self.dtype()),
-                "Executed vector dtype mismatch for {}",
-                self.encoding_id()
-            );
-        }
-
-        Ok(result)
-    }
 }
 
 impl<V: VTable> ArrayHash for ArrayAdapter<V> {