diff --git a/docs/polars.md b/docs/polars.md index 78de012ee..38713cf94 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. +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`. @@ -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..42b737268 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -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: + """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( diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 1b1985818..6b05c9c26 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}}, @@ -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(