From c7846e619459926f07cb5db6c1a342658bea2f54 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 22:55:44 +0900 Subject: [PATCH] Stop a closed PolarsDataFrameIterator for every reader kind 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 --- pyathena/polars/result_set.py | 6 ++++-- tests/pyathena/polars/test_result_set.py | 15 ++++++++++++++- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 0563ec819..039850751 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -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(()) + if isinstance(reader, GeneratorType): + reader.close() def iterrows(self) -> Iterator[tuple[int, dict[str, Any]]]: """Iterate over rows as (index, row_dict) tuples. diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 74dd1cd2b..e6828b54b 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -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 @@ -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) == []