Skip to content

Take buffered rows in bulk in fetchmany() and fetchall() #1187

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
laughingman7743 merged 6 commits into from
Oct 15, 2026
Merged

Take buffered rows in bulk in fetchmany() and fetchall() #1187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c5f2fc6
Take buffered rows in bulk in fetchmany() and fetchall()
laughingman7743 Oct 14, 2026
8bba6ce
Keep the rows taken before a close() from another thread
laughingman7743 Oct 14, 2026
4744d9b
Simplify the bulk take and keep fetchone() at its per-row cost
laughingman7743 Oct 14, 2026
2d63ab6
Document TypeError only for a size not replaced with arraysize
laughingman7743 Oct 14, 2026
8391a49
Document TypeError for any size not replaced with arraysize
laughingman7743 Oct 14, 2026
e817537
Merge remote-tracking branch 'origin/master' into perf/1171-bulk-fetch
laughingman7743 Oct 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 35 additions & 18 deletions pyathena/aio/result_set.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from __future__ import annotations

import logging
import operator
from typing import (
TYPE_CHECKING,
Any,
Expand Down Expand Up @@ -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]
Expand All @@ -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]
Expand All @@ -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
Expand Down
10 changes: 6 additions & 4 deletions pyathena/arrow/result_set.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()``.
Expand Down
42 changes: 2 additions & 40 deletions pyathena/pandas/result_set.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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.

Expand Down
42 changes: 2 additions & 40 deletions pyathena/polars/result_set.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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.

Expand Down
95 changes: 77 additions & 18 deletions pyathena/result_set.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.

@laughingman7743 laughingman7743 Oct 14, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. At c5f2fc6 the bulk path left rownumber None after close() (_count_rows()), while fetchone() turned None into 0 and counted. Repaired in 8bba6ce: the base, Arrow, S3FS, and asyncio fetchone() keep None as well. A None rownumber with buffered rows occurs only after close(), because rows are read only for a succeeded query, which starts at 0.

  2. Calling self._count_rows(1) in fetchone() added about 20 ns per row to plain iteration. Repaired in 4744d9b by inlining the same None-safe update: a fetchone() loop over 1 million rows measured 70.1 ns per row, as on 27ceea3.

rownumber = self._rownumber
if rownumber is not None:
self._rownumber = rownumber + 1
return row

@override
def fetchmany(
Expand All @@ -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))

@laughingman7743 laughingman7743 Oct 14, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round two (claims, caller compatibility, operational effects, evidence)
Base 27ceea3, head 4744d9b. Covered: the PR body, the three commit messages, and the changed docstrings and comments; callers of the fetch methods; docs/. Result: FINDINGS in the PR description, corrected; no code changes.

Corrected in the PR body:

  1. "Unchanged: fetchone() of every result set" stopped holding after 8bba6ce: the base, Arrow, S3FS, and asyncio fetchone() now leave rownumber None after close(). The body states this, and that the old None → 1 count happened only when a fetchone() took a row while another thread closed the result set, since close() empties the buffer.
  2. The benchmark table was from c5f2fc6. Re-measured at 4744d9b against 27ceea3 in one session, with the per-row fetchone() loop (70.5 vs 68.2 ns per row) and the fetchmany(1) per-call cost (414 vs 173 ns) added.
  3. "Athena does not return zero-column rows" was stated as fact; narrowed to what I know.
  4. "The new tests pass with the previous source" was from before the close() test: with 27ceea3's source, every new test passes except the fetchall() case of the close() test, which returns [] there.

Checked, holding:

  • Callers: size as an int subclass or NumPy integer works through operator.index(), as through range() before; a float or string still raises TypeError. Return shapes are unchanged. Nothing outside the result sets holds _rows, so replacing the deque object is safe; AthenaS3FSResultSet.__init__ copies it with list().
  • Concurrency: threadsafety = 2, so two fetches sharing one cursor are out of contract; close() from another thread is the case covered.
  • Comments: a row is a tuple or dict, never None (the iter(popleft, None) sentinel); list.extend() keeps what an iterator appended before an exception (exercised by the fetchmany() close() test); try costs nothing without an exception on Python 3.11+, the floor in pyproject.toml; islice() takes no size above sys.maxsize (pandas/Polars tests).
  • docs/: no prose describes per-row fetching or rownumber after close().
  • Evidence at 4744d9b: just lint; 392 passed, 10 skipped in the result set unit tests; 150 passed in the live AWS fetch tests of the sync and asyncio cursors (-n 1, local). The full suite is left to CI on Ready.
  • No AWS request count changes: pages are still requested only when the buffer is empty.


@override
def fetchall(
Expand All @@ -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]]:

@laughingman7743 laughingman7743 Oct 14, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deferred, with reasons:

  • The sync _take_rows() inherited by AthenaAioResultSet would append coroutines if called; nothing calls it, since AthenaAioResultSet overrides fetchmany() and fetchall().
  • A subclass whose fetchone() transforms the buffered rows would now get raw rows from fetchmany()/fetchall(). No PyAthena result set does this, and the buffer is private; a subclass whose fetchone() does not use the buffer keeps working (test_fetch_with_fetchone_override).
  • The sync and asyncio take loops differ only in await, following the existing sync/asyncio split; both share _take_buffered_rows().
  • Proposal, not in this PR: one refill hook (for example _fetch_next(): the base fetches while NextToken remains, Arrow and S3FS always) used by fetchone() and _take_rows() would remove the Arrow and S3FS fetchone() overrides. It changes the subclasses' method roles, so it needs the maintainer's approval.
  • Pre-existing, outside the diff (pyathena/arrow/result_set.py:199): ParquetDataset.read() of an UNLOAD output with an empty Parquet file gives a zero-row batch (checked locally with PyArrow 26.0.0: batches [2, 0, 1]), so Arrow's fetchone() returns None before the end and fetches stop early, before and after this PR. To be filed separately.

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

@laughingman7743 laughingman7743 Oct 14, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. contextlib.suppress(IndexError) cost about 150 ns per partial take; replaced with try in 4744d9b (noqa: SIM105 with the reason). Accepted: fetchmany(1) costs 425 ns per call against 177 ns on 27ceea3, a fixed per-call cost, while fetchmany(10) takes 58 ns per row against 97 ns. The SQLAlchemy dialect has no server-side cursors, so SQLAlchemy does not grow fetchmany() from size 1.

  2. Simplifications in 4744d9b: _take_buffered_rows() takes the total size instead of both callers computing the limit; one stubbed GetQueryResults page builder in tests/pyathena/util.py for the sync and asyncio tests; the closing buffer stub drops a guard that never applied.

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

@laughingman7743 laughingman7743 Oct 14, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round one (implementation behavior, public contracts, simplicity, regression coverage)
Base 27ceea3, initial head c5f2fc6; repairs 8bba6ce and 4744d9b. Covered: base, asyncio, pandas, Polars, Arrow, and S3FS fetch paths and the new tests; my own pass plus /code-review (high) and /simplify (reuse, simplification, efficiency, altitude). Result: FINDINGS, repaired as below.

  1. _take_buffered_rows() at c5f2fc6 read the buffer and then cleared it, or popped a count read earlier. A close() from another thread in between (for example, cursor.close() after an asyncio wait_for() timeout while asyncio.to_thread(result_set.fetchall) still runs) dropped the rows, or raised IndexError. Repaired in 8bba6ce: the buffer is replaced before its rows are read, and popleft() stops at an empty buffer, keeping the rows taken before. test_fetch_with_close_during_take fails on c5f2fc6 (IndexError; [] instead of the rows) and passes now.

Kept: two branches. A single islice(iter(popleft, None), None) branch measured 0.029 s against 0.009 s for the swap over 1.5 million rows.

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:
Expand Down
10 changes: 6 additions & 4 deletions pyathena/s3fs/result_set.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading
Loading