Add ExecutionPlan::apply_expressions()

[LiaCastaneda]

Which issue does this PR close?

• Closes Access DynamicFilterPhysicalExpr expressions from outside the plan #18296

Needed for datafusion-contrib/datafusion-distributed#180

Rationale for this change

Right now, there is no easy way to know if a given node in the plan holds Dynamic Filters or to traverse all physical expressions in an ExecutionPlan. This PR implements apply_expressions() that visits all PhysicalExprs inside an ExecutionPlan using a callback pattern, including DynamicFilterPhysicalExpr. This is similar to the existing apply_expressions() API for LogicalPlan.

What changes are included in this PR?

• Added apply_expressions() method to the ExecutionPlan trait with no default implementation, forcing all implementors to explicitly handle their expressions
• Uses a visitor pattern with FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion> to avoid allocations
• Implemented apply_expressions() for all ExecutionPlan implementations
• Also added apply_expressions() to FileSource and DataSource traits (required, no default)

Are these changes tested?

Yes, added a test that traverses the plan and discovers dynamic filters using apply_expressions().

Are there any user-facing changes?

Yes, the new API ExecutionPlan::apply_expressions(), FileSource::apply_expressions(), and DataSource::apply_expressions().

[adriangb]

This makes sense to me. It mirrors the APIs for Logical expressions, is clean and a relatively small change.

But since this is an API change let's leave this open for a couple of days and get at least 1 more approval from a committer before moving forward with it.

[adriangb]

Should this expr.apply() for nested expressions? Should it deduplicate Arc'ed copies?

[LiaCastaneda]

Should this expr.apply() for nested expressions?

iiuc the LogicalPlan counterpart returns just the top level expressions.

Should it deduplicate Arc'ed copies?

yeah deduping is a good idea

[adriangb]

I was referring to this helper function, not the general API. The general API should only expose top level expressions and do no deduplication.

[LiaCastaneda]

about deduping, the objective of this test is to prove how many times the Dynamic Filter appears in the plan and if each node is able count how many dynamic filters it contains, if we dedup then we would count it once only

[alamb]

This API forces an allocation and also cloning all the PhysicalExprs -- what would you think about adding apply_expressions and map_expressions methods to parallel the ones on LogicalPlan instead?

• https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.LogicalPlan.html#method.apply_expressions
• https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.LogicalPlan.html#method.map_expressions

Maybe you can start with just the apply_expressions one in this PR

I think we should probably also not provide a default implementation to force all implementations to properly visit the expressions

If we provide this default implementation, then downstream implementors will likely not implement the API and if something in the datafusion core depends on the API in the future it will be hard to debug what is going on

[LiaCastaneda]

I think we should probably also not provide a default implementation to force all implementations to properly visit the expressions
If we provide this default implementation, then downstream implementors will likely not implement the API and if something in the datafusion core depends on the API in the future it will be hard to debug what is going on

makes sense, I included a default implementation because didn't want to incroduce a breaking change but is better to be safe and force the implementation 👍

what would you think about adding apply_expressions and map_expressions methods to parallel the ones on LogicalPlan instead?

nice catch, I missed the allocation fact, I will give it a try

[alamb]

Thanks @LiaCastaneda and @adriangb

I am a little worried about the default implementation here --

I also think a slightly different API might be worth considering

[adriangb]

Thanks for reviewing Andrew - that's very good feedback that I missed in my review. I agree that apply_expressions(|expr: &Arc<dyn PhysicalExpr>| ...) would be a better API.

[LiaCastaneda]

Thanks both for the reviews! I will work on your suggestion @alamb

[adriangb]

Some minor comments but I think we can merge this whenever you think it's ready Lía

[LiaCastaneda]

There are some conflicts again, wil fix them...

[adriangb]

There are some conflicts again, wil fix them...

Thank you and sorry for the delays causing conflicts and bump to v54

[LiaCastaneda]

no worries, they were not too complex to solve, I added docs/source/library-user-guide/upgrading/54.0.0.md does it look ok? (it's the first time I add a file from scratch in DF)

if so, I think the PR is good to go

[askalt]

Hi! There is a patch #20009 that adds a more expressive API by splitting responsibilities into:

1. reading expressions
2. writing expressions

This approach not only helps to check for specific types of expressions in the plan but also enables replacing them, which extends the number of contexts where the API can be used. It looks a bit confusing to have all these methods together (apply_expressions, physical_expressions and with_physical_expressions), so with this, we can implement apply_expressions as a simple helper, like:

pub fn visit_expressions(
    plan: &dyn ExecutionPlan,
    f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
) -> Result<TreeNodeRecursion> {
    let mut tnr = TreeNodeRecursion::Continue;
    for expr in plan.physical_expressions() {
        tnr = tnr.visit_sibling(|| f(expr.as_ref()))?;
    }
    Ok(tnr)
}

[LiaCastaneda]

👋 Hey, I was not aware there was already an initiative to build a similar API. This PR implements apply_expressions, which mirrors LogicalPlan::apply_expressions and is intended to be read only and allocation free. Ideally, we should also implement map_expressions (mirroring LogicalPlan::map_expressions) to support modifying PhysicalExprs and rebuilding the node at the same time. Would both of these APIs cover your use case?

[askalt]

Would both of these APIs cover your use case?

Yes, it would be nice to have a writing API. The important property we need is that map_expressions should not recompute plan properties, assuming that they are not changed (user responsibility), i.e. we avoid a typical plan ::new() call in this case. Is there an issue or branch to track the implementation?

[LiaCastaneda]

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

[LiaCastaneda]

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? 🤔

[askalt]

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? 🤔

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<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>>,
    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<usize>,
	}

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]

lets continue this discussion in the issue