Skip to content

Require polars>=1.39.0 for chunked PolarsCursor reads - #828

Merged
laughingman7743 merged 3 commits into
masterfrom
fix/820-821-polars-floor
Sep 26, 2026
Merged

laughingman7743 merged 3 commits into
masterfrom
fix/820-821-polars-floor

Conversation

@laughingman7743

@laughingman7743 laughingman7743 commented Sep 26, 2026 •

Copy link
Copy Markdown
Member

WHAT

  • Raise the Polars requirement from polars>=1.0.0 to polars>=1.39.0 in the polars extra and the dev dependency group, and in the extra-packages tables of README.md and docs/introduction.md. uv.lock changes 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, which chunksize does not limit (bounded independently of the result size for CSV), and a failed chunk read raises OperationalError.
  • docs/polars.md (intro): only CSV results read without chunksize use 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.md and the start_with_temp_dir docstring in benchmarks/pyathena_bench/runner.py: the Polars file cache in POLARS_TEMP_DIR is described as behavior of Polars before 1.39.0.
  • New no-AWS unit tests in tests/pyathena/polars/test_result_set.py: a CSV and a Parquet read that fail after some batches must raise OperationalError from _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:

Polars Change Effect on PolarsCursor with chunksize
< 1.34.0 no LazyFrame.collect_batches chunked reads fail with OperationalError wrapping AttributeError (from the code, not run)
1.34.0–1.36.x Python BatchCollector a failure inside the streaming query ends the iterator normally (#820)
1.37.0 "Implement collect_batches properly in Rust" (pola-rs/polars#25918) the error reaches the caller
1.39.0 "Streaming cloud download for scan_csv" (pola-rs/polars#26637); file cache removed no full download to POLARS_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):

Polars CSV Parquet
1.33.1 no collect_batches no collect_batches
1.34.0 ends normally after 290,000 rows ends normally after 300,000 rows
1.36.1 ends normally after 290,000 rows ends normally after 300,000 rows
1.37.0, 1.37.1, 1.38.1, 1.39.0, 1.44.2 raises ComputeError raises ComputeError

So 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 fresh POLARS_TEMP_DIR per run:

Polars Result object Peak RSS Peak POLARS_TEMP_DIR Left after exit
1.36.1 CSV, 2.5M rows, 331 MB 547 MB 331 MB 331 MB
1.39.0 same 536 MB 0 0
1.44.2 same 304 MB 0 0
1.44.2 CSV, 10M rows, ~1.3 GB 844 MB 0 0
1.44.2 UNLOAD Parquet via PolarsCursor(unload=True, chunksize=100_000), 10M rows, one 676 MB file (6 row groups, up to 150 MB uncompressed) 715 MB 0 0
1.39.0 same Parquet object via scan_parquet(...).collect_batches(chunk_size=100_000) 983 MB 0 0

Memory is still not proportional to chunksize.
Polars 1.44.2 prefetches up to POLARS_CSV_CHUNK_PREFETCH_LIMIT chunks 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_batches API.
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, and just docs build.

  • just lint: passed (ruff, format, mypy 0 issues in 92 files, cfn-lint, license headers).
  • markdownlint-cli2 on the changed Markdown files (including benchmarks/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 with DID 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 compile of .[polars] with polars==1.38.1 is unsatisfiable (pyathena[polars] depend on polars>=1.39.0); with polars==1.39.0 it resolves.
  • Not run here: the 10M-row, 2.8 GB measurement on EC2 from the issues; it is covered by the Comprehensive cursor benchmark with memory and performance metrics #644 benchmark rerun, whose harness records POLARS_TEMP_DIR peak usage.

🤖 Generated with Claude Code

laughingman7743 and others added 2 commits September 26, 2026 10:37
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."""

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

Comment thread docs/polars.md
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.

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 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_DIR peak 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_LIMIT defaulting to num_pipelines * 2 in polars-stream/.../csv/builder.rs at 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/polars pass on 1.39.0 and 1.44.2 (104 each); polars==1.38.1 with .[polars] is unsatisfiable. Users pinned below 1.39.0 must upgrade even without chunksize (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:228 still 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())

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

Comment thread docs/polars.md
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.

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

@laughingman7743 laughingman7743 Sep 26, 2026 •

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 (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: with POLARS_FORCE_STREAMING=1 or POLARS_AUTO_STREAMING=1, Polars routes a non-chunked pl.read_csv through 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.

Comment thread benchmarks/README.md
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.

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

@laughingman7743
laughingman7743 marked this pull request as ready for review September 26, 2026 02:40
@laughingman7743
laughingman7743 merged commit ac6a3a1 into master Sep 26, 2026
20 checks passed
@laughingman7743
laughingman7743 deleted the fix/820-821-polars-floor branch September 26, 2026 02:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant