Skip to main content

vortex_array/optimizer/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! The optimizer applies metadata-only rewrite rules (`reduce` and `reduce_parent`) in a
5//! fixpoint loop until no more transformations are possible.
6//!
7//! Optimization runs between execution steps, which is what enables cross-step optimizations:
8//! after a child is decoded, new `reduce_parent` rules may match that were previously blocked.
9
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12
13use crate::ArrayRef;
14
15pub mod rules;
16
17/// Extension trait for optimizing array trees using reduce/reduce_parent rules.
18pub trait ArrayOptimizer {
19    /// Optimize the root array node only by running reduce and reduce_parent rules to fixpoint.
20    fn optimize(&self) -> VortexResult<ArrayRef>;
21
22    /// Optimize the entire array tree recursively (root and all descendants).
23    fn optimize_recursive(&self) -> VortexResult<ArrayRef>;
24}
25
26impl ArrayOptimizer for ArrayRef {
27    fn optimize(&self) -> VortexResult<ArrayRef> {
28        Ok(try_optimize(self)?.unwrap_or_else(|| self.clone()))
29    }
30
31    fn optimize_recursive(&self) -> VortexResult<ArrayRef> {
32        Ok(try_optimize_recursive(self)?.unwrap_or_else(|| self.clone()))
33    }
34}
35
36fn try_optimize(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
37    let mut current_array = array.clone();
38    let mut any_optimizations = false;
39
40    // Apply reduction rules to the current array until no more rules apply.
41    let mut loop_counter = 0;
42    'outer: loop {
43        if loop_counter > 100 {
44            vortex_bail!("Exceeded maximum optimization iterations (possible infinite loop)");
45        }
46        loop_counter += 1;
47
48        if let Some(new_array) = current_array.reduce()? {
49            current_array = new_array;
50            any_optimizations = true;
51            continue;
52        }
53
54        // Apply parent reduction rules to each slot in the context of the current array.
55        // Its important to take all slots here, as `current_array` can change inside the loop.
56        for (slot_idx, slot) in current_array.slots().iter().enumerate() {
57            let Some(child) = slot else { continue };
58            if let Some(new_array) = child.reduce_parent(&current_array, slot_idx)? {
59                // If the parent was replaced, then we attempt to reduce it again.
60                current_array = new_array;
61                any_optimizations = true;
62
63                // Continue to the start of the outer loop
64                continue 'outer;
65            }
66        }
67
68        // No more optimizations can be applied
69        break;
70    }
71
72    if any_optimizations {
73        Ok(Some(current_array))
74    } else {
75        Ok(None)
76    }
77}
78
79fn try_optimize_recursive(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
80    let mut current_array = array.clone();
81    let mut any_optimizations = false;
82
83    if let Some(new_array) = try_optimize(&current_array)? {
84        current_array = new_array;
85        any_optimizations = true;
86    }
87
88    let mut new_slots = Vec::with_capacity(current_array.slots().len());
89    let mut any_slot_optimized = false;
90    for slot in current_array.slots() {
91        match slot {
92            Some(child) => {
93                if let Some(new_child) = try_optimize_recursive(child)? {
94                    new_slots.push(Some(new_child));
95                    any_slot_optimized = true;
96                } else {
97                    new_slots.push(Some(child.clone()));
98                }
99            }
100            None => new_slots.push(None),
101        }
102    }
103
104    if any_slot_optimized {
105        current_array = current_array.with_slots(new_slots)?;
106        any_optimizations = true;
107    }
108
109    if any_optimizations {
110        Ok(Some(current_array))
111    } else {
112        Ok(None)
113    }
114}