Skip to content

Return PolarsCursor rows as tuples without a dictionary per row - #1177

Merged
laughingman7743 merged 9 commits into
masterfrom
perf/1170-polars-row-tuples
Oct 14, 2026
Merged

laughingman7743 merged 9 commits into
masterfrom
perf/1170-polars-row-tuples

Conversation

@laughingman7743

@laughingman7743 laughingman7743 commented Oct 14, 2026 •

Copy link
Copy Markdown
Member

WHAT

PolarsCursor returns rows without building a dictionary per row, as #1167 did for PandasCursor.

  • New public PolarsDataFrameIterator.itertuples() yields each row as a tuple across all chunks. Each chunk is converted in blocks of 10,000 rows with Series.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 on itertuples().
  • AthenaPolarsResultSet.fetchone() takes the next tuple; fetchmany() and fetchall() take slices with islice() instead of calling fetchone() per row. They keep the base class's size handling (None or not positive → arraysize, non-integer → TypeError) and count the rows taken before an exception in rownumber.
  • A converter that raises StopIteration raises RuntimeError, 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, fails Series.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 raises ValueError, or Polars' PanicException for a value within a struct. Before, iter_rows() returned the rows of the 512-row buffers before the value, then raised ValueError on Polars 2.0.0 and a PanicException on Polars 1.44.2.
  • rownumber is read once when the fetched rows are counted, so a close() from another thread between the check and the update no longer raises TypeError. AthenaPandasResultSet had 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:

Row output Time Rows/s
Before: iter_rows(named=True), dictionary, tuple (without the base class's per-row fetchone() loop) 4.82 s 311,167
After: PolarsDataFrameIterator.itertuples() 1.42 s 1,056,368

Both returned equal rows. Series.to_list() and iter_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_values checks 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, full tests/pyathena/); test-sqla and test-sqla-async skipped 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.py with 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 to iter_rows() across 24 dtypes, converters by column name (default converter skipped, missing name, extra column), a converter raising StopIteration, a date in year 0 after valid rows, directly and within a struct, rows without column names, a chunk larger than two blocks, and continuous iterrows() row numbers.
  • TestAthenaPolarsResultSet: fetches across chunks with lazy chunk reads and rownumber, empty chunks, a conversion error in each fetch method, rows of a started chunk after iter_chunks() runs out, close() during fetchall() in another thread, fetchmany() sizes above sys.maxsize and not integers, and fetches after close().
  • TestPolarsCursor.test_fetch_across_chunks: chunksize=2 with and without UNLOAD against Athena.

Not run locally: the full just test pyathena suite (CI ran it). The close() race in the rownumber update has no deterministic test.

🤖 Generated with Claude Code

laughingman7743 and others added 5 commits October 14, 2026 23:27
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(

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) — 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.schema built a Schema per column per chunk (O(columns²)); replaced with df.get_column(name, default=None) (e80c100), available from the Polars floor 1.39.0 through 2.0.0.
  • _identity only served as a sentinel after the change; a column without a converter is now None from dict.get() (2519fb8).
  • /code-review (high): the check-then-increment of rownumber could raise TypeError when another thread's close() ran between them; _count_rows() now reads it once, here and in AthenaPandasResultSet, 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 where iter_rows() gave empty ones; kept with repeat((), 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:

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.

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

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, callers, operations) — base a618cd57cd84d3997e46b48e30ec3c3f9e1d60e5, head 5ce55f6777c318981a4050ba38432d451f3316bd (full pass).

Claims checked:

  • PR body "converter applied unless identity (_to_default or the iterator's default)": stale after _identity was 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_values covers 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]:

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.

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.

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.

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]:

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.

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.

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.

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.

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.

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.

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.

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

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.

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>
@laughingman7743
laughingman7743 marked this pull request as ready for review October 14, 2026 15:18
@laughingman7743
laughingman7743 marked this pull request as draft October 14, 2026 15:29
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>
@laughingman7743
laughingman7743 marked this pull request as ready for review October 14, 2026 15:33
@laughingman7743
laughingman7743 merged commit 38dc8f8 into master Oct 14, 2026
12 checks passed
@laughingman7743
laughingman7743 deleted the perf/1170-polars-row-tuples branch October 14, 2026 15:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Return PolarsCursor rows without building a dictionary per row

1 participant