u

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-array/src/array/operator.rs

@@ -26,12 +26,6 @@ pub trait ArrayOperator: 'static + Send + Sync {
     /// If the array's implementation returns an invalid vector (wrong length, wrong type, etc.).
     fn execute_batch(&self, selection: &Mask, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector>;
 
-    /// Optimize the array by running the optimization rules.
-    fn reduce(&self) -> VortexResult<Option<ArrayRef>>;
-
-    /// Optimize the array by pushing down a parent array.
-    fn reduce_parent(&self, parent: &ArrayRef, child_idx: usize) -> VortexResult<Option<ArrayRef>>;
-
     /// Returns the array as a pipeline node, if supported.
     fn as_pipelined(&self) -> Option<&dyn PipelinedNode>;
 
@@ -48,14 +42,6 @@ impl ArrayOperator for Arc<dyn Array> {
         self.as_ref().execute_batch(selection, 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)
-    }
-
     fn as_pipelined(&self) -> Option<&dyn PipelinedNode> {
         self.as_ref().as_pipelined()
     }
@@ -96,14 +82,6 @@ impl<V: VTable> ArrayOperator for ArrayAdapter<V> {
         Ok(vector)
     }
 
-    fn reduce(&self) -> VortexResult<Option<ArrayRef>> {
-        <V::OperatorVTable as OperatorVTable<V>>::reduce(&self.0)
-    }
-
-    fn reduce_parent(&self, parent: &ArrayRef, child_idx: usize) -> VortexResult<Option<ArrayRef>> {
-        <V::OperatorVTable as OperatorVTable<V>>::reduce_parent(&self.0, parent, child_idx)
-    }
-
     fn as_pipelined(&self) -> Option<&dyn PipelinedNode> {
         <V::OperatorVTable as OperatorVTable<V>>::pipeline_node(&self.0)
     }

vortex-array/src/array/transform/optimizer.rs

@@ -79,7 +79,6 @@ impl ArrayOptimizer {
         // Now try to apply parent rules to each optimized child in the context of this array
         // Use the optimized_children list directly instead of re-fetching from array.children()
         // let mut transformed_children = Vec::with_capacity(optimized_children.len());
-        let rules_applied = false;
 
         for (idx, child) in optimized_children.iter().enumerate() {
             let child_id = child.encoding_id();

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

@@ -7,7 +7,7 @@ pub use array::DecimalArray;
 mod compute;
 
 mod vtable;
-pub use vtable::{DecimalEncoding, DecimalVTable};
+pub use vtable::{DecimalEncoding, DecimalMaskedValidityRule, DecimalVTable};
 
 mod utils;
 pub use utils::*;

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

@@ -17,10 +17,12 @@ use crate::{
 mod array;
 mod canonical;
 mod operations;
-mod operator;
+pub mod operator;
 mod validity;
 mod visitor;
 
+pub use operator::DecimalMaskedValidityRule;
+
 vtable!(Decimal);
 
 // The type of the values can be determined by looking at the type info...right?
@@ -42,7 +44,7 @@ impl VTable for DecimalVTable {
     type VisitorVTable = Self;
     type ComputeVTable = NotSupported;
     type EncodeVTable = NotSupported;
-    type OperatorVTable = NotSupported;
+    type OperatorVTable = Self;
 
     fn id(_encoding: &Self::Encoding) -> EncodingId {
         EncodingId::new_ref("vortex.decimal")

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

@@ -6,7 +6,8 @@ use vortex_dtype::{PrecisionScale, match_each_decimal_value_type};
 use vortex_error::VortexResult;
 use vortex_vector::decimal::DVector;
 
-use crate::arrays::{DecimalArray, DecimalVTable, MaskedVTable};
+use crate::array::transform::{ArrayParentReduceRule, ArrayRuleContext};
+use crate::arrays::{DecimalArray, DecimalVTable, MaskedArray, MaskedVTable};
 use crate::execution::{BatchKernelRef, BindCtx, kernel};
 use crate::vtable::{OperatorVTable, ValidityHelper};
 use crate::{ArrayRef, IntoArray};
@@ -36,30 +37,37 @@ impl OperatorVTable<DecimalVTable> for DecimalVTable {
             }))
         })
     }
+}
+
+/// 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.
+pub struct DecimalMaskedValidityRule;
 
+impl ArrayParentReduceRule<DecimalVTable, MaskedVTable> for DecimalMaskedValidityRule {
     fn reduce_parent(
+        &self,
         array: &DecimalArray,
-        parent: &ArrayRef,
+        parent: &MaskedArray,
         _child_idx: usize,
+        _ctx: &ArrayRuleContext,
     ) -> VortexResult<Option<ArrayRef>> {
-        // Push-down masking of `validity` from the parent `MaskedArray`.
-        if let Some(masked) = parent.as_opt::<MaskedVTable>() {
-            let masked_array = match_each_decimal_value_type!(array.values_type(), |D| {
-                // SAFETY: Since we are only flipping some bits in the validity, all invariants that
-                // were upheld are still upheld.
-                unsafe {
-                    DecimalArray::new_unchecked(
-                        array.buffer::<D>(),
-                        array.decimal_dtype(),
-                        array.validity().clone().and(masked.validity().clone()),
-                    )
-                }
-                .into_array()
-            });
-
-            return Ok(Some(masked_array));
-        }
+        // Merge the parent's validity mask into the child's validity
+        // TODO(joe): make this lazy
+        let masked_array = match_each_decimal_value_type!(array.values_type(), |D| {
+            // SAFETY: Since we are only flipping some bits in the validity, all invariants that
+            // were upheld are still upheld.
+            unsafe {
+                DecimalArray::new_unchecked(
+                    array.buffer::<D>(),
+                    array.decimal_dtype(),
+                    array.validity().clone().and(parent.validity().clone()),
+                )
+            }
+            .into_array()
+        });
 
-        Ok(None)
+        Ok(Some(masked_array))
     }
 }

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

@@ -8,7 +8,7 @@ mod compute;
 pub use compute::{IS_CONST_LANE_WIDTH, compute_is_constant};
 
 mod vtable;
-pub use vtable::{PrimitiveEncoding, PrimitiveVTable};
+pub use vtable::{PrimitiveEncoding, PrimitiveMaskedValidityRule, PrimitiveVTable};
 
 mod native_value;
 pub use native_value::NativeValue;

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

@@ -14,10 +14,12 @@ use crate::{EmptyMetadata, EncodingId, EncodingRef, vtable};
 mod array;
 mod canonical;
 mod operations;
-mod operator;
+pub mod operator;
 mod validity;
 mod visitor;
 
+pub use operator::PrimitiveMaskedValidityRule;
+
 vtable!(Primitive);
 
 impl VTable for PrimitiveVTable {

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

@@ -7,7 +7,8 @@ use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
 use vortex_vector::primitive::PVector;
 
-use crate::arrays::{MaskedVTable, PrimitiveArray, PrimitiveVTable};
+use crate::array::transform::{ArrayParentReduceRule, ArrayRuleContext};
+use crate::arrays::{MaskedArray, MaskedVTable, PrimitiveArray, PrimitiveVTable};
 use crate::execution::{BatchKernelRef, BindCtx, kernel};
 use crate::vtable::{OperatorVTable, ValidityHelper};
 use crate::{ArrayRef, IntoArray};
@@ -35,29 +36,36 @@ impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
             }))
         })
     }
+}
+
+/// Rule to push down validity masking from MaskedArray parent into PrimitiveArray child.
+///
+/// When a PrimitiveArray is wrapped by a MaskedArray, this rule merges the mask's validity
+/// with the PrimitiveArray's existing validity, eliminating the need for the MaskedArray wrapper.
+pub struct PrimitiveMaskedValidityRule;
 
+impl ArrayParentReduceRule<PrimitiveVTable, MaskedVTable> for PrimitiveMaskedValidityRule {
     fn reduce_parent(
+        &self,
         array: &PrimitiveArray,
-        parent: &ArrayRef,
+        parent: &MaskedArray,
         _child_idx: usize,
+        _ctx: &ArrayRuleContext,
     ) -> VortexResult<Option<ArrayRef>> {
-        // Push-down masking of `validity` from the parent `MaskedArray`.
-        if let Some(masked) = parent.as_opt::<MaskedVTable>() {
-            let masked_array = match_each_native_ptype!(array.ptype(), |T| {
-                // SAFETY: Since we are only flipping some bits in the validity, all invariants that
-                // were upheld are still upheld.
-                unsafe {
-                    PrimitiveArray::new_unchecked(
-                        Buffer::<T>::from_byte_buffer(array.byte_buffer().clone()),
-                        array.validity().clone().and(masked.validity().clone()),
-                    )
-                }
-                .into_array()
-            });
-
-            return Ok(Some(masked_array));
-        }
-
-        Ok(None)
+        // Merge the parent's validity mask into the child's validity
+        // TODO(joe): make this lazy
+        let masked_array = match_each_native_ptype!(array.ptype(), |T| {
+            // SAFETY: Since we are only flipping some bits in the validity, all invariants that
+            // were upheld are still upheld.
+            unsafe {
+                PrimitiveArray::new_unchecked(
+                    Buffer::<T>::from_byte_buffer(array.byte_buffer().clone()),
+                    array.validity().clone().and(parent.validity().clone()),
+                )
+            }
+            .into_array()
+        });
+
+        Ok(Some(masked_array))
     }
 }

vortex-array/src/lib.rs

@@ -179,6 +179,18 @@ impl Default for ArraySession {
             arrays::BoolMaskedValidityRule,
         );
 
+        session.register_parent_rule::<arrays::PrimitiveVTable, arrays::MaskedVTable, _>(
+            &PrimitiveEncoding,
+            &MaskedEncoding,
+            arrays::PrimitiveMaskedValidityRule,
+        );
+
+        session.register_parent_rule::<arrays::DecimalVTable, arrays::MaskedVTable, _>(
+            &DecimalEncoding,
+            &MaskedEncoding,
+            arrays::DecimalMaskedValidityRule,
+        );
+
         session.register_parent_rule::<arrays::StructVTable, arrays::ExprVTable, _>(
             &StructEncoding,
             &arrays::ExprEncoding,

vortex-array/src/optimizer.rs

@@ -1,43 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::sync::Arc;
-
 use vortex_error::VortexResult;
 
+use crate::ArrayRef;
 use crate::vtable::VTable;
-use crate::{Array, ArrayRef};
-
-impl dyn Array + '_ {
-    /// Optimize this array by applying optimization rules recursively to its children in a single
-    /// bottom-up pass.
-    pub fn optimize(&self) -> VortexResult<ArrayRef> {
-        let slf = self.to_array();
-        let children = self.children();
-
-        let mut new_children = Vec::with_capacity(children.len());
-        let mut children_modified = false;
-        for (idx, child) in children.iter().enumerate() {
-            let child = child.optimize()?;
-
-            // Check if the child can reduce us (its parent), and if so bail early.
-            if let Some(reduced) = child.reduce_parent(&slf, idx)? {
-                return Ok(reduced);
-            }
-
-            if !Arc::ptr_eq(&child, &children[idx]) {
-                children_modified = true;
-            }
-            new_children.push(child);
-        }
-
-        if children_modified {
-            return self.with_children(&new_children);
-        }
-
-        Ok(slf)
-    }
-}
 
 /// An optimizer rule that tries to reduce/replace a parent array where the implementer is a
 /// child array in the `CHILD_IDX` position of the parent array.
@@ -63,9 +30,11 @@ mod tests {
     use vortex_dtype::PTypeDowncast;
     use vortex_vector::VectorOps;
 
-    use crate::IntoArray;
     use crate::arrays::{BoolArray, MaskedArray, PrimitiveArray};
+    use crate::expr::session::ExprSession;
+    use crate::expr::transform::ExprOptimizer;
     use crate::validity::Validity;
+    use crate::{ArraySession, IntoArray};
 
     #[test]
     fn test_masked_pushdown() {
@@ -78,8 +47,16 @@ mod tests {
         )
         .unwrap();
 
-        let result = masked.optimize().unwrap();
-        assert_eq!(masked.dtype(), result.dtype());
+        let masked_dtype = masked.dtype().clone();
+
+        // Use the new ArrayOptimizer via ArraySession
+        let array_session = ArraySession::default();
+        let expr_session = ExprSession::default();
+        let expr_optimizer = ExprOptimizer::new(&expr_session);
+        let optimizer = array_session.optimizer(expr_optimizer);
+
+        let result = optimizer.optimize_array(masked.into_array()).unwrap();
+        assert_eq!(&masked_dtype, result.dtype());
         assert!(result.dtype().is_nullable());
 
         let vector = result.execute().unwrap().into_primitive().into_u32();