Read UNLOAD timestamps at microseconds

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/<date>/<uuid>/`.
 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.

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.
 

docs/polars.md

@@ -11,10 +11,11 @@ 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
-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.
+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
+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.
@@ -214,12 +215,22 @@ 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.
+Polars 1.x reads them correctly only through pyarrow:
+
+- Without the chunksize option and without Polars read options given to `execute()`, a result with timestamps is read with pyarrow when it is installed, through PyAthena's S3FileSystem, and a warning is logged.
+- With the chunksize option and without Polars read options, the chunks are read with pyarrow when it is installed.
+- Otherwise Polars 1.x reads the timestamps at nanoseconds, which gives wrong values outside 1677-09-21 to 2262-04-11; they are cast to microseconds, and a warning is logged.
+
 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`.
 
@@ -394,7 +405,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`.

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)
         ):

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}.")

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: