ARROW-530: [C++/Python] Provide subpools for better memory allocation …

cpp/src/arrow/memory_pool-test.cc

@@ -91,4 +91,25 @@ TEST(LoggingMemoryPool, Logging) {
 
   ASSERT_EQ(200, pool->max_memory());
 }
+
+TEST(ProxyMemoryPool, Logging) {
+  MemoryPool* pool = default_memory_pool();
+
+  ProxyMemoryPool pp = ProxyMemoryPool(pool);
+
+  uint8_t* data;
+  ASSERT_OK(pool->Allocate(100, &data));
+
+  uint8_t* data2;
+  ASSERT_OK(pp.Allocate(300, &data2));
+
+  ASSERT_EQ(400, pool->bytes_allocated());
+  ASSERT_EQ(300, pp.bytes_allocated());
+
+  pool->Free(data, 100);
+  pp.Free(data2, 300);
+
+  ASSERT_EQ(0, pool->bytes_allocated());
+  ASSERT_EQ(0, pp.bytes_allocated());
+}
 }  // namespace arrow

cpp/src/arrow/memory_pool.cc

@@ -201,4 +201,33 @@ int64_t LoggingMemoryPool::max_memory() const {
   std::cout << "max_memory: " << mem << std::endl;
   return mem;
 }
+
+ProxyMemoryPool::ProxyMemoryPool(MemoryPool* pool) : pool_(pool) {}
+
+Status ProxyMemoryPool::Allocate(int64_t size, uint8_t** out) {
+  Status s = pool_->Allocate(size, out);
+  bytes_allocated_ += size;
+  return s;
+}
+
+Status ProxyMemoryPool::Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) {
+  Status s = pool_->Reallocate(old_size, new_size, ptr);
+  bytes_allocated_ += new_size - old_size;
+  return s;
+}
+
+void ProxyMemoryPool::Free(uint8_t* buffer, int64_t size) {
+  pool_->Free(buffer, size);
+  bytes_allocated_ -= size;
+}
+
+int64_t ProxyMemoryPool::bytes_allocated() const {
+  int64_t nb_bytes = bytes_allocated_;
+  return nb_bytes;
+}
+
+int64_t ProxyMemoryPool::max_memory() const {
+  int64_t mem = pool_->max_memory();
+  return mem;
+}
 }  // namespace arrow

cpp/src/arrow/memory_pool.h

@@ -86,6 +86,25 @@ class ARROW_EXPORT LoggingMemoryPool : public MemoryPool {
   MemoryPool* pool_;
 };
 
+class ARROW_EXPORT ProxyMemoryPool : public MemoryPool {
+ public:
+  explicit ProxyMemoryPool(MemoryPool* pool);
+  ~ProxyMemoryPool() override = default;
+
+  Status Allocate(int64_t size, uint8_t** out) override;
+  Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override;
+
+  void Free(uint8_t* buffer, int64_t size) override;
+
+  int64_t bytes_allocated() const override;
+
+  int64_t max_memory() const override;
+
+ private:
+  MemoryPool* pool_;
+  int64_t bytes_allocated_ = 0;
+};
+
 ARROW_EXPORT MemoryPool* default_memory_pool();
 
 #ifdef ARROW_NO_DEFAULT_MEMORY_POOL

python/pyarrow/__init__.py

@@ -90,7 +90,7 @@ def parse_version(root):
 from pyarrow.lib import (Buffer, ResizableBuffer, foreign_buffer, py_buffer,
                          compress, decompress, allocate_buffer)
 
-from pyarrow.lib import (MemoryPool, total_allocated_bytes,
+from pyarrow.lib import (MemoryPool, ProxyMemoryPool, total_allocated_bytes,
                          set_memory_pool, default_memory_pool,
                          log_memory_allocations)
 

python/pyarrow/includes/libarrow.pxd

@@ -192,6 +192,10 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
     cdef cppclass CLoggingMemoryPool" arrow::LoggingMemoryPool"(CMemoryPool):
         CLoggingMemoryPool(CMemoryPool*)
 
+    cdef cppclass CProxyMemoryPool" arrow::ProxyMemoryPool"(CMemoryPool):
+        CProxyMemoryPool(CMemoryPool*)
+        int64_t bytes_allocated()
+
     cdef cppclass CBuffer" arrow::Buffer":
         CBuffer(const uint8_t* data, int64_t size)
         const uint8_t* data()

python/pyarrow/memory.pxi

@@ -44,6 +44,15 @@ cdef class LoggingMemoryPool(MemoryPool):
         self.init(self.logging_pool.get())
 
 
+cdef class ProxyMemoryPool(MemoryPool):
+    cdef:
+        unique_ptr[CProxyMemoryPool] proxy_pool
+
+    def __cinit__(self, MemoryPool pool):
+        self.proxy_pool.reset(new CProxyMemoryPool(pool.pool))
+        self.init(self.proxy_pool.get())
+
+
 def default_memory_pool():
     cdef:
         MemoryPool pool = MemoryPool()
@@ -58,6 +67,8 @@ def set_memory_pool(MemoryPool pool):
 cdef MemoryPool _default_memory_pool = default_memory_pool()
 cdef LoggingMemoryPool _logging_memory_pool = (
     LoggingMemoryPool(_default_memory_pool))
+cdef ProxyMemoryPool _proxy_memory_pool = (
+    ProxyMemoryPool(_default_memory_pool))
 
 
 def log_memory_allocations(enable=True):