diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py index 27642580a..b55627ab2 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, @@ -186,10 +187,12 @@ 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() + # 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 async def fetchmany( # type: ignore[override] @@ -203,17 +206,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 and not replaced with ``arraysize``. """ 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 +223,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, size) else: - break + row = await self.fetchone() + if row is None: + break + rows.append(row) return rows @override diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 3b9b4597c..53939437e 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -224,10 +224,12 @@ 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() + # As _count_rows() does, inline on the per-row path. + rownumber = self._rownumber + if rownumber is not None: + self._rownumber = rownumber + 1 + return row def _open_csv(self, location: str) -> Any: """Open the CSV result file for ``pyarrow.csv.read_csv()``. diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index e0d5af173..b52e60f2e 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 import threading from collections import abc from collections.abc import Callable, Iterable, Iterator @@ -987,32 +985,10 @@ 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. - - Args: - size: The maximum number of rows, or None for all remaining rows. + """Take the next rows from the row iterator in bulk. - 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: @@ -1023,20 +999,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 2bc77b264..68b8696ba 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 import threading from collections import abc from collections.abc import Callable, Iterable, Iterator @@ -631,32 +629,10 @@ 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. - - Args: - size: The maximum number of rows, or None for all remaining rows. + """Take the next rows from the row iterator in bulk. - 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: @@ -667,20 +643,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..a72271fab 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -4,7 +4,10 @@ import collections import logging +import operator +import sys from datetime import datetime +from itertools import islice from typing import ( TYPE_CHECKING, Any, @@ -491,10 +494,12 @@ 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() + # 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 fetchmany( @@ -508,17 +513,14 @@ def fetchmany( Returns: The rows, fewer than ``size`` when the result is exhausted. + + Raises: + TypeError: If ``size`` is not an integer and not replaced with ``arraysize``. """ 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 +531,72 @@ 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 the + result set 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, size) 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]], size: int | None + ) -> None: + """Move the buffered rows to ``rows`` and count them in ``rownumber``. + + Args: + rows: The list that receives the 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. + # 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. + 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``. + + ``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/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py index b5caaa632..1ae07684a 100644 --- a/pyathena/s3fs/result_set.py +++ b/pyathena/s3fs/result_set.py @@ -266,10 +266,12 @@ 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() + # 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: diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py index b34910afb..61c27e8f1 100644 --- a/tests/pyathena/aio/test_result_set.py +++ b/tests/pyathena/aio/test_result_set.py @@ -12,24 +12,22 @@ 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): - """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"}]}], - } - } + connection.client.get_query_results.side_effect = get_query_results_responses(*pages) query_execution = AthenaQueryExecution( { "QueryExecution": { @@ -40,7 +38,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 +59,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..c0e11238c 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,11 @@ 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 +from tests.pyathena.util import get_query_results_responses def _page(values, next_token=None): @@ -27,7 +31,140 @@ 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. + """ + connection = MagicMock() + connection.client.get_query_results.side_effect = get_query_results_responses(*pages) + 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 _ClosingDeque(collections.deque): + """A buffer that closes its result set when it is 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 + + @override + def popleft(self): + row = super().popleft() + self._result_set.close() + return row + + @override + def __iter__(self): + self._result_set.close() + return super().__iter__() + + 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 + + @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,)]) + 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. diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 2b2784f6d..e58a37f46 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -31,6 +31,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 overlapping_calls(calls, timeout=5): """Return a callable that records the peak number of its overlapping calls.