Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Consolidate TreeNode transform and rewrite APIs #8891

Merged
merged 54 commits into from
Mar 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
54 commits
Select commit Hold shift + click to select a range
212c489
refactor `TreeNode::rewrite()`
peter-toth Jan 13, 2024
c52f134
use handle_tree_recursion in `Expr`
peter-toth Jan 29, 2024
3fd2214
use macro for transform recursions
peter-toth Jan 29, 2024
efd613f
Merge commit '488cfe117c218729ac4009010d85ab6f77e83eea' into refactor…
peter-toth Jan 30, 2024
35f9006
fix api
peter-toth Jan 31, 2024
f6150be
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Jan 31, 2024
6f763dc
minor fixes
peter-toth Jan 31, 2024
84d91c6
fix
peter-toth Feb 1, 2024
3653aa6
don't trust `t.transformed` coming from transformation closures, keep…
peter-toth Feb 1, 2024
8295ae5
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Feb 2, 2024
5b94531
rephrase todo comment, always propagate up `t.transformed` from the t…
peter-toth Feb 2, 2024
277cef7
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Feb 2, 2024
2ede5e6
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Feb 5, 2024
5c2baef
Fix `TreeNodeRecursion` docs
peter-toth Feb 7, 2024
9c0ec41
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Feb 9, 2024
dcf0189
extend Skip (Prune) functionality to Jump as it is defined in https:/…
peter-toth Feb 10, 2024
6edd05f
fix Jump and add tests
peter-toth Feb 10, 2024
623e5fa
jump test fixes
peter-toth Feb 10, 2024
492777c
fix clippy
peter-toth Feb 10, 2024
384c3e0
unify "transform" traversals using macros, fix "visit" traversal jump…
peter-toth Feb 11, 2024
c224de7
fix macro rewrite
peter-toth Feb 11, 2024
87ad995
minor fixes
peter-toth Feb 15, 2024
a78a4ee
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Feb 15, 2024
8c34a9c
minor fix
peter-toth Feb 16, 2024
f2fa09a
refactor tests
peter-toth Feb 16, 2024
a2c791e
add transform tests
peter-toth Feb 16, 2024
b9a2afa
add apply, transform_down and transform_up tests
peter-toth Feb 16, 2024
83fe709
refactor tests
peter-toth Feb 19, 2024
e5dbc7e
test jump on both a and e nodes in both top-down and bottom-up traver…
peter-toth Feb 19, 2024
cbdc649
better transform/rewrite tests
peter-toth Feb 19, 2024
545d14a
minor fix
peter-toth Feb 19, 2024
af989a0
simplify tests
peter-toth Feb 19, 2024
79db28d
add stop tests, reorganize tests
peter-toth Feb 20, 2024
a2b4c48
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Feb 20, 2024
971b50b
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Feb 23, 2024
1475306
Merge commit 'ff36f6d99285404cf45f870ccc1c650411fdb335' into refactor…
peter-toth Feb 25, 2024
e1a499a
fix previous merges and remove leftover file
peter-toth Feb 26, 2024
271a0fd
Review TreeNode Refactor (#1)
berkaysynnada Feb 29, 2024
6f8f940
fix
peter-toth Feb 29, 2024
61272ea
minor fixes
peter-toth Feb 29, 2024
54a1b46
fix f_up call when f_down returns jump
peter-toth Feb 29, 2024
f780007
simplify code
peter-toth Feb 29, 2024
e75f2e3
minor fix
peter-toth Feb 29, 2024
8d34349
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Feb 29, 2024
60e76ac
revert unnecessary changes
peter-toth Feb 29, 2024
3af92ce
fix `DynTreeNode` and `ConcreteTreeNode` `transformed` and `tnr` prop…
peter-toth Feb 29, 2024
6b4b6dd
introduce TransformedResult helper
peter-toth Mar 1, 2024
171cd72
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Mar 1, 2024
08e3c7b
fix docs
peter-toth Mar 1, 2024
9f4b28a
restore transform as alias to trassform_up
peter-toth Mar 1, 2024
af6ab4a
restore transform as alias to trassform_up 2
peter-toth Mar 1, 2024
59dc032
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Mar 1, 2024
8306c9d
Merge branch 'main' into refactor-treenode-rewrite
peter-toth Mar 4, 2024
c944b81
Simplifications and comment improvements (#2)
ozankabak Mar 4, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 11 additions & 8 deletions datafusion-examples/examples/rewrite_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_common::{plan_err, Result, ScalarValue};
use datafusion_expr::{
AggregateUDF, Between, Expr, Filter, LogicalPlan, ScalarUDF, TableSource, WindowUDF,
Expand Down Expand Up @@ -95,14 +95,15 @@ impl MyAnalyzerRule {
Ok(match plan {
LogicalPlan::Filter(filter) => {
let predicate = Self::analyze_expr(filter.predicate.clone())?;
Transformed::Yes(LogicalPlan::Filter(Filter::try_new(
Transformed::yes(LogicalPlan::Filter(Filter::try_new(
predicate,
filter.input,
)?))
}
_ => Transformed::No(plan),
_ => Transformed::no(plan),
})
})
.data()
}

fn analyze_expr(expr: Expr) -> Result<Expr> {
Expand All @@ -111,13 +112,14 @@ impl MyAnalyzerRule {
Ok(match expr {
Expr::Literal(ScalarValue::Int64(i)) => {
// transform to UInt64
Transformed::Yes(Expr::Literal(ScalarValue::UInt64(
Transformed::yes(Expr::Literal(ScalarValue::UInt64(
i.map(|i| i as u64),
)))
}
_ => Transformed::No(expr),
_ => Transformed::no(expr),
})
})
.data()
}
}

Expand Down Expand Up @@ -175,14 +177,15 @@ fn my_rewrite(expr: Expr) -> Result<Expr> {
let low: Expr = *low;
let high: Expr = *high;
if negated {
Transformed::Yes(expr.clone().lt(low).or(expr.gt(high)))
Transformed::yes(expr.clone().lt(low).or(expr.gt(high)))
} else {
Transformed::Yes(expr.clone().gt_eq(low).and(expr.lt_eq(high)))
Transformed::yes(expr.clone().gt_eq(low).and(expr.lt_eq(high)))
}
}
_ => Transformed::No(expr),
_ => Transformed::no(expr),
})
})
.data()
}

#[derive(Default)]
Expand Down
Loading
Loading