Skip to content

Commit a9043d5

Browse files
joseph-isaacsjoseph-isaacs
committed
Use static parent-reduction rules for plan expressions
Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
1 parent 87243b8 commit a9043d5

9 files changed

Lines changed: 417 additions & 166 deletions

File tree

‎vortex-layout/src/plan/mod.rs‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ pub use plans::RowIdxPlan;
2929
pub use plans::RowIdxValuesPlan;
3030
pub use plans::StructPlan;
3131
use vortex_array::dtype::DType;
32-
use vortex_array::expr::Expression;
3332
use vortex_error::VortexResult;
3433
use vortex_error::vortex_bail;
3534

@@ -65,15 +64,6 @@ pub trait Plan: 'static + Send + Sync {
6564
/// domain.
6665
fn optimize(&self) -> VortexResult<PlanRef>;
6766

68-
/// Attempts to rewrite `expression` through this plan.
69-
///
70-
/// Returns `None` when this plan has no applicable expression rewrite. Implementations should
71-
/// preserve child slots that the rewrite does not change.
72-
fn optimize_expression(&self, expression: &Expression) -> VortexResult<Option<PlanRef>> {
73-
let _ = expression;
74-
Ok(None)
75-
}
76-
7767
/// Returns the dtype produced by this plan.
7868
fn dtype(&self) -> &DType;
7969

‎vortex-layout/src/plan/optimizer/mod.rs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,38 @@ pub use rules::DynPlanParentReduceRule;
99
pub use rules::PlanParentReduceRule;
1010
pub use rules::PlanParentReduceRuleAdapter;
1111
pub use rules::PlanParentRuleSet;
12+
use vortex_error::VortexResult;
13+
14+
use super::ChunkedPlan;
15+
use super::DictPlan;
16+
use super::PlanRef;
17+
use super::RowIdxPlan;
18+
use super::StructPlan;
19+
use super::plans::ExpressionChunkedRule;
20+
use super::plans::ExpressionDictRule;
21+
use super::plans::ExpressionRowIdxRule;
22+
use super::plans::ExpressionStructRule;
23+
24+
static EXPRESSION_CHUNKED_RULE: PlanParentReduceRuleAdapter<ChunkedPlan, ExpressionChunkedRule> =
25+
PlanParentReduceRuleAdapter::new(ExpressionChunkedRule);
26+
static EXPRESSION_DICT_RULE: PlanParentReduceRuleAdapter<DictPlan, ExpressionDictRule> =
27+
PlanParentReduceRuleAdapter::new(ExpressionDictRule);
28+
static EXPRESSION_ROW_IDX_RULE: PlanParentReduceRuleAdapter<RowIdxPlan, ExpressionRowIdxRule> =
29+
PlanParentReduceRuleAdapter::new(ExpressionRowIdxRule);
30+
static EXPRESSION_STRUCT_RULE: PlanParentReduceRuleAdapter<StructPlan, ExpressionStructRule> =
31+
PlanParentReduceRuleAdapter::new(ExpressionStructRule);
32+
33+
static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[
34+
&EXPRESSION_CHUNKED_RULE,
35+
&EXPRESSION_DICT_RULE,
36+
&EXPRESSION_ROW_IDX_RULE,
37+
&EXPRESSION_STRUCT_RULE,
38+
]);
39+
40+
/// Attempts a static rewrite for `parent` and its child at `child_idx`.
41+
pub(crate) fn reduce_parent(parent: &PlanRef, child_idx: usize) -> VortexResult<Option<PlanRef>> {
42+
let Some(child) = parent.child(child_idx)? else {
43+
return Ok(None);
44+
};
45+
PARENT_RULES.evaluate(&child, parent, child_idx)
46+
}

‎vortex-layout/src/plan/plans/chunked.rs‎

Lines changed: 35 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ use std::borrow::Cow;
55
use std::sync::Arc;
66

77
use vortex_array::dtype::DType;
8-
use vortex_array::expr::Expression;
98
use vortex_array::expr::label_tree;
109
use vortex_error::VortexResult;
1110

@@ -16,6 +15,7 @@ use crate::plan::LazyPlanChildren;
1615
use crate::plan::Plan;
1716
use crate::plan::PlanRef;
1817
use crate::plan::new_plan;
18+
use crate::plan::optimizer::PlanParentReduceRule;
1919

2020
/// A physical plan with one child per row chunk.
2121
pub struct ChunkedPlan {
@@ -63,26 +63,6 @@ impl Plan for ChunkedPlan {
6363
Ok(Arc::new(self.with_chunks(self.dtype.clone(), chunks)))
6464
}
6565

66-
fn optimize_expression(&self, expression: &Expression) -> VortexResult<Option<PlanRef>> {
67-
let references_row_idx = label_tree(
68-
expression,
69-
|node| node.is::<RowIdx>(),
70-
|acc, &child| acc | child,
71-
)
72-
.get(expression)
73-
.copied()
74-
.unwrap_or(false);
75-
if references_row_idx {
76-
return Ok(None);
77-
}
78-
79-
let dtype = expression.return_dtype(&self.dtype)?;
80-
let chunks = self
81-
.chunks
82-
.try_map(|_, chunk| ExpressionPlan::try_new(expression.clone(), chunk)?.optimize())?;
83-
Ok(Some(Arc::new(self.with_chunks(dtype, chunks))))
84-
}
85-
8666
fn dtype(&self) -> &DType {
8767
&self.dtype
8868
}
@@ -106,3 +86,37 @@ impl Plan for ChunkedPlan {
10686
Cow::Owned(format!("chunks[{index}]"))
10787
}
10888
}
89+
90+
/// Pushes an expression through every chunk of a chunked plan.
91+
#[derive(Debug)]
92+
pub(crate) struct ExpressionChunkedRule;
93+
94+
impl PlanParentReduceRule<ChunkedPlan> for ExpressionChunkedRule {
95+
type Parent = ExpressionPlan;
96+
97+
fn reduce_parent(
98+
&self,
99+
child: &ChunkedPlan,
100+
parent: &ExpressionPlan,
101+
_child_idx: usize,
102+
) -> VortexResult<Option<PlanRef>> {
103+
let expression = parent.expression();
104+
let references_row_idx = label_tree(
105+
expression,
106+
|node| node.is::<RowIdx>(),
107+
|acc, &child| acc | child,
108+
)
109+
.get(expression)
110+
.copied()
111+
.unwrap_or(false);
112+
if references_row_idx {
113+
return Ok(None);
114+
}
115+
116+
let dtype = expression.return_dtype(&child.dtype)?;
117+
let chunks = child
118+
.chunks
119+
.try_map(|_, chunk| ExpressionPlan::try_new(expression.clone(), chunk)?.optimize())?;
120+
Ok(Some(Arc::new(child.with_chunks(dtype, chunks))))
121+
}
122+
}

‎vortex-layout/src/plan/plans/dict.rs‎

Lines changed: 42 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
use std::borrow::Cow;
55
use std::sync::Arc;
66

7-
use vortex_array::expr::Expression;
87
use vortex_array::expr::is_root;
98
use vortex_array::expr::label_is_fallible;
109
use vortex_array::expr::label_strict;
@@ -17,6 +16,7 @@ use crate::plan::ExpressionPlan;
1716
use crate::plan::Plan;
1817
use crate::plan::PlanRef;
1918
use crate::plan::new_plan;
19+
use crate::plan::optimizer::PlanParentReduceRule;
2020

2121
/// A physical dictionary plan with children ordered as `[codes, values]`.
2222
pub struct DictPlan {
@@ -73,33 +73,6 @@ impl Plan for DictPlan {
7373
Ok(Arc::new(self.with_children(codes, values)))
7474
}
7575

76-
fn optimize_expression(&self, expression: &Expression) -> VortexResult<Option<PlanRef>> {
77-
if !expression.return_dtype(&self.dtype)?.is_boolean() {
78-
return Ok(None);
79-
}
80-
let references_root = label_tree(expression, is_root, |acc, &child| acc | child)
81-
.get(expression)
82-
.copied()
83-
.unwrap_or(false);
84-
let is_strict = label_strict(expression)
85-
.get(expression)
86-
.copied()
87-
.unwrap_or(false);
88-
let is_fallible = label_is_fallible(expression)
89-
.get(expression)
90-
.copied()
91-
.unwrap_or(true);
92-
if !references_root || !is_strict || is_fallible {
93-
return Ok(None);
94-
}
95-
96-
let values =
97-
ExpressionPlan::try_new(expression.clone(), Arc::clone(&self.values))?.optimize()?;
98-
Ok(Some(Arc::new(
99-
self.with_children(Arc::clone(&self.codes), values),
100-
)))
101-
}
102-
10376
fn dtype(&self) -> &vortex_array::dtype::DType {
10477
&self.dtype
10578
}
@@ -128,3 +101,44 @@ impl Plan for DictPlan {
128101
}
129102
}
130103
}
104+
105+
/// Pushes a safe boolean expression into dictionary values.
106+
#[derive(Debug)]
107+
pub(crate) struct ExpressionDictRule;
108+
109+
impl PlanParentReduceRule<DictPlan> for ExpressionDictRule {
110+
type Parent = ExpressionPlan;
111+
112+
fn reduce_parent(
113+
&self,
114+
child: &DictPlan,
115+
parent: &ExpressionPlan,
116+
_child_idx: usize,
117+
) -> VortexResult<Option<PlanRef>> {
118+
let expression = parent.expression();
119+
if !expression.return_dtype(&child.dtype)?.is_boolean() {
120+
return Ok(None);
121+
}
122+
let references_root = label_tree(expression, is_root, |acc, &child| acc | child)
123+
.get(expression)
124+
.copied()
125+
.unwrap_or(false);
126+
let is_strict = label_strict(expression)
127+
.get(expression)
128+
.copied()
129+
.unwrap_or(false);
130+
let is_fallible = label_is_fallible(expression)
131+
.get(expression)
132+
.copied()
133+
.unwrap_or(true);
134+
if !references_root || !is_strict || is_fallible {
135+
return Ok(None);
136+
}
137+
138+
let values =
139+
ExpressionPlan::try_new(expression.clone(), Arc::clone(&child.values))?.optimize()?;
140+
Ok(Some(Arc::new(
141+
child.with_children(Arc::clone(&child.codes), values),
142+
)))
143+
}
144+
}

‎vortex-layout/src/plan/plans/expression.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ use vortex_error::vortex_bail;
1515

1616
use crate::plan::Plan;
1717
use crate::plan::PlanRef;
18+
use crate::plan::optimizer::reduce_parent;
1819

1920
/// A physical plan that applies an expression to the output of `child`.
2021
pub struct ExpressionPlan {
@@ -64,10 +65,11 @@ impl Plan for ExpressionPlan {
6465
let expression = replace(expression, &root(), inner.expression.clone());
6566
return Self::try_new(expression, Arc::clone(&inner.child))?.optimize();
6667
}
67-
if let Some(rewritten) = child.optimize_expression(&expression)? {
68+
let parent: PlanRef = Arc::new(Self::try_new(expression, child)?);
69+
if let Some(rewritten) = reduce_parent(&parent, 0)? {
6870
return Ok(rewritten);
6971
}
70-
Ok(Arc::new(Self::try_new(expression, child)?))
72+
Ok(parent)
7173
}
7274

7375
fn dtype(&self) -> &DType {

‎vortex-layout/src/plan/plans/mod.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,15 @@ mod row_idx;
1010
mod struct_;
1111

1212
pub use chunked::ChunkedPlan;
13+
pub(crate) use chunked::ExpressionChunkedRule;
1314
pub use dict::DictPlan;
15+
pub(crate) use dict::ExpressionDictRule;
1416
pub use expression::ExpressionPlan;
1517
pub use flat::FlatPlan;
1618
pub use list::ListPlan;
19+
pub(crate) use row_idx::ExpressionRowIdxRule;
1720
pub use row_idx::RowIdxPartitionPlan;
1821
pub use row_idx::RowIdxPlan;
1922
pub use row_idx::RowIdxValuesPlan;
23+
pub(crate) use struct_::ExpressionStructRule;
2024
pub use struct_::StructPlan;

0 commit comments

Comments
 (0)