diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index b13b51ed3..e0d5af173 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,31 @@ 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 = iter(()) if self._close_requested else self._reader + try: + yield reader + finally: + with self._lock: + self._readers -= 1 + close = self._close_requested and not self._readers + if close: + try: + self.close() + except Exception: + _logger.exception("Failed to close the chunk reader after a read.") @override def __next__(self) -> DataFrame: @@ -348,7 +374,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 +395,27 @@ 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. Reads that start after the + close get no chunks, and an error from the deferred close is logged. + """ 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 +502,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..2bc77b264 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,31 @@ 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 = iter(()) if self._close_requested else self._reader + try: + yield reader + finally: + with self._lock: + self._readers -= 1 + close = self._close_requested and not self._readers + if close: + try: + self.close() + except Exception: + _logger.exception("Failed to close the chunk reader after a read.") @override def __next__(self) -> pl.DataFrame: @@ -110,7 +137,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 +157,18 @@ 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. Reads that start after the + close get no chunks, and an error from the deferred close is logged. + """ + 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..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 @@ -140,6 +141,46 @@ def test_as_pandas_single_dataframe(self): assert df_iter.as_pandas() is df + @pytest.mark.parametrize( + "read", + [PandasDataFrameIterator.__next__, PandasDataFrameIterator.get_chunk], + ids=["next", "get_chunk"], + ) + @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 = 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: + if close_error: + raise close_error + + csv_stream = io.StringIO() + 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) + assert reading.wait(10) + df_iter.close() + with pytest.raises(StopIteration): + next(df_iter) + assert inspect.getgeneratorstate(reader) == inspect.GEN_RUNNING + assert not csv_stream.closed + release.set() + assert future.result(10)["n"].tolist() == [2] + 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) == [] + _FS = sentinel.pyathena_fs _USER_FS = sentinel.user_fs @@ -281,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 72f8499d0..7894d615e 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 @@ -161,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]})]) @@ -187,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) @@ -629,3 +650,35 @@ def test_close_stops_iteration(self, reader): df_iter = PolarsDataFrameIterator(reader, {}, ["a"]) df_iter.close() assert list(df_iter) == [] + + @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 = 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: + if close_error: + raise close_error + + reader = chunks() + df_iter = PolarsDataFrameIterator(reader, {}, ["a"]) + next(df_iter) + with ThreadPoolExecutor(1) as executor: + future = executor.submit(next, df_iter) + assert reading.wait(10) + df_iter.close() + with pytest.raises(StopIteration): + next(df_iter) + assert inspect.getgeneratorstate(reader) == inspect.GEN_RUNNING + release.set() + assert future.result(10)["a"].to_list() == [2] + assert inspect.getgeneratorstate(reader) == inspect.GEN_CLOSED + assert ("Failed to close the chunk reader" in caplog.text) == bool(close_error) + assert list(df_iter) == []