Repository navigation
Remove the Polars 1.x pyarrow fallback for whole UNLOAD reads with timestamps (GH-1280) #1282
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
6e1e25f
9aa7419
259cdb3
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 |
|---|---|---|
|
|
@@ -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." | ||
|
|
||
|
|
||
|
|
@@ -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. | ||
|
|
||
|
|
@@ -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( | ||
|
|
@@ -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. | ||
|
|
||
|
|
@@ -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: | ||
|
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 (behavior and implementation), base Covered: Result: FINDINGS, repaired in
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), base Reviewer: Codex CLI 0.163.0, model 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:
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) on the repair Result: CLEAN. Reviewer's notes: both cases exercise the native Polars 1.x cast ( |
||
| """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: | ||
|
|
@@ -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 | ||
| ) | ||
|
|
@@ -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( | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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): | ||
|
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. Round 1, code-review candidates checked and not changed:
|
||
| """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) | ||
|
|
@@ -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}}, | ||
|
|
@@ -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"): | ||
|
|
@@ -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( | ||
|
|
||
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.
Self-review round 2 (claims, callers, operations), base
d7f6710f5b7a8e2bcb736f840eac13d102e87b7c, head9aa7419dbe963c0c5cc790d63a471c9e9e219cf4.Claims checked against evidence:
git tag --contains 4ac04db3(the Read UNLOAD timestamps at microseconds #1233 merge) lists no tag; the latest tag is v4.1.0.devgroup pinspolars>=2.0.0; python_version >= '3.14'andpolars<2.0.0; python_version < '3.14'(pyproject.toml), and the Draft run of this PR rantest / 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)._iter_parquet_chunks()takes theParquetChunkReaderpath exactly under that condition, and the reader opens files withcoerce_int96_timestamp_unit="us"(pyathena/arrow/reader.py:169), so the chunks are at microseconds without the cast or the warning, asdocs/polars.mdnow says.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 withtemporal.nested_query(): the schema isDatetime("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 bytest_unload_timestamps_polars_1[whole-no-options]._int96_read_options()returns the read schema and no warning is logged); the locked suites pass (179 passed, 4 skipped).HeadObject/rangedGetObjectset per file); Polars' object store reads the files instead, as for Polars 2.0. No retries or extra requests added.docs/polars.mdwere walked through against_read_parquet(),_iter_parquet_chunks(), and_csv_storage_options; the reference{ref}polars-unload-optionson line 14 still resolves to the rewritten section.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.xoverlay 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.