diff --git a/cpp/src/arrow/util/thread_pool.cc b/cpp/src/arrow/util/thread_pool.cc index 00e09b6573fc..6424a3497ab4 100644 --- a/cpp/src/arrow/util/thread_pool.cc +++ b/cpp/src/arrow/util/thread_pool.cc @@ -525,6 +525,23 @@ void ThreadPool::WaitForIdle() { state_->cv_idle_.wait(lk, [this] { return state_->tasks_queued_or_running_ == 0; }); } +void ThreadPool::WaitForIdle(ThreadPool& other) { + if (this == &other) { + WaitForIdle(); + return; + } + while (true) { + WaitForIdle(); + other.WaitForIdle(); + // A task can move between pools while each pool is checked on its own. + std::scoped_lock lock(state_->mutex_, other.state_->mutex_); + if (state_->tasks_queued_or_running_ == 0 && + other.state_->tasks_queued_or_running_ == 0) { + return; + } + } +} + ThreadPool::ThreadPool() : sp_state_(std::make_shared()), state_(sp_state_.get()), @@ -847,6 +864,12 @@ void ThreadPool::WaitForIdle() { } } +void ThreadPool::WaitForIdle(ThreadPool& other) { + while (!state_->task_queue.empty() || !other.state_->task_queue.empty()) { + RunTasksOnAllExecutors(); + } +} + Status ThreadPool::SetCapacity(int threads) { state_->max_tasks_running = threads; return Status::OK(); diff --git a/cpp/src/arrow/util/thread_pool.h b/cpp/src/arrow/util/thread_pool.h index 582505cda97c..026cf4cd7835 100644 --- a/cpp/src/arrow/util/thread_pool.h +++ b/cpp/src/arrow/util/thread_pool.h @@ -490,6 +490,10 @@ class ARROW_EXPORT ThreadPool : public Executor { // This is useful for sequencing tests void WaitForIdle(); + // Wait for both pools to be idle at the same time. Tasks may submit work to + // the other pool. The caller must prevent new work from other producers. + void WaitForIdle(ThreadPool& other); + void KeepAlive(std::shared_ptr resource) override; struct State; @@ -563,6 +567,7 @@ class ARROW_EXPORT ThreadPool : public SerialExecutor { // // This is useful for sequencing tests void WaitForIdle(); + void WaitForIdle(ThreadPool& other); protected: static std::shared_ptr MakeCpuThreadPool(); diff --git a/cpp/src/arrow/util/thread_pool_test.cc b/cpp/src/arrow/util/thread_pool_test.cc index 7c7498838a64..04253ddadc2b 100644 --- a/cpp/src/arrow/util/thread_pool_test.cc +++ b/cpp/src/arrow/util/thread_pool_test.cc @@ -576,6 +576,30 @@ TEST_F(TestThreadPool, ConstructDestruct) { // Correctness and stress tests using Spawn() and Shutdown() +TEST_F(TestThreadPool, WaitForBothPoolsIdle) { + auto first = MakeThreadPool(1); + auto second = MakeThreadPool(1); + auto finished = Future<>::Make(); + std::function submit; + submit = [&](ThreadPool* from, ThreadPool* to, int remaining) { + ASSERT_OK(from->Spawn([&, from, to, remaining] { + if (remaining == 0) { + finished.MarkFinished(); + } else { + submit(to, from, remaining - 1); + } + })); + }; + submit(first.get(), second.get(), 1000); + first->WaitForIdle(*second); + EXPECT_TRUE(finished.is_finished()); + EXPECT_EQ(first->GetNumTasks(), 0); + EXPECT_EQ(second->GetNumTasks(), 0); + ASSERT_FINISHES_OK(finished); + first->WaitForIdle(*first); + second->WaitForIdle(); +} + TEST_F(TestThreadPool, Spawn) { auto pool = this->MakeThreadPool(3); SpawnAdds(pool.get(), 7, task_add); diff --git a/python/pyarrow/__init__.py b/python/pyarrow/__init__.py index f2a874f56fa4..f61f4637bcef 100644 --- a/python/pyarrow/__init__.py +++ b/python/pyarrow/__init__.py @@ -29,6 +29,7 @@ For more information see the official page at https://arrow.apache.org """ +import atexit as _atexit import importlib as _importlib import os as _os import platform as _platform @@ -62,6 +63,12 @@ def parse_git(root, **kwargs): cpu_count, set_cpu_count, enable_signal_handlers, io_thread_count, is_opentelemetry_enabled, set_io_thread_count) +from pyarrow.lib import ( + _wait_for_thread_pools_at_exit as _wait_for_thread_pools_at_exit +) +_atexit.register(_wait_for_thread_pools_at_exit) +del _atexit, _wait_for_thread_pools_at_exit + def show_versions(): """ diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 296148f52816..74d532f5f595 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -1627,6 +1627,7 @@ cdef extern from "arrow/io/api.h" namespace "arrow::io" nogil: CIOContext(CStopToken) CIOContext(CMemoryPool*) CIOContext(CMemoryPool*, CStopToken) + CExecutor* executor() const CIOContext c_default_io_context "arrow::io::default_io_context"() int GetIOThreadPoolCapacity() @@ -2202,6 +2203,7 @@ cdef extern from "arrow/util/thread_pool.h" namespace "arrow::internal" nogil: cdef cppclass CThreadPool "arrow::internal::ThreadPool"(CExecutor): @staticmethod CResult[shared_ptr[CThreadPool]] Make(int threads) + void WaitForIdle(CThreadPool& other) CThreadPool* GetCpuThreadPool() diff --git a/python/pyarrow/lib.pyx b/python/pyarrow/lib.pyx index d949720ca8e5..abd4b063b13d 100644 --- a/python/pyarrow/lib.pyx +++ b/python/pyarrow/lib.pyx @@ -89,6 +89,13 @@ def set_cpu_count(int count): check_status(SetCpuThreadPoolCapacity(count)) +def _wait_for_thread_pools_at_exit(): + """Wait for CPU and I/O tasks before interpreter finalization.""" + cdef CThreadPool* io_pool = c_default_io_context().executor() + with nogil: + GetCpuThreadPool().WaitForIdle(deref(io_pool)) + + def is_opentelemetry_enabled() -> bool: """ Returns True if OpenTelemetry is enabled in libarrow. diff --git a/python/pyarrow/tests/test_misc.py b/python/pyarrow/tests/test_misc.py index 856873f441b4..a2b55a1753ee 100644 --- a/python/pyarrow/tests/test_misc.py +++ b/python/pyarrow/tests/test_misc.py @@ -137,6 +137,103 @@ def import_arrow(): subprocess.check_call([sys.executable, "-c", code]) +@pytest.mark.processes +def test_cpu_thread_pool_is_idle_before_interpreter_finalization(): + pytest.importorskip("fsspec") + code = """if 1: + import atexit + from pathlib import Path + import shutil + import tempfile + + root = Path(tempfile.mkdtemp()) + atexit.register(shutil.rmtree, root, ignore_errors=True) + + import fsspec + import pyarrow as pa + import pyarrow.dataset as ds + import pyarrow.fs as fs + import pyarrow.parquet as pq + + pa.set_cpu_count(24) + pq.write_table( + pa.table({"value": [1, 2, 3]}), root / "valid.parquet", compression="NONE" + ) + (root / "bad.parquet").write_bytes(b"not a parquet file") + filesystem = fs.PyFileSystem( + fs.FSSpecHandler(fsspec.filesystem("file")) + ) + paths = [str(root / name) for name in ( + "valid.parquet", "bad.parquet", "valid.parquet", "valid.parquet" + )] + dataset = ds.dataset( + paths, filesystem=filesystem, format="parquet", + schema=pa.schema([("value", pa.int64())]), + exclude_invalid_files=False + ) + try: + dataset.to_table(use_threads=True, fragment_readahead=4) + except pa.ArrowInvalid: + pass + else: + raise AssertionError("Expected the malformed-fragment error") + """ + subprocess.run([sys.executable, "-c", code], check=True, timeout=30) + + +@pytest.mark.processes +@pytest.mark.parametrize("use_threads", [False, True]) +def test_pending_prebuffer_reads_finish_before_interpreter_finalization( + tmp_path, use_threads): + code = """if 1: + import io + import sys + import threading + import time + + import pyarrow as pa + import pyarrow.dataset as ds + import pyarrow.parquet as pq + + path = sys.argv[1] + values = [b"x" * 16384] * 1024 + pq.write_table( + pa.table({f"c{i}": values for i in range(4)}), path, + compression="NONE", use_dictionary=False, + ) + pending = threading.Event() + + class RangeErrorFile(io.BufferedReader): + failed = False + + def read(self, size=-1): + if size > 65536 and not self.failed: + self.failed = True + raise OSError("range read failed") + if self.failed: + pending.set() + time.sleep(0.5) + return super().read(size) + + with RangeErrorFile(io.FileIO(path, "r")) as source: + fragment = ds.ParquetFileFormat().make_fragment(source) + reader = fragment.scanner( + use_threads=sys.argv[2] == "True", + batch_readahead=0, fragment_readahead=0, + ).to_reader() + try: + reader.read_next_batch() + except OSError as error: + assert str(error) == "range read failed" + else: + raise AssertionError("Expected the range-read error") + assert pending.wait(5), "No later prebuffer read started" + reader.close() + """ + subprocess.run([sys.executable, "-c", code, str(tmp_path / "data.parquet"), + str(use_threads)], check=True, timeout=30) + + @pytest.mark.skipif(sys.platform == "win32", reason="Path to timezone database is not configurable " "on non-Windows platforms")