Skip to content

Commit

Permalink
fix
Browse files Browse the repository at this point in the history
  • Loading branch information
@peter-toth
peter-toth committed Feb 1, 2024
1 parent 6f763dc commit 84d91c6
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 35 deletions.
32 changes: 16 additions & 16 deletions datafusion/common/src/tree_node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,9 +135,9 @@ pub trait TreeNode: Sized {
FD: FnMut(Self) -> Result<Transformed<Self>>,
FU: FnMut(Self) -> Result<Transformed<Self>>,
{
f_down(self)?.and_then_transform_children(|t| {
t.map_children(|node| node.transform(f_down, f_up))?
.and_then_transform_sibling(f_up)
f_down(self)?.and_then_transform_children(|n| {
n.map_children(|c| c.transform(f_down, f_up))?
.and_then_transform(f_up)
})
}

Expand All @@ -148,7 +148,7 @@ pub trait TreeNode: Sized {
where
F: Fn(Self) -> Result<Transformed<Self>>,
{
f(self)?.and_then_transform_children(|t| t.map_children(|n| n.transform_down(f)))
f(self)?.and_then_transform_children(|n| n.map_children(|c| c.transform_down(f)))
}

/// Convenience utils for writing optimizers rule: recursively apply the given 'op' to the node and all of its
Expand All @@ -159,7 +159,7 @@ pub trait TreeNode: Sized {
F: FnMut(Self) -> Result<Transformed<Self>>,
{
f(self)?
.and_then_transform_children(|t| t.map_children(|n| n.transform_down_mut(f)))
.and_then_transform_children(|n| n.map_children(|c| c.transform_down_mut(f)))
}

/// Convenience utils for writing optimizers rule: recursively apply the given 'op' first to all of its
Expand All @@ -169,8 +169,8 @@ pub trait TreeNode: Sized {
where
F: Fn(Self) -> Result<Transformed<Self>>,
{
self.map_children(|node| node.transform_up(f))?
.and_then_transform_sibling(f)
self.map_children(|c| c.transform_up(f))?
.and_then_transform(f)
}

/// Convenience utils for writing optimizers rule: recursively apply the given 'op' first to all of its
Expand All @@ -180,8 +180,8 @@ pub trait TreeNode: Sized {
where
F: FnMut(Self) -> Result<Transformed<Self>>,
{
self.map_children(|n| n.transform_up_mut(f))?
.and_then_transform_sibling(f)
self.map_children(|c| c.transform_up_mut(f))?
.and_then_transform(f)
}

/// Implements the [visitor pattern](https://en.wikipedia.org/wiki/Visitor_pattern) for
Expand Down Expand Up @@ -212,9 +212,9 @@ pub trait TreeNode: Sized {
self,
rewriter: &mut R,
) -> Result<Transformed<Self>> {
rewriter.f_down(self)?.and_then_transform_children(|t| {
t.map_children(|n| n.rewrite(rewriter))?
.and_then_transform_sibling(|t| rewriter.f_up(t))
rewriter.f_down(self)?.and_then_transform_children(|n| {
n.map_children(|c| c.rewrite(rewriter))?
.and_then_transform(|n| rewriter.f_up(n))
})
}

Expand Down Expand Up @@ -349,7 +349,7 @@ impl<T> Transformed<T> {
})
}

fn and_then_transform<F: FnOnce(T) -> Result<Transformed<T>>>(
fn and_then<F: FnOnce(T) -> Result<Transformed<T>>>(
self,
f: F,
children: bool,
Expand All @@ -375,18 +375,18 @@ impl<T> Transformed<T> {
})
}

pub fn and_then_transform_sibling<F: FnOnce(T) -> Result<Transformed<T>>>(
pub fn and_then_transform<F: FnOnce(T) -> Result<Transformed<T>>>(
self,
f: F,
) -> Result<Transformed<T>> {
self.and_then_transform(f, false)
self.and_then(f, false)
}

pub fn and_then_transform_children<F: FnOnce(T) -> Result<Transformed<T>>>(
self,
f: F,
) -> Result<Transformed<T>> {
self.and_then_transform(f, true)
self.and_then(f, true)
}
}

Expand Down
36 changes: 17 additions & 19 deletions datafusion/expr/src/tree_node/expr.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,7 +164,7 @@ impl TreeNode for Expr {
Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
transform_box(left, &mut f)?
.map_data(|new_left| (new_left, right))
.and_then_transform_sibling(|(new_left, right)| {
.and_then_transform(|(new_left, right)| {
Ok(transform_box(right, &mut f)?
.map_data(|new_right| (new_left, new_right)))
})?
Expand All @@ -180,7 +180,7 @@ impl TreeNode for Expr {
case_insensitive,
}) => transform_box(expr, &mut f)?
.map_data(|new_expr| (new_expr, pattern))
.and_then_transform_sibling(|(new_expr, pattern)| {
.and_then_transform(|(new_expr, pattern)| {
Ok(transform_box(pattern, &mut f)?
.map_data(|new_pattern| (new_expr, new_pattern)))
})?
Expand All @@ -201,7 +201,7 @@ impl TreeNode for Expr {
case_insensitive,
}) => transform_box(expr, &mut f)?
.map_data(|new_expr| (new_expr, pattern))
.and_then_transform_sibling(|(new_expr, pattern)| {
.and_then_transform(|(new_expr, pattern)| {
Ok(transform_box(pattern, &mut f)?
.map_data(|new_pattern| (new_expr, new_pattern)))
})?
Expand Down Expand Up @@ -241,11 +241,11 @@ impl TreeNode for Expr {
high,
}) => transform_box(expr, &mut f)?
.map_data(|new_expr| (new_expr, low, high))
.and_then_transform_sibling(|(new_expr, low, high)| {
.and_then_transform(|(new_expr, low, high)| {
Ok(transform_box(low, &mut f)?
.map_data(|new_low| (new_expr, new_low, high)))
})?
.and_then_transform_sibling(|(new_expr, new_low, high)| {
.and_then_transform(|(new_expr, new_low, high)| {
Ok(transform_box(high, &mut f)?
.map_data(|new_high| (new_expr, new_low, new_high)))
})?
Expand All @@ -258,13 +258,13 @@ impl TreeNode for Expr {
else_expr,
}) => transform_option_box(expr, &mut f)?
.map_data(|new_expr| (new_expr, when_then_expr, else_expr))
.and_then_transform_sibling(|(new_expr, when_then_expr, else_expr)| {
.and_then_transform(|(new_expr, when_then_expr, else_expr)| {
Ok(when_then_expr
.into_iter()
.map_till_continue_and_collect(|(when, then)| {
transform_box(when, &mut f)?
.map_data(|new_when| (new_when, then))
.and_then_transform_sibling(|(new_when, then)| {
.and_then_transform(|(new_when, then)| {
Ok(transform_box(then, &mut f)?
.map_data(|new_then| (new_when, new_then)))
})
Expand All @@ -273,13 +273,11 @@ impl TreeNode for Expr {
(new_expr, new_when_then_expr, else_expr)
}))
})?
.and_then_transform_sibling(
|(new_expr, new_when_then_expr, else_expr)| {
Ok(transform_option_box(else_expr, &mut f)?.map_data(
|new_else_expr| (new_expr, new_when_then_expr, new_else_expr),
))
},
)?
.and_then_transform(|(new_expr, new_when_then_expr, else_expr)| {
Ok(transform_option_box(else_expr, &mut f)?.map_data(
|new_else_expr| (new_expr, new_when_then_expr, new_else_expr),
))
})?
.map_data(|(new_expr, new_when_then_expr, new_else_expr)| {
Expr::Case(Case::new(new_expr, new_when_then_expr, new_else_expr))
}),
Expand Down Expand Up @@ -314,12 +312,12 @@ impl TreeNode for Expr {
window_frame,
}) => transform_vec(args, &mut f)?
.map_data(|new_args| (new_args, partition_by, order_by))
.and_then_transform_sibling(|(new_args, partition_by, order_by)| {
.and_then_transform(|(new_args, partition_by, order_by)| {
Ok(transform_vec(partition_by, &mut f)?.map_data(
|new_partition_by| (new_args, new_partition_by, order_by),
))
})?
.and_then_transform_sibling(|(new_args, new_partition_by, order_by)| {
.and_then_transform(|(new_args, new_partition_by, order_by)| {
Ok(transform_vec(order_by, &mut f)?.map_data(|new_order_by| {
(new_args, new_partition_by, new_order_by)
}))
Expand All @@ -341,11 +339,11 @@ impl TreeNode for Expr {
order_by,
}) => transform_vec(args, &mut f)?
.map_data(|new_args| (new_args, filter, order_by))
.and_then_transform_sibling(|(new_args, filter, order_by)| {
.and_then_transform(|(new_args, filter, order_by)| {
Ok(transform_option_box(filter, &mut f)?
.map_data(|new_filter| (new_args, new_filter, order_by)))
})?
.and_then_transform_sibling(|(new_args, new_filter, order_by)| {
.and_then_transform(|(new_args, new_filter, order_by)| {
Ok(transform_option_vec(order_by, &mut f)?
.map_data(|new_order_by| (new_args, new_filter, new_order_by)))
})?
Expand Down Expand Up @@ -390,7 +388,7 @@ impl TreeNode for Expr {
negated,
}) => transform_box(expr, &mut f)?
.map_data(|new_expr| (new_expr, list))
.and_then_transform_sibling(|(new_expr, list)| {
.and_then_transform(|(new_expr, list)| {
Ok(transform_vec(list, &mut f)?
.map_data(|new_list| (new_expr, new_list)))
})?
Expand Down

0 comments on commit 84d91c6

Please sign in to comment.