Repository navigation
Read CSV results through fsspec for Polars 2.0 - #1169
Conversation
Polars 2.0 reads a path in read_csv() with its native object store and rejects the fsspec storage options that PyAthena passed, so every CSV result read without chunksize failed. The result set now opens the file with fsspec.open() and the same options, and passes its contents to Polars, which keeps PyAthena's S3 filesystem and its read settings with Polars 1.x and 2.0. Polars 2.0 also names headerless columns from column_0 and requires new_columns to name every column, so the tests that pass those options through no longer depend on the Polars version. The dev lock tests Polars 2 on Python 3.14, which pull requests run, and keeps the newest Polars 1.x release on older versions. Closes #1168 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…otes Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
|
||
| try: | ||
| if self.output_location: | ||
| with fsspec.open(self.output_location, "rb", **(storage_options or {})) as file: |
There was a problem hiding this comment.
Self-review round one (implementation behavior): CLEAN
Base 6ea5587dde5aae5a5066a91da1f1231803d471e0, head 0c6a0b2de52fe1c3d8d03b79f01b2b3bf52c1ab9.
Covered: _read_csv() (S3 and GetQueryResults sources, timestamp re-read, new_columns renaming, error wrapping), _csv_storage_options, the chunked CSV and Parquet paths (unchanged, already native object store), the four changed tests, and the dev-group Polars markers and lock fork.
_read_kwargs()always setsstorage_options, so the unconditionalpop()is safe; aNonegiven toexecute()opens with no options, as Polars 1.x did (storage_options or {}).- The S3 read is now inside the
try, so S3 errors are still logged asFailed to read <location>and raised asOperationalError, as they were when Polars opened the file. - Polars 1.44.2
read_csv()opened a remote path withfsspec.open(file, **storage_options)(polars/io/_utils.pyprepare_file_arg) and read the whole binary file, so reading the bytes here keeps the same filesystem, options, and memory profile; the timestamp re-read now reuses the bytes instead of opening the object again. AioPolarsCursorandAsyncPolarsCursoruse the same result set.- Regression coverage:
test_csv_storage_options_replace_defaultsfails before the change on both Polars versions (Polars 2.0 rejectsanon=True; on 1.x the patchedfsspec.openis not called). - Validation: Polars suites (AWS) 147 passed on 3.14/Polars 2.0.0 and on 3.12/Polars 1.44.2; SQLAlchemy
-k polars6 passed on 3.14.
| They apply to the reads through PyAthena's filesystem: CSV results read without the chunksize option, which use `"bytes"` and 5 MiB blocks by default, and the Parquet files of UNLOAD results read with the chunksize option when pyarrow is installed and `execute()` is given no Polars read options. | ||
| Of the latter, a file of up to 32 MiB is read at once, and a larger one is read with the given `cache_type`, or without it with the `"none"` cache, which requests only the ranges that pyarrow reads. | ||
| With `storage_options` given to `execute()`, the given options replace PyAthena's S3 settings, and `block_size` and `cache_type` do not apply. | ||
| For a CSV result read without the chunksize option, the options are passed to `fsspec.open()`, which uses PyAthena's S3FileSystem; for the other reads, they are passed to Polars, which uses its native S3 access. |
There was a problem hiding this comment.
Self-review round two (claims and operations): FINDINGS, repaired in 6ec3e00
Base 6ea5587dde5aae5a5066a91da1f1231803d471e0, head 0c6a0b2de52fe1c3d8d03b79f01b2b3bf52c1ab9.
Claims checked against Polars 1.44.2 and 2.0.0 directly: Polars 2.0 read_csv() passes a path and storage_options to scan_csv() (native object store) and rejects non-string values; headerless names are column_1, column_2 in 1.44.2 and column_0, column_1 in 2.0.0; a shorter new_columns is accepted by 1.44.2 and raises SchemaError in 2.0.0. Also checked that _csv_storage_options carries default_block_size, default_cache_type, and max_workers, so the documented block_size/cache_type/s3_max_workers behavior still holds; PR CI runs the newest supported Python (3.14), which now locks Polars 2.0.0.
Findings:
- "Polars 2.0 requires
new_columnsto name every column" overstated it:extra_columns='ignore'accepts a shorter list and drops the remaining columns. Narrowed the test comments and the PR body. docs/polars.mdsaidstorage_optionsgiven toexecute()replace PyAthena's S3 settings but not where they go; with this change PyAthena, not Polars, passes them tofsspec.open()for a non-chunked CSV result. Added the routing for each read (this line).
Repair validation: just lint (3.14), just docs lint, and tests/pyathena/polars/test_result_set.py (18 passed) rerun; comment- and docs-only.
Polars 2.0 reads a Parquet map as a Map column, which get_athena_type() did not know, so UNLOAD results described map columns as string. Nested types are now matched by their outer type, which also describes a struct containing a list as row instead of array. The docs now say that Polars reads with fsspec when use_pyarrow=True is given. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| if base_dtype == pl.Struct: | ||
| return ("row", 0, 0) | ||
| # Polars 2.0 reads a Parquet map as Map; Polars 1.x reads it as a list of structs. | ||
| map_dtype = getattr(pl, "Map", None) |
There was a problem hiding this comment.
Independent review (relayed): FINDINGS
Reviewer: Codex CLI 0.162.1, model gpt-6-astra, model_reasoning_effort=high, sandbox read-only; session 01a13a7f-0062-7242-b5df-82faf4563415.
Base 6ea5587dde5aae5a5066a91da1f1231803d471e0, head 6ec3e0087590eec57e94b9a0c188038e7e45bb68, detached snapshot worktree (unchanged after the review).
Constraints: review only (no edits, builds, tests, or GitHub access); the prompt omitted the PR number, description, commit messages, and prior findings. It was a static review using installed Polars 1.44.2/2.0.0 sources and public sources; nothing was executed.
Reviewer output:
Covered:
- S3 CSV reads, filesystem registration/settings, explicit
storage_optionsincludingNone, reader kwargs/encoding, timestamp retries, and error wrapping. - GetQueryResults fallback and forwarding from all three cursor variants.
- Chunked CSV, native Parquet, PyArrow chunking, schema conversion, and fetch paths.
- Changed tests, comments/docstrings, documentation, dependency markers, lock, and CI Python selection.
- Installed Polars 1.44.2/2.0.0 sources and relevant public versioned sources, including the 1.39.0 batch API.
FINDINGS
-
P2 — Polars 2.0 UNLOAD maps receive incorrect DB-API metadata.
pyathena/polars/util.py:108With
unload=True, a query such asSELECT MAP(ARRAY['a'], ARRAY[1]) AS mproduces a nativepl.Mapunder Polars 2.0.get_athena_type()has no map case, soMap(String, Int32)falls through to"string"with varchar precision. Both whole-result and chunked UNLOAD paths replace their metadata using this function, affecting all three cursor variants. Polars 1.44.2 represented Arrow maps as lists; Polars 2.0 introduces native map conversion.Fix: Recognize the top-level Map type, guarded for Polars 1.x, and return Athena
"map"metadata. Add an explicit map-description assertion: the existing complex UNLOAD test only compares whole/chunked frames and misses this error. -
P3 — The added documentation overstates native S3 routing.
docs/polars.md:400Non-chunked UNLOAD with
execute(..., use_pyarrow=True, storage_options=...)reaches Polars’ PyArrow Parquet branch. That branch usesprepare_file_arg()and fsspec for S3 paths, rather than native object-store access. Consequently, the documented instruction can lead callers to supply native storage-option keys to an fsspec path.Fix: Qualify native routing as the default Polars reader behavior and identify the
use_pyarrow=Trueexception.
The CSV fix itself appears sound. The new test meaningfully checks replacement options and real parsing and would fail against the old Polars 2.0 path. The adjusted naming tests accommodate verified upstream changes, though they remove explicit partial-renaming-with-schema coverage on Polars 1.x. The dependency split matches the stated CI intent.
Static review only; no tests, builds, edits, or GitHub writes performed.
There was a problem hiding this comment.
Author verification and repair (finding 1): confirmed and fixed in a473e9c.
Live UNLOAD check of SELECT MAP(ARRAY['a'], ARRAY[1]) AS m: Polars 1.44.2 described m as array (List(Struct(key, value))); Polars 2.0.0 described it as string (Map(String, Int32)).
get_athena_type() now matches nested types by their outer type (List → array, Struct → row, Map → map when Polars has it). That also fixes an existing defect: the old substring match described Struct({'x': List(Int32)}) as array.
Polars 1.x maps stay array, since they are indistinguishable from a list of key-value structs. The value change (dict in 2.0, list of structs in 1.x) comes from Polars and is listed in the PR's release notes.
Tests: test_get_athena_type_nested, test_get_athena_type_map (skipped on 1.x), and test_nested_type_description_unload (whole and chunked UNLOAD) failed before the fix and pass after it. The Polars suites pass: 154 on 3.14/Polars 2.0.0; 153 plus 1 skipped on 3.12/Polars 1.44.2.
Deferred, with reasons: the reviewer noted that naming every column in the new_columns + schema_overrides tests drops partial-renaming coverage on Polars 1.x. Those options pass through to Polars unchanged, so partial renaming there is Polars' own behavior, and Polars 2.0 rejects it.
There was a problem hiding this comment.
Independent follow-up review (relayed): CLEAN
Reviewer: Codex CLI 0.162.1, model gpt-6-astra, model_reasoning_effort=high, sandbox read-only; session 01a13a8d-16c7-72e0-85b1-3449e37fe5eb.
Scope: repair 6ec3e0087590eec57e94b9a0c188038e7e45bb68..a473e9cd816c3b58d2d782080b49d90b2af3458c (base 6ea5587d unchanged), detached snapshot unchanged after the review. Same constraints as above; the prompt gave the two earlier findings and the affected contracts, without the PR description or author conclusions. Static review only.
Reviewer output:
Covered surfaces:
- Exact repair diff
6ec3e008..a473e9cd; confirmed checkout HEAD. get_athena_type()→to_column_info()→ UNLOAD metadata and DB APIdescription, including chunked and non-chunked paths for all three Polars cursor variants.- Installed Polars 1.44.2 and 2.0.0 datatype and S3 reader sources.
- Added tests,
CLAUDE.mdconventions, and changed comments/documentation.
CLEAN
Both earlier findings are resolved within the reviewed scope.
- Polars 2.0
Mapnow producesmap, including maps containing lists or structs. Lists containing maps remainarray. - Another description intentionally changes: structs containing lists—or field names containing
"list"—now correctly producerowinstead ofarray. Existing scalar and decimal mappings remain unchanged. CSV descriptions and other cursor families do not call this utility. - The guarded
getattr(pl, "Map", None)handles Polars 1.x without accessing a missing attribute. Its list-of-structs map representation remainsarray. Map-specific test expressions execute only inside the skipped test body. - The tests meaningfully distinguish the repair from the previous code through map and struct-with-list cases. Standalone helper tests, declarative datatype parameters, and class-based integration tests using indirect cursor fixtures follow repository conventions. Integration coverage includes both chunking modes.
- The documentation now accounts for fsspec in
read_parquet(use_pyarrow=True). Chunked scan paths continue using native S3 access.
Static review only; no builds, tests, or AWS calls were run. Polars 1.39 compatibility was assessed from the guarded API usage, not execution against that exact version.
| They apply to the reads through PyAthena's filesystem: CSV results read without the chunksize option, which use `"bytes"` and 5 MiB blocks by default, and the Parquet files of UNLOAD results read with the chunksize option when pyarrow is installed and `execute()` is given no Polars read options. | ||
| Of the latter, a file of up to 32 MiB is read at once, and a larger one is read with the given `cache_type`, or without it with the `"none"` cache, which requests only the ranges that pyarrow reads. | ||
| With `storage_options` given to `execute()`, the given options replace PyAthena's S3 settings, and `block_size` and `cache_type` do not apply. | ||
| For a CSV result read without the chunksize option, the options are passed to `fsspec.open()`, which uses PyAthena's S3FileSystem; for the other reads, they are passed to the Polars read function, which uses its native S3 access by default and fsspec with `use_pyarrow=True`. |
There was a problem hiding this comment.
Author verification and repair (finding 2): confirmed and fixed in a473e9c.
In both Polars 1.44.2 and 2.0.0, read_parquet(use_pyarrow=True) calls _read_parquet_with_pyarrow(), which opens the source with prepare_file_arg(..., storage_options=...), that is, with fsspec. The added line now says the Polars read function uses its native S3 access by default and fsspec with use_pyarrow=True. just docs lint passes.
| if "list" in dtype_str: | ||
| # Nested types are matched by their outer type, not by the types they contain. | ||
| if base_dtype == pl.List: | ||
| return ("array", 0, 0) |
There was a problem hiding this comment.
Self-review of the repair (both perspectives): CLEAN
Range: 6ec3e0087590eec57e94b9a0c188038e7e45bb68..a473e9cd816c3b58d2d782080b49d90b2af3458c (same base 6ea5587d).
Round-one perspective (behavior): get_athena_type() is used only by pyathena.polars.util.to_column_info() for UNLOAD/Parquet metadata (result_set.py:674). Comparing the old and new functions on Polars 1.44.2 for Array, Duration, Categorical, Null, Decimal, and List gives identical results; only Struct containing a list (array → row) and Polars 2.0 Map (string → map) change. Live UNLOAD values for map/array columns are unchanged by the description change (2.0: {'a': 1}; 1.44.2: [{'key': 'a', 'value': 1}]). getattr(pl, "Map", None) keeps the 1.39–1.x floor importable, and mypy passes against both Polars versions.
Round-two perspective (claims): the commit message, PR release notes, and docs line were checked against the live check above and the read_parquet(use_pyarrow=True) source in both versions. The tests fail before the fix (see the reply above).
Validation: Polars suites 154 passed (3.14/2.0.0) and 153 passed + 1 skipped (3.12/1.44.2); just lint, just docs lint.
Polars 2.0 converts an Arrow map to a Map column whose values are dicts; Polars 1.x converts it to a list of key-value structs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| assert df.width == 16 | ||
| # Polars 2.0 converts an Arrow map to Map; Polars 1.x to a list of key-value structs. | ||
| if hasattr(pl, "Map"): | ||
| map_dtype = pl.Map(pl.Int32, pl.Int32) |
There was a problem hiding this comment.
CI failure, repair, and self-review of the repair (both perspectives): CLEAN
Test run 38570714905 at a473e9cd816c3b58d2d782080b49d90b2af3458c (Python 3.14, Polars 2.0.0): test-sqla and test-sqla-async passed; test failed only TestArrowCursor::test_complex_unload_as_polars: At index 13 diff: Map(Int32, Int32) != List(Struct({'key': Int32, 'value': Int32})). The PR went back to Draft.
Cause, checked locally: ArrowCursor.as_polars() is pl.from_arrow(self._table). For an Arrow map<int32, int32>, Polars 2.0.0 returns Map(Int32, Int32) with {1: 2, 3: 4}, and Polars 1.44.2 returns List(Struct(key, value)) with [{'key': 1, 'value': 2}, ...]. This is Polars' conversion, so the test now states each major's dtype and value explicitly (repair a473e9cd..22de45412887431ca90c180b46b243abc5ef39c7, test-only).
Behavior perspective: PyAthena code is unchanged; no other test or doc describes the map representation of as_polars(). The full test job found no other failure on Polars 2.0.
Claims perspective: added a release-note item for the as_polars() map representation under Polars 2.0.
Validation: the test passes on 3.14/Polars 2.0.0 and 3.12/Polars 1.44.2; just lint passes.
There was a problem hiding this comment.
Independent follow-up review (relayed): CLEAN
Reviewer: Codex CLI 0.162.1, model gpt-6-astra, model_reasoning_effort=high, sandbox read-only; session 01a13a9d-20f1-74e2-8465-c0e7c368b874.
Scope: a473e9cd816c3b58d2d782080b49d90b2af3458c..22de45412887431ca90c180b46b243abc5ef39c7, detached snapshot unchanged after the review; the prompt gave the CI failure and the conversion path, without the PR description or author conclusions. Static review only.
Reviewer output:
Surfaces covered:
- Frozen diff
a473e9cd..22de4541and the fulltest_complex_unload_as_polarsassertions. CLAUDE.mdtest conventions.- Installed Polars 1.44.2 and 2.0.0 sources.
- Arrow conversion, Polars whole/chunked UNLOAD readers, metadata mapping, and related synchronous/asynchronous tests.
- Repository-wide searches for map-as-list-of-structs assumptions.
CLEAN
The expectations at test_cursor.py:629 correctly distinguish Polars 1.x’s list-of-structs representation from Polars 2.0’s Map(Int32, Int32) and Python dictionary. Both branches specify explicit expected types and values independently of the returned data.
The change preserves the class-based integration test, indirect fixture parametrization, and existing assertions, consistent with CLAUDE.md.
No missed unconditional Polars 1.x map assumption was found. Related metadata handling already guards Map at util.py:112, and the corresponding integration expectation handles both versions at test_cursor.py:789.
Static review only; no tests, builds, or AWS calls were run. Installed-source checks covered 1.44.2 and 2.0.0, not every supported release.
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>
WHAT
Make
PolarsCursorandAioPolarsCursorwork with Polars 2.0 while keeping Polars 1.x support.AthenaPolarsResultSet._read_csv()opens the CSV result file withfsspec.open()and the storage options it used to give Polars (PyAthena's options, or thestorage_optionsgiven toexecute()), and passes the file's bytes topl.read_csv().Polars 1.x did the same
fsspec.open()insideread_csv(); Polars 2.0 reads a path with its native object store and rejects these options (invalid value for 'connection' ... (expected str)).CSV results read without
chunksizetherefore keep using PyAthena's S3 filesystem withblock_size,cache_type, ands3_max_workers, as documented, on both Polars versions.The chunked CSV and Parquet paths already used Polars' native object store and are unchanged.
column_1... in 1.x,column_0... in 2.0).new_columnsgiven together withschema_overridesreach Polars unchanged, and Polars 2.0 rejects a list that does not name every column (unlessextra_columns='ignore'is given, which drops the remaining columns), so the tests name every column.storage_optionsgiven toexecute()reachfsspec.open()and not Polars, replacing the_read_csvcase oftest_storage_options_replace_defaults.pyathena.polars.util.get_athena_type()matches nested types by their outer type and recognizes Polars 2.0'sMap.Polars 2.0 reads a Parquet map of an UNLOAD result as
Map, which was described asstring; it is now described asmap.Polars 1.x reads such a map as a list of key-value structs, which stays
array.Matching by the outer type also describes a struct containing a list as
rowinstead ofarray(wrong before this change on any Polars version).TestArrowCursor::test_complex_unload_as_polarsstates the map dtype and values that each Polars major returns fromArrowCursor.as_polars()(pl.from_arrow()).docs/polars.mdnow states thatstorage_optionsfor a CSV result read withoutchunksizeare passed tofsspec.open(), and for the other reads to the Polars read function, which uses its native S3 access by default and fsspec withuse_pyarrow=True.The
polarsextra stays>=1.39.0.Behavior changes for the release notes:
new_columnsgiven toexecute()together withschema_overridesmust name every column, as Polars 2.0 requires for itsnew_columns; withoutschema_overrides, PyAthena renames the columns itself and a shorter list still works.mapand its values are dicts, as Polars 2.0 returns them; with Polars 1.x it staysarraywith lists of{"key": ..., "value": ...}structs.rowinstead ofarray.ArrowCursor.as_polars()returns a map column aspl.Mapwith dict values, aspl.from_arrow()converts it; with Polars 1.x it is a list of key-value structs.WHY
Closes #1168.
Polars 2.0.0 (released 2026-10-06) satisfies the unbounded
polars>=1.39.0requirement, and every non-chunked CSV read failed with it.This also unblocks Polars on Python 3.15 (#1163, #1164), where only Polars 2.0 works.
TEST
Tested commit: a473e9c (macOS arm64) unless noted
just linton Python 3.14 (Polars 2.0.0) andmypy .on Python 3.12 (Polars 1.44.2): passed.just docs lint: 0 errors.pytest -n 4 tests/pyathena/polars tests/pyathena/aio/polars(AWS): 154 passed on Python 3.14 with Polars 2.0.0; 153 passed and 1 skipped (test_get_athena_type_map, noMapin Polars 1.x) on Python 3.12 with Polars 1.44.2.test_nested_type_description_unloadandtest_get_athena_type_nestedfailed before theget_athena_type()change on both Polars versions, andtest_get_athena_type_mapon Polars 2.0.0.Before the change, master with Polars 2.0.0 failed 50 of these tests on 3.12 and 3.15.
pytest -n 4 tests/pyathena/sqlalchemy tests/pyathena/aio/sqlalchemy -k polarson 3.14 with Polars 2.0.0 at 0c6a0b2: 6 passed.test-sqlaandtest-sqla-asyncpassed;testfailed onlyTestArrowCursor::test_complex_unload_as_polars, which expected Polars 1.x's map conversion. Fixed in 22de454; that test passes on 3.14/Polars 2.0.0 and 3.12/Polars 1.44.2 locally.test,test-sqla, andtest-sqla-asyncall passed.🤖 Generated with Claude Code