Skip to content

Commit 26b4811

Browse files
peter-tothpeter-toth
committed
refactor TreeNode::apply and relatives
1 parent d28c79d commit 26b4811

File tree

18 files changed

+401
-500
lines changed

18 files changed

+401
-500
lines changed

‎datafusion/common/src/tree_node.rs‎

Lines changed: 224 additions & 91 deletions
Large diffs are not rendered by default.

‎datafusion/core/src/datasource/listing/helpers.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
5858
Expr::Column(Column { ref name, .. }) => {
5959
is_applicable &= col_names.contains(name);
6060
if is_applicable {
61-
VisitRecursion::Skip
61+
VisitRecursion::Prune
6262
} else {
6363
VisitRecursion::Stop
6464
}

‎datafusion/core/src/execution/context/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2036,9 +2036,9 @@ impl<'a> BadPlanVisitor<'a> {
20362036
}
20372037

20382038
impl<'a> TreeNodeVisitor for BadPlanVisitor<'a> {
2039-
type N = LogicalPlan;
2039+
type Node = LogicalPlan;
20402040

2041-
fn pre_visit(&mut self, node: &Self::N) -> Result<VisitRecursion> {
2041+
fn pre_visit(&mut self, node: &Self::Node) -> Result<VisitRecursion> {
20422042
match node {
20432043
LogicalPlan::Ddl(ddl) if !self.options.allow_ddl => {
20442044
plan_err!("DDL not supported: {}", ddl.name())

‎datafusion/core/src/physical_optimizer/enforce_distribution.rs‎

Lines changed: 7 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,9 @@ use crate::physical_plan::{
5050
};
5151

5252
use arrow::compute::SortOptions;
53-
use datafusion_common::tree_node::{Transformed, TreeNode, VisitRecursion};
53+
use datafusion_common::tree_node::{
54+
Transformed, TreeNode, TreeNodeIterator, VisitRecursion,
55+
};
5456
use datafusion_expr::logical_plan::JoinType;
5557
use datafusion_physical_expr::equivalence::EquivalenceProperties;
5658
use datafusion_physical_expr::expressions::{Column, NoOp};
@@ -1503,18 +1505,11 @@ impl DistributionContext {
15031505
}
15041506

15051507
impl TreeNode for DistributionContext {
1506-
fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
1508+
fn apply_children<F>(&self, f: &mut F) -> Result<VisitRecursion>
15071509
where
15081510
F: FnMut(&Self) -> Result<VisitRecursion>,
15091511
{
1510-
for child in self.children() {
1511-
match op(&child)? {
1512-
VisitRecursion::Continue => {}
1513-
VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
1514-
VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
1515-
}
1516-
}
1517-
Ok(VisitRecursion::Continue)
1512+
self.children().iter().for_each_continue_with(f)
15181513
}
15191514

15201515
fn map_children<F>(self, transform: F) -> Result<Self>
@@ -1593,20 +1588,11 @@ impl PlanWithKeyRequirements {
15931588
}
15941589

15951590
impl TreeNode for PlanWithKeyRequirements {
1596-
fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
1591+
fn apply_children<F>(&self, f: &mut F) -> Result<VisitRecursion>
15971592
where
15981593
F: FnMut(&Self) -> Result<VisitRecursion>,
15991594
{
1600-
let children = self.children();
1601-
for child in children {
1602-
match op(&child)? {
1603-
VisitRecursion::Continue => {}
1604-
VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
1605-
VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
1606-
}
1607-
}
1608-
1609-
Ok(VisitRecursion::Continue)
1595+
self.children().iter().for_each_continue_with(f)
16101596
}
16111597

16121598
fn map_children<F>(self, transform: F) -> Result<Self>

‎datafusion/core/src/physical_optimizer/enforce_sorting.rs‎

Lines changed: 7 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,9 @@ use crate::physical_plan::windows::{
5555
};
5656
use crate::physical_plan::{with_new_children_if_necessary, Distribution, ExecutionPlan};
5757

58-
use datafusion_common::tree_node::{Transformed, TreeNode, VisitRecursion};
58+
use datafusion_common::tree_node::{
59+
Transformed, TreeNode, TreeNodeIterator, VisitRecursion,
60+
};
5961
use datafusion_common::{plan_err, DataFusionError};
6062
use datafusion_physical_expr::utils::{
6163
ordering_satisfy, ordering_satisfy_requirement_concrete,
@@ -158,20 +160,11 @@ impl PlanWithCorrespondingSort {
158160
}
159161

160162
impl TreeNode for PlanWithCorrespondingSort {
161-
fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
163+
fn apply_children<F>(&self, f: &mut F) -> Result<VisitRecursion>
162164
where
163165
F: FnMut(&Self) -> Result<VisitRecursion>,
164166
{
165-
let children = self.children();
166-
for child in children {
167-
match op(&child)? {
168-
VisitRecursion::Continue => {}
169-
VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
170-
VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
171-
}
172-
}
173-
174-
Ok(VisitRecursion::Continue)
167+
self.children().iter().for_each_continue_with(f)
175168
}
176169

177170
fn map_children<F>(self, transform: F) -> Result<Self>
@@ -274,20 +267,11 @@ impl PlanWithCorrespondingCoalescePartitions {
274267
}
275268

276269
impl TreeNode for PlanWithCorrespondingCoalescePartitions {
277-
fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
270+
fn apply_children<F>(&self, f: &mut F) -> Result<VisitRecursion>
278271
where
279272
F: FnMut(&Self) -> Result<VisitRecursion>,
280273
{
281-
let children = self.children();
282-
for child in children {
283-
match op(&child)? {
284-
VisitRecursion::Continue => {}
285-
VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
286-
VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
287-
}
288-
}
289-
290-
Ok(VisitRecursion::Continue)
274+
self.children().iter().for_each_continue_with(f)
291275
}
292276

293277
fn map_children<F>(self, transform: F) -> Result<Self>

‎datafusion/core/src/physical_optimizer/pipeline_checker.rs‎

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,9 @@ use crate::physical_plan::joins::SymmetricHashJoinExec;
2828
use crate::physical_plan::{with_new_children_if_necessary, ExecutionPlan};
2929

3030
use datafusion_common::config::OptimizerOptions;
31-
use datafusion_common::tree_node::{Transformed, TreeNode, VisitRecursion};
31+
use datafusion_common::tree_node::{
32+
Transformed, TreeNode, TreeNodeIterator, VisitRecursion,
33+
};
3234
use datafusion_common::{plan_err, DataFusionError};
3335
use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supported};
3436

@@ -86,20 +88,14 @@ impl PipelineStatePropagator {
8688
}
8789

8890
impl TreeNode for PipelineStatePropagator {
89-
fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
91+
fn apply_children<F>(&self, f: &mut F) -> Result<VisitRecursion>
9092
where
9193
F: FnMut(&Self) -> Result<VisitRecursion>,
9294
{
93-
let children = self.plan.children();
94-
for child in children {
95-
match op(&PipelineStatePropagator::new(child))? {
96-
VisitRecursion::Continue => {}
97-
VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
98-
VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
99-
}
100-
}
101-
102-
Ok(VisitRecursion::Continue)
95+
self.plan
96+
.children()
97+
.into_iter()
98+
.for_each_continue_with(&mut |c| f(&PipelineStatePropagator::new(c)))
10399
}
104100

105101
fn map_children<F>(self, transform: F) -> Result<Self>

‎datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@ use crate::physical_plan::{with_new_children_if_necessary, ExecutionPlan};
3030
use super::utils::is_repartition;
3131

3232
use datafusion_common::config::ConfigOptions;
33-
use datafusion_common::tree_node::{Transformed, TreeNode, VisitRecursion};
33+
use datafusion_common::tree_node::{
34+
Transformed, TreeNode, TreeNodeIterator, VisitRecursion,
35+
};
3436
use datafusion_physical_expr::utils::ordering_satisfy;
3537
use datafusion_physical_plan::unbounded_output;
3638

@@ -119,18 +121,11 @@ impl OrderPreservationContext {
119121
}
120122

121123
impl TreeNode for OrderPreservationContext {
122-
fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
124+
fn apply_children<F>(&self, f: &mut F) -> Result<VisitRecursion>
123125
where
124126
F: FnMut(&Self) -> Result<VisitRecursion>,
125127
{
126-
for child in self.children() {
127-
match op(&child)? {
128-
VisitRecursion::Continue => {}
129-
VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
130-
VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
131-
}
132-
}
133-
Ok(VisitRecursion::Continue)
128+
self.children().iter().for_each_continue_with(f)
134129
}
135130

136131
fn map_children<F>(self, transform: F) -> Result<Self>

‎datafusion/core/src/physical_optimizer/sort_pushdown.rs‎

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,9 @@ use crate::physical_plan::repartition::RepartitionExec;
2828
use crate::physical_plan::sorts::sort::SortExec;
2929
use crate::physical_plan::{with_new_children_if_necessary, ExecutionPlan};
3030

31-
use datafusion_common::tree_node::{Transformed, TreeNode, VisitRecursion};
31+
use datafusion_common::tree_node::{
32+
Transformed, TreeNode, TreeNodeIterator, VisitRecursion,
33+
};
3234
use datafusion_common::{
3335
plan_datafusion_err, plan_err, DataFusionError, JoinSide, Result,
3436
};
@@ -85,20 +87,11 @@ impl SortPushDown {
8587
}
8688

8789
impl TreeNode for SortPushDown {
88-
fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
90+
fn apply_children<F>(&self, f: &mut F) -> Result<VisitRecursion>
8991
where
9092
F: FnMut(&Self) -> Result<VisitRecursion>,
9193
{
92-
let children = self.children();
93-
for child in children {
94-
match op(&child)? {
95-
VisitRecursion::Continue => {}
96-
VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
97-
VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
98-
}
99-
}
100-
101-
Ok(VisitRecursion::Continue)
94+
self.children().iter().for_each_continue_with(f)
10295
}
10396

10497
fn map_children<F>(mut self, transform: F) -> Result<Self>

‎datafusion/expr/src/logical_plan/display.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ impl<'a, 'b> IndentVisitor<'a, 'b> {
4949
}
5050

5151
impl<'a, 'b> TreeNodeVisitor for IndentVisitor<'a, 'b> {
52-
type N = LogicalPlan;
52+
type Node = LogicalPlan;
5353

5454
fn pre_visit(
5555
&mut self,
@@ -171,7 +171,7 @@ impl<'a, 'b> GraphvizVisitor<'a, 'b> {
171171
}
172172

173173
impl<'a, 'b> TreeNodeVisitor for GraphvizVisitor<'a, 'b> {
174-
type N = LogicalPlan;
174+
type Node = LogicalPlan;
175175

176176
fn pre_visit(
177177
&mut self,

0 commit comments

Comments
 (0)