python: jemalloc (#3108)

* python: jemalloc

* keep mimalloc on windows

polars/src/lib.rs

@@ -261,8 +261,8 @@
 //!
 //! ### Custom allocator
 //! A DataFrame library naturally does a lot of heap allocations. It is recommended to use a custom
-//! allocator. [Mimalloc](https://docs.rs/mimalloc/0.1.25/mimalloc/) for instance, shows a significant
-//! performance gain in runtime as well as memory usage. Especially with `MIMALLOC_LARGE_OS_PAGES=1`.
+//! allocator. [Mimalloc](https://docs.rs/mimalloc/0.1.25/mimalloc/) and [JeMalloc](https://crates.io/crates/jemallocator) for instance, show a significant
+//! performance gain in runtime as well as memory usage.
 //!
 //! #### Usage
 //! ```ignore

py-polars/.flake8

@@ -3,5 +3,5 @@
 max-line-length = 180
 # E203, W503: due to black fmt
 ignore = E203,W503
-exclude = legacy, docs, venv
+exclude = legacy, docs, venv, target
 

py-polars/Cargo.toml

@@ -15,12 +15,16 @@ description = "Blazingly fast DataFrame library"
 [workspace]
 # prevents package from thinking it's in the workspace
 
+[target.'cfg(windows)'.dependencies]
+mimalloc = { version = "*", default-features = false }
+
+[target.'cfg(not(windows))'.dependencies]
+jemallocator = { version = "0.3", features = ["disable_initial_exec_tls"] }
+
 [dependencies]
 ahash = "0.7"
 bincode = "1.3"
 libc = "0.2"
-libmimalloc-sys = { version = "*", features = ["extended"] }
-mimalloc = { version = "*", default-features = false }
 ndarray = "0.15"
 numpy = "0.16"
 polars-core = { path = "../polars/polars-core", default-features = false }

py-polars/docs/source/reference/config.rst

@@ -14,7 +14,5 @@ Config
     Config.set_tbl_cols
     Config.set_global_string_cache
     Config.unset_global_string_cache
-    Config.set_large_os_pages
-    Config.unset_large_os_pages
     toggle_string_cache
     StringCache

py-polars/polars/__init__.py

@@ -3,16 +3,13 @@
 
 try:
     from polars.polars import version
-
-    _DOCUMENTING = False
 except ImportError as e:  # pragma: no cover
 
     def version() -> str:
         return ""
 
     # this is only useful for documentation
     warnings.warn("polars binary missing!")
-    _DOCUMENTING = True
 
 import polars.testing as testing
 from polars.cfg import (  # flake8: noqa. We do not export in __all__
@@ -233,6 +230,3 @@ def version() -> str:
 import os
 
 os.environ["POLARS_ALLOW_EXTENSION"] = "true"
-
-if not _DOCUMENTING:
-    Config.set_large_os_pages()

py-polars/polars/cfg.py

@@ -3,13 +3,6 @@
 
 from polars.string_cache import toggle_string_cache
 
-try:
-    from polars.polars import enable_large_os_pages as _enable_large_os_pages
-
-    _DOCUMENTING = False
-except ImportError:  # pragma: no cover
-    _DOCUMENTING = True
-
 
 class Config:
     """
@@ -106,23 +99,3 @@ def unset_global_string_cache(cls) -> "Type[Config]":
         """
         toggle_string_cache(False)
         return cls
-
-    @classmethod
-    def unset_large_os_pages(cls) -> "Type[Config]":
-        """
-        Turn off large OS pages in `mimalloc` allocator.
-
-        This can be slower depending on your use case. It is on by default in polars.
-        """
-        _enable_large_os_pages(False)
-        return cls
-
-    @classmethod
-    def set_large_os_pages(cls) -> "Type[Config]":
-        """
-        Turn on large OS pages in `mimalloc` allocator.
-
-        It is on by default in polars.
-        """
-        _enable_large_os_pages(True)
-        return cls

py-polars/src/lib.rs

@@ -39,6 +39,9 @@ use crate::error::{
 use crate::file::get_either_file;
 use crate::prelude::{ClosedWindow, DataType, DatetimeArgs, Duration, DurationArgs, PyDataType};
 use dsl::ToExprs;
+#[cfg(not(target_os = "windows"))]
+use jemallocator::Jemalloc;
+#[cfg(target_os = "windows")]
 use mimalloc::MiMalloc;
 use polars::functions::{diag_concat_df, hor_concat_df};
 use polars::prelude::Null;
@@ -48,16 +51,12 @@ use polars_core::prelude::IntoSeries;
 use pyo3::types::{PyBool, PyDict, PyFloat, PyInt, PyString};
 
 #[global_allocator]
-static GLOBAL: MiMalloc = MiMalloc;
+#[cfg(not(target_os = "windows"))]
+static ALLOC: Jemalloc = Jemalloc;
 
-#[pyfunction]
-fn enable_large_os_pages(_py: Python, toggle: bool) {
-    // Safety
-    // holding the python gil makes this thread safe
-    unsafe {
-        libmimalloc_sys::mi_option_set_enabled(libmimalloc_sys::mi_option_large_os_pages, toggle)
-    }
-}
+#[global_allocator]
+#[cfg(target_os = "windows")]
+static ALLOC: MiMalloc = MiMalloc;
 
 #[pyfunction]
 fn col(name: &str) -> dsl::PyExpr {
@@ -477,7 +476,5 @@ fn polars(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(max_exprs)).unwrap();
     m.add_wrapped(wrap_pyfunction!(as_struct)).unwrap();
     m.add_wrapped(wrap_pyfunction!(repeat)).unwrap();
-    m.add_wrapped(wrap_pyfunction!(enable_large_os_pages))
-        .unwrap();
     Ok(())
 }