-
Notifications
You must be signed in to change notification settings - Fork 113
Require polars>=1.39.0 for chunked PolarsCursor reads #828
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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. | ||
|
|
@@ -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. | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Finding 2 — accepted, fixed in cc7135d. Reviewer: "Memory usage depends … not on Finding 3 — accepted, fixed in cc7135d (pre-existing text, folded as a contained fix). Reviewer:
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Independent follow-up (relayed Codex result) — Codex CLI 0.157.0, model Reviewer: the repair resolves findings 2–4 for the default read paths. Verdict FINDINGS (1):
Author decision — rejected, no change. Both variables are read only inside Reviewer's statically unverified items (Rust read-ahead, local writes, pre-1.39.0 cache behavior, 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. | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Self-review round 2 (claims, callers, operations) — base Claims checked:
Result: FINDINGS (2, fixed in 7e3acd3).
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: | ||
|
|
||
|
|
||
| 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.""" | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Self-review round 1 (implementation behavior) — base Covered: Result: FINDINGS (1, fixed).
Regression evidence: the two new tests pass on polars 1.44.2 and fail with Pre-existing, out of scope: after the chunk generator raises, it is finished, so a further |
||
| 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()) | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Independent review (relayed Codex result) — reviewer: Codex CLI 0.157.0, model 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: 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()) | ||
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
There was a problem hiding this comment.
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_dirdocstring inbenchmarks/pyathena_bench/runner.py:178still 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 lintpassed.