feat[vortex-array]: added debug to expr/array reduce rules (#5442)

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

Author: joseph-isaacs

vortex-array/src/array/session/rewrite.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::fmt::Debug;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
@@ -11,17 +12,17 @@ use crate::EncodingId;
 use crate::array::ArrayRef;
 use crate::array::transform::context::ArrayRuleContext;
 use crate::array::transform::rules::{
-    AnyParent, ArrayParentMatcher, ArrayParentReduceRule, ArrayReduceRule,
+    AnyArrayParent, ArrayParentMatcher, ArrayParentReduceRule, ArrayReduceRule,
 };
 use crate::vtable::VTable;
 
 /// Dynamic trait for array reduce rules
-pub trait DynArrayReduceRule: Send + Sync {
+pub trait DynArrayReduceRule: Debug + Send + Sync {
     fn reduce(&self, array: &ArrayRef, ctx: &ArrayRuleContext) -> VortexResult<Option<ArrayRef>>;
 }
 
 /// Dynamic trait for array parent reduce rules
-pub trait DynArrayParentReduceRule: Send + Sync {
+pub trait DynArrayParentReduceRule: Debug + Send + Sync {
     fn reduce_parent(
         &self,
         array: &ArrayRef,
@@ -37,12 +38,30 @@ struct ArrayReduceRuleAdapter<V: VTable, R> {
     _phantom: PhantomData<V>,
 }
 
+impl<V: VTable, R: Debug> Debug for ArrayReduceRuleAdapter<V, R> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ArrayReduceRuleAdapter")
+            .field("rule", &self.rule)
+            .finish()
+    }
+}
+
 /// Adapter for ArrayParentReduceRule
 struct ArrayParentReduceRuleAdapter<Child: VTable, Parent: ArrayParentMatcher, R> {
     rule: R,
     _phantom: PhantomData<(Child, Parent)>,
 }
 
+impl<Child: VTable, Parent: ArrayParentMatcher, R: Debug> Debug
+    for ArrayParentReduceRuleAdapter<Child, Parent, R>
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ArrayParentReduceRuleAdapter")
+            .field("rule", &self.rule)
+            .finish()
+    }
+}
+
 impl<V, R> DynArrayReduceRule for ArrayReduceRuleAdapter<V, R>
 where
     V: VTable,
@@ -81,7 +100,7 @@ where
 
 /// Inner struct that holds all the rule registries.
 /// Wrapped in a single Arc by ArrayRewriteRuleRegistry for efficient cloning.
-#[derive(Default)]
+#[derive(Default, Debug)]
 struct ArrayRewriteRuleRegistryInner {
     /// Reduce rules indexed by encoding ID
     reduce_rules: DashMap<EncodingId, Vec<Arc<dyn DynArrayReduceRule>>>,
@@ -91,25 +110,6 @@ struct ArrayRewriteRuleRegistryInner {
     any_parent_rules: DashMap<EncodingId, Vec<Arc<dyn DynArrayParentReduceRule>>>,
 }
 
-impl std::fmt::Debug for ArrayRewriteRuleRegistryInner {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("ArrayRewriteRuleRegistryInner")
-            .field(
-                "reduce_rules",
-                &format!("{} encodings", self.reduce_rules.len()),
-            )
-            .field(
-                "parent_rules",
-                &format!("{} pairs", self.parent_rules.len()),
-            )
-            .field(
-                "any_parent_rules",
-                &format!("{} encodings", self.any_parent_rules.len()),
-            )
-            .finish()
-    }
-}
-
 /// Registry of array rewrite rules.
 ///
 /// Stores rewrite rules indexed by the encoding ID they apply to.
@@ -135,8 +135,7 @@ impl ArrayRewriteRuleRegistry {
     pub fn register_reduce_rule<V, R>(&self, encoding: &V::Encoding, rule: R)
     where
         V: VTable,
-        R: 'static,
-        R: ArrayReduceRule<V>,
+        R: ArrayReduceRule<V> + 'static,
     {
         let adapter = ArrayReduceRuleAdapter {
             rule,
@@ -159,8 +158,7 @@ impl ArrayRewriteRuleRegistry {
     ) where
         Child: VTable,
         Parent: VTable,
-        R: 'static,
-        R: ArrayParentReduceRule<Child, Parent>,
+        R: ArrayParentReduceRule<Child, Parent> + 'static,
     {
         let adapter = ArrayParentReduceRuleAdapter {
             rule,
@@ -179,8 +177,7 @@ impl ArrayRewriteRuleRegistry {
     pub fn register_any_parent_rule<Child, R>(&self, child_encoding: &Child::Encoding, rule: R)
     where
         Child: VTable,
-        R: 'static,
-        R: ArrayParentReduceRule<Child, AnyParent>,
+        R: ArrayParentReduceRule<Child, AnyArrayParent> + 'static,
     {
         let adapter = ArrayParentReduceRuleAdapter {
             rule,

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

@@ -9,4 +9,4 @@ mod tests;
 
 pub use context::ArrayRuleContext;
 pub use optimizer::ArrayOptimizer;
-pub use rules::{AnyParent, ArrayParentMatcher, ArrayParentReduceRule, ArrayReduceRule};
+pub use rules::{AnyArrayParent, ArrayParentMatcher, ArrayParentReduceRule, ArrayReduceRule};

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

@@ -89,13 +89,7 @@ impl ArrayOptimizer {
                 Some(&parent_id),
                 |rules| -> VortexResult<Option<ArrayRef>> {
                     for rule in rules {
-                        println!("apply rule");
                         if let Some(new_array) = rule.reduce_parent(child, &array, idx, ctx)? {
-                            println!(
-                                "matched new: {}, old: {}",
-                                new_array.display_tree(),
-                                array.display_tree()
-                            );
                             return Ok(Some(new_array));
                         }
                     }

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

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::fmt::Debug;
+
 use vortex_error::VortexResult;
 
 use crate::array::ArrayRef;
@@ -16,9 +18,10 @@ pub trait ArrayParentMatcher: Send + Sync + 'static {
 }
 
 /// Matches any parent type (wildcard matcher)
-pub struct AnyParent;
+#[derive(Debug)]
+pub struct AnyArrayParent;
 
-impl ArrayParentMatcher for AnyParent {
+impl ArrayParentMatcher for AnyArrayParent {
     type View<'a> = &'a ArrayRef;
 
     fn try_match(parent: &ArrayRef) -> Option<Self::View<'_>> {
@@ -36,7 +39,7 @@ impl<V: VTable> ArrayParentMatcher for V {
 }
 
 /// A rewrite rule that transforms arrays based on the array itself and its children
-pub trait ArrayReduceRule<V: VTable>: Send + Sync {
+pub trait ArrayReduceRule<V: VTable>: Debug + Send + Sync + 'static {
     /// Attempt to rewrite this array.
     ///
     /// Returns:
@@ -47,7 +50,9 @@ pub trait ArrayReduceRule<V: VTable>: Send + Sync {
 }
 
 /// A rewrite rule that transforms arrays based on parent context
-pub trait ArrayParentReduceRule<Child: VTable, Parent: ArrayParentMatcher>: Send + Sync {
+pub trait ArrayParentReduceRule<Child: VTable, Parent: ArrayParentMatcher>:
+    Debug + Send + Sync + 'static
+{
     /// Attempt to rewrite this child array given information about its parent.
     ///
     /// Returns:

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

@@ -18,6 +18,7 @@ use crate::expr::transform::ExprOptimizer;
 use crate::validity::Validity;
 
 /// Test rule that unwraps single-chunk ChunkedArrays
+#[derive(Debug, Default)]
 struct UnwrapSingleChunkRule;
 
 impl ArrayReduceRule<ChunkedVTable> for UnwrapSingleChunkRule {
@@ -42,8 +43,9 @@ fn test_unwrap_single_chunk_rule() -> VortexResult<()> {
     let primitive = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
     let chunked = ChunkedArray::from_iter([primitive.clone()]);
 
-    let rule = UnwrapSingleChunkRule;
-    let result = rule.reduce(&chunked, &ctx)?.vortex_expect("transformed");
+    let result = UnwrapSingleChunkRule
+        .reduce(&chunked, &ctx)?
+        .vortex_expect("transformed");
 
     assert!(Arc::ptr_eq(&primitive, &result));
     Ok(())
@@ -60,8 +62,7 @@ fn test_unwrap_single_chunk_rule_no_op() -> VortexResult<()> {
         PrimitiveArray::from_iter([3i32, 4]).into_array(),
     ]);
 
-    let rule = UnwrapSingleChunkRule;
-    let result = rule.reduce(&chunked, &ctx)?;
+    let result = UnwrapSingleChunkRule.reduce(&chunked, &ctx)?;
 
     assert!(result.is_none());
     Ok(())
@@ -72,7 +73,10 @@ fn test_reduce_rules_traverse_whole_tree() -> VortexResult<()> {
     let array_session = ArraySession::default();
     let expr_session = ExprSession::default();
 
-    array_session.register_reduce_rule::<ChunkedVTable, _>(&ChunkedEncoding, UnwrapSingleChunkRule);
+    array_session.register_reduce_rule::<ChunkedVTable, UnwrapSingleChunkRule>(
+        &ChunkedEncoding,
+        UnwrapSingleChunkRule,
+    );
 
     let expr_optimizer = ExprOptimizer::new(&expr_session);
     let optimizer = array_session.optimizer(expr_optimizer);
@@ -121,6 +125,7 @@ fn test_reduce_rules_traverse_whole_tree() -> VortexResult<()> {
 }
 
 // Odd rule for testing
+#[derive(Debug, Default)]
 struct ConstantInStructRule;
 
 impl ArrayParentReduceRule<ConstantVTable, StructVTable> for ConstantInStructRule {
@@ -161,7 +166,7 @@ fn test_parent_rules_traverse_whole_tree() -> VortexResult<()> {
     let array_session = ArraySession::default();
     let expr_session = ExprSession::default();
 
-    array_session.register_parent_rule::<ConstantVTable, StructVTable, _>(
+    array_session.register_parent_rule::<ConstantVTable, StructVTable, ConstantInStructRule>(
         &ConstantEncoding,
         &crate::arrays::StructEncoding,
         ConstantInStructRule,

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

@@ -37,6 +37,7 @@ impl OperatorVTable<BoolVTable> for BoolVTable {
 ///
 /// 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<BoolVTable, MaskedVTable> for BoolMaskedValidityRule {

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

@@ -43,6 +43,7 @@ impl OperatorVTable<DecimalVTable> for DecimalVTable {
 ///
 /// 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<DecimalVTable, MaskedVTable> for DecimalMaskedValidityRule {

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

@@ -12,6 +12,7 @@ use crate::vtable::OperatorVTable;
 impl OperatorVTable<ExprVTable> for ExprVTable {}
 
 /// Rule to optimize expressions within ExprArrays.
+#[derive(Default, Debug)]
 pub struct ExprOptimizationRule;
 
 impl ArrayReduceRule<ExprVTable> for ExprOptimizationRule {

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

@@ -42,6 +42,7 @@ impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
 ///
 /// 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.
+#[derive(Default, Debug)]
 pub struct PrimitiveMaskedValidityRule;
 
 impl ArrayParentReduceRule<PrimitiveVTable, MaskedVTable> for PrimitiveMaskedValidityRule {

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

@@ -47,6 +47,7 @@ impl OperatorVTable<StructVTable> for StructVTable {
 ///
 /// This optimization pushes expression evaluation down to individual struct fields, enabling
 /// better field-level optimizations and potentially avoiding materialization of unused fields.
+#[derive(Default, Debug)]
 pub struct StructExprPartitionRule;
 
 impl ArrayParentReduceRule<StructVTable, ExprVTable> for StructExprPartitionRule {