build(python): Fix allocator features (#16365)

pola-rs · May 21, 2024 · 63f0fcd · 63f0fcd
1 parent 30ba423
commit 63f0fcd
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 21 deletions.
diff --git a/.github/workflows/release-python.yml b/.github/workflows/release-python.yml
@@ -144,7 +144,7 @@ jobs:
         if: matrix.architecture == 'x86-64'
         env:
           FEATURES: ${{ steps.features.outputs.features }}
-          CFG: ${{ matrix.package == 'polars-lts-cpu' && '--cfg default_allocator' || '' }}
+          CFG: ${{ matrix.package == 'polars-lts-cpu' && '--cfg allocator="default"' || '' }}
         run: echo "RUSTFLAGS=-C target-feature=${{ steps.features.outputs.features }} $CFG" >> $GITHUB_ENV
 
       - name: Set variables in CPU check module

diff --git a/py-polars/Cargo.toml b/py-polars/Cargo.toml
@@ -98,16 +98,13 @@ features = [
 [build-dependencies]
 built = { version = "0.7", features = ["chrono", "git2", "cargo-lock"], optional = true }
 
-[target.'cfg(all(any(not(target_family = "unix"), use_mimalloc), not(default_allocator)))'.dependencies]
+[target.'cfg(all(any(not(target_family = "unix"), allocator = "mimalloc"), not(allocator = "default")))'.dependencies]
 mimalloc = { version = "0.1", default-features = false }
 
-[target.'cfg(all(target_family = "unix", not(use_mimalloc), not(default_allocator)))'.dependencies]
+[target.'cfg(all(target_family = "unix", not(allocator = "mimalloc"), not(allocator = "default")))'.dependencies]
 jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] }
 
 [features]
-default_allocator = []
-use_mimalloc = []
-
 # Features below are only there to enable building a slim binary during development.
 avro = ["polars/avro"]
 parquet = ["polars/parquet", "polars-parquet"]

diff --git a/py-polars/build.rs b/py-polars/build.rs
@@ -1,6 +1,8 @@
 /// Build script using 'built' crate to generate build info.
 
 fn main() {
+    println!("cargo::rustc-check-cfg=cfg(allocator, values(\"default\", \"mimalloc\"))");
+
     #[cfg(feature = "build_info")]
     {
         println!("cargo:rerun-if-changed=build.rs");

diff --git a/py-polars/src/allocator.rs b/py-polars/src/allocator.rs
@@ -1,38 +1,38 @@
 #[cfg(all(
     target_family = "unix",
-    not(feature = "default_allocator"),
-    not(feature = "use_mimalloc"),
+    not(allocator = "default"),
+    not(allocator = "mimalloc"),
 ))]
 use jemallocator::Jemalloc;
 #[cfg(all(
     not(debug_assertions),
-    not(feature = "default_allocator"),
-    any(not(target_family = "unix"), feature = "use_mimalloc"),
+    not(allocator = "default"),
+    any(not(target_family = "unix"), allocator = "mimalloc"),
 ))]
 use mimalloc::MiMalloc;
 
 #[cfg(all(
     debug_assertions,
     target_family = "unix",
-    not(feature = "default_allocator"),
-    not(feature = "use_mimalloc"),
+    not(allocator = "default"),
+    not(allocator = "mimalloc"),
 ))]
 use crate::memory::TracemallocAllocator;
 
 #[global_allocator]
 #[cfg(all(
     not(debug_assertions),
-    not(feature = "use_mimalloc"),
-    not(feature = "default_allocator"),
+    not(allocator = "mimalloc"),
+    not(allocator = "default"),
     target_family = "unix",
 ))]
 static ALLOC: Jemalloc = Jemalloc;
 
 #[global_allocator]
 #[cfg(all(
     not(debug_assertions),
-    not(feature = "default_allocator"),
-    any(not(target_family = "unix"), feature = "use_mimalloc"),
+    not(allocator = "default"),
+    any(not(target_family = "unix"), allocator = "mimalloc"),
 ))]
 static ALLOC: MiMalloc = MiMalloc;
 
@@ -41,9 +41,5 @@ static ALLOC: MiMalloc = MiMalloc;
 // linking breaks on Windows if we use tracemalloc C APIs. So we only use this
 // on Unix for now.
 #[global_allocator]
-#[cfg(all(
-    debug_assertions,
-    not(feature = "default_allocator"),
-    target_family = "unix",
-))]
+#[cfg(all(debug_assertions, not(allocator = "default"), target_family = "unix",))]
 static ALLOC: TracemallocAllocator<Jemalloc> = TracemallocAllocator::new(Jemalloc);