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.encoding().as_dyn().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        for (idx, child) in current_array.children().iter().enumerate() {
50            if let Some(new_array) =
51                child
52                    .encoding()
53                    .as_dyn()
54                    .reduce_parent(child, &current_array, idx)?
55            {
56                // If the parent was replaced, then we attempt to reduce it again.
57                current_array = new_array;
58                any_optimizations = true;
59
60                // Continue to the start of the outer loop
61                continue 'outer;
62            }
63        }
64
65        // No more optimizations can be applied
66        break;
67    }
68
69    if any_optimizations {
70        Ok(Some(current_array))
71    } else {
72        Ok(None)
73    }
74}
75
76fn try_optimize_recursive(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
77    let mut current_array = array.clone();
78    let mut any_optimizations = false;
79
80    if let Some(new_array) = try_optimize(&current_array)? {
81        current_array = new_array;
82        any_optimizations = true;
83    }
84
85    let mut new_children = Vec::with_capacity(current_array.nchildren());
86    let mut any_child_optimized = false;
87    for child in current_array.children() {
88        if let Some(new_child) = try_optimize_recursive(&child)? {
89            new_children.push(new_child);
90            any_child_optimized = true;
91        } else {
92            new_children.push(child.clone());
93        }
94    }
95
96    if any_child_optimized {
97        current_array = current_array.with_children(new_children)?;
98        any_optimizations = true;
99    }
100
101    if any_optimizations {
102        Ok(Some(current_array))
103    } else {
104        Ok(None)
105    }
106}
vortex_array/optimizer/mod.rs

vortex_array/optimizer/
mod.rs
