Implement ExecutionPlan::map_expressions

[LiaCastaneda]

Is your feature request related to a problem or challenge?

ExecutionPlan::apply_expressions (#20337) added a read-only visitor API for physical expressions. To complete the API that mirrors LogicalPlan::map_expressions — we should also implement a write counterpart that allows transforming expressions and reconstructing the node.

Use cases:

• Part of Prepared physical plan reusage #14342
• Any transformation that needs to replace expressions across a plan tree

Describe the solution you'd like

Proposed API:

fn map_expressions<F>(
    self: Arc<Self>,
    f: &mut F,
) -> Result<Transformed<Arc<dyn ExecutionPlan>>>
where
    F: FnMut(Arc<dyn PhysicalExpr>) -> Result<Transformed<Arc<dyn PhysicalExpr>>>;

Each ExecutionPlan implementor is responsible for transforming its expressions and reconstructing itself using direct struct construction, so that the PlanProperties (schema, ordering, partitioning) are reused without recomputation -- matching the behavior of LogicalPlan::map_expressions.

Describe alternatives you've considered

No response

Additional context

No response

[LiaCastaneda]

Bringing some discussions from the apply_expressions PR:

I think we can reuse the properties of the rest of the plan (avoiding ::new()), similar to how LogicalPlan::map_expressions does it.

I created this issue #20899. I haven't started working on it yet and probably won't have much time this week, so I'll likely give it a try next week, but feel free to take it if you'd like'

Actually, now that I think about it, there are some cases where we would need to recompute properties right? for example, if a user changes an expression from a > something to a < something. How do we specify in this API whether we want to recompute properties or not? should map_expressions have a recompute_properties: bool argument? 🤔

From @askalt

Yes, it may be useful to explicitly ask for properties re-computation. And it seems for me that by default the safest option is to force properties to be re-computed.

Another way to satisfy it is to introduce "args struct" like:

struct MapExpressionsArgs<'a> {
f: &'a dyn FnMut(&Arc) -> Result<Arc>,
preserve_properties: bool,
}
Like is done here:

datafusion/datafusion/catalog/src/table.rs

Lines 366 to 372 in 8d9b080

/// Arguments for scanning a table with [TableProvider::scan_with_args].
#[derive(Debug, Clone, Default)]
pub struct ScanArgs<'a> {
filters: Option<&'a [Expr]>,
projection: Option<&'a [usize]>,
limit: Option,
}
to not add a bool argument each time when the method semantics is extended. But maybe this is overkill here and bool parameter will be enough.

[LiaCastaneda]

I think the only thing that feels weird to me is that implementors will receive preserve_properties and will have to do something different based on its value... 🤔

if preserve_properties {
    recreate node using new::
} else {
    recreate node not using new::
}

which adds boilerplate to every ExecutionPlan implementation since this method would be required.

I see other execution plan nodes have separate functions for explicitly not recomputing properties, for example with_new_children_and_same_properties

[askalt]

Implementing two methods with the similar semantic seems hard to maintain. What if we inject recompute properties behavior to the closure return type with a lightweight structure?

pub enum RecomputePropertiesBehavior {
    Preserve,
    Recompute,
}

pub struct MappedExpr {
    pub properties_behavior: RecomputePropertiesBehavior,
    pub expr: Arc<dyn PhysicalExpr>,
}

pub trait ExecutionPlan {
   ...

    fn map_expressions(
        self: Arc<Self>,
        f: &dyn FnMut(Arc<dyn PhysicalExpr>) -> Result<Transformed<MappedExpr>>,
    ) -> Result<Transformed<Arc<dyn ExecutionPlan>>>;
}

And then if at least one returned expression has RecomputePropertiesBehavior::Recompute -- the implementation will recompute them and leave as is otherwise.