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

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

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

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.

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