Stop a closed PolarsDataFrameIterator for every reader kind - #1029
Conversation
PolarsDataFrameIterator.close() only closed a generator reader, so an iterator over a whole DataFrame still yielded it after close() and kept the reference. close() now replaces the reader with an empty iterator first, as PandasDataFrameIterator.close() does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| if isinstance(self._reader, GeneratorType): | ||
| self._reader.close() | ||
| reader = self._reader | ||
| self._reader = iter(()) |
There was a problem hiding this comment.
Self-review round one (implementation behavior): CLEAN
Base de8cc52ac2a43cdba72bb4571c185883287741fc, head c7846e619459926f07cb5db6c1a342658bea2f54.
Covered: PolarsDataFrameIterator.close(), __next__() (which calls close() on StopIteration), __exit__, iterrows(), and the callers AthenaPolarsResultSet.close() and iter_chunks() (whole-DataFrame and chunk-generator readers).
- Whole DataFrame: after
close(), the reader isiter(()), so iteration ends and theiter([df])list iterator (with the DataFrame) is released. - Generator: the generator is still closed (
GeneratorExitat itsyield;_iter_*_chunkscatches onlyException), and the reader isiter(())instead of the closed generator. Both end iteration the same way. - A repeated
close()andclose()from__next__after exhaustion are no-ops. - The new offline test fails on the base for the
dataframereader. Thegeneratorcase guards the existing behavior.
| if isinstance(self._reader, GeneratorType): | ||
| self._reader.close() | ||
| reader = self._reader | ||
| self._reader = iter(()) |
There was a problem hiding this comment.
Self-review round two (claims, callers, operations): CLEAN
Base de8cc52ac2a43cdba72bb4571c185883287741fc, head c7846e619459926f07cb5db6c1a342658bea2f54. Full pass over the PR body, the commit message, and #1019.
- "as
PandasDataFrameIterator.close()does": it assignsself._reader = iter(())before closing aTextFileReader. - "whole DataFrame without
chunksize, with managed query result storage, or for failed or empty results":AthenaPolarsResultSet.__init__wrapsself._dfin those cases (since Keep the whole pandas and Polars results reusable #1005). - Callers: no signature change.
AthenaPolarsResultSet.close()(Keep the whole pandas and Polars results reusable #1005) callsself._df_iter.close()and then replaces the iterator, so its behavior does not change. No AWS requests are involved.
| if isinstance(self._reader, GeneratorType): | ||
| self._reader.close() | ||
| reader = self._reader | ||
| self._reader = iter(()) |
There was a problem hiding this comment.
Independent review (relayed): CLEAN
Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort high, sandbox read-only, session 01a1020d-2122-7b00-887a-6cae10cf0584. Static review only. Snapshot: detached worktree at head c7846e619459926f07cb5db6c1a342658bea2f54, base de8cc52ac2a43cdba72bb4571c185883287741fc. The prompt contained the literal diff and the intended behavior only. Afterwards, the snapshot and the PR worktree were clean.
Covered (reviewer): both reader kinds; __next__, repeated close(), __exit__, and iterrows(); result-set close() and iter_chunks(); the sync, async, and aio Polars cursors; CSV and Parquet generator cleanup (GeneratorExit bypasses except Exception); pandas parity; and the new test, whose dataframe case fails on the base. No regressions.
Pre-existing boundaries noted by the reviewer, not changed:
- A partly consumed
iterrows()generator can still yield the rest of its current frame afterclose(); a newiterrows()yields nothing. - A whole-DataFrame iterator returned by
iter_chunks()is independent, so closing the result set does not close it (the Keep the whole pandas and Polars results reusable #1005 design). Closing it explicitly now works.
WHAT
PolarsDataFrameIterator.close()replaces its reader with an empty iterator before closing a generator reader, asPandasDataFrameIterator.close()does.A closed iterator over a whole DataFrame (without
chunksize, with managed query result storage, or for failed or empty results) now yields nothing and no longer holds the DataFrame; closing a chunk generator behaves as before.Release note (4.0.0, fix): a closed
PolarsDataFrameIterator, such as one returned byiter_chunks(), yields nothing for every result kind.WHY
Closes #1019.
close()only handledGeneratorTypereaders, soiter([df])kept yielding afterclose().TEST
Tested commit: c7846e6
just lint: passed.TestPolarsDataFrameIterator::test_close_stops_iteration(whole DataFrame and generator readers):uv run --env-file <dummy env> pytest --noconftest tests/pyathena/polars/test_result_set.py: 4 passed. With the source change reverted, thedataframecase fails and thegeneratorcase passes.uv run --env-file .env pytest -n 1 tests/pyathena/polars/test_cursor.py -k "close or iter_chunks or whole_result": 9 passed.just test pyathenasuite; it runs in CI when the PR is marked Ready.🤖 Generated with Claude Code