From 63646b5b78e2f7cea6565c7b5b49a923f614fd3c Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Fri, 16 Oct 2026 00:16:38 +0900
Subject: [PATCH 1/3] Refill the result set buffer through _fetch() in
 fetchone() and the bulk fetches

The base _fetch() and the asyncio _async_fetch() now return without rows
when no NextToken remains, so every result set's _fetch() means "read the
next rows, if any". fetchone() calls it whenever the buffer is empty, which
removes the identical fetchone() overrides of the Arrow and S3FS result
sets. _take_rows() and _async_take_rows() refill the buffer with it directly
and take every row in bulk instead of taking the first row of each page
through fetchone().

Closes #1193

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/aio/result_set.py              | 29 +++++++++----------
 pyathena/arrow/result_set.py            | 15 ----------
 pyathena/result_set.py                  | 32 +++++++++++----------
 pyathena/s3fs/result_set.py             | 20 -------------
 tests/pyathena/aio/test_result_set.py   | 37 +++++++++++++++++++++++--
 tests/pyathena/arrow/test_result_set.py | 29 +++++++++++++++++++
 tests/pyathena/s3fs/test_result_set.py  | 12 ++++++++
 tests/pyathena/test_result_set.py       | 37 ++++++++-----------------
 8 files changed, 119 insertions(+), 92 deletions(-)

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index b55627ab2..a26546d3d 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -143,14 +143,17 @@ async def _async_get_query_results(
             return cast(dict[str, Any], response)
 
     async def _async_fetch(self) -> None:
-        """Fetch the next page of rows into the result set.
+        """Fetch the next page of rows, if any, into the buffer.
+
+        Nothing is fetched when no page remains or the result set is closed.
 
         Raises:
-            ProgrammingError: If there is no next page.
+            ProgrammingError: If the query ID is missing, the query has not
+                succeeded, or another thread closes the result set.
             OperationalError: If the request fails.
         """
         if not self._next_token:
-            raise ProgrammingError("NextToken is none or empty.")
+            return
         response = await self._async_get_query_results(self._arraysize, self._next_token)
         rows, self._next_token = self._parse_result_rows(response)
         self._process_rows(rows)
@@ -183,10 +186,10 @@ async def fetchone(  # type: ignore[override]
             The next row (a tuple, or a dict for ``AthenaAioDictResultSet``), or
             None if no more rows.
         """
-        if not self._rows and self._next_token:
-            await self._async_fetch()
         if not self._rows:
-            return None
+            await self._async_fetch()
+            if not self._rows:
+                return None
         row = self._rows.popleft()
         # As _count_rows() does, inline on the per-row path.
         rownumber = self._rownumber
@@ -230,8 +233,8 @@ async def _async_take_rows(
     ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
         """Take the next rows and count them in ``rownumber``.
 
-        The buffered rows are taken in bulk. When the buffer is empty, ``fetchone()``
-        takes the next row, so that it fetches the next page into the buffer.
+        The buffered rows are taken in bulk, and ``_async_fetch()`` refills the buffer
+        when it is empty.
 
         Args:
             size: The maximum number of rows, or None for all remaining rows.
@@ -241,13 +244,11 @@ async def _async_take_rows(
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
         while size is None or len(rows) < size:
-            if self._rows:
-                self._take_buffered_rows(rows, size)
-            else:
-                row = await self.fetchone()
-                if row is None:
+            if not self._rows:
+                await self._async_fetch()
+                if not self._rows:
                     break
-                rows.append(row)
+            self._take_buffered_rows(rows, size)
         return rows
 
     @override
diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py
index 53939437e..508f644d7 100644
--- a/pyathena/arrow/result_set.py
+++ b/pyathena/arrow/result_set.py
@@ -216,21 +216,6 @@ def _fetch(self) -> None:
                 processed_rows = list(zip(*columns, strict=False))
             self._rows.extend(processed_rows)
 
-    @override
-    def fetchone(
-        self,
-    ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None:
-        if not self._rows:
-            self._fetch()
-        if not self._rows:
-            return None
-        row = self._rows.popleft()
-        # As _count_rows() does, inline on the per-row path.
-        rownumber = self._rownumber
-        if rownumber is not None:
-            self._rownumber = rownumber + 1
-        return row
-
     def _open_csv(self, location: str) -> Any:
         """Open the CSV result file for ``pyarrow.csv.read_csv()``.
 
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index a72271fab..5251664c3 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -459,14 +459,18 @@ def _get_query_results(self, max_results: int, next_token: str | None = None) ->
             return cast(dict[str, Any], response)
 
     def _fetch(self) -> None:
-        """Fetch the next page of rows into the result set.
+        """Fetch the next rows, if any, into the buffer.
+
+        The default result set fetches the next page; it fetches nothing when no
+        page remains or the result set is closed.
 
         Raises:
-            ProgrammingError: If there is no next page.
+            ProgrammingError: If the query ID is missing, the query has not
+                succeeded, or another thread closes the result set.
             OperationalError: If the request fails.
         """
         if not self._next_token:
-            raise ProgrammingError("NextToken is none or empty.")
+            return
         response = self._get_query_results(self._arraysize, self._next_token)
         rows, self._next_token = self._parse_result_rows(response)
         self._process_rows(rows)
@@ -490,10 +494,10 @@ def _pre_fetch(self) -> None:
     def fetchone(
         self,
     ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None:
-        if not self._rows and self._next_token:
-            self._fetch()
         if not self._rows:
-            return None
+            self._fetch()
+            if not self._rows:
+                return None
         row = self._rows.popleft()
         # As _count_rows() does, inline on the per-row path.
         rownumber = self._rownumber
@@ -536,9 +540,9 @@ def fetchall(
     def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
         """Take the next rows and count them in ``rownumber``.
 
-        The buffered rows are taken in bulk. When the buffer is empty, ``fetchone()``
-        takes the next row, so that it fetches the next rows into the buffer as the
-        result set does.
+        The buffered rows are taken in bulk, and ``_fetch()`` refills the buffer when
+        it is empty. A subclass whose ``fetchone()`` does not use the buffer overrides
+        this method too.
 
         Args:
             size: The maximum number of rows, or None for all remaining rows.
@@ -548,13 +552,11 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
         while size is None or len(rows) < size:
-            if self._rows:
-                self._take_buffered_rows(rows, size)
-            else:
-                row = self.fetchone()
-                if row is None:
+            if not self._rows:
+                self._fetch()
+                if not self._rows:
                     break
-                rows.append(row)
+            self._take_buffered_rows(rows, size)
         return rows
 
     def _take_buffered_rows(
diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py
index 1ae07684a..7569a92f3 100644
--- a/pyathena/s3fs/result_set.py
+++ b/pyathena/s3fs/result_set.py
@@ -253,26 +253,6 @@ def _fetch(self) -> None:
             self._rows.append(converted_row)
             rows_fetched += 1
 
-    @override
-    def fetchone(
-        self,
-    ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None:
-        """Fetch the next row of the result set.
-
-        Returns:
-            A tuple representing the next row, or None if no more rows.
-        """
-        if not self._rows:
-            self._fetch()
-        if not self._rows:
-            return None
-        row = self._rows.popleft()
-        # As _count_rows() does, inline on the per-row path.
-        rownumber = self._rownumber
-        if rownumber is not None:
-            self._rownumber = rownumber + 1
-        return row
-
     @override
     def close(self) -> None:
         """Close the result set and release resources."""
diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py
index 61c27e8f1..ba5983797 100644
--- a/tests/pyathena/aio/test_result_set.py
+++ b/tests/pyathena/aio/test_result_set.py
@@ -10,6 +10,7 @@
 
 from pyathena.aio.result_set import AthenaAioDictResultSet, AthenaAioResultSet
 from pyathena.converter import DefaultTypeConverter
+from pyathena.error import OperationalError
 from pyathena.model import AthenaQueryExecution
 from pyathena.util import RetryConfig
 from tests.pyathena.util import get_query_results_responses
@@ -20,8 +21,9 @@ async def _create_result_set(result_set_class, pages=(("1", "2"),)):
 
     Args:
         result_set_class: The asyncio result set class to create.
-        pages: The values of an integer column ``a``, one tuple per page. The first
-            page is fetched at creation, the others on demand.
+        pages: The values of an integer column ``a``, one tuple per page, or an
+            exception that the request for the page raises. The first page is fetched
+            at creation, the others on demand.
 
     Returns:
         The result set.
@@ -90,3 +92,34 @@ async def test_fetchmany_size_not_integer(self):
         with pytest.raises(TypeError):
             await result_set.fetchmany(1.5)
         assert await result_set.fetchall() == [(1,), (2,)]
+
+    async def test_fetchone_across_pages(self):
+        result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)])
+        assert [await result_set.fetchone() for _ in range(4)] == [(1,), (2,), (3,), None]
+        assert result_set.rownumber == 3
+        assert result_set.connection.client.get_query_results.call_count == 2
+
+    async def test_fetch_without_next_page(self):
+        """_async_fetch() requests nothing once no page remains."""
+        result_set = await _create_result_set(AthenaAioResultSet)
+        await result_set._async_fetch()
+        assert list(result_set._rows) == [(1,), (2,)]
+        assert result_set.connection.client.get_query_results.call_count == 1
+
+    @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())])
+    async def test_fetch_counts_rows_taken_before_error(self, method, args):
+        """The rows taken before a failed page request are counted in rownumber."""
+        result_set = await _create_result_set(
+            AthenaAioResultSet, [("1", "2"), RuntimeError("failed")]
+        )
+        with pytest.raises(OperationalError):
+            await getattr(result_set, method)(*args)
+        assert result_set.rownumber == 2
+
+    async def test_fetch_after_close(self):
+        result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)])
+        result_set.close()
+        assert await result_set.fetchone() is None
+        assert await result_set.fetchmany() == []
+        assert await result_set.fetchall() == []
+        assert result_set.rownumber is None
diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py
index 2b8eddf35..46be7dc41 100644
--- a/tests/pyathena/arrow/test_result_set.py
+++ b/tests/pyathena/arrow/test_result_set.py
@@ -35,6 +35,35 @@ def test_fetch_after_close(self):
         assert result_set.fetchmany() == []
         assert result_set.fetchall() == []
 
+    def test_fetch_across_batches(self):
+        """The buffer is refilled one record batch at a time.
+
+        No AWS calls; the first page request and the result table are mocked.
+        """
+        with (
+            patch.object(AthenaArrowResultSet, "_pre_fetch"),
+            patch.object(AthenaArrowResultSet, "_create_s3_file_system"),
+            patch.object(
+                AthenaArrowResultSet, "_as_arrow", return_value=pa.table({"a": [1, 2, 3]})
+            ),
+        ):
+            result_set = AthenaArrowResultSet(
+                connection=MagicMock(),
+                converter=DefaultArrowTypeConverter(),
+                query_execution=MagicMock(
+                    state=AthenaQueryExecution.STATE_SUCCEEDED,
+                    output_location="s3://bucket/query.csv",
+                ),
+                arraysize=1,
+                retry_config=RetryConfig(),
+            )
+        assert result_set.fetchmany(2) == [(1,), (2,)]
+        assert result_set.fetchone() == (3,)
+        assert result_set.fetchmany() == []
+        assert result_set.fetchall() == []
+        assert result_set.fetchone() is None
+        assert result_set.rownumber == 3
+
     def test_read_csv_timestamps_with_the_same_name(self):
         """Timestamp columns are converted by position, also with names that repeat.
 
diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py
index 4bb1355e9..ed604331f 100644
--- a/tests/pyathena/s3fs/test_result_set.py
+++ b/tests/pyathena/s3fs/test_result_set.py
@@ -121,6 +121,18 @@ def test_reader_closes_at_end_of_results(self):
         assert result_set.fetchone() is None
         result_set.close()
 
+    def test_fetch_across_batches(self):
+        """The buffer is refilled one batch of arraysize rows at a time."""
+        result_set, stream = _result_set("null_col,empty_col,value_col\n,,a\n,,b\n,,c\n", None)
+        assert result_set.fetchmany(2) == [(None, None, "a"), (None, None, "b")]
+        assert result_set.fetchone() == (None, None, "c")
+        assert result_set.fetchmany() == []
+        assert result_set.fetchall() == []
+        assert result_set.fetchone() is None
+        assert result_set.rownumber == 3
+        assert stream.closed
+        result_set.close()
+
     def test_reader_failure_closes_file(self):
         """A failure to read the header row closes the file."""
         data = "null_col,empty_col,value_col\n"
diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py
index c0e11238c..7041db80e 100644
--- a/tests/pyathena/test_result_set.py
+++ b/tests/pyathena/test_result_set.py
@@ -54,27 +54,6 @@ def _fetching_result_set(*pages):
     )
 
 
-class _FetchoneResultSet(AthenaResultSet):
-    """A result set whose ``fetchone()`` takes rows from an iterator, not the buffer."""
-
-    def __init__(self, rows):
-        """Initialize the result set with its rows, skipping the base initialization.
-
-        Args:
-            rows: The rows that ``fetchone()`` returns.
-        """
-        self._rows = collections.deque()
-        self._rownumber = 0
-        self._row_iter = iter(rows)
-
-    @override
-    def fetchone(self):
-        row = next(self._row_iter, None)
-        if row is not None:
-            self._rownumber += 1
-        return row
-
-
 class _ClosingDeque(collections.deque):
     """A buffer that closes its result set when it is read, as another thread may."""
 
@@ -158,12 +137,18 @@ def test_fetch_with_close_during_take(self, method, args, expected):
         assert getattr(result_set, method)(*args) == expected
         assert result_set.rownumber is None
 
-    def test_fetch_with_fetchone_override(self):
-        """A subclass whose fetchone() does not use the buffer has its rows fetched."""
-        result_set = _FetchoneResultSet([(1,), (2,), (3,)])
-        assert result_set.fetchmany(2) == [(1,), (2,)]
-        assert result_set.fetchall() == [(3,)]
+    def test_fetchone_across_pages(self):
+        result_set = _fetching_result_set(("1", "2"), ("3",))
+        assert [result_set.fetchone() for _ in range(4)] == [(1,), (2,), (3,), None]
         assert result_set.rownumber == 3
+        assert result_set.connection.client.get_query_results.call_count == 2
+
+    def test_fetch_without_next_page(self):
+        """_fetch() requests nothing once no page remains."""
+        result_set = _fetching_result_set(("1", "2"))
+        result_set._fetch()
+        assert list(result_set._rows) == [(1,), (2,)]
+        assert result_set.connection.client.get_query_results.call_count == 1
 
     def test_fetch_all_rows_skips_column_labels_only_on_first_page(self):
         """A later page can start with a data row equal to the column labels.

From 0100221287e8329c273967a3b1923c2ce9dce287 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Fri, 16 Oct 2026 00:31:01 +0900
Subject: [PATCH 2/3] Describe when _fetch() fetches nothing and drop redundant
 refill tests

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/aio/result_set.py             | 4 ++--
 pyathena/result_set.py                 | 4 ++--
 tests/pyathena/aio/test_result_set.py  | 7 -------
 tests/pyathena/s3fs/test_result_set.py | 3 +--
 tests/pyathena/test_result_set.py      | 7 -------
 5 files changed, 5 insertions(+), 20 deletions(-)

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index a26546d3d..a9a1a139d 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -145,11 +145,11 @@ async def _async_get_query_results(
     async def _async_fetch(self) -> None:
         """Fetch the next page of rows, if any, into the buffer.
 
-        Nothing is fetched when no page remains or the result set is closed.
+        Nothing is fetched when no page remains, as after ``close()``.
 
         Raises:
             ProgrammingError: If the query ID is missing, the query has not
-                succeeded, or another thread closes the result set.
+                succeeded, or the result set is closed.
             OperationalError: If the request fails.
         """
         if not self._next_token:
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index 5251664c3..ae200be9d 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -462,11 +462,11 @@ def _fetch(self) -> None:
         """Fetch the next rows, if any, into the buffer.
 
         The default result set fetches the next page; it fetches nothing when no
-        page remains or the result set is closed.
+        page remains, as after ``close()``.
 
         Raises:
             ProgrammingError: If the query ID is missing, the query has not
-                succeeded, or another thread closes the result set.
+                succeeded, or the result set is closed.
             OperationalError: If the request fails.
         """
         if not self._next_token:
diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py
index ba5983797..90a635dcc 100644
--- a/tests/pyathena/aio/test_result_set.py
+++ b/tests/pyathena/aio/test_result_set.py
@@ -99,13 +99,6 @@ async def test_fetchone_across_pages(self):
         assert result_set.rownumber == 3
         assert result_set.connection.client.get_query_results.call_count == 2
 
-    async def test_fetch_without_next_page(self):
-        """_async_fetch() requests nothing once no page remains."""
-        result_set = await _create_result_set(AthenaAioResultSet)
-        await result_set._async_fetch()
-        assert list(result_set._rows) == [(1,), (2,)]
-        assert result_set.connection.client.get_query_results.call_count == 1
-
     @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())])
     async def test_fetch_counts_rows_taken_before_error(self, method, args):
         """The rows taken before a failed page request are counted in rownumber."""
diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py
index ed604331f..0c72aa5b7 100644
--- a/tests/pyathena/s3fs/test_result_set.py
+++ b/tests/pyathena/s3fs/test_result_set.py
@@ -123,14 +123,13 @@ def test_reader_closes_at_end_of_results(self):
 
     def test_fetch_across_batches(self):
         """The buffer is refilled one batch of arraysize rows at a time."""
-        result_set, stream = _result_set("null_col,empty_col,value_col\n,,a\n,,b\n,,c\n", None)
+        result_set, _ = _result_set("null_col,empty_col,value_col\n,,a\n,,b\n,,c\n", None)
         assert result_set.fetchmany(2) == [(None, None, "a"), (None, None, "b")]
         assert result_set.fetchone() == (None, None, "c")
         assert result_set.fetchmany() == []
         assert result_set.fetchall() == []
         assert result_set.fetchone() is None
         assert result_set.rownumber == 3
-        assert stream.closed
         result_set.close()
 
     def test_reader_failure_closes_file(self):
diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py
index 7041db80e..087506552 100644
--- a/tests/pyathena/test_result_set.py
+++ b/tests/pyathena/test_result_set.py
@@ -143,13 +143,6 @@ def test_fetchone_across_pages(self):
         assert result_set.rownumber == 3
         assert result_set.connection.client.get_query_results.call_count == 2
 
-    def test_fetch_without_next_page(self):
-        """_fetch() requests nothing once no page remains."""
-        result_set = _fetching_result_set(("1", "2"))
-        result_set._fetch()
-        assert list(result_set._rows) == [(1,), (2,)]
-        assert result_set.connection.client.get_query_results.call_count == 1
-
     def test_fetch_all_rows_skips_column_labels_only_on_first_page(self):
         """A later page can start with a data row equal to the column labels.
 

From 6cd9b5962dcc93ac668eff417e137641aef5e0c4 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Fri, 16 Oct 2026 00:39:55 +0900
Subject: [PATCH 3/3] Qualify when _fetch() raises ProgrammingError

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/aio/result_set.py | 4 ++--
 pyathena/result_set.py     | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index a9a1a139d..5e27e4b44 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -148,8 +148,8 @@ async def _async_fetch(self) -> None:
         Nothing is fetched when no page remains, as after ``close()``.
 
         Raises:
-            ProgrammingError: If the query ID is missing, the query has not
-                succeeded, or the result set is closed.
+            ProgrammingError: If the next page is requested and the query ID is
+                missing, the query has not succeeded, or the result set is closed.
             OperationalError: If the request fails.
         """
         if not self._next_token:
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index ae200be9d..3e29b938a 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -465,8 +465,8 @@ def _fetch(self) -> None:
         page remains, as after ``close()``.
 
         Raises:
-            ProgrammingError: If the query ID is missing, the query has not
-                succeeded, or the result set is closed.
+            ProgrammingError: If the next page is requested and the query ID is
+                missing, the query has not succeeded, or the result set is closed.
             OperationalError: If the request fails.
         """
         if not self._next_token: