From 1050454467e4d495409f375e8f0020b76e8dbfb6 Mon Sep 17 00:00:00 2001
From: mika <211269698+mikamikasuki@users.noreply.github.com>
Date: Sun, 18 Oct 2026 10:26:48 -0700
Subject: [PATCH] GH-50482: [Python] Wait for CPU tasks before interpreter
 finalization

---
 python/pyarrow/__init__.py           |  7 +++++
 python/pyarrow/includes/libarrow.pxd |  1 +
 python/pyarrow/lib.pyx               |  6 ++++
 python/pyarrow/tests/test_misc.py    | 44 ++++++++++++++++++++++++++++
 4 files changed, 58 insertions(+)

diff --git a/python/pyarrow/__init__.py b/python/pyarrow/__init__.py
index f2a874f56fa4..984ec5e987f2 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_cpu_thread_pool_at_exit as _wait_for_cpu_thread_pool_at_exit
+)
+_atexit.register(_wait_for_cpu_thread_pool_at_exit)
+del _atexit, _wait_for_cpu_thread_pool_at_exit
+
 
 def show_versions():
     """
diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd
index 296148f52816..bd77489aa92f 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -2202,6 +2202,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* GetCpuThreadPool()
 
diff --git a/python/pyarrow/lib.pyx b/python/pyarrow/lib.pyx
index d949720ca8e5..bc1a17c7083a 100644
--- a/python/pyarrow/lib.pyx
+++ b/python/pyarrow/lib.pyx
@@ -89,6 +89,12 @@ def set_cpu_count(int count):
     check_status(SetCpuThreadPoolCapacity(count))
 
 
+def _wait_for_cpu_thread_pool_at_exit():
+    """Wait for CPU tasks to finish before interpreter finalization."""
+    with nogil:
+        GetCpuThreadPool().WaitForIdle()
+
+
 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..cb02d0bdab0e 100644
--- a/python/pyarrow/tests/test_misc.py
+++ b/python/pyarrow/tests/test_misc.py
@@ -137,6 +137,50 @@ 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.skipif(sys.platform == "win32",
                     reason="Path to timezone database is not configurable "
                            "on non-Windows platforms")