Move optimizer rules onto Array vtable (#5712)

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

Author: gatesn

encodings/runend/src/array.rs

@@ -45,6 +45,7 @@ use vortex_scalar::PValue;
 use crate::compress::runend_decode_bools;
 use crate::compress::runend_decode_primitive;
 use crate::compress::runend_encode;
+use crate::rules::RULES;
 
 vtable!(RunEnd);
 
@@ -132,6 +133,14 @@ impl VTable for RunEndVTable {
 
         Ok(())
     }
+
+    fn reduce_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        RULES.evaluate(array, parent, child_idx)
+    }
 }
 
 #[derive(Clone, Debug)]

encodings/runend/src/lib.rs

@@ -24,17 +24,14 @@ 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(
@@ -69,10 +66,6 @@ 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)]

encodings/runend/src/rules.rs

@@ -8,23 +8,24 @@ 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_array::optimizer::rules::ParentRuleSet;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
 
 use crate::RunEndArray;
 use crate::RunEndVTable;
 
+pub(super) const RULES: ParentRuleSet<RunEndVTable> =
+    ParentRuleSet::new(&[ParentRuleSet::lift(&RunEndScalarFnRule)]);
+
 /// 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)
-    }
+impl ArrayParentReduceRule<RunEndVTable> for RunEndScalarFnRule {
+    type Parent = AnyScalarFn;
 
     fn parent(&self) -> AnyScalarFn {
         AnyScalarFn

vortex-array/src/array/mod.rs

@@ -54,6 +54,7 @@ use crate::hash;
 use crate::kernel::BindCtx;
 use crate::kernel::KernelRef;
 use crate::kernel::ValidateKernel;
+use crate::optimizer::ArrayOptimizer;
 use crate::stats::StatsSetRef;
 use crate::vtable::ArrayId;
 use crate::vtable::ArrayVTable;
@@ -196,6 +197,12 @@ pub trait Array:
 
     /// Invoke the batch execution function for the array to produce a canonical vector.
     fn bind_kernel(&self, ctx: &mut BindCtx) -> VortexResult<KernelRef>;
+
+    /// Reduce the array to a more simple representation, if possible.
+    fn reduce(&self) -> VortexResult<Option<ArrayRef>>;
+
+    /// Attempt to perform a reduction of the parent of this array.
+    fn reduce_parent(&self, parent: &ArrayRef, child_idx: usize) -> VortexResult<Option<ArrayRef>>;
 }
 
 impl Array for Arc<dyn Array> {
@@ -309,6 +316,14 @@ impl Array for Arc<dyn Array> {
     fn bind_kernel(&self, ctx: &mut BindCtx) -> VortexResult<KernelRef> {
         self.as_ref().bind_kernel(ctx)
     }
+
+    fn reduce(&self) -> VortexResult<Option<ArrayRef>> {
+        self.as_ref().reduce()
+    }
+
+    fn reduce_parent(&self, parent: &ArrayRef, child_idx: usize) -> VortexResult<Option<ArrayRef>> {
+        self.as_ref().reduce_parent(parent, child_idx)
+    }
 }
 
 /// A reference counted pointer to a dynamic [`Array`] trait object.
@@ -512,11 +527,15 @@ impl<V: VTable> Array for ArrayAdapter<V> {
 
     fn filter(&self, mask: Mask) -> VortexResult<ArrayRef> {
         vortex_ensure!(self.len() == mask.len(), "Filter mask length mismatch");
-        Ok(FilterArray::new(self.to_array(), mask).into_array())
+        FilterArray::new(self.to_array(), mask)
+            .into_array()
+            .optimize()
     }
 
     fn take(&self, indices: ArrayRef) -> VortexResult<ArrayRef> {
-        Ok(DictArray::try_new(indices, self.to_array())?.into_array())
+        DictArray::try_new(indices, self.to_array())?
+            .into_array()
+            .optimize()
     }
 
     fn scalar_at(&self, index: usize) -> Scalar {
@@ -658,6 +677,42 @@ impl<V: VTable> Array for ArrayAdapter<V> {
             Ok(kernel)
         }
     }
+
+    fn reduce(&self) -> VortexResult<Option<ArrayRef>> {
+        let Some(reduced) = V::reduce(&self.0)? else {
+            return Ok(None);
+        };
+        vortex_ensure!(reduced.len() == self.len(), "Reduced array length mismatch");
+        vortex_ensure!(
+            reduced.dtype() == self.dtype(),
+            "Reduced array dtype mismatch"
+        );
+        Ok(Some(reduced))
+    }
+
+    fn reduce_parent(&self, parent: &ArrayRef, child_idx: usize) -> VortexResult<Option<ArrayRef>> {
+        #[cfg(debug_assertions)]
+        vortex_ensure!(
+            Arc::as_ptr(&parent.children()[child_idx]) == self,
+            "Parent array's child at index {} does not match self",
+            child_idx
+        );
+
+        let Some(reduced) = V::reduce_parent(&self.0, parent, child_idx)? else {
+            return Ok(None);
+        };
+
+        vortex_ensure!(
+            reduced.len() == parent.len(),
+            "Reduced array length mismatch"
+        );
+        vortex_ensure!(
+            reduced.dtype() == parent.dtype(),
+            "Reduced array dtype mismatch"
+        );
+
+        Ok(Some(reduced))
+    }
 }
 
 impl<V: VTable> ArrayHash for ArrayAdapter<V> {

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

@@ -27,12 +27,13 @@ use crate::vtable::ValidityVTableFromValidityHelper;
 mod array;
 mod canonical;
 mod operations;
-pub mod operator;
+pub mod rules;
 mod validity;
 mod visitor;
 
-pub use operator::BoolMaskedValidityRule;
+pub use rules::BoolMaskedValidityRule;
 
+use crate::arrays::bool::vtable::rules::RULES;
 use crate::kernel::KernelRef;
 use crate::kernel::ready;
 use crate::vtable::ArrayId;
@@ -112,12 +113,6 @@ impl VTable for BoolVTable {
         BoolArray::try_new(bits, validity)
     }
 
-    fn bind_kernel(array: &Self::Array, _ctx: &mut BindCtx) -> VortexResult<KernelRef> {
-        Ok(ready(
-            BoolVector::new(array.bit_buffer().clone(), array.validity_mask()).into(),
-        ))
-    }
-
     fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
         vortex_ensure!(
             children.len() <= 1,
@@ -133,6 +128,20 @@ impl VTable for BoolVTable {
 
         Ok(())
     }
+
+    fn bind_kernel(array: &Self::Array, _ctx: &mut BindCtx) -> VortexResult<KernelRef> {
+        Ok(ready(
+            BoolVector::new(array.bit_buffer().clone(), array.validity_mask()).into(),
+        ))
+    }
+
+    fn reduce_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        RULES.evaluate(array, parent, child_idx)
+    }
 }
 
 #[derive(Debug)]

…array/src/arrays/bool/vtable/operator.rs …ex-array/src/arrays/bool/vtable/rules.rsvortex-array/src/arrays/bool/vtable/operator.rs renamed to vortex-array/src/arrays/bool/vtable/rules.rs vortex-array/src/arrays/bool/vtable/operator.rs renamed to vortex-array/src/arrays/bool/vtable/rules.rs

@@ -11,19 +11,21 @@ use crate::arrays::MaskedArray;
 use crate::arrays::MaskedVTable;
 use crate::optimizer::rules::ArrayParentReduceRule;
 use crate::optimizer::rules::Exact;
+use crate::optimizer::rules::ParentRuleSet;
 use crate::vtable::ValidityHelper;
 
+pub(super) const RULES: ParentRuleSet<BoolVTable> =
+    ParentRuleSet::new(&[ParentRuleSet::lift(&BoolMaskedValidityRule)]);
+
 /// Rule to push down validity masking from MaskedArray parent into BoolArray child.
 ///
 /// When a BoolArray is wrapped by a MaskedArray, this rule merges the mask's validity
 /// with the BoolArray's existing validity, eliminating the need for the MaskedArray wrapper.
 #[derive(Default, Debug)]
 pub struct BoolMaskedValidityRule;
 
-impl ArrayParentReduceRule<Exact<BoolVTable>, Exact<MaskedVTable>> for BoolMaskedValidityRule {
-    fn child(&self) -> Exact<BoolVTable> {
-        Exact::from(&BoolVTable)
-    }
+impl ArrayParentReduceRule<BoolVTable> for BoolMaskedValidityRule {
+    type Parent = Exact<MaskedVTable>;
 
     fn parent(&self) -> Exact<MaskedVTable> {
         Exact::from(&MaskedVTable)

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

@@ -16,7 +16,9 @@ use vortex_vector::VectorMut;
 use vortex_vector::VectorMutOps;
 
 use crate::ArrayRef;
+use crate::Canonical;
 use crate::EmptyMetadata;
+use crate::IntoArray;
 use crate::ToCanonical;
 use crate::arrays::ChunkedArray;
 use crate::arrays::PrimitiveArray;
@@ -175,6 +177,14 @@ impl VTable for ChunkedVTable {
             dtype: array.dtype.clone(),
         }))
     }
+
+    fn reduce(array: &Self::Array) -> VortexResult<Option<ArrayRef>> {
+        Ok(match array.chunks.len() {
+            0 => Some(Canonical::empty(array.dtype()).into_array()),
+            1 => Some(array.chunks[0].clone()),
+            _ => None,
+        })
+    }
 }
 
 #[derive(Debug)]

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

@@ -32,12 +32,13 @@ use crate::vtable::ValidityVTableFromValidityHelper;
 mod array;
 mod canonical;
 mod operations;
-pub mod operator;
+pub mod rules;
 mod validity;
 mod visitor;
 
-pub use operator::DecimalMaskedValidityRule;
+pub use rules::DecimalMaskedValidityRule;
 
+use crate::arrays::decimal::vtable::rules::RULES;
 use crate::kernel::KernelRef;
 use crate::kernel::kernel;
 use crate::vtable::ArrayId;
@@ -181,6 +182,14 @@ impl VTable for DecimalVTable {
             })
         })
     }
+
+    fn reduce_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        RULES.evaluate(array, parent, child_idx)
+    }
 }
 
 #[derive(Debug)]

…ay/src/arrays/decimal/vtable/operator.rs …array/src/arrays/decimal/vtable/rules.rsvortex-array/src/arrays/decimal/vtable/operator.rs renamed to vortex-array/src/arrays/decimal/vtable/rules.rs vortex-array/src/arrays/decimal/vtable/operator.rs renamed to vortex-array/src/arrays/decimal/vtable/rules.rs

@@ -12,21 +12,21 @@ use crate::arrays::MaskedArray;
 use crate::arrays::MaskedVTable;
 use crate::optimizer::rules::ArrayParentReduceRule;
 use crate::optimizer::rules::Exact;
+use crate::optimizer::rules::ParentRuleSet;
 use crate::vtable::ValidityHelper;
 
+pub(super) static RULES: ParentRuleSet<DecimalVTable> =
+    ParentRuleSet::new(&[ParentRuleSet::lift(&DecimalMaskedValidityRule)]);
+
 /// Rule to push down validity masking from MaskedArray parent into DecimalArray child.
 ///
 /// When a DecimalArray is wrapped by a MaskedArray, this rule merges the mask's validity
 /// with the DecimalArray's existing validity, eliminating the need for the MaskedArray wrapper.
 #[derive(Default, Debug)]
 pub struct DecimalMaskedValidityRule;
 
-impl ArrayParentReduceRule<Exact<DecimalVTable>, Exact<MaskedVTable>>
-    for DecimalMaskedValidityRule
-{
-    fn child(&self) -> Exact<DecimalVTable> {
-        Exact::from(&DecimalVTable)
-    }
+impl ArrayParentReduceRule<DecimalVTable> for DecimalMaskedValidityRule {
+    type Parent = Exact<MaskedVTable>;
 
     fn parent(&self) -> Exact<MaskedVTable> {
         Exact::from(&MaskedVTable)

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

@@ -28,12 +28,13 @@ use crate::vtable::ValidityVTableFromValidityHelper;
 mod array;
 mod canonical;
 mod operations;
-pub mod operator;
+pub mod rules;
 mod validity;
 mod visitor;
 
-pub use operator::PrimitiveMaskedValidityRule;
+pub use rules::PrimitiveMaskedValidityRule;
 
+use crate::arrays::primitive::vtable::rules::RULES;
 use crate::kernel::KernelRef;
 use crate::kernel::ready;
 use crate::vtable::ArrayId;
@@ -141,6 +142,14 @@ impl VTable for PrimitiveVTable {
 
         Ok(())
     }
+
+    fn reduce_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        RULES.evaluate(array, parent, child_idx)
+    }
 }
 
 #[derive(Debug)]