Repository navigation
Return PolarsCursor rows as tuples without a dictionary per row - #1177
Conversation
PolarsDataFrameIterator.itertuples() converts each column of a chunk in blocks of 10,000 rows with Series.to_list() and applies only the converters that are not the identity. fetchmany() and fetchall() take slices of these rows instead of calling fetchone() per row, and count the rows taken before an exception in rownumber. Closes #1170 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A close() from another thread between the check and the update of rownumber made the fetch raise TypeError; PandasCursor shares the same code since #1167. Rows of an iterator without column names stay empty tuples, as iter_rows() gave them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…h tests Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| columns: list[tuple[pl.Series | None, Callable[[Any], Any] | None]] = [] | ||
| for name in self._column_names: | ||
| converter = self._converters.get(name) | ||
| columns.append( |
There was a problem hiding this comment.
Self-review round one (implementation behavior) — base a618cd57cd84d3997e46b48e30ec3c3f9e1d60e5, initial head 043e8c9ae0473d01962d7b32c0ce12d9ac90b136.
Covered: PolarsDataFrameIterator.itertuples()/iterrows()/_iter_row_blocks(), AthenaPolarsResultSet fetch methods and close(), callers (sync PolarsCursor, AsyncPolarsCursor, aio AioPolarsCursor via asyncio.to_thread(result_set.fetchmany/fetchall), the awsathena+polars dialect), CSV / UNLOAD / chunked / non-chunked paths, values and types across dtypes on Polars 1.39.0, 1.44.2, 2.0.0, and the new tests.
Result: FINDINGS, repaired.
name in df.schemabuilt a Schema per column per chunk (O(columns²)); replaced withdf.get_column(name, default=None)(e80c100), available from the Polars floor 1.39.0 through 2.0.0._identityonly served as a sentinel after the change; a column without a converter is nowNonefromdict.get()(2519fb8)./code-review(high): the check-then-increment ofrownumbercould raiseTypeErrorwhen another thread'sclose()ran between them;_count_rows()now reads it once, here and inAthenaPandasResultSet, which has the same code since Return PandasCursor rows as tuples without a dictionary per row #1167 (98db6b3). Without column names,zip()gave no rows whereiter_rows()gave empty ones; kept withrepeat((), length)(98db6b3). The dtype test now covers all 24 dtypes the PR body names./simplify: test name/docstring cleanups (b418a88); reuse, efficiency (measured), and altitude passes clean.
Deferred: fetchone/fetchmany/fetchall/_take_rows/_count_rows duplicate AthenaPandasResultSet; sharing them is a class-structure change outside this PR (relates to #1171). Skipping identity converters is decided per chunk in the iterator, so the public converters property is unchanged. The close() race is not covered by a deterministic test.
| self._count_rows(len(rows)) | ||
| return rows | ||
|
|
||
| def _count_rows(self, count: int) -> None: |
There was a problem hiding this comment.
Round one repair: reading rownumber once removes the TypeError window; a close() that lands between the read and the write can still leave the count in place, as the docstring states. The same helper replaces the identical code in pyathena/pandas/result_set.py:989.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| """Iterate over the rows of a chunk in blocks of rows. | ||
|
|
||
| A block takes the values of a column at once with ``Series.to_list()`` and | ||
| applies the column's converter unless the column has none or has the default |
There was a problem hiding this comment.
Self-review round two (claims, callers, operations) — base a618cd57cd84d3997e46b48e30ec3c3f9e1d60e5, head 5ce55f6777c318981a4050ba38432d451f3316bd (full pass).
Claims checked:
- PR body "converter applied unless identity (
_to_defaultor the iterator's default)": stale after_identitywas removed; this docstring said "only when it is not the identity", but a custom identity converter is still called. Both reworded to "none or the default converter" (5ce55f6). - Speedup: re-measured at 5ce55f6, 4.75 s → 1.38 s with equal rows; the PR body now states that the before row excludes the base class's per-row
fetchone()loop. Local only; the end-to-end number belongs to the 4.2.0 combined benchmark. - 24 dtypes:
test_itertuples_valuescovers all of them and passes on Polars 1.39.0 (floor), 1.44.2, 2.0.0. - "rownumber stays None after a concurrent close()": narrowed in the body and
_count_rows()docstring — the TypeError is gone, but a close() between the read and the write can leave the count.
Callers: iterrows() output, fetch* return shapes and TypeError for a non-integer size, exceptions ending the iteration, and close() are unchanged; aio fetches reach the overrides through asyncio.to_thread. Removed names are private (_iterrows, _column_names_cache, module _identity). datetime.UTC in tests needs 3.11, the requires-python floor. docs/polars.md fetch/chunk statements remain true; docs/api/polars.rst documents itertuples() through :members:.
AWS operator: no change in API calls, retries, or S3 reads; the fetch methods read chunks lazily as before (test_fetch_across_chunks).
Evidence at 5ce55f6: lint; polars + aio polars + pandas result set 421 passed / 11 skipped (live AWS); pandas fetch integration 35 passed; SQLAlchemy polars/pandas 16 passed; Polars 1.39.0 and 2.0.0 unit 58 passed each. Not run: full just test pyathena; no deterministic test for the close() race.
Result: FINDINGS (wording/evidence), repaired; CLEAN otherwise.
A StopIteration from a converter ended the rows of its block through map(), dropping them without an error; a generator turns it into a RuntimeError, as the row iteration raised before. A value that Python does not hold, such as a date in year 0, failed to_list() for the whole block; the block's values are then taken one at a time, so the rows before it are returned before it raises. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| for i in range(block.len()): | ||
| yield block.slice(i, 1).to_list()[0] | ||
|
|
||
| def convert(converter: Callable[[Any], Any], values: Iterable[Any]) -> Iterator[Any]: |
There was a problem hiding this comment.
Independent review (relayed result) — reviewer: Codex CLI 0.162.1, model gpt-6-astra (OpenAI), session 01a13ae1-8a63-77f2-9340-34c31ac5b284; sandbox read-only, static review with read-only Polars 1.44.2 probes (no tests, builds, edits, or network). Base a618cd57cd84d3997e46b48e30ec3c3f9e1d60e5, head 5ce55f6777c318981a4050ba38432d451f3316bd, detached snapshot unchanged during the review. The prompt carried the diff and conventions only, without the PR number, description, commit messages, or prior findings.
Covered: PolarsCursor/AsyncPolarsCursor/AioPolarsCursor shared paths, CSV/UNLOAD, chunked/non-chunked, values/types, NULLs, nested values, renaming, missing/extra columns, empty results, fetch sizes, rownumber, conversion errors, close/concurrent close, the pandas change, test assertions. Result: FINDINGS (2). No further regression in the pandas change; its remaining concurrent-close window predates this diff. Polars 1.39 and 2.x were not executed by the reviewer.
Finding 1 (P2) — converter failures can silently drop rows (pyathena/polars/result_set.py:194 at 5ce55f6): a custom converter raising StopIteration inside map() ends the block as exhaustion. Single-column chunks ["good", "bad", "lost"] and ["later"] with a failure on "bad" return [("good",), ("later",)]; before, the row generator raised RuntimeError: generator raised StopIteration.
Verified by the author: reproduced at 5ce55f6 on Polars 1.44.2.
There was a problem hiding this comment.
Repaired in 737e1f7: the converter is applied by a generator (convert()), which turns its StopIteration into RuntimeError as the previous row iteration did; a first attempt with a generator expression late-bound converter to the last column and was caught by test_itertuples_converters, so the converter is passed as an argument. Regression test: test_itertuples_converter_stop_iteration (raises after ("good",), no rows of the later chunk). Both self-review perspectives applied to the repair: no hot-path change (benchmark 1.42 s vs 1.38 s, identity columns unaffected); the PR body states the behavior. Validation at 737e1f7: lint; Polars + aio Polars 197 passed, 1 skipped (live AWS); test_result_set.py 60 passed on Polars 1.39.0, 1.44.2, 2.0.0.
| (df.get_column(name, default=None), None if converter is _to_default else converter) | ||
| ) | ||
|
|
||
| def take_values(block: pl.Series) -> Iterator[Any]: |
There was a problem hiding this comment.
Finding 2 (P2) — eager conversion makes valid-prefix fetches fail (pyathena/polars/result_set.py:192 at 5ce55f6, relayed from the same review): a Date column with 512 rows of 1970-01-01 followed by 0000-01-01 (no chunking or chunksize >= 513). Before, fetchmany(512) returned the valid rows; now the first fetchone() converts the whole block with to_list(), raises ValueError: year 0 is out of range, and leaves rownumber at 0.
Verified by the author: reproduced at 5ce55f6. The old iter_rows() returned 512 rows, then raised ValueError on Polars 2.0.0 and a PanicException (BaseException) on Polars 1.44.2.
There was a problem hiding this comment.
Repaired in 737e1f7: when to_list() fails for a block, that column's values are taken one at a time with slice(i, 1).to_list()[0] (take_values(); Series.item() would return a Series for List columns), so the rows before the value are returned and the value raises ValueError — finer than the old 512-row buffers and without the PanicException of Polars 1.44.2. except Exception leaves BaseException such as a Rust panic to propagate. Regression test: test_itertuples_value_python_does_not_hold (two rows with a following column, then ValueError, then no rows), passing on Polars 1.39.0, 1.44.2, 2.0.0.
There was a problem hiding this comment.
Independent follow-up (Codex CLI, gpt-6-astra, session 01a13af2-c33d-7eb3-b8dd-7ffdc1a61a16, read-only, static with Polars 1.44.2 probes; range 5ce55f67..737e1f72): FINDINGS (1, P2) — except Exception at pyathena/polars/result_set.py:212 missed a year-0 Date inside a Struct, which Polars raises from to_list() as PanicException (a BaseException), so 600 valid struct rows before it were not returned. The reviewer found finding 1 fixed with its cause retained, the direct-Date case fixed, no late binding, misalignment, value/type change, or added per-row work for columns without converters, and both new tests able to fail without the repair.
Verified by the author on Polars 1.39.0, 1.44.2, 2.0.0 (whole-block and single-value to_list() both panic). Repaired in 0f2afc5: the fallback also catches pl.exceptions.PanicException; the value still raises it at its own row. Regression test: test_itertuples_nested_value_python_does_not_hold. Both self-review perspectives on the repair: only the failure path changes; a panic unrelated to a value would raise again on the first per-value read. Validation at 0f2afc5: lint; Polars + aio Polars 198 passed, 1 skipped (live AWS); test_result_set.py 61 passed on Polars 1.39.0, 1.44.2, 2.0.0.
There was a problem hiding this comment.
Independent follow-up 2 (Codex CLI, gpt-6-astra, session 01a13afb-ad35-7e30-b2bd-48940e4681e5, read-only, static with one Polars 1.44.2 probe; range 737e1f72..0f2afc51, snapshot unchanged): CLEAN. Covered _iter_row_blocks(), itertuples(), result-set wiring, all fetch methods, row counting, and the new test. The bulk PanicException is retried per value without replacing or coercing values; the invalid value still raises and ends the generator; fetch methods count the rows consumed before it; the new test fails without the repair. Other Polars versions were not executed by the reviewer (author ran the unit tests on 1.39.0, 1.44.2, 2.0.0).
There was a problem hiding this comment.
CI (run 38588175656, test / run (3.15), live AWS at 0f2afc5): 1 failed, 3206 passed, 20 skipped. The only failure was test_itertuples_value_python_does_not_hold: Python 3.15 reports a date in year 0 as year must be in 1..9999, not 0, so match="year 0" did not match; the behavior itself was as expected (rows before the value, then ValueError). The PR was returned to Draft. Repaired in 90c94d9: match="year", checked against the messages of Python 3.11 (year 0 is out of range) and 3.15. Independent follow-up (Codex CLI, gpt-6-astra, session 01a13b09-9333-7492-8ba1-a3ffdd997721, read-only, static; range 0f2afc51..90c94d96): CLEAN — the test still requires the two exact rows, the ValueError on the third, and exhaustion despite a later valid row.
Within a struct, Polars raises a value that Python does not hold, such as a date in year 0, as a PanicException, which is a BaseException, so the block's values were not taken one at a time and the rows before it were not returned. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Python 3.15 reports a date in year 0 as "year must be in 1..9999, not 0". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
WHAT
PolarsCursor returns rows without building a dictionary per row, as #1167 did for PandasCursor.
PolarsDataFrameIterator.itertuples()yields each row as a tuple across all chunks. Each chunk is converted in blocks of 10,000 rows withSeries.to_list()per column; a column's converter is applied value by value as the rows are taken, unless the column has none or has the default converter (_to_default), which returns the values as they are.iterrows()keeps its(row_index, dict)output and is built onitertuples().AthenaPolarsResultSet.fetchone()takes the next tuple;fetchmany()andfetchall()take slices withislice()instead of callingfetchone()per row. They keep the base class'ssizehandling (None or not positive →arraysize, non-integer →TypeError) and count the rows taken before an exception inrownumber.StopIterationraisesRuntimeError, as the previous row iteration did, instead of ending its block's rows. A value that Python does not hold, such as a date in year 0, failsSeries.to_list()for its whole block, so that block's values are then taken one at a time: the rows before the value are returned, and the value raisesValueError, or Polars'PanicExceptionfor a value within a struct. Before,iter_rows()returned the rows of the 512-row buffers before the value, then raisedValueErroron Polars 2.0.0 and aPanicExceptionon Polars 1.44.2.rownumberis read once when the fetched rows are counted, so aclose()from another thread between the check and the update no longer raisesTypeError.AthenaPandasResultSethad the same check-then-increment since Return PandasCursor rows as tuples without a dictionary per row #1167 and uses the same counting now.Unchanged: the tuples follow the iterator's column names, so a name that a chunk does not have still gives None to its converter, a frame column outside the names is left out, and an iterator without column names yields empty rows, as
iter_rows(named=True)did before; chunked reads (chunksize),iter_chunks(), and the converted values and their types.WHY
Closes #1170.
The per-row dictionary dominated the row output. Locally at 737e1f7, whose row output path is unchanged by the last commit except on a failed
to_list()(Apple M1 Pro, Python 3.13.1, Polars 1.44.2), on 1.5 million rows × 15 columns shaped like the benchmark's flat projection (12 strings, 8 with 5% NULL, 3 integers), median of 3 runs:iter_rows(named=True), dictionary, tuple (without the base class's per-rowfetchone()loop)PolarsDataFrameIterator.itertuples()Both returned equal rows.
Series.to_list()anditer_rows()give the same values with the same types for 24 dtypes, including Decimal, time-zone-aware and nanosecond Datetime, Time, Duration, Binary, List, Array, Struct, nested types, Null, Categorical, Enum, Object, UInt64, and Int128;test_itertuples_valueschecks all 24 and passes on Polars 1.39.0 (the floor), 1.44.2, and 2.0.0. The end-to-end re-measurement belongs to the combined 4.2.0 benchmark run.TEST
Tested at 0f2afc5 unless noted; 90c94d9 changes only the error-message match of one test.
CI at 90c94d9 (run 38590313837):
test / run (3.15)3215 passed, 20 skipped (live AWS, fulltests/pyathena/);test-sqlaandtest-sqla-asyncskipped by the workflow's path filter, which this PR's paths do not match.just lint: passed.uv run --env-file .env pytest -n 4 tests/pyathena/polars tests/pyathena/aio/polars: 198 passed, 1 skipped (live AWS).uv run --env-file .env --with 'polars==1.39.0'and--with 'polars==2.0.0'pytest -n 4 tests/pyathena/polars/test_result_set.py: 61 passed each; 61 passed on the locked Polars 1.44.2.At 5ce55f6, before the last repair, which changes only
PolarsDataFrameIterator:tests/pyathena/pandas/test_result_set.pywith the Polars suites, 421 passed, 11 skipped;tests/pyathena/pandas/test_cursor.py tests/pyathena/aio/pandas -k "fetch or iterator or rownumber", 35 passed;tests/pyathena/sqlalchemy/test_base.py tests/pyathena/aio/sqlalchemy/test_base.py -k "polars or pandas", 16 passed.New tests:
TestPolarsDataFrameIterator: values and types equal toiter_rows()across 24 dtypes, converters by column name (default converter skipped, missing name, extra column), a converter raisingStopIteration, a date in year 0 after valid rows, directly and within a struct, rows without column names, a chunk larger than two blocks, and continuousiterrows()row numbers.TestAthenaPolarsResultSet: fetches across chunks with lazy chunk reads andrownumber, empty chunks, a conversion error in each fetch method, rows of a started chunk afteriter_chunks()runs out,close()duringfetchall()in another thread,fetchmany()sizes abovesys.maxsizeand not integers, and fetches afterclose().TestPolarsCursor.test_fetch_across_chunks:chunksize=2with and without UNLOAD against Athena.Not run locally: the full
just test pyathenasuite (CI ran it). Theclose()race in therownumberupdate has no deterministic test.🤖 Generated with Claude Code