From 9526552cb92e210137fba80959c5c12463b91281 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 00:11:28 +0900 Subject: [PATCH 1/5] Defer closing a chunk reader while another thread reads it PandasDataFrameIterator and PolarsDataFrameIterator closed their chunk reader directly, so closing a cursor while another thread read the next chunk raised "ValueError: generator already executing" for a generator reader and left the result set partly closed. AioPandasCursor and AioPolarsCursor do this themselves when a cancelled fetch keeps reading in its thread and the cursor is closed. The iterators now count the reads in progress. A close during a read records the request, and the reading thread closes the reader, and the CSV stream for pandas, when its read completes. Fixes #1172 Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 53 +++++++++++++++++++----- pyathena/polars/result_set.py | 41 ++++++++++++++++-- tests/pyathena/pandas/test_result_set.py | 30 ++++++++++++++ tests/pyathena/polars/test_result_set.py | 26 ++++++++++++ 4 files changed, 136 insertions(+), 14 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index b13b51ed3..f54f8ed64 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -7,10 +7,11 @@ import logging import operator import sys +import threading from collections import abc from collections.abc import Callable, Iterable, Iterator from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from functools import partial from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper from itertools import chain, islice @@ -336,6 +337,28 @@ def __init__( self._reader = reader self._trunc_date = trunc_date self._csv_stream = csv_stream + self._lock = threading.Lock() + self._readers = 0 + self._close_requested = False + + @contextmanager + def _read(self) -> Iterator[TextFileReader | Iterator[DataFrame]]: + """Track a chunk read and run a close that was requested during it. + + Yields: + The reader to read the chunk from. + """ + with self._lock: + self._readers += 1 + reader = self._reader + try: + yield reader + finally: + with self._lock: + self._readers -= 1 + close = self._close_requested and not self._readers + if close: + self.close() @override def __next__(self) -> DataFrame: @@ -348,7 +371,8 @@ def __next__(self) -> DataFrame: StopIteration: When no more chunks are available. """ try: - df = next(self._reader) + with self._read() as reader: + df = next(reader) return self._trunc_date(df) except BaseException: self.close() @@ -368,17 +392,26 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: self.close() def close(self) -> None: - """Close the iterator and release resources.""" + """Close the iterator and release resources. + + If another thread is reading a chunk, the read returns that chunk, and the + reader and the CSV stream are closed when it completes. + """ from pandas.io.parsers import TextFileReader - reader = self._reader - self._reader = iter(()) + with self._lock: + if self._readers: + self._close_requested = True + return + self._close_requested = False + reader, self._reader = self._reader, iter(()) + csv_stream, self._csv_stream = self._csv_stream, None try: if isinstance(reader, (TextFileReader, GeneratorType)): reader.close() finally: - if self._csv_stream is not None: - self._csv_stream.close() + if csv_stream is not None: + csv_stream.close() def iterrows(self) -> Iterator[tuple[int, dict[str, Any]]]: """Iterate over rows as (index, row_dict) tuples. @@ -465,9 +498,9 @@ def get_chunk(self, size: int | None = None) -> DataFrame: from pandas.io.parsers import TextFileReader try: - if isinstance(self._reader, TextFileReader): - return self._trunc_date(self._reader.get_chunk(size)) - return self._trunc_date(next(self._reader)) + with self._read() as reader: + df = reader.get_chunk(size) if isinstance(reader, TextFileReader) else next(reader) + return self._trunc_date(df) except BaseException: self.close() raise diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index da67e95ba..acd9355a4 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -14,8 +14,10 @@ import logging import operator import sys +import threading from collections import abc from collections.abc import Callable, Iterable, Iterator +from contextlib import contextmanager from io import BytesIO, StringIO from itertools import islice, repeat from multiprocessing import cpu_count @@ -98,6 +100,28 @@ def __init__( self._reader = reader self._converters = converters self._column_names = column_names + self._lock = threading.Lock() + self._readers = 0 + self._close_requested = False + + @contextmanager + def _read(self) -> Iterator[Iterator[pl.DataFrame]]: + """Track a chunk read and run a close that was requested during it. + + Yields: + The reader to read the chunk from. + """ + with self._lock: + self._readers += 1 + reader = self._reader + try: + yield reader + finally: + with self._lock: + self._readers -= 1 + close = self._close_requested and not self._readers + if close: + self.close() @override def __next__(self) -> pl.DataFrame: @@ -110,7 +134,8 @@ def __next__(self) -> pl.DataFrame: StopIteration: When no more chunks are available. """ try: - return next(self._reader) + with self._read() as reader: + return next(reader) except StopIteration: self.close() raise @@ -129,9 +154,17 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: self.close() def close(self) -> None: - """Close the iterator and release resources.""" - reader = self._reader - self._reader = iter(()) + """Close the iterator and release resources. + + If another thread is reading a chunk, the read returns that chunk, and the + reader is closed when it completes. + """ + with self._lock: + if self._readers: + self._close_requested = True + return + self._close_requested = False + reader, self._reader = self._reader, iter(()) if isinstance(reader, GeneratorType): reader.close() diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index a4fc2421d..83ae82589 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -140,6 +140,36 @@ def test_as_pandas_single_dataframe(self): assert df_iter.as_pandas() is df + @pytest.mark.parametrize("read", ["next", "get_chunk"]) + def test_close_during_read(self, read): + """A close during a read in another thread is deferred until the read completes.""" + reading, release, released = threading.Event(), threading.Event(), threading.Event() + + def chunks(): + try: + yield pd.DataFrame({"n": [1]}) + reading.set() + release.wait(10) + yield pd.DataFrame({"n": [2]}) + yield pd.DataFrame({"n": [3]}) + finally: + released.set() + + csv_stream = io.StringIO() + df_iter = PandasDataFrameIterator(chunks(), _no_trunc_date, csv_stream) + next(df_iter) + with ThreadPoolExecutor(1) as executor: + future = executor.submit(getattr(df_iter, "__next__" if read == "next" else read)) + assert reading.wait(10) + df_iter.close() + assert not released.is_set() + assert not csv_stream.closed + release.set() + assert future.result(10)["n"].tolist() == [2] + assert released.is_set() + assert csv_stream.closed + assert list(df_iter) == [] + _FS = sentinel.pyathena_fs _USER_FS = sentinel.user_fs diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 72f8499d0..728b1e7f8 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -629,3 +629,29 @@ def test_close_stops_iteration(self, reader): df_iter = PolarsDataFrameIterator(reader, {}, ["a"]) df_iter.close() assert list(df_iter) == [] + + def test_close_during_read(self): + """A close during a read in another thread is deferred until the read completes.""" + reading, release, released = threading.Event(), threading.Event(), threading.Event() + + def chunks(): + try: + yield pl.DataFrame({"a": [1]}) + reading.set() + release.wait(10) + yield pl.DataFrame({"a": [2]}) + yield pl.DataFrame({"a": [3]}) + finally: + released.set() + + df_iter = PolarsDataFrameIterator(chunks(), {}, ["a"]) + next(df_iter) + with ThreadPoolExecutor(1) as executor: + future = executor.submit(next, df_iter) + assert reading.wait(10) + df_iter.close() + assert not released.is_set() + release.set() + assert future.result(10)["a"].to_list() == [2] + assert released.is_set() + assert list(df_iter) == [] From f966192c4323e69fe6b44ef0cb5688174e1235f9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 00:21:36 +0900 Subject: [PATCH 2/5] Pass the read method to the pandas close-during-read test directly Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_result_set.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 83ae82589..d64487586 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -140,7 +140,11 @@ def test_as_pandas_single_dataframe(self): assert df_iter.as_pandas() is df - @pytest.mark.parametrize("read", ["next", "get_chunk"]) + @pytest.mark.parametrize( + "read", + [PandasDataFrameIterator.__next__, PandasDataFrameIterator.get_chunk], + ids=["next", "get_chunk"], + ) def test_close_during_read(self, read): """A close during a read in another thread is deferred until the read completes.""" reading, release, released = threading.Event(), threading.Event(), threading.Event() @@ -159,7 +163,7 @@ def chunks(): df_iter = PandasDataFrameIterator(chunks(), _no_trunc_date, csv_stream) next(df_iter) with ThreadPoolExecutor(1) as executor: - future = executor.submit(getattr(df_iter, "__next__" if read == "next" else read)) + future = executor.submit(read, df_iter) assert reading.wait(10) df_iter.close() assert not released.is_set() From b8be349fb6c8a0754b740fe6845510de146ffc54 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 00:28:00 +0900 Subject: [PATCH 3/5] Give no chunks after a deferred close and log its errors A read that starts after a close was deferred now gets an empty reader instead of the reader that another thread is still reading. An error from the deferred close, which no caller waits for, is logged instead of replacing the chunk that the reading thread has read. Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 10 +++++++--- pyathena/polars/result_set.py | 10 +++++++--- tests/pyathena/pandas/test_result_set.py | 8 +++++++- tests/pyathena/polars/test_result_set.py | 8 +++++++- 4 files changed, 28 insertions(+), 8 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index f54f8ed64..e0d5af173 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -350,7 +350,7 @@ def _read(self) -> Iterator[TextFileReader | Iterator[DataFrame]]: """ with self._lock: self._readers += 1 - reader = self._reader + reader = iter(()) if self._close_requested else self._reader try: yield reader finally: @@ -358,7 +358,10 @@ def _read(self) -> Iterator[TextFileReader | Iterator[DataFrame]]: self._readers -= 1 close = self._close_requested and not self._readers if close: - self.close() + try: + self.close() + except Exception: + _logger.exception("Failed to close the chunk reader after a read.") @override def __next__(self) -> DataFrame: @@ -395,7 +398,8 @@ def close(self) -> None: """Close the iterator and release resources. If another thread is reading a chunk, the read returns that chunk, and the - reader and the CSV stream are closed when it completes. + reader and the CSV stream are closed when it completes. Reads that start after the + close get no chunks, and an error from the deferred close is logged. """ from pandas.io.parsers import TextFileReader diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index acd9355a4..2bc77b264 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -113,7 +113,7 @@ def _read(self) -> Iterator[Iterator[pl.DataFrame]]: """ with self._lock: self._readers += 1 - reader = self._reader + reader = iter(()) if self._close_requested else self._reader try: yield reader finally: @@ -121,7 +121,10 @@ def _read(self) -> Iterator[Iterator[pl.DataFrame]]: self._readers -= 1 close = self._close_requested and not self._readers if close: - self.close() + try: + self.close() + except Exception: + _logger.exception("Failed to close the chunk reader after a read.") @override def __next__(self) -> pl.DataFrame: @@ -157,7 +160,8 @@ def close(self) -> None: """Close the iterator and release resources. If another thread is reading a chunk, the read returns that chunk, and the - reader is closed when it completes. + reader is closed when it completes. Reads that start after the + close get no chunks, and an error from the deferred close is logged. """ with self._lock: if self._readers: diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index d64487586..2d4cbfcb5 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -145,7 +145,8 @@ def test_as_pandas_single_dataframe(self): [PandasDataFrameIterator.__next__, PandasDataFrameIterator.get_chunk], ids=["next", "get_chunk"], ) - def test_close_during_read(self, read): + @pytest.mark.parametrize("close_error", [None, RuntimeError("close failed")]) + def test_close_during_read(self, caplog, read, close_error): """A close during a read in another thread is deferred until the read completes.""" reading, release, released = threading.Event(), threading.Event(), threading.Event() @@ -158,6 +159,8 @@ def chunks(): yield pd.DataFrame({"n": [3]}) finally: released.set() + if close_error: + raise close_error csv_stream = io.StringIO() df_iter = PandasDataFrameIterator(chunks(), _no_trunc_date, csv_stream) @@ -166,11 +169,14 @@ def chunks(): future = executor.submit(read, df_iter) assert reading.wait(10) df_iter.close() + with pytest.raises(StopIteration): + next(df_iter) assert not released.is_set() assert not csv_stream.closed release.set() assert future.result(10)["n"].tolist() == [2] assert released.is_set() + assert ("Failed to close the chunk reader" in caplog.text) == bool(close_error) assert csv_stream.closed assert list(df_iter) == [] diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 728b1e7f8..2c432ba2e 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -630,7 +630,8 @@ def test_close_stops_iteration(self, reader): df_iter.close() assert list(df_iter) == [] - def test_close_during_read(self): + @pytest.mark.parametrize("close_error", [None, RuntimeError("close failed")]) + def test_close_during_read(self, caplog, close_error): """A close during a read in another thread is deferred until the read completes.""" reading, release, released = threading.Event(), threading.Event(), threading.Event() @@ -643,6 +644,8 @@ def chunks(): yield pl.DataFrame({"a": [3]}) finally: released.set() + if close_error: + raise close_error df_iter = PolarsDataFrameIterator(chunks(), {}, ["a"]) next(df_iter) @@ -650,8 +653,11 @@ def chunks(): future = executor.submit(next, df_iter) assert reading.wait(10) df_iter.close() + with pytest.raises(StopIteration): + next(df_iter) assert not released.is_set() release.set() assert future.result(10)["a"].to_list() == [2] assert released.is_set() + assert ("Failed to close the chunk reader" in caplog.text) == bool(close_error) assert list(df_iter) == [] From 3daaf6475a17aab4c9d74d85c05295d5400fe163 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 00:30:31 +0900 Subject: [PATCH 4/5] Check the generator state in the close-during-read tests Also test that closing a pandas result set during a fetch from a generator reader does not raise and closes the generator once the fetch completes. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_result_set.py | 32 ++++++++++++++++++++---- tests/pyathena/polars/test_result_set.py | 11 ++++---- 2 files changed, 33 insertions(+), 10 deletions(-) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 2d4cbfcb5..61ae75d21 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT import csv +import inspect import io import threading from collections import defaultdict @@ -148,7 +149,7 @@ def test_as_pandas_single_dataframe(self): @pytest.mark.parametrize("close_error", [None, RuntimeError("close failed")]) def test_close_during_read(self, caplog, read, close_error): """A close during a read in another thread is deferred until the read completes.""" - reading, release, released = threading.Event(), threading.Event(), threading.Event() + reading, release = threading.Event(), threading.Event() def chunks(): try: @@ -158,12 +159,12 @@ def chunks(): yield pd.DataFrame({"n": [2]}) yield pd.DataFrame({"n": [3]}) finally: - released.set() if close_error: raise close_error csv_stream = io.StringIO() - df_iter = PandasDataFrameIterator(chunks(), _no_trunc_date, csv_stream) + reader = chunks() + df_iter = PandasDataFrameIterator(reader, _no_trunc_date, csv_stream) next(df_iter) with ThreadPoolExecutor(1) as executor: future = executor.submit(read, df_iter) @@ -171,11 +172,11 @@ def chunks(): df_iter.close() with pytest.raises(StopIteration): next(df_iter) - assert not released.is_set() + assert inspect.getgeneratorstate(reader) == inspect.GEN_RUNNING assert not csv_stream.closed release.set() assert future.result(10)["n"].tolist() == [2] - assert released.is_set() + assert inspect.getgeneratorstate(reader) == inspect.GEN_CLOSED assert ("Failed to close the chunk reader" in caplog.text) == bool(close_error) assert csv_stream.closed assert list(df_iter) == [] @@ -321,6 +322,27 @@ def __next__(self): assert future.result(10) == [(0,), (1,), (2,), (3,)] assert result_set.rownumber is None + def test_close_during_fetchall_generator(self): + """A close() during a fetch from a generator, such as a chunked UNLOAD, does not raise.""" + reading, release = threading.Event(), threading.Event() + + def chunks(): + yield pd.DataFrame({"n": [0, 1]}) + reading.set() + release.wait(10) + yield pd.DataFrame({"n": [2, 3]}, index=[2, 3]) + yield pd.DataFrame({"n": [4]}, index=[4]) + + reader = chunks() + result_set = _fetching_result_set(reader, arraysize=3) + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(result_set.fetchall) + assert reading.wait(10) + result_set.close() + release.set() + assert future.result(10) == [(0,), (1,), (2,), (3,)] + assert inspect.getgeneratorstate(reader) == inspect.GEN_CLOSED + def test_fetchmany_size_above_maxsize(self): result_set = _fetching_result_set(pd.DataFrame({"n": [1, 2]}), arraysize=3) assert result_set.fetchmany(10**100) == [(1,), (2,)] diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 2c432ba2e..c61f10872 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -5,6 +5,7 @@ # # SPDX-License-Identifier: MIT +import inspect import threading from concurrent.futures import ThreadPoolExecutor from datetime import UTC, date, datetime, time, timedelta @@ -633,7 +634,7 @@ def test_close_stops_iteration(self, reader): @pytest.mark.parametrize("close_error", [None, RuntimeError("close failed")]) def test_close_during_read(self, caplog, close_error): """A close during a read in another thread is deferred until the read completes.""" - reading, release, released = threading.Event(), threading.Event(), threading.Event() + reading, release = threading.Event(), threading.Event() def chunks(): try: @@ -643,11 +644,11 @@ def chunks(): yield pl.DataFrame({"a": [2]}) yield pl.DataFrame({"a": [3]}) finally: - released.set() if close_error: raise close_error - df_iter = PolarsDataFrameIterator(chunks(), {}, ["a"]) + reader = chunks() + df_iter = PolarsDataFrameIterator(reader, {}, ["a"]) next(df_iter) with ThreadPoolExecutor(1) as executor: future = executor.submit(next, df_iter) @@ -655,9 +656,9 @@ def chunks(): df_iter.close() with pytest.raises(StopIteration): next(df_iter) - assert not released.is_set() + assert inspect.getgeneratorstate(reader) == inspect.GEN_RUNNING release.set() assert future.result(10)["a"].to_list() == [2] - assert released.is_set() + assert inspect.getgeneratorstate(reader) == inspect.GEN_CLOSED assert ("Failed to close the chunk reader" in caplog.text) == bool(close_error) assert list(df_iter) == [] From 0c5268f1bcb72c11c02854c73935f13fad7e4b13 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 07:25:37 +0900 Subject: [PATCH 5/5] Test closing a Polars result set during a fetch from a generator Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/polars/test_result_set.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index c61f10872..7894d615e 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -162,7 +162,6 @@ def test_close_during_fetchall(self): requested = threading.Event() closed = threading.Event() - # Not a generator, which close() cannot close while it runs (#1172). class Chunks: def __init__(self): self._chunks = iter([pl.DataFrame({"n": [0, 1]}), pl.DataFrame({"n": [2, 3]})]) @@ -188,6 +187,27 @@ def __next__(self): assert future.result(10) == [(0,), (1,), (2,), (3,)] assert result_set.rownumber is None + def test_close_during_fetchall_generator(self): + """A close() during a fetch from a generator, such as a chunked UNLOAD, does not raise.""" + reading, release = threading.Event(), threading.Event() + + def chunks(): + yield pl.DataFrame({"n": [0, 1]}) + reading.set() + release.wait(10) + yield pl.DataFrame({"n": [2, 3]}) + yield pl.DataFrame({"n": [4]}) + + reader = chunks() + result_set = _fetching_result_set(reader, ["n"], arraysize=3) + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(result_set.fetchall) + assert reading.wait(10) + result_set.close() + release.set() + assert future.result(10) == [(0,), (1,), (2,), (3,)] + assert inspect.getgeneratorstate(reader) == inspect.GEN_CLOSED + def test_fetchmany_size_above_maxsize(self): """A size above sys.maxsize returns the remaining rows.""" result_set = _fetching_result_set([pl.DataFrame({"n": [1, 2]})], ["n"], arraysize=3)