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::DynArray;
14use crate::array::ArrayRef;
15
16pub mod rules;
17
18/// Extension trait for optimizing array trees using reduce/reduce_parent rules.
19pub trait ArrayOptimizer {
20    /// Optimize the root array node only by running reduce and reduce_parent rules to fixpoint.
21    fn optimize(&self) -> VortexResult<ArrayRef>;
22
23    /// Optimize the entire array tree recursively (root and all descendants).
24    fn optimize_recursive(&self) -> VortexResult<ArrayRef>;
25}
26
27impl ArrayOptimizer for ArrayRef {
28    fn optimize(&self) -> VortexResult<ArrayRef> {
29        Ok(try_optimize(self)?.unwrap_or_else(|| self.clone()))
30    }
31
32    fn optimize_recursive(&self) -> VortexResult<ArrayRef> {
33        Ok(try_optimize_recursive(self)?.unwrap_or_else(|| self.clone()))
34    }
35}
36
37fn try_optimize(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
38    let mut current_array = array.clone();
39    let mut any_optimizations = false;
40
41    // Apply reduction rules to the current array until no more rules apply.
42    let mut loop_counter = 0;
43    'outer: loop {
44        if loop_counter > 100 {
45            vortex_bail!("Exceeded maximum optimization iterations (possible infinite loop)");
46        }
47        loop_counter += 1;
48
49        if let Some(new_array) = current_array.vtable().reduce(&current_array)? {
50            current_array = new_array;
51            any_optimizations = true;
52            continue;
53        }
54
55        // Apply parent reduction rules to each child in the context of the current array.
56        // Its important to take all children here, as `current_array` can change inside the loop.
57        for (idx, child) in current_array.children().iter().enumerate() {
58            if let Some(new_array) = child.vtable().reduce_parent(child, &current_array, 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_children = Vec::with_capacity(current_array.nchildren());
89    let mut any_child_optimized = false;
90    for child in current_array.children() {
91        if let Some(new_child) = try_optimize_recursive(&child)? {
92            new_children.push(new_child);
93            any_child_optimized = true;
94        } else {
95            new_children.push(child.clone());
96        }
97    }
98
99    if any_child_optimized {
100        current_array = current_array.with_children(new_children)?;
101        any_optimizations = true;
102    }
103
104    if any_optimizations {
105        Ok(Some(current_array))
106    } else {
107        Ok(None)
108    }
109}