Use C++ API, rather than multiprocessing, in pyarrow.{cpu_count,set_cpu_count}

python/pyarrow/_config.pyx

@@ -1,57 +0,0 @@
-# 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
-
-import numpy as np
-import multiprocessing
-import os
-
-cdef extern from 'arrow/python/init.h':
-    int arrow_init_numpy() except -1
-
-cdef extern from 'arrow/python/config.h' namespace 'arrow::py':
-    void set_numpy_nan(object o)
-
-arrow_init_numpy()
-
-set_numpy_nan(np.nan)
-
-cdef int CPU_COUNT = int(
-    os.environ.get('OMP_NUM_THREADS',
-                   max(multiprocessing.cpu_count() // 2, 1)))
-
-
-def cpu_count():
-    """
-    Returns
-    -------
-    count : Number of CPUs to use by default in parallel operations. Default is
-      max(1, multiprocessing.cpu_count() / 2), but can be overridden by the
-      OMP_NUM_THREADS environment variable. For the default, we divide the CPU
-      count by 2 because most modern computers have hyperthreading turned on,
-      so doubling the CPU count beyond the number of physical cores does not
-      help.
-    """
-    return CPU_COUNT
-
-
-def set_cpu_count(count):
-    global CPU_COUNT
-    CPU_COUNT = max(int(count), 1)

python/pyarrow/includes/libarrow.pxd

@@ -1045,3 +1045,8 @@ cdef extern from 'arrow/util/compression.h' namespace 'arrow' nogil:
                          int64_t* output_length)
 
         int64_t MaxCompressedLen(int64_t input_len, const uint8_t* input)
+
+
+cdef extern from 'arrow/util/thread-pool.h' namespace 'arrow' nogil:
+    size_t GetCPUThreadPoolCapacity()
+    CStatus SetCPUThreadPoolCapacity(size_t threads)

python/pyarrow/lib.pyx

@@ -36,28 +36,26 @@ cimport cpython as cp
 arrow_init_numpy()
 set_numpy_nan(np.nan)
 
-cdef int CPU_COUNT = int(
-    os.environ.get('OMP_NUM_THREADS',
-                   max(multiprocessing.cpu_count() // 2, 1)))
-
 
 def cpu_count():
     """
-    Returns
-    -------
-    count : Number of CPUs to use by default in parallel operations. Default is
-      max(1, multiprocessing.cpu_count() / 2), but can be overridden by the
-      OMP_NUM_THREADS environment variable. For the default, we divide the CPU
-      count by 2 because most modern computers have hyperthreading turned on,
-      so doubling the CPU count beyond the number of physical cores does not
-      help.
+    Return the number of threads to use in parallel operations.
+
+    The number of threads is determined at startup by inspecting the
+    OMP_NUM_THREADS and OMP_THREAD_LIMIT environment variables.  If neither
+    is present, it will default to the number of hardware threads on the
+    system.  It can be modified at runtime by calling set_cpu_count().
     """
-    return CPU_COUNT
+    return GetCPUThreadPoolCapacity()
 
 
-def set_cpu_count(count):
-    global CPU_COUNT
-    CPU_COUNT = max(int(count), 1)
+def set_cpu_count(size_t count):
+    """
+    Set the number of threads to use in parallel operations.
+    """
+    if count < 1:
+        raise ValueError("CPU count must be strictly positive")
+    check_status(SetCPUThreadPoolCapacity(count))
 
 
 Type_NA = _Type_NA

python/pyarrow/tests/test_misc.py

@@ -23,3 +23,13 @@ import pyarrow as pa
 def test_get_include():
     include_dir = pa.get_include()
     assert os.path.exists(os.path.join(include_dir, 'arrow', 'api.h'))
+
+
+def test_cpu_count():
+    n = pa.cpu_count()
+    assert n > 0
+    try:
+        pa.set_cpu_count(n + 5)
+        assert pa.cpu_count() == n + 5
+    finally:
+        pa.set_cpu_count(n)