TreeNode Expr Implementation for LogicalPlan's

[berkaysynnada]

Is your feature request related to a problem or challenge?

The recent update in #8891 has made significant improvements, making many TreeNode implementations simpler and more straightforward. However, the map_children method in Expr enum appears to be quite complex. I think there's room for improvement in making this design more user-friendly, which would not only enhance user comprehension, but also help debugging.

Take, for instance, the Expr::Case pattern:

Expr::Case(Case {
    expr,
    when_then_expr,
    else_expr,
}) => transform_option_box(expr, &mut f)?
    .update_data(|new_expr| (new_expr, when_then_expr, else_expr))
    .try_transform_node(|(new_expr, when_then_expr, else_expr)| {
        Ok(when_then_expr
            .into_iter()
            .map_until_stop_and_collect(|(when, then)| {
                transform_box(when, &mut f)?
                    .update_data(|new_when| (new_when, then))
                    .try_transform_node(|(new_when, then)| {
                        Ok(transform_box(then, &mut f)?
                            .update_data(|new_then| (new_when, new_then)))
                    })
            })?
            .update_data(|new_when_then_expr| {
                (new_expr, new_when_then_expr, else_expr)
            }))
    })?
    .try_transform_node(|(new_expr, new_when_then_expr, else_expr)| {
        Ok(transform_option_box(else_expr, &mut f)?.update_data(
            |new_else_expr| (new_expr, new_when_then_expr, new_else_expr),
        ))
    })?
    .update_data(|(new_expr, new_when_then_expr, new_else_expr)| {
        Expr::Case(Case::new(new_expr, new_when_then_expr, new_else_expr))
    }),

As shown, the function employs multiple nested and chained utility functions, which can be quite daunting. Simplifying these would significantly enhance readability and understanding.

Describe the solution you'd like

New implementations of TreeNode for Option and Box may be implemented, or new simplifying utils may be written with a more understandable style.

Describe alternatives you've considered

No response

Additional context

No response

[berkaysynnada]

@peter-toth I'm curious about your thoughts and solution suggestions on this issue.

[peter-toth]

Thanks for pinging me @berkaysynnada.

I was thinking about using macros or helper functions instead of the current chaining Transfomeds (update_data() to create tuples + try_transform_node()s to map tuple elements). In those we still need to handle tnr == Stop from the result of the previous transformation, but we could avoid building tuples ((new_expr, when_then_expr, else_expr)) at least.
But this week is a bit hectic for me so I'm not sure I can open a PR this week. If you have any other idea or want to experiment with the above feel free to do it and ping me for review.

[peter-toth]

I was thinking of a macro that accepts a sequence of expressions and their transformations:

let t = transform_siblings!(
    expr, |e| transform_option_box(e, &mut f),
    when_then_expr, |wte| wte.into_iter().map_until_stop_and_collect(|(when, then)| {
        transform_siblings!(
            when, |w| transform_box(w, &mut f),
            then, |t| transform_box(t, &mut f),
        ),
    }),
    else_expr, |ee| transform_option_box(ee, &mut f)
)
t.update_data(|(new_expr, new_when_then_expr, new_else_expr)| {
    Expr::Case(Case::new(new_expr, new_when_then_expr, new_else_expr))
})

Inside the macro we can build up the same update_data + try_transform_node chain or maybe avoid that to make it easier to debug.