alamb commented on code in PR #9876:
URL: https://github.com/apache/arrow-datafusion/pull/9876#discussion_r1548510420
##########
datafusion/common/src/tree_node.rs:
##########
@@ -532,7 +532,7 @@ impl<T> Transformed<T> {
}
}
-/// Transformation helper to process tree nodes that are siblings.
+/// Transformation helper to process sequence of iterable tree nodes that are
siblings.
Review Comment:
I found the reason / logic for continue / jump quite subtle when trying to
make TreeNodeMutator. I have some suggested comments below to explain the
rationale that might help
##########
datafusion/common/src/tree_node.rs:
##########
@@ -532,7 +532,7 @@ impl<T> Transformed<T> {
}
}
-/// Transformation helper to process tree nodes that are siblings.
+/// Transformation helper to process sequence of iterable tree nodes that are
siblings.
pub trait TransformedIterator: Iterator {
fn map_until_stop_and_collect<
Review Comment:
I think explaining the intended behavior of this method would have helped me
previously. Something like this perhaps (similarly for the next function too)
```suggestion
/// Apples `f` to each item in this iterator
///
/// Visits all items in the iterator unless
/// `f` returns an error or `f` returns TreeNodeRecursion::stop.
///
/// # Returns
/// Error if f returns an error
///
/// Ok(Transformed) such that:
/// 1. `transformed` is true if any return from `f` had transformed true
/// 2. `data` from the last invocation of `f`
/// 3. `tnr` from the last invocation of `f` or `Continue` if the
iterator is empty
fn map_until_stop_and_collect<
```
##########
datafusion/expr/src/tree_node/expr.rs:
##########
@@ -320,48 +309,39 @@ impl TreeNode for Expr {
order_by,
window_frame,
null_treatment,
- }) => transform_vec(args, &mut f)?
- .update_data(|new_args| (new_args, partition_by, order_by))
- .try_transform_node(|(new_args, partition_by, order_by)| {
- Ok(transform_vec(partition_by, &mut f)?.update_data(
- |new_partition_by| (new_args, new_partition_by,
order_by),
- ))
- })?
- .try_transform_node(|(new_args, new_partition_by, order_by)| {
- Ok(
- transform_vec(order_by, &mut
f)?.update_data(|new_order_by| {
- (new_args, new_partition_by, new_order_by)
- }),
- )
- })?
- .update_data(|(new_args, new_partition_by, new_order_by)| {
- Expr::WindowFunction(WindowFunction::new(
- fun,
- new_args,
- new_partition_by,
- new_order_by,
- window_frame,
- null_treatment,
- ))
- }),
+ }) => map_until_stop_and_collect!(
Review Comment:
this is very clever
##########
datafusion/common/src/tree_node.rs:
##########
@@ -551,22 +551,48 @@ impl<I: Iterator> TransformedIterator for I {
) -> Result<Transformed<Vec<Self::Item>>> {
let mut tnr = TreeNodeRecursion::Continue;
let mut transformed = false;
- let data = self
- .map(|item| match tnr {
- TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
- f(item).map(|result| {
- tnr = result.tnr;
- transformed |= result.transformed;
- result.data
- })
- }
- TreeNodeRecursion::Stop => Ok(item),
- })
- .collect::<Result<Vec<_>>>()?;
- Ok(Transformed::new(data, transformed, tnr))
+ self.map(|item| match tnr {
+ TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
+ f(item).map(|result| {
+ tnr = result.tnr;
+ transformed |= result.transformed;
+ result.data
+ })
+ }
+ TreeNodeRecursion::Stop => Ok(item),
+ })
+ .collect::<Result<Vec<_>>>()
+ .map(|data| Transformed::new(data, transformed, tnr))
}
}
+/// Transformation helper to process sequence of tree node containing
expressions.
+/// This macro is very similar to
[TransformedIterator::map_until_stop_and_collect] to
+/// process nodes that are siblings, but it accepts an initial transformation
and a
+/// sequence of pairs of an expression and its transformation.
+#[macro_export]
+macro_rules! map_until_stop_and_collect {
+ ($TRANSFORMED_EXPR_0:expr, $($EXPR:expr, $TRANSFORMED_EXPR:expr),*) => {{
+ $TRANSFORMED_EXPR_0.and_then(|Transformed { data: data0, mut
transformed, mut tnr }| {
+ let data = (
Review Comment:
Maybe we could name this something other than `data` so it was clearer what
type it is (maybe `all_datas`?)
##########
datafusion/common/src/tree_node.rs:
##########
@@ -551,22 +551,48 @@ impl<I: Iterator> TransformedIterator for I {
) -> Result<Transformed<Vec<Self::Item>>> {
let mut tnr = TreeNodeRecursion::Continue;
let mut transformed = false;
- let data = self
- .map(|item| match tnr {
- TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
- f(item).map(|result| {
- tnr = result.tnr;
- transformed |= result.transformed;
- result.data
- })
- }
- TreeNodeRecursion::Stop => Ok(item),
- })
- .collect::<Result<Vec<_>>>()?;
- Ok(Transformed::new(data, transformed, tnr))
+ self.map(|item| match tnr {
+ TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
+ f(item).map(|result| {
+ tnr = result.tnr;
+ transformed |= result.transformed;
+ result.data
+ })
+ }
+ TreeNodeRecursion::Stop => Ok(item),
+ })
+ .collect::<Result<Vec<_>>>()
+ .map(|data| Transformed::new(data, transformed, tnr))
}
}
+/// Transformation helper to process sequence of tree node containing
expressions.
+/// This macro is very similar to
[TransformedIterator::map_until_stop_and_collect] to
+/// process nodes that are siblings, but it accepts an initial transformation
and a
+/// sequence of pairs of an expression and its transformation.
Review Comment:
This is quite clever. I have some suggestions on naming that would have
helped me:
Use `F` to mirror the nomenclature of `TransformedIterator`. I thought
`$TRANSFORMED_EXPR` was an actually expr (but it is actually a closure that is
invoked I think 🤔 )
```
$TRANSFORMED_EXPR_0 --> F0
$TRANSFORMED_EXPR --> F.
```
I think it would help to document each parameter and the return value
(specifically it looks like it returns `Transformed<(data, ..., data)>`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
