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

[laughingman7743]

WHAT

fetchmany() and fetchall() of the result sets that buffer converted rows take those rows in bulk instead of calling fetchone() once per row.

• AthenaResultSet.fetchmany() and fetchall() call the new _take_rows(). It moves the buffered rows to the result in bulk (all of them, or as many as size allows) and counts them in rownumber once per batch. When the buffer is empty, it calls fetchone() once, so that the next rows are fetched by the result set's own fetchone(): the base class requests the next page while NextToken remains, and ArrowCursor's and S3FSCursor's result sets read their next batch. A subclass whose fetchone() does not use the buffer gets its rows from fetchone() one at a time, as before.
• AthenaAioResultSet (the default asyncio cursor) takes its rows the same way, awaiting fetchone() only when the buffer is empty.
• The pandas and Polars result sets now override only _take_rows() (their islice() over the row tuples since Return PandasCursor rows as tuples without a dictionary per row #1167 and Return PolarsCursor rows as tuples without a dictionary per row #1177). Their identical fetchmany(), fetchall(), and _count_rows() moved to AthenaResultSet.
• A close() from another thread while the rows are taken, for example cursor.close() after an asyncio wait_for() timeout while asyncio.to_thread(result_set.fetchall) still runs, keeps the rows taken before it and leaves rownumber None. fetchone() of the base, Arrow, S3FS, and asyncio result sets also leaves rownumber None after close() now, as the pandas and Polars result sets did; before, a fetchone() that took a row while another thread closed the result set turned it into 1.

Unchanged: size handling (None or not positive → arraysize, then any other non-integer → TypeError without consuming rows); rows taken before an exception stay counted in rownumber; fetches after close() return [].

One difference: the previous loops stopped at the first falsy row (if row:), so an empty tuple or dictionary would have ended fetchmany()/fetchall() early. Rows now end only at None, as fetchone() signals. An empty row needs a result without columns, and I know of no Athena statement whose result has rows but no columns.

WHY

Closes #1171.
Measured locally at 4744d9b and at its base 27ceea3 in one session (Apple M1 Pro, Python 3.15.0), 1.5 million already converted 15-column rows in the buffer, median of 3 runs, without the 0.034 s it took to fill the buffer:

Operation	Before	After
Sync fetchall()	0.120 s	0.026 s
Sync fetchmany(10000) until exhausted	0.121 s	0.035 s
Asyncio fetchall()	0.178 s	0.028 s
Asyncio fetchmany(10000) until exhausted	0.180 s	0.032 s

On 1 million one-column rows, median of 5 runs: a fetchone() loop takes 70.5 ns per row against 68.2 ns before; fetchmany(10) takes 56.6 ns per row against 95.8 ns; fetchmany(1) takes 414 ns per call against 173 ns, a fixed per-call cost. The SQLAlchemy dialect has no server-side cursors, so SQLAlchemy does not grow fetchmany() sizes from 1.

As #1171 notes, the gain is small next to the work that produces the rows: GetQueryResults paging for the default cursors, to_pylist() conversion for ArrowCursor, and per-value converters for S3FSCursor. The pandas and Polars row paths are unchanged.

TEST

Tested at 4744d9b unless noted; 2d63ab6 and 8391a49 change only two docstrings.

• CI at e817537, the merge of master 4bc1d3f (run 38750369809, live AWS): test / run (3.15) 4211 passed, 22 skipped; test-sqla / run (3.15) and test-sqla-async / run (3.15) 589 passed, 759 skipped each. The merge resolved import conflicts with Defer closing a chunk reader while another thread reads it #1182 in the pandas and Polars result sets and kept both helpers in tests/pyathena/util.py; Defer closing a chunk reader while another thread reads it #1182 changed only the chunk readers' close, not the fetch methods. After the merge: just lint passed; the result set unit tests with tests/pyathena/filesystem/test_s3_executor.py, 424 passed, 10 skipped.

• CI at 8391a49 (run 38653687756): test 4156 passed, 22 skipped; test-sqla and test-sqla-async 589 passed, 759 skipped each.

• just lint: passed.

• uv run --env-file .env pytest tests/pyathena/test_result_set.py tests/pyathena/aio/test_result_set.py tests/pyathena/{pandas,polars,arrow,s3fs}/test_result_set.py tests/pyathena/aio/test_common.py: 392 passed, 10 skipped.

• uv run --env-file .env pytest -n 1 -k "fetch or iter or arraysize or dict" over the sync and asyncio cursor tests of the default, Arrow, S3FS, pandas, and Polars cursors (live AWS): 150 passed.

New tests:

• TestAthenaResultSet: fetchmany() and fetchall() across GetQueryResults pages with rownumber and the number of requests, sizes above sys.maxsize and not integers, rows counted before a failed page request, fetches after close(), a close() while the buffered rows are taken (both cases fail on c5f2fc6, the first commit, with IndexError and []), and a subclass whose fetchone() does not use the buffer. With the pyathena/ source of 27ceea3, every new test except the fetchall() case of the close() test passes (that case returns [] there), which shows that the tests pin the existing fetch semantics.
• TestAthenaAioResultSet: fetchmany() (tuple and dictionary rows) and fetchall() across pages, and a non-integer size.

Not run locally: the full just test pyathena suite (CI ran it). The remaining close() window inside fetchone(), between its buffer check and popleft(), predates this PR and has no test.

🤖 Generated with Claude Code

[laughingman7743]

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.

[laughingman7743]

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

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

[laughingman7743]

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

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

[laughingman7743]

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.

[laughingman7743]

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.

[laughingman7743]

Independent review (relayed): Codex CLI 0.162.1, model gpt-6-astra, reasoning effort max, codex exec -s read-only, session 01a13ca6-cc98-72f1-9271-64db706e376c. Static review of a detached snapshot at 4744d9b against merge-base 27ceea3, given the literal diff and the repository conventions, without the PR text or prior findings; no edits, builds, tests, or GitHub access.

Covered (reviewer): base/dict and asyncio fetching (row order, sizes, rownumber, exhaustion, close interleavings, page failures); Arrow/S3FS refills, pandas/Polars iterator overrides, sync and asyncio cursor callers, SQLAlchemy adapters; the changed tests, response helper, docstrings, and comments. Noted limits: the close injection covers selected interleavings and does not establish general thread safety; the new tests do not directly cover asyncio failure/close cases or reduced fetchone() call counts.

Result: FINDINGS, one P3, no introduced runtime regression identified.

• P3: the new Raises: TypeError (here and pyathena/aio/result_set.py:211) overstated validation. With arraysize=2 and three buffered rows, fetchmany(-1.5) returns two rows, because a size that is not positive is replaced with arraysize before it is checked. Behavior pre-existing; the description new.

Verified: -1.5 and 0.0 return arraysize rows; 1.5 and "3" raise TypeError; True takes one row. Repaired in 2d63ab6: "If size is not an integer and not replaced with arraysize." in both docstrings; the PR body's size handling says the same.

[laughingman7743]

Self-review of the repair 4744d9b..2d63ab6 (both perspectives). Behavior: docstrings only, no code or test change; just lint passed. Claims: the new wording matches the measured cases (-1.5 and 0.0 → arraysize rows; 1.5 and "3" → TypeError), for the base and the asyncio fetchmany(), which share the same normalization; the PR body's size handling was corrected the same way. CLEAN.

[laughingman7743]

Correction: 2d63ab6 committed only the intermediate wording "If size is positive and not an integer." (the amended docstring stayed uncommitted). The Codex follow-up (session 01a13cab-f015-7ed3-b667-b1e1117081c6, static, read-only) on 4744d9b..2d63ab6 found that it omits a nonempty string such as "2", which raises TypeError in the comparison with zero. Repaired in 8391a49: "If size is not an integer and not replaced with arraysize." in both docstrings, which covers a positive float (integer check), a nonempty string (comparison), and an empty string, negative float, None, or 0 (replaced). just lint passed; docstrings only.

[laughingman7743]

Independent follow-up on 2d63ab6..8391a49: Codex (gpt-6-astra, codex exec -s read-only, session 01a13cad-ae18-7770-8fa3-1625f5c96dab), static review of a detached snapshot at 8391a49. Result: CLEAN. The wording holds for None, 0, negative int, positive and negative non-integers, bools, and empty and nonempty strings in both methods; the repair changes only the two Raises sentences.