Own the AioS3FileSystem pool and keep aio DataFrame reads off the event loop

Part of #1205 (step 5; see R5 and model 5 there).

Problem

AioS3FileSystem._thread_pool (pyathena/filesystem/s3_async.py:145) is never shut down. A cached instance keeps its threads for the life of the process.
AioPandasCursor.as_pandas() returns the synchronous PandasDataFrameIterator (pyathena/aio/pandas/cursor.py:239-248). Its chunk reads, and with automatic chunks the whole read, run on the event loop thread. AioPolarsCursor.as_polars() / as_arrow() follow the same pattern (pyathena/aio/polars/cursor.py:214-234, source inspection only).

Reproduction

The stub reproductions are summarized in #1205. They use no AWS calls.

Environment

PyAthena master at b2d7aa2
Python 3.15

Proposed fix (optional)

Give AioS3FileSystem an owned pool lifecycle, through a close() or a release at garbage collection, coordinated with the executor cleanup in Free the data of failed S3 requests and clean up async copies at loop shutdown #1192.
Make the aio cursors' as_pandas(), as_polars(), and as_arrow() read through threads, or document that the returned iterator blocks.

Decisions (2026-10-18)

Agreed with the maintainer after design consultations with Codex and claude-fable-5-1. Documentation alone was ruled out: docs/aio.md already describes the blocking.

AioS3FileSystem.close() releases the threads of its pool, without ending the filesystem.
- S3AioThreadPool drops its ThreadPoolExecutor (shutdown(wait=False)), and the next call creates a new one, as after a fork. Calls that were submitted before still run.
- There is no closed state: the filesystem stays usable, with no error after close(), and it stays in fsspec's instance cache, so other holders of a cached instance are not broken.
- The cleanup of a multipart copy that is running at close() (the case in the comment above) still schedules its completion and abort.
- with and async with call close() on exit.
- Rejected: a terminal close(), which would have raised RuntimeError afterwards and evicted the instance from the cache. It would break other holders of a shared cached instance and needs closed checks at every I/O entry point. weakref.finalize was rejected too: collection already releases the threads of an uncached instance, and a cached one is never collected.
Async reads on the aio cursors; the synchronous as_* methods stay unchanged.
- AioPandasCursor.iter_chunks() and AioPolarsCursor.iter_chunks() return async generators (async for df in cursor.iter_chunks()). They advance one result-set chunk iterator in threads, bound to the result set at the start. A cancellation follows the fetch methods.
- await AioPandasCursor.read_pandas() and await AioPolarsCursor.read_polars() / read_arrow() return the whole remaining result. They run the existing collectors, including the categorical restoration of pandas, in a thread.
- AioArrowCursor is unchanged: its as_arrow() never reads S3.
- Rejected: making as_* coroutines, which is breaking and still leaves a synchronous chunk iterator. Also rejected: documenting pd.concat / pl.concat over iter_chunks() as the whole-read path. It fails on an empty result, concatenates on the event loop, and skips the categorical restoration.

Revisions during the review of #1286 (2026-10-18)

Pool release. S3AioThreadPool.release() keeps the executor until no submitted call runs or waits, and then shuts it down. Dropping a busy executor at once would stack executors beyond max_workers when a shared filesystem is closed repeatedly.
Whole reads. read_pandas() / read_polars() / read_arrow() are dropped. The as_*() methods of the aio cursors become coroutines instead, which follows the aio naming of PyAthena: the same name as the synchronous method, awaited. This covers AioPandasCursor.as_pandas(), AioPolarsCursor.as_polars() / as_arrow(), and AioArrowCursor.as_arrow() / as_polars(). No base class defines them, and nothing in PyAthena calls them.
- AioPandasCursor.as_pandas() always returns a DataFrame. With chunks, it joins the chunks that have not yet been read.
- This is a breaking change for 4.2.0.
Chunk iterators after close. The synchronous chunk iterators of AthenaPandasResultSet and AthenaPolarsResultSet used to end silently once their result set was closed. They now raise ProgrammingError, matching the aio iterator, and this replaces the behavior documented in docs/usage.md. The iterators share a closed flag with their result set, so they hold no reference to it.
Async iter_chunks(). It binds the result set when called. Leaving the loop early keeps the unread chunks with the result set.

A related case for the pool ownership in item 1, found in the review of #1275 (#1230). It predates #1275: master at 7030514 behaves the same way.

Once AioS3FileSystem._thread_pool has been shut down, the multipart copy cannot clean up its upload. AioS3FileSystem._copy_object_with_multipart_upload() runs its cleanup with await self._thread_pool.run_shielded(_abort) in its except BaseException branch. S3AioThreadPool.submit() then raises RuntimeError: cannot schedule new futures after shutdown. As a result:

On a cancellation, that RuntimeError replaces the CancelledError.
If the pool is shut down while a part copy runs, the completion cannot be scheduled either, so the cleanup raises a second RuntimeError.
In both cases, neither CompleteMultipartUpload nor AbortMultipartUpload is sent. The created multipart upload is left incomplete.

Nothing in PyAthena shuts this pool down today: item 1 notes that it is never shut down. The case arises only once the pool gets a lifecycle, through a close() or a release at garbage collection, or when a caller shuts it down. That lifecycle should keep the cleanup of a running copy able to run, or run it outside the pool.

Reproduction with stubbed requests (no AWS calls):

import asyncio
import threading
from types import SimpleNamespace
from unittest import mock

from pyathena.filesystem.s3_async import AioS3FileSystem
from pyathena.filesystem.s3_core import S3Core
from pyathena.filesystem.s3_object import S3MultipartUpload
from pyathena.filesystem.s3_path import S3Path


async def main(cancel):
    fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True)
    core = fs._sync_fs.core
    core.create_multipart_upload = mock.MagicMock(
        return_value=S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "u"})
    )
    started, release = threading.Event(), threading.Event()

    def upload_part_copy(**kw):
        started.set()
        release.wait(10)
        return SimpleNamespace(etag='"e"', part_number=kw["part_number"])

    core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy)
    core.complete_multipart_upload = mock.MagicMock()
    core.abort_multipart_upload = mock.MagicMock()
    # The HeadObject of a source with 2 parts.
    size = 2 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE
    fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock(return_value={"ContentLength": size})
    task = asyncio.ensure_future(
        fs._copy_object_with_multipart_upload(
            S3Path("bucket", "src"),
            S3Path("bucket", "dst"),
            MetadataDirective="REPLACE",
            TaggingDirective="REPLACE",
            AnnotationDirective="EXCLUDE",
        )
    )
    assert await asyncio.to_thread(started.wait, 10)
    fs._thread_pool.shutdown(wait=False)
    if cancel:
        task.cancel()
        await asyncio.sleep(0.1)
    release.set()
    try:
        await task
    except BaseException as e:
        print(f"{type(e).__name__}: {e}")
    await asyncio.sleep(0.5)
    print(core.complete_multipart_upload.call_count, core.abort_multipart_upload.call_count)


asyncio.run(main(cancel=True))
asyncio.run(main(cancel=False))

Output, on master at 7030514 and on the #1275 branch at aaed141:

RuntimeError: cannot schedule new futures after shutdown
0 0
RuntimeError: cannot schedule new futures after shutdown
0 0

Release the pool threads when idle and bind iter_chunks() at its call (…

Share the idle release step of S3AioThreadPool and trim docstrings (G…

Describe the idle release and the aio chunk iterator after close (GH-…

Track pending futures for the idle release and correct the aio read d…

Check that close() does not wait for a running copy part (GH-1209)

Make the aio as_*() methods coroutines and raise from closed chunk it…

Prune the pending futures only in release() and correct close-time do…

Refer to the pool weakly from future callbacks and narrow close-time …

Drop the retry claim of a cancelled chunk read and settle the pool GC…

Merge pull request #1286 from pyathena-dev/feat/1209-aio-pool-close-a…

Navigation Menu

Repository navigation

Own the AioS3FileSystem pool and keep aio DataFrame reads off the event loop #1209

Problem

Reproduction

Environment

Proposed fix (optional)

Decisions (2026-10-18)

Revisions during the review of #1286 (2026-10-18)

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Own the AioS3FileSystem pool and keep aio DataFrame reads off the event loop #1209

Description

Problem

Reproduction

Environment

Proposed fix (optional)

Decisions (2026-10-18)

Revisions during the review of #1286 (2026-10-18)

Activity

laughingman7743 commented on Oct 18, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions