Repository navigation
Pass the open CSV result file to Polars instead of its bytes (GH-1307) - #1313
Conversation
Since GH-1169, an unchunked PolarsCursor CSV read passed the whole result file to pl.read_csv() as bytes. Polars 1.x wraps bytes in a BytesIO and calls getbuffer() on it in _check_empty(), which makes CPython copy the shared buffer, so the read held a second copy of the file. Polars reads an open file object with one read() call and wraps the returned bytes without copying, in 1.x and 2.0, which is also what Polars 1.x did with the path and fsspec storage options before GH-1169. The result set now passes the file opened with fsspec.open() and seeks back to its start before the second read that parses timestamps as text. Rows fetched with GetQueryResults are still passed as bytes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
fsspec.open() creates the filesystem when it is called, so an error in creating it, such as a profile in the storage options that does not exist, escaped as the botocore error instead of OperationalError. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With a cache type such as "all" or "readahead", the open file can keep the bytes it read until it is closed, so the file is now open only for the reads. The second read seeks to the start when a file was opened, and the filesystem test checks that the error comes from the profile. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
After seek(0), the second read of the same file missed fsspec's BytesCache, which then joined the fetched range onto its tail and held two copies of the file. Each read now opens the file, so the second one starts with an empty cache, and the file is closed before the timestamps are converted without seeking or a separate open block. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ff20ab1 to
0f76f07
Compare
| if strict or not timestamp_dtypes: | ||
| try: | ||
| df = pl.read_csv(source, **kwargs) | ||
| with open_source() as data: |
There was a problem hiding this comment.
Self-review round 1 — implementation behavior (perspective: behavior/failure paths, data/resource boundaries, simplicity, test quality)
Scope: git diff adf5aec568cd536e76e1568afed76f0a4caca4a1..0f76f074a4fd805552f7e8813f4f769704e8d1f7 (merge-base adf5aec568cd536e76e1568afed76f0a4caca4a1, head 0f76f074a4fd805552f7e8813f4f769704e8d1f7); inventory: pyathena/polars/result_set.py (_read_csv(), _csv_storage_options docstring) and tests/pyathena/polars/test_result_set.py. Callers traced: AthenaPolarsResultSet.__init__ (whole read via _as_polars(), and the GetQueryResults path without an output location); the sync, async, and aio Polars cursors share this result set. Chunked reads (scan_csv) are untouched.
Result: FINDINGS, all repaired (the round was run on the first commits and followed by /code-review and /simplify):
fsspec.open()creates the filesystem when it is called, and the first commit called it before thetry, so an error such as aprofilethat does not exist escaped asProfileNotFoundinstead ofOperationalError. Repaired: the call is inside thetry, throughopen_source(); regression testtest_csv_read_raises_when_file_system_fails(fails with the call outside thetry, passes on master), asserting theProfileNotFoundcause.- With cache types such as
"all"or"readahead", the open file can keep the bytes it read until it is closed, and the first commit kept it open through_to_datetimes(). Repaired: eachwith open_source()covers only onepl.read_csv(). - The second (timestamps-as-text) read used
seek(0)on the same file; fsspec'sBytesCachekeeps only the tail after the first full read, so the re-read joined the fetched range onto it (self.cache = new + self.cache, fsspeccaching.py), holding two copies. Repaired: each read opens the file again (verified in fsspec source). isinstance(data, bytes)decided the seek from the value's type; removed with item 3.
Deferred with reason: the fallback second read downloads the file again (4.1.0 did too; keeping the bytes would restore the copy this PR removes). BytesIO(rows) for the GetQueryResults path was declined: Polars 1.x copies a BytesIO built from bytes in getbuffer().
Checks: just lint; Polars tests on Polars 2.0.0 (196 passed, 7 skipped) and 1.44.2 (195 passed, 8 skipped); just benchmark test (256 passed, 1 skipped).
| # Polars 2.0 reads a path with its native object store, which does not take | ||
| # fsspec options, so the result file is read through fsspec here. | ||
| # fsspec options, so the result file is opened through fsspec here and given | ||
| # to Polars as a file, which Polars 1.x does not copy as it does bytes. Each |
There was a problem hiding this comment.
Self-review round 2 — claims, callers, operations (merge-base adf5aec568cd536e76e1568afed76f0a4caca4a1, head 0f76f074a4fd805552f7e8813f4f769704e8d1f7)
Claims checked against sources and measurements:
- Polars 1.44.2 copies bytes:
read_csv()→prepare_file_arg()→_check_empty(BytesIO(file))callsgetbuffer(); a tracemalloc check on CPython 3.13 showsBytesIO(b).getbuffer()allocates the full size. Confirmed. - File objects are read with one
read()and wrapped without copying in Polars 1.44.2 and 2.0.0:PyFileLikeObject.to_buffer()incrates/polars-python/src/file.rsatpy-1.44.2andpy-2.0.0. Confirmed; the measured peak with an S3 file equals the 4.1.0 path read (1,893–1,897 vs 1,897–1,904 MiB). - Polars 2.0 does not copy bytes:
read_csv()→scan_csv()→get_python_scan_source_input()wrapsPyBytes. Confirmed; bytes vs file peaks overlap. - The PR body's peak table comes from the Linux runs in the description; the head changed only after those runs by commits measured at ff20ab1 (1 trial) and a rebase that leaves the Polars files unchanged (
git diff ff20ab10 0f76f074a4fd805552f7e8813f4f769704e8d1f7 -- pyathena/polars tests/pyathena/polarsis empty).
Callers: read options use_pyarrow, columns, n_rows, encoding="utf8-lossy" read an fsspec file the same way on both Polars versions (local check). An encoding other than UTF-8 is not decoded for a file object by Polars (both versions), unlike bytes; it fails with OperationalError, as with 4.1.0. The maintainer decided this Polars behavior needs no handling; it is stated in the PR body.
AWS operations: the normal path makes the same requests as master (one open, ranged GETs). The fallback opens the file a second time (another HEAD and full GET), as the 4.1.0 path read did; stated in the PR body.
Result: CLEAN after correcting the PR body (tested commit, new test, behavior differences, and the unmeasured second-read peak).
…H-1307) Unpacking storage_options that are not a mapping raised TypeError outside the try, so it escaped without OperationalError. The opener is now created inside it, and the test covers both storage options that fail to create the filesystem and ones that are not a mapping. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| if self.output_location: | ||
| with fsspec.open(self.output_location, "rb", **(storage_options or {})) as file: | ||
| source = file.read() | ||
| open_source: Callable[[], AbstractContextManager[Any]] = ( |
There was a problem hiding this comment.
Independent review (relayed): Codex — codex exec -s read-only (default model and effort), session 01a159ec-818d-7be3-b255-be98a740ef0c. Static review only (no builds, tests, or network); detached snapshot of 0f76f074a4fd805552f7e8813f4f769704e8d1f7, merge-base adf5aec568cd536e76e1568afed76f0a4caca4a1; the prompt held the diff and repository conventions, without the PR number, description, commits, or prior findings. Snapshot and PR worktree unchanged afterwards.
Covered (per reviewer): exception wrapping, fsspec/S3 resource closing, timestamp retries, managed-result bytes; read-option forwarding across sync, thread-pool, and asyncio cursors; Polars 1.39.0, 1.44.2, and 2.0.0 sources (encoding and PyArrow paths); context-manager typing; memory/cache lifetime and repeated I/O; the added test against docs/testing.md (judged to protect a distinct caller-visible contract).
Result: FINDINGS (2)
- P2,
pyathena/polars/result_set.py:882: a file object bypasses Polars' decoding of a non-UTF-8encoding; e.g.encoding="ascii"on a result withénow succeeds instead of raisingOperationalErrorfromUnicodeDecodeError. Rejected: this is Polars' handling of file objects (prepare_file_arg()decodes only bytes, BytesIO, Path, and str), the same as with the 4.1.0 path read, and Athena writes UTF-8; the maintainer decided it needs no handling. Stated in the PR body. - P3,
pyathena/polars/result_set.py:866:**(storage_options or {})was unpacked outside thetry, sostorage_options="invalid"raised a rawTypeError. Accepted and repaired inba06c2eab2285fc865617087f49bda7f73248f69(see the reply).
Reviewer note (not a finding): the timestamp retry opens the file again and reads it from S3 a second time. Already stated in the PR body as a deliberate difference.
There was a problem hiding this comment.
Repair of finding 2 in ba06c2eab2285fc865617087f49bda7f73248f69: the opener (partial(fsspec.open, ...) / partial(nullcontext, source)) is created inside the try, so storage options that are not a mapping raise OperationalError caused by the TypeError. The test is now test_csv_read_raises_with_invalid_storage_options[file-system|not-a-mapping], asserting the cause (ProfileNotFound / TypeError); not-a-mapping fails at 0f76f074 and both pass on master and at ba06c2eab2285fc865617087f49bda7f73248f69.
Self-review of the repair: round 1 (behavior) — the GetQueryResults fetch and the empty return stay outside the try as on master; each read still opens the file in its own with. Round 2 (claims) — the PR body is updated with the tested commit, the renamed test, and the test counts. Both CLEAN.
Validation: just lint; Polars tests on 2.0.0 (197 passed, 7 skipped) and 1.44.2 (196 passed, 8 skipped).
Independent follow-up (relayed): Codex codex exec -s read-only, session 01a159f8-ed92-7f30-88dc-42560593ab6d, static review of 0f76f074..ba06c2eab2285fc865617087f49bda7f73248f69 on a detached snapshot of ba06c2eab2285fc865617087f49bda7f73248f69: CLEAN — the TypeError is wrapped with its cause; fetch/empty return/nullcontext/reopening/closing unchanged; the not-a-mapping case fails without the repair and asserts caller-visible behavior; no typing problem.
There was a problem hiding this comment.
Rename in 7bf315c1b7e811fa84459c08502ebe04dceffe85 (maintainer's request): the local factory open_source (which also reads as "open source") is now open_csv, and the sentence about reopening moved into a comment directly above it, which also says why it is created and called inside the try. No behavior change.
The name was chosen after consulting Codex (codex exec -s read-only) and Fable (claude-fable-5-1, max profile, effort high, plan mode, read-only tools) with the same neutral package. Both recommended open_csv (reads naturally as with open_csv() as data: and matches AthenaArrowResultSet._open_csv(), pyathena/arrow/result_set.py:346), keeping the local factory over a method or inlined branches, and placing the factory comment above its definition.
Self-review of the change: round 1 — pure rename, both uses updated, open_source no longer appears; round 2 — the comments match the code (each read opens and closes the file; fsspec's BytesCache reread joins the fetched range onto its tail; the factory is created and called inside the try). Both CLEAN. Validation: just lint; tests/pyathena/polars/test_result_set.py 72 passed on Polars 2.0.0 and on 1.44.2.
Independent follow-up (relayed): Codex codex exec -s read-only, session 01a15a18-7c30-7eb1-9b83-2fdb879c5421, static review of ba06c2ea..7bf315c1b7e811fa84459c08502ebe04dceffe85 on a detached snapshot of 7bf315c1b7e811fa84459c08502ebe04dceffe85: CLEAN — pure rename and comment change; both uses updated; the comments are accurate.
The local factory was named open_source, which also reads as "open source". It is now open_csv, the name the Arrow result set uses for the same job, and the comment above it says why each read opens the file again and why it is created inside the try. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
WHAT
An unchunked PolarsCursor read of a CSV result gives
pl.read_csv()the result file opened withfsspec.open()instead of the file's bytes. Each read opens the file again: the second read, which parses the timestamp columns as text after the first one fails, starts with a new file and its own cache, and the file is closed before the timestamps are converted. The file is opened inside the read's error handling, so storage options that cannot create the filesystem or are not a mapping still raiseOperationalError. Rows fetched with GetQueryResults are still passed as bytes.With Polars 1.x, the read no longer holds a second copy of the result file, so its peak memory returns to that of 4.1.0. With Polars 2.0, the peak does not change.
WHY
Closes #1307.
Since #1169,
_read_csv()read the whole file withfile.read()and passed the bytes to Polars.read_csv()passes bytes toprepare_file_arg(), which returnsBytesIO(file)after_check_empty()callsgetbuffer()on it.getbuffer()makes CPython copy the buffer that the BytesIO shares with the bytes (polars/io/_utils.py). PyAthena still referenced the bytes, so both copies were alive during the read.read()once and wrap the returned bytes without copying (PyFileLikeObject.to_buffer()incrates/polars-python/src/file.rs). Before Read CSV results through fsspec for Polars 2.0 #1169, Polars 1.x opened the path withfsspec.open(path, **storage_options)inprepare_file_arg()and read it the same way.read_csv()passes bytes toscan_csv(), which wraps them without copying, so the change does not reduce its peak.BytesCache, which keeps only the tail of the first read, and the cache would join the fetched range onto that tail, holding two copies of the file. A new file avoids this.Polars 2.0 still reads the result through fsspec, because it reads a path with its native object store, which rejects fsspec options (#1169). This PR does not change that.
Behavior differences, the same as with the 4.1.0 path read:
encodingother than UTF-8 for bytes, but not for a file object, so such anencodinggiven toexecute()makes the read fail withOperationalErrorinstead of decoding. Athena writes CSV results in UTF-8.TEST
Tested commit: ba06c2e (rebased on adf5aec); 7bf315c only renames the local opener to
open_csvand moves a comment, after whichtests/pyathena/polars/test_result_set.pypassed again on both Polars versions (72 passed). At the measured commits, the normal read is the same: one open of the result file and onepl.read_csv()of it.just lint: passed.uv run --env-file .env pytest -n 4 tests/pyathena/polars tests/pyathena/aio/polarswith Python 3.14 and Polars 2.0.0: 197 passed, 7 skipped.uv run --env-file .env --python 3.13 --isolated --with polars==1.44.2: 196 passed, 8 skipped.just benchmark test: 256 passed, 1 skipped.test_csv_read_raises_with_invalid_storage_options: a whole CSV read with aprofilethat does not exist instorage_optionsraisesOperationalErrorcaused byProfileNotFound, and withstorage_optionsthat are not a mapping, caused byTypeError.file-systemfails whenfsspec.open()is called before thetry(9c81f29),not-a-mappingfails when the opener is created before it (0f76f07), and both pass on master.test_read_csv_truncates_timestamps; reading the consumed file again fails[None-kwargs0-expected0]and[None-kwargs3-expected3]withNoDataError: empty CSV.use_pyarrow,columns,n_rows, andencoding="utf8-lossy"read options read an fsspec file the same way with Polars 1.44.2 and 2.0.0 (local check).Peak RSS (
ru_maxrss) on Linux (Docker on macOS, aarch64, Python 3.13, PyAthena's S3 filesystem), for a 4-million-row query result of 513 MiB of CSV in S3:PolarsCursor.execute()+as_polars()Each master/9c81f29a pair ran back to back, and the Polars 2.0.0 pairs agree within 4 MiB. The same S3 file read directly with
pl.read_csv()gives the same picture: Polars 1.44.2 2,403–2,407 MiB with bytes, 1,893–1,897 MiB with the open file, and 1,897–1,904 MiB with the path andstorage_options(the 4.1.0 read); Polars 2.0.0 2,234–2,267 MiB with bytes and 2,237–2,299 MiB with the open file.Not run: the benchmark harness on EC2 (the #1307 table), the peak of the second read, and AWS CI, which runs when the PR is marked Ready.
🤖 Generated with Claude Code