diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index 41a6bb69b25..c5b154fc9e7 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -22,12 +22,17 @@ use vortex_error::vortex_bail; use vortex_session::VortexSession; use crate::ArrayRef; +use crate::optimizer::kernels::ArrayKernels; use crate::optimizer::kernels::ArrayKernelsExt; use crate::trace_op; pub mod kernels; pub mod rules; +/// Last zero-based fixpoint pass attempted before treating continued rewrites as an infinite loop. +/// Increasing this value permits longer rewrite chains but delays detection of cyclic rules. +const MAX_OPTIMIZER_REWRITE_PASS: usize = 100; + /// Extension trait for optimizing array trees using reduce/reduce_parent rules. pub trait ArrayOptimizer { /// Optimize the root array node by running reduce and reduce_parent rules to fixpoint. @@ -72,12 +77,11 @@ fn try_optimize( ) -> VortexResult> { let mut current_array = array.clone(); let mut any_optimizations = false; - let array_ref = session.map(|s| s.kernels()); + let session_kernels = session.map(|session| session.kernels()); trace_op!(record_optimize_start(array, session.is_some())); - // Apply reduction rules to the current array until no more rules apply. - for _ in 0..=100 { + for _ in 0..=MAX_OPTIMIZER_REWRITE_PASS { trace_op!(record_optimize_loop_start(¤t_array)); if let Some(new_array) = current_array.reduce()? { @@ -89,50 +93,29 @@ fn try_optimize( trace_op!(record_optimize_reduce_none(¤t_array)); - // Apply parent reduction rules to each slot in the context of the current array. - // Its important to take all slots here, as `current_array` can change inside the loop. - let mut parent_reduced = None; + // Try children in order; the first parent rewrite restarts the fixpoint loop. + let mut reduced_parent = None; for (slot_idx, slot) in current_array.slots().iter().enumerate() { - let Some(child) = slot else { continue }; - - // Session kernels take precedence over the child encoding's static PARENT_RULES. - if let Some(array_ref) = &array_ref - && let Some(plugins) = - array_ref.find_reduce_parent(current_array.encoding_id(), child.encoding_id()) + let Some(child) = slot else { + continue; + }; + + // Session kernels take precedence over the child's static parent-reduce rules. + if let Some(session_kernels) = &session_kernels + && let Some(new_array) = + try_session_parent_reduce(session_kernels, ¤t_array, child, slot_idx)? { - #[allow(clippy::unused_enumerate_index)] - for (_plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { - if let Some(new_array) = plugin(child, ¤t_array, slot_idx)? { - trace_op!(record_session_parent_reduce_applied( - ¤t_array, - child, - slot_idx, - _plugin_idx, - &new_array, - )); - parent_reduced = Some(new_array); - break; - } - trace_op!(record_session_parent_reduce_declined( - ¤t_array, - child, - slot_idx, - _plugin_idx, - )); - } - if parent_reduced.is_some() { - break; - } + reduced_parent = Some(new_array); + break; } if let Some(new_array) = child.reduce_parent(¤t_array, slot_idx)? { - parent_reduced = Some(new_array); + reduced_parent = Some(new_array); break; } } - if let Some(new_array) = parent_reduced { - // If the parent was replaced, then we attempt to reduce it again. + if let Some(new_array) = reduced_parent { current_array = new_array; any_optimizations = true; trace_op!(record_optimize_loop_end()); @@ -142,19 +125,51 @@ fn try_optimize( trace_op!(record_optimize_parent_reduce_none(¤t_array)); trace_op!(record_optimize_loop_end()); - // No more optimizations can be applied trace_op!(record_optimize_done(¤t_array, any_optimizations)); - if any_optimizations { - return Ok(Some(current_array)); - } else { - return Ok(None); - } + return Ok(any_optimizations.then_some(current_array)); } vortex_bail!("Exceeded maximum optimization iterations (possible infinite loop)"); } +fn try_session_parent_reduce( + kernels: &ArrayKernels, + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, +) -> VortexResult> { + let Some(reduce_parent_fns) = + kernels.find_reduce_parent(parent.encoding_id(), child.encoding_id()) + else { + return Ok(None); + }; + + #[allow(clippy::unused_enumerate_index)] + for (_kernel_idx, reduce_parent) in reduce_parent_fns.iter().enumerate() { + if let Some(new_array) = reduce_parent(child, parent, slot_idx)? { + trace_op!(record_session_parent_reduce_applied( + parent, + child, + slot_idx, + _kernel_idx, + &new_array, + )); + + return Ok(Some(new_array)); + } + + trace_op!(record_session_parent_reduce_declined( + parent, + child, + slot_idx, + _kernel_idx, + )); + } + + Ok(None) +} + fn try_optimize_recursive( array: &ArrayRef, session: &VortexSession,