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
14 changes: 5 additions & 9 deletions docs/polars.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,8 @@ 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; for UNLOAD results with timestamps read with Polars 1.x, see {ref}`polars-unload-options`.
CSV results read without the chunksize option, UNLOAD results read with it when pyarrow is
installed and `execute()` is given no Polars read options, and UNLOAD results with timestamps read
without it with Polars 1.x use PyAthena's own S3FileSystem (fsspec compatible), and other reads use
CSV results read without the chunksize option, and UNLOAD results read with it when pyarrow is
installed and `execute()` is given no Polars read options, use PyAthena's own S3FileSystem (fsspec compatible), and other reads use
Polars' native S3 access, so s3fs is also not required.

Options passed to `execute()`, other than PolarsCursor's own options such as `chunksize`, are passed to the Polars read function.
Expand Down Expand Up @@ -225,11 +224,8 @@ See {ref}`arrow-unload-options` for more information.

Timestamps of UNLOAD results are read at microseconds, as `Datetime("us")`.
Polars 2.0 and later read them with Polars.
Polars 1.x reads them correctly only through pyarrow:

- Without the chunksize option and without Polars read options given to `execute()`, a result with timestamps is read with pyarrow when it is installed, through PyAthena's S3FileSystem, and a warning is logged.
- With the chunksize option and without Polars read options, the chunks are read with pyarrow when it is installed.
- Otherwise Polars 1.x reads the timestamps at nanoseconds, which gives wrong values outside 1677-09-21 to 2262-04-11; they are cast to microseconds, and a warning is logged.
Polars 1.x reads them at nanoseconds, which gives wrong values outside 1677-09-21 to 2262-04-11; they are cast to microseconds, and a warning is logged.
The exception is a read with the chunksize option, with pyarrow installed and without Polars read options given to `execute()`: the chunks are read with pyarrow at microseconds, without the cast or the warning.

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 d7f6710f5b7a8e2bcb736f840eac13d102e87b7c, head 9aa7419dbe963c0c5cc790d63a471c9e9e219cf4.

Claims checked against evidence:

  • "Nothing of this is released": git tag --contains 4ac04db3 (the Read UNLOAD timestamps at microseconds #1233 merge) lists no tag; the latest tag is v4.1.0.
  • "CI runs Polars 2.0": the dev group pins polars>=2.0.0; python_version >= '3.14' and polars<2.0.0; python_version < '3.14' (pyproject.toml), and the Draft run of this PR ran test / run (3.15). The Polars 1.x paths are therefore exercised only by the local matrix recorded in the PR body (1.44.2 and 1.39.0 on Python 3.13).
  • The warning's remedy "read the result in chunks with pyarrow installed and no read options": _iter_parquet_chunks() takes the ParquetChunkReader path exactly under that condition, and the reader opens files with coerce_int96_timestamp_unit="us" (pyathena/arrow/reader.py:169), so the chunks are at microseconds without the cast or the warning, as docs/polars.md now says.
  • Existing callers on Polars 1.x: a whole UNLOAD read without options now goes through pl.read_parquet() with the frozen-credential storage options, as every other Polars read does, instead of PyAthena's S3FileSystem; the output dtypes are unchanged. Live check on Polars 1.44.2 with temporal.nested_query(): the schema is Datetime("us") at every nesting level (top level, List, Struct, List of key-value Structs, nested), and the warning above is logged once. The values of years 1, 1600, and 9999 are wrong there, as documented; in-range values are asserted by test_unload_timestamps_polars_1[whole-no-options].
  • Existing callers on Polars 2.0: unaffected (_int96_read_options() returns the read schema and no warning is logged); the locked suites pass (179 passed, 4 skipped).
  • AWS operations: the Polars 1.x whole read no longer makes the S3 requests of PyAthena's filesystem (one HeadObject/ranged GetObject set per file); Polars' object store reads the files instead, as for Polars 2.0. No retries or extra requests added.
  • Documentation: the three changed passages of docs/polars.md were walked through against _read_parquet(), _iter_parquet_chunks(), and _csv_storage_options; the reference {ref}polars-unload-options on line 14 still resolves to the rewritten section.
  • Evidence: all results in the PR body were obtained at 9aa7419d (lint, docs lint, the suites on 2.0.0/1.44.2/1.39.0, and the mutation check of [whole-no-options]), and the skipped cases are listed with their reasons.

Result: CLEAN within this scope. One limit: the --with polars==1.x overlay on the Python 3.15 environment (whose base holds Polars 2.0.0) is unusable for the matrix; the isolated Python 3.13 runs replace it, as the PR body records.


The unload option can be enabled by specifying it in the `cursor_kwargs` argument of the connect method or as an argument to the cursor method.
An argument to the cursor method takes precedence over the same key in `cursor_kwargs`.
Expand Down Expand Up @@ -405,7 +401,7 @@ cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
df = cursor.execute("SELECT * FROM many_rows").as_polars()
```

They apply to the reads through PyAthena's filesystem: CSV results read without the chunksize option and, with Polars 1.x, the Parquet files of UNLOAD results with timestamps read without it, 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.
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`.
Expand Down
71 changes: 10 additions & 61 deletions pyathena/polars/result_set.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,6 @@

_logger = logging.getLogger(__name__)

# The consequence of a Polars 1.x read of UNLOAD timestamps at nanoseconds.
_NANOSECOND_READ_CONSEQUENCE = (
"so those outside 1677-09-21 to 2262-04-11 are wrong. Upgrade to Polars 2.0 or "
"later, or install pyarrow and give no read options to execute(), to read them "
"correctly."
)
_READ_FAILED_MESSAGE = "PolarsDataFrameIterator cannot be read after a failed read."


Expand Down Expand Up @@ -759,9 +753,8 @@ def _read_parquet(self, schema: pl.Schema) -> pl.DataFrame:
"""Read query results from Parquet files in S3.

Timestamps are read at microseconds; see ``_int96_read_options()``. Polars 1.x
cannot read them so; when ``_reads_with_pyarrow()``, a result with timestamps
is read with pyarrow instead, and otherwise Polars 1.x reads them at
nanoseconds, and they are cast to microseconds.
cannot read them so: it reads them at nanoseconds, and they are cast to
microseconds.

``_prepare_parquet_location()`` must have found the unload location first.

Expand All @@ -780,12 +773,6 @@ def _read_parquet(self, schema: pl.Schema) -> pl.DataFrame:
raise ProgrammingError("unload_location is not available.")

options = self._int96_read_options(schema)
if options is None and self._reads_with_pyarrow():
self._warn_nanosecond_read(
"so they are read with pyarrow instead. Upgrade to Polars 2.0 or later "
"to read them with Polars."
)
return self._read_parquet_with_pyarrow()
dtypes = self._output_dtypes(schema, options)
try:
kwargs = self._read_kwargs(
Expand All @@ -796,33 +783,9 @@ def _read_parquet(self, schema: pl.Schema) -> pl.DataFrame:
_logger.exception(f"Failed to read {self._unload_location}.")
raise OperationalError(*e.args) from e
if options is None and dtypes.keys() & set(df.columns):
self._warn_nanosecond_read(_NANOSECOND_READ_CONSEQUENCE)
self._warn_nanosecond_read()
return df

def _read_parquet_with_pyarrow(self) -> pl.DataFrame:
"""Read the files of the data manifest with pyarrow, with INT96 at microseconds.

Returns:
Polars DataFrame containing the Parquet data.

Raises:
OperationalError: If reading the Parquet files fails.
"""
import polars as pl
from pyarrow import parquet

try:
table = parquet.ParquetDataset(
[path.removeprefix("s3://") for path in self._data_manifest],
filesystem=S3FileSystem(**self._csv_storage_options),
coerce_int96_timestamp_unit="us",
).read(use_threads=True)
# Not rechunked, as Polars does not rechunk its own reads.
return cast("pl.DataFrame", pl.from_arrow(table, rechunk=False))
except Exception as e:
_logger.exception(f"Failed to read {self._unload_location}.")
raise OperationalError(*e.args) from e

def _read_parquet_schema(self) -> pl.Schema:
"""Read the schema of the Parquet files.

Expand Down Expand Up @@ -922,30 +885,16 @@ def _cast(df: pl.DataFrame, dtypes: dict[str, Any]) -> pl.DataFrame:
casts = [df.get_column(name).cast(dtype) for name, dtype in dtypes.items() if name in df]
return df.with_columns(casts) if casts else df

def _reads_with_pyarrow(self) -> bool:
"""Whether the Parquet files are read with pyarrow through PyAthena's filesystem.

They are without read options given to ``execute()``, and with pyarrow
installed: always in chunks, and with Polars 1.x when the files hold
timestamps.

Returns:
True if they are read with pyarrow when pyarrow is chosen for them.
"""
return not self._kwargs and importlib.util.find_spec("pyarrow") is not None

@staticmethod
def _warn_nanosecond_read(consequence: str) -> None:
"""Warn that Polars 1.x reads the timestamps of UNLOAD at nanoseconds.

Args:
consequence: The rest of the message, after "nanoseconds, ".
"""
def _warn_nanosecond_read() -> None:

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 (behavior and implementation), base d7f6710f5b7a8e2bcb736f840eac13d102e87b7c, head 9aa7419dbe963c0c5cc790d63a471c9e9e219cf4 (findings recorded against 6e1e25f4ec37c296d47d6ae294f3761f940bafd8, repaired in 9aa7419d).

Covered: pyathena/polars/result_set.py _read_parquet() (the removed pyarrow branch, the native read, cast, and warning for Polars 1.x), the removed _read_parquet_with_pyarrow() and _reads_with_pyarrow(), _iter_parquet_chunks() / _read_parquet_chunks() / _scan_parquet_chunks() (the chunked pyarrow read stays), _warn_nanosecond_read(), the docstrings of _int96_read_options() and _output_dtypes(), the aio Polars cursors (same result set), docs/polars.md (the filesystem reads, the UNLOAD timestamps section, the S3 read options section), and the three changed tests. Live matrix on Polars 2.0.0 (locked, Python 3.15), 1.44.2, and 1.39.0 (Python 3.13). Also ran /code-review high on the PR (7 candidates).

Result: FINDINGS, repaired in 9aa7419d:

  • R1-1 (code-review) _warn_nanosecond_read(): the whole read and the chunked scan logged different remedies, and the whole read's omitted the chunked read with pyarrow that docs/polars.md describes as reading the timestamps correctly. One message now names both remedies, assembled in one place.
  • R1-2 (code-review) tests/pyathena/polars/test_cursor.py: the new whole-read case had taken the existing ID whole, against the rule to preserve test IDs. The existing cases keep their IDs and the new case is whole-no-options; the mutation check was repeated (fails with the fallback restored).
  • R1-3 (code-review) _reads_with_pyarrow() had one caller left and a docstring that said "in chunks" without checking chunksize; inlined at the caller in _iter_parquet_chunks().
  • R1-4 (code-review) docs/polars.md did not say that the chunked read with pyarrow has neither the cast nor the warning; a sentence says so.

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), base d7f6710f5b7a8e2bcb736f840eac13d102e87b7c, head 9aa7419dbe963c0c5cc790d63a471c9e9e219cf4.

Reviewer: Codex CLI 0.163.0, model gpt-6-astra, reasoning effort max, session 01a14e45-c79d-75b3-9301-7561bc25f209, read-only sandbox, stdin closed, in a detached checkout of the head. The package held the literal diff, AGENTS.md, and the testing guide's "Choose what to test" section; no PR number, PR body, commit messages, prior findings, or agent memory. Static review only: no builds, tests, or GitHub access. The snapshot and the PR worktree were unchanged afterwards.

Covered (reviewer's list): the diff; whole and chunked reads across Polars 1.x/2.0, optional pyarrow, and read options; scalar and nested timestamps, schema overrides, warnings, manifest selection, and error propagation through the sync, async, and aio cursors; the changed and surrounding documentation, temporal fixtures, test coverage and skips, and removed-path references and imports.

Result: FINDINGS (4), verified by the author:

  1. "Whole reads can omit partitions" (_prepare_parquet_location() derives the unload location from the first manifest entry when the formatter gave none; a user-written partitioned UNLOAD would lose other partitions; before, the Polars 1.x pyarrow fallback read every manifest entry). Rejected as a regression: the derivation is the existing logic of this result set, of ArrowResultSet._read_parquet() (pyathena/arrow/result_set.py:396), and of the Polars 2.0 and read-option paths; PyAthena's unload option wraps a SELECT in its own UNLOAD without partitioned_by, so the scenario is not a supported one, and this PR makes Polars 1.x read as Polars 2.0 does.
  2. schema=None given to execute() disables the timestamp handling and the warning ("schema" in self._kwargs). Pre-existing, outside this PR's scope; reported to the maintainer for a decision.
  3. use_pyarrow=True on Polars 2.0 fails on the schema= that _int96_read_options() adds. Pre-existing and known: recorded in Pre-existing UNLOAD read defects found in the #1214 review #1235 as reachable only once read_parquet(use_pyarrow=True) cannot read a remote directory or a list of remote files pola-rs/polars#30038 is fixed, since the prefix read fails first.
  4. The skips left the Polars 1.x native cast of nested timestamps without a test with in-range values. Accepted: test_unload_nested_timestamps_polars_1 (whole and scanned chunks) asserts Datetime("us") inside arrays, rows, and maps and the values of a 2020 timestamp and NULL; with the nested cast disabled both cases fail (2/2). Repaired in 259cdb33. The repair was checked under both self-review perspectives: it adds a test only, its docstring claim (values within 1677-2262) matches the query, and it passes on Polars 1.44.2 and 1.39.0 and is skipped on 2.0.0. The loss of the UNLOAD date-boundary coverage on Polars 1.x from the test_temporal_types[unload] skip stays as recorded in round 1 (CI runs Polars 2.0).

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) on the repair 9aa7419dbe963c0c5cc790d63a471c9e9e219cf4..259cdb334756717b4ec82c108ba606811707596f (the added test only). Reviewer: Codex CLI 0.163.0, model gpt-6-astra, reasoning effort max, session 01a14e53-bde4-74c0-9b61-da2efc9bac23, read-only sandbox, detached checkout of 259cdb33, given the literal patch and the testing guide's section; static review only.

Result: CLEAN. Reviewer's notes: both cases exercise the native Polars 1.x cast (low_memory=True keeps the scan off the pyarrow chunk reader); the nested schema assertions fail if _cast() or the recursive to_microseconds() is removed, even where the Python values would still compare equal; the value assertions also check the scaling and NULLs; the skip condition and docstring are correct; the existing nested coverage skips native Polars 1.x reads, so the test adds distinct evidence.

"""Warn that Polars 1.x reads the timestamps of UNLOAD at nanoseconds."""
import polars as pl

_logger.warning(
f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at "
f"nanoseconds, {consequence}"
"nanoseconds, so those outside 1677-09-21 to 2262-04-11 are wrong. Upgrade to "
"Polars 2.0 or later, or read the result in chunks with pyarrow installed and "
"no read options given to execute(), to read them correctly."
)

def _as_polars(self) -> pl.DataFrame:
Expand Down Expand Up @@ -1158,7 +1107,7 @@ def _iter_parquet_chunks(self, schema: pl.Schema) -> Iterator[pl.DataFrame]:
if self._unload_location is None:
raise ProgrammingError("unload_location is not available.")

if self._reads_with_pyarrow():
if not self._kwargs and importlib.util.find_spec("pyarrow") is not None:
return self._read_parquet_chunks(
S3FileSystem(**self._csv_storage_options), self._data_manifest
)
Expand Down Expand Up @@ -1218,7 +1167,7 @@ def _scan_parquet_chunks(
options = self._int96_read_options(schema)
dtypes = self._output_dtypes(schema, options)
if options is None and dtypes:
self._warn_nanosecond_read(_NANOSECOND_READ_CONSEQUENCE)
self._warn_nanosecond_read()
lazy_df = pl.scan_parquet(
location,
**self._read_kwargs(
Expand Down
70 changes: 60 additions & 10 deletions tests/pyathena/polars/test_cursor.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,15 +127,21 @@ def test_many_as_polars(self, polars_cursor):
("polars_cursor", "unload"),
[
({"cursor_kwargs": {"unload": False}}, False),
({"cursor_kwargs": {"unload": True}}, True),
pytest.param(
{"cursor_kwargs": {"unload": True}},
True,
marks=pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x reads at nanoseconds."),
),
],
indirect=["polars_cursor"],
ids=["csv", "unload"],
)
def test_temporal_types(self, polars_cursor, unload):

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.

Round 1, code-review candidates checked and not changed:

  • test_temporal_types[unload] is skipped as a whole on Polars 1.x rather than per column: the query's years outside 1677-2262 are out of Polars 1.x's reach by design, and a per-column exception would write Polars 1.x's defect into the contract test. CI runs Polars 2.0 (the dev group pins polars>=2.0.0 on Python 3.14+ and the Test workflow runs 3.15), so the UNLOAD coverage of dates and text fallbacks stays in CI; test_unload_timestamps_polars_1 covers the Polars 1.x UNLOAD read with in-range timestamps.
  • test_unload_nested_timestamps[chunks] has no pyarrow mark and would fail on Polars 1.x without pyarrow: pre-existing, and pyarrow is in the dev dependency group, so the test environment always has it.

"""Temporal values convert as the contract of #1212 defines.

Types that Polars cannot represent stay Athena's text.
Types that Polars cannot represent stay Athena's text. Polars 1.x reads the
UNLOAD timestamps at nanoseconds, which the years outside 1677-2262 of the
query exceed; see test_unload_timestamps_polars_1 for its read.
"""
types = {"timestamp": pl.Datetime("us"), "date": pl.Date()}
columns = temporal.columns(unload)
Expand Down Expand Up @@ -167,7 +173,11 @@ def test_temporal_types_chunked(self, polars_cursor):
@pytest.mark.parametrize(
("polars_cursor", "kwargs"),
[
({"cursor_kwargs": {"unload": True}}, {}),
pytest.param(
{"cursor_kwargs": {"unload": True}},
{},
marks=pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x reads at nanoseconds."),
),
({"cursor_kwargs": {"unload": True, "chunksize": 1}}, {}),
pytest.param(
{"cursor_kwargs": {"unload": True, "chunksize": 1}},
Expand Down Expand Up @@ -217,17 +227,17 @@ def test_unload_nested_timestamps(self, polars_cursor, kwargs):
True,
),
({"cursor_kwargs": {"unload": True}}, {"columns": ["id"]}, ["id"], False),
({"cursor_kwargs": {"unload": True}}, {}, ["id", "ts"], True),
],
indirect=["polars_cursor"],
ids=["whole", "scanned-chunks", "no-timestamps"],
ids=["whole", "scanned-chunks", "no-timestamps", "whole-no-options"],
)
def test_unload_timestamps_with_read_options_polars_1(
self, polars_cursor, caplog, kwargs, columns, warned
):
"""With read options, Polars 1.x reads the timestamps itself and a warning is logged.
def test_unload_timestamps_polars_1(self, polars_cursor, caplog, kwargs, columns, warned):
"""Polars 1.x reads the timestamps itself at nanoseconds, and a warning is logged.

The timestamps are cast to microseconds; the warning is logged only when the read
has timestamps.
The timestamps are cast to microseconds; the warning, which says that values
outside 1677-09-21 to 2262-04-11 are wrong, is logged only when the read has
timestamps.
"""
query = "SELECT 1 AS id, TIMESTAMP '2020-01-02 03:04:05.123' AS ts"
with caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"):
Expand All @@ -243,6 +253,46 @@ def test_unload_timestamps_with_read_options_polars_1(
if r.name == "pyathena.polars.result_set" and r.levelno == logging.WARNING
]
assert bool(warnings) is warned
assert all("outside 1677-09-21 to 2262-04-11 are wrong" in r.getMessage() for r in warnings)

@pytest.mark.skipif(_POLARS_2, reason="Polars 2.0 reads INT96 at microseconds.")
@pytest.mark.parametrize(
("polars_cursor", "kwargs"),
[
({"cursor_kwargs": {"unload": True}}, {}),
({"cursor_kwargs": {"unload": True, "chunksize": 1}}, {"low_memory": True}),
],
indirect=["polars_cursor"],
ids=["whole", "scanned-chunks"],
)
def test_unload_nested_timestamps_polars_1(self, polars_cursor, kwargs):
"""Polars 1.x casts the timestamps inside arrays, rows, and maps to microseconds too.

The values are within 1677-2262, which Polars 1.x reads correctly; see
test_unload_timestamps_polars_1 for the warning.
"""
polars_cursor.execute(
"SELECT id, ts, ARRAY[ts] AS arr, CAST(ROW(ts) AS ROW(t TIMESTAMP)) AS st, "
"MAP(ARRAY['k'], ARRAY[ts]) AS mp "
"FROM (VALUES (1, TIMESTAMP '2020-01-02 03:04:05.123'), "
"(2, CAST(NULL AS TIMESTAMP))) AS t (id, ts)",
**kwargs,
)
df = pl.concat(polars_cursor.iter_chunks()).sort("id")
us = pl.Datetime("us")
assert df.drop("id").schema == pl.Schema(
{
"ts": us,
"arr": pl.List(us),
"st": pl.Struct({"t": us}),
"mp": pl.List(pl.Struct({"key": pl.String, "value": us})),
}
)
values = [datetime(2020, 1, 2, 3, 4, 5, 123000), None]
assert df["ts"].to_list() == values
assert df["arr"].to_list() == [[v] for v in values]
assert df["st"].to_list() == [{"t": v} for v in values]
assert df["mp"].to_list() == [[{"key": "k", "value": v}] for v in values]

@pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x reads INT96 at nanoseconds.")
@pytest.mark.parametrize(
Expand Down
Loading