feat(rust,python): Add a config option to specify the default engine to use during lazyframe collect calls

[Matt711]

Closes #19797.

[ion-elgreco]

Wouldn't it be more practical if you can set any engine then as default?

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 79.27%. Comparing base (ebb513c) to head (2922103).
Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #20717   +/-   ##
=======================================
  Coverage   79.27%   79.27%           
=======================================
  Files        1585     1586    +1     
  Lines      226132   226169   +37     
  Branches     2592     2594    +2     
=======================================
+ Hits       179258   179288   +30     
- Misses      46279    46285    +6     
- Partials      595      596    +1     

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

[orlp]

I think this isn't as simple as it seems, because we do also call Polars methods in the Python code internally which are only supported on, for example, the eager engine.

[wence-]

I think this isn't as simple as it seems, because we do also call Polars methods in the Python code internally which are only supported on, for example, the eager engine.

Can you expand on what you mean here? Effectively what this is trying to do is add a global config option for the default value in LazyFrame.collect(..., engine="cpu") (to make that literal "cpu" string optionally "gpu" instead). The Python-side logic that then perhaps nonetheless turns off the gpu engine still exists.

Why would hooking this up via a config option not work?

[orlp]

@wence- I have no idea what the full scope of such examples are throughout the codebase, but one example is pl.align_frames dispatching to other calls. Such a function might suddenly stop working if you change the default engine, because not all engines have the exact same supported feature set.

[wence-]

OK, thanks, I think I see. I think there's some amount of tension here between making this "easy" to configure globally and I guess some things breaking at a distance from where the user might expect.

[Matt711]

Ok tests are passing now. Gentle ping for reviews from @ritchie46 @orlp et al

[ritchie46]

I think this is useful, but I think we should do it a bit differently.

We don't have to check the flag in python, but rather, check in rust in the main optimizer loop. I also think we should call it `"engine affinity" because we don't guarantee anything.

[ritchie46]

I don't think this should dispatch here, but rather in fn optimize.

I also think we should do all engines at once. Where we call it "engine_affinity" as we cannot guarantee it will on the preferred engine.

[wence-]

OK, so that would be moving the current python-based engine-selection logic into optimize? Note how below we need to know here whether or not we have a "gpu" engine to deliver the appropriate post_opt_callback, I think.

[Matt711]

I opened up #21102 to track dispatching in the main optimizer loop. We can move the discussion over there. I don't think it should block this PR though.

[ritchie46]

Hmm, right this only works for streaming/in-memory engine. For the post-opt-callback this indeed needs to be in python. I will take a better look to better understand where we can have that logic.