From c5f2fc663635b4e419a3d8f554258022e43179f8 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 15 Oct 2026 07:34:11 +0900
Subject: [PATCH 1/5] Take buffered rows in bulk in fetchmany() and fetchall()

AthenaResultSet.fetchmany() and fetchall() called fetchone() once per
row. They now move the buffered rows to the result in bulk and update
rownumber once per batch, calling fetchone() only when the buffer is
empty, so that a subclass's fetchone() still fetches the next rows.
AthenaAioResultSet takes its rows the same way.

The pandas and Polars result sets now override only _take_rows(), with
fetchmany(), fetchall(), and _count_rows() moved to the base class.

Closes #1171

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/aio/result_set.py            |  43 ++++++----
 pyathena/pandas/result_set.py         |  35 +-------
 pyathena/polars/result_set.py         |  35 +-------
 pyathena/result_set.py                |  77 ++++++++++++++----
 tests/pyathena/aio/test_result_set.py |  60 +++++++++++---
 tests/pyathena/test_result_set.py     | 113 +++++++++++++++++++++++++-
 6 files changed, 258 insertions(+), 105 deletions(-)

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index 27642580a..95dcbbe3d 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import operator
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -203,17 +204,13 @@ async def fetchmany(  # type: ignore[override]
 
         Returns:
             The rows, fewer than ``size`` when the result is exhausted.
+
+        Raises:
+            TypeError: If ``size`` is not an integer.
         """
         if not size or size <= 0:
             size = self._arraysize
-        rows = []
-        for _ in range(size):
-            row = await self.fetchone()
-            if row:
-                rows.append(row)
-            else:
-                break
-        return rows
+        return await self._async_take_rows(operator.index(size))
 
     @override
     async def fetchall(  # type: ignore[override]
@@ -224,13 +221,31 @@ async def fetchall(  # type: ignore[override]
         Returns:
             The remaining rows.
         """
-        rows = []
-        while True:
-            row = await self.fetchone()
-            if row:
-                rows.append(row)
+        return await self._async_take_rows(None)
+
+    async def _async_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 page into the buffer.
+
+        Args:
+            size: The maximum number of rows, or None for all remaining rows.
+
+        Returns:
+            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 self._rows:
+                self._take_buffered_rows(rows, None if size is None else size - len(rows))
             else:
-                break
+                row = await self.fetchone()
+                if row is None:
+                    break
+                rows.append(row)
         return rows
 
     @override
diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py
index b13b51ed3..586c39dba 100644
--- a/pyathena/pandas/result_set.py
+++ b/pyathena/pandas/result_set.py
@@ -5,8 +5,6 @@
 import csv
 import importlib
 import logging
-import operator
-import sys
 from collections import abc
 from collections.abc import Callable, Iterable, Iterator
 from concurrent.futures import ThreadPoolExecutor
@@ -950,26 +948,11 @@ def fetchone(
         return row
 
     @override
-    def fetchmany(
-        self, size: int | None = None
-    ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
-        if not size or size <= 0:
-            size = self._arraysize
-        # A size that is not an integer raises TypeError, as with the base class, and
-        # islice() takes no size above sys.maxsize.
-        return self._take_rows(min(operator.index(size), sys.maxsize))
-
-    @override
-    def fetchall(
-        self,
-    ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
-        return self._take_rows(None)
-
     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 rows taken before an exception are counted too, as the base class's
-        ``fetchone()`` per row counted them.
+        The rows are taken from the row iterator in bulk. The rows taken before an
+        exception are counted too.
 
         Args:
             size: The maximum number of rows, or None for all remaining rows.
@@ -986,20 +969,6 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
                 self._count_rows(len(rows))
         return rows
 
-    def _count_rows(self, count: int) -> None:
-        """Count the rows that a fetch took in ``rownumber``.
-
-        ``rownumber`` stays None for a result set that another thread closed before
-        the update. It is read once, so that a ``close()`` during the update raises no
-        TypeError, though it may leave the count in place.
-
-        Args:
-            count: The number of rows taken.
-        """
-        rownumber = self._rownumber
-        if rownumber is not None:
-            self._rownumber = rownumber + count
-
     def _read_csv(self) -> TextFileReader | DataFrame:
         """Read the CSV result file, or the GetQueryResults rows as one without it.
 
diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py
index da67e95ba..17fc35e7a 100644
--- a/pyathena/polars/result_set.py
+++ b/pyathena/polars/result_set.py
@@ -12,8 +12,6 @@
 import csv
 import importlib.util
 import logging
-import operator
-import sys
 from collections import abc
 from collections.abc import Callable, Iterable, Iterator
 from io import BytesIO, StringIO
@@ -594,26 +592,11 @@ def fetchone(
         return row
 
     @override
-    def fetchmany(
-        self, size: int | None = None
-    ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
-        if not size or size <= 0:
-            size = self._arraysize
-        # A size that is not an integer raises TypeError, as with the base class, and
-        # islice() takes no size above sys.maxsize.
-        return self._take_rows(min(operator.index(size), sys.maxsize))
-
-    @override
-    def fetchall(
-        self,
-    ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
-        return self._take_rows(None)
-
     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 rows taken before an exception are counted too, as the base class's
-        ``fetchone()`` per row counted them.
+        The rows are taken from the row iterator in bulk. The rows taken before an
+        exception are counted too.
 
         Args:
             size: The maximum number of rows, or None for all remaining rows.
@@ -630,20 +613,6 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
                 self._count_rows(len(rows))
         return rows
 
-    def _count_rows(self, count: int) -> None:
-        """Count the rows that a fetch took in ``rownumber``.
-
-        ``rownumber`` stays None for a result set that another thread closed before
-        the update. It is read once, so that a ``close()`` during the update raises no
-        TypeError, though it may leave the count in place.
-
-        Args:
-            count: The number of rows taken.
-        """
-        rownumber = self._rownumber
-        if rownumber is not None:
-            self._rownumber = rownumber + count
-
     def _is_csv_readable(self) -> bool:
         """Check if CSV output is available and can be read.
 
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index c0800858c..61c3916e8 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -4,6 +4,8 @@
 
 import collections
 import logging
+import operator
+import sys
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
@@ -508,17 +510,14 @@ def fetchmany(
 
         Returns:
             The rows, fewer than ``size`` when the result is exhausted.
+
+        Raises:
+            TypeError: If ``size`` is not an integer.
         """
         if not size or size <= 0:
             size = self._arraysize
-        rows = []
-        for _ in range(size):
-            row = self.fetchone()
-            if row:
-                rows.append(row)
-            else:
-                break
-        return rows
+        # islice() in the subclasses' _take_rows() takes no size above sys.maxsize.
+        return self._take_rows(min(operator.index(size), sys.maxsize))
 
     @override
     def fetchall(
@@ -529,15 +528,65 @@ def fetchall(
         Returns:
             The remaining rows.
         """
-        rows = []
-        while True:
-            row = self.fetchone()
-            if row:
-                rows.append(row)
+        return self._take_rows(None)
+
+    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 a
+        subclass that overrides it does.
+
+        Args:
+            size: The maximum number of rows, or None for all remaining rows.
+
+        Returns:
+            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 self._rows:
+                self._take_buffered_rows(rows, None if size is None else size - len(rows))
             else:
-                break
+                row = self.fetchone()
+                if row is None:
+                    break
+                rows.append(row)
         return rows
 
+    def _take_buffered_rows(
+        self, rows: list[tuple[Any | None, ...] | dict[Any, Any | None]], limit: int | None
+    ) -> None:
+        """Move the buffered rows to ``rows`` and count them in ``rownumber``.
+
+        Args:
+            rows: The list that receives the rows.
+            limit: The maximum number of rows, or None for all buffered rows.
+        """
+        count = len(self._rows)
+        if limit is not None and limit < count:
+            count = limit
+            popleft = self._rows.popleft
+            rows.extend([popleft() for _ in range(count)])
+        else:
+            rows.extend(self._rows)
+            self._rows.clear()
+        self._count_rows(count)
+
+    def _count_rows(self, count: int) -> None:
+        """Count the rows that a fetch took in ``rownumber``.
+
+        ``rownumber`` stays None for a result set that another thread closed before
+        the update. It is read once, so that a ``close()`` during the update raises no
+        TypeError, though it may leave the count in place.
+
+        Args:
+            count: The number of rows taken.
+        """
+        rownumber = self._rownumber
+        if rownumber is not None:
+            self._rownumber = rownumber + count
+
     def _process_metadata(self, response: dict[str, Any]) -> None:
         result_set = response.get("ResultSet")
         if not result_set:
diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py
index b34910afb..cc8434e0a 100644
--- a/tests/pyathena/aio/test_result_set.py
+++ b/tests/pyathena/aio/test_result_set.py
@@ -14,22 +14,31 @@
 from pyathena.util import RetryConfig
 
 
-async def _create_result_set(result_set_class):
-    """Create a result set of a succeeded query from a stubbed ``GetQueryResults``.
+async def _create_result_set(result_set_class, pages=(("1", "2"),)):
+    """Create a result set of a succeeded query from stubbed ``GetQueryResults`` pages.
 
     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.
 
     Returns:
-        The result set, holding the rows 1 and 2 of an integer column ``a``.
+        The result set.
     """
-    connection = MagicMock()
-    connection.client.get_query_results.return_value = {
-        "ResultSet": {
-            "ResultSetMetadata": {"ColumnInfo": [{"Name": "a", "Type": "integer"}]},
-            "Rows": [{"Data": [{"VarCharValue": "1"}]}, {"Data": [{"VarCharValue": "2"}]}],
+    responses = []
+    for i, page in enumerate(pages):
+        response = {
+            "ResultSet": {"Rows": [{"Data": [{"VarCharValue": v}]} for v in page]},
         }
-    }
+        if i == 0:
+            response["ResultSet"]["ResultSetMetadata"] = {
+                "ColumnInfo": [{"Name": "a", "Type": "integer"}]
+            }
+        if i + 1 < len(pages):
+            response["NextToken"] = f"token{i + 1}"
+        responses.append(response)
+    connection = MagicMock()
+    connection.client.get_query_results.side_effect = responses
     query_execution = AthenaQueryExecution(
         {
             "QueryExecution": {
@@ -40,7 +49,7 @@ async def _create_result_set(result_set_class):
         }
     )
     return await result_set_class.create(
-        connection, DefaultTypeConverter(), query_execution, 1000, RetryConfig()
+        connection, DefaultTypeConverter(), query_execution, 2, RetryConfig()
     )
 
 
@@ -61,3 +70,34 @@ async def test_sync_iteration_raises(self, result_set_class):
     async def test_async_iteration(self, result_set_class, expected):
         result_set = await _create_result_set(result_set_class)
         assert [row async for row in result_set] == expected
+
+    @pytest.mark.parametrize(
+        ("result_set_class", "expected"),
+        [
+            (AthenaAioResultSet, [[(1,)], [(2,), (3,), (4,)], [(5,)]]),
+            (AthenaAioDictResultSet, [[{"a": 1}], [{"a": 2}, {"a": 3}, {"a": 4}], [{"a": 5}]]),
+        ],
+    )
+    async def test_fetchmany_across_pages(self, result_set_class, expected):
+        result_set = await _create_result_set(result_set_class, [("1", "2"), ("3", "4"), ("5",)])
+        assert await result_set.fetchmany(1) == expected[0]
+        assert result_set.rownumber == 1
+        assert await result_set.fetchmany(3) == expected[1]
+        assert result_set.rownumber == 4
+        assert await result_set.fetchmany() == expected[2]
+        assert await result_set.fetchmany() == []
+        assert result_set.rownumber == 5
+
+    async def test_fetchall_across_pages(self):
+        result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3", "4"), ("5",)])
+        assert await result_set.fetchone() == (1,)
+        assert await result_set.fetchall() == [(2,), (3,), (4,), (5,)]
+        assert result_set.rownumber == 5
+        assert await result_set.fetchall() == []
+
+    async def test_fetchmany_size_not_integer(self):
+        """A size that is not an integer raises TypeError without consuming rows."""
+        result_set = await _create_result_set(AthenaAioResultSet)
+        with pytest.raises(TypeError):
+            await result_set.fetchmany(1.5)
+        assert await result_set.fetchall() == [(1,), (2,)]
diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py
index 27d7c3b1a..01dfbbfdf 100644
--- a/tests/pyathena/test_result_set.py
+++ b/tests/pyathena/test_result_set.py
@@ -4,6 +4,8 @@
 # See LICENSE or https://opensource.org/licenses/MIT.
 #
 # SPDX-License-Identifier: MIT
+import collections
+import sys
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -11,9 +13,10 @@
 from pyathena.aio.common import AioBaseCursor, WithAsyncFetch
 from pyathena.common import BaseCursor, CursorIterator
 from pyathena.converter import DefaultTypeConverter
+from pyathena.error import OperationalError
 from pyathena.model import AthenaQueryExecution
 from pyathena.result_set import AthenaResultSet, WithFetch, WithResultSet
-from pyathena.util import RetryConfig
+from pyathena.util import RetryConfig, override
 
 
 def _page(values, next_token=None):
@@ -27,7 +30,115 @@ def _row(*values):
     return {"Data": [{} if v is None else {"VarCharValue": v} for v in values]}
 
 
+def _fetching_result_set(*pages):
+    """Create a result set of a succeeded query from stubbed ``GetQueryResults`` pages.
+
+    No AWS calls; the first page is fetched at creation, the others on demand.
+
+    Args:
+        *pages: The values of an integer column ``a``, one tuple per page, or an
+            exception that the request for the page raises.
+
+    Returns:
+        The result set.
+    """
+    responses = []
+    for i, page in enumerate(pages):
+        if isinstance(page, Exception):
+            responses.append(page)
+            continue
+        response = _page(page, f"token{i + 1}" if i + 1 < len(pages) else None)
+        if i == 0:
+            response["ResultSet"]["ResultSetMetadata"] = {
+                "ColumnInfo": [{"Name": "a", "Type": "integer"}]
+            }
+        responses.append(response)
+    connection = MagicMock()
+    connection.client.get_query_results.side_effect = responses
+    return AthenaResultSet(
+        connection=connection,
+        converter=DefaultTypeConverter(),
+        query_execution=MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED),
+        arraysize=2,
+        retry_config=RetryConfig(),
+    )
+
+
+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 TestAthenaResultSet:
+    def test_fetchmany_across_pages(self):
+        result_set = _fetching_result_set(("1", "2"), ("3", "4"), ("5",))
+        assert result_set.fetchmany(1) == [(1,)]
+        assert result_set.rownumber == 1
+        assert result_set.fetchmany(3) == [(2,), (3,), (4,)]
+        assert result_set.rownumber == 4
+        assert result_set.fetchmany() == [(5,)]
+        assert result_set.fetchmany() == []
+        assert result_set.rownumber == 5
+        assert result_set.connection.client.get_query_results.call_count == 3
+
+    def test_fetchall_across_pages(self):
+        result_set = _fetching_result_set(("1", "2"), ("3", "4"), ("5",))
+        assert result_set.fetchone() == (1,)
+        assert result_set.fetchall() == [(2,), (3,), (4,), (5,)]
+        assert result_set.rownumber == 5
+        assert result_set.fetchall() == []
+        assert result_set.fetchone() is None
+
+    def test_fetchmany_size_above_maxsize(self):
+        result_set = _fetching_result_set(("1", "2"), ("3",))
+        assert result_set.fetchmany(sys.maxsize + 1) == [(1,), (2,), (3,)]
+        assert result_set.rownumber == 3
+
+    def test_fetchmany_size_not_integer(self):
+        """A size that is not an integer raises TypeError without consuming rows."""
+        result_set = _fetching_result_set(("1", "2"))
+        with pytest.raises(TypeError):
+            result_set.fetchmany(1.5)
+        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."""
+        result_set = _fetching_result_set(("1", "2"), RuntimeError("failed"))
+        with pytest.raises(OperationalError):
+            getattr(result_set, method)(*args)
+        assert result_set.rownumber == 2
+
+    def test_fetch_after_close(self):
+        result_set = _fetching_result_set(("1", "2"), ("3",))
+        result_set.close()
+        assert result_set.fetchmany() == []
+        assert result_set.fetchall() == []
+        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,)]
+        assert result_set.rownumber == 3
+
     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 8bba6ceda0dab74513693b0e0a32161ed4dc2067 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 15 Oct 2026 07:44:13 +0900
Subject: [PATCH 2/5] Keep the rows taken before a close() from another thread

The buffered rows were read and then cleared, or popped for a count
read earlier, so a close() from another thread meanwhile raised
IndexError or dropped the rows. The buffer is now replaced before its
rows are read, and popleft() stops at an empty buffer.

fetchone() of the base, Arrow, S3FS, and asyncio result sets counts its
row with _count_rows() too, so that rownumber stays None after close()
whichever fetch method takes the rows.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/aio/result_set.py        |  7 +++--
 pyathena/arrow/result_set.py      |  7 +++--
 pyathena/pandas/result_set.py     | 11 ++------
 pyathena/polars/result_set.py     | 11 ++------
 pyathena/result_set.py            | 31 +++++++++++++---------
 pyathena/s3fs/result_set.py       |  7 +++--
 tests/pyathena/test_result_set.py | 43 +++++++++++++++++++++++++++++++
 7 files changed, 75 insertions(+), 42 deletions(-)

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index 95dcbbe3d..8211b2745 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -187,10 +187,9 @@ async def fetchone(  # type: ignore[override]
             await self._async_fetch()
         if not self._rows:
             return None
-        if self._rownumber is None:
-            self._rownumber = 0
-        self._rownumber += 1
-        return self._rows.popleft()
+        row = self._rows.popleft()
+        self._count_rows(1)
+        return row
 
     @override
     async def fetchmany(  # type: ignore[override]
diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py
index 3b9b4597c..4ed050ac1 100644
--- a/pyathena/arrow/result_set.py
+++ b/pyathena/arrow/result_set.py
@@ -224,10 +224,9 @@ def fetchone(
             self._fetch()
         if not self._rows:
             return None
-        if self._rownumber is None:
-            self._rownumber = 0
-        self._rownumber += 1
-        return self._rows.popleft()
+        row = self._rows.popleft()
+        self._count_rows(1)
+        return row
 
     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 586c39dba..81093b6a0 100644
--- a/pyathena/pandas/result_set.py
+++ b/pyathena/pandas/result_set.py
@@ -949,16 +949,9 @@ def fetchone(
 
     @override
     def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
-        """Take the next rows and count them in ``rownumber``.
+        """Take the next rows from the row iterator in bulk.
 
-        The rows are taken from the row iterator in bulk. The rows taken before an
-        exception are counted too.
-
-        Args:
-            size: The maximum number of rows, or None for all remaining rows.
-
-        Returns:
-            The rows, fewer than ``size`` when the result is exhausted.
+        The rows taken before an exception are counted in ``rownumber`` too.
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
         try:
diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py
index 17fc35e7a..2c63c1ba6 100644
--- a/pyathena/polars/result_set.py
+++ b/pyathena/polars/result_set.py
@@ -593,16 +593,9 @@ def fetchone(
 
     @override
     def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
-        """Take the next rows and count them in ``rownumber``.
+        """Take the next rows from the row iterator in bulk.
 
-        The rows are taken from the row iterator in bulk. The rows taken before an
-        exception are counted too.
-
-        Args:
-            size: The maximum number of rows, or None for all remaining rows.
-
-        Returns:
-            The rows, fewer than ``size`` when the result is exhausted.
+        The rows taken before an exception are counted in ``rownumber`` too.
         """
         rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = []
         try:
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index 61c3916e8..ca84144cf 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -3,10 +3,12 @@
 from __future__ import annotations
 
 import collections
+import contextlib
 import logging
 import operator
 import sys
 from datetime import datetime
+from itertools import islice
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -493,10 +495,9 @@ def fetchone(
             self._fetch()
         if not self._rows:
             return None
-        if self._rownumber is None:
-            self._rownumber = 0
-        self._rownumber += 1
-        return self._rows.popleft()
+        row = self._rows.popleft()
+        self._count_rows(1)
+        return row
 
     @override
     def fetchmany(
@@ -559,19 +560,25 @@ def _take_buffered_rows(
     ) -> None:
         """Move the buffered rows to ``rows`` and count them in ``rownumber``.
 
+        A ``close()`` from another thread meanwhile empties the buffer, and the rows
+        moved before it are kept.
+
         Args:
             rows: The list that receives the rows.
             limit: The maximum number of rows, or None for all buffered rows.
         """
-        count = len(self._rows)
-        if limit is not None and limit < count:
-            count = limit
-            popleft = self._rows.popleft
-            rows.extend([popleft() for _ in range(count)])
+        start = len(rows)
+        if limit is not None and limit < len(self._rows):
+            # A row is never None, so popleft() runs until the limit or an empty
+            # buffer, and extend() keeps the rows that it appended before that.
+            with contextlib.suppress(IndexError):
+                rows.extend(islice(iter(self._rows.popleft, None), limit))
         else:
-            rows.extend(self._rows)
-            self._rows.clear()
-        self._count_rows(count)
+            # The buffer is replaced, not cleared after it is read, so that a close()
+            # in between empties only the new buffer.
+            buffered, self._rows = self._rows, collections.deque()
+            rows.extend(buffered)
+        self._count_rows(len(rows) - start)
 
     def _count_rows(self, count: int) -> None:
         """Count the rows that a fetch took in ``rownumber``.
diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py
index b5caaa632..89e5a1f32 100644
--- a/pyathena/s3fs/result_set.py
+++ b/pyathena/s3fs/result_set.py
@@ -266,10 +266,9 @@ def fetchone(
             self._fetch()
         if not self._rows:
             return None
-        if self._rownumber is None:
-            self._rownumber = 0
-        self._rownumber += 1
-        return self._rows.popleft()
+        row = self._rows.popleft()
+        self._count_rows(1)
+        return row
 
     @override
     def close(self) -> None:
diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py
index 01dfbbfdf..d99aecdce 100644
--- a/tests/pyathena/test_result_set.py
+++ b/tests/pyathena/test_result_set.py
@@ -85,6 +85,38 @@ def fetchone(self):
         return row
 
 
+class _ClosingDeque(collections.deque):
+    """A buffer that closes its result set when it is first read, as another thread may."""
+
+    def __init__(self, rows, result_set):
+        """Initialize the buffer.
+
+        Args:
+            rows: The buffered rows.
+            result_set: The result set to close.
+        """
+        super().__init__(rows)
+        self._result_set = result_set
+        self._closed = False
+
+    def _close_once(self):
+        """Close the result set on the first read."""
+        if not self._closed:
+            self._closed = True
+            self._result_set.close()
+
+    @override
+    def popleft(self):
+        row = super().popleft()
+        self._close_once()
+        return row
+
+    @override
+    def __iter__(self):
+        self._close_once()
+        return super().__iter__()
+
+
 class TestAthenaResultSet:
     def test_fetchmany_across_pages(self):
         result_set = _fetching_result_set(("1", "2"), ("3", "4"), ("5",))
@@ -132,6 +164,17 @@ def test_fetch_after_close(self):
         assert result_set.fetchall() == []
         assert result_set.rownumber is None
 
+    @pytest.mark.parametrize(
+        ("method", "args", "expected"),
+        [("fetchmany", (2,), [(1,)]), ("fetchall", (), [(1,), (2,), (3,)])],
+    )
+    def test_fetch_with_close_during_take(self, method, args, expected):
+        """A close() while the buffered rows are taken keeps the rows taken before it."""
+        result_set = _fetching_result_set(("1", "2", "3"))
+        result_set._rows = _ClosingDeque(result_set._rows, result_set)
+        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,)])

From 4744d9ba3fdb6eb93e837e74c4448c57cb2f5ff0 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 15 Oct 2026 07:48:59 +0900
Subject: [PATCH 3/5] Simplify the bulk take and keep fetchone() at its per-row
 cost

- _take_buffered_rows() takes the total size and works out its limit.
- A try statement replaces contextlib.suppress(), which cost about
  150 ns per partial take.
- fetchone() updates rownumber inline, None-safe as _count_rows() is,
  since the method call added about 20 ns per row to iteration.
- The tests share one builder of stubbed GetQueryResults pages, and
  the closing buffer stub drops a guard that never applied.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/aio/result_set.py            |  7 +++++--
 pyathena/arrow/result_set.py          |  5 ++++-
 pyathena/result_set.py                | 25 ++++++++++++++-----------
 pyathena/s3fs/result_set.py           |  5 ++++-
 tests/pyathena/aio/test_result_set.py | 15 ++-------------
 tests/pyathena/test_result_set.py     | 27 +++++----------------------
 tests/pyathena/util.py                | 27 +++++++++++++++++++++++++++
 7 files changed, 61 insertions(+), 50 deletions(-)

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index 8211b2745..cf0256cf6 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -188,7 +188,10 @@ async def fetchone(  # type: ignore[override]
         if not self._rows:
             return None
         row = self._rows.popleft()
-        self._count_rows(1)
+        # 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
@@ -239,7 +242,7 @@ 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, None if size is None else size - len(rows))
+                self._take_buffered_rows(rows, size)
             else:
                 row = await self.fetchone()
                 if row is None:
diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py
index 4ed050ac1..53939437e 100644
--- a/pyathena/arrow/result_set.py
+++ b/pyathena/arrow/result_set.py
@@ -225,7 +225,10 @@ def fetchone(
         if not self._rows:
             return None
         row = self._rows.popleft()
-        self._count_rows(1)
+        # 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:
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index ca84144cf..edabfc217 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import collections
-import contextlib
 import logging
 import operator
 import sys
@@ -496,7 +495,10 @@ def fetchone(
         if not self._rows:
             return None
         row = self._rows.popleft()
-        self._count_rows(1)
+        # 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
@@ -535,8 +537,8 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
         """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 a
-        subclass that overrides it does.
+        takes the next row, so that it fetches the next rows into the buffer as the
+        result set does.
 
         Args:
             size: The maximum number of rows, or None for all remaining rows.
@@ -547,7 +549,7 @@ 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, None if size is None else size - len(rows))
+                self._take_buffered_rows(rows, size)
             else:
                 row = self.fetchone()
                 if row is None:
@@ -556,23 +558,24 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any
         return rows
 
     def _take_buffered_rows(
-        self, rows: list[tuple[Any | None, ...] | dict[Any, Any | None]], limit: int | None
+        self, rows: list[tuple[Any | None, ...] | dict[Any, Any | None]], size: int | None
     ) -> None:
         """Move the buffered rows to ``rows`` and count them in ``rownumber``.
 
-        A ``close()`` from another thread meanwhile empties the buffer, and the rows
-        moved before it are kept.
-
         Args:
             rows: The list that receives the rows.
-            limit: The maximum number of rows, or None for all buffered rows.
+            size: The maximum length of ``rows``, or None for all buffered rows.
         """
         start = len(rows)
+        limit = None if size is None else size - start
         if limit is not None and limit < len(self._rows):
             # A row is never None, so popleft() runs until the limit or an empty
             # buffer, and extend() keeps the rows that it appended before that.
-            with contextlib.suppress(IndexError):
+            # try costs nothing without an exception, unlike contextlib.suppress().
+            try:  # noqa: SIM105
                 rows.extend(islice(iter(self._rows.popleft, None), limit))
+            except IndexError:
+                pass
         else:
             # The buffer is replaced, not cleared after it is read, so that a close()
             # in between empties only the new buffer.
diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py
index 89e5a1f32..1ae07684a 100644
--- a/pyathena/s3fs/result_set.py
+++ b/pyathena/s3fs/result_set.py
@@ -267,7 +267,10 @@ def fetchone(
         if not self._rows:
             return None
         row = self._rows.popleft()
-        self._count_rows(1)
+        # 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
diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py
index cc8434e0a..61c27e8f1 100644
--- a/tests/pyathena/aio/test_result_set.py
+++ b/tests/pyathena/aio/test_result_set.py
@@ -12,6 +12,7 @@
 from pyathena.converter import DefaultTypeConverter
 from pyathena.model import AthenaQueryExecution
 from pyathena.util import RetryConfig
+from tests.pyathena.util import get_query_results_responses
 
 
 async def _create_result_set(result_set_class, pages=(("1", "2"),)):
@@ -25,20 +26,8 @@ async def _create_result_set(result_set_class, pages=(("1", "2"),)):
     Returns:
         The result set.
     """
-    responses = []
-    for i, page in enumerate(pages):
-        response = {
-            "ResultSet": {"Rows": [{"Data": [{"VarCharValue": v}]} for v in page]},
-        }
-        if i == 0:
-            response["ResultSet"]["ResultSetMetadata"] = {
-                "ColumnInfo": [{"Name": "a", "Type": "integer"}]
-            }
-        if i + 1 < len(pages):
-            response["NextToken"] = f"token{i + 1}"
-        responses.append(response)
     connection = MagicMock()
-    connection.client.get_query_results.side_effect = responses
+    connection.client.get_query_results.side_effect = get_query_results_responses(*pages)
     query_execution = AthenaQueryExecution(
         {
             "QueryExecution": {
diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py
index d99aecdce..c0e11238c 100644
--- a/tests/pyathena/test_result_set.py
+++ b/tests/pyathena/test_result_set.py
@@ -17,6 +17,7 @@
 from pyathena.model import AthenaQueryExecution
 from pyathena.result_set import AthenaResultSet, WithFetch, WithResultSet
 from pyathena.util import RetryConfig, override
+from tests.pyathena.util import get_query_results_responses
 
 
 def _page(values, next_token=None):
@@ -42,19 +43,8 @@ def _fetching_result_set(*pages):
     Returns:
         The result set.
     """
-    responses = []
-    for i, page in enumerate(pages):
-        if isinstance(page, Exception):
-            responses.append(page)
-            continue
-        response = _page(page, f"token{i + 1}" if i + 1 < len(pages) else None)
-        if i == 0:
-            response["ResultSet"]["ResultSetMetadata"] = {
-                "ColumnInfo": [{"Name": "a", "Type": "integer"}]
-            }
-        responses.append(response)
     connection = MagicMock()
-    connection.client.get_query_results.side_effect = responses
+    connection.client.get_query_results.side_effect = get_query_results_responses(*pages)
     return AthenaResultSet(
         connection=connection,
         converter=DefaultTypeConverter(),
@@ -86,7 +76,7 @@ def fetchone(self):
 
 
 class _ClosingDeque(collections.deque):
-    """A buffer that closes its result set when it is first read, as another thread may."""
+    """A buffer that closes its result set when it is read, as another thread may."""
 
     def __init__(self, rows, result_set):
         """Initialize the buffer.
@@ -97,23 +87,16 @@ def __init__(self, rows, result_set):
         """
         super().__init__(rows)
         self._result_set = result_set
-        self._closed = False
-
-    def _close_once(self):
-        """Close the result set on the first read."""
-        if not self._closed:
-            self._closed = True
-            self._result_set.close()
 
     @override
     def popleft(self):
         row = super().popleft()
-        self._close_once()
+        self._result_set.close()
         return row
 
     @override
     def __iter__(self):
-        self._close_once()
+        self._result_set.close()
         return super().__iter__()
 
 
diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py
index 50f6712e5..ff6a16d6c 100644
--- a/tests/pyathena/util.py
+++ b/tests/pyathena/util.py
@@ -30,6 +30,33 @@
 )
 
 
+def get_query_results_responses(*pages):
+    """Build stubbed ``GetQueryResults`` responses of an integer column ``a``.
+
+    Args:
+        *pages: The values of the column, one sequence per page, or an exception that
+            the request for the page raises.
+
+    Returns:
+        The responses, the first with the result set metadata and each but the last
+        with a ``NextToken``.
+    """
+    responses = []
+    for i, page in enumerate(pages):
+        if isinstance(page, Exception):
+            responses.append(page)
+            continue
+        response = {"ResultSet": {"Rows": [{"Data": [{"VarCharValue": v}]} for v in page]}}
+        if i == 0:
+            response["ResultSet"]["ResultSetMetadata"] = {
+                "ColumnInfo": [{"Name": "a", "Type": "integer"}]
+            }
+        if i + 1 < len(pages):
+            response["NextToken"] = f"token{i + 1}"
+        responses.append(response)
+    return responses
+
+
 def read_query(name, **kwargs):
     template = _queries.get_template(name)
     return [q.strip() for q in template.render(**kwargs).split(";") if q and q.strip()]

From 2d63ab60a505af265025c80b4da22116049a0fdb Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 15 Oct 2026 08:06:35 +0900
Subject: [PATCH 4/5] Document TypeError only for a size not replaced with
 arraysize

A size that is not positive, such as -1.5, is replaced with arraysize
before it is checked, as before; a string still raises TypeError.

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

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index cf0256cf6..1dddd6968 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -208,7 +208,7 @@ async def fetchmany(  # type: ignore[override]
             The rows, fewer than ``size`` when the result is exhausted.
 
         Raises:
-            TypeError: If ``size`` is not an integer.
+            TypeError: If ``size`` is positive and not an integer.
         """
         if not size or size <= 0:
             size = self._arraysize
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index edabfc217..2139bfea7 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -515,7 +515,7 @@ def fetchmany(
             The rows, fewer than ``size`` when the result is exhausted.
 
         Raises:
-            TypeError: If ``size`` is not an integer.
+            TypeError: If ``size`` is positive and not an integer.
         """
         if not size or size <= 0:
             size = self._arraysize

From 8391a497a98a6cb4de697202c9b5a7a683da5416 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 15 Oct 2026 08:09:18 +0900
Subject: [PATCH 5/5] Document TypeError for any size not replaced with
 arraysize

A nonempty string fails the comparison with zero, and a positive
non-integer fails the integer check; both raise TypeError.

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

diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py
index 1dddd6968..b55627ab2 100644
--- a/pyathena/aio/result_set.py
+++ b/pyathena/aio/result_set.py
@@ -208,7 +208,7 @@ async def fetchmany(  # type: ignore[override]
             The rows, fewer than ``size`` when the result is exhausted.
 
         Raises:
-            TypeError: If ``size`` is positive and not an integer.
+            TypeError: If ``size`` is not an integer and not replaced with ``arraysize``.
         """
         if not size or size <= 0:
             size = self._arraysize
diff --git a/pyathena/result_set.py b/pyathena/result_set.py
index 2139bfea7..a72271fab 100644
--- a/pyathena/result_set.py
+++ b/pyathena/result_set.py
@@ -515,7 +515,7 @@ def fetchmany(
             The rows, fewer than ``size`` when the result is exhausted.
 
         Raises:
-            TypeError: If ``size`` is positive and not an integer.
+            TypeError: If ``size`` is not an integer and not replaced with ``arraysize``.
         """
         if not size or size <= 0:
             size = self._arraysize