feat(python): Expose allocator to capsule

[ruihe774]

Currently, pyo3 extensions and plugins are statically linked to their own copy of allocator (jemalloc/mimalloc). This does not only bloat the code size, but also increase memory usage and fragmentation: arenas of allocators are not shared, and dirty pages may not be purged after use. (purging only happens on new allocation; when an extension/plugin finishes its work, it does not allocate anymore, so the dirty pages allocated by its allocator cannot be purged.)

This PR adds the allocator capsule (polars.polars._allocator), which exposes the allocator used by python polars. pyo3 extensions and plugins can import this capsule and relay allocations to it. All allocations can now be managed by one single jemalloc/mimalloc.

[ruihe774]

The corresponding PR to pyo3-polars: pola-rs/pyo3-polars#88

[codecov]

Codecov Report

Attention: Patch coverage is 50.00000% with 16 lines in your changes missing coverage. Please review.

Project coverage is 80.49%. Comparing base (9eec688) to head (ae0c33c).

Files	Patch %	Lines
py-polars/src/allocator.rs	46.66%	16 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #17817   +/-   ##
=======================================
  Coverage   80.49%   80.49%           
=======================================
  Files        1503     1504    +1     
  Lines      197027   197059   +32     
  Branches     2795     2795           
=======================================
+ Hits       158589   158616   +27     
- Misses      37918    37923    +5     
  Partials      520      520           

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

[orlp]

I've thought about pros and cons and I think overall this is worth doing. Neat idea! The main downside I can think of is that each allocation call in plugins will involve an extra (properly predicted) branch and call, but a big pro to me is that in the future we can add (with a feature flag) extra memory usage tracking in our global allocator and also track the memory usage of plugins.

[ruihe774]

The main downside I can think of is that each allocation call in plugins will involve an extra (properly predicted) branch and call

Yes, it involves a branch and an indirect call. However, imo the overhead is trivial and is not measurable in benchmarks.

we can add (with a feature flag) extra memory usage tracking in our global allocator and also track the memory usage of plugins.

Yes, it is a selling point.

[s-banach]

I'm learning to write plugins from @MarcoGorelli's cookie cutter repo. Other than deleting the jemalloc stuff from Cargo.toml, will I have to make any other changes?

[ruihe774]

I'm learning to write plugins from @MarcoGorelli's cookie cutter repo. Other than deleting the jemalloc stuff from Cargo.toml, will I have to make any other changes?

I've updated the examples in pyo3-polars: pola-rs/pyo3-polars@b5fce50. Please have a look.