ARROW-12760: [C++][Python][R] Allow setting I/O thread pool size

This adds functions to change the pool size at runtime, but doesn't adjust the default size of 8.

Closes #10316 from lidavidm/arrow-12760

Authored-by: David Li <li.davidm96@gmail.com>
Signed-off-by: Antoine Pitrou <antoine@python.org>

Author: lidavidm

cpp/src/arrow/io/interfaces.cc

@@ -29,6 +29,7 @@
 
 #include "arrow/buffer.h"
 #include "arrow/io/concurrency.h"
+#include "arrow/io/type_fwd.h"
 #include "arrow/io/util_internal.h"
 #include "arrow/result.h"
 #include "arrow/status.h"
@@ -53,6 +54,12 @@ IOContext::IOContext(MemoryPool* pool, StopToken stop_token)
 
 const IOContext& default_io_context() { return g_default_io_context; }
 
+int GetIOThreadPoolCapacity() { return internal::GetIOThreadPool()->GetCapacity(); }
+
+Status SetIOThreadPoolCapacity(int threads) {
+  return internal::GetIOThreadPool()->SetCapacity(threads);
+}
+
 FileInterface::~FileInterface() = default;
 
 Status FileInterface::Abort() { return Close(); }

cpp/src/arrow/io/memory_test.cc

@@ -861,5 +861,15 @@ TEST(CacheOptions, Basics) {
   check(CacheOptions::MakeFromNetworkMetrics(5, 500, .75, 5), 2.5, 5);
 }
 
+TEST(IOThreadPool, Capacity) {
+  // Simple sanity check
+  auto pool = internal::GetIOThreadPool();
+  int capacity = pool->GetCapacity();
+  ASSERT_GT(capacity, 0);
+  ASSERT_EQ(GetIOThreadPoolCapacity(), capacity);
+  ASSERT_OK(SetIOThreadPoolCapacity(capacity + 1));
+  ASSERT_EQ(GetIOThreadPoolCapacity(), capacity + 1);
+}
+
 }  // namespace io
 }  // namespace arrow

cpp/src/arrow/io/type_fwd.h

@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include "arrow/type_fwd.h"
 #include "arrow/util/visibility.h"
 
 namespace arrow {
@@ -33,6 +34,23 @@ struct CacheOptions;
 ARROW_EXPORT
 const IOContext& default_io_context();
 
+/// \brief Get the capacity of the global I/O thread pool
+///
+/// Return the number of worker threads in the thread pool to which
+/// Arrow dispatches various I/O-bound tasks.  This is an ideal number,
+/// not necessarily the exact number of threads at a given point in time.
+///
+/// You can change this number using SetIOThreadPoolCapacity().
+ARROW_EXPORT int GetIOThreadPoolCapacity();
+
+/// \brief Set the capacity of the global I/O thread pool
+///
+/// Set the number of worker threads in the thread pool to which
+/// Arrow dispatches various I/O-bound tasks.
+///
+/// The current number is returned by GetIOThreadPoolCapacity().
+ARROW_EXPORT Status SetIOThreadPoolCapacity(int threads);
+
 class FileInterface;
 class Seekable;
 class Writable;

python/pyarrow/__init__.py

@@ -67,7 +67,8 @@ def parse_git(root, **kwargs):
 from pyarrow.lib import (BuildInfo, RuntimeInfo, VersionInfo,
                          cpp_build_info, cpp_version, cpp_version_info,
                          runtime_info, cpu_count, set_cpu_count,
-                         enable_signal_handlers)
+                         enable_signal_handlers,
+                         io_thread_count, set_io_thread_count)
 
 
 def show_versions():

python/pyarrow/includes/libarrow.pxd

@@ -1163,6 +1163,8 @@ cdef extern from "arrow/io/api.h" namespace "arrow::io" nogil:
         CIOContext(CMemoryPool*, CStopToken)
 
     CIOContext c_default_io_context "arrow::io::default_io_context"()
+    int GetIOThreadPoolCapacity()
+    CStatus SetIOThreadPoolCapacity(int threads)
 
     cdef cppclass FileStatistics:
         int64_t size

python/pyarrow/io.pxi

@@ -42,6 +42,46 @@ cdef extern from "Python.h":
         char *v, Py_ssize_t len) except NULL
 
 
+def io_thread_count():
+    """
+    Return the number of threads to use for I/O operations.
+
+    Many operations, such as scanning a dataset, will implicitly make
+    use of this pool. The number of threads is set to a fixed value at
+    startup. It can be modified at runtime by calling
+    :func:`set_io_thread_count()`.
+
+    See Also
+    --------
+    set_io_thread_count : Modify the size of this pool.
+    cpu_count : The analogous function for the CPU thread pool.
+    """
+    return GetIOThreadPoolCapacity()
+
+
+def set_io_thread_count(int count):
+    """
+    Set the number of threads to use for I/O operations.
+
+    Many operations, such as scanning a dataset, will implicitly make
+    use of this pool.
+
+    Parameters
+    ----------
+    count : int
+        The max number of threads that may be used for I/O.
+        Must be positive.
+
+    See Also
+    --------
+    io_thread_count : Get the size of this pool.
+    set_cpu_count : The analogous function for the CPU thread pool.
+    """
+    if count < 1:
+        raise ValueError("IO thread count must be strictly positive")
+    check_status(SetIOThreadPoolCapacity(count))
+
+
 cdef class NativeFile(_Weakrefable):
     """
     The base class for all Arrow streams.

python/pyarrow/lib.pyx

@@ -48,13 +48,23 @@ def cpu_count():
     If neither is present, it will default to the number of hardware threads
     on the system. It can be modified at runtime by calling
     :func:`set_cpu_count()`.
+
+    See Also
+    --------
+    set_cpu_count : Modify the size of this pool.
+    io_thread_count : The analogous function for the I/O thread pool.
     """
     return GetCpuThreadPoolCapacity()
 
 
 def set_cpu_count(int count):
     """
     Set the number of threads to use in parallel operations.
+
+    See Also
+    --------
+    cpu_count : Get the size of this pool.
+    set_io_thread_count : The analogous function for the I/O thread pool.
     """
     if count < 1:
         raise ValueError("CPU count must be strictly positive")

python/pyarrow/tests/test_misc.py

@@ -45,6 +45,16 @@ def test_cpu_count():
         pa.set_cpu_count(n)
 
 
+def test_io_thread_count():
+    n = pa.io_thread_count()
+    assert n > 0
+    try:
+        pa.set_io_thread_count(n + 5)
+        assert pa.io_thread_count() == n + 5
+    finally:
+        pa.set_io_thread_count(n)
+
+
 def test_build_info():
     assert isinstance(pa.cpp_build_info, pa.BuildInfo)
     assert isinstance(pa.cpp_version_info, pa.VersionInfo)

r/NAMESPACE

@@ -229,6 +229,7 @@ export(int16)
 export(int32)
 export(int64)
 export(int8)
+export(io_thread_count)
 export(is_in)
 export(large_binary)
 export(large_list_of)
@@ -261,6 +262,7 @@ export(record_batch)
 export(s3_bucket)
 export(schema)
 export(set_cpu_count)
+export(set_io_thread_count)
 export(starts_with)
 export(string)
 export(struct)