Skip to content

Commit

Permalink
apacheGH-45006: [CI][Python] Fix test_memory failures
Browse files Browse the repository at this point in the history
`test_memory.py` has started failing on some builds after apache#44951 was merged
  • Loading branch information
pitrou committed Dec 12, 2024
1 parent 9016a83 commit 42f05c2
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 26 deletions.
4 changes: 2 additions & 2 deletions ci/scripts/cpp_build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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} \
Expand Down
40 changes: 16 additions & 24 deletions python/pyarrow/tests/test_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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):
Expand Down

0 comments on commit 42f05c2

Please sign in to comment.