feat(python): Add post-optimization callback #15972
Conversation
ritchie46
requested review from
stinodego,
c-peters,
alexander-beedie,
MarcoGorelli,
reswqa and
orlp
as code owners
github-actions
bot
added
enhancement
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #15972 +/- ##
==========================================
+ Coverage 80.89% 80.95% +0.06%
==========================================
Files 1384 1384
Lines 178020 178106 +86
Branches 3043 3043
==========================================
+ Hits 144003 144183 +180
+ Misses 33535 33440 -95
- Partials 482 483 +1
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
| // Pass the node visitor which allows the python callback to replace parts of the query plan. | ||
| // Remove "cuda" or specify better once we have multiple post-opt callbacks. | ||
| lambda.call1(py, (nt,)).map_err( | ||
| |e| polars_err!(ComputeError: "'cuda' conversion failed: {}", e), | ||
| )?; |
There was a problem hiding this comment.
OK so the contract here is that the callback should modify the NodeTraverser via set_udf to substitute in the subtree it can execute. But that callback isn't allowed to hold on to the NodeTraverser, since when the callback actually runs the arenas inside it are garbage.
There was a problem hiding this comment.
Indeed. Any state that needs to be preserved should be in the callback itself.
This allows a post-optimization python callback to be inserted.