ARROW-458: Python: Expose jemalloc MemoryPool

ci/travis_script_python.sh

@@ -85,9 +85,10 @@ python_version_tests() {
   # Other stuff pip install
   pip install -r requirements.txt
 
-  python setup.py build_ext --inplace --with-parquet
+  python setup.py build_ext --inplace --with-parquet --with-jemalloc
 
   python -c "import pyarrow.parquet"
+  python -c "import pyarrow.jemalloc"
 
   python -m pytest -vv -r sxX pyarrow
 

python/CMakeLists.txt

@@ -53,6 +53,9 @@ if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}")
   option(PYARROW_BUILD_PARQUET
     "Build the PyArrow Parquet integration"
     OFF)
+  option(PYARROW_BUILD_JEMALLOC
+    "Build the PyArrow jemalloc integration"
+    OFF)
 endif()
 
 if(NOT PYARROW_BUILD_TESTS)
@@ -447,6 +450,17 @@ if (PYARROW_BUILD_PARQUET)
     _parquet)
 endif()
 
+if (PYARROW_BUILD_JEMALLOC)
+  ADD_THIRDPARTY_LIB(arrow_jemalloc
+    SHARED_LIB ${ARROW_JEMALLOC_SHARED_LIB})
+  set(LINK_LIBS
+    ${LINK_LIBS}
+    arrow_jemalloc)
+  set(CYTHON_EXTENSIONS
+    ${CYTHON_EXTENSIONS}
+    jemalloc)
+endif()
+
 add_library(pyarrow SHARED
   ${PYARROW_SRCS})
 target_link_libraries(pyarrow ${LINK_LIBS})

python/cmake_modules/FindArrow.cmake

@@ -52,11 +52,17 @@ find_library(ARROW_IPC_LIB_PATH NAMES arrow_ipc
   ${ARROW_SEARCH_LIB_PATH}
   NO_DEFAULT_PATH)
 
+find_library(ARROW_JEMALLOC_LIB_PATH NAMES arrow_jemalloc
+  PATHS
+  ${ARROW_SEARCH_LIB_PATH}
+  NO_DEFAULT_PATH)
+
 if (ARROW_INCLUDE_DIR AND ARROW_LIB_PATH)
   set(ARROW_FOUND TRUE)
   set(ARROW_LIB_NAME libarrow)
   set(ARROW_IO_LIB_NAME libarrow_io)
   set(ARROW_IPC_LIB_NAME libarrow_ipc)
+  set(ARROW_JEMALLOC_LIB_NAME libarrow_jemalloc)
 
   set(ARROW_LIBS ${ARROW_SEARCH_LIB_PATH})
   set(ARROW_STATIC_LIB ${ARROW_SEARCH_LIB_PATH}/${ARROW_LIB_NAME}.a)
@@ -68,10 +74,14 @@ if (ARROW_INCLUDE_DIR AND ARROW_LIB_PATH)
   set(ARROW_IPC_STATIC_LIB ${ARROW_SEARCH_LIB_PATH}/${ARROW_IPC_LIB_NAME}.a)
   set(ARROW_IPC_SHARED_LIB ${ARROW_LIBS}/${ARROW_IPC_LIB_NAME}${CMAKE_SHARED_LIBRARY_SUFFIX})
 
+  set(ARROW_JEMALLOC_STATIC_LIB ${ARROW_SEARCH_LIB_PATH}/${ARROW_JEMALLOC_LIB_NAME}.a)
+  set(ARROW_JEMALLOC_SHARED_LIB ${ARROW_LIBS}/${ARROW_JEMALLOC_LIB_NAME}${CMAKE_SHARED_LIBRARY_SUFFIX})
+
   if (NOT Arrow_FIND_QUIETLY)
     message(STATUS "Found the Arrow core library: ${ARROW_LIB_PATH}")
     message(STATUS "Found the Arrow IO library: ${ARROW_IO_LIB_PATH}")
     message(STATUS "Found the Arrow IPC library: ${ARROW_IPC_LIB_PATH}")
+    message(STATUS "Found the Arrow jemalloc library: ${ARROW_JEMALLOC_LIB_PATH}")
   endif ()
 else ()
   if (NOT Arrow_FIND_QUIETLY)
@@ -94,4 +104,8 @@ mark_as_advanced(
   ARROW_SHARED_LIB
   ARROW_IO_STATIC_LIB
   ARROW_IO_SHARED_LIB
+  ARROW_IPC_STATIC_LIB
+  ARROW_IPC_SHARED_LIB
+  ARROW_JEMALLOC_STATIC_LIB
+  ARROW_JEMALLOC_SHARED_LIB
 )

python/pyarrow/includes/libarrow_jemalloc.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.
+
+# distutils: language = c++
+
+from pyarrow.includes.common cimport *
+from pyarrow.includes.libarrow cimport *
+    
+cdef extern from "arrow/jemalloc/memory_pool.h" namespace "arrow::jemalloc" nogil:
+    cdef cppclass CJemallocMemoryPool" arrow::jemalloc::MemoryPool":
+        int64_t bytes_allocated()
+        @staticmethod
+        CMemoryPool* default_pool()

python/pyarrow/jemalloc.pyx

@@ -0,0 +1,28 @@
+# 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_jemalloc cimport CJemallocMemoryPool
+from pyarrow.memory cimport MemoryPool
+
+def default_pool():
+    cdef MemoryPool pool = MemoryPool()
+    pool.init(CJemallocMemoryPool.default_pool())
+    return pool

python/pyarrow/memory.pyx

@@ -26,6 +26,9 @@ cdef class MemoryPool:
     cdef init(self, CMemoryPool* pool):
         self.pool = pool
 
+    def bytes_allocated(self):
+        return self.pool.bytes_allocated()
+
 cdef CMemoryPool* maybe_unbox_memory_pool(MemoryPool memory_pool):
     if memory_pool is None:
         return get_memory_pool()

python/pyarrow/tests/test_jemalloc.py

@@ -0,0 +1,56 @@
+# 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 gc
+import pytest
+
+try:
+    import pyarrow.jemalloc
+    HAVE_JEMALLOC = True
+except ImportError:
+    HAVE_JEMALLOC = False
+
+jemalloc = pytest.mark.skipif(not HAVE_JEMALLOC,
+                              reason='jemalloc support not built')
+
+
+@jemalloc
+def test_different_memory_pool():
+    gc.collect()
+    bytes_before_default = pyarrow.total_allocated_bytes()
+    bytes_before_jemalloc = pyarrow.jemalloc.default_pool().bytes_allocated()
+    array = pyarrow.from_pylist([1, None, 3, None],
+                        memory_pool=pyarrow.jemalloc.default_pool())
+    gc.collect()
+    assert pyarrow.total_allocated_bytes() == bytes_before_default
+    assert pyarrow.jemalloc.default_pool().bytes_allocated() > bytes_before_jemalloc
+
+@jemalloc
+def test_default_memory_pool():
+    gc.collect()
+    bytes_before_default = pyarrow.total_allocated_bytes()
+    bytes_before_jemalloc = pyarrow.jemalloc.default_pool().bytes_allocated()
+
+    old_memory_pool = pyarrow.memory.default_pool()
+    pyarrow.memory.set_default_pool(pyarrow.jemalloc.default_pool())
+    array = pyarrow.from_pylist([1, None, 3, None])
+    pyarrow.memory.set_default_pool(old_memory_pool)
+    gc.collect()
+
+    assert pyarrow.total_allocated_bytes() == bytes_before_default
+    assert pyarrow.jemalloc.default_pool().bytes_allocated() > bytes_before_jemalloc
+

python/setup.py

@@ -80,7 +80,8 @@ class build_ext(_build_ext):
     description = "Build the C-extensions for arrow"
     user_options = ([('extra-cmake-args=', None, 'extra arguments for CMake'),
                      ('build-type=', None, 'build type (debug or release)'),
-                     ('with-parquet', None, 'build the Parquet extension')] +
+                     ('with-parquet', None, 'build the Parquet extension'),
+                     ('with-jemalloc', None, 'build the jemalloc extension')] +
                     _build_ext.user_options)
 
     def initialize_options(self):
@@ -88,12 +89,14 @@ class build_ext(_build_ext):
         self.extra_cmake_args = os.environ.get('PYARROW_CMAKE_OPTIONS', '')
         self.build_type = os.environ.get('PYARROW_BUILD_TYPE', 'debug').lower()
         self.with_parquet = False
+        self.with_jemalloc = False
 
     CYTHON_MODULE_NAMES = [
         'array',
         'config',
         'error',
         'io',
+        'jemalloc',
         'memory',
         '_parquet',
         'scalar',
@@ -136,6 +139,9 @@ class build_ext(_build_ext):
         if self.with_parquet:
             cmake_options.append('-DPYARROW_BUILD_PARQUET=on')
 
+        if self.with_jemalloc:
+            cmake_options.append('-DPYARROW_BUILD_JEMALLOC=on')
+
         if sys.platform != 'win32':
             cmake_options.append('-DCMAKE_BUILD_TYPE={0}'
                                  .format(self.build_type))
@@ -217,6 +223,8 @@ class build_ext(_build_ext):
     def _failure_permitted(self, name):
         if name == '_parquet' and not self.with_parquet:
             return True
+        if name == 'jemalloc' and not self.with_jemalloc:
+            return True
         return False
 
     def _get_inplace_dir(self):