Skip to content

Stop a closed PolarsDataFrameIterator for every reader kind - #1029

Merged
laughingman7743 merged 1 commit into
masterfrom
fix/1019-polars-iterator-close
Oct 3, 2026
Merged

laughingman7743 merged 1 commit into
masterfrom
fix/1019-polars-iterator-close

Conversation

@laughingman7743

Copy link
Copy Markdown
Member

WHAT

PolarsDataFrameIterator.close() replaces its reader with an empty iterator before closing a generator reader, as PandasDataFrameIterator.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 by iter_chunks(), yields nothing for every result kind.

WHY

Closes #1019.

close() only handled GeneratorType readers, so iter([df]) kept yielding after close().

TEST

Tested commit: c7846e6

  • just lint: passed.
  • New offline test 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, the dataframe case fails and the generator case 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.
  • Not run locally: the full just test pyathena suite; it runs in CI when the PR is marked Ready.

🤖 Generated with Claude Code

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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 is iter(()), so iteration ends and the iter([df]) list iterator (with the DataFrame) is released.
  • Generator: the generator is still closed (GeneratorExit at its yield; _iter_*_chunks catches only Exception), and the reader is iter(()) instead of the closed generator. Both end iteration the same way.
  • A repeated close() and close() from __next__ after exhaustion are no-ops.
  • The new offline test fails on the base for the dataframe reader. The generator case guards the existing behavior.

if isinstance(self._reader, GeneratorType):
self._reader.close()
reader = self._reader
self._reader = iter(())

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 assigns self._reader = iter(()) before closing a TextFileReader.
  • "whole DataFrame without chunksize, with managed query result storage, or for failed or empty results": AthenaPolarsResultSet.__init__ wraps self._df in 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) calls self._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(())

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 after close(); a new iterrows() 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

PolarsDataFrameIterator.close() does not stop a whole-DataFrame iterator

1 participant