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

cpp/src/arrow/memory_pool.cc

@@ -201,4 +201,45 @@ int64_t LoggingMemoryPool::max_memory() const {
   std::cout << "max_memory: " << mem << std::endl;
   return mem;
 }
+
+ProxyMemoryPool::ProxyMemoryPool(MemoryPool* pool) : pool_(pool), proxy_bytes_allocated_(0) {}
+
+Status ProxyMemoryPool::Allocate(int64_t size, uint8_t** out) {
+  proxy_bytes_allocated_ += size;
+  Status s = pool_->Allocate(size, out);
+  std::cout << "Allocate: size = " << size << std::endl;
+  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);
+  proxy_bytes_allocated_ += new_size - old_size;
+  std::cout << "Reallocate: old_size = " << old_size << " - new_size = " << new_size
+            << " - proxy_allocated - " << proxy_bytes_allocated_ << std::endl;
+  return s;
+}
+
+void ProxyMemoryPool::Free(uint8_t* buffer, int64_t size) {
+  pool_->Free(buffer, size);
+  proxy_bytes_allocated_ -= size;
+  std::cout << "Free: size = " << size << std::endl;
+}
+
+int64_t ProxyMemoryPool::bytes_allocated() const {
+  int64_t nb_bytes = pool_->bytes_allocated();
+  std::cout << "bytes_allocated: " << nb_bytes << std::endl;
+  return nb_bytes;
+}
+
+int64_t ProxyMemoryPool::proxy_bytes_allocated() const {
+  int64_t nb_bytes = proxy_bytes_allocated_;
+  std::cout << "bytes_allocated: " << nb_bytes << std::endl;
+  return nb_bytes;
+}
+
+int64_t ProxyMemoryPool::max_memory() const {
+  int64_t mem = pool_->max_memory();
+  std::cout << "max_memory: " << mem << std::endl;
+  return mem;
+}
 }  // namespace arrow

cpp/src/arrow/memory_pool.h

@@ -86,6 +86,27 @@ 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;
+
+  int64_t proxy_bytes_allocated_;
+  int64_t proxy_bytes_allocated() const;
+
+ private:
+  MemoryPool* pool_;
+};
+
 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 proxy_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,18 @@ 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 proxy_bytes_allocated(self):
+        return self.proxy_pool.get().proxy_bytes_allocated()
+
+
 def default_memory_pool():
     cdef:
         MemoryPool pool = MemoryPool()
@@ -58,6 +70,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):