Defer closing a chunk reader while another thread reads it

[laughingman7743]

WHAT

PandasDataFrameIterator and PolarsDataFrameIterator no longer close their chunk reader while another thread is reading a chunk from it.

• Each iterator counts the reads in progress under a lock (_read()), covering __next__() and, for pandas, get_chunk().
• A close() during a read records the request and returns at once. The read returns the chunk it was reading, and the reading thread then closes the reader, and for pandas the CSV stream, when its read completes.
• A read that starts after such a close gets no chunks, as after an immediate close.
• An error from the deferred close is logged, because no caller waits for it, instead of replacing the chunk that the reading thread has read.
• A close with no read in progress behaves as before.

WHY

Fixes #1172.

Closing a PandasCursor or PolarsCursor result set while another thread read the next chunk from a generator reader (chunked UNLOAD for PandasCursor; chunked CSV and UNLOAD for PolarsCursor) raised ValueError: generator already executing.
AthenaPandasResultSet.close() and AthenaPolarsResultSet.close() stopped at that point, after the base class had cleared the query execution, and the error propagated out of the cursor's close() or __aexit__().
AioPandasCursor and AioPolarsCursor do this themselves: a fetch runs in asyncio.to_thread() and keeps reading after its task is cancelled, while leaving async with closes the cursor through close_in_thread().

For a chunked pandas CSV result, close previously closed the TextFileReader and the CSV stream under the running read; that close is now deferred the same way, so the read finishes its chunk before they are closed.

A lock held around each read, which would make close() wait for the read, was considered instead. It does not stop the reading thread from taking the lock again for the next chunk before the closer, and an interrupted wait would leave the reader unclosed.

TEST

Tested commit: 0c5268f (rebased onto 38dc8f8).

• just lint: passed.
• uv run --env-file .env pytest -q -p no:cacheprovider tests/pyathena/pandas/test_result_set.py tests/pyathena/polars/test_result_set.py: 295 passed, 10 skipped.
• New tests:
  • TestPandasDataFrameIterator.test_close_during_read (__next__() and get_chunk()) and TestPolarsDataFrameIterator.test_close_during_read: a close during a read in a worker thread does not raise; a read started after it gets no chunk; the generator and the CSV stream stay open until the read completes; the read returns its chunk; the generator is closed afterwards; an error from the deferred close is logged.
  • test_close_during_fetchall_generator in both result set test classes: close() during a fetchall() from a generator does not raise, the fetch returns the rows of the chunk it was reading, and the generator is closed.
  • With the source changes reverted, all of them fail with ValueError: generator already executing.
• AWS: uv run --env-file .env pytest -q -p no:cacheprovider -n 4 tests/pyathena/pandas tests/pyathena/polars tests/pyathena/aio/pandas tests/pyathena/aio/polars at 0c5268f: 778 passed, 10 skipped.
• Follow-up: Closing a PolarsCursor result set during the first chunked CSV read makes that read raise ProgrammingError #1188 (a Polars chunked CSV read that is still in its first next() when the result set closes raises ProgrammingError; pre-existing, found by the independent review).
• Not covered: the aio cancellation path is traced in code (pyathena/aio/common.py fetch and __aexit__), not reproduced against Athena. The deferral for a chunked CSV TextFileReader has no deterministic test, because the reader buffers 256 KiB per read; it uses the same _read() path that the tests cover. The rest of just test pyathena and the SQLAlchemy suites run in the AWS CI when the PR is marked Ready.

🤖 Generated with Claude Code

[laughingman7743]

Self-review round one: behavior and implementation (FINDINGS, repaired)

• Initial pass: 301ae52f6ff9be1db33be61f16282a8edb3066ff..6c60ea60a67d1213bee3ffdb4805c87eebd66909, then a rebase onto 38dc8f860ece434b36779f0da0291d505f4e42a1 (Return PolarsCursor rows as tuples without a dictionary per row #1177 import conflicts only). Current head: 0c5268f1.
• Covered: PandasDataFrameIterator and PolarsDataFrameIterator read and close paths (__next__, get_chunk, itertuples/iterrows/as_pandas/as_polars through __next__); both result sets' close(); the aio fetch (asyncio.to_thread) and close (close_in_thread) paths in pyathena/aio/common.py; the new tests.
• Method: a manual pass, /code-review (high), and /simplify (reuse, simplification, efficiency, altitude). A Codex read-only consultation covered the design choice between a deferred close and a lock held around each read.
• Findings and repairs are in the threads on the related lines. Every self._reader access now goes through _read() or close(). _read() runs once per chunk, not per row: about 1 µs plus two uncontended lock round trips.

[laughingman7743]

Self-review round one finding (repaired in b8be349): in the initial head, a deferred close() left self._reader in place. A read that started after close() had returned still received the live generator, and raised ValueError: generator already executing while another thread was reading it. For a TextFileReader, it would read the same parser concurrently. _read() now yields iter(()) once a close has been requested; the same change is in pyathena/polars/result_set.py:116. The tests now call next() during the deferred close and expect StopIteration.

[laughingman7743]

Self-review round one finding (repaired in b8be349): the deferred close runs in the reading thread's finally. An error from reader.close() or csv_stream.close() replaced the chunk that the read had already returned, and no caller of close() waited for that error. The error is now logged; the same change is in pyathena/polars/result_set.py:127. test_close_during_read[...-close_error1] covers it.

[laughingman7743]

Self-review round one reasoned deferrals

• The count is incremented before the try in _read(). An asynchronous exception in that window leaves the count stuck. The window is a lock release and the try setup, as in any with lock: followed by try:, so no change.
• For a chunked CSV TextFileReader, close now waits for the running chunk instead of closing the stream under it. pandas' C tokenizer can run without the GIL, so closing it concurrently is unsafe. The cost is bounded by one chunk.
• AthenaPolarsResultSet.close() clears the query execution before the iterator is closed. That ordering is the same on 301ae52f, so a generator still in its first read can observe it either way. It is pre-existing and not changed here.
• A lock held around each read, with a blocking close(), was rejected. Python locks are not fair, so fetchall() can take the lock again for the next chunk before the closer. An interrupted wait in close() would also leave the reader unclosed after the result set is already marked closed (pyathena/result_set.py:1264).
• The _read()/close() protocol is duplicated in the two iterator classes. Sharing it needs a new shared helper or base, which is a design change for the maintainer to approve.

[laughingman7743]

Self-review round one test findings (repaired)

• f966192: the method under test is passed as a parameter instead of being looked up by name.
• 3daaf64, 0c5268f: the tests check the generator state (GEN_RUNNING during the deferred close, GEN_CLOSED after it) instead of a separate Event. New test_close_during_fetchall_generator tests in both result set classes drive the Closing a PandasCursor or PolarsCursor result set while another thread reads the next chunk raises ValueError #1172 path through result_set.close() during fetchall().
• 0c5268f: a comment in the Polars test_close_during_fetchall, which said a generator cannot be closed while it runs (Closing a PandasCursor or PolarsCursor result set while another thread reads the next chunk raises ValueError #1172), is removed.
• With the source changes reverted, every new test fails with ValueError: generator already executing.
• Not covered: TextFileReader deferral has no deterministic test, because the reader buffers 256 KiB per read.

[laughingman7743]

Self-review round two: claims, callers, and evidence (FINDINGS, repaired in the PR description)

Full pass over 38dc8f860ece434b36779f0da0291d505f4e42a1..0c5268f1bcb72c11c02854c73935f13fad7e4b13.

Claims checked:

• Generator readers. Pandas chunked UNLOAD uses _iter_parquet_chunks(). Polars chunked CSV and UNLOAD use _iter_csv_chunks() and _iter_parquet_chunks(), selected in _create_dataframe_iterator(). Pandas chunked CSV uses a TextFileReader with the CSV stream.
• The ValueError and partial close. Reproduced on 301ae52f with the issue's script. Both result sets' close() call super().close() before self._df_iter.close().
• The aio path. Traced in code, not reproduced against Athena. Fetches use asyncio.to_thread(result_set.fetchone/fetchmany/fetchall) (pyathena/aio/common.py), and __aexit__ uses close_in_thread(self.close).
• The rejected lock design. Cursor close() skips a result set that is already closed (pyathena/result_set.py:1264). fetchall() takes rows continuously through islice(self._row_iter, size). Python documents no fairness for threading.Lock.
• The docstrings. They state the deferred close, empty reads after it, and the logged error. The tests assert each of these.

Callers:

• With no read in progress, close() behaves as before.
• Readers that are not generators (iter([df])) are unaffected.
• The docs have no statement on closing during a read (searched docs/).

Finding: the initial PR description did not mention that reads after a deferred close get no chunks, or that deferred close errors are logged. It also reported tests from before the round-one repairs. The description is updated.

[laughingman7743]

Self-review round two evidence

• Local unit tests at 0c5268f1bcb72c11c02854c73935f13fad7e4b13: tests/pyathena/pandas/test_result_set.py tests/pyathena/polars/test_result_set.py: 295 passed, 10 skipped.
• The new tests fail on the reverted source with the issue's ValueError.
• The AWS run of tests/pyathena/pandas tests/pyathena/polars tests/pyathena/aio/pandas tests/pyathena/aio/polars at the pre-rebase head 6c60ea60 passed: 728 passed, 10 skipped. The run at the current head, serialized with -n 1 because another session's AWS run overlaps, will be recorded in this thread.
• Static only: aio cancellation against Athena, and TextFileReader deferral.

[laughingman7743]

AWS run at 0c5268f1: uv run --env-file .env pytest -q -p no:cacheprovider -n 4 tests/pyathena/pandas tests/pyathena/polars tests/pyathena/aio/pandas tests/pyathena/aio/polars: 778 passed, 10 skipped. The serialized -n 1 run was stopped at 18% after the overlapping run ended, and replaced by this run.

[laughingman7743]

Independent review (relayed result from Codex)

• Reviewer: Codex CLI (codex exec -s read-only), a different model from the author. Session 01a13c89-9d67-72d0-a83b-80aea2fb93d5. Static review only: no tests or builds were run.
• Scope: 38dc8f860ece434b36779f0da0291d505f4e42a1..0c5268f1bcb72c11c02854c73935f13fad7e4b13. The source was a detached snapshot at the head. The prompt contained no PR number, description, commit messages, or prior findings. The snapshot and the PR worktree were unchanged afterwards.
• Covered: all four changed files, read/close interleavings including free-threaded Python, CSV and UNLOAD resource ownership, exception paths, sync and asyncio callers, fetch semantics, compatibility, tests, and conventions.
• Verdict: FINDINGS. There is one pre-existing gap and no regression introduced by the diff.
  • P2, pre-existing (pyathena/polars/result_set.py:1064, :933): AthenaPolarsResultSet.close() clears the query execution in super().close(). If the first chunk read is still in _is_csv_readable()'s head_object when the cursor closes, the generator then sees output_location is None and raises ProgrammingError instead of returning its chunk. The race exists at the merge-base.
  • The reviewer found that the tests catch the original failure and assert the returned data, the eventual close, exhaustion afterwards, and the logged close error. Their synthetic generators bypass the Polars generator's own preamble, and no test exercises an active pandas TextFileReader.
• Author verification: confirmed in the code. _iter_csv_chunks() checks self.output_location after _is_csv_readable(), and close() calls super().close() first; the same ordering is on 301ae52f. Scope decision pending with the maintainer.

[laughingman7743]

Scope decision (maintainer): the pre-existing P2 is split into #1188 and not fixed in this PR. The test-coverage limits (synthetic generators; no active TextFileReader test) stay as recorded in the PR description. No change to the reviewed diff, so no follow-up review is needed.