vortex_array/optimizer/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6
7use crate::Array;
8use crate::array::ArrayRef;
9
10pub mod rules;
11
12pub trait ArrayOptimizer {
13    /// Optimize the root array node only.
14    fn optimize(&self) -> VortexResult<ArrayRef>;
15
16    /// Optimize the entire array tree recursively.
17    fn optimize_recursive(&self) -> VortexResult<ArrayRef>;
18}
19
20impl ArrayOptimizer for ArrayRef {
21    fn optimize(&self) -> VortexResult<ArrayRef> {
22        Ok(try_optimize(self)?.unwrap_or_else(|| self.clone()))
23    }
24
25    fn optimize_recursive(&self) -> VortexResult<ArrayRef> {
26        Ok(try_optimize_recursive(self)?.unwrap_or_else(|| self.clone()))
27    }
28}
29
30fn try_optimize(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
31    let mut current_array = array.clone();
32    let mut any_optimizations = false;
33
34    // Apply reduction rules to the current array until no more rules apply.
35    let mut loop_counter = 0;
36    'outer: loop {
37        if loop_counter > 100 {
38            vortex_bail!("Exceeded maximum optimization iterations (possible infinite loop)");
39        }
40        loop_counter += 1;
41
42        if let Some(new_array) = current_array.vtable().reduce(&current_array)? {
43            current_array = new_array;
44            any_optimizations = true;
45            continue;
46        }
47
48        // Apply parent reduction rules to each child in the context of the current array.
49        // Its important to take all children here, as `current_array` can change inside the loop.
50        for (idx, child) in current_array.children().iter().enumerate() {
51            if let Some(new_array) = child.vtable().reduce_parent(child, &current_array, idx)? {
52                // If the parent was replaced, then we attempt to reduce it again.
53                current_array = new_array;
54                any_optimizations = true;
55
56                // Continue to the start of the outer loop
57                continue 'outer;
58            }
59        }
60
61        // No more optimizations can be applied
62        break;
63    }
64
65    if any_optimizations {
66        Ok(Some(current_array))
67    } else {
68        Ok(None)
69    }
70}
71
72fn try_optimize_recursive(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
73    let mut current_array = array.clone();
74    let mut any_optimizations = false;
75
76    if let Some(new_array) = try_optimize(&current_array)? {
77        current_array = new_array;
78        any_optimizations = true;
79    }
80
81    let mut new_children = Vec::with_capacity(current_array.nchildren());
82    let mut any_child_optimized = false;
83    for child in current_array.children() {
84        if let Some(new_child) = try_optimize_recursive(&child)? {
85            new_children.push(new_child);
86            any_child_optimized = true;
87        } else {
88            new_children.push(child.clone());
89        }
90    }
91
92    if any_child_optimized {
93        current_array = current_array.with_children(new_children)?;
94        any_optimizations = true;
95    }
96
97    if any_optimizations {
98        Ok(Some(current_array))
99    } else {
100        Ok(None)
101    }
102}
vortex_array/optimizer/mod.rs

vortex_array/optimizer/
mod.rs
