Skip to main content

vortex_layout/plan/physical/
optimize.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Plan optimization.
5//!
6//! The optimizer applies static rewrites top-down, optimizes children, then retries rewrites
7//! exposed by the optimized children.
8
9use vortex_error::VortexResult;
10
11use crate::plan::physical::PlanRef;
12use crate::plan::physical::optimizer::reduce_parent;
13use crate::plan::physical::optimizer::reduce_plan;
14
15fn reduce(plan: &PlanRef) -> VortexResult<Option<PlanRef>> {
16    if let Some(rewritten) = reduce_plan(plan)? {
17        return Ok(Some(rewritten));
18    }
19    for child_idx in 0..plan.child_count() {
20        if let Some(rewritten) = reduce_parent(plan, child_idx)? {
21            return Ok(Some(rewritten));
22        }
23    }
24    Ok(None)
25}
26
27/// Optimizes `plan`, preserving its dtype and row domain.
28pub fn optimize(plan: PlanRef) -> VortexResult<PlanRef> {
29    if let Some(rewritten) = reduce(&plan)? {
30        return optimize(rewritten);
31    }
32
33    let mut children: Option<Vec<PlanRef>> = None;
34    for (index, child) in plan.children().iter_refs().enumerate() {
35        let child = child?;
36        let optimized = optimize(child.clone())?;
37        if let Some(children) = &mut children {
38            children.push(optimized);
39        } else if !PlanRef::ptr_eq(child, &optimized) {
40            let mut optimized_children = Vec::with_capacity(plan.child_count());
41            for previous in plan.children().iter_refs().take(index) {
42                optimized_children.push(previous?.clone());
43            }
44            optimized_children.push(optimized);
45            children = Some(optimized_children);
46        }
47    }
48
49    let Some(children) = children else {
50        return Ok(plan);
51    };
52
53    let plan = plan.with_children(children)?;
54    if let Some(rewritten) = reduce(&plan)? {
55        return optimize(rewritten);
56    }
57    Ok(plan)
58}