Skip to content

Commit a84e5f8

Browse files
peter-tothberkaysynnadamustafasrepoozankabak
authored
Consolidate TreeNode transform and rewrite APIs (apache#8891)
* refactor `TreeNode::rewrite()` * use handle_tree_recursion in `Expr` * use macro for transform recursions * fix api * minor fixes * fix * don't trust `t.transformed` coming from transformation closures, keep the old way of detecting if changes were made * rephrase todo comment, always propagate up `t.transformed` from the transformation closure, fix projection pushdown closure * Fix `TreeNodeRecursion` docs * extend Skip (Prune) functionality to Jump as it is defined in https://synnada.notion.site/synnada/TreeNode-Design-Proposal-bceac27d18504a2085145550e267c4c1 * fix Jump and add tests * jump test fixes * fix clippy * unify "transform" traversals using macros, fix "visit" traversal jumps, add visit jump tests, ensure consistent naming `f` instead of `op`, `f_down` instead of `pre_visit` and `f_up` instead of `post_visit` * fix macro rewrite * minor fixes * minor fix * refactor tests * add transform tests * add apply, transform_down and transform_up tests * refactor tests * test jump on both a and e nodes in both top-down and bottom-up traversals * better transform/rewrite tests * minor fix * simplify tests * add stop tests, reorganize tests * fix previous merges and remove leftover file * Review TreeNode Refactor (#1) * Minor changes * Jump doesn't ignore f_up * update test * Update rewriter * LogicalPlan visit update and propagate from children flags * Update tree_node.rs * Update map_children's --------- Co-authored-by: Mustafa Akur <[email protected]> * fix * minor fixes * fix f_up call when f_down returns jump * simplify code * minor fix * revert unnecessary changes * fix `DynTreeNode` and `ConcreteTreeNode` `transformed` and `tnr` propagation * introduce TransformedResult helper * fix docs * restore transform as alias to trassform_up * restore transform as alias to trassform_up 2 * Simplifications and comment improvements (#2) --------- Co-authored-by: Berkay Şahin <[email protected]> Co-authored-by: Mustafa Akur <[email protected]> Co-authored-by: Mehmet Ozan Kabak <[email protected]>
1 parent 684b4fa commit a84e5f8

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

63 files changed

+3074
-1579
lines changed

datafusion-examples/examples/rewrite_expr.rs

Lines changed: 11 additions & 8 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::{Transformed, TransformedResult, TreeNode};
2121
use datafusion_common::{plan_err, Result, ScalarValue};
2222
use datafusion_expr::{
2323
AggregateUDF, Between, Expr, Filter, LogicalPlan, ScalarUDF, TableSource, WindowUDF,
@@ -95,14 +95,15 @@ 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(
98+
Transformed::yes(LogicalPlan::Filter(Filter::try_new(
9999
predicate,
100100
filter.input,
101101
)?))
102102
}
103-
_ => Transformed::No(plan),
103+
_ => Transformed::no(plan),
104104
})
105105
})
106+
.data()
106107
}
107108

108109
fn analyze_expr(expr: Expr) -> Result<Expr> {
@@ -111,13 +112,14 @@ impl MyAnalyzerRule {
111112
Ok(match expr {
112113
Expr::Literal(ScalarValue::Int64(i)) => {
113114
// transform to UInt64
114-
Transformed::Yes(Expr::Literal(ScalarValue::UInt64(
115+
Transformed::yes(Expr::Literal(ScalarValue::UInt64(
115116
i.map(|i| i as u64),
116117
)))
117118
}
118-
_ => Transformed::No(expr),
119+
_ => Transformed::no(expr),
119120
})
120121
})
122+
.data()
121123
}
122124
}
123125

@@ -175,14 +177,15 @@ fn my_rewrite(expr: Expr) -> Result<Expr> {
175177
let low: Expr = *low;
176178
let high: Expr = *high;
177179
if negated {
178-
Transformed::Yes(expr.clone().lt(low).or(expr.gt(high)))
180+
Transformed::yes(expr.clone().lt(low).or(expr.gt(high)))
179181
} else {
180-
Transformed::Yes(expr.clone().gt_eq(low).and(expr.lt_eq(high)))
182+
Transformed::yes(expr.clone().gt_eq(low).and(expr.lt_eq(high)))
181183
}
182184
}
183-
_ => Transformed::No(expr),
185+
_ => Transformed::no(expr),
184186
})
185187
})
188+
.data()
186189
}
187190

188191
#[derive(Default)]

0 commit comments

Comments
 (0)