Skip to content

Stop a closed PolarsDataFrameIterator for every reader kind #1029

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
laughingman7743 merged 1 commit into from
Oct 3, 2026
Merged

Stop a closed PolarsDataFrameIterator for every reader kind #1029

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions pyathena/polars/result_set.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,8 +121,10 @@ def close(self) -> None:
"""Close the iterator and release resources."""
from types import GeneratorType

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

@laughingman7743 laughingman7743 Oct 3, 2026

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.

@laughingman7743 laughingman7743 Oct 3, 2026

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.

@laughingman7743 laughingman7743 Oct 3, 2026

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.

if isinstance(reader, GeneratorType):
reader.close()

def iterrows(self) -> Iterator[tuple[int, dict[str, Any]]]:
"""Iterate over rows as (index, row_dict) tuples.
Expand Down
15 changes: 14 additions & 1 deletion tests/pyathena/polars/test_result_set.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
import pytest

from pyathena.error import OperationalError
from pyathena.polars.result_set import AthenaPolarsResultSet
from pyathena.polars.result_set import AthenaPolarsResultSet, PolarsDataFrameIterator

_ROWS_BEFORE_FAILURE = 300_000

Expand Down Expand Up @@ -80,3 +80,16 @@ def test_iter_parquet_chunks_raises_when_read_fails_partway(self, tmp_path):
pytest.raises(OperationalError),
):
list(result_set._iter_parquet_chunks())


class TestPolarsDataFrameIterator:
@pytest.mark.parametrize(
"reader",
[pl.DataFrame({"a": [1, 2]}), (df for df in [pl.DataFrame({"a": [1]})] * 2)],
ids=["dataframe", "generator"],
)
def test_close_stops_iteration(self, reader):
"""A closed iterator yields nothing for either reader kind."""
df_iter = PolarsDataFrameIterator(reader, {}, ["a"])
df_iter.close()
assert list(df_iter) == []
Loading