Cancelled chunked UNLOAD reads keep retrying in ParquetChunkReader's own thread pool #1296

Problem

Cancelling an asyncio read of a chunked UNLOAD result does not stop PyAthena's retries of the requests that open the Parquet files.

ParquetChunkReader._open_files() (pyathena/arrow/reader.py) opens each file in a ThreadPoolExecutor of its own, with executor.submit(self._open, path). That submission does not copy the caller's context. As a result, the stop scope of pyathena.util.stop_retries_on(), which an asyncio cancellation sets, does not reach _open(). Its HeadObject/GetObject retries run to the end of the RetryConfig, and the cancelled read waits for them in ahead.result(). Only then does the result set close.

S3ThreadPoolExecutor.submit() copies the context for the same reason (#1203).

Affected paths, all with chunksize on UNLOAD results:

• AioPandasCursor and AioPolarsCursor: the fetch methods, iter_chunks(), and as_pandas() / as_polars() / as_arrow(). AthenaPandasResultSet and AthenaPolarsResultSet both read chunked UNLOAD results through ParquetChunkReader.

Reproduction

No AWS calls. A stub filesystem retries a throttled request with retry_api_call(), and the read runs in a thread inside a stop scope, as the aio cursors run it:

import threading
import time

from botocore.exceptions import ClientError

from pyathena.arrow.reader import ParquetChunkReader
from pyathena.util import RetryConfig, retry_api_call, stop_retries_on


class FS:
    def __init__(self):
        self.calls = 0
        self.started = threading.Event()

    def size(self, path):
        def throttled():
            self.calls += 1
            self.started.set()
            raise ClientError(
                {"Error": {"Code": "ThrottlingException", "Message": "slow down"}}, "GetObject"
            )

        return retry_api_call(throttled, RetryConfig(attempt=4, multiplier=1, max_delay=1))


def run(label, read):
    fs, stop = FS(), threading.Event()

    def target():
        with stop_retries_on(stop):
            try:
                read(fs)
            except ClientError:
                pass

    thread = threading.Thread(target=target)
    thread.start()
    fs.started.wait(10)
    stopped = time.monotonic()
    stop.set()
    thread.join(30)
    print(f"{label}: {fs.calls} attempts, read ended {time.monotonic() - stopped:.1f}s after the stop")


run("size() in the reading thread", lambda fs: fs.size("s3://bucket/a.parquet"))
run(
    "ParquetChunkReader",
    lambda fs: next(ParquetChunkReader(fs, ["s3://bucket/a.parquet"], chunksize=10)),
)

Output on master at 44a50dc:

size() in the reading thread: 1 attempts, read ended 0.0s after the stop
ParquetChunkReader: 4 attempts, read ended 3.9s after the stop

With the default RetryConfig, the wait after a cancellation can reach about 127 s plus jitter, as described in #1205 (R1).

Environment

• PyAthena master at 44a50dc
• Python 3.15.0

Proposed fix (optional)

Submit _open() in a copy of the caller's context in ParquetChunkReader._open_files(), as S3ThreadPoolExecutor.submit() does: executor.submit(contextvars.copy_context().run, self._open, path). The file opened ahead is then also opened in the scope of the read that requested it. A regression test can use the stub above.

Found in the independent review of #1286. Part of #1205 (model 1, cooperative stop).

Scope extended (2026-10-19)

Copying the context alone (PR #1297, first commit) is not enough. The open of file k+1 is submitted in the read that obtains file k, but the read that waits for it in ahead.result() is often a later one. The aio cursors run each fetch, and each step of iter_chunks(), in run_in_thread_or_close() with a new stop event, and retry_api_call() takes its stop events when it starts. A cancellation of the later read therefore does not stop the open that it waits for. A stub with two scopes, where the second read is stopped while it waits for the file opened ahead, prints 4 attempts, read 2 ended 4.1s after its stop with the context copy applied.

_open_files() has the same wait when it finishes: ahead.exception() waits for an open still retrying, so ParquetChunkReader.close() between chunks, also from the sync cursors, waits for the whole retry policy. On the aio cancellation path this close runs under unstoppable_retries(), so stopping only the waiting read would not cover it.

Decisions (2026-10-19)

After consulting Codex and Claude Fable, which agreed:

1. Add a public pyathena.util.wait_with_retry_stop(future, *, stop). The work of future runs in stop_retries_on(stop). While the future is not done, the function sets stop once a stop event of the caller's current scopes is set, and then keeps waiting. It returns the result of the future or raises its exception. Without an enclosing stop scope it waits without polling; otherwise it polls with concurrent.futures.wait() at the existing interval, so that a TimeoutError of the work is raised as such.
2. ParquetChunkReader creates one threading.Event per file it opens, and opens the file in a copy of the requesting read's context inside stop_retries_on() of that event. The read that receives the file waits with wait_with_retry_stop(). A per-file event keeps a stopped open from stopping the later opens.
3. When _open_files() finishes with an open not yet received, it sets the event of the open before waiting for it, so that close() and a failed or stopped read do not wait for its retries. A file opened anyway is closed as before.
4. Tests without AWS: a stop of a later read that waits for the file opened ahead; close() while that open retries; and the contract of wait_with_retry_stop() (no scope, nested scopes, unstoppable_retries(), the work finishing first, a TimeoutError of the work).

Behavior change for the release notes: closing a chunked UNLOAD reader, or its result set, while the next file is still being opened ahead no longer waits for the retries of that open.