Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ Extra packages:
| AioSQLAlchemy | `pip install PyAthena[AioSQLAlchemy]` | >=2.0.0 |
| Pandas | `pip install PyAthena[Pandas]` | >=1.3.0 |
| Arrow | `pip install PyAthena[Arrow]` | >=10.0.0 |
| Polars | `pip install PyAthena[Polars]` | >=1.0.0 |
| Polars | `pip install PyAthena[Polars]` | >=1.39.0 |

## Usage

Expand Down
4 changes: 2 additions & 2 deletions benchmarks/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -225,8 +225,8 @@ API row conversion and native DataFrame/Table access appear as separate cases.
Sampled RSS can miss short peaks, and the constructor suite's RSS includes its subsequent validation read.
If the operating system denies access to per-thread CPU times, `thread_cpu_available` is false; thread counts and RSS are still recorded.
Each trial process receives its own `POLARS_TEMP_DIR`, which the parent samples with RSS and removes after the trial.
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.


| Capability | Treatment |
| --- | --- |
Expand Down
7 changes: 4 additions & 3 deletions benchmarks/pyathena_bench/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,9 +175,10 @@ def directory_bytes(path: Path) -> int:
def start_with_temp_dir(process: Any, temp_dir: Path) -> None:
"""Start a spawned trial process with its own Polars temporary directory.

Polars downloads cloud objects for lazy scans into a file cache under
``POLARS_TEMP_DIR`` and keeps the files after the process exits. A
per-trial directory lets the parent measure and remove that storage.
Polars before 1.39.0 downloads cloud objects for lazy CSV scans into a
file cache under ``POLARS_TEMP_DIR`` and keeps the files after the process
exits. A per-trial directory lets the parent measure and remove any
storage Polars writes there.

Args:
process: Unstarted multiprocessing process; it inherits the environment.
Expand Down
2 changes: 1 addition & 1 deletion docs/introduction.md
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ Extra packages:
| AioSQLAlchemy | `pip install PyAthena[AioSQLAlchemy]` | >=2.0.0 |
| Pandas | `pip install PyAthena[Pandas]` | >=1.3.0 |
| Arrow | `pip install PyAthena[Arrow]` | >=10.0.0 |
| Polars | `pip install PyAthena[Polars]` | >=1.0.0 |
| Polars | `pip install PyAthena[Polars]` | >=1.39.0 |

(features)=

Expand Down
12 changes: 8 additions & 4 deletions docs/polars.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@ This cursor downloads the CSV file after executing the query and loads it into a
Performance is better than fetching data with Cursor.

PolarsCursor uses [Polars](https://pola.rs/) native reading capabilities (`pl.read_csv`, `pl.read_parquet`) and
does not require PyArrow as a dependency. PyAthena's own S3FileSystem (fsspec compatible)
is used for S3 access, so s3fs is also not required.
does not require PyArrow as a dependency. CSV results read without the chunksize option use
PyAthena's own S3FileSystem (fsspec compatible), and other reads use Polars' native S3 access,
so s3fs is also not required.

You can use the PolarsCursor by specifying the `cursor_class`
with the connect method or connection object.
Expand Down Expand Up @@ -310,8 +311,11 @@ for chunk in cursor.iter_chunks():
print(f"Processed chunk with {chunk.height} rows")
```

This method uses Polars' `scan_csv()` and `scan_parquet()` with `collect_batches()`
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.

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.

If reading a chunk fails, iteration raises `OperationalError` instead of ending early.

The chunked iteration also works with the unload option:

Expand Down
4 changes: 2 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ arrow = [
"pyarrow>=22.0.0; python_version>='3.14'",
]
polars = [
"polars>=1.0.0",
"polars>=1.39.0",
]

[dependency-groups]
Expand All @@ -77,7 +77,7 @@ dev = [
"numpy>=2.3.0; python_version>='3.14'",
"pyarrow>=10.0.0; python_version<'3.14'",
"pyarrow>=22.0.0; python_version>='3.14'",
"polars>=1.0.0",
"polars>=1.39.0",
"Jinja2>=3.1.0",
"mypy>=0.900",
"pytest>=3.5",
Expand Down
83 changes: 83 additions & 0 deletions tests/pyathena/polars/test_result_set.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
# Copyright 2026 The PyAthena authors
#
# Licensed under the MIT License.
# See LICENSE or https://opensource.org/licenses/MIT.
#
# SPDX-License-Identifier: MIT

from unittest.mock import PropertyMock, patch

import polars as pl
import pytest

from pyathena.error import OperationalError
from pyathena.polars.result_set import AthenaPolarsResultSet

_ROWS_BEFORE_FAILURE = 300_000


def _chunked_result_set() -> AthenaPolarsResultSet:
"""Create a chunked result set without calling Athena.

Returns:
An AthenaPolarsResultSet with only the attributes used by the chunk readers.
"""
result_set = AthenaPolarsResultSet.__new__(AthenaPolarsResultSet) # bypass __init__
result_set._chunksize = 10_000
result_set._kwargs = {}
return result_set


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.

path = tmp_path / "result.csv"
path.write_text(
"a\n" + "".join(f"{i}\n" for i in range(_ROWS_BEFORE_FAILURE)) + "not-a-number\n"
)
result_set = _chunked_result_set()
with (
patch.object(
AthenaPolarsResultSet,
"output_location",
new_callable=PropertyMock,
return_value=str(path),
),
patch.object(
AthenaPolarsResultSet,
"dtypes",
new_callable=PropertyMock,
return_value={"a": pl.Int64},
),
patch.object(
AthenaPolarsResultSet,
"_parquet_storage_options",
new_callable=PropertyMock,
return_value={},
),
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.


def test_iter_parquet_chunks_raises_when_read_fails_partway(self, tmp_path):
"""A Parquet read that fails partway through the data raises instead of ending early."""
pl.DataFrame({"a": range(_ROWS_BEFORE_FAILURE)}).write_parquet(
tmp_path / "0.parquet", row_group_size=10_000
)
valid = (tmp_path / "0.parquet").read_bytes()
# Keep the footer magic but corrupt the metadata so the second file fails to read.
(tmp_path / "1.parquet").write_bytes(valid[: len(valid) // 2] + b"\x00" * 64 + valid[-8:])
result_set = _chunked_result_set()
result_set._unload_location = f"{tmp_path}/"
with (
patch.object(
AthenaPolarsResultSet,
"_parquet_storage_options",
new_callable=PropertyMock,
return_value={},
),
patch.object(AthenaPolarsResultSet, "_prepare_parquet_location", return_value=True),
pytest.raises(OperationalError),
):
list(result_set._iter_parquet_chunks())
4 changes: 2 additions & 2 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading