ARROW-457: Python: Better control over memory pool

python/CMakeLists.txt

@@ -412,6 +412,7 @@ set(CYTHON_EXTENSIONS
   config
   error
   io
+  memory
   scalar
   schema
   table

python/pyarrow/__init__.py

@@ -30,7 +30,6 @@ from pyarrow.config import cpu_count, set_cpu_count
 
 from pyarrow.array import (Array,
                            from_pandas_series, from_pylist,
-                           total_allocated_bytes,
                            NumericArray, IntegerArray, FloatingPointArray,
                            BooleanArray,
                            Int8Array, UInt8Array,
@@ -48,6 +47,8 @@ from pyarrow.io import (HdfsFile, NativeFile, PythonFileInterface,
 
 from pyarrow.ipc import FileReader, FileWriter, StreamReader, StreamWriter
 
+from pyarrow.memory import MemoryPool, total_allocated_bytes
+
 from pyarrow.scalar import (ArrayValue, Scalar, NA, NAType,
                             BooleanValue,
                             Int8Value, Int16Value, Int32Value, Int64Value,

python/pyarrow/_parquet.pxd

@@ -19,7 +19,7 @@
 
 from pyarrow.includes.common cimport *
 from pyarrow.includes.libarrow cimport (CArray, CSchema, CStatus,
-                                        CTable, MemoryPool)
+                                        CTable, CMemoryPool)
 from pyarrow.includes.libarrow_io cimport ReadableFileInterface, OutputStream
 
 
@@ -204,13 +204,13 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil:
 
 cdef extern from "parquet/arrow/reader.h" namespace "parquet::arrow" nogil:
     CStatus OpenFile(const shared_ptr[ReadableFileInterface]& file,
-                     MemoryPool* allocator,
+                     CMemoryPool* allocator,
                      const ReaderProperties& properties,
                      const shared_ptr[CFileMetaData]& metadata,
                      unique_ptr[FileReader]* reader)
 
     cdef cppclass FileReader:
-        FileReader(MemoryPool* pool, unique_ptr[ParquetFileReader] reader)
+        FileReader(CMemoryPool* pool, unique_ptr[ParquetFileReader] reader)
         CStatus ReadColumn(int i, shared_ptr[CArray]* out);
         CStatus ReadTable(shared_ptr[CTable]* out);
         CStatus ReadTable(const vector[int]& column_indices,
@@ -229,7 +229,7 @@ cdef extern from "parquet/arrow/schema.h" namespace "parquet::arrow" nogil:
 
 cdef extern from "parquet/arrow/writer.h" namespace "parquet::arrow" nogil:
     cdef CStatus WriteTable(
-        const CTable* table, MemoryPool* pool,
+        const CTable* table, CMemoryPool* pool,
         const shared_ptr[OutputStream]& sink,
         int64_t chunk_size,
         const shared_ptr[WriterProperties]& properties)

python/pyarrow/_parquet.pyx

@@ -32,6 +32,7 @@ from pyarrow.compat import tobytes, frombytes
 from pyarrow.error import ArrowException
 from pyarrow.error cimport check_status
 from pyarrow.io import NativeFile
+from pyarrow.memory cimport MemoryPool, maybe_unbox_memory_pool
 from pyarrow.table cimport Table
 
 from pyarrow.io cimport NativeFile, get_reader, get_writer
@@ -342,13 +343,13 @@ cdef logical_type_name_from_enum(ParquetLogicalType type_):
 cdef class ParquetReader:
     cdef:
         object source
-        MemoryPool* allocator
+        CMemoryPool* allocator
         unique_ptr[FileReader] reader
         column_idx_map
         FileMetaData _metadata
 
-    def __cinit__(self):
-        self.allocator = default_memory_pool()
+    def __cinit__(self, MemoryPool memory_pool=None):
+        self.allocator = maybe_unbox_memory_pool(memory_pool)
         self._metadata = None
 
     def open(self, object source, FileMetaData metadata=None):
@@ -471,6 +472,7 @@ cdef class ParquetWriter:
     cdef:
         shared_ptr[WriterProperties] properties
         shared_ptr[OutputStream] sink
+        CMemoryPool* allocator
 
     cdef readonly:
         object use_dictionary
@@ -479,7 +481,7 @@ cdef class ParquetWriter:
         int row_group_size
 
     def __cinit__(self, where, use_dictionary=None, compression=None,
-                  version=None):
+                  version=None, MemoryPool memory_pool=None):
         cdef shared_ptr[FileOutputStream] filestream
 
         if isinstance(where, six.string_types):
@@ -487,6 +489,7 @@ cdef class ParquetWriter:
             self.sink = <shared_ptr[OutputStream]> filestream
         else:
             get_writer(where, &self.sink)
+        self.allocator = maybe_unbox_memory_pool(memory_pool)
 
         self.use_dictionary = use_dictionary
         self.compression = compression
@@ -540,6 +543,6 @@ cdef class ParquetWriter:
 
         cdef int c_row_group_size = row_group_size
         with nogil:
-            check_status(WriteTable(ctable, default_memory_pool(),
+            check_status(WriteTable(ctable, self.allocator,
                                         self.sink, c_row_group_size,
                                         self.properties))

python/pyarrow/array.pyx

@@ -29,6 +29,7 @@ import pyarrow.config
 
 from pyarrow.compat import frombytes, tobytes
 from pyarrow.error cimport check_status
+from pyarrow.memory cimport MemoryPool, maybe_unbox_memory_pool
 
 cimport pyarrow.scalar as scalar
 from pyarrow.scalar import NA
@@ -44,11 +45,6 @@ cdef _pandas():
     return pd
 
 
-def total_allocated_bytes():
-    cdef MemoryPool* pool = pyarrow.get_memory_pool()
-    return pool.bytes_allocated()
-
-
 cdef class Array:
 
     cdef init(self, const shared_ptr[CArray]& sp_array):
@@ -58,7 +54,7 @@ cdef class Array:
         self.type.init(self.sp_array.get().type())
 
     @staticmethod
-    def from_pandas(obj, mask=None, timestamps_to_ms=False, Field field=None):
+    def from_pandas(obj, mask=None, timestamps_to_ms=False, Field field=None, MemoryPool memory_pool=None):
         """
         Convert pandas.Series to an Arrow Array.
 
@@ -74,6 +70,9 @@ cdef class Array:
             compatibility with other functionality like Parquet I/O which
             only supports milliseconds.
 
+        memory_pool: MemoryPool, optional
+            Specific memory pool to use to allocate the resulting Arrow array.
+
         Notes
         -----
         Localized timestamps will currently be returned as UTC (pandas's native representation).
@@ -107,6 +106,7 @@ cdef class Array:
         cdef:
             shared_ptr[CArray] out
             shared_ptr[CField] c_field
+            CMemoryPool* pool
 
         pd = _pandas()
 
@@ -121,20 +121,20 @@ cdef class Array:
         if isinstance(series_values, pd.Categorical):
             return DictionaryArray.from_arrays(series_values.codes,
                                                series_values.categories.values,
-                                               mask=mask)
+                                               mask=mask, memory_pool=memory_pool)
         else:
             if series_values.dtype.type == np.datetime64 and timestamps_to_ms:
                 series_values = series_values.astype('datetime64[ms]')
 
+            pool = maybe_unbox_memory_pool(memory_pool)
             with nogil:
                 check_status(pyarrow.PandasToArrow(
-                    pyarrow.get_memory_pool(), series_values, mask,
-                    c_field, &out))
+                    pool, series_values, mask, c_field, &out))
 
             return box_arrow_array(out)
 
     @staticmethod
-    def from_list(object list_obj, DataType type=None):
+    def from_list(object list_obj, DataType type=None, MemoryPool memory_pool=None):
         """
         Convert Python list to Arrow array
 
@@ -147,10 +147,12 @@ cdef class Array:
         pyarrow.array.Array
         """
         cdef:
-            shared_ptr[CArray] sp_array
+           shared_ptr[CArray] sp_array
+           CMemoryPool* pool
 
+        pool = maybe_unbox_memory_pool(memory_pool)
         if type is None:
-            check_status(pyarrow.ConvertPySequence(list_obj, &sp_array))
+            check_status(pyarrow.ConvertPySequence(list_obj, pool, &sp_array))
         else:
             raise NotImplementedError()
 
@@ -330,7 +332,7 @@ cdef class BinaryArray(Array):
 cdef class DictionaryArray(Array):
 
     @staticmethod
-    def from_arrays(indices, dictionary, mask=None):
+    def from_arrays(indices, dictionary, mask=None, MemoryPool memory_pool=None):
         """
         Construct Arrow DictionaryArray from array of indices (must be
         non-negative integers) and corresponding array of dictionary values
@@ -352,8 +354,8 @@ cdef class DictionaryArray(Array):
             shared_ptr[CDataType] c_type
             shared_ptr[CArray] c_result
 
-        arrow_indices = Array.from_pandas(indices, mask=mask)
-        arrow_dictionary = Array.from_pandas(dictionary)
+        arrow_indices = Array.from_pandas(indices, mask=mask, memory_pool=memory_pool)
+        arrow_dictionary = Array.from_pandas(dictionary, memory_pool=memory_pool)
 
         if not isinstance(arrow_indices, IntegerArray):
             raise ValueError('Indices must be integer type')

python/pyarrow/includes/libarrow.pxd

@@ -90,7 +90,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
     shared_ptr[CDataType] timestamp(TimeUnit unit)
 
-    cdef cppclass MemoryPool" arrow::MemoryPool":
+    cdef cppclass CMemoryPool" arrow::MemoryPool":
         int64_t bytes_allocated()
 
     cdef cppclass CBuffer" arrow::Buffer":
@@ -104,9 +104,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
     cdef cppclass PoolBuffer(ResizableBuffer):
         PoolBuffer()
-        PoolBuffer(MemoryPool*)
+        PoolBuffer(CMemoryPool*)
 
-    cdef MemoryPool* default_memory_pool()
+    cdef CMemoryPool* default_memory_pool()
 
     cdef cppclass CListType" arrow::ListType"(CDataType):
         CListType(const shared_ptr[CDataType]& value_type)

python/pyarrow/includes/libarrow_io.pxd

@@ -82,7 +82,7 @@ cdef extern from "arrow/io/file.h" namespace "arrow::io" nogil:
         CStatus Open(const c_string& path, shared_ptr[ReadableFile]* file)
 
         @staticmethod
-        CStatus Open(const c_string& path, MemoryPool* memory_pool,
+        CStatus Open(const c_string& path, CMemoryPool* memory_pool,
                      shared_ptr[ReadableFile]* file)
 
         int file_descriptor()

python/pyarrow/includes/libarrow_ipc.pxd

@@ -18,8 +18,7 @@
 # distutils: language = c++
 
 from pyarrow.includes.common cimport *
-from pyarrow.includes.libarrow cimport (MemoryPool, CArray, CSchema,
-                                        CRecordBatch)
+from pyarrow.includes.libarrow cimport (CArray, CSchema, CRecordBatch)
 from pyarrow.includes.libarrow_io cimport (InputStream, OutputStream,
                                            ReadableFileInterface)
 

python/pyarrow/includes/pyarrow.pxd

@@ -20,7 +20,7 @@
 from pyarrow.includes.common cimport *
 from pyarrow.includes.libarrow cimport (CArray, CBuffer, CColumn, CField,
                                         CTable, CDataType, CStatus, Type,
-                                        MemoryPool, TimeUnit)
+                                        CMemoryPool, TimeUnit)
 
 cimport pyarrow.includes.libarrow_io as arrow_io
 
@@ -28,9 +28,9 @@ cimport pyarrow.includes.libarrow_io as arrow_io
 cdef extern from "pyarrow/api.h" namespace "pyarrow" nogil:
     shared_ptr[CDataType] GetPrimitiveType(Type type)
     shared_ptr[CDataType] GetTimestampType(TimeUnit unit)
-    CStatus ConvertPySequence(object obj, shared_ptr[CArray]* out)
+    CStatus ConvertPySequence(object obj, CMemoryPool* pool, shared_ptr[CArray]* out)
 
-    CStatus PandasToArrow(MemoryPool* pool, object ao, object mo,
+    CStatus PandasToArrow(CMemoryPool* pool, object ao, object mo,
                           shared_ptr[CField] field,
                           shared_ptr[CArray]* out)
 
@@ -43,7 +43,8 @@ cdef extern from "pyarrow/api.h" namespace "pyarrow" nogil:
     CStatus ConvertTableToPandas(const shared_ptr[CTable]& table,
                                  int nthreads, PyObject** out)
 
-    MemoryPool* get_memory_pool()
+    void set_default_memory_pool(CMemoryPool* pool)
+    CMemoryPool* get_memory_pool()
 
 
 cdef extern from "pyarrow/common.h" namespace "pyarrow" nogil:

python/pyarrow/io.pyx

@@ -33,6 +33,7 @@ cimport pyarrow.includes.pyarrow as pyarrow
 
 from pyarrow.compat import frombytes, tobytes, encode_file_path
 from pyarrow.error cimport check_status
+from pyarrow.memory cimport MemoryPool, maybe_unbox_memory_pool
 from pyarrow.schema cimport Schema
 from pyarrow.table cimport (RecordBatch, batch_from_cbatch,
                             table_from_ctable)
@@ -372,7 +373,7 @@ cdef class OSFile(NativeFile):
     cdef:
         object path
 
-    def __cinit__(self, path, mode='r'):
+    def __cinit__(self, path, mode='r', MemoryPool memory_pool=None):
         self.path = path
 
         cdef:
@@ -383,7 +384,7 @@ cdef class OSFile(NativeFile):
         self.is_readable = self.is_writeable = 0
 
         if mode in ('r', 'rb'):
-            self._open_readable(c_path)
+            self._open_readable(c_path, maybe_unbox_memory_pool(memory_pool))
         elif mode in ('w', 'wb'):
             self._open_writeable(c_path)
         else:
@@ -391,12 +392,11 @@ cdef class OSFile(NativeFile):
 
         self.is_open = True
 
-    cdef _open_readable(self, c_string path):
+    cdef _open_readable(self, c_string path, CMemoryPool* pool):
         cdef shared_ptr[ReadableFile] handle
 
         with nogil:
-            check_status(ReadableFile.Open(path, pyarrow.get_memory_pool(),
-                                           &handle))
+            check_status(ReadableFile.Open(path, pool, &handle))
 
         self.is_readable = 1
         self.rd_file = <shared_ptr[ReadableFileInterface]> handle
@@ -450,9 +450,9 @@ cdef class Buffer:
             self.buffer.get().size())
 
 
-cdef shared_ptr[PoolBuffer] allocate_buffer():
+cdef shared_ptr[PoolBuffer] allocate_buffer(CMemoryPool* pool):
     cdef shared_ptr[PoolBuffer] result
-    result.reset(new PoolBuffer(pyarrow.get_memory_pool()))
+    result.reset(new PoolBuffer(pool))
     return result
 
 
@@ -461,8 +461,8 @@ cdef class InMemoryOutputStream(NativeFile):
     cdef:
         shared_ptr[PoolBuffer] buffer
 
-    def __cinit__(self):
-        self.buffer = allocate_buffer()
+    def __cinit__(self, MemoryPool memory_pool=None):
+        self.buffer = allocate_buffer(maybe_unbox_memory_pool(memory_pool))
         self.wr_file.reset(new BufferOutputStream(
             <shared_ptr[ResizableBuffer]> self.buffer))
         self.is_readable = 0

python/pyarrow/memory.pxd

@@ -0,0 +1,27 @@
+# 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.
+
+from pyarrow.includes.libarrow cimport CMemoryPool
+
+
+cdef class MemoryPool:
+    cdef:
+        CMemoryPool* pool
+
+    cdef init(self, CMemoryPool* pool)
+
+cdef CMemoryPool* maybe_unbox_memory_pool(MemoryPool memory_pool)

python/pyarrow/memory.pyx

@@ -0,0 +1,46 @@
+# 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.
+
+# cython: profile=False
+# distutils: language = c++
+# cython: embedsignature = True
+
+from pyarrow.includes.libarrow cimport CMemoryPool
+from pyarrow.includes.pyarrow cimport set_default_memory_pool, get_memory_pool
+
+cdef class MemoryPool:
+    cdef init(self, CMemoryPool* pool):
+        self.pool = pool
+
+cdef CMemoryPool* maybe_unbox_memory_pool(MemoryPool memory_pool):
+    if memory_pool is None:
+        return get_memory_pool()
+    else:
+        return memory_pool.pool
+
+def default_pool():
+    cdef: 
+        MemoryPool pool = MemoryPool()
+    pool.init(get_memory_pool())
+    return pool
+
+def set_default_pool(MemoryPool pool):
+    set_default_memory_pool(pool.pool)
+
+def total_allocated_bytes():
+    cdef CMemoryPool* pool = get_memory_pool()
+    return pool.bytes_allocated()

python/setup.py

@@ -94,6 +94,7 @@ class build_ext(_build_ext):
         'config',
         'error',
         'io',
+        'memory',
         '_parquet',
         'scalar',
         'schema',

python/src/pyarrow/adapters/builtin.cc

@@ -29,6 +29,7 @@
 
 using arrow::ArrayBuilder;
 using arrow::DataType;
+using arrow::MemoryPool;
 using arrow::Status;
 using arrow::Type;
 
@@ -495,7 +496,8 @@ Status ListConverter::Init(const std::shared_ptr<ArrayBuilder>& builder) {
   return Status::OK();
 }
 
-Status ConvertPySequence(PyObject* obj, std::shared_ptr<arrow::Array>* out) {
+Status ConvertPySequence(
+    PyObject* obj, MemoryPool* pool, std::shared_ptr<arrow::Array>* out) {
   std::shared_ptr<DataType> type;
   int64_t size;
   PyDateTime_IMPORT;
@@ -516,7 +518,7 @@ Status ConvertPySequence(PyObject* obj, std::shared_ptr<arrow::Array>* out) {
 
   // Give the sequence converter an array builder
   std::shared_ptr<ArrayBuilder> builder;
-  RETURN_NOT_OK(arrow::MakeBuilder(get_memory_pool(), type, &builder));
+  RETURN_NOT_OK(arrow::MakeBuilder(pool, type, &builder));
   converter->Init(builder);
 
   RETURN_NOT_OK(converter->AppendData(obj));

python/src/pyarrow/adapters/builtin.h

@@ -38,7 +38,8 @@ class Status;
 namespace pyarrow {
 
 PYARROW_EXPORT
-arrow::Status ConvertPySequence(PyObject* obj, std::shared_ptr<arrow::Array>* out);
+arrow::Status ConvertPySequence(
+    PyObject* obj, arrow::MemoryPool* pool, std::shared_ptr<arrow::Array>* out);
 
 }  // namespace pyarrow
 

python/src/pyarrow/common.cc

@@ -28,58 +28,21 @@ using arrow::Status;
 
 namespace pyarrow {
 
-class PyArrowMemoryPool : public arrow::MemoryPool {
- public:
-  PyArrowMemoryPool() : bytes_allocated_(0) {}
-  virtual ~PyArrowMemoryPool() {}
+static std::mutex memory_pool_mutex;
+static arrow::MemoryPool* default_pyarrow_pool = nullptr;
 
-  Status Allocate(int64_t size, uint8_t** out) override {
-    std::lock_guard<std::mutex> guard(pool_lock_);
-    *out = static_cast<uint8_t*>(std::malloc(size));
-    if (*out == nullptr) {
-      std::stringstream ss;
-      ss << "malloc of size " << size << " failed";
-      return Status::OutOfMemory(ss.str());
-    }
-
-    bytes_allocated_ += size;
-
-    return Status::OK();
-  }
-
-  Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override {
-    *ptr = reinterpret_cast<uint8_t*>(std::realloc(*ptr, new_size));
-
-    if (*ptr == NULL) {
-      std::stringstream ss;
-      ss << "realloc of size " << new_size << " failed";
-      return Status::OutOfMemory(ss.str());
-    }
-
-    bytes_allocated_ += new_size - old_size;
-
-    return Status::OK();
-  }
-
-  int64_t bytes_allocated() const override {
-    std::lock_guard<std::mutex> guard(pool_lock_);
-    return bytes_allocated_;
-  }
-
-  void Free(uint8_t* buffer, int64_t size) override {
-    std::lock_guard<std::mutex> guard(pool_lock_);
-    std::free(buffer);
-    bytes_allocated_ -= size;
-  }
-
- private:
-  mutable std::mutex pool_lock_;
-  int64_t bytes_allocated_;
-};
+void set_default_memory_pool(arrow::MemoryPool* pool) {
+  std::lock_guard<std::mutex> guard(memory_pool_mutex);
+  default_pyarrow_pool = pool;
+}
 
 arrow::MemoryPool* get_memory_pool() {
-  static PyArrowMemoryPool memory_pool;
-  return &memory_pool;
+  std::lock_guard<std::mutex> guard(memory_pool_mutex);
+  if (default_pyarrow_pool) {
+    return default_pyarrow_pool;
+  } else {
+    return arrow::default_memory_pool();
+  }
 }
 
 // ----------------------------------------------------------------------

python/src/pyarrow/common.h

@@ -98,6 +98,7 @@ struct PyObjectStringify {
   }
 
 // Return the common PyArrow memory pool
+PYARROW_EXPORT void set_default_memory_pool(arrow::MemoryPool* pool);
 PYARROW_EXPORT arrow::MemoryPool* get_memory_pool();
 
 class PYARROW_EXPORT NumPyBuffer : public arrow::Buffer {