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,