From 8f98061145ddfdbb9c18a2c3de69a7b33f814786 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 12:26:36 +0900 Subject: [PATCH 01/12] Read UNLOAD timestamps at microseconds Athena's UNLOAD writes timestamps as Parquet INT96, which the readers decoded at nanoseconds, so values outside 1677-09-21 to 2262-04-11 were wrapped or saturated, and the dtype was ns instead of the us of CSV results. PandasCursor and ArrowCursor read INT96 at microseconds with pyarrow's coerce_int96_timestamp_unit, which an execute() option given to pandas.read_parquet() still replaces. ParquetChunkReader does the same for the chunked reads of PandasCursor and PolarsCursor. Polars 2.0 reads INT96 at the unit of the read schema, so PolarsCursor reads with the file schema in microseconds. Map values still saturate in Polars 2.0, so maps are read as lists of key-value structs and cast back. Polars 1.x reads INT96 only at nanoseconds, so a result with timestamps is read with pyarrow and a warning recommends Polars 2.0; it raises NotSupportedError without pyarrow or with Polars read options, except use_pyarrow=True, which gets coerce_int96_timestamp_unit in pyarrow_options. Co-Authored-By: Claude Opus 5.5 --- docs/arrow.md | 2 +- docs/pandas.md | 1 + docs/polars.md | 11 +- pyathena/arrow/reader.py | 5 +- pyathena/arrow/result_set.py | 14 +- pyathena/pandas/result_set.py | 24 +- pyathena/polars/result_set.py | 214 +++++++++++++++--- pyathena/polars/util.py | 50 ++++ tests/pyathena/arrow/test_cursor.py | 4 +- tests/pyathena/arrow/test_reader.py | 16 +- tests/pyathena/arrow/test_result_set.py | 34 +++ tests/pyathena/pandas/test_result_set.py | 49 +++- tests/pyathena/polars/test_result_set.py | 148 +++++++++++- tests/pyathena/polars/test_util.py | 56 ++++- tests/pyathena/util.py | 14 ++ .../resources/unload/int96_timestamps.parquet | Bin 0 -> 1254 bytes 16 files changed, 595 insertions(+), 47 deletions(-) create mode 100644 tests/resources/unload/int96_timestamps.parquet diff --git a/docs/arrow.md b/docs/arrow.md index cec6bd1bd..649290dba 100644 --- a/docs/arrow.md +++ b/docs/arrow.md @@ -199,7 +199,7 @@ ArrowCursor supports the unload option. When this option is enabled, queries that start with SELECT or WITH are automatically converted to unload statements and executed to Athena, and the results are output in Parquet format (Snappy compressed) under `{s3_staging_dir}unload///`. This option requires `s3_staging_dir`; without it, `execute()` raises `ProgrammingError`. -The cursor reads the output Parquet file directly. +The cursor reads the output Parquet file directly, with timestamps at microseconds, as `timestamp[us]`. The output of query results with the unload statement is faster than normal query execution. In addition, the output Parquet file is split and can be read faster than a CSV file. diff --git a/docs/pandas.md b/docs/pandas.md index afd591f81..67b8fd3ca 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -604,6 +604,7 @@ With `storage_options` given to `execute()`, or `filesystem` for an UNLOAD resul PandasCursor also supports the unload option, as does {ref}`arrow-cursor`. Reading the unloaded Parquet files requires pyarrow (`pip install PyAthena[Pandas,Arrow]`), and `engine` must be `auto` or `pyarrow`. +Timestamps are read at microseconds, as `datetime64[us]`, unless `coerce_int96_timestamp_unit` is given to `execute()`. See {ref}`arrow-unload-options` for more information. diff --git a/docs/polars.md b/docs/polars.md index b72f86dc2..022f6ab27 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -11,7 +11,8 @@ This cursor downloads the CSV file after executing the query and loads it into a 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. CSV results read without the chunksize option, and +does not require PyArrow as a dependency, except for UNLOAD results with timestamps read with Polars 1.x; see {ref}`polars-unload-options`. +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. @@ -214,12 +215,20 @@ cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", If the unload option is enabled, the Parquet file itself has a schema, so the conversion is done to the dtypes according to that schema, and the `types` setting of the Converter class is not used. +(polars-unload-options)= + ### Unload Options PolarsCursor supports the unload option, as does {ref}`arrow-cursor`. 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. +With Polars 1.x, a result with timestamps is read with pyarrow, through PyAthena's S3FileSystem, and a warning is logged. +This needs pyarrow, and no Polars read options given to `execute()`; otherwise the read raises `NotSupportedError`. +With `use_pyarrow=True` given to `execute()`, a result read without the chunksize option is read by Polars with pyarrow, with `coerce_int96_timestamp_unit="us"` added to `pyarrow_options` unless they set it. + 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/arrow/reader.py b/pyathena/arrow/reader.py index e73a73860..cbc7467ca 100644 --- a/pyathena/arrow/reader.py +++ b/pyathena/arrow/reader.py @@ -41,7 +41,8 @@ class ParquetChunkReader(abc.Iterator): # type: ignore[type-arg] coalesced ranges, which PyAthena's filesystem splits into parallel requests, and holds them until the row group is decoded. - The files must have the same schema, as the files of one UNLOAD query do. + The files must have the same schema, as the files of one UNLOAD query do. INT96 + timestamps, which Athena writes, are read at microseconds. Example: >>> with fs.open(manifest) as f: @@ -165,7 +166,7 @@ def _read_file(self, source: Any) -> Generator[RecordBatch]: """ import pyarrow.parquet as pq - parquet_file = pq.ParquetFile(source, pre_buffer=True) + parquet_file = pq.ParquetFile(source, pre_buffer=True, coerce_int96_timestamp_unit="us") if self._columns is not None and ( missing := set(self._columns) - set(parquet_file.schema_arrow.names) ): diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 6f7742b60..3bb2117e8 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -379,6 +379,14 @@ def _read_csv(self) -> Table: raise OperationalError(*e.args) from e def _read_parquet(self) -> Table: + """Read the UNLOAD result, with timestamps at microseconds. + + Returns: + The table of the result, or an empty table without data files. + + Raises: + OperationalError: If reading the Parquet files fails. + """ import pyarrow as pa from pyarrow import parquet @@ -390,7 +398,11 @@ def _read_parquet(self) -> Table: bucket, key = parse_output_location(self._unload_location) try: - dataset = parquet.ParquetDataset(f"{bucket}/{key}", filesystem=self._fs) + # Athena writes timestamps as INT96, whose default unit, nanoseconds, cannot + # hold years before 1677 or after 2262. + dataset = parquet.ParquetDataset( + f"{bucket}/{key}", filesystem=self._fs, coerce_int96_timestamp_unit="us" + ) return dataset.read(use_threads=True) except Exception as e: _logger.exception(f"Failed to read {bucket}/{key}.") diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index a8d26bca4..ce0cdb12f 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -1469,6 +1469,22 @@ def _open_binary_csv_stream( return stream def _read_parquet(self, engine) -> DataFrame: + """Read the UNLOAD result whole with ``pandas.read_parquet()``. + + The options given to ``execute()`` are passed to ``pandas.read_parquet()``. + Timestamps are read at microseconds unless ``coerce_int96_timestamp_unit`` is + given. + + Args: + engine: The Parquet engine, which must be ``"pyarrow"``. + + Returns: + The DataFrame of the result, or an empty DataFrame without data files. + + Raises: + OperationalError: If reading the Parquet files fails. + ProgrammingError: If the engine is not ``"pyarrow"``. + """ import pandas as pd self._data_manifest = self._read_data_manifest() @@ -1478,7 +1494,13 @@ def _read_parquet(self, engine) -> DataFrame: self._unload_location = "/".join(self._data_manifest[0].split("/")[:-1]) + "/" if engine == "pyarrow": - kwargs: dict[str, Any] = {"use_threads": True, **self._kwargs} + kwargs: dict[str, Any] = { + "use_threads": True, + # Athena writes timestamps as INT96, whose default unit, nanoseconds, + # cannot hold years before 1677 or after 2262. + "coerce_int96_timestamp_unit": "us", + **self._kwargs, + } # Given storage_options, even None, pandas opens the files itself, # as for CSV results. if "filesystem" not in kwargs and "storage_options" not in kwargs: diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 3d29a7acb..51bd34cf2 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -33,11 +33,11 @@ from pyathena import OperationalError from pyathena.arrow.reader import ParquetChunkReader from pyathena.converter import Converter, _to_default -from pyathena.error import ProgrammingError +from pyathena.error import NotSupportedError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.polars.converter import _to_datetimes -from pyathena.polars.util import to_column_info +from pyathena.polars.util import holds_datetimes, to_column_info, to_microseconds from pyathena.result_set import AthenaResultSet from pyathena.util import RetryConfig, override @@ -636,8 +636,9 @@ def _create_dataframe_iterator(self) -> PolarsDataFrameIterator: if not self.is_unload: reader = self._iter_csv_chunks() elif self._prepare_parquet_location(): - self._metadata = self._read_parquet_schema() - reader = self._iter_parquet_chunks() + schema = self._read_parquet_schema() + self._metadata = to_column_info(schema) + reader = self._iter_parquet_chunks(schema) else: self._metadata = () reader = iter(()) @@ -763,34 +764,88 @@ def _read_csv(self) -> pl.DataFrame: _logger.exception(f"Failed to read {location}.") raise OperationalError(*e.args) from e - def _read_parquet(self) -> pl.DataFrame: + 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 without ``use_pyarrow=True``, and a result with timestamps + is read with pyarrow instead, which needs pyarrow and no read options given to + ``execute()``. + + ``_prepare_parquet_location()`` must have found the unload location first. + + Args: + schema: The schema of the Parquet files, from ``_read_parquet_schema()``. + Returns: Polars DataFrame containing the Parquet data. Raises: + NotSupportedError: If Polars 1.x would read timestamps at nanoseconds. OperationalError: If reading the Parquet files fails. """ import polars as pl - if not self._prepare_parquet_location(): - return pl.DataFrame() - if self._unload_location is None: raise ProgrammingError("unload_location is not available.") - try: - return pl.read_parquet( - self._unload_location, - **self._read_kwargs(lambda: self._get_parquet_storage_options(self.connection)), + options = self._int96_read_options(schema) + if options is None: + if self._kwargs or importlib.util.find_spec("pyarrow") is None: + raise self._int96_read_error() + _logger.warning( + f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at " + "nanoseconds, so they are read with pyarrow instead. Upgrade to " + "Polars 2.0 or later to read them with Polars." + ) + df = self._read_parquet_with_pyarrow() + else: + kwargs = self._read_kwargs( + lambda: self._get_parquet_storage_options(self.connection), **options ) + if "pyarrow_options" in options: + # The options given to execute() are in it. + kwargs["pyarrow_options"] = options["pyarrow_options"] + try: + df = pl.read_parquet(self._unload_location, **kwargs) + except Exception as e: + _logger.exception(f"Failed to read {self._unload_location}.") + raise OperationalError(*e.args) from e + return self._restore_maps(df, schema, options) + + 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) + return cast("pl.DataFrame", pl.from_arrow(table)) except Exception as e: _logger.exception(f"Failed to read {self._unload_location}.") raise OperationalError(*e.args) from e - def _read_parquet_schema(self) -> tuple[dict[str, Any], ...]: - """Read schema from Parquet files for metadata.""" + def _read_parquet_schema(self) -> pl.Schema: + """Read the schema of the Parquet files. + + Returns: + The schema, with INT96 timestamps as nanosecond Datetimes. + + Raises: + ProgrammingError: If the unload location is not set. + OperationalError: If reading the schema fails. + """ import polars as pl if not self._unload_location: @@ -804,12 +859,88 @@ def _read_parquet_schema(self) -> tuple[dict[str, Any], ...]: lambda: self._get_parquet_storage_options(self.connection) ), ) - schema = lazy_df.collect_schema() - return to_column_info(schema) + return lazy_df.collect_schema() except Exception as e: _logger.exception(f"Failed to read schema from {self._unload_location}.") raise OperationalError(*e.args) from e + def _int96_read_options(self, schema: pl.Schema) -> dict[str, Any] | None: + """Get the Polars read arguments that read the INT96 timestamps of UNLOAD at microseconds. + + Polars reads INT96 at nanoseconds by default, which gives wrong values outside + 1677-09-21 to 2262-04-11. With ``use_pyarrow=True`` given to ``execute()``, + pyarrow reads it at microseconds unless ``pyarrow_options`` sets + ``coerce_int96_timestamp_unit``. Polars 2.0 and later read INT96 at the time + unit of the read schema, except in a Map, which the read schema therefore + makes a List of Structs; ``_restore_maps()`` restores it. Polars 1.x reads + INT96 at nanoseconds whatever the schema. + + Args: + schema: The schema of the Parquet files. + + Returns: + No arguments if the files hold no timestamps, ``pyarrow_options`` with + ``use_pyarrow=True``, the ``schema`` argument with Polars 2.0 and later, or + None with Polars 1.x, which cannot read them. + """ + import polars as pl + + if not holds_datetimes(pl.Struct(schema)): + return {} + if self._kwargs.get("use_pyarrow"): + return { + "pyarrow_options": { + "coerce_int96_timestamp_unit": "us", + **(self._kwargs.get("pyarrow_options") or {}), + } + } + if int(pl.__version__.split(".")[0]) < 2: + return None + return { + "schema": {name: to_microseconds(dtype, maps=False) for name, dtype in schema.items()} + } + + @staticmethod + def _int96_read_error() -> NotSupportedError: + """Get the error of a read of timestamps that Polars 1.x cannot do with pyarrow. + + Returns: + The error, which suggests the ways to read the timestamps. + """ + import polars as pl + + return NotSupportedError( + f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at " + "nanoseconds, which gives wrong values outside 1677-09-21 to 2262-04-11. " + "Upgrade to Polars 2.0 or later, or install pyarrow and give no read " + "options to execute()." + ) + + def _restore_maps( + self, df: pl.DataFrame, schema: pl.Schema, options: dict[str, Any] | None + ) -> pl.DataFrame: + """Restore the Map dtypes that the read schema of ``_int96_read_options()`` replaced. + + Args: + df: The DataFrame that was read. + schema: The schema of the Parquet files. + options: The read arguments from ``_int96_read_options()``. + + Returns: + The DataFrame with the Map dtypes, unless ``execute()`` was given its own + ``schema``. + """ + import polars as pl + + if "schema" not in (options or {}) or "schema" in self._kwargs: + return df + casts = [ + pl.col(name).cast(dtype) + for name in df.columns + if name in schema and (dtype := to_microseconds(schema[name])) != df.schema[name] + ] + return df.with_columns(casts) if casts else df + def _as_polars(self) -> pl.DataFrame: """Load query results as a Polars DataFrame. @@ -818,14 +949,16 @@ def _as_polars(self) -> pl.DataFrame: Returns: Polars DataFrame containing the query results. """ - if self.is_unload: - df = self._read_parquet() - if df.is_empty(): - self._metadata = () - else: - self._metadata = self._read_parquet_schema() - else: - df = self._read_csv() + import polars as pl + + if not self.is_unload: + return self._read_csv() + if not self._prepare_parquet_location(): + self._metadata = () + return pl.DataFrame() + schema = self._read_parquet_schema() + df = self._read_parquet(schema) + self._metadata = () if df.is_empty() else to_column_info(schema) return df def as_polars(self) -> pl.DataFrame: @@ -987,24 +1120,28 @@ def _scan_csv_chunks( _logger.exception(f"Failed to read {location}.") raise OperationalError(*e.args) from e - def _iter_parquet_chunks(self) -> Iterator[pl.DataFrame]: + def _iter_parquet_chunks(self, schema: pl.Schema) -> Iterator[pl.DataFrame]: """Create an iterator that reads the Parquet files in chunks. Without read options given to ``execute()``, and with pyarrow installed, the files of the data manifest are read one at a time through PyAthena's filesystem with ``ParquetChunkReader``, so that memory use depends on the chunk size and on the row group being read, not on the size of the result. - Otherwise Polars scans - the unload location with the given options, decoding whole row groups ahead - of the chunks it returns. + Otherwise Polars scans the unload location with the given options, decoding + whole row groups ahead of the chunks it returns. Both read timestamps at + microseconds; Polars 1.x cannot scan them so. ``_prepare_parquet_location()`` must have found the unload location first. + Args: + schema: The schema of the Parquet files, from ``_read_parquet_schema()``. + Returns: An iterator of a Polars DataFrame for each chunk, read lazily. Raises: ProgrammingError: If the unload location is not set. + NotSupportedError: If Polars 1.x would read timestamps at nanoseconds. """ if self._unload_location is None: raise ProgrammingError("unload_location is not available.") @@ -1013,7 +1150,10 @@ def _iter_parquet_chunks(self) -> Iterator[pl.DataFrame]: return self._read_parquet_chunks( S3FileSystem(**self._csv_storage_options), self._data_manifest ) - return self._scan_parquet_chunks(self._unload_location, self.connection) + options = self._int96_read_options(schema) + if options is None: + raise self._int96_read_error() + return self._scan_parquet_chunks(self._unload_location, self.connection, schema, options) def _read_parquet_chunks(self, fs: S3FileSystem, paths: list[str]) -> Iterator[pl.DataFrame]: """Read the files of the data manifest in chunks with ``ParquetChunkReader``. @@ -1044,7 +1184,11 @@ def _read_parquet_chunks(self, fs: S3FileSystem, paths: list[str]) -> Iterator[p raise OperationalError(*e.args) from e def _scan_parquet_chunks( - self, location: str, connection: Connection[Any] + self, + location: str, + connection: Connection[Any], + schema: pl.Schema, + options: dict[str, Any], ) -> Iterator[pl.DataFrame]: """Scan the unload location in chunks with Polars and the read options. @@ -1052,6 +1196,8 @@ def _scan_parquet_chunks( location: The unload location. connection: The connection whose credentials the read uses. They are frozen when the first chunk is read. + schema: The schema of the Parquet files. + options: The read arguments from ``_int96_read_options()``. Yields: Polars DataFrame for each chunk. @@ -1064,9 +1210,13 @@ def _scan_parquet_chunks( try: lazy_df = pl.scan_parquet( location, - **self._read_kwargs(lambda: self._get_parquet_storage_options(connection)), + **self._read_kwargs( + lambda: self._get_parquet_storage_options(connection), + **options, + ), ) - yield from lazy_df.collect_batches(chunk_size=self._chunksize) + for batch in lazy_df.collect_batches(chunk_size=self._chunksize): + yield self._restore_maps(batch, schema, options) except Exception as e: _logger.exception(f"Failed to read {location}.") raise OperationalError(*e.args) from e diff --git a/pyathena/polars/util.py b/pyathena/polars/util.py index 98976a087..f5b749f4e 100644 --- a/pyathena/polars/util.py +++ b/pyathena/polars/util.py @@ -123,3 +123,53 @@ def get_athena_type(dtype: Any) -> tuple[str, int, int]: return ("decimal", 38, 9) # Default precision and scale return ("string", 2147483647, 0) + + +def to_microseconds(dtype: Any, maps: bool = True) -> Any: + """Get a dtype whose nanosecond Datetimes, at any nesting level, are in microseconds. + + Args: + dtype: The dtype. + maps: Whether to keep Map dtypes. If False, a Map is a List of Structs of its + key and value, which Polars 2.0 needs to read INT96 timestamps in a Parquet + map at microseconds. + + Returns: + The dtype with microsecond Datetimes. + """ + import polars as pl + + if isinstance(dtype, pl.Datetime): + return pl.Datetime("us", dtype.time_zone) if dtype.time_unit == "ns" else dtype + if isinstance(dtype, pl.List): + return pl.List(to_microseconds(dtype.inner, maps)) + if isinstance(dtype, pl.Struct): + return pl.Struct({f.name: to_microseconds(f.dtype, maps) for f in dtype.fields}) + map_dtype = getattr(pl, "Map", None) + if map_dtype is not None and isinstance(dtype, map_dtype): + key, value = to_microseconds(dtype.key, maps), to_microseconds(dtype.value, maps) + return map_dtype(key, value) if maps else pl.List(pl.Struct({"key": key, "value": value})) + return dtype + + +def holds_datetimes(dtype: Any) -> bool: + """Whether a dtype is a Datetime or holds one at any nesting level. + + Args: + dtype: The dtype. + + Returns: + True if the dtype is or holds a Datetime. + """ + import polars as pl + + if isinstance(dtype, pl.Datetime): + return True + if isinstance(dtype, pl.List): + return holds_datetimes(dtype.inner) + if isinstance(dtype, pl.Struct): + return any(holds_datetimes(f.dtype) for f in dtype.fields) + map_dtype = getattr(pl, "Map", None) + if map_dtype is not None and isinstance(dtype, map_dtype): + return holds_datetimes(dtype.key) or holds_datetimes(dtype.value) + return False diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index d4774e46b..08ead4b05 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -517,7 +517,7 @@ def test_complex_unload_as_arrow(self, arrow_cursor): pa.field("col_double", pa.float64()), pa.field("col_string", pa.string()), pa.field("col_varchar", pa.string()), - pa.field("col_timestamp", pa.timestamp("ns")), + pa.field("col_timestamp", pa.timestamp("us")), pa.field("col_date", pa.date32()), pa.field("col_binary", pa.binary()), pa.field("col_array", pa.list_(pa.field("array_element", pa.int32()))), @@ -701,7 +701,7 @@ def test_complex_unload_as_polars(self, arrow_cursor): pl.Float64, pl.String, pl.String, - pl.Datetime("ns"), + pl.Datetime("us"), pl.Date, pl.Binary, pl.List(pl.Int32), diff --git a/tests/pyathena/arrow/test_reader.py b/tests/pyathena/arrow/test_reader.py index 7569a4d2b..ab64df9c7 100644 --- a/tests/pyathena/arrow/test_reader.py +++ b/tests/pyathena/arrow/test_reader.py @@ -15,7 +15,7 @@ from fsspec.spec import AbstractBufferedFile from pyathena.arrow.reader import ParquetChunkReader -from tests.pyathena.util import ImmediateExecutor +from tests.pyathena.util import INT96_UNLOAD_FILE, INT96_VALUES, ImmediateExecutor class _RecordingFile(AbstractBufferedFile): @@ -208,6 +208,20 @@ def test_streamed_file_row_groups(self, tmp_path, monkeypatch, cache_type): for (_, end), (_, span_end) in zip(ranges[:2], spans[:2], strict=True) ) + def test_int96_timestamps(self, max_prefetch_size): + """INT96 timestamps that Athena writes are read at microseconds at any level.""" + with ParquetChunkReader(LocalFileSystem(), [str(INT96_UNLOAD_FILE)], 3) as reader: + table = pa.concat_tables(reader) + timestamp = pa.timestamp("us") + assert table.schema.types[1:] == [ + timestamp, + pa.list_(timestamp), + pa.struct([("t", timestamp)]), + pa.map_(pa.string(), timestamp), + ] + assert table.column("ts").to_pylist() == INT96_VALUES + assert table.column("mp").to_pylist() == [[("k", v)] for v in INT96_VALUES] + def test_no_files(self): assert list(ParquetChunkReader(LocalFileSystem(), [], 4)) == [] diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index cd214b6b4..45d578908 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -5,18 +5,35 @@ # # SPDX-License-Identifier: MIT import os +import shutil import sys from datetime import datetime from unittest.mock import MagicMock, patch import pyarrow as pa import pytest +from pyarrow import fs as pyarrow_fs from pyathena.arrow.converter import DefaultArrowTypeConverter from pyathena.arrow.result_set import AthenaArrowResultSet from pyathena.connection import Connection from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig +from tests.pyathena.util import INT96_UNLOAD_FILE, INT96_VALUES + + +def _read_unload_result(tmp_path): + """Read the Parquet file that Athena wrote for UNLOAD through the result set, without AWS.""" + unload = tmp_path / "bucket" / "unload" + unload.mkdir(parents=True) + shutil.copy(INT96_UNLOAD_FILE, unload / "0.parquet") + result_set = AthenaArrowResultSet.__new__(AthenaArrowResultSet) # bypass __init__ + result_set._fs = pyarrow_fs.SubTreeFileSystem(str(tmp_path), pyarrow_fs.LocalFileSystem()) + result_set._unload_location = "s3://bucket/unload/" + with patch.object( + AthenaArrowResultSet, "_read_data_manifest", return_value=["s3://bucket/unload/0.parquet"] + ): + return result_set._read_parquet() class TestAthenaArrowResultSet: @@ -201,3 +218,20 @@ def test_open_csv_closes_file_when_stream_fails(self): with pytest.raises(OSError, match="HEAD failed"): result_set._open_csv("bucket/key.csv") file.close.assert_called_once_with() + + def test_read_parquet_int96(self, tmp_path): + """INT96 timestamps that Athena writes are read at microseconds at any level.""" + table = _read_unload_result(tmp_path).sort_by("id") + timestamp = pa.timestamp("us") + assert table.schema.types[1:] == [ + timestamp, + pa.list_(timestamp), + pa.struct([("t", timestamp)]), + pa.map_(pa.string(), timestamp), + ] + assert table.drop_columns("id").to_pydict() == { + "ts": INT96_VALUES, + "arr": [[v] for v in INT96_VALUES], + "st": [{"t": v} for v in INT96_VALUES], + "mp": [[("k", v)] for v in INT96_VALUES], + } diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 186bb8842..b68e637ce 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -9,6 +9,7 @@ import csv import inspect import io +import shutil import threading from collections import defaultdict from concurrent.futures import ThreadPoolExecutor @@ -34,7 +35,7 @@ _read_csv_with_pyarrow, ) from pyathena.util import RetryConfig, retry_api_call, stop_retries_on -from tests.pyathena.util import ThrottledOperation +from tests.pyathena.util import INT96_UNLOAD_FILE, INT96_VALUES, ThrottledOperation def _typed_values(rows): @@ -554,10 +555,34 @@ def test_read_parquet_filesystem(self, execute_kwargs, path, filesystem_kwargs): assert read_parquet.call_args.kwargs == { "engine": "pyarrow", "use_threads": True, + "coerce_int96_timestamp_unit": "us", **execute_kwargs, **filesystem_kwargs, } + @pytest.mark.parametrize( + ("chunksize", "kwargs"), + [(None, {}), (None, {"dtype_backend": "pyarrow"}), (3, {})], + ids=["whole", "whole-pyarrow-dtypes", "chunks"], + ) + def test_read_parquet_int96(self, tmp_path, chunksize, kwargs): + """INT96 timestamps that Athena writes are read at microseconds at any level.""" + df = _read_int96_unload_result(tmp_path, chunksize, **kwargs).sort_values("id") + expected_dtype = "timestamp[us][pyarrow]" if kwargs else "datetime64[us]" + assert str(df["ts"].dtype) == expected_dtype + assert [None if pd.isna(v) else v for v in df["ts"]] == INT96_VALUES + # The numpy backend holds arrays as numpy arrays, and the pyarrow one as lists. + arrays = [v.tolist() if isinstance(v, np.ndarray) else v for v in df["arr"]] + assert arrays == [[v] for v in INT96_VALUES] + assert df["st"].tolist() == [{"t": v} for v in INT96_VALUES] + assert df["mp"].tolist() == [[("k", v)] for v in INT96_VALUES] + + def test_read_parquet_int96_unit_option(self, tmp_path): + """coerce_int96_timestamp_unit given to execute() replaces microseconds.""" + df = _read_int96_unload_result(tmp_path, None, coerce_int96_timestamp_unit="ms") + assert df["ts"].dtype == "datetime64[ms]" + assert [None if pd.isna(v) else v for v in df.sort_values("id")["ts"]] == INT96_VALUES + def test_iter_parquet_chunks_null_in_later_file(self, tmp_path): """Top-level integer and boolean columns with NULL anywhere convert alike in every chunk.""" chunks, _ = _read_unload_result( @@ -997,6 +1022,28 @@ def _unload_result_set(tmp_path, tables, chunksize, write_statistics=True, **kwa return result_set, paths +def _read_int96_unload_result(tmp_path, chunksize, **kwargs): + """Read the Parquet file that Athena wrote for UNLOAD through the result set, without AWS. + + Returns: + The whole result, or its chunks joined. + """ + result_set, _ = _unload_result_set(tmp_path, [], chunksize, **kwargs) + path = tmp_path / "0.parquet" + shutil.copy(INT96_UNLOAD_FILE, path) + with ( + patch.object(AthenaPandasResultSet, "_read_data_manifest", return_value=[str(path)]), + patch.object( + AthenaPandasResultSet, "is_unload", new_callable=PropertyMock, return_value=True + ), + ): + if chunksize is None: + # Without a filesystem, pandas reads the local unload location. + result_set._kwargs = {"filesystem": None, **kwargs} + return result_set._read_parquet("pyarrow") + return pd.concat(list(result_set._as_pandas()), ignore_index=True) + + def _read_unload_result(tmp_path, tables, chunksize, write_statistics=True, **kwargs): """Read Parquet files in chunks through the result set, without AWS. diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index c54fd4496..bc36aa485 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -6,6 +6,8 @@ # SPDX-License-Identifier: MIT import inspect +import logging +import shutil import threading from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor @@ -19,11 +21,12 @@ from fsspec.implementations.local import LocalFileSystem from pyathena.converter import _to_date, _to_default -from pyathena.error import OperationalError +from pyathena.error import NotSupportedError, OperationalError from pyathena.model import AthenaQueryExecution from pyathena.polars.converter import DefaultPolarsTypeConverter from pyathena.polars.result_set import AthenaPolarsResultSet, PolarsDataFrameIterator from pyathena.util import RetryConfig +from tests.pyathena.util import INT96_UNLOAD_FILE, INT96_VALUES _ROWS_BEFORE_FAILURE = 300_000 # A timestamp column of a CSV result file with 12 fractional digits and NULL, as text and @@ -32,6 +35,22 @@ _TRUNCATED_TIMESTAMPS = pl.Series( [datetime(2020, 1, 2, 3, 4, 5, 123456), None], dtype=pl.Datetime("us") ) +# Polars 2.0 and later read Parquet maps as Map, and INT96 at the unit of the schema. +_POLARS_2 = int(pl.__version__.split(".")[0]) >= 2 +_MAP = ( + pl.Map(pl.String, pl.Datetime("us")) + if _POLARS_2 + else pl.List(pl.Struct({"key": pl.String, "value": pl.Datetime("us")})) +) +# The ways to read an UNLOAD result: whole, whole by Polars with pyarrow, in chunks +# with ParquetChunkReader, and in chunks that Polars scans with the read options. +_UNLOAD_READS = [ + pytest.param(None, {}, id="whole"), + pytest.param(None, {"use_pyarrow": True}, id="whole-use-pyarrow"), + pytest.param(3, {}, id="chunks"), + pytest.param(3, {"low_memory": True}, id="scanned-chunks"), +] +_UPGRADE = "Upgrade to Polars 2.0 or later" def _chunked_result_set() -> AthenaPolarsResultSet: @@ -94,6 +113,38 @@ def _fetching_result_set(chunks, names, arraysize, converters=None) -> AthenaPol ) +def _read_unload_result(tmp_path, chunksize, kwargs, file=INT96_UNLOAD_FILE): + """Read a Parquet file of an UNLOAD result through the result set, without AWS. + + Returns: + The DataFrames that were read: the whole result, or its chunks. + """ + path = tmp_path / "0.parquet" + shutil.copy(file, path) + result_set = _chunked_result_set() + result_set._chunksize = chunksize + result_set._kwargs = dict(kwargs) + result_set._unload_location = f"{tmp_path}/" + with ( + patch.object(AthenaPolarsResultSet, "_read_data_manifest", return_value=[str(path)]), + patch.object( + AthenaPolarsResultSet, "is_unload", new_callable=PropertyMock, return_value=True + ), + patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), + patch.object( + AthenaPolarsResultSet, + "_csv_storage_options", + new_callable=PropertyMock, + return_value={}, + ), + patch("pyathena.polars.result_set.S3FileSystem", lambda **_: LocalFileSystem()), + ): + if chunksize is None: + return [result_set._as_polars()] + assert result_set._prepare_parquet_location() + return list(result_set._iter_parquet_chunks(result_set._read_parquet_schema())) + + class TestAthenaPolarsResultSet: def test_fetch_across_chunks(self): """The fetch methods continue across chunks, reading a chunk when its rows are needed.""" @@ -388,7 +439,7 @@ def test_iter_parquet_chunks_raises_when_read_fails_partway(self, tmp_path, read pytest.raises(OperationalError) as error, ): # extend() keeps the chunks read before the failure. - chunks.extend(result_set._iter_parquet_chunks()) + chunks.extend(result_set._iter_parquet_chunks(pl.Schema({"a": pl.Int64}))) # The error comes from the corrupt file, after the valid one was read. assert not isinstance(error.value.__cause__, AttributeError) assert scan_parquet.called is bool(read_options) @@ -504,7 +555,9 @@ def test_storage_options_replace_defaults(self, reader, function): patch(f"polars.{function}") as read, patch("pyathena.polars.result_set.to_column_info"), ): - result = getattr(result_set, reader)() + # The Parquet readers take the schema of the files. + args = (pl.Schema(),) if reader in ("_read_parquet", "_iter_parquet_chunks") else () + result = getattr(result_set, reader)(*args) if not isinstance(result, (pl.DataFrame, tuple)): list(result) assert read.call_args.kwargs["storage_options"] == {"anon": True} @@ -645,6 +698,95 @@ def test_read_csv_ignore_errors(self, tmp_path, reader): pl.Series([datetime(2020, 1, 2, 3, 4, 5, 123456), None], dtype=pl.Datetime("us")), check_dtypes=True, ) + @pytest.mark.parametrize(("chunksize", "kwargs"), _UNLOAD_READS) + def test_read_parquet_int96(self, tmp_path, caplog, chunksize, kwargs): + """INT96 timestamps that Athena writes are read at microseconds at any level. + + Polars 1.x reads them with pyarrow and logs a warning, as it reads INT96 only + at nanoseconds; it cannot scan them. + """ + scanned = chunksize is not None and kwargs + if not _POLARS_2 and scanned: + with pytest.raises(NotSupportedError, match=_UPGRADE): + _read_unload_result(tmp_path, chunksize, kwargs) + return + with caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"): + df = pl.concat(_read_unload_result(tmp_path, chunksize, kwargs)).sort("id") + timestamp = pl.Datetime("us") + assert df.drop("id").schema == pl.Schema( + { + "ts": timestamp, + "arr": pl.List(timestamp), + "st": pl.Struct({"t": timestamp}), + "mp": _MAP, + } + ) + assert df["ts"].to_list() == INT96_VALUES + assert df["arr"].to_list() == [[v] for v in INT96_VALUES] + assert df["st"].to_list() == [{"t": v} for v in INT96_VALUES] + assert df["mp"].to_list() == ( + [{"k": v} for v in INT96_VALUES] + if _POLARS_2 + else [[{"key": "k", "value": v}] for v in INT96_VALUES] + ) + assert (_UPGRADE in caplog.text) is (not _POLARS_2 and chunksize is None and not kwargs) + + @pytest.mark.skipif(_POLARS_2, reason="Polars 2.0 reads INT96 without pyarrow.") + @pytest.mark.parametrize( + ("chunksize", "kwargs", "match"), + [ + (None, {}, "install pyarrow"), + (3, {}, "install pyarrow"), + (None, {"low_memory": True}, "give no read options"), + ], + ids=["whole", "chunks", "whole-read-options"], + ) + def test_read_parquet_int96_polars_1(self, tmp_path, chunksize, kwargs, match): + """Polars 1.x reads INT96 at microseconds only with pyarrow and no read options.""" + with ( + patch("pyathena.polars.result_set.importlib.util.find_spec", return_value=None), + pytest.raises(NotSupportedError, match=match), + ): + _read_unload_result(tmp_path, chunksize, kwargs) + + def test_read_parquet_int96_pyarrow_options(self, tmp_path): + """coerce_int96_timestamp_unit in pyarrow_options replaces microseconds.""" + (df,) = _read_unload_result( + tmp_path, + None, + {"use_pyarrow": True, "pyarrow_options": {"coerce_int96_timestamp_unit": "ms"}}, + ) + assert df.schema["ts"] == pl.Datetime("ms") + assert df.sort("id")["ts"].to_list() == INT96_VALUES + + @pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x reads INT96 at nanoseconds.") + @pytest.mark.parametrize("chunksize", [None, 3], ids=["whole", "scanned-chunks"]) + def test_read_parquet_int96_schema_option(self, tmp_path, chunksize): + """A schema given to execute() replaces the one that reads INT96 at microseconds.""" + timestamp = pl.Datetime("us") + schema = { + "id": pl.Int32, + "ts": pl.Datetime("ms"), + "arr": pl.List(timestamp), + "st": pl.Struct({"t": timestamp}), + "mp": pl.List(pl.Struct({"key": pl.String, "value": timestamp})), + } + kwargs = {"schema": schema} if chunksize is None else {"schema": schema, "low_memory": True} + df = pl.concat(_read_unload_result(tmp_path, chunksize, kwargs)).sort("id") + assert df.schema == pl.Schema(schema) + assert df["ts"].to_list() == INT96_VALUES + assert df["mp"].to_list() == [[{"key": "k", "value": v}] for v in INT96_VALUES] + + def test_read_parquet_without_timestamps(self, tmp_path, caplog): + """A result without timestamps is read by Polars with read options and no warning.""" + file = tmp_path / "source.parquet" + pl.DataFrame({"i": [1, None]}).write_parquet(file) + unload = tmp_path / "unload" + unload.mkdir() + with caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"): + (df,) = _read_unload_result(unload, None, {"low_memory": True}, file=file) + assert df["i"].to_list() == [1, None] + assert not caplog.text class TestPolarsDataFrameIterator: diff --git a/tests/pyathena/polars/test_util.py b/tests/pyathena/polars/test_util.py index 800387ccc..601b2c394 100644 --- a/tests/pyathena/polars/test_util.py +++ b/tests/pyathena/polars/test_util.py @@ -8,7 +8,10 @@ import polars as pl import pytest -from pyathena.polars.util import get_athena_type +from pyathena.polars.util import get_athena_type, holds_datetimes, to_microseconds + +_MAP = getattr(pl, "Map", None) +_requires_map = pytest.mark.skipif(_MAP is None, reason="Polars 1.x has no Map type") @pytest.mark.parametrize( @@ -24,8 +27,57 @@ def test_get_athena_type_nested(dtype, expected): assert get_athena_type(dtype) == expected -@pytest.mark.skipif(not hasattr(pl, "Map"), reason="Polars 1.x has no Map type") +@_requires_map def test_get_athena_type_map(): assert get_athena_type(pl.Map(pl.String, pl.Int32)) == ("map", 0, 0) assert get_athena_type(pl.Map(pl.String, pl.List(pl.Int32))) == ("map", 0, 0) assert get_athena_type(pl.List(pl.Map(pl.String, pl.Int32))) == ("array", 0, 0) + + +_NS, _US = pl.Datetime("ns"), pl.Datetime("us") + + +@pytest.mark.parametrize( + ("dtype", "expected"), + [ + (_NS, _US), + (pl.Datetime("ns", "UTC"), pl.Datetime("us", "UTC")), + (pl.Datetime("ms"), pl.Datetime("ms")), + (pl.Int64, pl.Int64), + (pl.List(_NS), pl.List(_US)), + (pl.Struct({"a": _NS, "b": pl.List(_NS)}), pl.Struct({"a": _US, "b": pl.List(_US)})), + ], +) +def test_to_microseconds(dtype, expected): + assert to_microseconds(dtype) == expected + assert to_microseconds(dtype, maps=False) == expected + + +@_requires_map +def test_to_microseconds_map(): + """A Map is kept, or with maps=False becomes a List of Structs of its key and value.""" + dtype = pl.List(pl.Map(_NS, pl.Struct({"t": _NS}))) + assert to_microseconds(dtype) == pl.List(pl.Map(_US, pl.Struct({"t": _US}))) + assert to_microseconds(dtype, maps=False) == pl.List( + pl.List(pl.Struct({"key": _US, "value": pl.Struct({"t": _US})})) + ) + + +@pytest.mark.parametrize( + ("dtype", "expected"), + [ + (_US, True), + (pl.Date, False), + (pl.List(pl.Struct({"a": pl.Int8, "b": pl.List(_NS)})), True), + (pl.Struct({"a": pl.Int8, "b": pl.List(pl.String)}), False), + ], +) +def test_holds_datetimes(dtype, expected): + assert holds_datetimes(dtype) is expected + + +@_requires_map +def test_holds_datetimes_map(): + assert holds_datetimes(pl.Map(_US, pl.Int8)) + assert holds_datetimes(pl.Map(pl.String, pl.List(_NS))) + assert not holds_datetimes(pl.Map(pl.String, pl.Int8)) diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 95d4a053c..0ff742ce8 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -665,3 +665,17 @@ def stub_multipart_copy(stubber, fail_list=False, fail_part=False, fail_annotati ) return stubber.add_response("put_object_annotation", {}, put) + + +# A Parquet file that Athena's UNLOAD wrote, with INT96 timestamps and no Arrow schema: +# UNLOAD (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, ...), ...) AS t (id, ts)) +INT96_UNLOAD_FILE = Path(__file__).parents[1] / "resources" / "unload" / "int96_timestamps.parquet" +# The timestamps of its rows, by id: years 1 to 9999 that nanoseconds cannot hold, and +# NULL. Each row's arr, st, and mp hold its timestamp. +INT96_VALUES = [ + datetime(1, 1, 1), + datetime(1600, 1, 2, 3, 4, 5, 123000), + datetime(9999, 12, 31, 23, 59, 59, 999000), + None, +] diff --git a/tests/resources/unload/int96_timestamps.parquet b/tests/resources/unload/int96_timestamps.parquet new file mode 100644 index 0000000000000000000000000000000000000000..daa6737b3fa63ea3fabc0d4d92646b7c3b05aa1e GIT binary patch literal 1254 zcmb_cO=uHA6rP!Er`s5?VrSfipldWs3Q29+rc^D8Kh+)Hn+K6X@gg3ihk_T|t0EpmLe>v;2n!wHkvD4RAQUuwxW=jXKp z-47z};*&}>+q?2*yg0P__3lR3X-b#x)naHla{J`bumECRY&ay9wk0i2<*!-@)oSje z=q~Xk*!cz@ZD~3bNzZC)&uR~^-N)$H#nk$~UE<{IRD-kC*dlO#%$q{w(A-G>o;qQa z^%;!(Wop4R-&1s_m^j>4ckW=*x$Tui7x3DNkt2I|NyUFVN7N6avA}mq<%#O#wC!0e ziWo@}g*b&|!dTg_34>`E392YZl!)c~AeL|{Lnv*F~!{(U8 zWqyfI>+-N7!Du5z!mx6b+VY0fIE*g{f|C$W0WVR&2UUT8iupIJ|HOJFK1h(7hEyq( zB8jbmNdjajlr5w-B#HtrF#nGA&k_9-*{OnHO@Sde5U?f8H?Y3mroWc{4U4IQn0S?@ zQ=vnmLV>@;`X<$vXk5WFaD#BM@2sMwSl!+&I^N0p+xG1n;;WR)gU>s_QKj>S?*z8r z6F8pTbE{mkmB{#&b4WQ}vHa_^c98M>j2J0x&i1QLxuo^>^ks9IzJbAkY-Xz1Z{_TK XK}+Rl9k-yBj#=Gk0lv@}{NerpVPpId literal 0 HcmV?d00001 From 5d1284d471202a1642f47ee4286e320c31ef347c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 12:30:19 +0900 Subject: [PATCH 02/12] Keep the chunks of the pyarrow read of Polars 1.x pl.from_arrow() rechunks by default, which added a copy of the result to the peak memory of the fallback; Polars does not rechunk its own reads. The Polars 1.x test now also covers read options with pyarrow installed. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 3 ++- tests/pyathena/polars/test_result_set.py | 19 +++++++++---------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 51bd34cf2..82e1f8d34 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -831,7 +831,8 @@ def _read_parquet_with_pyarrow(self) -> pl.DataFrame: filesystem=S3FileSystem(**self._csv_storage_options), coerce_int96_timestamp_unit="us", ).read(use_threads=True) - return cast("pl.DataFrame", pl.from_arrow(table)) + # 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 diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index bc36aa485..ea3ee9dae 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -733,19 +733,18 @@ def test_read_parquet_int96(self, tmp_path, caplog, chunksize, kwargs): @pytest.mark.skipif(_POLARS_2, reason="Polars 2.0 reads INT96 without pyarrow.") @pytest.mark.parametrize( - ("chunksize", "kwargs", "match"), - [ - (None, {}, "install pyarrow"), - (3, {}, "install pyarrow"), - (None, {"low_memory": True}, "give no read options"), - ], - ids=["whole", "chunks", "whole-read-options"], + ("chunksize", "kwargs", "pyarrow"), + [(None, {}, False), (3, {}, False), (None, {"low_memory": True}, True)], + ids=["whole-without-pyarrow", "chunks-without-pyarrow", "whole-read-options"], ) - def test_read_parquet_int96_polars_1(self, tmp_path, chunksize, kwargs, match): + def test_read_parquet_int96_polars_1(self, tmp_path, chunksize, kwargs, pyarrow): """Polars 1.x reads INT96 at microseconds only with pyarrow and no read options.""" with ( - patch("pyathena.polars.result_set.importlib.util.find_spec", return_value=None), - pytest.raises(NotSupportedError, match=match), + patch( + "pyathena.polars.result_set.importlib.util.find_spec", + return_value=MagicMock() if pyarrow else None, + ), + pytest.raises(NotSupportedError, match=_UPGRADE), ): _read_unload_result(tmp_path, chunksize, kwargs) From 7633e22ad445d89ed2d17f09235bb9b96b207e4d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 12:33:57 +0900 Subject: [PATCH 03/12] Expect the UNLOAD temporal types to pass The UNLOAD timestamps of ArrowCursor and PolarsCursor now match the contract; PandasCursor's UNLOAD dates remain #1219. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_cursor.py | 6 +----- tests/pyathena/pandas/test_cursor.py | 2 +- tests/pyathena/polars/test_cursor.py | 6 +----- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index 08ead4b05..a650f959e 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -179,11 +179,7 @@ def test_invalid_arraysize(self, arrow_cursor): ("arrow_cursor", "unload"), [ pytest.param({"cursor_kwargs": {"unload": False}}, False), - pytest.param( - {"cursor_kwargs": {"unload": True}}, - True, - marks=pytest.mark.xfail(strict=True, reason="#1214"), - ), + ({"cursor_kwargs": {"unload": True}}, True), ], indirect=["arrow_cursor"], ids=["csv", "unload"], diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index e4fd7fc9f..bcb762d3a 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -532,7 +532,7 @@ def test_invalid_arraysize(self, pandas_cursor): pytest.param( {"cursor_kwargs": {"unload": True}}, True, - marks=pytest.mark.xfail(strict=True, reason="#1214, #1219"), + marks=pytest.mark.xfail(strict=True, reason="#1219"), ), ], indirect=["pandas_cursor"], diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index adf7bea6c..ebe3e10ee 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -153,11 +153,7 @@ def test_many_as_polars(self, polars_cursor): ("polars_cursor", "unload"), [ ({"cursor_kwargs": {"unload": False}}, False), - pytest.param( - {"cursor_kwargs": {"unload": True}}, - True, - marks=pytest.mark.xfail(strict=True, reason="#1214"), - ), + ({"cursor_kwargs": {"unload": True}}, True), ], indirect=["polars_cursor"], ids=["csv", "unload"], From b9287d7e8c09f28e3c69e88ca15bcd31ea0e0a6c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 12:38:45 +0900 Subject: [PATCH 04/12] List the Polars 1.x UNLOAD reads that use PyAthena's filesystem Co-Authored-By: Claude Opus 5.5 --- docs/polars.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/polars.md b/docs/polars.md index 022f6ab27..4cae3ec26 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -12,10 +12,10 @@ 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, except for UNLOAD results with timestamps read with Polars 1.x; see {ref}`polars-unload-options`. -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. +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 +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. Column types for a CSV result are given as `schema_overrides`, which replaces PyAthena's column types. @@ -403,7 +403,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, 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 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. 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`. From 7871a6cb47e1a58e72dda3b4c5fcfc8219a1e415 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 12:57:56 +0900 Subject: [PATCH 05/12] Read Polars 1.x UNLOAD timestamps natively with a warning Polars 1.x raised NotSupportedError for a result with timestamps read with Polars read options, such as storage_options, which read in-range timestamps correctly before. It now reads them natively at nanoseconds, casts them to microseconds, and warns that values outside 1677-09-21 to 2262-04-11 are wrong; the warning follows the columns actually read. Only nanosecond Datetimes select the INT96 handling, and only maps that hold them are read as lists of structs. Columns are cast by name with Series, as pl.col() reads names such as "*" as selectors. Co-Authored-By: Claude Opus 5.5 --- docs/polars.md | 11 +- pyathena/polars/result_set.py | 141 ++++++++++++----------- pyathena/polars/util.py | 33 +----- tests/pyathena/polars/test_result_set.py | 88 +++++++++++--- tests/pyathena/polars/test_util.py | 25 +--- 5 files changed, 162 insertions(+), 136 deletions(-) diff --git a/docs/polars.md b/docs/polars.md index 4cae3ec26..4b4506117 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -11,7 +11,7 @@ This cursor downloads the CSV file after executing the query and loads it into a 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, except for UNLOAD results with timestamps read with Polars 1.x; see {ref}`polars-unload-options`. +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 @@ -225,9 +225,12 @@ 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. -With Polars 1.x, a result with timestamps is read with pyarrow, through PyAthena's S3FileSystem, and a warning is logged. -This needs pyarrow, and no Polars read options given to `execute()`; otherwise the read raises `NotSupportedError`. -With `use_pyarrow=True` given to `execute()`, a result read without the chunksize option is read by Polars with pyarrow, with `coerce_int96_timestamp_unit="us"` added to `pyarrow_options` unless they set it. +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. +- With `use_pyarrow=True` given to `execute()`, a result read without the chunksize option is read by Polars with pyarrow, with `coerce_int96_timestamp_unit="us"` added to `pyarrow_options` unless they set it. +- 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. 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 82e1f8d34..a269dcbf0 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -33,11 +33,11 @@ from pyathena import OperationalError from pyathena.arrow.reader import ParquetChunkReader from pyathena.converter import Converter, _to_default -from pyathena.error import NotSupportedError, ProgrammingError +from pyathena.error import ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.polars.converter import _to_datetimes -from pyathena.polars.util import holds_datetimes, to_column_info, to_microseconds +from pyathena.polars.util import to_column_info, to_microseconds from pyathena.result_set import AthenaResultSet from pyathena.util import RetryConfig, override @@ -768,9 +768,10 @@ 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 without ``use_pyarrow=True``, and a result with timestamps - is read with pyarrow instead, which needs pyarrow and no read options given to - ``execute()``. + reads them so only with ``use_pyarrow=True``. Without read options given to + ``execute()``, and with pyarrow installed, a result with timestamps is + therefore read with pyarrow instead; otherwise Polars 1.x reads them at + nanoseconds, and ``_output_casts()`` casts them to microseconds. ``_prepare_parquet_location()`` must have found the unload location first. @@ -781,7 +782,6 @@ def _read_parquet(self, schema: pl.Schema) -> pl.DataFrame: Polars DataFrame containing the Parquet data. Raises: - NotSupportedError: If Polars 1.x would read timestamps at nanoseconds. OperationalError: If reading the Parquet files fails. """ import polars as pl @@ -790,28 +790,30 @@ 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: - if self._kwargs or importlib.util.find_spec("pyarrow") is None: - raise self._int96_read_error() + if options is None and not self._kwargs and importlib.util.find_spec("pyarrow"): _logger.warning( f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at " "nanoseconds, so they are read with pyarrow instead. Upgrade to " "Polars 2.0 or later to read them with Polars." ) - df = self._read_parquet_with_pyarrow() - else: - kwargs = self._read_kwargs( - lambda: self._get_parquet_storage_options(self.connection), **options - ) - if "pyarrow_options" in options: - # The options given to execute() are in it. - kwargs["pyarrow_options"] = options["pyarrow_options"] - try: - df = pl.read_parquet(self._unload_location, **kwargs) - except Exception as e: - _logger.exception(f"Failed to read {self._unload_location}.") - raise OperationalError(*e.args) from e - return self._restore_maps(df, schema, options) + return self._read_parquet_with_pyarrow() + kwargs = self._read_kwargs( + lambda: self._get_parquet_storage_options(self.connection), **(options or {}) + ) + if options and "pyarrow_options" in options: + # The options given to execute() are in it. + kwargs["pyarrow_options"] = options["pyarrow_options"] + try: + df = pl.read_parquet(self._unload_location, **kwargs) + casts = self._output_casts(df, schema, options) + if casts: + df = df.with_columns(casts) + except Exception as e: + _logger.exception(f"Failed to read {self._unload_location}.") + raise OperationalError(*e.args) from e + if casts and options is None: + 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. @@ -873,20 +875,20 @@ def _int96_read_options(self, schema: pl.Schema) -> dict[str, Any] | None: pyarrow reads it at microseconds unless ``pyarrow_options`` sets ``coerce_int96_timestamp_unit``. Polars 2.0 and later read INT96 at the time unit of the read schema, except in a Map, which the read schema therefore - makes a List of Structs; ``_restore_maps()`` restores it. Polars 1.x reads + makes a List of Structs; ``_output_casts()`` restores it. Polars 1.x reads INT96 at nanoseconds whatever the schema. Args: schema: The schema of the Parquet files. Returns: - No arguments if the files hold no timestamps, ``pyarrow_options`` with - ``use_pyarrow=True``, the ``schema`` argument with Polars 2.0 and later, or - None with Polars 1.x, which cannot read them. + No arguments if the files hold no nanosecond timestamps, ``pyarrow_options`` + with ``use_pyarrow=True``, the ``schema`` argument with Polars 2.0 and + later, or None with Polars 1.x, which cannot read them. """ import polars as pl - if not holds_datetimes(pl.Struct(schema)): + if all(to_microseconds(dtype) == dtype for dtype in schema.values()): return {} if self._kwargs.get("use_pyarrow"): return { @@ -901,26 +903,14 @@ def _int96_read_options(self, schema: pl.Schema) -> dict[str, Any] | None: "schema": {name: to_microseconds(dtype, maps=False) for name, dtype in schema.items()} } - @staticmethod - def _int96_read_error() -> NotSupportedError: - """Get the error of a read of timestamps that Polars 1.x cannot do with pyarrow. - - Returns: - The error, which suggests the ways to read the timestamps. - """ - import polars as pl - - return NotSupportedError( - f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at " - "nanoseconds, which gives wrong values outside 1677-09-21 to 2262-04-11. " - "Upgrade to Polars 2.0 or later, or install pyarrow and give no read " - "options to execute()." - ) - - def _restore_maps( + def _output_casts( self, df: pl.DataFrame, schema: pl.Schema, options: dict[str, Any] | None - ) -> pl.DataFrame: - """Restore the Map dtypes that the read schema of ``_int96_read_options()`` replaced. + ) -> list[pl.Series]: + """Get the columns of a read that need casting to their output dtypes. + + Polars 1.x reads timestamps at nanoseconds, which are cast to microseconds. + The Maps that the read schema of ``_int96_read_options()`` replaced are + restored. Nothing is cast if ``execute()`` was given its own ``schema``. Args: df: The DataFrame that was read. @@ -928,19 +918,34 @@ def _restore_maps( options: The read arguments from ``_int96_read_options()``. Returns: - The DataFrame with the Map dtypes, unless ``execute()`` was given its own - ``schema``. + The cast columns, keeping their names. """ + if "schema" in self._kwargs: + return [] + if options is None: + dtypes = {name: to_microseconds(dtype) for name, dtype in df.schema.items()} + elif "schema" in options: + dtypes = {name: to_microseconds(schema[name]) for name in df.columns if name in schema} + else: + return [] + # Series, not pl.col(), so that a name such as "*" is not a selector. + return [ + df.get_column(name).cast(dtype) + for name, dtype in dtypes.items() + if dtype != df.schema[name] + ] + + @staticmethod + def _warn_nanosecond_read() -> None: + """Warn that Polars 1.x read timestamps at nanoseconds.""" import polars as pl - if "schema" not in (options or {}) or "schema" in self._kwargs: - return df - casts = [ - pl.col(name).cast(dtype) - for name in df.columns - if name in schema and (dtype := to_microseconds(schema[name])) != df.schema[name] - ] - return df.with_columns(casts) if casts else df + _logger.warning( + f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at " + "nanoseconds, 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." + ) def _as_polars(self) -> pl.DataFrame: """Load query results as a Polars DataFrame. @@ -1130,7 +1135,8 @@ def _iter_parquet_chunks(self, schema: pl.Schema) -> Iterator[pl.DataFrame]: chunk size and on the row group being read, not on the size of the result. Otherwise Polars scans the unload location with the given options, decoding whole row groups ahead of the chunks it returns. Both read timestamps at - microseconds; Polars 1.x cannot scan them so. + microseconds, except that Polars 1.x scans them at nanoseconds and casts them, + as ``_read_parquet()`` does. ``_prepare_parquet_location()`` must have found the unload location first. @@ -1142,7 +1148,6 @@ def _iter_parquet_chunks(self, schema: pl.Schema) -> Iterator[pl.DataFrame]: Raises: ProgrammingError: If the unload location is not set. - NotSupportedError: If Polars 1.x would read timestamps at nanoseconds. """ if self._unload_location is None: raise ProgrammingError("unload_location is not available.") @@ -1151,10 +1156,9 @@ def _iter_parquet_chunks(self, schema: pl.Schema) -> Iterator[pl.DataFrame]: return self._read_parquet_chunks( S3FileSystem(**self._csv_storage_options), self._data_manifest ) - options = self._int96_read_options(schema) - if options is None: - raise self._int96_read_error() - return self._scan_parquet_chunks(self._unload_location, self.connection, schema, options) + return self._scan_parquet_chunks( + self._unload_location, self.connection, schema, self._int96_read_options(schema) + ) def _read_parquet_chunks(self, fs: S3FileSystem, paths: list[str]) -> Iterator[pl.DataFrame]: """Read the files of the data manifest in chunks with ``ParquetChunkReader``. @@ -1189,7 +1193,7 @@ def _scan_parquet_chunks( location: str, connection: Connection[Any], schema: pl.Schema, - options: dict[str, Any], + options: dict[str, Any] | None, ) -> Iterator[pl.DataFrame]: """Scan the unload location in chunks with Polars and the read options. @@ -1208,16 +1212,21 @@ def _scan_parquet_chunks( """ import polars as pl + warned = False try: lazy_df = pl.scan_parquet( location, **self._read_kwargs( lambda: self._get_parquet_storage_options(connection), - **options, + **(options or {}), ), ) for batch in lazy_df.collect_batches(chunk_size=self._chunksize): - yield self._restore_maps(batch, schema, options) + casts = self._output_casts(batch, schema, options) + if casts and options is None and not warned: + self._warn_nanosecond_read() + warned = True + yield batch.with_columns(casts) if casts else batch except Exception as e: _logger.exception(f"Failed to read {location}.") raise OperationalError(*e.args) from e diff --git a/pyathena/polars/util.py b/pyathena/polars/util.py index f5b749f4e..71414c45e 100644 --- a/pyathena/polars/util.py +++ b/pyathena/polars/util.py @@ -130,9 +130,9 @@ def to_microseconds(dtype: Any, maps: bool = True) -> Any: Args: dtype: The dtype. - maps: Whether to keep Map dtypes. If False, a Map is a List of Structs of its - key and value, which Polars 2.0 needs to read INT96 timestamps in a Parquet - map at microseconds. + maps: Whether to keep Map dtypes. If False, a Map that holds nanosecond + Datetimes is a List of Structs of its key and value, which Polars 2.0 needs + to read INT96 timestamps in a Parquet map at microseconds. Returns: The dtype with microsecond Datetimes. @@ -148,28 +148,7 @@ def to_microseconds(dtype: Any, maps: bool = True) -> Any: map_dtype = getattr(pl, "Map", None) if map_dtype is not None and isinstance(dtype, map_dtype): key, value = to_microseconds(dtype.key, maps), to_microseconds(dtype.value, maps) - return map_dtype(key, value) if maps else pl.List(pl.Struct({"key": key, "value": value})) + if maps or (key, value) == (dtype.key, dtype.value): + return map_dtype(key, value) + return pl.List(pl.Struct({"key": key, "value": value})) return dtype - - -def holds_datetimes(dtype: Any) -> bool: - """Whether a dtype is a Datetime or holds one at any nesting level. - - Args: - dtype: The dtype. - - Returns: - True if the dtype is or holds a Datetime. - """ - import polars as pl - - if isinstance(dtype, pl.Datetime): - return True - if isinstance(dtype, pl.List): - return holds_datetimes(dtype.inner) - if isinstance(dtype, pl.Struct): - return any(holds_datetimes(f.dtype) for f in dtype.fields) - map_dtype = getattr(pl, "Map", None) - if map_dtype is not None and isinstance(dtype, map_dtype): - return holds_datetimes(dtype.key) or holds_datetimes(dtype.value) - return False diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index ea3ee9dae..324448d43 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -17,11 +17,13 @@ from unittest.mock import MagicMock, PropertyMock, patch import polars as pl +import pyarrow as pa +import pyarrow.parquet as pq import pytest from fsspec.implementations.local import LocalFileSystem from pyathena.converter import _to_date, _to_default -from pyathena.error import NotSupportedError, OperationalError +from pyathena.error import OperationalError from pyathena.model import AthenaQueryExecution from pyathena.polars.converter import DefaultPolarsTypeConverter from pyathena.polars.result_set import AthenaPolarsResultSet, PolarsDataFrameIterator @@ -50,7 +52,9 @@ pytest.param(3, {}, id="chunks"), pytest.param(3, {"low_memory": True}, id="scanned-chunks"), ] -_UPGRADE = "Upgrade to Polars 2.0 or later" +# The warnings of Polars 1.x reads of timestamps with pyarrow and at nanoseconds. +_PYARROW_READ = "so they are read with pyarrow instead" +_NANOSECOND_READ = "so those outside 1677-09-21 to 2262-04-11 are wrong" def _chunked_result_set() -> AthenaPolarsResultSet: @@ -702,14 +706,10 @@ def test_read_csv_ignore_errors(self, tmp_path, reader): def test_read_parquet_int96(self, tmp_path, caplog, chunksize, kwargs): """INT96 timestamps that Athena writes are read at microseconds at any level. - Polars 1.x reads them with pyarrow and logs a warning, as it reads INT96 only - at nanoseconds; it cannot scan them. + Polars 1.x reads INT96 only at nanoseconds: it reads a whole result without + read options with pyarrow, and scans with them at nanoseconds, giving wrong + values outside 1677-09-21 to 2262-04-11; both log a warning. """ - scanned = chunksize is not None and kwargs - if not _POLARS_2 and scanned: - with pytest.raises(NotSupportedError, match=_UPGRADE): - _read_unload_result(tmp_path, chunksize, kwargs) - return with caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"): df = pl.concat(_read_unload_result(tmp_path, chunksize, kwargs)).sort("id") timestamp = pl.Datetime("us") @@ -721,6 +721,13 @@ def test_read_parquet_int96(self, tmp_path, caplog, chunksize, kwargs): "mp": _MAP, } ) + scanned_by_polars_1 = not _POLARS_2 and chunksize is not None and bool(kwargs) + assert (_NANOSECOND_READ in caplog.text) is scanned_by_polars_1 + assert (_PYARROW_READ in caplog.text) is ( + not _POLARS_2 and chunksize is None and not kwargs + ) + if scanned_by_polars_1: + return assert df["ts"].to_list() == INT96_VALUES assert df["arr"].to_list() == [[v] for v in INT96_VALUES] assert df["st"].to_list() == [{"t": v} for v in INT96_VALUES] @@ -729,24 +736,69 @@ def test_read_parquet_int96(self, tmp_path, caplog, chunksize, kwargs): if _POLARS_2 else [[{"key": "k", "value": v}] for v in INT96_VALUES] ) - assert (_UPGRADE in caplog.text) is (not _POLARS_2 and chunksize is None and not kwargs) - @pytest.mark.skipif(_POLARS_2, reason="Polars 2.0 reads INT96 without pyarrow.") + @pytest.mark.skipif(_POLARS_2, reason="Polars 2.0 reads INT96 at microseconds.") @pytest.mark.parametrize( - ("chunksize", "kwargs", "pyarrow"), - [(None, {}, False), (3, {}, False), (None, {"low_memory": True}, True)], - ids=["whole-without-pyarrow", "chunks-without-pyarrow", "whole-read-options"], + ("chunksize", "kwargs", "pyarrow", "warned"), + [ + (None, {}, False, True), + (3, {}, False, True), + (None, {"low_memory": True}, True, True), + (None, {"columns": ["id"]}, True, False), + ], + ids=["whole-without-pyarrow", "chunks-without-pyarrow", "read-options", "no-timestamps"], ) - def test_read_parquet_int96_polars_1(self, tmp_path, chunksize, kwargs, pyarrow): - """Polars 1.x reads INT96 at microseconds only with pyarrow and no read options.""" + def test_read_parquet_int96_polars_1( + self, tmp_path, caplog, chunksize, kwargs, pyarrow, warned + ): + """Polars 1.x reads timestamps at nanoseconds without pyarrow or with read options. + + They are cast to microseconds, and a warning is logged if the read has any. + """ + timestamp = datetime(2020, 1, 2, 3, 4, 5, 123000) + file = tmp_path / "source.parquet" + pq.write_table( + pa.table({"id": [1], "ts": pa.array([timestamp], pa.timestamp("us"))}), + file, + use_deprecated_int96_timestamps=True, + store_schema=False, + ) + unload = tmp_path / "unload" + unload.mkdir() with ( patch( "pyathena.polars.result_set.importlib.util.find_spec", return_value=MagicMock() if pyarrow else None, ), - pytest.raises(NotSupportedError, match=_UPGRADE), + caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"), ): - _read_unload_result(tmp_path, chunksize, kwargs) + df = pl.concat(_read_unload_result(unload, chunksize, kwargs, file=file)) + if warned: + assert df.schema["ts"] == pl.Datetime("us") + assert df["ts"].to_list() == [timestamp] + assert df["id"].to_list() == [1] + assert (_NANOSECOND_READ in caplog.text) is warned + + @pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x has no Map type.") + @pytest.mark.parametrize( + ("chunksize", "kwargs"), [(None, {}), (3, {"low_memory": True})], ids=["whole", "scanned"] + ) + def test_read_parquet_int96_map_named_like_a_selector(self, tmp_path, chunksize, kwargs): + """A Map is restored by its name, which Polars would read as a selector in pl.col().""" + timestamp = datetime(1600, 1, 2, 3, 4, 5, 123000) + file = tmp_path / "source.parquet" + table = pa.table( + { + "id": [1], + "*": pa.array([[("k", timestamp)]], pa.map_(pa.string(), pa.timestamp("us"))), + } + ) + pq.write_table(table, file, use_deprecated_int96_timestamps=True, store_schema=False) + unload = tmp_path / "unload" + unload.mkdir() + df = pl.concat(_read_unload_result(unload, chunksize, kwargs, file=file)) + assert df.schema == pl.Schema({"id": pl.Int64, "*": _MAP}) + assert df.row(0) == (1, {"k": timestamp}) def test_read_parquet_int96_pyarrow_options(self, tmp_path): """coerce_int96_timestamp_unit in pyarrow_options replaces microseconds.""" diff --git a/tests/pyathena/polars/test_util.py b/tests/pyathena/polars/test_util.py index 601b2c394..0175b19ee 100644 --- a/tests/pyathena/polars/test_util.py +++ b/tests/pyathena/polars/test_util.py @@ -8,7 +8,7 @@ import polars as pl import pytest -from pyathena.polars.util import get_athena_type, holds_datetimes, to_microseconds +from pyathena.polars.util import get_athena_type, to_microseconds _MAP = getattr(pl, "Map", None) _requires_map = pytest.mark.skipif(_MAP is None, reason="Polars 1.x has no Map type") @@ -61,23 +61,6 @@ def test_to_microseconds_map(): assert to_microseconds(dtype, maps=False) == pl.List( pl.List(pl.Struct({"key": _US, "value": pl.Struct({"t": _US})})) ) - - -@pytest.mark.parametrize( - ("dtype", "expected"), - [ - (_US, True), - (pl.Date, False), - (pl.List(pl.Struct({"a": pl.Int8, "b": pl.List(_NS)})), True), - (pl.Struct({"a": pl.Int8, "b": pl.List(pl.String)}), False), - ], -) -def test_holds_datetimes(dtype, expected): - assert holds_datetimes(dtype) is expected - - -@_requires_map -def test_holds_datetimes_map(): - assert holds_datetimes(pl.Map(_US, pl.Int8)) - assert holds_datetimes(pl.Map(pl.String, pl.List(_NS))) - assert not holds_datetimes(pl.Map(pl.String, pl.Int8)) + # A Map without nanosecond Datetimes is kept. + for kept in (pl.Map(pl.String, pl.Int8), pl.Map(pl.String, _US)): + assert to_microseconds(kept, maps=False) == kept From 4552f5899094c14f81f57dc80f368ab3c9d0839e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 13:08:01 +0900 Subject: [PATCH 06/12] Work out the Polars output dtypes from the file schema once The casts after a read follow from the file schema and the read arguments, so they are worked out once per read instead of per chunk from each DataFrame's schema, and a schema given to execute() now selects no read arguments. The whole and chunked reads share the pyarrow condition and the warning, and the tests share the fixture's Arrow types and write their own tables in place. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 128 ++++++++++++++--------- tests/pyathena/arrow/test_reader.py | 15 ++- tests/pyathena/arrow/test_result_set.py | 10 +- tests/pyathena/polars/test_result_set.py | 36 +++---- tests/pyathena/util.py | 8 ++ 5 files changed, 109 insertions(+), 88 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index a269dcbf0..4b9eaaa14 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -49,6 +49,13 @@ _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." +) + class PolarsDataFrameIterator(abc.Iterator): # type: ignore[type-arg] """Iterator for chunked DataFrame results from Athena queries. @@ -768,10 +775,9 @@ 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 - reads them so only with ``use_pyarrow=True``. Without read options given to - ``execute()``, and with pyarrow installed, a result with timestamps is - therefore read with pyarrow instead; otherwise Polars 1.x reads them at - nanoseconds, and ``_output_casts()`` casts them to microseconds. + reads them so only with ``use_pyarrow=True``. When ``_reads_with_pyarrow()``, + a result with timestamps is therefore read with pyarrow instead; otherwise + Polars 1.x reads them at nanoseconds, and they are cast to microseconds. ``_prepare_parquet_location()`` must have found the unload location first. @@ -790,11 +796,10 @@ 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 not self._kwargs and importlib.util.find_spec("pyarrow"): - _logger.warning( - f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at " - "nanoseconds, so they are read with pyarrow instead. Upgrade to " - "Polars 2.0 or later to read them with Polars." + 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() kwargs = self._read_kwargs( @@ -803,16 +808,14 @@ def _read_parquet(self, schema: pl.Schema) -> pl.DataFrame: if options and "pyarrow_options" in options: # The options given to execute() are in it. kwargs["pyarrow_options"] = options["pyarrow_options"] + dtypes = self._output_dtypes(schema, options) try: - df = pl.read_parquet(self._unload_location, **kwargs) - casts = self._output_casts(df, schema, options) - if casts: - df = df.with_columns(casts) + df = self._cast(pl.read_parquet(self._unload_location, **kwargs), dtypes) except Exception as e: _logger.exception(f"Failed to read {self._unload_location}.") raise OperationalError(*e.args) from e - if casts and options is None: - self._warn_nanosecond_read() + if options is None and dtypes.keys() & set(df.columns): + self._warn_nanosecond_read(_NANOSECOND_READ_CONSEQUENCE) return df def _read_parquet_with_pyarrow(self) -> pl.DataFrame: @@ -875,16 +878,17 @@ def _int96_read_options(self, schema: pl.Schema) -> dict[str, Any] | None: pyarrow reads it at microseconds unless ``pyarrow_options`` sets ``coerce_int96_timestamp_unit``. Polars 2.0 and later read INT96 at the time unit of the read schema, except in a Map, which the read schema therefore - makes a List of Structs; ``_output_casts()`` restores it. Polars 1.x reads + makes a List of Structs; ``_output_dtypes()`` restores it. Polars 1.x reads INT96 at nanoseconds whatever the schema. Args: schema: The schema of the Parquet files. Returns: - No arguments if the files hold no nanosecond timestamps, ``pyarrow_options`` - with ``use_pyarrow=True``, the ``schema`` argument with Polars 2.0 and - later, or None with Polars 1.x, which cannot read them. + No arguments if the files hold no nanosecond timestamps or ``execute()`` + was given a ``schema``, ``pyarrow_options`` with ``use_pyarrow=True``, the + ``schema`` argument with Polars 2.0 and later, or None with Polars 1.x, + which cannot read them. """ import polars as pl @@ -897,54 +901,80 @@ def _int96_read_options(self, schema: pl.Schema) -> dict[str, Any] | None: **(self._kwargs.get("pyarrow_options") or {}), } } + if "schema" in self._kwargs: + return {} if int(pl.__version__.split(".")[0]) < 2: return None return { "schema": {name: to_microseconds(dtype, maps=False) for name, dtype in schema.items()} } - def _output_casts( - self, df: pl.DataFrame, schema: pl.Schema, options: dict[str, Any] | None - ) -> list[pl.Series]: - """Get the columns of a read that need casting to their output dtypes. + @staticmethod + def _output_dtypes(schema: pl.Schema, options: dict[str, Any] | None) -> dict[str, Any]: + """Get the output dtypes of the columns that a read gives other dtypes. - Polars 1.x reads timestamps at nanoseconds, which are cast to microseconds. - The Maps that the read schema of ``_int96_read_options()`` replaced are - restored. Nothing is cast if ``execute()`` was given its own ``schema``. + Polars 1.x reads timestamps at nanoseconds, which are cast to microseconds, + and the read schema of ``_int96_read_options()`` makes Maps Lists of Structs, + which are restored. Args: - df: The DataFrame that was read. schema: The schema of the Parquet files. options: The read arguments from ``_int96_read_options()``. Returns: - The cast columns, keeping their names. + The output dtypes by column name. """ - if "schema" in self._kwargs: - return [] if options is None: - dtypes = {name: to_microseconds(dtype) for name, dtype in df.schema.items()} + read_dtypes = dict(schema) elif "schema" in options: - dtypes = {name: to_microseconds(schema[name]) for name in df.columns if name in schema} + read_dtypes = options["schema"] else: - return [] + return {} + return { + name: dtype + for name, read_dtype in read_dtypes.items() + if (dtype := to_microseconds(schema[name])) != read_dtype + } + + @staticmethod + def _cast(df: pl.DataFrame, dtypes: dict[str, Any]) -> pl.DataFrame: + """Cast the columns of a DataFrame that are in ``dtypes`` to their dtypes. + + Args: + df: The DataFrame. + dtypes: The dtypes by column name, from ``_output_dtypes()``. + + Returns: + The DataFrame with the cast columns. + """ # Series, not pl.col(), so that a name such as "*" is not a selector. - return [ - df.get_column(name).cast(dtype) - for name, dtype in dtypes.items() - if dtype != df.schema[name] - ] + 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() -> None: - """Warn that Polars 1.x read timestamps at nanoseconds.""" + 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, ". + """ import polars as pl _logger.warning( f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at " - "nanoseconds, 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." + f"nanoseconds, {consequence}" ) def _as_polars(self) -> pl.DataFrame: @@ -1152,7 +1182,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 not self._kwargs and importlib.util.find_spec("pyarrow") is not None: + if self._reads_with_pyarrow(): return self._read_parquet_chunks( S3FileSystem(**self._csv_storage_options), self._data_manifest ) @@ -1212,7 +1242,9 @@ def _scan_parquet_chunks( """ import polars as pl - warned = False + dtypes = self._output_dtypes(schema, options) + if options is None and dtypes: + self._warn_nanosecond_read(_NANOSECOND_READ_CONSEQUENCE) try: lazy_df = pl.scan_parquet( location, @@ -1222,11 +1254,7 @@ def _scan_parquet_chunks( ), ) for batch in lazy_df.collect_batches(chunk_size=self._chunksize): - casts = self._output_casts(batch, schema, options) - if casts and options is None and not warned: - self._warn_nanosecond_read() - warned = True - yield batch.with_columns(casts) if casts else batch + yield self._cast(batch, dtypes) except Exception as e: _logger.exception(f"Failed to read {location}.") raise OperationalError(*e.args) from e diff --git a/tests/pyathena/arrow/test_reader.py b/tests/pyathena/arrow/test_reader.py index ab64df9c7..3fd0e8d61 100644 --- a/tests/pyathena/arrow/test_reader.py +++ b/tests/pyathena/arrow/test_reader.py @@ -15,7 +15,12 @@ from fsspec.spec import AbstractBufferedFile from pyathena.arrow.reader import ParquetChunkReader -from tests.pyathena.util import INT96_UNLOAD_FILE, INT96_VALUES, ImmediateExecutor +from tests.pyathena.util import ( + INT96_ARROW_TYPES, + INT96_UNLOAD_FILE, + INT96_VALUES, + ImmediateExecutor, +) class _RecordingFile(AbstractBufferedFile): @@ -212,13 +217,7 @@ def test_int96_timestamps(self, max_prefetch_size): """INT96 timestamps that Athena writes are read at microseconds at any level.""" with ParquetChunkReader(LocalFileSystem(), [str(INT96_UNLOAD_FILE)], 3) as reader: table = pa.concat_tables(reader) - timestamp = pa.timestamp("us") - assert table.schema.types[1:] == [ - timestamp, - pa.list_(timestamp), - pa.struct([("t", timestamp)]), - pa.map_(pa.string(), timestamp), - ] + assert table.schema.types[1:] == INT96_ARROW_TYPES assert table.column("ts").to_pylist() == INT96_VALUES assert table.column("mp").to_pylist() == [[("k", v)] for v in INT96_VALUES] diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index 45d578908..91c9ccf88 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -19,7 +19,7 @@ from pyathena.connection import Connection from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig -from tests.pyathena.util import INT96_UNLOAD_FILE, INT96_VALUES +from tests.pyathena.util import INT96_ARROW_TYPES, INT96_UNLOAD_FILE, INT96_VALUES def _read_unload_result(tmp_path): @@ -222,13 +222,7 @@ def test_open_csv_closes_file_when_stream_fails(self): def test_read_parquet_int96(self, tmp_path): """INT96 timestamps that Athena writes are read at microseconds at any level.""" table = _read_unload_result(tmp_path).sort_by("id") - timestamp = pa.timestamp("us") - assert table.schema.types[1:] == [ - timestamp, - pa.list_(timestamp), - pa.struct([("t", timestamp)]), - pa.map_(pa.string(), timestamp), - ] + assert table.schema.types[1:] == INT96_ARROW_TYPES assert table.drop_columns("id").to_pydict() == { "ts": INT96_VALUES, "arr": [[v] for v in INT96_VALUES], diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 324448d43..403810677 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -117,14 +117,21 @@ def _fetching_result_set(chunks, names, arraysize, converters=None) -> AthenaPol ) -def _read_unload_result(tmp_path, chunksize, kwargs, file=INT96_UNLOAD_FILE): +def _read_unload_result(tmp_path, chunksize, kwargs, table=None): """Read a Parquet file of an UNLOAD result through the result set, without AWS. + Args: + table: The table to write as Athena writes it, with INT96 timestamps and no + Arrow schema, or None for the file that Athena wrote, ``INT96_UNLOAD_FILE``. + Returns: The DataFrames that were read: the whole result, or its chunks. """ path = tmp_path / "0.parquet" - shutil.copy(file, path) + if table is None: + shutil.copy(INT96_UNLOAD_FILE, path) + else: + pq.write_table(table, path, use_deprecated_int96_timestamps=True, store_schema=False) result_set = _chunked_result_set() result_set._chunksize = chunksize result_set._kwargs = dict(kwargs) @@ -756,15 +763,7 @@ def test_read_parquet_int96_polars_1( They are cast to microseconds, and a warning is logged if the read has any. """ timestamp = datetime(2020, 1, 2, 3, 4, 5, 123000) - file = tmp_path / "source.parquet" - pq.write_table( - pa.table({"id": [1], "ts": pa.array([timestamp], pa.timestamp("us"))}), - file, - use_deprecated_int96_timestamps=True, - store_schema=False, - ) - unload = tmp_path / "unload" - unload.mkdir() + table = pa.table({"id": [1], "ts": pa.array([timestamp], pa.timestamp("us"))}) with ( patch( "pyathena.polars.result_set.importlib.util.find_spec", @@ -772,7 +771,7 @@ def test_read_parquet_int96_polars_1( ), caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"), ): - df = pl.concat(_read_unload_result(unload, chunksize, kwargs, file=file)) + df = pl.concat(_read_unload_result(tmp_path, chunksize, kwargs, table=table)) if warned: assert df.schema["ts"] == pl.Datetime("us") assert df["ts"].to_list() == [timestamp] @@ -786,17 +785,13 @@ def test_read_parquet_int96_polars_1( def test_read_parquet_int96_map_named_like_a_selector(self, tmp_path, chunksize, kwargs): """A Map is restored by its name, which Polars would read as a selector in pl.col().""" timestamp = datetime(1600, 1, 2, 3, 4, 5, 123000) - file = tmp_path / "source.parquet" table = pa.table( { "id": [1], "*": pa.array([[("k", timestamp)]], pa.map_(pa.string(), pa.timestamp("us"))), } ) - pq.write_table(table, file, use_deprecated_int96_timestamps=True, store_schema=False) - unload = tmp_path / "unload" - unload.mkdir() - df = pl.concat(_read_unload_result(unload, chunksize, kwargs, file=file)) + df = pl.concat(_read_unload_result(tmp_path, chunksize, kwargs, table=table)) assert df.schema == pl.Schema({"id": pl.Int64, "*": _MAP}) assert df.row(0) == (1, {"k": timestamp}) @@ -830,12 +825,9 @@ def test_read_parquet_int96_schema_option(self, tmp_path, chunksize): def test_read_parquet_without_timestamps(self, tmp_path, caplog): """A result without timestamps is read by Polars with read options and no warning.""" - file = tmp_path / "source.parquet" - pl.DataFrame({"i": [1, None]}).write_parquet(file) - unload = tmp_path / "unload" - unload.mkdir() + table = pa.table({"i": [1, None]}) with caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"): - (df,) = _read_unload_result(unload, None, {"low_memory": True}, file=file) + (df,) = _read_unload_result(tmp_path, None, {"low_memory": True}, table=table) assert df["i"].to_list() == [1, None] assert not caplog.text diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 0ff742ce8..4b572b820 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -15,6 +15,7 @@ from unittest.mock import patch from zoneinfo import ZoneInfo +import pyarrow as pa from botocore.config import Config from botocore.exceptions import ClientError from botocore.response import StreamingBody @@ -671,6 +672,13 @@ def stub_multipart_copy(stubber, fail_list=False, fail_part=False, fail_annotati # UNLOAD (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, ...), ...) AS t (id, ts)) INT96_UNLOAD_FILE = Path(__file__).parents[1] / "resources" / "unload" / "int96_timestamps.parquet" +# The Arrow types of its columns ts, arr, st, and mp, with timestamps at microseconds. +INT96_ARROW_TYPES = [ + pa.timestamp("us"), + pa.list_(pa.timestamp("us")), + pa.struct([("t", pa.timestamp("us"))]), + pa.map_(pa.string(), pa.timestamp("us")), +] # The timestamps of its rows, by id: years 1 to 9999 that nanoseconds cannot hold, and # NULL. Each row's arr, st, and mp hold its timestamp. INT96_VALUES = [ From b71c65773d242aed38363c70667b99474a02e8f6 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 13:21:49 +0900 Subject: [PATCH 07/12] Translate read preparation errors and drop the use_pyarrow branch The storage options of the whole Polars read, which can refresh credentials, were computed outside the OperationalError boundary; the chunked scan now also works out its read arguments inside its own. use_pyarrow=True cannot read an UNLOAD prefix from S3 (Polars opens it as one file through fsspec), so the pyarrow_options merge for it is removed with its documentation. Co-Authored-By: Claude Opus 5.5 --- docs/polars.md | 1 - pyathena/polars/result_set.py | 56 +++++++++--------------- tests/pyathena/polars/test_result_set.py | 36 +++++++++------ 3 files changed, 43 insertions(+), 50 deletions(-) diff --git a/docs/polars.md b/docs/polars.md index 4b4506117..762fad79c 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -229,7 +229,6 @@ 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. -- With `use_pyarrow=True` given to `execute()`, a result read without the chunksize option is read by Polars with pyarrow, with `coerce_int96_timestamp_unit="us"` added to `pyarrow_options` unless they set it. - 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. 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. diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 4b9eaaa14..451e5fc6b 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -775,9 +775,9 @@ 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 - reads them so only with ``use_pyarrow=True``. When ``_reads_with_pyarrow()``, - a result with timestamps is therefore read with pyarrow instead; otherwise - Polars 1.x reads them at nanoseconds, and they are cast to microseconds. + 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. ``_prepare_parquet_location()`` must have found the unload location first. @@ -802,14 +802,11 @@ def _read_parquet(self, schema: pl.Schema) -> pl.DataFrame: "to read them with Polars." ) return self._read_parquet_with_pyarrow() - kwargs = self._read_kwargs( - lambda: self._get_parquet_storage_options(self.connection), **(options or {}) - ) - if options and "pyarrow_options" in options: - # The options given to execute() are in it. - kwargs["pyarrow_options"] = options["pyarrow_options"] dtypes = self._output_dtypes(schema, options) try: + kwargs = self._read_kwargs( + lambda: self._get_parquet_storage_options(self.connection), **(options or {}) + ) df = self._cast(pl.read_parquet(self._unload_location, **kwargs), dtypes) except Exception as e: _logger.exception(f"Failed to read {self._unload_location}.") @@ -874,34 +871,24 @@ def _int96_read_options(self, schema: pl.Schema) -> dict[str, Any] | None: """Get the Polars read arguments that read the INT96 timestamps of UNLOAD at microseconds. Polars reads INT96 at nanoseconds by default, which gives wrong values outside - 1677-09-21 to 2262-04-11. With ``use_pyarrow=True`` given to ``execute()``, - pyarrow reads it at microseconds unless ``pyarrow_options`` sets - ``coerce_int96_timestamp_unit``. Polars 2.0 and later read INT96 at the time - unit of the read schema, except in a Map, which the read schema therefore - makes a List of Structs; ``_output_dtypes()`` restores it. Polars 1.x reads - INT96 at nanoseconds whatever the schema. + 1677-09-21 to 2262-04-11. Polars 2.0 and later read INT96 at the time unit of + the read schema, except in a Map, which the read schema therefore makes a List + of Structs; ``_output_dtypes()`` restores it. Polars 1.x reads INT96 at + nanoseconds whatever the schema. Args: schema: The schema of the Parquet files. Returns: No arguments if the files hold no nanosecond timestamps or ``execute()`` - was given a ``schema``, ``pyarrow_options`` with ``use_pyarrow=True``, the - ``schema`` argument with Polars 2.0 and later, or None with Polars 1.x, - which cannot read them. + was given a ``schema``, the ``schema`` argument with Polars 2.0 and later, + or None with Polars 1.x, which cannot read them. """ import polars as pl - if all(to_microseconds(dtype) == dtype for dtype in schema.values()): - return {} - if self._kwargs.get("use_pyarrow"): - return { - "pyarrow_options": { - "coerce_int96_timestamp_unit": "us", - **(self._kwargs.get("pyarrow_options") or {}), - } - } - if "schema" in self._kwargs: + if "schema" in self._kwargs or all( + to_microseconds(dtype) == dtype for dtype in schema.values() + ): return {} if int(pl.__version__.split(".")[0]) < 2: return None @@ -1186,9 +1173,7 @@ def _iter_parquet_chunks(self, schema: pl.Schema) -> Iterator[pl.DataFrame]: return self._read_parquet_chunks( S3FileSystem(**self._csv_storage_options), self._data_manifest ) - return self._scan_parquet_chunks( - self._unload_location, self.connection, schema, self._int96_read_options(schema) - ) + return self._scan_parquet_chunks(self._unload_location, self.connection, schema) def _read_parquet_chunks(self, fs: S3FileSystem, paths: list[str]) -> Iterator[pl.DataFrame]: """Read the files of the data manifest in chunks with ``ParquetChunkReader``. @@ -1223,7 +1208,6 @@ def _scan_parquet_chunks( location: str, connection: Connection[Any], schema: pl.Schema, - options: dict[str, Any] | None, ) -> Iterator[pl.DataFrame]: """Scan the unload location in chunks with Polars and the read options. @@ -1232,7 +1216,6 @@ def _scan_parquet_chunks( connection: The connection whose credentials the read uses. They are frozen when the first chunk is read. schema: The schema of the Parquet files. - options: The read arguments from ``_int96_read_options()``. Yields: Polars DataFrame for each chunk. @@ -1242,10 +1225,11 @@ def _scan_parquet_chunks( """ import polars as pl - dtypes = self._output_dtypes(schema, options) - if options is None and dtypes: - self._warn_nanosecond_read(_NANOSECOND_READ_CONSEQUENCE) try: + 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) lazy_df = pl.scan_parquet( location, **self._read_kwargs( diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 403810677..8b3e2d438 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -44,11 +44,10 @@ if _POLARS_2 else pl.List(pl.Struct({"key": pl.String, "value": pl.Datetime("us")})) ) -# The ways to read an UNLOAD result: whole, whole by Polars with pyarrow, in chunks -# with ParquetChunkReader, and in chunks that Polars scans with the read options. +# The ways to read an UNLOAD result: whole, in chunks with ParquetChunkReader, and in +# chunks that Polars scans with the read options. _UNLOAD_READS = [ pytest.param(None, {}, id="whole"), - pytest.param(None, {"use_pyarrow": True}, id="whole-use-pyarrow"), pytest.param(3, {}, id="chunks"), pytest.param(3, {"low_memory": True}, id="scanned-chunks"), ] @@ -795,16 +794,6 @@ def test_read_parquet_int96_map_named_like_a_selector(self, tmp_path, chunksize, assert df.schema == pl.Schema({"id": pl.Int64, "*": _MAP}) assert df.row(0) == (1, {"k": timestamp}) - def test_read_parquet_int96_pyarrow_options(self, tmp_path): - """coerce_int96_timestamp_unit in pyarrow_options replaces microseconds.""" - (df,) = _read_unload_result( - tmp_path, - None, - {"use_pyarrow": True, "pyarrow_options": {"coerce_int96_timestamp_unit": "ms"}}, - ) - assert df.schema["ts"] == pl.Datetime("ms") - assert df.sort("id")["ts"].to_list() == INT96_VALUES - @pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x reads INT96 at nanoseconds.") @pytest.mark.parametrize("chunksize", [None, 3], ids=["whole", "scanned-chunks"]) def test_read_parquet_int96_schema_option(self, tmp_path, chunksize): @@ -823,6 +812,27 @@ def test_read_parquet_int96_schema_option(self, tmp_path, chunksize): assert df["ts"].to_list() == INT96_VALUES assert df["mp"].to_list() == [[{"key": "k", "value": v}] for v in INT96_VALUES] + @pytest.mark.parametrize(("chunksize", "kwargs"), [(None, {}), (3, {"low_memory": True})]) + def test_read_parquet_storage_options_error(self, chunksize, kwargs): + """A failure to get the storage options, such as of credentials, is OperationalError.""" + result_set = _chunked_result_set() + result_set._chunksize = chunksize + result_set._kwargs = dict(kwargs) + result_set._unload_location = "s3://bucket/unload/" + schema = pl.Schema({"i": pl.Int64}) + with patch.object( + AthenaPolarsResultSet, + "_get_parquet_storage_options", + side_effect=RuntimeError("credentials"), + ): + if chunksize is None: + with pytest.raises(OperationalError, match="credentials"): + result_set._read_parquet(schema) + else: + chunks = result_set._iter_parquet_chunks(schema) + with pytest.raises(OperationalError, match="credentials"): + next(chunks) + def test_read_parquet_without_timestamps(self, tmp_path, caplog): """A result without timestamps is read by Polars with read options and no warning.""" table = pa.table({"i": [1, None]}) From 9c0b3f218e7f74b3c78828a6b6c1df54c8898b5e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 13:28:00 +0900 Subject: [PATCH 08/12] Drop the use_pyarrow route from the Polars S3 options docs Polars' use_pyarrow=True opens an UNLOAD prefix from S3 as one file and fails, and the chunked CSV and UNLOAD scans do not take it. Co-Authored-By: Claude Opus 5.5 --- docs/polars.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/polars.md b/docs/polars.md index 762fad79c..92bda6e1b 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -408,7 +408,7 @@ 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. 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`. +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. (async-polars-cursor)= From b08940fae0a7066a1c53bb2880431d0469cf4195 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 13:30:24 +0900 Subject: [PATCH 09/12] Revert "Drop the use_pyarrow route from the Polars S3 options docs" The sentence describes where Polars sends storage_options, which holds for use_pyarrow=True; that use_pyarrow=True cannot read an UNLOAD prefix is a separate, pre-existing defect. Co-Authored-By: Claude Opus 5.5 --- docs/polars.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/polars.md b/docs/polars.md index 92bda6e1b..762fad79c 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -408,7 +408,7 @@ 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. 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. +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`. (async-polars-cursor)= From bc55d4bbd39ededc34c26abfc7f980bf6a145b2d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 16:30:04 +0900 Subject: [PATCH 10/12] Separate the Polars tests that the rebase joined Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/polars/test_result_set.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 8b3e2d438..9bccae721 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -708,6 +708,7 @@ def test_read_csv_ignore_errors(self, tmp_path, reader): pl.Series([datetime(2020, 1, 2, 3, 4, 5, 123456), None], dtype=pl.Datetime("us")), check_dtypes=True, ) + @pytest.mark.parametrize(("chunksize", "kwargs"), _UNLOAD_READS) def test_read_parquet_int96(self, tmp_path, caplog, chunksize, kwargs): """INT96 timestamps that Athena writes are read at microseconds at any level. From 043e295b6cfacb3b9ba6df086b8a66c63e8b45d6 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 17:52:58 +0900 Subject: [PATCH 11/12] Test UNLOAD timestamps through the cursors against Athena The tests that patched private result set methods, asserted helper dtypes, or read a stored file are replaced by queries through each cursor's public API. They check the timestamps that Athena's UNLOAD writes inside arrays, rows, and a map named "*", in whole and chunked reads, and the documented options and warning: coerce_int96_timestamp_unit for PandasCursor, and for PolarsCursor a schema given to execute() on Polars 2.0 and the warning of a read with read options on Polars 1.x. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_cursor.py | 12 ++ tests/pyathena/arrow/test_reader.py | 15 +- tests/pyathena/arrow/test_result_set.py | 28 --- tests/pyathena/pandas/test_cursor.py | 26 +++ tests/pyathena/pandas/test_result_set.py | 48 +---- tests/pyathena/polars/test_cursor.py | 62 ++++++ tests/pyathena/polars/test_result_set.py | 194 ------------------ tests/pyathena/polars/test_util.py | 39 +--- tests/pyathena/temporal.py | 21 ++ tests/pyathena/util.py | 22 -- .../resources/unload/int96_timestamps.parquet | Bin 1254 -> 0 bytes 11 files changed, 125 insertions(+), 342 deletions(-) delete mode 100644 tests/resources/unload/int96_timestamps.parquet diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index a650f959e..6c8e6da21 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -202,6 +202,18 @@ def test_temporal_types(self, arrow_cursor, unload): assert row == temporal.fetched_row(unload) assert temporal.zones(row) == temporal.zones(temporal.fetched_row(unload)) + @pytest.mark.parametrize( + "arrow_cursor", [{"cursor_kwargs": {"unload": True}}], indirect=["arrow_cursor"] + ) + def test_unload_nested_timestamps(self, arrow_cursor): + """UNLOAD timestamps inside arrays, rows, and maps are read at microseconds.""" + table = arrow_cursor.execute(temporal.nested_query()).as_arrow().sort_by("id") + assert table.drop_columns("id").to_pydict() == { + "arr": [[v] for v in temporal.NESTED_VALUES], + "st": [{"t": v} for v in temporal.NESTED_VALUES], + "*": [[("k", v)] for v in temporal.NESTED_VALUES], + } + def test_complex(self, arrow_cursor): arrow_cursor.execute( """ diff --git a/tests/pyathena/arrow/test_reader.py b/tests/pyathena/arrow/test_reader.py index 3fd0e8d61..7569a4d2b 100644 --- a/tests/pyathena/arrow/test_reader.py +++ b/tests/pyathena/arrow/test_reader.py @@ -15,12 +15,7 @@ from fsspec.spec import AbstractBufferedFile from pyathena.arrow.reader import ParquetChunkReader -from tests.pyathena.util import ( - INT96_ARROW_TYPES, - INT96_UNLOAD_FILE, - INT96_VALUES, - ImmediateExecutor, -) +from tests.pyathena.util import ImmediateExecutor class _RecordingFile(AbstractBufferedFile): @@ -213,14 +208,6 @@ def test_streamed_file_row_groups(self, tmp_path, monkeypatch, cache_type): for (_, end), (_, span_end) in zip(ranges[:2], spans[:2], strict=True) ) - def test_int96_timestamps(self, max_prefetch_size): - """INT96 timestamps that Athena writes are read at microseconds at any level.""" - with ParquetChunkReader(LocalFileSystem(), [str(INT96_UNLOAD_FILE)], 3) as reader: - table = pa.concat_tables(reader) - assert table.schema.types[1:] == INT96_ARROW_TYPES - assert table.column("ts").to_pylist() == INT96_VALUES - assert table.column("mp").to_pylist() == [[("k", v)] for v in INT96_VALUES] - def test_no_files(self): assert list(ParquetChunkReader(LocalFileSystem(), [], 4)) == [] diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index 91c9ccf88..cd214b6b4 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -5,35 +5,18 @@ # # SPDX-License-Identifier: MIT import os -import shutil import sys from datetime import datetime from unittest.mock import MagicMock, patch import pyarrow as pa import pytest -from pyarrow import fs as pyarrow_fs from pyathena.arrow.converter import DefaultArrowTypeConverter from pyathena.arrow.result_set import AthenaArrowResultSet from pyathena.connection import Connection from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig -from tests.pyathena.util import INT96_ARROW_TYPES, INT96_UNLOAD_FILE, INT96_VALUES - - -def _read_unload_result(tmp_path): - """Read the Parquet file that Athena wrote for UNLOAD through the result set, without AWS.""" - unload = tmp_path / "bucket" / "unload" - unload.mkdir(parents=True) - shutil.copy(INT96_UNLOAD_FILE, unload / "0.parquet") - result_set = AthenaArrowResultSet.__new__(AthenaArrowResultSet) # bypass __init__ - result_set._fs = pyarrow_fs.SubTreeFileSystem(str(tmp_path), pyarrow_fs.LocalFileSystem()) - result_set._unload_location = "s3://bucket/unload/" - with patch.object( - AthenaArrowResultSet, "_read_data_manifest", return_value=["s3://bucket/unload/0.parquet"] - ): - return result_set._read_parquet() class TestAthenaArrowResultSet: @@ -218,14 +201,3 @@ def test_open_csv_closes_file_when_stream_fails(self): with pytest.raises(OSError, match="HEAD failed"): result_set._open_csv("bucket/key.csv") file.close.assert_called_once_with() - - def test_read_parquet_int96(self, tmp_path): - """INT96 timestamps that Athena writes are read at microseconds at any level.""" - table = _read_unload_result(tmp_path).sort_by("id") - assert table.schema.types[1:] == INT96_ARROW_TYPES - assert table.drop_columns("id").to_pydict() == { - "ts": INT96_VALUES, - "arr": [[v] for v in INT96_VALUES], - "st": [{"t": v} for v in INT96_VALUES], - "mp": [[("k", v)] for v in INT96_VALUES], - } diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index bcb762d3a..1d370a8d1 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -562,6 +562,32 @@ def test_temporal_types(self, pandas_cursor, unload): row = pandas_cursor.fetchone() assert row == temporal.fetched_row(unload) + @pytest.mark.parametrize( + "pandas_cursor", + [{"cursor_kwargs": {"unload": True}}, {"cursor_kwargs": {"unload": True, "chunksize": 1}}], + indirect=["pandas_cursor"], + ids=["whole", "chunks"], + ) + def test_unload_nested_timestamps(self, pandas_cursor): + """UNLOAD timestamps inside arrays, rows, and maps are read at microseconds.""" + pandas_cursor.execute(temporal.nested_query()) + df = pd.concat(pandas_cursor.iter_chunks()).sort_values("id") + # The arrays are numpy arrays of datetime64[us]. + assert [v.tolist() for v in df["arr"]] == [[v] for v in temporal.NESTED_VALUES] + assert df["st"].tolist() == [{"t": v} for v in temporal.NESTED_VALUES] + assert df["*"].tolist() == [[("k", v)] for v in temporal.NESTED_VALUES] + + @pytest.mark.parametrize( + "pandas_cursor", [{"cursor_kwargs": {"unload": True}}], indirect=["pandas_cursor"] + ) + def test_unload_int96_unit_option(self, pandas_cursor): + """coerce_int96_timestamp_unit given to execute() replaces microseconds.""" + df = pandas_cursor.execute( + "SELECT TIMESTAMP '2020-01-02 03:04:05.123' AS ts", coerce_int96_timestamp_unit="ms" + ).as_pandas() + assert df["ts"].dtype == "datetime64[ms]" + assert df["ts"].tolist() == [pd.Timestamp(2020, 1, 2, 3, 4, 5, 123000)] + @pytest.mark.parametrize( ("pandas_cursor", "chunksize"), [ diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index b68e637ce..5440941f1 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -9,7 +9,6 @@ import csv import inspect import io -import shutil import threading from collections import defaultdict from concurrent.futures import ThreadPoolExecutor @@ -35,7 +34,7 @@ _read_csv_with_pyarrow, ) from pyathena.util import RetryConfig, retry_api_call, stop_retries_on -from tests.pyathena.util import INT96_UNLOAD_FILE, INT96_VALUES, ThrottledOperation +from tests.pyathena.util import ThrottledOperation def _typed_values(rows): @@ -560,29 +559,6 @@ def test_read_parquet_filesystem(self, execute_kwargs, path, filesystem_kwargs): **filesystem_kwargs, } - @pytest.mark.parametrize( - ("chunksize", "kwargs"), - [(None, {}), (None, {"dtype_backend": "pyarrow"}), (3, {})], - ids=["whole", "whole-pyarrow-dtypes", "chunks"], - ) - def test_read_parquet_int96(self, tmp_path, chunksize, kwargs): - """INT96 timestamps that Athena writes are read at microseconds at any level.""" - df = _read_int96_unload_result(tmp_path, chunksize, **kwargs).sort_values("id") - expected_dtype = "timestamp[us][pyarrow]" if kwargs else "datetime64[us]" - assert str(df["ts"].dtype) == expected_dtype - assert [None if pd.isna(v) else v for v in df["ts"]] == INT96_VALUES - # The numpy backend holds arrays as numpy arrays, and the pyarrow one as lists. - arrays = [v.tolist() if isinstance(v, np.ndarray) else v for v in df["arr"]] - assert arrays == [[v] for v in INT96_VALUES] - assert df["st"].tolist() == [{"t": v} for v in INT96_VALUES] - assert df["mp"].tolist() == [[("k", v)] for v in INT96_VALUES] - - def test_read_parquet_int96_unit_option(self, tmp_path): - """coerce_int96_timestamp_unit given to execute() replaces microseconds.""" - df = _read_int96_unload_result(tmp_path, None, coerce_int96_timestamp_unit="ms") - assert df["ts"].dtype == "datetime64[ms]" - assert [None if pd.isna(v) else v for v in df.sort_values("id")["ts"]] == INT96_VALUES - def test_iter_parquet_chunks_null_in_later_file(self, tmp_path): """Top-level integer and boolean columns with NULL anywhere convert alike in every chunk.""" chunks, _ = _read_unload_result( @@ -1022,28 +998,6 @@ def _unload_result_set(tmp_path, tables, chunksize, write_statistics=True, **kwa return result_set, paths -def _read_int96_unload_result(tmp_path, chunksize, **kwargs): - """Read the Parquet file that Athena wrote for UNLOAD through the result set, without AWS. - - Returns: - The whole result, or its chunks joined. - """ - result_set, _ = _unload_result_set(tmp_path, [], chunksize, **kwargs) - path = tmp_path / "0.parquet" - shutil.copy(INT96_UNLOAD_FILE, path) - with ( - patch.object(AthenaPandasResultSet, "_read_data_manifest", return_value=[str(path)]), - patch.object( - AthenaPandasResultSet, "is_unload", new_callable=PropertyMock, return_value=True - ), - ): - if chunksize is None: - # Without a filesystem, pandas reads the local unload location. - result_set._kwargs = {"filesystem": None, **kwargs} - return result_set._read_parquet("pyarrow") - return pd.concat(list(result_set._as_pandas()), ignore_index=True) - - def _read_unload_result(tmp_path, tables, chunksize, write_statistics=True, **kwargs): """Read Parquet files in chunks through the result set, without AWS. diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index ebe3e10ee..e814ffd30 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT import contextlib +import logging import random import string import time @@ -34,6 +35,9 @@ cached_file_systems, ) +# Polars 2.0 reads INT96 at the time unit of the read schema; Polars 1.x at nanoseconds. +_POLARS_2 = int(pl.__version__.split(".")[0]) >= 2 + class TestPolarsCursor: @pytest.mark.parametrize( @@ -190,6 +194,64 @@ def test_temporal_types_chunked(self, polars_cursor): c.value if c.athena_type in types else c.text for c in columns ] + @pytest.mark.parametrize( + ("polars_cursor", "kwargs"), + [ + ({"cursor_kwargs": {"unload": True}}, {}), + ({"cursor_kwargs": {"unload": True, "chunksize": 1}}, {}), + pytest.param( + {"cursor_kwargs": {"unload": True, "chunksize": 1}}, + {"low_memory": True}, + marks=pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x scans at nanoseconds."), + ), + ], + indirect=["polars_cursor"], + ids=["whole", "chunks", "scanned-chunks"], + ) + def test_unload_nested_timestamps(self, polars_cursor, kwargs): + """UNLOAD timestamps inside arrays, rows, and maps are read at microseconds.""" + polars_cursor.execute(temporal.nested_query(), **kwargs) + df = pl.concat(polars_cursor.iter_chunks()).sort("id") + assert df["arr"].to_list() == [[v] for v in temporal.NESTED_VALUES] + assert df["st"].to_list() == [{"t": v} for v in temporal.NESTED_VALUES] + # Polars 2.0 reads a map as Map, and Polars 1.x as a List of key-value Structs. + entries = pl.List(pl.Struct({"key": pl.String, "value": pl.Datetime("us")})) + assert df.get_column("*").cast(entries).to_list() == [ + [{"key": "k", "value": v}] for v in temporal.NESTED_VALUES + ] + + @pytest.mark.skipif(_POLARS_2, reason="Polars 2.0 reads INT96 at microseconds.") + @pytest.mark.parametrize( + "polars_cursor", [{"cursor_kwargs": {"unload": True}}], indirect=["polars_cursor"] + ) + def test_unload_timestamps_with_read_options_polars_1(self, polars_cursor, caplog): + """With read options, Polars 1.x reads the timestamps itself and a warning is logged.""" + with caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"): + df = polars_cursor.execute( + "SELECT TIMESTAMP '2020-01-02 03:04:05.123' AS ts", low_memory=True + ).as_polars() + assert df.schema == pl.Schema({"ts": pl.Datetime("us")}) + assert df["ts"].to_list() == [datetime(2020, 1, 2, 3, 4, 5, 123000)] + assert [r.levelno for r in caplog.records if r.name == "pyathena.polars.result_set"] == [ + logging.WARNING + ] + + @pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x reads INT96 at nanoseconds.") + @pytest.mark.parametrize( + "polars_cursor", [{"cursor_kwargs": {"unload": True}}], indirect=["polars_cursor"] + ) + def test_unload_schema_option(self, polars_cursor): + """A schema given to execute() is used as given, a map as a List of Structs too.""" + schema = {"mp": pl.List(pl.Struct({"key": pl.String, "value": pl.Datetime("ms")}))} + df = polars_cursor.execute( + "SELECT MAP(ARRAY['k'], ARRAY[TIMESTAMP '2020-01-02 03:04:05.123']) AS mp", + schema=schema, + ).as_polars() + assert df.schema == pl.Schema(schema) + assert df["mp"].to_list() == [ + [{"key": "k", "value": datetime(2020, 1, 2, 3, 4, 5, 123000)}] + ] + def test_complex(self, polars_cursor): polars_cursor.execute( """ diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 9bccae721..8684aa6c7 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -6,8 +6,6 @@ # SPDX-License-Identifier: MIT import inspect -import logging -import shutil import threading from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor @@ -17,8 +15,6 @@ from unittest.mock import MagicMock, PropertyMock, patch import polars as pl -import pyarrow as pa -import pyarrow.parquet as pq import pytest from fsspec.implementations.local import LocalFileSystem @@ -28,7 +24,6 @@ from pyathena.polars.converter import DefaultPolarsTypeConverter from pyathena.polars.result_set import AthenaPolarsResultSet, PolarsDataFrameIterator from pyathena.util import RetryConfig -from tests.pyathena.util import INT96_UNLOAD_FILE, INT96_VALUES _ROWS_BEFORE_FAILURE = 300_000 # A timestamp column of a CSV result file with 12 fractional digits and NULL, as text and @@ -37,23 +32,6 @@ _TRUNCATED_TIMESTAMPS = pl.Series( [datetime(2020, 1, 2, 3, 4, 5, 123456), None], dtype=pl.Datetime("us") ) -# Polars 2.0 and later read Parquet maps as Map, and INT96 at the unit of the schema. -_POLARS_2 = int(pl.__version__.split(".")[0]) >= 2 -_MAP = ( - pl.Map(pl.String, pl.Datetime("us")) - if _POLARS_2 - else pl.List(pl.Struct({"key": pl.String, "value": pl.Datetime("us")})) -) -# The ways to read an UNLOAD result: whole, in chunks with ParquetChunkReader, and in -# chunks that Polars scans with the read options. -_UNLOAD_READS = [ - pytest.param(None, {}, id="whole"), - pytest.param(3, {}, id="chunks"), - pytest.param(3, {"low_memory": True}, id="scanned-chunks"), -] -# The warnings of Polars 1.x reads of timestamps with pyarrow and at nanoseconds. -_PYARROW_READ = "so they are read with pyarrow instead" -_NANOSECOND_READ = "so those outside 1677-09-21 to 2262-04-11 are wrong" def _chunked_result_set() -> AthenaPolarsResultSet: @@ -116,45 +94,6 @@ def _fetching_result_set(chunks, names, arraysize, converters=None) -> AthenaPol ) -def _read_unload_result(tmp_path, chunksize, kwargs, table=None): - """Read a Parquet file of an UNLOAD result through the result set, without AWS. - - Args: - table: The table to write as Athena writes it, with INT96 timestamps and no - Arrow schema, or None for the file that Athena wrote, ``INT96_UNLOAD_FILE``. - - Returns: - The DataFrames that were read: the whole result, or its chunks. - """ - path = tmp_path / "0.parquet" - if table is None: - shutil.copy(INT96_UNLOAD_FILE, path) - else: - pq.write_table(table, path, use_deprecated_int96_timestamps=True, store_schema=False) - result_set = _chunked_result_set() - result_set._chunksize = chunksize - result_set._kwargs = dict(kwargs) - result_set._unload_location = f"{tmp_path}/" - with ( - patch.object(AthenaPolarsResultSet, "_read_data_manifest", return_value=[str(path)]), - patch.object( - AthenaPolarsResultSet, "is_unload", new_callable=PropertyMock, return_value=True - ), - patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch("pyathena.polars.result_set.S3FileSystem", lambda **_: LocalFileSystem()), - ): - if chunksize is None: - return [result_set._as_polars()] - assert result_set._prepare_parquet_location() - return list(result_set._iter_parquet_chunks(result_set._read_parquet_schema())) - - class TestAthenaPolarsResultSet: def test_fetch_across_chunks(self): """The fetch methods continue across chunks, reading a chunk when its rows are needed.""" @@ -709,139 +648,6 @@ def test_read_csv_ignore_errors(self, tmp_path, reader): check_dtypes=True, ) - @pytest.mark.parametrize(("chunksize", "kwargs"), _UNLOAD_READS) - def test_read_parquet_int96(self, tmp_path, caplog, chunksize, kwargs): - """INT96 timestamps that Athena writes are read at microseconds at any level. - - Polars 1.x reads INT96 only at nanoseconds: it reads a whole result without - read options with pyarrow, and scans with them at nanoseconds, giving wrong - values outside 1677-09-21 to 2262-04-11; both log a warning. - """ - with caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"): - df = pl.concat(_read_unload_result(tmp_path, chunksize, kwargs)).sort("id") - timestamp = pl.Datetime("us") - assert df.drop("id").schema == pl.Schema( - { - "ts": timestamp, - "arr": pl.List(timestamp), - "st": pl.Struct({"t": timestamp}), - "mp": _MAP, - } - ) - scanned_by_polars_1 = not _POLARS_2 and chunksize is not None and bool(kwargs) - assert (_NANOSECOND_READ in caplog.text) is scanned_by_polars_1 - assert (_PYARROW_READ in caplog.text) is ( - not _POLARS_2 and chunksize is None and not kwargs - ) - if scanned_by_polars_1: - return - assert df["ts"].to_list() == INT96_VALUES - assert df["arr"].to_list() == [[v] for v in INT96_VALUES] - assert df["st"].to_list() == [{"t": v} for v in INT96_VALUES] - assert df["mp"].to_list() == ( - [{"k": v} for v in INT96_VALUES] - if _POLARS_2 - else [[{"key": "k", "value": v}] for v in INT96_VALUES] - ) - - @pytest.mark.skipif(_POLARS_2, reason="Polars 2.0 reads INT96 at microseconds.") - @pytest.mark.parametrize( - ("chunksize", "kwargs", "pyarrow", "warned"), - [ - (None, {}, False, True), - (3, {}, False, True), - (None, {"low_memory": True}, True, True), - (None, {"columns": ["id"]}, True, False), - ], - ids=["whole-without-pyarrow", "chunks-without-pyarrow", "read-options", "no-timestamps"], - ) - def test_read_parquet_int96_polars_1( - self, tmp_path, caplog, chunksize, kwargs, pyarrow, warned - ): - """Polars 1.x reads timestamps at nanoseconds without pyarrow or with read options. - - They are cast to microseconds, and a warning is logged if the read has any. - """ - timestamp = datetime(2020, 1, 2, 3, 4, 5, 123000) - table = pa.table({"id": [1], "ts": pa.array([timestamp], pa.timestamp("us"))}) - with ( - patch( - "pyathena.polars.result_set.importlib.util.find_spec", - return_value=MagicMock() if pyarrow else None, - ), - caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"), - ): - df = pl.concat(_read_unload_result(tmp_path, chunksize, kwargs, table=table)) - if warned: - assert df.schema["ts"] == pl.Datetime("us") - assert df["ts"].to_list() == [timestamp] - assert df["id"].to_list() == [1] - assert (_NANOSECOND_READ in caplog.text) is warned - - @pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x has no Map type.") - @pytest.mark.parametrize( - ("chunksize", "kwargs"), [(None, {}), (3, {"low_memory": True})], ids=["whole", "scanned"] - ) - def test_read_parquet_int96_map_named_like_a_selector(self, tmp_path, chunksize, kwargs): - """A Map is restored by its name, which Polars would read as a selector in pl.col().""" - timestamp = datetime(1600, 1, 2, 3, 4, 5, 123000) - table = pa.table( - { - "id": [1], - "*": pa.array([[("k", timestamp)]], pa.map_(pa.string(), pa.timestamp("us"))), - } - ) - df = pl.concat(_read_unload_result(tmp_path, chunksize, kwargs, table=table)) - assert df.schema == pl.Schema({"id": pl.Int64, "*": _MAP}) - assert df.row(0) == (1, {"k": timestamp}) - - @pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x reads INT96 at nanoseconds.") - @pytest.mark.parametrize("chunksize", [None, 3], ids=["whole", "scanned-chunks"]) - def test_read_parquet_int96_schema_option(self, tmp_path, chunksize): - """A schema given to execute() replaces the one that reads INT96 at microseconds.""" - timestamp = pl.Datetime("us") - schema = { - "id": pl.Int32, - "ts": pl.Datetime("ms"), - "arr": pl.List(timestamp), - "st": pl.Struct({"t": timestamp}), - "mp": pl.List(pl.Struct({"key": pl.String, "value": timestamp})), - } - kwargs = {"schema": schema} if chunksize is None else {"schema": schema, "low_memory": True} - df = pl.concat(_read_unload_result(tmp_path, chunksize, kwargs)).sort("id") - assert df.schema == pl.Schema(schema) - assert df["ts"].to_list() == INT96_VALUES - assert df["mp"].to_list() == [[{"key": "k", "value": v}] for v in INT96_VALUES] - - @pytest.mark.parametrize(("chunksize", "kwargs"), [(None, {}), (3, {"low_memory": True})]) - def test_read_parquet_storage_options_error(self, chunksize, kwargs): - """A failure to get the storage options, such as of credentials, is OperationalError.""" - result_set = _chunked_result_set() - result_set._chunksize = chunksize - result_set._kwargs = dict(kwargs) - result_set._unload_location = "s3://bucket/unload/" - schema = pl.Schema({"i": pl.Int64}) - with patch.object( - AthenaPolarsResultSet, - "_get_parquet_storage_options", - side_effect=RuntimeError("credentials"), - ): - if chunksize is None: - with pytest.raises(OperationalError, match="credentials"): - result_set._read_parquet(schema) - else: - chunks = result_set._iter_parquet_chunks(schema) - with pytest.raises(OperationalError, match="credentials"): - next(chunks) - - def test_read_parquet_without_timestamps(self, tmp_path, caplog): - """A result without timestamps is read by Polars with read options and no warning.""" - table = pa.table({"i": [1, None]}) - with caplog.at_level(logging.WARNING, logger="pyathena.polars.result_set"): - (df,) = _read_unload_result(tmp_path, None, {"low_memory": True}, table=table) - assert df["i"].to_list() == [1, None] - assert not caplog.text - class TestPolarsDataFrameIterator: @pytest.mark.parametrize( diff --git a/tests/pyathena/polars/test_util.py b/tests/pyathena/polars/test_util.py index 0175b19ee..800387ccc 100644 --- a/tests/pyathena/polars/test_util.py +++ b/tests/pyathena/polars/test_util.py @@ -8,10 +8,7 @@ import polars as pl import pytest -from pyathena.polars.util import get_athena_type, to_microseconds - -_MAP = getattr(pl, "Map", None) -_requires_map = pytest.mark.skipif(_MAP is None, reason="Polars 1.x has no Map type") +from pyathena.polars.util import get_athena_type @pytest.mark.parametrize( @@ -27,40 +24,8 @@ def test_get_athena_type_nested(dtype, expected): assert get_athena_type(dtype) == expected -@_requires_map +@pytest.mark.skipif(not hasattr(pl, "Map"), reason="Polars 1.x has no Map type") def test_get_athena_type_map(): assert get_athena_type(pl.Map(pl.String, pl.Int32)) == ("map", 0, 0) assert get_athena_type(pl.Map(pl.String, pl.List(pl.Int32))) == ("map", 0, 0) assert get_athena_type(pl.List(pl.Map(pl.String, pl.Int32))) == ("array", 0, 0) - - -_NS, _US = pl.Datetime("ns"), pl.Datetime("us") - - -@pytest.mark.parametrize( - ("dtype", "expected"), - [ - (_NS, _US), - (pl.Datetime("ns", "UTC"), pl.Datetime("us", "UTC")), - (pl.Datetime("ms"), pl.Datetime("ms")), - (pl.Int64, pl.Int64), - (pl.List(_NS), pl.List(_US)), - (pl.Struct({"a": _NS, "b": pl.List(_NS)}), pl.Struct({"a": _US, "b": pl.List(_US)})), - ], -) -def test_to_microseconds(dtype, expected): - assert to_microseconds(dtype) == expected - assert to_microseconds(dtype, maps=False) == expected - - -@_requires_map -def test_to_microseconds_map(): - """A Map is kept, or with maps=False becomes a List of Structs of its key and value.""" - dtype = pl.List(pl.Map(_NS, pl.Struct({"t": _NS}))) - assert to_microseconds(dtype) == pl.List(pl.Map(_US, pl.Struct({"t": _US}))) - assert to_microseconds(dtype, maps=False) == pl.List( - pl.List(pl.Struct({"key": _US, "value": pl.Struct({"t": _US})})) - ) - # A Map without nanosecond Datetimes is kept. - for kept in (pl.Map(pl.String, pl.Int8), pl.Map(pl.String, _US)): - assert to_microseconds(kept, maps=False) == kept diff --git a/tests/pyathena/temporal.py b/tests/pyathena/temporal.py index bc7eb6aa2..e3082bc2c 100644 --- a/tests/pyathena/temporal.py +++ b/tests/pyathena/temporal.py @@ -190,3 +190,24 @@ def fetched_row(unload: bool = False) -> tuple[Any, ...]: def zones(row: tuple[Any, ...]) -> list[Any]: """The time zone of each aware value: a fixed-offset timezone or a ZoneInfo.""" return [v.tzinfo for v in row if isinstance(v, datetime | time) and v.tzinfo is not None] + + +# UNLOAD writes the timestamps inside arrays, rows, and maps as INT96 too. Each row +# holds one timestamp, ordered by id; the map column is named "*", which Polars would +# read as a selector in pl.col(). +_NESTED = [ + c + for name in ("ts_min", "ts_1600", "ts_max", "ts_null") + for c in UNLOAD_COLUMNS + if c.name == name +] +NESTED_VALUES = [c.value for c in _NESTED] + + +def nested_query() -> str: + """The SELECT query of the timestamps of NESTED_VALUES in an array, a row, and a map.""" + rows = ", ".join(f"({i}, {c.sql})" for i, c in enumerate(_NESTED)) + return ( + "SELECT id, ARRAY[ts] AS arr, CAST(ROW(ts) AS ROW(t TIMESTAMP)) AS st, " + f"MAP(ARRAY['k'], ARRAY[ts]) AS \"*\" FROM (VALUES {rows}) AS t (id, ts)" + ) diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 4b572b820..95d4a053c 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -15,7 +15,6 @@ from unittest.mock import patch from zoneinfo import ZoneInfo -import pyarrow as pa from botocore.config import Config from botocore.exceptions import ClientError from botocore.response import StreamingBody @@ -666,24 +665,3 @@ def stub_multipart_copy(stubber, fail_list=False, fail_part=False, fail_annotati ) return stubber.add_response("put_object_annotation", {}, put) - - -# A Parquet file that Athena's UNLOAD wrote, with INT96 timestamps and no Arrow schema: -# UNLOAD (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, ...), ...) AS t (id, ts)) -INT96_UNLOAD_FILE = Path(__file__).parents[1] / "resources" / "unload" / "int96_timestamps.parquet" -# The Arrow types of its columns ts, arr, st, and mp, with timestamps at microseconds. -INT96_ARROW_TYPES = [ - pa.timestamp("us"), - pa.list_(pa.timestamp("us")), - pa.struct([("t", pa.timestamp("us"))]), - pa.map_(pa.string(), pa.timestamp("us")), -] -# The timestamps of its rows, by id: years 1 to 9999 that nanoseconds cannot hold, and -# NULL. Each row's arr, st, and mp hold its timestamp. -INT96_VALUES = [ - datetime(1, 1, 1), - datetime(1600, 1, 2, 3, 4, 5, 123000), - datetime(9999, 12, 31, 23, 59, 59, 999000), - None, -] diff --git a/tests/resources/unload/int96_timestamps.parquet b/tests/resources/unload/int96_timestamps.parquet deleted file mode 100644 index daa6737b3fa63ea3fabc0d4d92646b7c3b05aa1e..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1254 zcmb_cO=uHA6rP!Er`s5?VrSfipldWs3Q29+rc^D8Kh+)Hn+K6X@gg3ihk_T|t0EpmLe>v;2n!wHkvD4RAQUuwxW=jXKp z-47z};*&}>+q?2*yg0P__3lR3X-b#x)naHla{J`bumECRY&ay9wk0i2<*!-@)oSje z=q~Xk*!cz@ZD~3bNzZC)&uR~^-N)$H#nk$~UE<{IRD-kC*dlO#%$q{w(A-G>o;qQa z^%;!(Wop4R-&1s_m^j>4ckW=*x$Tui7x3DNkt2I|NyUFVN7N6avA}mq<%#O#wC!0e ziWo@}g*b&|!dTg_34>`E392YZl!)c~AeL|{Lnv*F~!{(U8 zWqyfI>+-N7!Du5z!mx6b+VY0fIE*g{f|C$W0WVR&2UUT8iupIJ|HOJFK1h(7hEyq( zB8jbmNdjajlr5w-B#HtrF#nGA&k_9-*{OnHO@Sde5U?f8H?Y3mroWc{4U4IQn0S?@ zQ=vnmLV>@;`X<$vXk5WFaD#BM@2sMwSl!+&I^N0p+xG1n;;WR)gU>s_QKj>S?*z8r z6F8pTbE{mkmB{#&b4WQ}vHa_^c98M>j2J0x&i1QLxuo^>^ks9IzJbAkY-Xz1Z{_TK XK}+Rl9k-yBj#=Gk0lv@}{NerpVPpId From 1dc9b28e0dd9a2b39c314d7fb3ab2e5b1a010e6a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 18:08:11 +0900 Subject: [PATCH 12/12] Check the units, deeper nesting, and option paths of UNLOAD timestamps Milliseconds hold the values of the timestamps that UNLOAD writes, so the nested tests now assert the microsecond types before the values, and the Polars Map types before casting them. The query adds the top-level timestamp and a map in a row in an array. The Polars 1.x test uses the columns option, whose effect shows which reader ran, covers scanned chunks, and shows no warning for a read without timestamps; the schema option is also checked for scanned chunks. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_cursor.py | 10 +++- tests/pyathena/pandas/test_cursor.py | 10 +++- tests/pyathena/polars/test_cursor.py | 80 ++++++++++++++++++++++------ tests/pyathena/temporal.py | 13 +++-- 4 files changed, 88 insertions(+), 25 deletions(-) diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index 6c8e6da21..f50e1e66b 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -206,12 +206,20 @@ def test_temporal_types(self, arrow_cursor, unload): "arrow_cursor", [{"cursor_kwargs": {"unload": True}}], indirect=["arrow_cursor"] ) def test_unload_nested_timestamps(self, arrow_cursor): - """UNLOAD timestamps inside arrays, rows, and maps are read at microseconds.""" + """UNLOAD timestamps, alone and inside arrays, rows, and maps, are microseconds.""" table = arrow_cursor.execute(temporal.nested_query()).as_arrow().sort_by("id") + us = pa.timestamp("us") + assert table.column("ts").type == us + assert table.column("arr").type.value_type == us + assert table.column("st").type.field("t").type == us + assert table.column("*").type.item_type == us + assert table.column("deep").type.value_type.field("m").type.item_type == us assert table.drop_columns("id").to_pydict() == { + "ts": temporal.NESTED_VALUES, "arr": [[v] for v in temporal.NESTED_VALUES], "st": [{"t": v} for v in temporal.NESTED_VALUES], "*": [[("k", v)] for v in temporal.NESTED_VALUES], + "deep": [[{"m": [("k", v)]}] for v in temporal.NESTED_VALUES], } def test_complex(self, arrow_cursor): diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 1d370a8d1..da7b142ef 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -569,13 +569,19 @@ def test_temporal_types(self, pandas_cursor, unload): ids=["whole", "chunks"], ) def test_unload_nested_timestamps(self, pandas_cursor): - """UNLOAD timestamps inside arrays, rows, and maps are read at microseconds.""" + """UNLOAD timestamps, alone and inside arrays, rows, and maps, are microseconds.""" pandas_cursor.execute(temporal.nested_query()) df = pd.concat(pandas_cursor.iter_chunks()).sort_values("id") - # The arrays are numpy arrays of datetime64[us]. + assert df["ts"].dtype == "datetime64[us]" + assert [None if pd.isna(v) else v for v in df["ts"]] == temporal.NESTED_VALUES + # The arrays are numpy arrays, and the rows and maps Python objects. + assert {v.dtype for v in df["arr"]} == {np.dtype("datetime64[us]")} assert [v.tolist() for v in df["arr"]] == [[v] for v in temporal.NESTED_VALUES] assert df["st"].tolist() == [{"t": v} for v in temporal.NESTED_VALUES] assert df["*"].tolist() == [[("k", v)] for v in temporal.NESTED_VALUES] + assert [v.tolist() for v in df["deep"]] == [ + [{"m": [("k", v)]}] for v in temporal.NESTED_VALUES + ] @pytest.mark.parametrize( "pandas_cursor", [{"cursor_kwargs": {"unload": True}}], indirect=["pandas_cursor"] diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index e814ffd30..164d785af 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -209,44 +209,90 @@ def test_temporal_types_chunked(self, polars_cursor): ids=["whole", "chunks", "scanned-chunks"], ) def test_unload_nested_timestamps(self, polars_cursor, kwargs): - """UNLOAD timestamps inside arrays, rows, and maps are read at microseconds.""" + """UNLOAD timestamps, alone and inside arrays, rows, and maps, are microseconds.""" polars_cursor.execute(temporal.nested_query(), **kwargs) df = pl.concat(polars_cursor.iter_chunks()).sort("id") + us = pl.Datetime("us") + entries = pl.List(pl.Struct({"key": pl.String, "value": us})) + # Polars 2.0 reads a map as Map, and Polars 1.x as a List of key-value Structs. + mp = pl.Map(pl.String, us) if _POLARS_2 else entries + assert df.drop("id").schema == pl.Schema( + { + "ts": us, + "arr": pl.List(us), + "st": pl.Struct({"t": us}), + "*": mp, + "deep": pl.List(pl.Struct({"m": mp})), + } + ) + assert df["ts"].to_list() == temporal.NESTED_VALUES assert df["arr"].to_list() == [[v] for v in temporal.NESTED_VALUES] assert df["st"].to_list() == [{"t": v} for v in temporal.NESTED_VALUES] - # Polars 2.0 reads a map as Map, and Polars 1.x as a List of key-value Structs. - entries = pl.List(pl.Struct({"key": pl.String, "value": pl.Datetime("us")})) assert df.get_column("*").cast(entries).to_list() == [ [{"key": "k", "value": v}] for v in temporal.NESTED_VALUES ] + assert df.get_column("deep").cast(pl.List(pl.Struct({"m": entries}))).to_list() == [ + [{"m": [{"key": "k", "value": v}]}] for v in temporal.NESTED_VALUES + ] @pytest.mark.skipif(_POLARS_2, reason="Polars 2.0 reads INT96 at microseconds.") @pytest.mark.parametrize( - "polars_cursor", [{"cursor_kwargs": {"unload": True}}], indirect=["polars_cursor"] + ("polars_cursor", "kwargs", "columns", "warned"), + [ + ({"cursor_kwargs": {"unload": True}}, {"columns": ["ts"]}, ["ts"], True), + ( + {"cursor_kwargs": {"unload": True, "chunksize": 1}}, + {"low_memory": True}, + ["id", "ts"], + True, + ), + ({"cursor_kwargs": {"unload": True}}, {"columns": ["id"]}, ["id"], False), + ], + indirect=["polars_cursor"], + ids=["whole", "scanned-chunks", "no-timestamps"], ) - def test_unload_timestamps_with_read_options_polars_1(self, polars_cursor, caplog): - """With read options, Polars 1.x reads the timestamps itself and a warning is logged.""" + 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. + + The timestamps are cast to microseconds; the warning 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"): - df = polars_cursor.execute( - "SELECT TIMESTAMP '2020-01-02 03:04:05.123' AS ts", low_memory=True - ).as_polars() - assert df.schema == pl.Schema({"ts": pl.Datetime("us")}) - assert df["ts"].to_list() == [datetime(2020, 1, 2, 3, 4, 5, 123000)] - assert [r.levelno for r in caplog.records if r.name == "pyathena.polars.result_set"] == [ - logging.WARNING + polars_cursor.execute(query, **kwargs) + df = pl.concat(polars_cursor.iter_chunks()) + assert df.columns == columns + if "ts" in columns: + assert df.schema["ts"] == pl.Datetime("us") + assert df["ts"].to_list() == [datetime(2020, 1, 2, 3, 4, 5, 123000)] + warnings = [ + r + for r in caplog.records + if r.name == "pyathena.polars.result_set" and r.levelno == logging.WARNING ] + assert bool(warnings) is warned @pytest.mark.skipif(not _POLARS_2, reason="Polars 1.x reads INT96 at nanoseconds.") @pytest.mark.parametrize( - "polars_cursor", [{"cursor_kwargs": {"unload": True}}], indirect=["polars_cursor"] + ("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_schema_option(self, polars_cursor): + def test_unload_schema_option(self, polars_cursor, kwargs): """A schema given to execute() is used as given, a map as a List of Structs too.""" schema = {"mp": pl.List(pl.Struct({"key": pl.String, "value": pl.Datetime("ms")}))} - df = polars_cursor.execute( + polars_cursor.execute( "SELECT MAP(ARRAY['k'], ARRAY[TIMESTAMP '2020-01-02 03:04:05.123']) AS mp", schema=schema, - ).as_polars() + **kwargs, + ) + df = pl.concat(polars_cursor.iter_chunks()) assert df.schema == pl.Schema(schema) assert df["mp"].to_list() == [ [{"key": "k", "value": datetime(2020, 1, 2, 3, 4, 5, 123000)}] diff --git a/tests/pyathena/temporal.py b/tests/pyathena/temporal.py index e3082bc2c..1bf07afec 100644 --- a/tests/pyathena/temporal.py +++ b/tests/pyathena/temporal.py @@ -193,8 +193,9 @@ def zones(row: tuple[Any, ...]) -> list[Any]: # UNLOAD writes the timestamps inside arrays, rows, and maps as INT96 too. Each row -# holds one timestamp, ordered by id; the map column is named "*", which Polars would -# read as a selector in pl.col(). +# holds one timestamp, ordered by id, in ts, alone in arr, st, and the map "*", and in +# deep, a map in a row in an array. "*" is a name that Polars would read as a selector +# in pl.col(). _NESTED = [ c for name in ("ts_min", "ts_1600", "ts_max", "ts_null") @@ -205,9 +206,11 @@ def zones(row: tuple[Any, ...]) -> list[Any]: def nested_query() -> str: - """The SELECT query of the timestamps of NESTED_VALUES in an array, a row, and a map.""" + """The SELECT query of the timestamps of NESTED_VALUES, alone and in containers.""" rows = ", ".join(f"({i}, {c.sql})" for i, c in enumerate(_NESTED)) return ( - "SELECT id, ARRAY[ts] AS arr, CAST(ROW(ts) AS ROW(t TIMESTAMP)) AS st, " - f"MAP(ARRAY['k'], ARRAY[ts]) AS \"*\" FROM (VALUES {rows}) AS t (id, ts)" + "SELECT id, ts, ARRAY[ts] AS arr, CAST(ROW(ts) AS ROW(t TIMESTAMP)) AS st, " + "MAP(ARRAY['k'], ARRAY[ts]) AS \"*\", " + "ARRAY[CAST(ROW(MAP(ARRAY['k'], ARRAY[ts])) AS ROW(m MAP(VARCHAR, TIMESTAMP)))] AS deep " + f"FROM (VALUES {rows}) AS t (id, ts)" )