From 42f05c2766845dcfefdebbedd87ab9997e2f3868 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <pitrou@free.fr>
Date: Thu, 12 Dec 2024 10:56:58 +0100
Subject: [PATCH] GH-45006: [CI][Python] Fix test_memory failures

`test_memory.py` has started failing on some builds after https://github.com/apache/arrow/pull/44951 was merged
---
 ci/scripts/cpp_build.sh             |  4 +--
 python/pyarrow/tests/test_memory.py | 40 ++++++++++++-----------------
 2 files changed, 18 insertions(+), 26 deletions(-)

diff --git a/ci/scripts/cpp_build.sh b/ci/scripts/cpp_build.sh
index e70f5da85ae2e..c1e7adf6a05e0 100755
--- a/ci/scripts/cpp_build.sh
+++ b/ci/scripts/cpp_build.sh
@@ -171,10 +171,10 @@ else
     -DARROW_GCS=${ARROW_GCS:-OFF} \
     -DARROW_HDFS=${ARROW_HDFS:-ON} \
     -DARROW_INSTALL_NAME_RPATH=${ARROW_INSTALL_NAME_RPATH:-ON} \
-    -DARROW_JEMALLOC=${ARROW_JEMALLOC:-ON} \
+    -DARROW_JEMALLOC=${ARROW_JEMALLOC:-OFF} \
     -DARROW_JSON=${ARROW_JSON:-ON} \
     -DARROW_LARGE_MEMORY_TESTS=${ARROW_LARGE_MEMORY_TESTS:-OFF} \
-    -DARROW_MIMALLOC=${ARROW_MIMALLOC:-OFF} \
+    -DARROW_MIMALLOC=${ARROW_MIMALLOC:-ON} \
     -DARROW_ORC=${ARROW_ORC:-OFF} \
     -DARROW_PARQUET=${ARROW_PARQUET:-OFF} \
     -DARROW_RUNTIME_SIMD_LEVEL=${ARROW_RUNTIME_SIMD_LEVEL:-MAX} \
diff --git a/python/pyarrow/tests/test_memory.py b/python/pyarrow/tests/test_memory.py
index b1eef176665af..87594776bf88a 100644
--- a/python/pyarrow/tests/test_memory.py
+++ b/python/pyarrow/tests/test_memory.py
@@ -30,15 +30,19 @@
 pytestmark = pytest.mark.processes
 
 possible_backends = ["system", "jemalloc", "mimalloc"]
+# Backends which are expected to be present in all builds of PyArrow,
+# except if the user manually recompiled Arrow C++.
+mandatory_backends = ["system", "mimalloc"]
 
-should_have_jemalloc = (sys.platform == "linux" and platform.machine() == 'x86_64')
-should_have_mimalloc = sys.platform == "win32"
+
+def backend_factory(backend_name):
+    return getattr(pa, f"{backend_name}_memory_pool")
 
 
 def supported_factories():
     yield pa.default_memory_pool
-    for backend in pa.supported_memory_backends():
-        yield getattr(pa, f"{backend}_memory_pool")
+    for backend_name in pa.supported_memory_backends():
+        yield backend_factory(backend_name)
 
 
 @contextlib.contextmanager
@@ -149,17 +153,12 @@ def check_env_var(name, expected, *, expect_warning=False):
 
 
 def test_env_var():
-    check_env_var("system", ["system"])
-    if should_have_jemalloc:
-        check_env_var("jemalloc", ["jemalloc"])
-    if should_have_mimalloc:
-        check_env_var("mimalloc", ["mimalloc"])
+    for backend_name in mandatory_backends:
+        check_env_var(backend_name, [backend_name])
     check_env_var("nonexistent", possible_backends, expect_warning=True)
 
 
-def test_specific_memory_pools():
-    specific_pools = set()
-
+def test_memory_pool_factories():
     def check(factory, name, *, can_fail=False):
         if can_fail:
             try:
@@ -169,23 +168,16 @@ def check(factory, name, *, can_fail=False):
         else:
             pool = factory()
         assert pool.backend_name == name
-        specific_pools.add(pool)
 
-    check(pa.system_memory_pool, "system")
-    check(pa.jemalloc_memory_pool, "jemalloc",
-          can_fail=not should_have_jemalloc)
-    check(pa.mimalloc_memory_pool, "mimalloc",
-          can_fail=not should_have_mimalloc)
+    for backend_name in possible_backends:
+        check(backend_factory(backend_name), backend_name,
+              can_fail=backend_name not in mandatory_backends)
 
 
 def test_supported_memory_backends():
     backends = pa.supported_memory_backends()
-
-    assert "system" in backends
-    if should_have_jemalloc:
-        assert "jemalloc" in backends
-    if should_have_mimalloc:
-        assert "mimalloc" in backends
+    assert set(backends) >= set(mandatory_backends)
+    assert set(backends) <= set(possible_backends)
 
 
 def run_debug_memory_pool(pool_factory, env_value):