Skip to main content

vortex_layout/plan/
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::PlanRef;
12use crate::plan::optimizer::reduce_parent;
13use crate::plan::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 = Vec::with_capacity(plan.child_count());
34    let mut changed = false;
35    for child in plan.children().iter() {
36        let child = child?;
37        let optimized = optimize(child.clone())?;
38        changed |= !PlanRef::ptr_eq(&child, &optimized);
39        children.push(optimized);
40    }
41
42    if !changed {
43        return Ok(plan);
44    }
45
46    let plan = plan.with_children(children)?;
47    if let Some(rewritten) = reduce(&plan)? {
48        return optimize(rewritten);
49    }
50    Ok(plan)
51}