From 8a868f4168efd3ba06647c183dcd8fbc166d904f Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 17 Oct 2026 00:15:04 +0900
Subject: [PATCH 1/4] Keep rows of a failed fetch, skip empty Arrow batches,
 move the clamp

- fetchmany() and fetchall() of the result set and its aio variant
  return the rows that they took before a failed page request to the
  front of the buffer, instead of dropping them while counting them in
  rownumber, so that a later fetch returns them.
- The Arrow result set skips record batches without rows, which ended
  the fetches early, such as the batch of an empty result file.
- The size clamp to sys.maxsize moves from the base fetchmany() into the
  islice() of the pandas and Polars result sets, which need it.
- The test helper _page is replaced with get_query_results_responses().

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/aio/result_set.py              | 16 +++++++---
 pyathena/arrow/result_set.py            | 35 +++++++++++----------
 pyathena/pandas/result_set.py           |  6 ++--
 pyathena/polars/result_set.py           |  6 ++--
 pyathena/result_set.py                  | 36 ++++++++++++++++-----
 tests/pyathena/aio/test_result_set.py   |  7 +++--
 tests/pyathena/arrow/test_result_set.py | 42 +++++++++++++++++++++++++
 tests/pyathena/test_result_set.py       | 16 ++++------
 8 files changed, 119 insertions(+), 45 deletions(-)

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index 40bb2ad3c..765261d9c 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -241,6 +241,7 @@ async def _async_take_rows(
 
         The buffered rows are taken in bulk, and ``_async_fetch()`` refills the buffer
         when it is empty.
+        If a refill fails, the rows taken are returned to the buffer.
 
         Args:
             size: The maximum number of rows, or None for all remaining rows.
@@ -249,12 +250,17 @@ async def _async_take_rows(
             The rows, fewer than ``size`` when the result is exhausted.
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
-        while size is None or len(rows) < size:
-            if not self._rows:
-                await self._async_fetch()
+        try:
+            while size is None or len(rows) < size:
                 if not self._rows:
-                    break
-            self._take_buffered_rows(rows, size)
+                    await self._async_fetch()
+                    if not self._rows:
+                        break
+                self._take_buffered_rows(rows, size)
+        except BaseException:
+            # Also on cancellation, which a later fetch can continue from.
+            self._return_rows(rows)
+            raise
         return rows
 
     @override
diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py
index 508f644d7..cd471a103 100644
--- a/pyathena/arrow/result_set.py
+++ b/pyathena/arrow/result_set.py
@@ -197,24 +197,27 @@ def converters(self) -> dict[str, Callable[[str | None], Any | None]]:
 
     @override
     def _fetch(self) -> None:
-        try:
-            rows = next(self._batches)
-        except StopIteration:
+        """Buffer the rows of the next record batch that has any.
+
+        A table can hold empty record batches between others, such as one of
+        an empty result file, which would otherwise end the fetches early.
+        """
+        rows = next((batch for batch in self._batches if batch.num_rows), None)
+        if rows is None:
             return
+        # Read the columns and their converters by position; to_pydict() and the
+        # converters property keep one column per name.
+        columns = [column.to_pylist() for column in rows.columns]
+        description = self.description if self.description else []
+        converters = [self._converter.get(d[1]) for d in description]
+        if any(convert is not _to_default for convert in converters):
+            processed_rows = [
+                tuple(convert(v) for convert, v in zip(converters, row, strict=False))
+                for row in zip(*columns, strict=False)
+            ]
         else:
-            # Read the columns and their converters by position; to_pydict() and the
-            # converters property keep one column per name.
-            columns = [column.to_pylist() for column in rows.columns]
-            description = self.description if self.description else []
-            converters = [self._converter.get(d[1]) for d in description]
-            if any(convert is not _to_default for convert in converters):
-                processed_rows = [
-                    tuple(convert(v) for convert, v in zip(converters, row, strict=False))
-                    for row in zip(*columns, strict=False)
-                ]
-            else:
-                processed_rows = list(zip(*columns, strict=False))
-            self._rows.extend(processed_rows)
+            processed_rows = list(zip(*columns, strict=False))
+        self._rows.extend(processed_rows)
 
     def _open_csv(self, location: str) -> Any:
         """Open the CSV result file for ``pyarrow.csv.read_csv()``.
diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py
index 2762b9b86..8592c7aa1 100644
--- a/pyathena/pandas/result_set.py
+++ b/pyathena/pandas/result_set.py
@@ -5,6 +5,7 @@
 import csv
 import importlib
 import logging
+import sys
 import threading
 from collections import abc
 from collections.abc import Callable, Iterable, Iterator
@@ -992,8 +993,9 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
         try:
-            # extend() keeps the rows that it appended before an exception.
-            rows.extend(islice(self._row_iter, size))
+            # extend() keeps the rows that it appended before an exception, and
+            # islice() takes no size above sys.maxsize.
+            rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize)))
         finally:
             if rows:
                 self._count_rows(len(rows))
diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py
index 74a2e099f..b79802547 100644
--- a/pyathena/polars/result_set.py
+++ b/pyathena/polars/result_set.py
@@ -12,6 +12,7 @@
 import csv
 import importlib.util
 import logging
+import sys
 import threading
 from collections import abc
 from collections.abc import Callable, Iterable, Iterator
@@ -647,8 +648,9 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
         try:
-            # extend() keeps the rows that it appended before an exception.
-            rows.extend(islice(self._row_iter, size))
+            # extend() keeps the rows that it appended before an exception, and
+            # islice() takes no size above sys.maxsize.
+            rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize)))
         finally:
             if rows:
                 self._count_rows(len(rows))
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index c0a7cf9e3..7bf314dc7 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -5,7 +5,6 @@
 import collections
 import logging
 import operator
-import sys
 from datetime import datetime
 from itertools import islice
 from typing import (
@@ -529,8 +528,7 @@ def fetchmany(
         """
         if not size or size <= 0:
             size = self._arraysize
-        # islice() in the subclasses' _take_rows() takes no size above sys.maxsize.
-        return self._take_rows(min(operator.index(size), sys.maxsize))
+        return self._take_rows(operator.index(size))
 
     @override
     def fetchall(
@@ -549,6 +547,7 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
         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.
+        If a refill fails, the rows taken are returned to the buffer.
 
         Args:
             size: The maximum number of rows, or None for all remaining rows.
@@ -557,12 +556,16 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
             The rows, fewer than ``size`` when the result is exhausted.
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
-        while size is None or len(rows) < size:
-            if not self._rows:
-                self._fetch()
+        try:
+            while size is None or len(rows) < size:
                 if not self._rows:
-                    break
-            self._take_buffered_rows(rows, size)
+                    self._fetch()
+                    if not self._rows:
+                        break
+                self._take_buffered_rows(rows, size)
+        except BaseException:
+            self._return_rows(rows)
+            raise
         return rows
 
     def _take_buffered_rows(
@@ -591,6 +594,23 @@ def _take_buffered_rows(
             rows.extend(buffered)
         self._count_rows(len(rows) - start)
 
+    def _return_rows(self, rows: list[tuple[Any | None, ...] | dict[Any, Any | None]]) -> None:
+        """Return the rows that a failed fetch took to the front of the buffer.
+
+        The rows are no longer counted in ``rownumber``, so that a later fetch
+        returns them, as ``fetchone()`` leaves the rows after a failed page
+        request in the buffer. A result set that is closed keeps its empty
+        buffer.
+
+        Args:
+            rows: The rows that the fetch took, in their order.
+        """
+        rownumber = self._rownumber
+        if rownumber is None or not rows:
+            return
+        self._rows.extendleft(reversed(rows))
+        self._rownumber = rownumber - len(rows)
+
     def _count_rows(self, count: int) -> None:
         """Count the rows that a fetch took in ``rownumber``.
 
diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py
index dc0292ba8..0312f6b33 100644
--- a/tests/pyathena/aio/test_result_set.py
+++ b/tests/pyathena/aio/test_result_set.py
@@ -100,13 +100,16 @@ async def test_fetchone_across_pages(self):
         assert result_set.connection.client.get_query_results.call_count == 2
 
     @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."""
+    async def test_fetch_returns_rows_taken_before_error(self, method, args):
+        """The rows taken before a failed page request are returned to the buffer."""
+        # GH-1186: they used to be dropped while 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 == 0
+        assert await result_set.fetchmany(2) == [(1,), (2,)]
         assert result_set.rownumber == 2
 
     async def test_fetch_after_close(self):
diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py
index 46be7dc41..c1eb15fa4 100644
--- a/tests/pyathena/arrow/test_result_set.py
+++ b/tests/pyathena/arrow/test_result_set.py
@@ -64,6 +64,48 @@ def test_fetch_across_batches(self):
         assert result_set.fetchone() is None
         assert result_set.rownumber == 3
 
+    @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall", "iter"])
+    def test_fetch_across_empty_batches(self, method):
+        """An empty record batch between others does not end the fetches.
+
+        GH-1186: a table read from result files that include an empty one has
+        such a batch, and the fetches used to stop at it. No AWS calls; the
+        first page request and the result table are mocked.
+        """
+        table = pa.concat_tables(
+            [
+                pa.table({"a": [1, 2, 3]}),
+                pa.table({"a": pa.array([], pa.int64())}),
+                pa.table({"a": [4]}),
+            ]
+        )
+        with (
+            patch.object(AthenaArrowResultSet, "_pre_fetch"),
+            patch.object(AthenaArrowResultSet, "_create_s3_file_system"),
+            patch.object(AthenaArrowResultSet, "_as_arrow", return_value=table),
+        ):
+            result_set = AthenaArrowResultSet(
+                connection=MagicMock(),
+                converter=DefaultArrowTypeConverter(),
+                query_execution=MagicMock(
+                    state=AthenaQueryExecution.STATE_SUCCEEDED,
+                    output_location="s3://bucket/query.csv",
+                ),
+                arraysize=2,
+                retry_config=RetryConfig(),
+            )
+        assert 0 in [batch.num_rows for batch in table.to_batches(2)]
+        if method == "fetchone":
+            rows = list(iter(result_set.fetchone, None))
+        elif method == "fetchmany":
+            rows = result_set.fetchmany(10)
+        elif method == "fetchall":
+            rows = result_set.fetchall()
+        else:
+            rows = list(result_set)
+        assert rows == [(1,), (2,), (3,), (4,)]
+        assert result_set.rownumber == 4
+
     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/test_result_set.py b/tests/pyathena/test_result_set.py
index 83df71187..3790237a8 100644
--- a/tests/pyathena/test_result_set.py
+++ b/tests/pyathena/test_result_set.py
@@ -20,13 +20,6 @@
 from tests.pyathena.util import close_during_next_request, get_query_results_responses
 
 
-def _page(values, next_token=None):
-    response = {"ResultSet": {"Rows": [{"Data": [{"VarCharValue": v}]} for v in values]}}
-    if next_token:
-        response["NextToken"] = next_token
-    return response
-
-
 def _row(*values):
     return {"Data": [{} if v is None else {"VarCharValue": v} for v in values]}
 
@@ -112,11 +105,14 @@ def test_fetchmany_size_not_integer(self):
         assert result_set.fetchall() == [(1,), (2,)]
 
     @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())])
-    def test_fetch_counts_rows_taken_before_error(self, method, args):
-        """The rows taken before a failed page request are counted in rownumber."""
+    def test_fetch_returns_rows_taken_before_error(self, method, args):
+        """The rows taken before a failed page request are returned to the buffer."""
+        # GH-1186: they used to be dropped while counted in rownumber.
         result_set = _fetching_result_set(("1", "2"), RuntimeError("failed"))
         with pytest.raises(OperationalError):
             getattr(result_set, method)(*args)
+        assert result_set.rownumber == 0
+        assert result_set.fetchmany(2) == [(1,), (2,)]
         assert result_set.rownumber == 2
 
     def test_fetch_after_close(self):
@@ -180,7 +176,7 @@ def test_fetch_all_rows_skips_column_labels_only_on_first_page(self):
         result_set._process_metadata(
             {"ResultSet": {"ResultSetMetadata": {"ColumnInfo": [{"Name": "a", "Type": "varchar"}]}}}
         )
-        pages = [_page(["a", "1"], "token"), _page(["a", "2"])]
+        pages = get_query_results_responses(["a", "1"], ["a", "2"])
         with patch.object(result_set, "_get_query_results", side_effect=pages):
             assert result_set._fetch_all_rows() == [("1",), ("a",), ("2",)]
 

From 6b8612c380ad849961a906bdd9b856a904d897ea Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 17 Oct 2026 01:05:32 +0900
Subject: [PATCH 2/4] Request a page again after its rows fail to convert

The result sets advanced the token of the next page before converting
the rows of the fetched one, so a fetch after a failed conversion
skipped that page; now that fetchmany() and fetchall() return the rows
taken before a failure to the buffer, a later fetch would continue past
it. The token is advanced once the rows are converted.

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

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index 765261d9c..f99f1fd18 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -156,8 +156,11 @@ async def _async_fetch(self) -> None:
         if not self._next_token:
             return
         response = await self._async_get_query_results(self._arraysize, self._next_token)
-        rows, self._next_token = self._parse_result_rows(response)
+        rows, next_token = self._parse_result_rows(response)
         self._process_rows(rows)
+        # Advanced once the rows are converted, so that a fetch after a failed
+        # conversion requests the page again instead of skipping it.
+        self._next_token = next_token
         # close() clears the connection before the rest, so checking after the
         # assignments catches a close() during the request or this method.
         if self.is_closed:
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index 7bf314dc7..8d235844f 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -472,8 +472,11 @@ def _fetch(self) -> None:
         if not self._next_token:
             return
         response = self._get_query_results(self._arraysize, self._next_token)
-        rows, self._next_token = self._parse_result_rows(response)
+        rows, next_token = self._parse_result_rows(response)
         self._process_rows(rows)
+        # Advanced once the rows are converted, so that a fetch after a failed
+        # conversion requests the page again instead of skipping it.
+        self._next_token = next_token
         # close() clears the connection before the rest, so checking after the
         # assignments catches a close() during the request or this method.
         if self.is_closed:
diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py
index 0312f6b33..632f381f7 100644
--- a/tests/pyathena/aio/test_result_set.py
+++ b/tests/pyathena/aio/test_result_set.py
@@ -112,6 +112,16 @@ async def test_fetch_returns_rows_taken_before_error(self, method, args):
         assert await result_set.fetchmany(2) == [(1,), (2,)]
         assert result_set.rownumber == 2
 
+    async def test_fetch_after_failed_conversion(self):
+        """A page whose rows fail to convert is requested again by a later fetch."""
+        # GH-1186: the token of the next page used to be advanced first, so the
+        # page was skipped.
+        result_set = await _create_result_set(AthenaAioResultSet, [("1",), ("x",), ("2",)])
+        assert await result_set.fetchone() == (1,)
+        with pytest.raises(ValueError, match="x"):
+            await result_set.fetchall()
+        assert result_set._next_token == "token1"
+
     async def test_fetch_after_close(self):
         result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)])
         result_set.close()
diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py
index 3790237a8..7bf5d98d3 100644
--- a/tests/pyathena/test_result_set.py
+++ b/tests/pyathena/test_result_set.py
@@ -115,6 +115,17 @@ def test_fetch_returns_rows_taken_before_error(self, method, args):
         assert result_set.fetchmany(2) == [(1,), (2,)]
         assert result_set.rownumber == 2
 
+    @pytest.mark.parametrize(("method", "args"), [("fetchone", ()), ("fetchall", ())])
+    def test_fetch_after_failed_conversion(self, method, args):
+        """A page whose rows fail to convert is requested again by a later fetch."""
+        # GH-1186: the token of the next page used to be advanced first, so the
+        # page was skipped.
+        result_set = _fetching_result_set(("1",), ("x",), ("2",))
+        assert result_set.fetchone() == (1,)
+        with pytest.raises(ValueError, match="x"):
+            getattr(result_set, method)(*args)
+        assert result_set._next_token == "token1"
+
     def test_fetch_after_close(self):
         result_set = _fetching_result_set(("1", "2"), ("3",))
         result_set.close()

From 58ced3eef74e352d0bf1b525684ffd358962d417 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 17 Oct 2026 08:20:41 +0900
Subject: [PATCH 3/4] Revert "Request a page again after its rows fail to
 convert"

This reverts commit 6b8612c380ad849961a906bdd9b856a904d897ea.
---
 pyathena/aio/result_set.py            |  5 +----
 pyathena/result_set.py                |  5 +----
 tests/pyathena/aio/test_result_set.py | 10 ----------
 tests/pyathena/test_result_set.py     | 11 -----------
 4 files changed, 2 insertions(+), 29 deletions(-)

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index f99f1fd18..765261d9c 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -156,11 +156,8 @@ async def _async_fetch(self) -> None:
         if not self._next_token:
             return
         response = await self._async_get_query_results(self._arraysize, self._next_token)
-        rows, next_token = self._parse_result_rows(response)
+        rows, self._next_token = self._parse_result_rows(response)
         self._process_rows(rows)
-        # Advanced once the rows are converted, so that a fetch after a failed
-        # conversion requests the page again instead of skipping it.
-        self._next_token = next_token
         # close() clears the connection before the rest, so checking after the
         # assignments catches a close() during the request or this method.
         if self.is_closed:
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index 8d235844f..7bf314dc7 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -472,11 +472,8 @@ def _fetch(self) -> None:
         if not self._next_token:
             return
         response = self._get_query_results(self._arraysize, self._next_token)
-        rows, next_token = self._parse_result_rows(response)
+        rows, self._next_token = self._parse_result_rows(response)
         self._process_rows(rows)
-        # Advanced once the rows are converted, so that a fetch after a failed
-        # conversion requests the page again instead of skipping it.
-        self._next_token = next_token
         # close() clears the connection before the rest, so checking after the
         # assignments catches a close() during the request or this method.
         if self.is_closed:
diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py
index 632f381f7..0312f6b33 100644
--- a/tests/pyathena/aio/test_result_set.py
+++ b/tests/pyathena/aio/test_result_set.py
@@ -112,16 +112,6 @@ async def test_fetch_returns_rows_taken_before_error(self, method, args):
         assert await result_set.fetchmany(2) == [(1,), (2,)]
         assert result_set.rownumber == 2
 
-    async def test_fetch_after_failed_conversion(self):
-        """A page whose rows fail to convert is requested again by a later fetch."""
-        # GH-1186: the token of the next page used to be advanced first, so the
-        # page was skipped.
-        result_set = await _create_result_set(AthenaAioResultSet, [("1",), ("x",), ("2",)])
-        assert await result_set.fetchone() == (1,)
-        with pytest.raises(ValueError, match="x"):
-            await result_set.fetchall()
-        assert result_set._next_token == "token1"
-
     async def test_fetch_after_close(self):
         result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)])
         result_set.close()
diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py
index 7bf5d98d3..3790237a8 100644
--- a/tests/pyathena/test_result_set.py
+++ b/tests/pyathena/test_result_set.py
@@ -115,17 +115,6 @@ def test_fetch_returns_rows_taken_before_error(self, method, args):
         assert result_set.fetchmany(2) == [(1,), (2,)]
         assert result_set.rownumber == 2
 
-    @pytest.mark.parametrize(("method", "args"), [("fetchone", ()), ("fetchall", ())])
-    def test_fetch_after_failed_conversion(self, method, args):
-        """A page whose rows fail to convert is requested again by a later fetch."""
-        # GH-1186: the token of the next page used to be advanced first, so the
-        # page was skipped.
-        result_set = _fetching_result_set(("1",), ("x",), ("2",))
-        assert result_set.fetchone() == (1,)
-        with pytest.raises(ValueError, match="x"):
-            getattr(result_set, method)(*args)
-        assert result_set._next_token == "token1"
-
     def test_fetch_after_close(self):
         result_set = _fetching_result_set(("1", "2"), ("3",))
         result_set.close()

From 9809e3312b116d50e0cb98d177b848e32de99c8f Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 17 Oct 2026 08:24:08 +0900
Subject: [PATCH 4/4] Keep the fetches after an error as they are

The behavior of a fetch after an exception is left unspecified, as
PEP 249 leaves it, and the other result sets could not follow a
recovery contract of the GetQueryResults ones, so returning the rows
of a failed fetch to the buffer is removed with the token reordering.
The empty-batch fix of the Arrow result set is tested with empty
batches before and between others, and the clamp and test helper
changes stay.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/aio/result_set.py              | 16 ++++--------
 pyathena/result_set.py                  | 32 ++++--------------------
 tests/pyathena/aio/test_result_set.py   |  7 ++----
 tests/pyathena/arrow/test_result_set.py | 33 ++++++++++++++-----------
 tests/pyathena/test_result_set.py       |  7 ++----
 5 files changed, 33 insertions(+), 62 deletions(-)

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index 765261d9c..40bb2ad3c 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -241,7 +241,6 @@ async def _async_take_rows(
 
         The buffered rows are taken in bulk, and ``_async_fetch()`` refills the buffer
         when it is empty.
-        If a refill fails, the rows taken are returned to the buffer.
 
         Args:
             size: The maximum number of rows, or None for all remaining rows.
@@ -250,17 +249,12 @@ async def _async_take_rows(
             The rows, fewer than ``size`` when the result is exhausted.
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
-        try:
-            while size is None or len(rows) < size:
+        while size is None or len(rows) < size:
+            if not self._rows:
+                await self._async_fetch()
                 if not self._rows:
-                    await self._async_fetch()
-                    if not self._rows:
-                        break
-                self._take_buffered_rows(rows, size)
-        except BaseException:
-            # Also on cancellation, which a later fetch can continue from.
-            self._return_rows(rows)
-            raise
+                    break
+            self._take_buffered_rows(rows, size)
         return rows
 
     @override
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index 7bf314dc7..1e1445641 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -547,7 +547,6 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
         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.
-        If a refill fails, the rows taken are returned to the buffer.
 
         Args:
             size: The maximum number of rows, or None for all remaining rows.
@@ -556,16 +555,12 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
             The rows, fewer than ``size`` when the result is exhausted.
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
-        try:
-            while size is None or len(rows) < size:
+        while size is None or len(rows) < size:
+            if not self._rows:
+                self._fetch()
                 if not self._rows:
-                    self._fetch()
-                    if not self._rows:
-                        break
-                self._take_buffered_rows(rows, size)
-        except BaseException:
-            self._return_rows(rows)
-            raise
+                    break
+            self._take_buffered_rows(rows, size)
         return rows
 
     def _take_buffered_rows(
@@ -594,23 +589,6 @@ def _take_buffered_rows(
             rows.extend(buffered)
         self._count_rows(len(rows) - start)
 
-    def _return_rows(self, rows: list[tuple[Any | None, ...] | dict[Any, Any | None]]) -> None:
-        """Return the rows that a failed fetch took to the front of the buffer.
-
-        The rows are no longer counted in ``rownumber``, so that a later fetch
-        returns them, as ``fetchone()`` leaves the rows after a failed page
-        request in the buffer. A result set that is closed keeps its empty
-        buffer.
-
-        Args:
-            rows: The rows that the fetch took, in their order.
-        """
-        rownumber = self._rownumber
-        if rownumber is None or not rows:
-            return
-        self._rows.extendleft(reversed(rows))
-        self._rownumber = rownumber - len(rows)
-
     def _count_rows(self, count: int) -> None:
         """Count the rows that a fetch took in ``rownumber``.
 
diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py
index 0312f6b33..dc0292ba8 100644
--- a/tests/pyathena/aio/test_result_set.py
+++ b/tests/pyathena/aio/test_result_set.py
@@ -100,16 +100,13 @@ async def test_fetchone_across_pages(self):
         assert result_set.connection.client.get_query_results.call_count == 2
 
     @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())])
-    async def test_fetch_returns_rows_taken_before_error(self, method, args):
-        """The rows taken before a failed page request are returned to the buffer."""
-        # GH-1186: they used to be dropped while counted in rownumber.
+    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 == 0
-        assert await result_set.fetchmany(2) == [(1,), (2,)]
         assert result_set.rownumber == 2
 
     async def test_fetch_after_close(self):
diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py
index c1eb15fa4..cd214b6b4 100644
--- a/tests/pyathena/arrow/test_result_set.py
+++ b/tests/pyathena/arrow/test_result_set.py
@@ -65,20 +65,24 @@ def test_fetch_across_batches(self):
         assert result_set.rownumber == 3
 
     @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall", "iter"])
-    def test_fetch_across_empty_batches(self, method):
-        """An empty record batch between others does not end the fetches.
+    @pytest.mark.parametrize(
+        ("chunks", "expected"),
+        [
+            ([[1, 2, 3], [], [4]], [1, 2, 3, 4]),
+            ([[], [1, 2]], [1, 2]),
+            ([[1], [], [], [2]], [1, 2]),
+        ],
+        ids=["middle", "leading", "consecutive"],
+    )
+    def test_fetch_across_empty_batches(self, method, chunks, expected):
+        """Empty record batches before or between others do not end the fetches.
 
-        GH-1186: a table read from result files that include an empty one has
-        such a batch, and the fetches used to stop at it. No AWS calls; the
-        first page request and the result table are mocked.
+        A table read from result files that include an empty one has such
+        batches; the fetches used to stop at the first one. to_batches() yields
+        no empty batch after the last rows. No AWS calls; the first page
+        request and the result table are mocked.
         """
-        table = pa.concat_tables(
-            [
-                pa.table({"a": [1, 2, 3]}),
-                pa.table({"a": pa.array([], pa.int64())}),
-                pa.table({"a": [4]}),
-            ]
-        )
+        table = pa.concat_tables([pa.table({"a": pa.array(chunk, pa.int64())}) for chunk in chunks])
         with (
             patch.object(AthenaArrowResultSet, "_pre_fetch"),
             patch.object(AthenaArrowResultSet, "_create_s3_file_system"),
@@ -103,8 +107,9 @@ def test_fetch_across_empty_batches(self, method):
             rows = result_set.fetchall()
         else:
             rows = list(result_set)
-        assert rows == [(1,), (2,), (3,), (4,)]
-        assert result_set.rownumber == 4
+        assert rows == [(value,) for value in expected]
+        assert result_set.rownumber == len(expected)
+        assert result_set.fetchone() is None
 
     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/test_result_set.py b/tests/pyathena/test_result_set.py
index 3790237a8..ec647d1bf 100644
--- a/tests/pyathena/test_result_set.py
+++ b/tests/pyathena/test_result_set.py
@@ -105,14 +105,11 @@ def test_fetchmany_size_not_integer(self):
         assert result_set.fetchall() == [(1,), (2,)]
 
     @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())])
-    def test_fetch_returns_rows_taken_before_error(self, method, args):
-        """The rows taken before a failed page request are returned to the buffer."""
-        # GH-1186: they used to be dropped while counted in rownumber.
+    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 = _fetching_result_set(("1", "2"), RuntimeError("failed"))
         with pytest.raises(OperationalError):
             getattr(result_set, method)(*args)
-        assert result_set.rownumber == 0
-        assert result_set.fetchmany(2) == [(1,), (2,)]
         assert result_set.rownumber == 2
 
     def test_fetch_after_close(self):