Skip to content

Pass the open CSV result file to Polars instead of its bytes (GH-1307) - #1313

Merged
laughingman7743 merged 6 commits into
masterfrom
fix/1307-polars-csv-file-object
Oct 20, 2026
Merged

laughingman7743 merged 6 commits into
masterfrom
fix/1307-polars-csv-file-object

Conversation

@laughingman7743

@laughingman7743 laughingman7743 commented Oct 19, 2026 •

Copy link
Copy Markdown
Member

WHAT

An unchunked PolarsCursor read of a CSV result gives pl.read_csv() the result file opened with fsspec.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 raise OperationalError. 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 with file.read() and passed the bytes to Polars.

  • Polars 1.44.2 read_csv() passes bytes to prepare_file_arg(), which returns BytesIO(file) after _check_empty() calls getbuffer() 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.
  • For a Python file object that is not a local file, Polars 1.44.2 and 2.0.0 call read() once and wrap the returned bytes without copying (PyFileLikeObject.to_buffer() in crates/polars-python/src/file.rs). Before Read CSV results through fsspec for Polars 2.0 #1169, Polars 1.x opened the path with fsspec.open(path, **storage_options) in prepare_file_arg() and read it the same way.
  • Polars 2.0 read_csv() passes bytes to scan_csv(), which wraps them without copying, so the change does not reduce its peak.
  • Seeking the same file back to its start for the second read would miss fsspec's 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:

  • When the first read fails to parse the timestamps, the second read downloads the result file again; master parsed the bytes it held twice.
  • Polars decodes an encoding other than UTF-8 for bytes, but not for a file object, so such an encoding given to execute() makes the read fail with OperationalError instead of decoding. Athena writes CSV results in UTF-8.

TEST

Tested commit: ba06c2e (rebased on adf5aec); 7bf315c only renames the local opener to open_csv and moves a comment, after which tests/pyathena/polars/test_result_set.py passed again on both Polars versions (72 passed). At the measured commits, the normal read is the same: one open of the result file and one pl.read_csv() of it.

  • just lint: passed.
  • uv run --env-file .env pytest -n 4 tests/pyathena/polars tests/pyathena/aio/polars with Python 3.14 and Polars 2.0.0: 197 passed, 7 skipped.
  • The same tests with 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.
  • New test test_csv_read_raises_with_invalid_storage_options: a whole CSV read with a profile that does not exist in storage_options raises OperationalError caused by ProfileNotFound, and with storage_options that are not a mapping, caused by TypeError. file-system fails when fsspec.open() is called before the try (9c81f29), not-a-mapping fails when the opener is created before it (0f76f07), and both pass on master.
  • The second read is covered by the existing test_read_csv_truncates_timestamps; reading the consumed file again fails [None-kwargs0-expected0] and [None-kwargs3-expected3] with NoDataError: empty CSV.
  • The use_pyarrow, columns, n_rows, and encoding="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() Polars 1.44.2 Polars 2.0.0
master (0a58be3), 3 trials 2,055–2,058 MiB 1,918–1,999 MiB
this PR (9c81f29), 3 trials 1,535–1,541 MiB 1,922–1,999 MiB
this PR (ff20ab1), 1 trial 1,540 MiB 1,942 MiB

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 and storage_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

laughingman7743 and others added 4 commits October 21, 2026 00:20
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>
@laughingman7743
laughingman7743 force-pushed the fix/1307-polars-csv-file-object branch from ff20ab1 to 0f76f07 Compare October 20, 2026 15:20
Comment thread pyathena/polars/result_set.py Outdated
if strict or not timestamp_dtypes:
try:
df = pl.read_csv(source, **kwargs)
with open_source() as data:

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 (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):

  1. fsspec.open() creates the filesystem when it is called, and the first commit called it before the try, so an error such as a profile that does not exist escaped as ProfileNotFound instead of OperationalError. Repaired: the call is inside the try, through open_source(); regression test test_csv_read_raises_when_file_system_fails (fails with the call outside the try, passes on master), asserting the ProfileNotFound cause.
  2. 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: each with open_source() covers only one pl.read_csv().
  3. The second (timestamps-as-text) read used seek(0) on the same file; fsspec's BytesCache keeps only the tail after the first full read, so the re-read joined the fetched range onto it (self.cache = new + self.cache, fsspec caching.py), holding two copies. Repaired: each read opens the file again (verified in fsspec source).
  4. 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).

Comment thread pyathena/polars/result_set.py Outdated
# 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

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 (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)) calls getbuffer(); a tracemalloc check on CPython 3.13 shows BytesIO(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() in crates/polars-python/src/file.rs at py-1.44.2 and py-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() wraps PyBytes. 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/polars is 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>
Comment thread pyathena/polars/result_set.py Outdated
if self.output_location:
with fsspec.open(self.output_location, "rb", **(storage_options or {})) as file:
source = file.read()
open_source: Callable[[], AbstractContextManager[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): 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)

  1. P2, pyathena/polars/result_set.py:882: a file object bypasses Polars' decoding of a non-UTF-8 encoding; e.g. encoding="ascii" on a result with é now succeeds instead of raising OperationalError from UnicodeDecodeError. 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.
  2. P3, pyathena/polars/result_set.py:866: **(storage_options or {}) was unpacked outside the try, so storage_options="invalid" raised a raw TypeError. Accepted and repaired in ba06c2eab2285fc865617087f49bda7f73248f69 (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.

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.

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.

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.

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.

@laughingman7743
laughingman7743 marked this pull request as ready for review October 20, 2026 15:43
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>
@laughingman7743
laughingman7743 merged commit c0a938e into master Oct 20, 2026
7 checks passed
@laughingman7743
laughingman7743 deleted the fix/1307-polars-csv-file-object branch October 20, 2026 23:40
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.

PolarsCursor holds a copy of the CSV result file while reading it whole with Polars 1.x

1 participant