diff --git a/cpp/src/arrow/memory_pool.cc b/cpp/src/arrow/memory_pool.cc index 3ace2c8f23ab0..42ca539204ffd 100644 --- a/cpp/src/arrow/memory_pool.cc +++ b/cpp/src/arrow/memory_pool.cc @@ -266,6 +266,8 @@ class DebugAllocator { } } + static void PrintStats() { WrappedAllocator::PrintStats(); } + private: static Result RawSize(int64_t size) { if (ARROW_PREDICT_FALSE(internal::AddWithOverflow(size, kOverhead, &size))) { @@ -378,6 +380,12 @@ class SystemAllocator { // The return value of malloc_trim is not an error but to inform // you if memory was actually released or not, which we do not care about here ARROW_UNUSED(malloc_trim(0)); +#endif + } + + static void PrintStats() { +#ifdef __GLIBC__ + malloc_stats(); #endif } }; @@ -430,6 +438,8 @@ class MimallocAllocator { mi_free(ptr); } } + + static void PrintStats() { mi_stats_print_out(nullptr, nullptr); } }; #endif // defined(ARROW_MIMALLOC) @@ -512,6 +522,8 @@ class BaseMemoryPoolImpl : public MemoryPool { void ReleaseUnused() override { Allocator::ReleaseUnused(); } + void PrintStats() override { Allocator::PrintStats(); } + int64_t bytes_allocated() const override { return stats_.bytes_allocated(); } int64_t max_memory() const override { return stats_.max_memory(); } @@ -724,6 +736,10 @@ void LoggingMemoryPool::Free(uint8_t* buffer, int64_t size, int64_t alignment) { std::cout << "Free: size = " << size << ", alignment = " << alignment << std::endl; } +void LoggingMemoryPool::ReleaseUnused() { pool_->ReleaseUnused(); } + +void LoggingMemoryPool::PrintStats() { pool_->PrintStats(); } + int64_t LoggingMemoryPool::bytes_allocated() const { int64_t nb_bytes = pool_->bytes_allocated(); std::cout << "bytes_allocated: " << nb_bytes << std::endl; @@ -775,6 +791,14 @@ class ProxyMemoryPool::ProxyMemoryPoolImpl { stats_.DidFreeBytes(size); } + void ReleaseUnused() { pool_->ReleaseUnused(); } + + void PrintStats() { + // XXX these are the allocation stats for the underlying allocator, not + // the subset allocated through the ProxyMemoryPool + pool_->PrintStats(); + } + int64_t bytes_allocated() const { return stats_.bytes_allocated(); } int64_t max_memory() const { return stats_.max_memory(); } @@ -809,6 +833,10 @@ void ProxyMemoryPool::Free(uint8_t* buffer, int64_t size, int64_t alignment) { return impl_->Free(buffer, size, alignment); } +void ProxyMemoryPool::ReleaseUnused() { impl_->ReleaseUnused(); } + +void ProxyMemoryPool::PrintStats() { impl_->PrintStats(); } + int64_t ProxyMemoryPool::bytes_allocated() const { return impl_->bytes_allocated(); } int64_t ProxyMemoryPool::max_memory() const { return impl_->max_memory(); } diff --git a/cpp/src/arrow/memory_pool.h b/cpp/src/arrow/memory_pool.h index 98c6dc3e211b8..19a938c33601b 100644 --- a/cpp/src/arrow/memory_pool.h +++ b/cpp/src/arrow/memory_pool.h @@ -151,6 +151,12 @@ class ARROW_EXPORT MemoryPool { /// unable to fulfill the request due to fragmentation. virtual void ReleaseUnused() {} + /// Print statistics + /// + /// Print allocation statistics on stderr. The output format is + /// implementation-specific. Not all memory pools implement this method. + virtual void PrintStats() {} + /// The number of bytes that were allocated and not yet free'd through /// this allocator. virtual int64_t bytes_allocated() const = 0; @@ -187,6 +193,8 @@ class ARROW_EXPORT LoggingMemoryPool : public MemoryPool { Status Reallocate(int64_t old_size, int64_t new_size, int64_t alignment, uint8_t** ptr) override; void Free(uint8_t* buffer, int64_t size, int64_t alignment) override; + void ReleaseUnused() override; + void PrintStats() override; int64_t bytes_allocated() const override; @@ -219,6 +227,8 @@ class ARROW_EXPORT ProxyMemoryPool : public MemoryPool { Status Reallocate(int64_t old_size, int64_t new_size, int64_t alignment, uint8_t** ptr) override; void Free(uint8_t* buffer, int64_t size, int64_t alignment) override; + void ReleaseUnused() override; + void PrintStats() override; int64_t bytes_allocated() const override; diff --git a/cpp/src/arrow/memory_pool_internal.h b/cpp/src/arrow/memory_pool_internal.h index 01500b3c1eae1..dc90d5680b665 100644 --- a/cpp/src/arrow/memory_pool_internal.h +++ b/cpp/src/arrow/memory_pool_internal.h @@ -44,6 +44,7 @@ class JemallocAllocator { uint8_t** ptr); static void DeallocateAligned(uint8_t* ptr, int64_t size, int64_t alignment); static void ReleaseUnused(); + static void PrintStats(); }; #endif // defined(ARROW_JEMALLOC) diff --git a/cpp/src/arrow/memory_pool_jemalloc.cc b/cpp/src/arrow/memory_pool_jemalloc.cc index 239d83b81bc67..1342fe13f7eed 100644 --- a/cpp/src/arrow/memory_pool_jemalloc.cc +++ b/cpp/src/arrow/memory_pool_jemalloc.cc @@ -131,6 +131,10 @@ void JemallocAllocator::ReleaseUnused() { mallctl("arena." ARROW_STRINGIFY(MALLCTL_ARENAS_ALL) ".purge", NULL, NULL, NULL, 0); } +void JemallocAllocator::PrintStats() { + malloc_stats_print(nullptr, nullptr, /*opts=*/""); +} + } // namespace internal } // namespace memory_pool diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index b2edeb0b4192f..021c1c782c6e5 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -323,8 +323,11 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil: cdef cppclass CMemoryPool" arrow::MemoryPool": int64_t bytes_allocated() int64_t max_memory() + int64_t total_bytes_allocated() + int64_t num_allocations() c_string backend_name() void ReleaseUnused() + void PrintStats() cdef cppclass CLoggingMemoryPool" arrow::LoggingMemoryPool"(CMemoryPool): CLoggingMemoryPool(CMemoryPool*) diff --git a/python/pyarrow/memory.pxi b/python/pyarrow/memory.pxi index 1ddcb01ccb6ab..fdd5b991a8859 100644 --- a/python/pyarrow/memory.pxi +++ b/python/pyarrow/memory.pxi @@ -58,6 +58,13 @@ cdef class MemoryPool(_Weakrefable): """ return self.pool.bytes_allocated() + def total_bytes_allocated(self): + """ + Return the total number of bytes that have been allocated from this + memory pool. + """ + return self.pool.total_bytes_allocated() + def max_memory(self): """ Return the peak memory allocation in this memory pool. @@ -69,6 +76,23 @@ cdef class MemoryPool(_Weakrefable): ret = self.pool.max_memory() return ret if ret >= 0 else None + def num_allocations(self): + """ + Return the number of allocations or reallocations that were made + using this memory pool. + """ + return self.pool.num_allocations() + + def print_stats(self): + """ + Print statistics about this memory pool. + + The output format is implementation-specific. Not all memory pools + implement this method. + """ + with nogil: + self.pool.PrintStats() + @property def backend_name(self): """ @@ -83,6 +107,7 @@ cdef class MemoryPool(_Weakrefable): f"bytes_allocated={self.bytes_allocated()} " f"max_memory={self.max_memory()}>") + cdef CMemoryPool* maybe_unbox_memory_pool(MemoryPool memory_pool): if memory_pool is None: return c_get_memory_pool() diff --git a/python/pyarrow/tests/test_memory.py b/python/pyarrow/tests/test_memory.py index 6ed999db42cee..28f7cd51df71e 100644 --- a/python/pyarrow/tests/test_memory.py +++ b/python/pyarrow/tests/test_memory.py @@ -67,12 +67,17 @@ def check_allocated_bytes(pool): """ allocated_before = pool.bytes_allocated() max_mem_before = pool.max_memory() + num_allocations_before = pool.num_allocations() with allocate_bytes(pool, 512): assert pool.bytes_allocated() == allocated_before + 512 new_max_memory = pool.max_memory() assert pool.max_memory() >= max_mem_before + num_allocations_after = pool.num_allocations() + assert num_allocations_after > num_allocations_before + assert num_allocations_after < num_allocations_before + 5 assert pool.bytes_allocated() == allocated_before assert pool.max_memory() == new_max_memory + assert pool.num_allocations() == num_allocations_after def test_default_allocated_bytes(): @@ -271,3 +276,20 @@ def test_debug_memory_pool_unknown(pool_factory): "Valid values are 'abort', 'trap', 'warn', 'none'." ) check_debug_memory_pool_disabled(pool_factory, env_value, msg) + + +@pytest.mark.parametrize('pool_factory', supported_factories()) +def test_print_stats(pool_factory): + code = f"""if 1: + import pyarrow as pa + + pool = pa.{pool_factory.__name__}() + buf = pa.allocate_buffer(64, memory_pool=pool) + pool.print_stats() + """ + res = subprocess.run([sys.executable, "-c", code], check=True, + universal_newlines=True, stdout=subprocess.PIPE, + stderr=subprocess.PIPE) + if sys.platform == "linux": + # On Linux at least, all memory pools should emit statistics + assert res.stderr.strip() != ""