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..762fad79c 100644 --- a/docs/polars.md +++ b/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`. 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..451e5fc6b 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -37,7 +37,7 @@ 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 to_column_info, to_microseconds from pyathena.result_set import AthenaResultSet from pyathena.util import RetryConfig, override @@ -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. @@ -636,8 +643,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,9 +771,19 @@ 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; 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. + + Args: + schema: The schema of the Parquet files, from ``_read_parquet_schema()``. + Returns: Polars DataFrame containing the Parquet data. @@ -774,23 +792,63 @@ def _read_parquet(self) -> pl.DataFrame: """ 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.") + options = self._int96_read_options(schema) + if options is None and self._reads_with_pyarrow(): + self._warn_nanosecond_read( + "so they are read with pyarrow instead. Upgrade to Polars 2.0 or later " + "to read them with Polars." + ) + return self._read_parquet_with_pyarrow() + dtypes = self._output_dtypes(schema, options) try: - return pl.read_parquet( - self._unload_location, - **self._read_kwargs(lambda: self._get_parquet_storage_options(self.connection)), + 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}.") + raise OperationalError(*e.args) from e + 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: + """Read the files of the data manifest with pyarrow, with INT96 at microseconds. + + Returns: + Polars DataFrame containing the Parquet data. + + Raises: + OperationalError: If reading the Parquet files fails. + """ + import polars as pl + from pyarrow import parquet + + try: + table = parquet.ParquetDataset( + [path.removeprefix("s3://") for path in self._data_manifest], + filesystem=S3FileSystem(**self._csv_storage_options), + coerce_int96_timestamp_unit="us", + ).read(use_threads=True) + # Not rechunked, as Polars does not rechunk its own reads. + return cast("pl.DataFrame", pl.from_arrow(table, rechunk=False)) except Exception as e: _logger.exception(f"Failed to read {self._unload_location}.") raise OperationalError(*e.args) from e - def _read_parquet_schema(self) -> 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 +862,108 @@ 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. 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``, the ``schema`` argument with Polars 2.0 and later, + or None with Polars 1.x, which cannot read them. + """ + import polars as pl + + 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 + return { + "schema": {name: to_microseconds(dtype, maps=False) for name, dtype in schema.items()} + } + + @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, + and the read schema of ``_int96_read_options()`` makes Maps Lists of Structs, + which are restored. + + Args: + schema: The schema of the Parquet files. + options: The read arguments from ``_int96_read_options()``. + + Returns: + The output dtypes by column name. + """ + if options is None: + read_dtypes = dict(schema) + elif "schema" in options: + read_dtypes = options["schema"] + else: + 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. + casts = [df.get_column(name).cast(dtype) for name, dtype in dtypes.items() if name in df] + return df.with_columns(casts) if casts else df + + def _reads_with_pyarrow(self) -> bool: + """Whether the Parquet files are read with pyarrow through PyAthena's filesystem. + + They are without read options given to ``execute()``, and with pyarrow + installed: always in chunks, and with Polars 1.x when the files hold + timestamps. + + Returns: + True if they are read with pyarrow when pyarrow is chosen for them. + """ + return not self._kwargs and importlib.util.find_spec("pyarrow") is not None + + @staticmethod + def _warn_nanosecond_read(consequence: str) -> None: + """Warn that Polars 1.x reads the timestamps of UNLOAD at nanoseconds. + + Args: + consequence: The rest of the message, after "nanoseconds, ". + """ + import polars as pl + + _logger.warning( + f"Polars {pl.__version__} reads the Parquet timestamps of UNLOAD at " + f"nanoseconds, {consequence}" + ) + def _as_polars(self) -> pl.DataFrame: """Load query results as a Polars DataFrame. @@ -818,14 +972,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,19 +1143,23 @@ 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, 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. + Args: + schema: The schema of the Parquet files, from ``_read_parquet_schema()``. + Returns: An iterator of a Polars DataFrame for each chunk, read lazily. @@ -1009,11 +1169,11 @@ def _iter_parquet_chunks(self) -> 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 ) - return self._scan_parquet_chunks(self._unload_location, self.connection) + 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``. @@ -1044,7 +1204,10 @@ 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, ) -> Iterator[pl.DataFrame]: """Scan the unload location in chunks with Polars and the read options. @@ -1052,6 +1215,7 @@ 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. Yields: Polars DataFrame for each chunk. @@ -1062,11 +1226,19 @@ def _scan_parquet_chunks( import polars as pl 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(lambda: self._get_parquet_storage_options(connection)), + **self._read_kwargs( + lambda: self._get_parquet_storage_options(connection), + **(options or {}), + ), ) - yield from lazy_df.collect_batches(chunk_size=self._chunksize) + for batch in lazy_df.collect_batches(chunk_size=self._chunksize): + 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/pyathena/polars/util.py b/pyathena/polars/util.py index 98976a087..71414c45e 100644 --- a/pyathena/polars/util.py +++ b/pyathena/polars/util.py @@ -123,3 +123,32 @@ 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 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. + """ + 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) + 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 diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index d4774e46b..f50e1e66b 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"], @@ -206,6 +202,26 @@ 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, 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): arrow_cursor.execute( """ @@ -517,7 +533,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 +717,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/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index e4fd7fc9f..da7b142ef 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"], @@ -562,6 +562,38 @@ 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, 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") + 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"] + ) + 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 186bb8842..5440941f1 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -554,6 +554,7 @@ 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, } diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index adf7bea6c..164d785af 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( @@ -153,11 +157,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"], @@ -194,6 +194,110 @@ 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, 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] + 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", "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, 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"): + 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", "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, 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")}))} + polars_cursor.execute( + "SELECT MAP(ARRAY['k'], ARRAY[TIMESTAMP '2020-01-02 03:04:05.123']) AS mp", + schema=schema, + **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)}] + ] + 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 c54fd4496..8684aa6c7 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -388,7 +388,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 +504,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} diff --git a/tests/pyathena/temporal.py b/tests/pyathena/temporal.py index bc7eb6aa2..1bf07afec 100644 --- a/tests/pyathena/temporal.py +++ b/tests/pyathena/temporal.py @@ -190,3 +190,27 @@ 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, 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") + 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, alone and in containers.""" + rows = ", ".join(f"({i}, {c.sql})" for i, c in enumerate(_NESTED)) + return ( + "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)" + )