Share the buffer refill between fetchone() and the bulk fetches of the result sets

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, as after ``close()``.
 
         Raises:
-            ProgrammingError: If there is no next page.
+            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:
-            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

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

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, as after ``close()``.
 
         Raises:
-            ProgrammingError: If there is no next page.
+            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:
-            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(

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."""

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,27 @@ 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
+
+    @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

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.
 

tests/pyathena/s3fs/test_result_set.py

@@ -121,6 +121,17 @@ 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, _ = _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
+        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"

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,11 @@ 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_all_rows_skips_column_labels_only_on_first_page(self):
         """A later page can start with a data row equal to the column labels.