GH-50482: [Python] Wait for CPU tasks before interpreter finalization

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<ThreadPool::State>()),
       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();

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<Executor::Resource> 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<ThreadPool> MakeCpuThreadPool();

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<void(ThreadPool*, ThreadPool*, int)> 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<int>);

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():
     """

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()
 

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 = <CThreadPool*>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.

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")