Skip to content

Commit 6094565

Browse files
committed
remove Transformed
1 parent 36877f3 commit 6094565

37 files changed

+247
-334
lines changed

datafusion-examples/examples/rewrite_expr.rs

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
1919
use datafusion_common::config::ConfigOptions;
20-
use datafusion_common::tree_node::{Transformed, TreeNode};
20+
use datafusion_common::tree_node::TreeNode;
2121
use datafusion_common::{plan_err, DataFusionError, Result, ScalarValue};
2222
use datafusion_expr::{
2323
AggregateUDF, Between, Expr, Filter, LogicalPlan, ScalarUDF, TableSource, WindowUDF,
@@ -95,12 +95,9 @@ impl MyAnalyzerRule {
9595
Ok(match plan {
9696
LogicalPlan::Filter(filter) => {
9797
let predicate = Self::analyze_expr(filter.predicate.clone())?;
98-
Transformed::Yes(LogicalPlan::Filter(Filter::try_new(
99-
predicate,
100-
filter.input,
101-
)?))
98+
LogicalPlan::Filter(Filter::try_new(predicate, filter.input)?)
10299
}
103-
_ => Transformed::No(plan),
100+
_ => plan,
104101
})
105102
})
106103
}
@@ -111,11 +108,9 @@ impl MyAnalyzerRule {
111108
Ok(match expr {
112109
Expr::Literal(ScalarValue::Int64(i)) => {
113110
// transform to UInt64
114-
Transformed::Yes(Expr::Literal(ScalarValue::UInt64(
115-
i.map(|i| i as u64),
116-
)))
111+
Expr::Literal(ScalarValue::UInt64(i.map(|i| i as u64)))
117112
}
118-
_ => Transformed::No(expr),
113+
_ => expr,
119114
})
120115
})
121116
}
@@ -175,12 +170,12 @@ fn my_rewrite(expr: Expr) -> Result<Expr> {
175170
let low: Expr = *low;
176171
let high: Expr = *high;
177172
if negated {
178-
Transformed::Yes(expr.clone().lt(low).or(expr.gt(high)))
173+
expr.clone().lt(low).or(expr.gt(high))
179174
} else {
180-
Transformed::Yes(expr.clone().gt_eq(low).and(expr.lt_eq(high)))
175+
expr.clone().gt_eq(low).and(expr.lt_eq(high))
181176
}
182177
}
183-
_ => Transformed::No(expr),
178+
_ => expr,
184179
})
185180
})
186181
}

datafusion/common/src/tree_node.rs

Lines changed: 9 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ pub trait TreeNode: Sized + Clone {
113113
/// The default tree traversal direction is transform_up(Postorder Traversal).
114114
fn transform<F>(self, op: &F) -> Result<Self>
115115
where
116-
F: Fn(Self) -> Result<Transformed<Self>>,
116+
F: Fn(Self) -> Result<Self>,
117117
{
118118
self.transform_up(op)
119119
}
@@ -123,9 +123,9 @@ pub trait TreeNode: Sized + Clone {
123123
/// When the `op` does not apply to a given node, it is left unchanged.
124124
fn transform_down<F>(self, op: &F) -> Result<Self>
125125
where
126-
F: Fn(Self) -> Result<Transformed<Self>>,
126+
F: Fn(Self) -> Result<Self>,
127127
{
128-
let after_op = op(self)?.into();
128+
let after_op = op(self)?;
129129
after_op.map_children(|node| node.transform_down(op))
130130
}
131131

@@ -134,9 +134,9 @@ pub trait TreeNode: Sized + Clone {
134134
/// When the `op` does not apply to a given node, it is left unchanged.
135135
fn transform_down_mut<F>(self, op: &mut F) -> Result<Self>
136136
where
137-
F: FnMut(Self) -> Result<Transformed<Self>>,
137+
F: FnMut(Self) -> Result<Self>,
138138
{
139-
let after_op = op(self)?.into();
139+
let after_op = op(self)?;
140140
after_op.map_children(|node| node.transform_down_mut(op))
141141
}
142142

@@ -145,11 +145,11 @@ pub trait TreeNode: Sized + Clone {
145145
/// When the `op` does not apply to a given node, it is left unchanged.
146146
fn transform_up<F>(self, op: &F) -> Result<Self>
147147
where
148-
F: Fn(Self) -> Result<Transformed<Self>>,
148+
F: Fn(Self) -> Result<Self>,
149149
{
150150
let after_op_children = self.map_children(|node| node.transform_up(op))?;
151151

152-
let new_node = op(after_op_children)?.into();
152+
let new_node = op(after_op_children)?;
153153
Ok(new_node)
154154
}
155155

@@ -158,11 +158,11 @@ pub trait TreeNode: Sized + Clone {
158158
/// When the `op` does not apply to a given node, it is left unchanged.
159159
fn transform_up_mut<F>(self, op: &mut F) -> Result<Self>
160160
where
161-
F: FnMut(Self) -> Result<Transformed<Self>>,
161+
F: FnMut(Self) -> Result<Self>,
162162
{
163163
let after_op_children = self.map_children(|node| node.transform_up_mut(op))?;
164164

165-
let new_node = op(after_op_children)?.into();
165+
let new_node = op(after_op_children)?;
166166
Ok(new_node)
167167
}
168168

@@ -314,29 +314,6 @@ pub enum VisitRecursion {
314314
Stop,
315315
}
316316

317-
pub enum Transformed<T> {
318-
/// The item was transformed / rewritten somehow
319-
Yes(T),
320-
/// The item was not transformed
321-
No(T),
322-
}
323-
324-
impl<T> Transformed<T> {
325-
pub fn into(self) -> T {
326-
match self {
327-
Transformed::Yes(t) => t,
328-
Transformed::No(t) => t,
329-
}
330-
}
331-
332-
pub fn into_pair(self) -> (T, bool) {
333-
match self {
334-
Transformed::Yes(t) => (t, true),
335-
Transformed::No(t) => (t, false),
336-
}
337-
}
338-
}
339-
340317
/// Helper trait for implementing [`TreeNode`] that have children stored as Arc's
341318
///
342319
/// If some trait object, such as `dyn T`, implements this trait,

datafusion/core/src/physical_optimizer/coalesce_batches.rs

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ use crate::{
2727
repartition::RepartitionExec, Partitioning,
2828
},
2929
};
30-
use datafusion_common::tree_node::{Transformed, TreeNode};
30+
use datafusion_common::tree_node::TreeNode;
3131
use std::sync::Arc;
3232

3333
/// Optimizer rule that introduces CoalesceBatchesExec to avoid overhead with small batches that
@@ -71,12 +71,9 @@ impl PhysicalOptimizerRule for CoalesceBatches {
7171
})
7272
.unwrap_or(false);
7373
if wrap_in_coalesce {
74-
Ok(Transformed::Yes(Arc::new(CoalesceBatchesExec::new(
75-
plan,
76-
target_batch_size,
77-
))))
74+
Ok(Arc::new(CoalesceBatchesExec::new(plan, target_batch_size)))
7875
} else {
79-
Ok(Transformed::No(plan))
76+
Ok(plan)
8077
}
8178
})
8279
}

datafusion/core/src/physical_optimizer/combine_partial_final_agg.rs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ use crate::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGro
2626
use crate::physical_plan::ExecutionPlan;
2727

2828
use datafusion_common::config::ConfigOptions;
29-
use datafusion_common::tree_node::{Transformed, TreeNode};
29+
use datafusion_common::tree_node::TreeNode;
3030
use datafusion_physical_expr::expressions::Column;
3131
use datafusion_physical_expr::{AggregateExpr, PhysicalExpr};
3232

@@ -109,9 +109,9 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
109109
});
110110

111111
Ok(if let Some(transformed) = transformed {
112-
Transformed::Yes(transformed)
112+
transformed
113113
} else {
114-
Transformed::No(plan)
114+
plan
115115
})
116116
})
117117
}
@@ -185,9 +185,9 @@ fn discard_column_index(group_expr: Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalEx
185185
None => None,
186186
};
187187
Ok(if let Some(normalized_form) = normalized_form {
188-
Transformed::Yes(normalized_form)
188+
normalized_form
189189
} else {
190-
Transformed::No(expr)
190+
expr
191191
})
192192
})
193193
.unwrap_or(group_expr)

datafusion/core/src/physical_optimizer/enforce_distribution.rs

Lines changed: 20 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ use crate::physical_plan::{
4848
};
4949

5050
use arrow::compute::SortOptions;
51-
use datafusion_common::tree_node::{Transformed, TreeNode};
51+
use datafusion_common::tree_node::TreeNode;
5252
use datafusion_expr::logical_plan::JoinType;
5353
use datafusion_physical_expr::expressions::{Column, NoOp};
5454
use datafusion_physical_expr::utils::map_columns_before_projection;
@@ -205,9 +205,7 @@ impl PhysicalOptimizerRule for EnforceDistribution {
205205
adjusted.plan
206206
} else {
207207
// Run a bottom-up process
208-
plan.transform_up(&|plan| {
209-
Ok(Transformed::Yes(reorder_join_keys_to_inputs(plan)?))
210-
})?
208+
plan.transform_up(&reorder_join_keys_to_inputs)?
211209
};
212210

213211
let distribution_context = DistributionContext::new(adjusted);
@@ -272,7 +270,7 @@ impl PhysicalOptimizerRule for EnforceDistribution {
272270
///
273271
fn adjust_input_keys_ordering(
274272
mut requirements: PlanWithKeyRequirements,
275-
) -> Result<Transformed<PlanWithKeyRequirements>> {
273+
) -> Result<PlanWithKeyRequirements> {
276274
let parent_required = requirements.required_key_ordering.clone();
277275
let plan_any = requirements.plan.as_any();
278276

@@ -309,7 +307,6 @@ fn adjust_input_keys_ordering(
309307
vec![],
310308
&join_constructor,
311309
)
312-
.map(Transformed::Yes)
313310
}
314311
PartitionMode::CollectLeft => {
315312
let new_right_request = match join_type {
@@ -329,13 +326,11 @@ fn adjust_input_keys_ordering(
329326
// Push down requirements to the right side
330327
requirements.children[1].required_key_ordering =
331328
new_right_request.unwrap_or(vec![]);
332-
Ok(Transformed::Yes(requirements))
329+
Ok(requirements)
333330
}
334331
PartitionMode::Auto => {
335332
// Can not satisfy, clear the current requirements and generate new empty requirements
336-
Ok(Transformed::Yes(PlanWithKeyRequirements::new(
337-
requirements.plan,
338-
)))
333+
Ok(PlanWithKeyRequirements::new(requirements.plan))
339334
}
340335
}
341336
} else if let Some(CrossJoinExec { left, .. }) =
@@ -345,7 +340,7 @@ fn adjust_input_keys_ordering(
345340
// Push down requirements to the right side
346341
requirements.children[1].required_key_ordering =
347342
shift_right_required(&parent_required, left_columns_len).unwrap_or_default();
348-
Ok(Transformed::Yes(requirements))
343+
Ok(requirements)
349344
} else if let Some(SortMergeJoinExec {
350345
left,
351346
right,
@@ -375,23 +370,19 @@ fn adjust_input_keys_ordering(
375370
sort_options.clone(),
376371
&join_constructor,
377372
)
378-
.map(Transformed::Yes)
379373
} else if let Some(aggregate_exec) = plan_any.downcast_ref::<AggregateExec>() {
380374
if !parent_required.is_empty() {
381375
match aggregate_exec.mode() {
382376
AggregateMode::FinalPartitioned => reorder_aggregate_keys(
383377
requirements.plan.clone(),
384378
&parent_required,
385379
aggregate_exec,
386-
)
387-
.map(Transformed::Yes),
388-
_ => Ok(Transformed::Yes(PlanWithKeyRequirements::new(
389-
requirements.plan,
390-
))),
380+
),
381+
_ => Ok(PlanWithKeyRequirements::new(requirements.plan)),
391382
}
392383
} else {
393384
// Keep everything unchanged
394-
Ok(Transformed::No(requirements))
385+
Ok(requirements)
395386
}
396387
} else if let Some(proj) = plan_any.downcast_ref::<ProjectionExec>() {
397388
let expr = proj.expr();
@@ -401,26 +392,22 @@ fn adjust_input_keys_ordering(
401392
let new_required = map_columns_before_projection(&parent_required, expr);
402393
if new_required.len() == parent_required.len() {
403394
requirements.children[0].required_key_ordering = new_required;
404-
Ok(Transformed::Yes(requirements))
395+
Ok(requirements)
405396
} else {
406397
// Can not satisfy, clear the current requirements and generate new empty requirements
407-
Ok(Transformed::Yes(PlanWithKeyRequirements::new(
408-
requirements.plan,
409-
)))
398+
Ok(PlanWithKeyRequirements::new(requirements.plan))
410399
}
411400
} else if plan_any.downcast_ref::<RepartitionExec>().is_some()
412401
|| plan_any.downcast_ref::<CoalescePartitionsExec>().is_some()
413402
|| plan_any.downcast_ref::<WindowAggExec>().is_some()
414403
{
415-
Ok(Transformed::Yes(PlanWithKeyRequirements::new(
416-
requirements.plan,
417-
)))
404+
Ok(PlanWithKeyRequirements::new(requirements.plan))
418405
} else {
419406
// By default, push down the parent requirements to children
420407
requirements.children.iter_mut().for_each(|child| {
421408
child.required_key_ordering = parent_required.clone();
422409
});
423-
Ok(Transformed::Yes(requirements))
410+
Ok(requirements)
424411
}
425412
}
426413

@@ -1139,11 +1126,11 @@ fn add_sort_preserving_partitions(
11391126
fn ensure_distribution(
11401127
dist_context: DistributionContext,
11411128
config: &ConfigOptions,
1142-
) -> Result<Transformed<DistributionContext>> {
1129+
) -> Result<DistributionContext> {
11431130
let dist_context = dist_context.update_children()?;
11441131

11451132
if dist_context.plan.children().is_empty() {
1146-
return Ok(Transformed::No(dist_context));
1133+
return Ok(dist_context);
11471134
}
11481135

11491136
let target_partitions = config.execution.target_partitions;
@@ -1333,7 +1320,7 @@ fn ensure_distribution(
13331320
children_nodes,
13341321
};
13351322

1336-
Ok(Transformed::Yes(new_distribution_context))
1323+
Ok(new_distribution_context)
13371324
}
13381325

13391326
/// A struct to keep track of distribution changing operators
@@ -1395,7 +1382,7 @@ impl DistributionContext {
13951382
.collect::<Vec<_>>();
13961383

13971384
Ok(Self {
1398-
plan: with_new_children_if_necessary(self.plan, children_plans)?.into(),
1385+
plan: with_new_children_if_necessary(self.plan, children_plans)?,
13991386
distribution_connection: false,
14001387
children_nodes: self.children_nodes,
14011388
})
@@ -1420,8 +1407,7 @@ impl TreeNode for DistributionContext {
14201407
self.plan = with_new_children_if_necessary(
14211408
self.plan,
14221409
self.children_nodes.iter().map(|c| c.plan.clone()).collect(),
1423-
)?
1424-
.into();
1410+
)?;
14251411
}
14261412
Ok(self)
14271413
}
@@ -1484,8 +1470,7 @@ impl TreeNode for PlanWithKeyRequirements {
14841470
self.plan = with_new_children_if_necessary(
14851471
self.plan,
14861472
self.children.iter().map(|c| c.plan.clone()).collect(),
1487-
)?
1488-
.into();
1473+
)?;
14891474
}
14901475
Ok(self)
14911476
}
@@ -1912,7 +1897,7 @@ pub(crate) mod tests {
19121897
config.optimizer.repartition_file_scans = false;
19131898
config.optimizer.repartition_file_min_size = 1024;
19141899
config.optimizer.prefer_existing_sort = prefer_existing_sort;
1915-
ensure_distribution(distribution_context, &config).map(|item| item.into().plan)
1900+
ensure_distribution(distribution_context, &config).map(|item| item.plan)
19161901
}
19171902

19181903
/// Test whether plan matches with expected plan

0 commit comments

Comments
 (0)