From 6e1e25f4ec37c296d47d6ae294f3761f940bafd8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 17:36:01 +0900 Subject: [PATCH 1/3] Remove the Polars 1.x pyarrow fallback for whole UNLOAD reads with timestamps (GH-1280) For a whole UNLOAD read on Polars 1.x without read options, a result with timestamps was read with pyarrow through PyAthena's S3FileSystem to work around Polars 1.x reading INT96 at nanoseconds. PyAthena does not switch readers on the user's behalf for a dependency's limitation: Polars 1.x now reads the result itself in every case, at nanoseconds, cast to microseconds, with the warning that recommends Polars 2.0. The chunked read through ParquetChunkReader stays, as it bounds memory use. Co-Authored-By: Claude Fable 5.1 --- docs/polars.md | 14 +++---- pyathena/polars/result_set.py | 63 +++++++--------------------- tests/pyathena/polars/test_cursor.py | 31 +++++++++----- 3 files changed, 42 insertions(+), 66 deletions(-) diff --git a/docs/polars.md b/docs/polars.md index 78de012ee..d53ef6ec3 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -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. @@ -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. +With the chunksize option and without Polars read options given to `execute()`, the chunks are read with pyarrow when it is installed, which reads the timestamps at microseconds. 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`. @@ -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`. diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index d4017a402..d6bb80e75 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -47,11 +47,12 @@ _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." +# The remedies of a Polars 1.x read of UNLOAD timestamps at nanoseconds; see +# _warn_nanosecond_read(). +_NANOSECOND_READ_REMEDY = "Upgrade to Polars 2.0 or later to read them correctly." +_NANOSECOND_SCAN_REMEDY = ( + "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 +760,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 +780,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 +790,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(_NANOSECOND_READ_REMEDY) 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. @@ -923,29 +893,28 @@ def _cast(df: pl.DataFrame, dtypes: dict[str, Any]) -> pl.DataFrame: 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. + """Whether the Parquet files are read in chunks 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. + installed. Returns: - True if they are read with pyarrow when pyarrow is chosen for them. + True if the chunks are read with pyarrow. """ return not self._kwargs and importlib.util.find_spec("pyarrow") is not None @staticmethod - def _warn_nanosecond_read(consequence: str) -> None: + def _warn_nanosecond_read(remedy: str) -> None: """Warn that Polars 1.x reads the timestamps of UNLOAD at nanoseconds. Args: - consequence: The rest of the message, after "nanoseconds, ". + remedy: The last sentence of the message, after the consequence. """ import polars as pl _logger.warning( f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at " - f"nanoseconds, {consequence}" + f"nanoseconds, so those outside 1677-09-21 to 2262-04-11 are wrong. {remedy}" ) def _as_polars(self) -> pl.DataFrame: @@ -1218,7 +1187,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(_NANOSECOND_SCAN_REMEDY) lazy_df = pl.scan_parquet( location, **self._read_kwargs( diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 1b1985818..e413e0911 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -127,7 +127,11 @@ 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"], @@ -135,7 +139,9 @@ def test_many_as_polars(self, polars_cursor): def test_temporal_types(self, polars_cursor, unload): """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}}, @@ -209,6 +219,7 @@ def test_unload_nested_timestamps(self, polars_cursor, kwargs): @pytest.mark.parametrize( ("polars_cursor", "kwargs", "columns", "warned"), [ + ({"cursor_kwargs": {"unload": True}}, {}, ["id", "ts"], True), ({"cursor_kwargs": {"unload": True}}, {"columns": ["ts"]}, ["ts"], True), ( {"cursor_kwargs": {"unload": True, "chunksize": 1}}, @@ -219,15 +230,14 @@ def test_unload_nested_timestamps(self, polars_cursor, kwargs): ({"cursor_kwargs": {"unload": True}}, {"columns": ["id"]}, ["id"], False), ], indirect=["polars_cursor"], - ids=["whole", "scanned-chunks", "no-timestamps"], + ids=["whole", "whole-read-options", "scanned-chunks", "no-timestamps"], ) - 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,7 @@ 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(not _POLARS_2, reason="Polars 1.x reads INT96 at nanoseconds.") @pytest.mark.parametrize( From 9aa7419dbe963c0c5cc790d63a471c9e9e219cf4 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 17:59:51 +0900 Subject: [PATCH 2/3] Give the nanosecond warning one remedy and keep the test IDs (GH-1280) The whole read and the chunked scan logged different remedies, and the whole read's left out the chunked read with pyarrow, which the documentation describes. The warning now names both remedies in one place. The predicate of the chunked pyarrow read is inlined at its only caller, the documentation says that the chunked read with pyarrow has neither the cast nor the warning, and the new test case takes a new ID. Co-Authored-By: Claude Fable 5.1 --- docs/polars.md | 2 +- pyathena/polars/result_set.py | 36 +++++++--------------------- tests/pyathena/polars/test_cursor.py | 4 ++-- 3 files changed, 11 insertions(+), 31 deletions(-) diff --git a/docs/polars.md b/docs/polars.md index d53ef6ec3..38713cf94 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -225,7 +225,7 @@ 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 at nanoseconds, which gives wrong values outside 1677-09-21 to 2262-04-11; they are cast to microseconds, and a warning is logged. -With the chunksize option and without Polars read options given to `execute()`, the chunks are read with pyarrow when it is installed, which reads the timestamps at microseconds. +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. 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`. diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index d6bb80e75..42b737268 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -47,13 +47,6 @@ _logger = logging.getLogger(__name__) -# The remedies of a Polars 1.x read of UNLOAD timestamps at nanoseconds; see -# _warn_nanosecond_read(). -_NANOSECOND_READ_REMEDY = "Upgrade to Polars 2.0 or later to read them correctly." -_NANOSECOND_SCAN_REMEDY = ( - "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." @@ -790,7 +783,7 @@ 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_REMEDY) + self._warn_nanosecond_read() return df def _read_parquet_schema(self) -> pl.Schema: @@ -892,29 +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 in chunks with pyarrow through PyAthena's filesystem. - - They are without read options given to ``execute()``, and with pyarrow - installed. - - Returns: - True if the chunks are read with pyarrow. - """ - return not self._kwargs and importlib.util.find_spec("pyarrow") is not None - @staticmethod - def _warn_nanosecond_read(remedy: str) -> None: - """Warn that Polars 1.x reads the timestamps of UNLOAD at nanoseconds. - - Args: - remedy: The last sentence of the message, after the consequence. - """ + def _warn_nanosecond_read() -> None: + """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, so those outside 1677-09-21 to 2262-04-11 are wrong. {remedy}" + "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: @@ -1127,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 ) @@ -1187,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_SCAN_REMEDY) + self._warn_nanosecond_read() lazy_df = pl.scan_parquet( location, **self._read_kwargs( diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index e413e0911..783ccc668 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -219,7 +219,6 @@ def test_unload_nested_timestamps(self, polars_cursor, kwargs): @pytest.mark.parametrize( ("polars_cursor", "kwargs", "columns", "warned"), [ - ({"cursor_kwargs": {"unload": True}}, {}, ["id", "ts"], True), ({"cursor_kwargs": {"unload": True}}, {"columns": ["ts"]}, ["ts"], True), ( {"cursor_kwargs": {"unload": True, "chunksize": 1}}, @@ -228,9 +227,10 @@ 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", "whole-read-options", "scanned-chunks", "no-timestamps"], + ids=["whole", "scanned-chunks", "no-timestamps", "whole-no-options"], ) 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. From 259cdb334756717b4ec82c108ba606811707596f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 18:22:03 +0900 Subject: [PATCH 3/3] Test the Polars 1.x cast of nested UNLOAD timestamps with in-range values (GH-1280) The nested timestamp test is skipped on Polars 1.x, whose native read is now the only whole read, so the cast of timestamps inside arrays, rows, and maps to microseconds had no test there. Co-Authored-By: Claude Fable 5.1 --- tests/pyathena/polars/test_cursor.py | 39 ++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 783ccc668..6b05c9c26 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -255,6 +255,45 @@ def test_unload_timestamps_polars_1(self, polars_cursor, caplog, kwargs, columns 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( ("polars_cursor", "kwargs"),