ARROW-2808: [Python] Add MemoryPool tests

Also expand the exposed API a bit.

python/pyarrow/__init__.py

@@ -96,9 +96,10 @@ from pyarrow.lib import (null, bool_,
 from pyarrow.lib import (Buffer, ResizableBuffer, foreign_buffer, py_buffer,
                          compress, decompress, allocate_buffer)
 
-from pyarrow.lib import (MemoryPool, ProxyMemoryPool, total_allocated_bytes,
-                         set_memory_pool, default_memory_pool,
-                         log_memory_allocations)
+from pyarrow.lib import (MemoryPool, LoggingMemoryPool, ProxyMemoryPool,
+                         total_allocated_bytes, set_memory_pool,
+                         default_memory_pool, logging_memory_pool,
+                         proxy_memory_pool, log_memory_allocations)
 
 from pyarrow.lib import (HdfsFile, NativeFile, PythonFile,
                          FixedSizeBufferWriter,

python/pyarrow/includes/libarrow.pxd

@@ -187,6 +187,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
     cdef cppclass CMemoryPool" arrow::MemoryPool":
         int64_t bytes_allocated()
+        int64_t max_memory()
 
     cdef cppclass CLoggingMemoryPool" arrow::LoggingMemoryPool"(CMemoryPool):
         CLoggingMemoryPool(CMemoryPool*)

python/pyarrow/memory.pxi

@@ -30,8 +30,23 @@ cdef class MemoryPool:
         self.pool = pool
 
     def bytes_allocated(self):
+        """
+        Return the number of bytes that are currently allocated from this
+        memory pool.
+        """
         return self.pool.bytes_allocated()
 
+    def max_memory(self):
+        """
+        Return the peak memory allocation in this memory pool.
+        This can be an approximate number in multi-threaded applications.
+
+        None is returned if the pool implementation doesn't know how to
+        compute this number.
+        """
+        ret = self.pool.max_memory()
+        return ret if ret >= 0 else None
+
 
 cdef CMemoryPool* maybe_unbox_memory_pool(MemoryPool memory_pool):
     if memory_pool is None:
@@ -76,8 +91,8 @@ def default_memory_pool():
 
 def proxy_memory_pool(MemoryPool parent):
     """
-    Derived MemoryPool class that tracks the number of bytes and
-    maximum memory allocated through its direct calls.
+    Create and return a MemoryPool instance that redirects to the
+    *parent*, but with separate allocation statistics.
     """
     cdef ProxyMemoryPool out = ProxyMemoryPool.__new__(ProxyMemoryPool)
     out.proxy_pool.reset(new CProxyMemoryPool(parent.pool))
@@ -85,13 +100,21 @@ def proxy_memory_pool(MemoryPool parent):
     return out
 
 
+def logging_memory_pool(MemoryPool parent):
+    """
+    Create and return a MemoryPool instance that redirects to the
+    *parent*, but also dumps allocation logs on stderr.
+    """
+    return LoggingMemoryPool.__new__(LoggingMemoryPool, parent)
+
+
 def set_memory_pool(MemoryPool pool):
     c_set_default_memory_pool(pool.pool)
 
 
 cdef MemoryPool _default_memory_pool = default_memory_pool()
-cdef LoggingMemoryPool _logging_memory_pool = LoggingMemoryPool.__new__(
-    LoggingMemoryPool, _default_memory_pool)
+cdef LoggingMemoryPool _logging_memory_pool = logging_memory_pool(
+    _default_memory_pool)
 
 
 def log_memory_allocations(enable=True):
@@ -110,5 +133,9 @@ def log_memory_allocations(enable=True):
 
 
 def total_allocated_bytes():
+    """
+    Return the currently allocated bytes from the default memory pool.
+    Other memory pools may not be accounted for.
+    """
     cdef CMemoryPool* pool = c_get_memory_pool()
     return pool.bytes_allocated()

python/pyarrow/tests/test_io.py

@@ -256,18 +256,6 @@ def test_python_file_closing():
     assert bio.closed
 
 
-# ----------------------------------------------------------------------
-# MemoryPool
-
-
-def test_memory_pool_cannot_use_ctor():
-    with pytest.raises(TypeError):
-        pa.MemoryPool()
-
-    with pytest.raises(TypeError):
-        pa.ProxyMemoryPool()
-
-
 # ----------------------------------------------------------------------
 # Buffers
 

python/pyarrow/tests/test_memory.py

@@ -0,0 +1,98 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import contextlib
+
+import pytest
+
+import pyarrow as pa
+
+
+@contextlib.contextmanager
+def allocate_bytes(pool, nbytes):
+    """
+    Temporarily allocate *nbytes* from the given *pool*.
+    """
+    arr = pa.array([b"x" * nbytes], type=pa.binary(), memory_pool=pool)
+    # Fetch the values buffer from the varbinary array and release the rest,
+    # to get the desired allocation amount
+    buf = arr.buffers()[2]
+    arr = None
+    assert len(buf) == nbytes
+    try:
+        yield
+    finally:
+        buf = None
+
+
+def check_allocated_bytes(pool):
+    """
+    Check allocation stats on *pool*.
+    """
+    allocated_before = pool.bytes_allocated()
+    max_mem_before = pool.max_memory()
+    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
+    assert pool.bytes_allocated() == allocated_before
+    assert pool.max_memory() == new_max_memory
+
+
+def test_cannot_use_constructors():
+    with pytest.raises(TypeError):
+        pa.MemoryPool()
+
+    with pytest.raises(TypeError):
+        pa.ProxyMemoryPool()
+
+    with pytest.raises(TypeError):
+        pa.LoggingMemoryPool()
+
+
+def test_default_allocated_bytes():
+    pool = pa.default_memory_pool()
+    with allocate_bytes(pool, 1024):
+        check_allocated_bytes(pool)
+        assert pool.bytes_allocated() == pa.total_allocated_bytes()
+
+
+def test_proxy_memory_pool():
+    pool = pa.proxy_memory_pool(pa.default_memory_pool())
+    check_allocated_bytes(pool)
+
+
+def test_logging_memory_pool(capfd):
+    pool = pa.logging_memory_pool(pa.default_memory_pool())
+    check_allocated_bytes(pool)
+    out, err = capfd.readouterr()
+    assert err == ""
+    assert out.count("Allocate:") > 0
+    assert out.count("Allocate:") == out.count("Free:")
+
+
+def test_set_memory_pool():
+    old_pool = pa.default_memory_pool()
+    pool = pa.proxy_memory_pool(old_pool)
+    pa.set_memory_pool(pool)
+    try:
+        allocated_before = pool.bytes_allocated()
+        with allocate_bytes(None, 512):
+            assert pool.bytes_allocated() == allocated_before + 512
+        assert pool.bytes_allocated() == allocated_before
+    finally:
+        pa.set_memory_pool(old_pool)