Defer closing a chunk reader while another thread reads it

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

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

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,)]

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) == []