Repository navigation
Skip empty Arrow record batches and move the fetchmany() clamp - #1222
Conversation
- fetchmany() and fetchall() of the result set and its aio variant return the rows that they took before a failed page request to the front of the buffer, instead of dropping them while counting them in rownumber, so that a later fetch returns them. - The Arrow result set skips record batches without rows, which ended the fetches early, such as the batch of an empty result file. - The size clamp to sys.maxsize moves from the base fetchmany() into the islice() of the pandas and Polars result sets, which need it. - The test helper _page is replaced with get_query_results_responses(). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The result sets advanced the token of the next page before converting the rows of the fetched one, so a fetch after a failed conversion skipped that page; now that fetchmany() and fetchall() return the rows taken before a failure to the buffer, a later fetch would continue past it. The token is advanced once the rows are converted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This reverts commit 6b8612c.
The behavior of a fetch after an exception is left unspecified, as PEP 249 leaves it, and the other result sets could not follow a recovery contract of the GetQueryResults ones, so returning the rows of a failed fetch to the buffer is removed with the token reordering. The empty-batch fix of the Arrow result set is tested with empty batches before and between others, and the clamp and test helper changes stay. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| A table can hold empty record batches between others, such as one of | ||
| an empty result file, which would otherwise end the fetches early. | ||
| """ | ||
| rows = next((batch for batch in self._batches if batch.num_rows), None) |
There was a problem hiding this comment.
Self-review round 1 (behavior, failure paths, simplicity, tests). Base b2d7aa24b00a204e9e150952c708e3bb1126f9e2, head 9809e3312b116d50e0cb98d177b848e32de99c8f. Inputs: my pass and /code-review high on the earlier head 6b8612c3 (7 findings). Result: FINDINGS, resolved by narrowing.
- The
/code-reviewfindings on the rows returned to the buffer (an interrupt between taking and counting rows, a concurrentclose()) and on the token reordering (Arrow/S3FS not covered, a deterministic conversion error repeating its request, a test checking only the private token) were verified from the source. The maintainer chose, after a consultation with Codex andclaude-fable-5-1that disagreed (unspecified vs. retryable), to leave the behavior after a fetch error unspecified here and to decide one contract for all result sets in Define the behavior of result set fetches after an error #1225; both changes are removed (58ced3e, 9809e33). - The finding that test comments cited an unrelated issue (Free failed part bodies without the GC and abort an async copy's upload created during loop shutdown #1186) is resolved: the remaining test cites none.
- The empty-batch fix:
_batchesis an iterator, so the generator expression resumes after the batches it consumed;close()replaces it with an empty iterator. The test asserts that the table has an empty batch before each case (pyarrow'sto_batches()yields none after the last rows, so leading, middle, and consecutive cases are covered) and fails on master in all 12 cases.
| # extend() keeps the rows that it appended before an exception. | ||
| rows.extend(islice(self._row_iter, size)) | ||
| # extend() keeps the rows that it appended before an exception, and | ||
| # islice() takes no size above sys.maxsize. |
There was a problem hiding this comment.
Self-review round 2 (claims, callers, docs, evidence). Base b2d7aa24b00a204e9e150952c708e3bb1126f9e2, head 9809e3312b116d50e0cb98d177b848e32de99c8f. Result: CLEAN.
- Claim that only the islice-based result sets need the clamp: the base
_take_buffered_rows()callsislice()only when the limit is below the buffer length, and the aiofetchmany()never clamped; the existingtest_fetchmany_size_above_maxsizeof the base, pandas, and Polars result sets pass. - Claim that the Arrow defect is older than Take buffered rows in bulk in fetchmany() and fetchall() #1187: the parent of Take buffered rows in bulk in fetchmany() and fetchall() #1187 also returned
Nonefromfetchone()at an empty batch and built the bulk fetches on it. - Callers:
ArrowCursor,AsyncArrowCursor, andAioArrowCursoruseAthenaArrowResultSet; no API changes. No documentation describes the fetch of empty batches. Not run: the AWS cursor suites, which run when the PR is Ready.
| A table can hold empty record batches between others, such as one of | ||
| an empty result file, which would otherwise end the fetches early. | ||
| """ | ||
| rows = next((batch for batch in self._batches if batch.num_rows), None) |
There was a problem hiding this comment.
Independent review (relayed Codex result). Codex CLI 0.162.1, model gpt-6-astra, codex exec -s read-only, session 01a14709-346e-7690-9844-1f0d8d2bbe66. Static review of a clean detached snapshot of 9809e3312b116d50e0cb98d177b848e32de99c8f against base b2d7aa24b00a204e9e150952c708e3bb1126f9e2, prompt without PR text or prior findings; no edits, builds, tests, or network. Covered: "Arrow fetch methods and iteration, including empty batches, exhaustion, conversion, close(), rownumber, and arraysize. Base, pandas, Polars, aio, and both async Arrow cursor paths; very large and invalid fetch sizes. Regression tests, helper replacement, and changed comments/docstrings." Result: CLEAN — "The 12 new cases explicitly check that empty batches exist and would expose the original early stop. Their setup has no timing or randomness dependency. The clamp relocation preserves size validation and keeps islice() arguments within bounds." Snapshot unchanged afterwards.
WHAT
AthenaArrowResultSetskips record batches without rows, such as the batch of an empty result file before or between others (for exampleParquetDataset.read()of an UNLOAD with an empty file).fetchone()returnedNoneat such a batch, so the fetches and iteration stopped there and silently dropped the rest of the result.sys.maxsizeclamp offetchmany()moves from the base class into theislice()of the pandas and Polars result sets, which need it; the base class takes buffered rows withpopleft()up to the size, bounded by the buffer._pageis replaced with the sharedget_query_results_responses().WHY
The Arrow defect is older than #1187 (checked against its parent) and loses rows without an error, whatever the cursor (
ArrowCursor,AsyncArrowCursor,AioArrowCursor). The other two are cleanups of #1187 found by the same review.An earlier revision also returned the rows of a failed fetch to the buffer and reordered the page token after a failed conversion. Review showed that this covered only the GetQueryResults result sets, miscounted
rownumberwhen an interrupt arrived mid-take, could refill a closed result set, and made a deterministic conversion error repeat its request; the behavior of a fetch after an error is now proposed as one contract for all result sets in #1225.Release note:
ArrowCursorfetches no longer stop at an empty record batch.TEST
Tested at the head of this PR on Python 3.15 (macOS).
just lint: passed.uv run --env-file .env pytest -n 4 tests/pyathena/test_result_set.py tests/pyathena/aio/test_result_set.py tests/pyathena/arrow/test_result_set.py tests/pyathena/pandas/test_result_set.py tests/pyathena/polars/test_result_set.py: 354 passed, 10 skipped.masterin all 12 cases:TestAthenaArrowResultSet.test_fetch_across_empty_batches[middle|leading|consecutive]×fetchone/fetchmany/fetchall/iteration, which checks that the table really has an empty batch.to_batches()yields no empty batch after the last rows.test_fetchmany_size_above_maxsizetests of the base, pandas, and Polars result sets cover the clamp.The AWS jobs run when the PR is marked Ready.
🤖 Generated with Claude Code