feat: use mimalloc as the Python wheel's global allocator - #2545
Draft
james-willis wants to merge 1 commit into
Draft
feat: use mimalloc as the Python wheel's global allocator#2545james-willis wants to merge 1 commit into
james-willis wants to merge 1 commit into
Conversation
Gold Data ReportNotes
Commit Information
Summary
DetailsGold Data Metrics
|
james-willis
marked this pull request as ready for review
September 3, 2026 23:45
sail-cli sets a mimalloc #[global_allocator], but that static lives in its bin target (main.rs) and is not linked into the sail-python cdylib. So the pysail wheel -- the artifact most users actually run -- falls back to the system allocator (glibc), whose per-arena free lists retain freed memory across queries in a long-lived server (see lakehq#1916). Set the mimalloc static at the cdylib root so the wheel gets the same allocator as the CLI. Pin mimalloc to v2: the crate currently defaults to v3.3.2, which segfaults under multi-threaded aligned allocation (microsoft/mimalloc#1287) and SIGABRTs on macOS when statically linked into a second shared library (microsoft/mimalloc#1301). local_dynamic_tls selects the local-dynamic TLS model so the static loads cleanly from the extension. Carves the mimalloc change out of the stalled lakehq#2235 so it can land on its own.
james-willis
force-pushed
the
jw/mimalloc-python-wheel
branch
from
September 3, 2026 23:47
1628901 to
424ec47
Compare
Spark 3.5.9 Test ReportCommit Information
Test Summary
Test DetailsError CountsPassed Tests Diff(empty) Failed Tests |
Spark 4.2.0 Test ReportCommit Information
Test Summary
Test DetailsError CountsPassed Tests Diff(empty) Failed Tests(truncated) |
Ibis Test ReportCommit Information
Test Summary
Test DetailsError CountsPassed Tests Diff--- before.txt 2026-09-03 23:56:23.176230347 +0000
+++ after.txt 2026-09-03 23:56:23.390234642 +0000
@@ -19,2 +18,0 @@
-ibis/backends/pyspark/tests/test_client.py::test_create_table_no_catalog
-ibis/backends/pyspark/tests/test_client.py::test_insert_bug
@@ -28,4 +25,0 @@
-ibis/backends/pyspark/tests/test_ddl.py::test_drop_database_exists
-ibis/backends/pyspark/tests/test_ddl.py::test_drop_non_empty_database
-ibis/backends/pyspark/tests/test_ddl.py::test_drop_table_not_exist
-ibis/backends/pyspark/tests/test_ddl.py::test_drop_view
@@ -358,2 +351,0 @@
-ibis/backends/tests/test_client.py::test_create_database[pyspark-None]
-ibis/backends/tests/test_client.py::test_create_database[pyspark-current_catalog]
@@ -370,6 +361,0 @@
-ibis/backends/tests/test_client.py::test_create_table_in_memory[pyspark-memtable pandas]
-ibis/backends/tests/test_client.py::test_create_table_in_memory[pyspark-memtable pyarrow]
-ibis/backends/tests/test_client.py::test_create_table_in_memory[pyspark-memtable_list]
-ibis/backends/tests/test_client.py::test_create_table_in_memory[pyspark-pandas]
-ibis/backends/tests/test_client.py::test_create_table_in_memory[pyspark-pyarrow table]
-ibis/backends/tests/test_client.py::test_create_table_in_memory[pyspark-pyarrow_single_batch]
@@ -377 +362,0 @@
-ibis/backends/tests/test_client.py::test_cross_database_join[pyspark]
@@ -401,4 +385,0 @@
-ibis/backends/tests/test_client.py::test_insert_using_col_name_not_position[pyspark-auto generated cols]
-ibis/backends/tests/test_client.py::test_insert_using_col_name_not_position[pyspark-column order matching]
-ibis/backends/tests/test_client.py::test_insert_using_col_name_not_position[pyspark-column order reversed]
-ibis/backends/tests/test_client.py::test_insert_with_database_specified[pyspark]
@@ -417 +397,0 @@
-ibis/backends/tests/test_client.py::test_no_accidental_cross_database_table_load[pyspark]
@@ -533 +512,0 @@
-ibis/backends/tests/test_export.py::test_table_to_csv[pyspark]
@@ -1366 +1344,0 @@
-ibis/backends/tests/test_struct.py::test_keyword_fields[pyspark-nullable]Failed Tests |
Codecov Report✅ All modified and coverable lines are covered by tests. @@ Coverage Diff @@
## main #2545 +/- ##
==========================================
- Coverage 78.77% 78.76% -0.01%
==========================================
Files 988 988
Lines 203785 203785
==========================================
- Hits 160522 160521 -1
- Misses 43263 43264 +1
*This pull request uses carry forward flags. Click here to find out more.
... and 1 file with indirect coverage changes 🚀 New features to boost your workflow:
|
Contributor
|
@james-willis We've been blocked by purpleprotocol/mimalloc_rust#169 At this point, I think we should maybe just fork the repo so that we don't have to downgrade to v2 |
Contributor
|
Thanks for raising this issue! I'm going to move this PR to draft as we need to decide what's the best approach here given that the mimalloc binding is not actively maintained. |
linhr
marked this pull request as draft
September 4, 2026 06:26
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
This PR carves the Python-wheel allocator change out of #2235 (@shehabgamin) so it can land on its own, independently of the perf/ClickBench work still in flight there. The approach here — including the mimalloc
v2pin — follows that PR.The problem
Sail sets a mimalloc
#[global_allocator]insail-cli, but a#[global_allocator]only takes effect in the crate that roots the final artifact. That static lives insail-cli's binary target (main.rs), which isn't linked into thesail-pythoncdylib. So thepysailwheel silently falls back to the system allocator (glibc).glibc keeps freed chunks on per-arena free lists and doesn't return them to the OS without an explicit
malloc_trim. In a long-lived Spark Connect server, that reads as memory that never comes back between queries — matching the report in #1916 ("spark.stop() is not releasing memory to os… taking more memory than apache spark").On workloads that run several memory-intensive queries back to back in one long-lived server, that retention can accumulate until the process is OOM-killed, even though no single query's peak comes close to the limit. The work around is making explicit
malloc_trimcalls between queries which most users wont be savvy enough for.SedonaDB vs Sail
SedonaDB is the same shape — a Rust engine shipped as a pyo3 cdylib wheel — and it sets its mimalloc
#[global_allocator]in both roots (its CLI andpython/sedonadb/src/lib.rs). Because Sail set it only in the CLI, two otherwise-similar engines behaved very differently on the same long-running workload: SedonaDB's memory plateaus across queries; Sail's climbs. This PR closes that gap by giving the wheel the allocator the CLI already has.The change
#[global_allocator]to thesail-pythoncdylib root (+ themimallocfeature, enabled via the maturin build).local_dynamic_tlsselects the local-dynamic TLS model so the static loads cleanly from the extension module. Once alibmimalloc-sys > 0.1.49vendors the 1287 and 1301 fixes, thev2pin can be dropped.