fix: Fix join_where incorrectly dropping transformations on RHS of equality expressions

[nameexhaustion]

• Fixes Incorrect results from join_where - transformations on RHS of equality expressions are dropped #21066

This restores the original code to before #20865 with an added clone step, and serves as a proper fix for the original issue at:

• Selecting columns from LazyFrame after join_where fails with DuplicateError #20831

I had previously misunderstood the cause of the original issue, but I believe I now have the correct understanding of what happened. Here is an illustration to show what was happening in the original issue:

Essentially, the bug is an in-place mutation of a Node in the expr_arena by collapse_joins, where the same Node was also being referenced by a projection in another place.

The previous PR (#20865) resolved the above issue by always inserting a new ColumnNode. However, it made the incorrect assumption that the RHS of the equality was always a col(), and would drop any extra transformations being applied on the RHS (e.g. col(A_right) + 1 gets incorrectly turned into col(A_right)).

The proper fix in this PR is to restore the original logic, with an added deep-clone before the mutation takes place to ensure the underlying Node being mutated isn't being referenced by e.g. a Select {} IR somewhere else.

Arena::<T>::get_mut()

In light of this issue, a quick search across the codebase shows collapse_joins as the only place that uses Arena::get_mut() on expressions arenas:

All other uses are generally with logical plan (LP / IR) arenas. I haven't checked in detail whether they are used properly in those places as we don't seem to have any bugs relating to logical plans being incorrectly mutated.

[nameexhaustion]

moved to test_joins.py below

[codecov]

Codecov Report

Attention: Patch coverage is 93.33333% with 3 lines in your changes missing coverage. Please review.

Project coverage is 79.26%. Comparing base (7f4dc50) to head (82cd89e).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
.../polars-plan/src/plans/optimizer/collapse_joins.rs	92.50%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #21067      +/-   ##
==========================================
- Coverage   79.27%   79.26%   -0.02%     
==========================================
  Files        1586     1586              
  Lines      225627   225669      +42     
  Branches     2588     2588              
==========================================
+ Hits       178877   178885       +8     
- Misses      46160    46194      +34     
  Partials      590      590              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ritchie46]

I think this should debug_assert that the current output_name also is lhs.

[nameexhaustion]

added

[ritchie46]

Instead of deep cloning can't we ensure we don't mutate? E.g. by adding a new node instead. The remove_suffix looks very suspect to me and I think we should just create a new node expression there.

I like to prevent going back into expr as there might be information lost and we do a lot of heap allocations. If we do that all over the optimizer it would get pretty expensive (and incorrect because of the information loss).

[nameexhaustion]

Currently, I'm reasoning that any node within an expression tree may be referenced at another location, so if any node in the tree is mutated, then all of the parent nodes must be cloned to a unique expression tree.

I made this assumption very conservative as I don't think we have node-sharing behavior of an Arena formally defined in the codebase 😶. E.g. if we define that only AExpr::Column nodes can be shared in multiple locations, then it's safe to just mutate via replacing the column node within an expression without cloning the entire expression tree.

But for the cloning itself I can work on adding a way to do it without converting to Expr 👍

[ritchie46]

We have a visitor that rewrites the Aexpr and then assigns that whole new rewritten to new slots in the arena's. Maybe we can use that instead of mutating?