Require polars>=1.39.0 for chunked PolarsCursor reads - #828
Conversation
Polars 1.34.0-1.36.x end LazyFrame.collect_batches() normally when the streaming query fails, so a chunked read could return zero or partial rows as success (#820). Polars 1.37.0 propagates the error to the caller. Polars before 1.39.0 downloads a cloud CSV into its local file cache before yielding batches and leaves the copy behind (#821). From 1.39.0, scan_csv streams from S3 with a bounded read-ahead. Raise the floor to 1.39.0, add no-AWS regression tests for partial CSV and Parquet read failures, and document the chunked read behavior. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With Polars 1.39.0, scan_parquet(...).collect_batches() read a whole 676 MB UNLOAD object before the first batch, so a result-size-independent memory bound holds only for CSV. The benchmark README note about the Polars file cache now applies only to Polars before 1.39.0. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
|
||
| class TestAthenaPolarsResultSet: | ||
| def test_iter_csv_chunks_raises_when_read_fails_partway(self, tmp_path): | ||
| """A CSV read that fails partway through the data raises instead of ending early.""" |
There was a problem hiding this comment.
Self-review round 1 (implementation behavior) — base 617dcb8e81d41491ae003ab21fca6bc071409cae, head 4f3f711191a1c93d620139a2c7f7638da37c13d4 (initial full pass), repair verified at 7e3acd345aba317337f6cef0dd875b2f23c1b623.
Covered: pyproject.toml/uv.lock floor, README/introduction tables, docs/polars.md chunksize text, new tests/pyathena/polars/test_result_set.py; traced the error path _iter_csv_chunks/_iter_parquet_chunks → PolarsDataFrameIterator.__next__ (catches only StopIteration) → iterrows → AthenaPolarsResultSet.fetchone (catches only StopIteration), as_polars()/iter_chunks(), and the thread/aio Polars cursors (no exception handling in between), so the OperationalError reaches every public caller.
Result: FINDINGS (1, fixed).
- The test docstrings said the read fails "after some batches", but with Polars 1.39.0 the CSV reproduction raises before any batch is yielded. Reworded to "fails partway through the data" in 7e3acd3; the assertions are unchanged.
Regression evidence: the two new tests pass on polars 1.44.2 and fail with DID NOT RAISE OperationalError on 1.36.1.
Pre-existing, out of scope: after the chunk generator raises, it is finished, so a further fetchone() returns None instead of raising again. This is unchanged by this PR.
| This method uses Polars' `scan_csv()` and `scan_parquet()` with `collect_batches()`, | ||
| which read the result from S3 without writing it to local storage. | ||
| Memory usage depends on how far Polars reads ahead, not on `chunksize`. | ||
| For CSV results, Polars limits the read-ahead by the number of CPU cores, so memory usage does not grow with the result size. |
There was a problem hiding this comment.
Self-review round 2 (claims, callers, operations) — base 617dcb8e81d41491ae003ab21fca6bc071409cae, head 7e3acd345aba317337f6cef0dd875b2f23c1b623, full pass.
Claims checked:
- 1.34.0–1.36.x swallow errors, 1.37.0+ raise: measured with the issue reproduction on 1.33.1/1.34.0/1.36.1/1.37.0/1.37.1/1.38.1/1.39.0/1.44.2, CSV and Parquet.
- No local file cache from 1.39.0:
POLARS_TEMP_DIRpeak 0 for S3 CSV on 1.39.0/1.44.2 and Parquet on 1.39.0/1.44.2; 331 MB left behind on 1.36.1. - CSV read-ahead bound:
POLARS_CSV_CHUNK_PREFETCH_LIMITdefaulting tonum_pipelines * 2inpolars-stream/.../csv/builder.rsat both py-1.39.0 and py-1.44.2; a 1.3 GB CSV peaked at 844 MB RSS on 1.44.2; limit 2 lowered the 331 MB read to about 250–280 MB. - Python floor: Polars 1.39.0
requires_python >=3.10, same as PyAthena. - Existing callers: full
tests/pyathena/polars+tests/pyathena/aio/polarspass on 1.39.0 and 1.44.2 (104 each);polars==1.38.1with.[polars]is unsatisfiable. Users pinned below 1.39.0 must upgrade even withoutchunksize(release note).
Result: FINDINGS (2, fixed in 7e3acd3).
- The first docs text claimed memory does not grow with the result size for both paths. On 1.39.0
scan_parquet(...).collect_batches()read the whole 676 MB UNLOAD object (6 row groups) before the first batch (peak RSS 983 MB). The claim is now limited to CSV; UNLOAD only states memory depends on Polars' read-ahead. benchmarks/README.md:228still said lazy CSV scans download the whole object into the file cache; it now applies to Polars before 1.39.0.
Not measured here: the 10M-row / 2.8 GB EC2 case; left to the #644 fleet rerun.
chunksize still sets how many rows each batch buffers, so it affects memory alongside Polars' read-ahead. Only CSV results read without chunksize go through PyAthena's fsspec S3FileSystem; the other reads use Polars' native S3 access. The benchmark notes about the Polars file cache now describe it as a pre-1.39.0 behavior. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), | ||
| pytest.raises(OperationalError, match="not-a-number"), | ||
| ): | ||
| list(result_set._iter_csv_chunks()) |
There was a problem hiding this comment.
Independent review (relayed Codex result) — reviewer: Codex CLI 0.157.0, model gpt-6-sol, reasoning effort medium, session 01a0db72-bf5b-7241-a070-1869b5c071c1; static review only (read-only sandbox, no tests/builds/network) of a clean detached snapshot, base 617dcb8e81d41491ae003ab21fca6bc071409cae, head 7e3acd345aba317337f6cef0dd875b2f23c1b623. Prompt had the diff, relevant sources, and repo conventions, without PR number, description, commit messages, or prior findings. Snapshot and PR worktree were unchanged afterwards. Verdict: FINDINGS (4).
Covered (reviewer): the diff; Polars dependency metadata and supported Python range; sync, thread-based async, and asyncio cursor paths; chunked and non-chunked CSV/Parquet readers; Polars 1.44.2 Python sources; affected tests; docs and benchmark text.
Finding 1 — rejected. Reviewer: test_result_set.py:61 and :83 do not assert that a valid chunk was yielded before the error, so the failure could happen at scan setup and not exercise a mid-iteration failure.
Author verification: both tests fail with DID NOT RAISE OperationalError on polars 1.36.1 (the #820 defect), so the error does occur inside collect_batches() iteration where the old iterator swallowed it. Asserting at least one batch would break the tests on the supported floor: with 1.39.0 the same CSV and Parquet inputs raise before any batch is yielded (0 rows), while 1.44.2 raises after 260,000 / 300,000 rows. No change.
Reviewer's unverified items: runtime on Polars 1.39.0, all Python versions, live S3 for the three cursor types, the CPU-core read-ahead limit, and absence of local writes. These were measured by the author outside the review (see self-review round 2 and the PR's TEST section), not by the reviewer.
| for efficient lazy evaluation, minimizing memory usage when processing large datasets. | ||
| This method uses Polars' `scan_csv()` and `scan_parquet()` with `collect_batches()`, | ||
| which read the result from S3 without writing it to local storage. | ||
| Memory usage depends on the size of each chunk and on how far Polars reads ahead, which `chunksize` does not limit. |
There was a problem hiding this comment.
Finding 2 — accepted, fixed in cc7135d. Reviewer: "Memory usage depends … not on chunksize" is wrong because Polars' chunk_size is the number of rows buffered before a batch is yielded, so a very large chunksize increases batch memory.
Verified against LazyFrame.collect_batches docs in Polars 1.44.2 ("The number of rows that are buffered before a chunk is given"). The sentence now names both the chunk size and the read-ahead, and says chunksize does not limit the read-ahead.
Finding 3 — accepted, fixed in cc7135d (pre-existing text, folded as a contained fix). Reviewer: docs/polars.md:13 said PyAthena's fsspec S3FileSystem is used for S3 access, but chunked CSV/Parquet scans use Polars' native cloud access.
Verified: _read_csv passes fsspec storage_options to pl.read_csv, which only dispatches to the native scan for hf:// paths or forced streaming in Polars 1.44.2; _read_parquet, _read_parquet_schema, _iter_csv_chunks, and _iter_parquet_chunks pass object_store credentials. The paragraph now says only CSV results read without chunksize use the fsspec filesystem.
There was a problem hiding this comment.
Independent follow-up (relayed Codex result) — Codex CLI 0.157.0, model gpt-6-sol, reasoning effort medium; static, read-only review of the repair diff 7e3acd345aba317337f6cef0dd875b2f23c1b623..cc7135db513ca5c7a8c8bfac39308c9d113ad68b on a clean detached snapshot of cc7135db513ca5c7a8c8bfac39308c9d113ad68b, checked against pyathena/polars/result_set.py and Polars 1.44.2 sources. Snapshot and PR worktree unchanged afterwards.
Reviewer: the repair resolves findings 2–4 for the default read paths. Verdict FINDINGS (1):
docs/polars.md:14: withPOLARS_FORCE_STREAMING=1orPOLARS_AUTO_STREAMING=1, Polars routes a non-chunkedpl.read_csvthrough its native scan, so that read would not use PyAthena's fsspec S3FileSystem; suggested adding "by default".
Author decision — rejected, no change. Both variables are read only inside polars/io/csv/functions.py (lines 513–514 in 1.44.2); they are not exposed through pl.Config or documented in the read_csv docstring. The PyAthena docs describe behavior under Polars' documented configuration. With those variables set, PyAthena's fsspec-style storage_options (a connection object) would reach the native reader, which is a separate, pre-existing incompatibility, not a documentation inaccuracy.
Reviewer's statically unverified items (Rust read-ahead, local writes, pre-1.39.0 cache behavior, /tmp mount type) were measured by the author outside the review, except the Amazon Linux 2023 /tmp tmpfs fact, which comes from the #644 EC2 runs recorded in #821.
Independent review status: initial review FINDINGS (1 rejected, 3 fixed in cc7135d); follow-up on the repair FINDINGS (1 rejected with the reason above); no open actionable findings.
| Polars lazy CSV scans of S3 objects download the whole object into a file cache in this directory before yielding batches; in the recorded runs, those files remained after the process exited. | ||
| When the temporary directory is a tmpfs, as `/tmp` is on Amazon Linux 2023, this storage uses memory that RSS does not include. | ||
| Polars before 1.39.0 downloaded the whole S3 object of a lazy CSV scan into a file cache in this directory before yielding batches and left the file after the process exited; PyAthena requires Polars 1.39.0 or later, whose scans do not write that cache. | ||
| When the temporary directory is a tmpfs, as `/tmp` is on Amazon Linux 2023, files written there use memory that RSS does not include. |
There was a problem hiding this comment.
Finding 4 — accepted, fixed in cc7135d. Reviewer: the tmpfs sentence here and the start_with_temp_dir docstring in benchmarks/pyathena_bench/runner.py:178 still described the Polars file cache as current behavior.
The tmpfs sentence now refers to any files written to the directory, and the docstring describes the file cache as Polars-before-1.39.0 behavior. just benchmark lint passed.
WHAT
polars>=1.0.0topolars>=1.39.0in thepolarsextra and thedevdependency group, and in the extra-packages tables ofREADME.mdanddocs/introduction.md.uv.lockchanges only the specifiers; the resolved version stays 1.44.2.docs/polars.md(Chunksize Options): chunked reads do not write the result to local storage, memory depends on the chunk size and on Polars' read-ahead, whichchunksizedoes not limit (bounded independently of the result size for CSV), and a failed chunk read raisesOperationalError.docs/polars.md(intro): only CSV results read withoutchunksizeuse PyAthena's fsspec S3FileSystem; the other reads use Polars' native S3 access (the old text said fsspec was used for all S3 access).benchmarks/README.mdand thestart_with_temp_dirdocstring inbenchmarks/pyathena_bench/runner.py: the Polars file cache inPOLARS_TEMP_DIRis described as behavior of Polars before 1.39.0.tests/pyathena/polars/test_result_set.py: a CSV and a Parquet read that fail after some batches must raiseOperationalErrorfrom_iter_csv_chunks()/_iter_parquet_chunks().No code change in
pyathena/; both defects are in the Polars versions the old floor allowed.WHY
Closes #820.
Closes #821.
Both were found in the #644 benchmark with polars 1.36.1. The Polars changes that fix them:
chunksizeLazyFrame.collect_batchesOperationalErrorwrappingAttributeError(from the code, not run)BatchCollectorscan_csv" (pola-rs/polars#26637); file cache removedPOLARS_TEMP_DIR(#821)Measured locally (no AWS for the first table): the #820 reproduction from the issue (300,000 integers + one unparsable row,
chunk_size=10_000) and a Parquet variant (valid file + file with corrupted metadata in one directory):collect_batchescollect_batchesComputeErrorComputeErrorSo the UNLOAD path (
_iter_parquet_chunks) had the same #820 exposure.#821,
scan_csv(...).collect_batches(chunk_size=100_000)on Athena CSV results in S3 read from a laptop (Apple Silicon, 10 cores), with a freshPOLARS_TEMP_DIRper run:POLARS_TEMP_DIRPolarsCursor(unload=True, chunksize=100_000), 10M rows, one 676 MB file (6 row groups, up to 150 MB uncompressed)scan_parquet(...).collect_batches(chunk_size=100_000)Memory is still not proportional to
chunksize.Polars 1.44.2 prefetches up to
POLARS_CSV_CHUNK_PREFETCH_LIMITchunks of 32 MiB, defaulting to twice the number of pipelines (20 chunks, 640 MiB here); with the limit set to 2 the 331 MB read peaked at about 250–280 MB.The 1.3 GB read peaking below the object size is consistent with that bound.
The docs therefore describe the bound as Polars' read-ahead, not
chunksize, and do not name the unstable environment variable.The Parquet (UNLOAD) path is not bounded the same way on every allowed version: with 1.39.0 the first batch arrived only after the whole 676 MB object had been read (peak RSS 983 MB, including the pyarrow import used to read the footer), while 1.44.2 yielded batches as row groups arrived.
The docs claim a result-size-independent bound only for CSV; for UNLOAD they state only that memory depends on Polars' read-ahead.
Raising the floor removes the silent truncation (#820) and the local file cache (#821) for every supported installation instead of adding version checks or a PyAthena-side producer thread over the unstable
collect_batchesAPI.Polars 1.39.0 requires Python >= 3.10, the same as PyAthena.
This is a dependency change for users who pin Polars below 1.39.0, including users who never set
chunksize, and needs a release note.TEST
Tested commits: 4f3f711 (code and tests); 7e3acd3 and cc7135d change only docs, docstrings, and test docstrings, rechecked with
just lint,just benchmark lint, markdownlint, andjust docs build.just lint: passed (ruff, format, mypy 0 issues in 92 files, cfn-lint, license headers).markdownlint-cli2on the changed Markdown files (includingbenchmarks/README.md): 0 errors;just docs build: passed.uv run --env-file .env pytest -n 1 tests/pyathena/polars/test_result_set.py: 2 passed on polars 1.44.2; the same tests with--with polars==1.36.1: 2 failed withDID NOT RAISE OperationalError, so they detect PolarsCursor with chunksize returns 0 rows instead of raising when the CSV read fails #820.uv run --env-file .env pytest -n 2 tests/pyathena/polars tests/pyathena/aio/polars: 104 passed (polars 1.44.2, against Athena).--with polars==1.39.0 pytest -n 2 tests/pyathena/polars tests/pyathena/aio/polars: 104 passed, so the new floor works with the sync, thread, and aio Polars cursors.uv pip compileof.[polars]withpolars==1.38.1is unsatisfiable (pyathena[polars] depend on polars>=1.39.0); withpolars==1.39.0it resolves.POLARS_TEMP_DIRpeak usage.🤖 Generated with Claude Code