chore: use custom global alloc

Signed-off-by: ion-elgreco <[email protected]>

Author: ion-elgreco

python/.cargo/config.toml

@@ -8,4 +8,8 @@ rustflags = [
 rustflags = [
   "-C", "link-arg=-undefined",
   "-C", "link-arg=dynamic_lookup",
-]
+]
+
+# Custom jemalloc conf, ref https://github.com/jemalloc/jemalloc/issues/2688
+[env]
+JEMALLOC_SYS_WITH_MALLOC_CONF = "dirty_decay_ms:500,muzzy_decay_ms:-1"

python/Cargo.toml

@@ -47,6 +47,18 @@ reqwest = { version = "*", features = ["native-tls-vendored"] }
 
 deltalake-mount = { path = "../crates/mount" }
 
+# Non-unix or emscripten os
+[target.'cfg(any(not(target_family = "unix"), target_os = "emscripten"))'.dependencies]
+mimalloc = { version = "0.1", default-features = false }
+
+# Unix (excluding macOS & emscripten) → jemalloc 
+[target.'cfg(all(target_family = "unix", not(target_os = "macos"), not(target_os = "emscripten")))'.dependencies]
+jemallocator = { version = "0.5", features = ["disable_initial_exec_tls", "background_threads"] }
+
+# macOS → jemalloc (without background_threads) (https://github.com/jemalloc/jemalloc/issues/843)
+[target.'cfg(all(target_family = "unix", target_os = "macos"))'.dependencies]
+jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] }
+
 [dependencies.pyo3]
 version = "0.22.6"
 features = ["extension-module", "abi3", "abi3-py39", "gil-refs"]

python/src/lib.rs

@@ -79,6 +79,20 @@ use crate::query::PyQueryBuilder;
 use crate::schema::{schema_to_pyobject, Field};
 use crate::utils::rt;
 
+#[cfg(all(target_family = "unix", not(target_os = "emscripten")))]
+use jemallocator::Jemalloc;
+
+#[cfg(all(any(not(target_family = "unix"), target_os = "emscripten")))]
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+#[cfg(all(target_family = "unix", not(target_os = "emscripten")))]
+static ALLOC: Jemalloc = Jemalloc;
+
+#[global_allocator]
+#[cfg(all(any(not(target_family = "unix"), target_os = "emscripten")))]
+static ALLOC: MiMalloc = MiMalloc;
+
 #[derive(FromPyObject)]
 enum PartitionFilterValue {
     Single(PyBackedStr),