diff --git a/docs/pandas.md b/docs/pandas.md index 88a16b404..92dfd5cc4 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -220,7 +220,9 @@ for row in cursor: print(row) ``` -The DATE and TIMESTAMP of Athena's data type are returned as [pandas.Timestamp](https://pandas.pydata.org/docs/reference/api/pandas.Timestamp.html) type. +In a DataFrame, `timestamp` and `date` columns are `datetime64[us]`, and `timestamp with time zone` columns are `datetime64[us, UTC]`, with fractional digits beyond microseconds truncated. +The fetch methods return their values as `datetime`, `date`, and aware `datetime` objects in UTC, and NULL as `None`. +A value that cannot be converted, such as a `timestamp with time zone` value with an unknown zone or with a local time that its zone repeats or skips, raises `DataError` naming the column. ```python from pyathena import connect @@ -231,7 +233,7 @@ cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", cursor_class=PandasCursor).cursor() cursor.execute("SELECT col_timestamp FROM one_row_complex") -print(type(cursor.fetchone()[0])) # +print(type(cursor.fetchone()[0])) # ``` Execution information of the query can also be retrieved. @@ -390,6 +392,7 @@ awsathena+pandas://:@athena.{region_name}.amazonaws.com:443/{schema_name}?s3_sta When this option is used, the object returned by the as_pandas method is a `PandasDataFrameIterator` object. This object is an iterator of DataFrames, like the pandas `TextFileReader` object. It also provides `get_chunk()`, `as_pandas()`, `iterrows()`, `itertuples()`, and `close()`, and can be used as a context manager. +Its `iterrows()` and `itertuples()` return the values of temporal columns as the fetch methods do. ```python from pyathena import connect @@ -566,7 +569,7 @@ PandasCursor reads a `json` column as text and applies the converter mapping's ` A `json` column with NULL in the result or chunk holds objects, with `None` for NULL; otherwise it has the dtype that pandas infers from the converted values. Converters passed to `execute()` replace the converter mapping and are applied by pandas. -With `engine="pyarrow"`, PandasCursor uses the PyArrow engine only when the result is not a tab-separated `.txt` file, pyarrow is installed, no chunksize is set (explicitly or by `auto_optimize_chunksize`), `quoting` is the default, the result has no columns that need a converter (`boolean`, `decimal`, `varbinary`, `json`, `time with time zone`, and `timestamp with time zone` with the default converter), and the result file is at least `AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES` bytes. +With `engine="pyarrow"`, PandasCursor uses the PyArrow engine only when the result is not a tab-separated `.txt` file, pyarrow is installed, no chunksize is set (explicitly or by `auto_optimize_chunksize`), `quoting` is the default, the result has no columns that need a converter (`boolean`, `decimal`, `varbinary`, `json`, `time with time zone`, columns of a type that the converter maps to a conversion function, and `timestamp with time zone` columns that PyAthena does not convert after reading), no `converters` are given to `execute()`, and the result file is at least `AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES` bytes. Otherwise, it falls back to the C engine. With `engine="pyarrow"`, tab-separated `.txt` results from DDL statements such as `SHOW TABLES`, `SHOW COLUMNS`, and `DESCRIBE` use the C engine to preserve leading zeros, exponent notation, and padding in string values. With the PyArrow engine, when `keep_default_na` and `na_values` are the defaults and the only pandas.read_csv() options passed to `execute()` are `dtype` as a mapping and `parse_dates` as a list, PyAthena reads the file with `pyarrow.csv`, allowing newlines in quoted values, and converts it to a DataFrame as `pandas.read_csv(engine="pyarrow")` does, except that columns with a string dtype are read as text, and without `future.infer_string`, a column with dtype `str` keeps NULL as a missing value. @@ -575,6 +578,7 @@ It also changes the values of a column with a string dtype whose values all look Apart from PandasCursor's own options such as `engine` and `chunksize`, an option passed here replaces the value PyAthena sets for the same pandas.read_csv() argument. For example, `dtype` replaces the whole column type mapping, and `parse_dates` replaces the list of date, time, and timestamp columns. +With `parse_dates`, `dtype`, or `converters` passed here, PyAthena does not convert the `date`, `timestamp`, and `timestamp with time zone` columns after reading. pandas reads them with these options and the values PyAthena sets for the others: `parse_dates` for date, time, and timestamp columns, and the converter mapping, which converts `timestamp with time zone` values in their zones. The fetch methods return their values as pandas holds them. A conversion function that the converter maps to one of these types also replaces PyAthena's conversion of the columns of that type. (pandas-s3-read-options)= @@ -604,7 +608,8 @@ 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()`. +Timestamps are read at microseconds, as `datetime64[us]`, and dates as `datetime64[us]`, as from CSV results. +With `coerce_int96_timestamp_unit` or `dtype_backend` given to `execute()`, the columns have the dtypes that these options give, and the fetch methods return their values as pandas holds them. See {ref}`arrow-unload-options` for more information. @@ -768,7 +773,9 @@ print(df.describe()) print(df.head()) ``` -The DATE and TIMESTAMP of Athena's data type are returned as [pandas.Timestamp](https://pandas.pydata.org/docs/reference/api/pandas.Timestamp.html) type. +In a DataFrame, `timestamp` and `date` columns are `datetime64[us]`, and `timestamp with time zone` columns are `datetime64[us, UTC]`, with fractional digits beyond microseconds truncated. +The fetch methods return their values as `datetime`, `date`, and aware `datetime` objects in UTC, and NULL as `None`. +A value that cannot be converted, such as a `timestamp with time zone` value with an unknown zone or with a local time that its zone repeats or skips, raises `DataError` naming the column. ```python from pyathena import connect @@ -780,7 +787,7 @@ cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", query_id, future = cursor.execute("SELECT col_timestamp FROM one_row_complex") result_set = future.result() -print(type(result_set.fetchone()[0])) # +print(type(result_set.fetchone()[0])) # ``` As with AsyncCursor, you need a query ID to cancel a query. diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 30348f9c6..9cfc0cd3a 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -7,10 +7,11 @@ import logging import sys from collections import abc -from collections.abc import Callable, Iterable, Iterator +from collections.abc import Callable, Iterable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack from contextvars import copy_context +from datetime import datetime, timezone from functools import partial from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper from itertools import chain, islice @@ -28,8 +29,8 @@ from pyathena import OperationalError from pyathena.arrow.reader import ParquetChunkReader from pyathena.arrow.util import to_column_info -from pyathena.converter import Converter -from pyathena.error import ProgrammingError +from pyathena.converter import Converter, _parse_time_zone, _to_datetime_with_tz +from pyathena.error import DataError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.pandas.reader import _BINARY_NULL, BinaryCSVReader @@ -37,7 +38,7 @@ from pyathena.util import RetryConfig, override, parse_output_location if TYPE_CHECKING: - from pandas import DataFrame, Index, Series + from pandas import DataFrame, Series from pandas.io.parsers import TextFileReader from pyarrow import Table from pyarrow.parquet import FileMetaData @@ -48,6 +49,12 @@ _READ_FAILED_MESSAGE = "PandasDataFrameIterator cannot be read after a failed read." +# The Athena types whose CSV text AthenaPandasResultSet reads as text and converts after +# reading; see _TemporalConverter. +_CSV_TEMPORAL_TYPES = ("timestamp", "date", "timestamp with time zone") +# A timestamp's text up to its sixth fractional digit, for years of four digits. +_TIMESTAMP_TEXT_LENGTH = len("2000-01-01 00:00:00.000000") + def _convert_binary_csv(converter: Callable[[str | None], Any], value: str) -> Any: return converter(None if value == _BINARY_NULL else value) @@ -195,6 +202,32 @@ def _read_csv_with_pyarrow(source: str | IOBase, read_csv_kwargs: dict[str, Any] return df +def _replace_by_label(df: DataFrame, name: Any, convert: Callable[[Series], Series]) -> None: + """Replace the values of a column or an index level of a DataFrame, given its label. + + Args: + df: The DataFrame, updated in place. + name: The label of a column, or the name of an index level or of the index. A + partial label of MultiIndex columns is not replaced. + convert: Returns the new values of the given values, with their index. + """ + import pandas as pd + + index = df.index + if name in df.columns: + values = df[name] + if isinstance(values, pd.Series): + # Not a partial label of MultiIndex columns. + df[name] = convert(values) + elif isinstance(index, pd.MultiIndex) and name in index.names: + levels = [index.get_level_values(i) for i in range(index.nlevels)] + i = index.names.index(name) + levels[i] = pd.Index(convert(pd.Series(levels[i])), name=name) + df.index = pd.MultiIndex.from_arrays(levels, names=index.names) + elif index.name == name: + df.index = pd.Index(convert(pd.Series(index)), name=name) + + class _JSONConverter: """A json converter for the text of a json column that ``pandas.read_csv()`` read. @@ -235,29 +268,14 @@ def convert(self, df: DataFrame, name: Any) -> None: df: The DataFrame or chunk, updated in place. name: The label of the json column, which can also be an index level. """ - import pandas as pd + _replace_by_label(df, name, self._convert_values) - index = df.index - if name in df.columns: - values = df[name] - if isinstance(values, pd.Series): - # Not a partial label of MultiIndex columns. - df[name] = self._convert_values(values, index) - elif isinstance(index, pd.MultiIndex) and name in index.names: - levels = [index.get_level_values(i) for i in range(index.nlevels)] - i = index.names.index(name) - levels[i] = pd.Index(self._convert_values(levels[i]), name=name) - df.index = pd.MultiIndex.from_arrays(levels, names=index.names) - elif index.name == name: - df.index = pd.Index(self._convert_values(index), name=name) - - def _convert_values(self, values: Series | Index, index: Index | None = None) -> Series: + def _convert_values(self, values: Series) -> Series: """Convert json text values. Args: values: The text values, with NaN for a value that pandas read as missing. Values that pandas already parsed, such as an index or dates, are kept. - index: The index of the returned values, or None for a default index. Returns: The converted values: objects with None if there is NULL, as a NULL would @@ -267,7 +285,7 @@ def _convert_values(self, values: Series | Index, index: Index | None = None) -> if values.dtype.kind in "mM": # pandas parsed the text as dates, as parse_dates given to execute() asks. - return pd.Series(values.array, index=index) + return pd.Series(values.array, index=values.index) # Each value is one of three kinds: # - text, which pandas kept as read: the JSON text, converted here; # - a missing value: NULL, converted from an empty string, as pandas passes @@ -281,12 +299,165 @@ def _convert_values(self, values: Series | Index, index: Index | None = None) -> else (self._converter("") if pd.isna(v) else v) for v in values.to_numpy() ], - index=index, + index=values.index, dtype=object, ) return converted if converted.isna().any() else converted.infer_objects() +class _TemporalConverter: + """A converter for the text of a temporal column that ``pandas.read_csv()`` read. + + ``timestamp``, ``date``, and ``timestamp with time zone`` columns are read as text + and converted as whole columns after reading, so that their dtypes depend only on + the Athena type: ``datetime64[us]``, ``datetime64[us]``, and ``datetime64[us, UTC]``. + Fractional digits beyond microseconds are truncated. A column that pandas already + parsed, as with ``parse_dates`` when a ``dtype`` given to ``execute()`` replaces + PyAthena's, gets the same unit. + """ + + __slots__ = ("_type",) + + def __init__(self, type_: str) -> None: + """Create the converter of a column type. + + Args: + type_: The Athena type of the column, one of ``_CSV_TEMPORAL_TYPES``. + """ + self._type = type_ + + @property + def type(self) -> str: + """The Athena type of the column.""" + return self._type + + def convert(self, df: DataFrame, name: Any) -> None: + """Convert the text of a column in a DataFrame that ``read_csv()`` returned. + + Args: + df: The DataFrame or chunk, updated in place. + name: The label of the column, which can also be an index level. + + Raises: + DataError: If a value fails to convert, such as a ``timestamp with time + zone`` value with an unknown zone or a local time without one instant in + its zone. The message names the column. + """ + try: + _replace_by_label(df, name, self._convert_values) + except Exception as e: + raise DataError(f"Column {name!r}: {e}") from e + + def _convert_values(self, values: Series) -> Series: + """Convert the values of a column. + + Args: + values: The text values, with missing values for NULL, or values that + pandas already parsed as dates. + + Returns: + The converted values, with the index of ``values``. + """ + if self._type == "timestamp with time zone": + return _to_utc_timestamps(values) + return _to_timestamps(values, date=self._type == "date") + + +def _to_timestamps(values: Series, date: bool = False) -> Series: + """Convert timestamp or date text to ``datetime64[us]``. + + Args: + values: The text values, with missing values for NULL, or datetime values. + date: Whether the values are dates, which have no time to truncate. + + Returns: + The values as ``datetime64[us]``, with NaT for NULL. + """ + import pandas as pd + + if values.dtype.kind == "M": + return values.dt.as_unit("us") + if isinstance(values.dtype, pd.StringDtype) and values.dtype.storage == "pyarrow": + # pyarrow parses its strings without a Python object per value, as + # pandas.to_datetime() would make. + import pyarrow as pa + import pyarrow.compute as pc + + array = pa.array(values.array) + if date: + parsed = pc.cast(pc.cast(array, pa.date32()), pa.timestamp("us")) + else: + # Truncate the fraction to microseconds, as Athena's years have four digits. + text = pc.utf8_slice_codeunits(array, 0, _TIMESTAMP_TEXT_LENGTH) + parsed = pc.cast(text, pa.timestamp("us")) + return pd.Series(parsed.to_numpy(zero_copy_only=False), index=values.index) + text = values.astype("string") + if not date: + text = text.str.slice(0, _TIMESTAMP_TEXT_LENGTH) + return pd.to_datetime(text, format="ISO8601").dt.as_unit("us") + + +def _to_utc_timestamps(values: Series) -> Series: + """Convert ``timestamp with time zone`` text to ``datetime64[us, UTC]``. + + The local times of each zone are localized together. As for the fetch methods of + the other cursors, a local time that the zone repeats or skips is not converted. + + Args: + values: The text values, a local time followed by a zone name or a ``+HH:MM`` + or ``-HH:MM`` offset, with missing values for NULL, or aware datetime + values. + + Returns: + The instants as ``datetime64[us, UTC]``, with NaT for NULL. + + Raises: + DataError: If a zone is unknown, or a local time does not identify one instant + in its zone. + """ + import numpy as np + import pandas as pd + + if values.dtype.kind == "M": + return values.dt.tz_convert("UTC").dt.as_unit("us") + text = values if isinstance(values.dtype, pd.StringDtype) else values.astype("string") + # The local time before the last space, and the zone after it. pandas runs these + # regular expressions in pyarrow for pyarrow-backed strings, without an object + # per value. + local = _to_timestamps(text.str.replace(r" [^ ]*$", "", regex=True)).to_numpy( + dtype="datetime64[us]" + ) + # Integer codes of the zones, -1 for NULL, and the positions of each zone's values, + # grouped once by a stable sort of the codes. + codes, zone_names = pd.factorize(text.str.replace(r"^.* ", "", regex=True)) + order = np.argsort(codes, kind="stable") + bounds = np.searchsorted(codes[order], np.arange(-1, len(zone_names) + 1)) + utc = np.full(len(values), np.datetime64("NaT", "us")) + for code, zone_name in enumerate(zone_names): + positions = order[bounds[code + 1] : bounds[code + 2]] + zone = _parse_time_zone(zone_name) + if isinstance(zone, timezone): + utc[positions] = local[positions] - np.timedelta64(zone.utcoffset(None), "us") + continue + localized = pd.Series(local[positions]).dt.tz_localize( + zone, ambiguous="NaT", nonexistent="NaT" + ) + utc[positions] = ( + localized.dt.tz_convert("UTC").dt.tz_localize(None).to_numpy(dtype="datetime64[us]") + ) + # pandas also reports local times before a zone's first transition as + # nonexistent, so the values that it does not localize are converted one at a + # time, which raises the DataError of a value without one instant. + failed = positions[np.isnat(utc[positions])] + if len(failed): + offsets = [ + cast(datetime, _to_datetime_with_tz(value)).utcoffset() + for value in text.iloc[failed] + ] + utc[failed] = local[failed] - np.array(offsets, dtype="timedelta64[us]") + return pd.Series(utc, index=values.index).dt.tz_localize("UTC") + + class PandasDataFrameIterator(abc.Iterator): # type: ignore[type-arg] """Iterator for chunked DataFrame results from Athena queries. @@ -321,6 +492,7 @@ def __init__( reader: TextFileReader | Iterator[DataFrame] | DataFrame, trunc_date: Callable[[DataFrame], DataFrame], csv_stream: IOBase | None = None, + fetch_types: Mapping[Any, str] | None = None, ) -> None: """Initialize the iterator. @@ -329,6 +501,9 @@ def __init__( single DataFrame. trunc_date: Function to apply to each chunk, such as date truncation. csv_stream: Optional CSV stream owned and closed by this iterator. + fetch_types: The labels of the datetime columns whose values the row + iterators return as ``"datetime"`` or ``"date"`` objects, with None + for NaT, instead of pandas ``Timestamp`` values. """ from pandas import DataFrame @@ -338,6 +513,7 @@ def __init__( self._reader = reader self._trunc_date = trunc_date self._csv_stream = csv_stream + self._fetch_types: dict[Any, str] = dict(fetch_types or {}) # Whether a read raised, after which the position of the next chunk is unknown. self._read_failed = False @@ -445,9 +621,11 @@ def _iter_row_blocks(self, df: DataFrame) -> Iterator[Iterator[tuple[Any, ...]]] A block takes the values of a column at once with ``tolist()`` where that gives the values that iterating over the column gives: for columns of a - NumPy array and of a pandas string dtype. The other columns, such as - nullable integer, datetime, and ``pyarrow``-backed columns, are iterated - over. Blocks bound the values held at once for a large chunk. + NumPy array and of a pandas string dtype. The datetime columns of + ``fetch_types`` become ``datetime`` or ``date`` objects, with None for NaT. + The other columns, such as nullable integer, datetime, and + ``pyarrow``-backed columns, are iterated over. Blocks bound the values held + at once for a large chunk. Args: df: The DataFrame chunk. @@ -457,30 +635,53 @@ def _iter_row_blocks(self, df: DataFrame) -> Iterator[Iterator[tuple[Any, ...]]] """ import pandas as pd - # Each column's values, and whether a slice of them is converted with tolist(). - columns: list[tuple[Any, bool]] = [] + # Each column's values, and how a slice of them is converted: "list" with + # tolist(), "datetime" or "date" to objects, or None to be iterated over. + columns: list[tuple[Any, str | None]] = [] for i in range(df.shape[1]): series = df.iloc[:, i] array = series.array - if isinstance(array.dtype, pd.StringDtype): - columns.append((array, True)) + fetch_type = self._fetch_types.get(df.columns[i]) + if fetch_type is not None and series.dtype.kind == "M": + columns.append((array, fetch_type)) + elif isinstance(array.dtype, pd.StringDtype): + columns.append((array, "list")) elif isinstance(array, pd.arrays.NumpyExtensionArray): # NumpyExtensionArray.to_numpy() would scan the whole column for missing values. - columns.append((series.to_numpy(), True)) + columns.append((series.to_numpy(), "list")) else: - columns.append((array, False)) + columns.append((array, None)) size = self._ROW_BLOCK_SIZE for start in range(0, len(df), size): yield zip( - *( - values[start : start + size].tolist() - if to_list - else values[start : start + size] - for values, to_list in columns - ), + *(self._take_block(values, mode, start, start + size) for values, mode in columns), strict=True, ) + @staticmethod + def _take_block(values: Any, mode: str | None, start: int, stop: int) -> Any: + """Take a block of a column's values for the row iterators. + + Args: + values: The column's array. + mode: How the block is converted: ``"list"`` with ``tolist()``, + ``"datetime"`` or ``"date"`` to objects with None for NaT, or None to + be iterated over. + start: The position of the first value. + stop: The position after the last value. + + Returns: + The values of the block. + """ + block = values[start:stop] + if mode is None: + return block + if mode == "list": + return block.tolist() + objects = block.date if mode == "date" else block.to_pydatetime() + objects[block.isna()] = None + return objects.tolist() + def get_chunk(self, size: int | None = None) -> DataFrame: """Get a chunk of specified size. @@ -715,6 +916,9 @@ def __init__( self._csv_stream: IOBase | None = None # The converters of the json columns that _read_csv() reads as text, keyed by label. self._csv_json_converters: dict[Any, _JSONConverter] = {} + # The converters of the temporal columns that _read_csv() reads as text, keyed by + # label; see _get_csv_temporal_converters(). + self._csv_temporal_converters: dict[Any, _TemporalConverter] = {} # Cache time column names for efficient _trunc_date processing. _read_csv() # replaces them with the labels of the columns that it reads. @@ -728,23 +932,32 @@ def __init__( if self.state == AthenaQueryExecution.STATE_SUCCEEDED: if self.output_location: result = self._as_pandas() - trunc_date = _no_trunc_date if self.is_unload else self._finish_csv_frame + trunc_date = ( + self._finish_parquet_frame if self.is_unload else self._finish_csv_frame + ) else: # Without a result file, as with managed query result storage, the rows # from GetQueryResults are read as a CSV result file, but not in chunks. result = self._read_csv() trunc_date = self._finish_csv_frame + # The datetime columns that the fetch methods return as datetime or date. + self._fetch_types = self._get_fetch_types() if isinstance(result, pd.DataFrame): self._df = trunc_date(result) else: - self._df_iter = PandasDataFrameIterator(result, trunc_date, self._csv_stream) + self._df_iter = PandasDataFrameIterator( + result, trunc_date, self._csv_stream, self._fetch_types + ) else: + self._fetch_types = {} self._df = pd.DataFrame() if self._df is not None: # A shallow copy keeps assignments to the DataFrame from as_pandas() # out of the rows that the fetch methods return. Mutable values in its # cells, such as lists from JSON columns, are still shared. - self._df_iter = PandasDataFrameIterator(self._df.copy(deep=False), _no_trunc_date) + self._df_iter = PandasDataFrameIterator( + self._df.copy(deep=False), _no_trunc_date, fetch_types=self._fetch_types + ) self._row_iter: Iterator[tuple[Any, ...]] = self._df_iter.itertuples() def _get_parquet_engine(self) -> str: @@ -787,7 +1000,13 @@ def _get_csv_engine( is_compatible = ( effective_chunksize is None and self._quoting == 1 - and not self.converters + # pandas' PyArrow engine takes no converters. Those of the temporal + # columns that are converted after reading are not used. + and self._kwargs.get("converters") is None + and not self.converters.keys() + - self._get_csv_temporal_converters( + [(d[0], d) for d in self.description or []] + ).keys() # The pyarrow engine does not rename columns with the same name, and # the column labels cannot be resolved for it (see # _get_csv_column_labels()). @@ -923,7 +1142,7 @@ def _get_column_names(self) -> list[Any]: def _finish_csv_frame(self, df: DataFrame) -> DataFrame: """Finish a DataFrame read from the CSV result file. - Converts the json columns and truncates the time columns. + Converts the json and temporal columns and truncates the time columns. Args: df: The DataFrame or chunk that ``pandas.read_csv()`` returned. @@ -933,6 +1152,8 @@ def _finish_csv_frame(self, df: DataFrame) -> DataFrame: Raises: OperationalError: If converting a json column fails. + DataError: If a ``timestamp with time zone`` value has an unknown zone or a + local time without one instant, naming the column. """ for name, converter in self._csv_json_converters.items(): try: @@ -940,8 +1161,91 @@ def _finish_csv_frame(self, df: DataFrame) -> DataFrame: except Exception as e: _logger.exception(f"Failed to convert the json column {name!r}.") raise OperationalError(*e.args) from e + for name, temporal_converter in self._csv_temporal_converters.items(): + temporal_converter.convert(df, name) return self._trunc_date(df) + def _finish_parquet_frame(self, df: DataFrame) -> DataFrame: + """Finish a DataFrame or chunk read from the Parquet files of an UNLOAD result. + + ``date`` columns, which pyarrow reads as ``date`` objects, become + ``datetime64[us]``, as from a CSV result. With ``dtype_backend`` or + ``coerce_int96_timestamp_unit`` given to ``execute()``, the columns keep the + dtypes that these options give. + + Args: + df: The DataFrame or chunk. + + Returns: + The same DataFrame. + """ + if self._kwargs.keys() & {"dtype_backend", "coerce_int96_timestamp_unit"}: + return df + dates = {d[0] for d in self.description or [] if d[1] == "date"} + for i, column in enumerate(df.columns): + if column in dates and df.dtypes.iloc[i].kind == "O": + df.isetitem(i, df.iloc[:, i].astype("datetime64[us]")) + return df + + def _get_fetch_types(self) -> dict[Any, str]: + """Get the datetime columns that the fetch methods return as datetime or date. + + These are the temporal columns that PyAthena converts: those of + ``_get_csv_temporal_converters()`` for a CSV result, and the ``timestamp`` and + ``date`` columns of an UNLOAD result, unless ``dtype_backend`` or + ``coerce_int96_timestamp_unit`` given to ``execute()`` set their dtypes. + Columns that ``parse_dates`` given to ``execute()`` parsed keep pandas' values. + + Returns: + ``"datetime"`` or ``"date"``, keyed by column label. + """ + if not self.is_unload: + return { + label: "date" if converter.type == "date" else "datetime" + for label, converter in self._csv_temporal_converters.items() + } + if self._kwargs.keys() & {"dtype_backend", "coerce_int96_timestamp_unit"}: + return {} + return { + d[0]: "date" if d[1] == "date" else "datetime" + for d in self.description or [] + if d[1] in ("timestamp", "date") + } + + def _get_csv_temporal_converters( + self, columns: list[tuple[Any, tuple[Any, ...]]] + ) -> dict[Any, _TemporalConverter]: + """Get the converters of the temporal columns that are converted after reading. + + PyAthena converts the ``timestamp``, ``date``, and ``timestamp with time zone`` + columns after reading, unless ``parse_dates``, ``dtype``, or ``converters`` is + given to ``execute()``, with which pandas reads them, or the converter maps the + type to a conversion function other than the default ``timestamp with time + zone`` one. + + Args: + columns: The label and the description of each column that is read. + + Returns: + The converters keyed by column label. + """ + # None is the default of each option, as if it were not given. + if any(self._kwargs.get(key) is not None for key in ("parse_dates", "dtype", "converters")): + return {} + mappings = self._converter.mappings + temporal_converters = {} + for label, d in columns: + type_ = d[1] + if type_ not in _CSV_TEMPORAL_TYPES: + continue + if type_ == "timestamp with time zone": + if mappings.get(type_) is not _to_datetime_with_tz: + continue + elif type_ in mappings: + continue + temporal_converters[label] = _TemporalConverter(type_) + return temporal_converters + def _get_csv_converter(self, type_: str) -> Callable[[str | None], Any] | _JSONConverter: """Get the converter for a column type of the CSV result file. @@ -1174,7 +1478,19 @@ def _reads_csv_with_pyarrow(self) -> bool: ) def _get_csv_read_options(self, csv_engine: str, chunksize: int | None) -> dict[str, Any]: - """Build pandas options for Athena CSV or tab-separated results.""" + """Build pandas options for Athena CSV or tab-separated results. + + The temporal columns of ``_get_csv_temporal_converters()`` are read as text, and + ``_finish_csv_frame()`` converts them. The options given to ``execute()`` + replace the ones built here. + + Args: + csv_engine: The CSV engine that reads the file. + chunksize: The number of rows of each chunk, or None for a whole read. + + Returns: + The options for ``pandas.read_csv()``. + """ if self.output_location and self.output_location.endswith(".txt"): sep = "\t" header = None @@ -1184,17 +1500,27 @@ def _get_csv_read_options(self, csv_engine: str, chunksize: int | None) -> dict[ header = 0 names = None + description = self.description or [] + self._csv_temporal_converters = self._get_csv_temporal_converters( + [(d[0], d) for d in description] + ) + # The temporal columns that are converted after reading are read as text. + # pandas' PyArrow engine turns NULL into "None" with the NumPy str dtype, which + # str means without future.infer_string, so the PyArrow engine reads pyarrow + # strings. + temporal = self._csv_temporal_converters + text_dtype = "string[pyarrow]" if csv_engine == "pyarrow" else "str" read_csv_kwargs: dict[str, Any] = { "sep": sep, "header": header, "names": names, - "dtype": self.dtypes, + "dtype": {**self.dtypes, **dict.fromkeys(temporal, text_dtype)}, "converters": { d[0]: self._get_csv_converter(d[1]) - for d in self.description or [] - if d[1] in self._converter.mappings + for d in description + if d[1] in self._converter.mappings and d[0] not in self._csv_temporal_converters }, - "parse_dates": self.parse_dates, + "parse_dates": [name for name in self.parse_dates if name not in temporal], "skip_blank_lines": False, "keep_default_na": self._keep_default_na, "na_values": self._na_values, @@ -1337,7 +1663,8 @@ def _key_csv_columns_by_labels( Columns with the same name keep their own converters and date parsing, and options that rename or select columns get the types of the columns they read. The ``dtype``, ``converters``, and ``parse_dates`` given to - ``execute()`` are kept as they are. + ``execute()`` are kept as they are. The temporal columns that are converted + after reading are read as text. Args: read_csv_kwargs: The options for ``pandas.read_csv()``, updated in place. @@ -1347,21 +1674,26 @@ def _key_csv_columns_by_labels( columns = [ (label, d) for label, d in zip(labels, description, strict=True) if label is not None ] + self._csv_temporal_converters = self._get_csv_temporal_converters(columns) + temporal = self._csv_temporal_converters if "dtype" not in self._kwargs: read_csv_kwargs["dtype"] = { - label: dtype - for label, d in columns - if (dtype := self._converter.get_dtype(d[1], d[4], d[5])) is not None + **{ + label: dtype + for label, d in columns + if (dtype := self._converter.get_dtype(d[1], d[4], d[5])) is not None + }, + **dict.fromkeys(temporal, "str"), } if "converters" not in self._kwargs: read_csv_kwargs["converters"] = { label: self._get_csv_converter(d[1]) for label, d in columns - if d[1] in self._converter.mappings + if d[1] in self._converter.mappings and label not in self._csv_temporal_converters } if "parse_dates" not in self._kwargs: read_csv_kwargs["parse_dates"] = [ - label for label, d in columns if d[1] in self._PARSE_DATES + label for label, d in columns if d[1] in self._PARSE_DATES and label not in temporal ] self._time_columns = [label for label, d in columns if d[1] == "time"] @@ -1795,7 +2127,7 @@ def iter_chunks(self) -> PandasDataFrameIterator: ... process(df) # Single DataFrame with all data """ if self._df is not None: - return PandasDataFrameIterator(self._df, _no_trunc_date) + return PandasDataFrameIterator(self._df, _no_trunc_date, fetch_types=self._fetch_types) return self._df_iter @override diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 20ec7f3a4..ae2cd7d98 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -6,7 +6,7 @@ import string import time from concurrent.futures import ThreadPoolExecutor -from datetime import datetime +from datetime import date, datetime from decimal import Decimal from unittest.mock import MagicMock, patch @@ -450,16 +450,8 @@ def test_invalid_arraysize(self, pandas_cursor): @pytest.mark.parametrize( ("pandas_cursor", "unload"), [ - pytest.param( - {"cursor_kwargs": {"unload": False}}, - False, - marks=pytest.mark.xfail(strict=True, reason="#1219"), - ), - pytest.param( - {"cursor_kwargs": {"unload": True}}, - True, - marks=pytest.mark.xfail(strict=True, reason="#1219"), - ), + pytest.param({"cursor_kwargs": {"unload": False}}, False), + pytest.param({"cursor_kwargs": {"unload": True}}, True), ], indirect=["pandas_cursor"], ids=["csv", "unload"], @@ -487,6 +479,7 @@ def test_temporal_types(self, pandas_cursor, unload): # The fetch methods read the same DataFrame, whose aware values are in UTC. row = pandas_cursor.fetchone() assert row == temporal.fetched_row(unload) + assert [type(v) for v in row] == [type(v) for v in temporal.fetched_row(unload)] @pytest.mark.parametrize( "pandas_cursor", @@ -588,9 +581,9 @@ def test_complex(self, pandas_cursor, chunksize): 0.25, "a string", "varchar", - pd.Timestamp(2017, 1, 1, 0, 0, 0), + datetime(2017, 1, 1, 0, 0, 0), datetime(2017, 1, 1, 0, 0, 0).time(), - pd.Timestamp(2017, 1, 2), + date(2017, 1, 2), b"123", "[1, 2]", [1, 2], @@ -697,8 +690,8 @@ def test_complex_unload_pyarrow(self, pandas_cursor, parquet_engine): 0.25, "a string", "varchar", - pd.Timestamp(2017, 1, 1, 0, 0, 0), - datetime(2017, 1, 2).date(), + datetime(2017, 1, 1, 0, 0, 0), + date(2017, 1, 2), b"123", # ValueError: The truth value of an array with more than one element is ambiguous. # Use a.any() or a.all() @@ -1034,7 +1027,7 @@ def test_complex_unload_as_pandas_pyarrow(self, pandas_cursor, parquet_engine): str, str, np.datetime64, - np.object_, + np.datetime64, np.object_, np.object_, np.object_, @@ -1074,7 +1067,7 @@ def test_complex_unload_as_pandas_pyarrow(self, pandas_cursor, parquet_engine): "a string", "varchar", pd.Timestamp(2017, 1, 1, 0, 0, 0), - datetime(2017, 1, 2).date(), + pd.Timestamp(2017, 1, 2), b"123", # ValueError: The truth value of an array with more than one element is ambiguous. # Use a.any() or a.all() @@ -1585,12 +1578,11 @@ def test_fetch_all_rows(self, pandas_cursor): assert pandas_cursor.fetchall() == [CONVERTED_VALUES_ROW] def test_ambiguous_time_zone_value(self, pandas_cursor): - """A local time that its zone repeats raises instead of choosing one. + """A local time that its zone repeats raises DataError instead of choosing one. - A whole read converts the values in execute() and wraps the converter's - DataError in OperationalError. + A whole read converts the values in execute(). """ - with pytest.raises(DatabaseError, match="is ambiguous"): + with pytest.raises(DataError, match="col_timestamp_tz"): pandas_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) @pytest.mark.parametrize( @@ -1599,7 +1591,7 @@ def test_ambiguous_time_zone_value(self, pandas_cursor): def test_ambiguous_time_zone_value_chunked(self, pandas_cursor): """A chunked read raises DataError when it reads the chunk with the value.""" pandas_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) - with pytest.raises(DataError, match="is ambiguous"): + with pytest.raises(DataError, match="col_timestamp_tz"): pandas_cursor.fetchall() @pytest.mark.skipif(not ENV.managed_work_group, reason="AWS_ATHENA_MANAGED_WORKGROUP not set") diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index dbc4361b3..793c3cf7a 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -11,8 +11,9 @@ import re import threading from collections import defaultdict -from datetime import UTC, datetime +from datetime import UTC, date, datetime from unittest.mock import MagicMock, PropertyMock, patch +from zoneinfo import ZoneInfo import boto3 import numpy as np @@ -23,8 +24,9 @@ from botocore.exceptions import ClientError from botocore.response import StreamingBody from pandas.testing import assert_frame_equal +from pyarrow import csv as pyarrow_csv -from pyathena.error import OperationalError, ProgrammingError +from pyathena.error import DataError, OperationalError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.pandas.converter import DefaultPandasTypeConverter @@ -35,6 +37,7 @@ _read_csv_with_pyarrow, ) from pyathena.util import RetryConfig, retry_api_call, stop_retries_on +from tests.pyathena import temporal from tests.pyathena.util import ( LocalBucketFileSystem, ThrottledOperation, @@ -48,6 +51,14 @@ _TEXT_CSV = '"v"\n"a"\n"b"\n' +# Temporal values as pandas reads them with PyAthena's parse_dates and converters. +_TIMESTAMP_TEXT = "2020-01-02 03:04:05.123456789" +_PARSED = pd.Timestamp(_TIMESTAMP_TEXT) +_DATE = pd.Timestamp("2020-01-02") +_TOKYO_TEXT = "2020-01-02 03:04:05.000 Asia/Tokyo" +_TOKYO = datetime(2020, 1, 2, 3, 4, 5, tzinfo=ZoneInfo("Asia/Tokyo")) + + def _typed_values(rows): """Express the rows by the type and repr of each value, so that NaN and NA compare.""" return [[(type(value), repr(value)) for value in row] for row in rows] @@ -888,6 +899,400 @@ def test_read_csv_json_custom_converters(self, tmp_path, converters, expected): ) assert df["j"].tolist() == expected + @pytest.mark.parametrize( + ("kwargs", "reader"), + [ + pytest.param({"engine": "c"}, "c", id="c"), + pytest.param({"engine": "python"}, "python", id="python"), + pytest.param({"engine": "pyarrow"}, "pyathena-pyarrow", id="pyarrow"), + # An option that PyAthena's PyArrow reader does not take: pandas reads the file. + pytest.param( + {"engine": "pyarrow", "encoding": "utf-8"}, "pyarrow", id="pandas-pyarrow" + ), + pytest.param({"chunksize": 1}, "c", id="chunks"), + ], + ) + @pytest.mark.parametrize("infer_string", [True, False]) + def test_temporal_types(self, tmp_path, kwargs, reader, infer_string): + """Temporal columns have the dtype of their Athena type in every CSV read path. + + Fractions beyond microseconds are truncated, and aware values are held in UTC. + The fetch methods return datetime, date, and time objects, and None for NULL. + With ``future.infer_string``, the text is read as pyarrow-backed strings, and + without it, as objects. The PyArrow engine reads results without the ``time + with time zone`` columns, whose converter it does not take. + """ + columns = temporal.columns() + if reader.endswith("pyarrow"): + columns = [c for c in columns if c.athena_type != "time with time zone"] + with ( + pd.option_context("future.infer_string", infer_string), + patch("pyarrow.csv.read_csv", wraps=pyarrow_csv.read_csv) as pyarrow_read, + patch("pandas.read_csv", wraps=pd.read_csv) as pandas_read, + ): + self._check_temporal_types(tmp_path, columns, kwargs) + # pandas' PyArrow engine also reads with pyarrow.csv. + readers = {call.kwargs["engine"] for call in pandas_read.call_args_list} or ( + {"pyathena-pyarrow"} if pyarrow_read.called else set() + ) + assert readers == {reader} + + @staticmethod + def _check_temporal_types(tmp_path, columns, kwargs): + """Check the dtypes, values, and fetched values of temporal corpus columns. + + Args: + tmp_path: The directory of the result files. + columns: The columns of ``tests/pyathena/temporal.py`` in the result. + kwargs: The result set's options. + """ + data = _temporal_csv(columns) + assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES + types = [(c.name, c.athena_type) for c in columns] + df = _local_result_set(LocalBucketFileSystem(tmp_path / "df"), types, data, **kwargs) + df = pd.concat(list(df.iter_chunks())) + assert {name: str(dtype) for name, dtype in df.dtypes.items()} == { + c.name: _TEMPORAL_DTYPES[c.athena_type] for c in columns + } + assert [None if pd.isna(v) else v for v in df.iloc[0]] == [ + None + if c.value is None + else pd.Timestamp(c.value).tz_convert("UTC") + if c.athena_type == "timestamp with time zone" + else pd.Timestamp(c.value) + if c.athena_type in ("timestamp", "date") + else c.value + for c in columns + ] + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path / "rows"), types, data, **kwargs + ) + row = result_set.fetchone() + expected = tuple(c.value for c in columns) + assert row == expected + assert [type(v) for v in row] == [type(v) for v in expected] + + @pytest.mark.parametrize("infer_string", [True, False]) + @pytest.mark.parametrize("chunksize", [None, 2]) + def test_temporal_values_in_one_column(self, tmp_path, chunksize, infer_string): + """The dtype of a temporal column does not depend on its values or chunks. + + A column holds 0 to 12 fractional digits, years 1 and 9999, and NULL, and a time + zone column holds names and offsets. + """ + data = ( + '"ts","tstz"\n' + '"2020-01-02 03:04:05","2020-01-02 03:04:05 UTC"\n' + '"0001-01-01 00:00:00.123456789012","2020-01-02 03:04:05.123 Asia/Tokyo"\n' + ",\n" + '"9999-12-31 23:59:59.999999999999","2020-01-02 03:04:05.123456789 +05:30"\n' + ) + with pd.option_context("future.infer_string", infer_string): + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("ts", "timestamp"), ("tstz", "timestamp with time zone")], + data, + chunksize=chunksize, + ) + chunks = list(result_set.iter_chunks()) + assert {(str(c["ts"].dtype), str(c["tstz"].dtype)) for c in chunks} == { + ("datetime64[us]", "datetime64[us, UTC]") + } + df = pd.concat(chunks) + assert df["ts"].tolist()[:2] + df["ts"].tolist()[3:] == [ + pd.Timestamp(2020, 1, 2, 3, 4, 5), + pd.Timestamp(1, 1, 1, 0, 0, 0, 123456), + pd.Timestamp(9999, 12, 31, 23, 59, 59, 999999), + ] + assert df["tstz"].tolist()[:2] + df["tstz"].tolist()[3:] == [ + pd.Timestamp(2020, 1, 2, 3, 4, 5, tz="UTC"), + pd.Timestamp(2020, 1, 1, 18, 4, 5, 123000, tz="UTC"), + pd.Timestamp(2020, 1, 1, 21, 34, 5, 123456, tz="UTC"), + ] + assert pd.isna(df["ts"].iloc[2]) + assert pd.isna(df["tstz"].iloc[2]) + + @pytest.mark.parametrize("chunksize", [None, 1]) + def test_unload_temporal_types(self, tmp_path, chunksize): + """UNLOAD dates are datetime64[us], as CSV dates are, and fetch returns date objects.""" + tables = [ + pa.table( + { + "ts": pa.array([datetime(1, 1, 1, 0, 0, 0, 123000), None], pa.timestamp("us")), + "d": pa.array([date(9999, 12, 31), None], pa.date32()), + } + ) + ] + df = _local_result_set( + LocalBucketFileSystem(tmp_path / "df"), tables=tables, chunksize=chunksize + ) + df = pd.concat(list(df.iter_chunks())) + assert df.dtypes.tolist() == ["datetime64[us]", "datetime64[us]"] + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path / "rows"), tables=tables, chunksize=chunksize + ) + rows = result_set.fetchall() + assert rows == [(datetime(1, 1, 1, 0, 0, 0, 123000), date(9999, 12, 31)), (None, None)] + assert [type(v) for v in rows[0]] == [datetime, date] + + @pytest.mark.parametrize( + "kwargs", + [ + {"dtype_backend": "pyarrow"}, + {"coerce_int96_timestamp_unit": "ms"}, + ], + ) + def test_unload_temporal_options(self, tmp_path, kwargs): + """With read options that set the dtypes of UNLOAD dates, they are kept as given.""" + tables = [pa.table({"d": pa.array([date(2020, 1, 2)], pa.date32())})] + result_set = _local_result_set(LocalBucketFileSystem(tmp_path), tables=tables, **kwargs) + assert result_set.as_pandas()["d"].dtype != "datetime64[us]" + + @pytest.mark.parametrize("chunksize", [None, 1]) + @pytest.mark.parametrize( + ("value", "reason"), + [ + ("2021-11-07 01:30:00.000 America/New_York", "is ambiguous"), + ("2021-03-14 02:30:00.000 America/New_York", "does not exist"), + ("2021-03-14 02:30:00.000 Mars/Base", "Unknown time zone"), + ], + ids=["ambiguous", "nonexistent", "unknown_zone"], + ) + def test_time_zone_errors(self, tmp_path, chunksize, value, reason): + """A value without one instant raises DataError naming the column (#1217). + + A whole read raises in the constructor, and a chunked read when it reads the + chunk with the value. + """ + data = f'"tstz"\n"2020-01-02 03:04:05.000 UTC"\n"{value}"\n' + + def read(): + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("tstz", "timestamp with time zone")], + data, + chunksize=chunksize, + ) + return list(result_set.iter_chunks()) + + with pytest.raises(DataError) as error: + read() + assert str(error.value).startswith("Column 'tstz': ") + assert reason in str(error.value) + + @pytest.mark.parametrize( + ("kwargs", "converter_types", "expected"), + [ + # parse_dates given to execute() replaces PyAthena's: pandas parses the column. + pytest.param( + {"parse_dates": ["ts"]}, + None, + pd.Timestamp("2020-01-02 03:04:05.123456789"), + id="parse_dates", + ), + # pandas parses the dates of PyAthena's parse_dates after a dtype or converter + # that names the column, without PyAthena's truncation. + pytest.param( + {"dtype": {"ts": "str"}}, + None, + pd.Timestamp("2020-01-02 03:04:05.123456789"), + id="dtype", + ), + # pandas keeps a value that parse_dates cannot parse as it is. + pytest.param( + {"converters": {"ts": lambda value: "converted"}}, + None, + "converted", + id="converters", + ), + pytest.param({}, {"timestamp": lambda value: "converted"}, "converted", id="converter"), + ], + ) + def test_temporal_options_take_precedence(self, tmp_path, kwargs, converter_types, expected): + """Read options and conversion functions that name a temporal column decide its values. + + The fetch methods return the values as pandas holds them. + """ + converter = DefaultPandasTypeConverter() + for type_, function in (converter_types or {}).items(): + converter.set(type_, function) + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("ts", "timestamp")], + '"ts"\n"2020-01-02 03:04:05.123456789"\n', + converter=converter, + **kwargs, + ) + assert result_set.fetchone() == (expected,) + assert type(result_set.as_pandas()["ts"].iloc[0]) is type(expected) + + def test_time_zone_before_first_transition(self, tmp_path): + """Local times before a zone's first transition convert by the zone's offset then. + + pandas reports them as nonexistent; the zone's local mean time gives their + instant. The other values of their zones keep the zones' later offsets. + """ + data = ( + '"tstz"\n' + '"1500-01-01 00:00:00.000 America/New_York"\n' + '"2020-06-01 12:00:00.000 America/New_York"\n' + '"0001-01-01 00:00:00.000 Asia/Tokyo"\n' + '"2020-01-02 03:04:05.000 Asia/Tokyo"\n' + ) + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), [("tstz", "timestamp with time zone")], data + ) + df = result_set.as_pandas() + assert str(df["tstz"].dtype) == "datetime64[us, UTC]" + values = df["tstz"].dt.tz_localize(None).to_numpy() + # America/New_York's local mean time is -04:56:02 and Asia/Tokyo's +09:18:59. + expected = [ + "1500-01-01T04:56:02", + "2020-06-01T16:00:00", + "0000-12-31T14:41:01", + "2020-01-01T18:04:05", + ] + np.testing.assert_array_equal(values, np.array(expected, dtype="datetime64[us]")) + + @pytest.mark.parametrize("infer_string", [True, False]) + @pytest.mark.parametrize("type_", ["timestamp", "date", "timestamp with time zone"]) + def test_temporal_conversion_error(self, tmp_path, type_, infer_string): + """A value that fails to convert raises DataError naming the column.""" + value = { + "timestamp": "2020-13-01 00:00:00.000", + "date": "2020-13-01", + "timestamp with time zone": "2020-13-01 00:00:00.000 UTC", + }[type_] + with ( + pd.option_context("future.infer_string", infer_string), + pytest.raises(DataError) as error, + ): + _local_result_set(LocalBucketFileSystem(tmp_path), [("x", type_)], f'"x"\n"{value}"\n') + assert str(error.value).startswith("Column 'x': ") + + @pytest.mark.filterwarnings("ignore:Both a converter and dtype:pandas.errors.ParserWarning") + @pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"converters": {"s": str.upper}}, id="converters"), + pytest.param({"dtype": {"tstz": "str"}}, id="dtype_entry"), + pytest.param({"dtype": "str"}, id="dtype"), + ], + ) + def test_time_zone_results_do_not_depend_on_engine(self, tmp_path, kwargs): + """engine="pyarrow" gives the values of the C engine with options for the columns.""" + data = '"tstz","s"\n' + '"2020-01-02 03:04:05.000 Asia/Tokyo","a"\n' * 5 + types = [("tstz", "timestamp with time zone"), ("s", "varchar")] + frames, rows = [], [] + for engine in ("c", "pyarrow"): + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path / engine), types, data, engine=engine, **kwargs + ) + frames.append(result_set.as_pandas()) + rows.append(result_set.fetchall()) + assert_frame_equal(frames[1], frames[0]) + assert rows[1] == rows[0] + + @pytest.mark.parametrize("engine", ["c", "pyarrow"]) + @pytest.mark.parametrize("key", ["parse_dates", "dtype", "converters"]) + def test_temporal_options_none(self, tmp_path, engine, key): + """parse_dates, dtype, or converters of None, their default, is as if not given.""" + data = '"d","n"\n' + '"2020-01-02","1"\n' * 10 + assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("d", "date"), ("n", "integer")], + data, + engine=engine, + **{key: None}, + ) + assert str(result_set.as_pandas()["d"].dtype) == "datetime64[us]" + assert result_set.fetchone() == (date(2020, 1, 2), 1) + + @pytest.mark.parametrize("engine", ["c", "python", "pyarrow"]) + @pytest.mark.parametrize( + ("kwargs", "expected"), + [ + # PyAthena's parse_dates and converters read the temporal columns. + pytest.param({"dtype": {"n": "Int64"}}, (_PARSED, _DATE, _TOKYO, 1), id="dtype"), + # The converters replace the timestamp with time zone conversion. + pytest.param({"converters": {}}, (_PARSED, _DATE, _TOKYO_TEXT, 1), id="converters"), + # parse_dates replaces PyAthena's. + pytest.param( + {"parse_dates": ["ts"]}, + (_PARSED, "2020-01-02", _TOKYO, 1), + id="parse_dates", + ), + ], + ) + def test_temporal_options_read_by_pandas(self, tmp_path, engine, kwargs, expected): + """With parse_dates, dtype, or converters, pandas reads every temporal column. + + PyAthena's own options read them as before its conversion, and the fetch + methods return the values as pandas holds them, with any engine. + """ + data = '"ts","d","tstz","n"\n' + ( + f'"{_TIMESTAMP_TEXT}","2020-01-02","{_TOKYO_TEXT}","1"\n' * 2 + ) + assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [ + ("ts", "timestamp"), + ("d", "date"), + ("tstz", "timestamp with time zone"), + ("n", "integer"), + ], + data, + engine=engine, + **kwargs, + ) + row = result_set.fetchone() + assert row == expected + # pandas keeps the zone of the time zone values. + assert [v.utcoffset() for v in row if isinstance(v, datetime) and v.tzinfo] == [ + v.utcoffset() for v in expected if isinstance(v, datetime) and v.tzinfo + ] + + @pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"dtype": {"n": "Int64"}}, id="dtype"), + # pyarrow reads the dates of a column that parse_dates does not name as date + # objects, as pandas' PyArrow engine does. + pytest.param({"parse_dates": ["ts", "d"]}, id="parse_dates"), + # converters of None, their default, do not rule out the PyArrow engine. + pytest.param({"parse_dates": ["ts", "d"], "converters": None}, id="converters_none"), + ], + ) + def test_temporal_options_read_by_pandas_with_pyarrow(self, tmp_path, kwargs): + """The PyArrow engine reads temporal columns given these options as the C engine.""" + data = '"ts","d","n"\n' + f'"{_TIMESTAMP_TEXT}","2020-01-02","1"\n' * 3 + types = [("ts", "timestamp"), ("d", "date"), ("n", "integer")] + frames, rows = [], [] + with patch("pyarrow.csv.read_csv", wraps=pyarrow_csv.read_csv) as pyarrow_read: + for engine in ("c", "pyarrow"): + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path / engine), types, data, engine=engine, **kwargs + ) + frames.append(result_set.as_pandas()) + rows.append(result_set.fetchall()) + assert pyarrow_read.called + assert_frame_equal(frames[1], frames[0]) + assert rows[1] == rows[0] + + def test_custom_time_zone_converter(self, tmp_path): + """A conversion function that the converter maps to timestamp with time zone wins.""" + converter = DefaultPandasTypeConverter() + converter.set("timestamp with time zone", lambda value: value and value.lower()) + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("tstz", "timestamp with time zone")], + '"tstz"\n"2020-01-02 03:04:05.000 UTC"\n', + converter=converter, + ) + assert result_set.fetchone() == ("2020-01-02 03:04:05.000 utc",) + # A CSV result of Athena with each type the PyArrow engine reads, then a row of NULLs. _TYPES_CSV = ( @@ -1049,14 +1454,16 @@ def answer(self, operation, params): } -def _local_result_set(fs, types=(), data=None, tables=None, name="result.csv", **kwargs): +def _local_result_set( + fs, types=(), data=None, tables=None, name="result.csv", converter=None, **kwargs +): """Create a result set of a succeeded query over local result files, without AWS. The CSV ``data`` is the result file ``s3://bucket/``; ``tables`` are the files of an UNLOAD result, which the connection's S3 client lists in the data manifest. ``fs`` holds the files and stands in for PyAthena's filesystem, ``types`` lists the - (name, Athena type) pairs of the GetQueryResults metadata, and ``kwargs`` are the - result set's options. + (name, Athena type) pairs of the GetQueryResults metadata, ``converter`` replaces + the default converter, and ``kwargs`` are the result set's options. """ unload = tables is not None location = f"s3://bucket/{name}" @@ -1089,7 +1496,7 @@ def _local_result_set(fs, types=(), data=None, tables=None, name="result.csv", * with patch("pyathena.pandas.result_set.S3FileSystem", return_value=fs): return AthenaPandasResultSet( connection=connection, - converter=DefaultPandasTypeConverter(), + converter=converter or DefaultPandasTypeConverter(), query_execution=query_execution, arraysize=1000, retry_config=RetryConfig(), @@ -1098,13 +1505,49 @@ def _local_result_set(fs, types=(), data=None, tables=None, name="result.csv", * ) +def _temporal_csv(columns): + """Write the temporal corpus as Athena writes a CSV result file: quoted text, NULL empty. + + Args: + columns: The columns of ``tests/pyathena/temporal.py``. + + Returns: + The text of the result file, with one row. + """ + + def text(column): + if column.value is None: + return "" + # The text that Athena writes is the column's text, or else its SQL literal. + value = column.text or re.search(r"'([^']*)'", column.sql)[1] + return f'"{value}"' + + header = ",".join(f'"{c.name}"' for c in columns) + return f"{header}\n{','.join(text(c) for c in columns)}\n" + + +# The dtypes of the temporal columns of a PandasCursor result (#1212). +_TEMPORAL_DTYPES = { + "timestamp": "datetime64[us]", + "date": "datetime64[us]", + "timestamp with time zone": "datetime64[us, UTC]", + "time": "object", + "time with time zone": "object", +} + + def _string_series(values, infer_string): """Express the expected string dtype for the active pandas option.""" return pd.Series(values, dtype="str" if infer_string else object) -def _types_frame(infer_string, parse_time=True): - """Build the typed literal expectation for the all-types CSV.""" +def _types_frame(infer_string, parse_time=True, temporal_text=False): + """Build the typed literal expectation for the all-types CSV. + + With ``temporal_text``, the date and timestamp columns are pyarrow strings, as + PyAthena reads them for the PyArrow engine with its default options before + converting them. + """ missing = float("nan") return pd.DataFrame( { @@ -1120,12 +1563,24 @@ def _types_frame(infer_string, parse_time=True): "arr": _string_series(["[1, 2]", missing], infer_string), "m": _string_series(["{k=1}", missing], infer_string), "rw": _string_series(["{a=1, b=x}", missing], infer_string), - "dt": pd.Series([pd.Timestamp("2024-02-29"), pd.NaT], dtype="datetime64[us]"), - "ts": pd.Series( - [pd.Timestamp("2024-02-29 23:59:58.123"), pd.NaT], dtype="datetime64[ns]" + "dt": ( + pd.Series(["2024-02-29", None], dtype="string[pyarrow]") + if temporal_text + else pd.Series([pd.Timestamp("2024-02-29"), pd.NaT], dtype="datetime64[us]") + ), + "ts": ( + pd.Series(["2024-02-29 23:59:58.123", None], dtype="string[pyarrow]") + if temporal_text + else pd.Series( + [pd.Timestamp("2024-02-29 23:59:58.123"), pd.NaT], dtype="datetime64[ns]" + ) ), - "ts6": pd.Series( - [pd.Timestamp("2024-02-29 23:59:58.123456"), pd.NaT], dtype="datetime64[ns]" + "ts6": ( + pd.Series(["2024-02-29 23:59:58.123456", None], dtype="string[pyarrow]") + if temporal_text + else pd.Series( + [pd.Timestamp("2024-02-29 23:59:58.123456"), pd.NaT], dtype="datetime64[ns]" + ) ), # pandas supplies today's date when it parses a time without a date. "tm": ( @@ -1151,8 +1606,9 @@ def _types_frame(infer_string, parse_time=True): _TYPES_CSV, _TYPES, {}, - _types_frame, - [0, 1, 2, 3, 4, 5, 12, 13, 14, 15, 16, 17, 18], + # The date and timestamp columns are read as text and converted afterwards. + lambda infer: _types_frame(infer, temporal_text=True), + [0, 1, 2, 3, 4, 5, 15, 16, 17, 18], id="types", ), pytest.param( @@ -1202,13 +1658,11 @@ def _types_frame(infer_string, parse_time=True): [ pd.Series([1, None], name="x", dtype="Int64"), pd.Series([2, None], name="x", dtype="Int64"), - pd.Series( - [pd.Timestamp("2024-01-01"), pd.NaT], name="d", dtype="datetime64[us]" - ), + pd.Series(["2024-01-01", None], name="d", dtype="string[pyarrow]"), ], axis=1, ), - [0, 1, 2], + [0, 1], id="duplicate_names", ), pytest.param( @@ -1303,13 +1757,10 @@ def _types_frame(infer_string, parse_time=True): { "a": _string_series(["x", float("nan"), "y y"], infer), "b": pd.Series([1, None, 3], dtype="Int64"), - "c": pd.Series( - [pd.Timestamp("2024-01-01"), pd.NaT, pd.Timestamp("2024-01-02")], - dtype="datetime64[us]", - ), + "c": pd.Series(["2024-01-01", None, "2024-01-02"], dtype="string[pyarrow]"), } ), - [1, 2], + [1], id="tab_separated", ), ],