From 0967010828438a34c445c9615bfa40c327c02896 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 12:03:38 +0900 Subject: [PATCH 01/13] Give PandasCursor's temporal columns one dtype per Athena type PandasCursor reads timestamp, date, and timestamp with time zone columns as text and converts each column at once after reading: timestamp and date become datetime64[us], fractions beyond six digits truncated, and timestamp with time zone becomes datetime64[us, UTC], localized once per zone. The dtypes no longer depend on the values, the engine, chunking, or managed results, and UNLOAD date columns are datetime64[us] too. pyarrow-backed text is parsed with pyarrow compute, so no Python object is made per value. The fetch methods return datetime, date, aware datetime in UTC, and None for NULL for the columns that PyAthena converts; columns that a caller's parse_dates made keep pandas' values. Time zone errors raise DataError naming the column, in whole and chunked reads. A caller's parse_dates, dtype or converters entry, or converter mapping still replaces PyAthena's conversion. Closes #1219 Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 401 +++++++++++++++++++++-- tests/pyathena/pandas/test_cursor.py | 36 +- tests/pyathena/pandas/test_result_set.py | 317 ++++++++++++++++-- 3 files changed, 673 insertions(+), 81 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 30348f9c6..276d82601 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 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 @@ -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) @@ -287,6 +294,167 @@ def _convert_values(self, values: Series | Index, index: Index | None = None) -> 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 ``timestamp with time zone`` value has an unknown zone or a + local time without one instant in its zone. The message names the + column. + """ + import pandas as pd + + try: + 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) + 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(pd.Series(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(pd.Series(index)), name=name) + except DataError 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. + codes, zone_names = pd.factorize(text.str.replace(r"^.* ", "", regex=True)) + utc = np.full(len(values), np.datetime64("NaT", "us")) + for code, zone_name in enumerate(zone_names): + mask = codes == code + zone = _parse_time_zone(zone_name) + if isinstance(zone, timezone): + offset = np.timedelta64(zone.utcoffset(None), "us") + utc[mask] = local[mask] - offset + continue + try: + localized = pd.Series(local[mask]).dt.tz_localize( + zone, ambiguous="raise", nonexistent="raise" + ) + except ValueError: + # Raise the DataError of the first value without one instant. + for value in text.to_numpy(dtype=object)[mask]: + _to_datetime_with_tz(value) + raise + utc[mask] = ( + localized.dt.tz_convert("UTC").dt.tz_localize(None).to_numpy(dtype="datetime64[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 +489,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 +498,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 +510,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 +618,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 +632,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 +913,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 +929,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 +997,12 @@ def _get_csv_engine( is_compatible = ( effective_chunksize is None and self._quoting == 1 - and not self.converters + # The default timestamp with time zone conversion runs after reading. + and not { + name + for name, converter in self.converters.items() + if converter is not _to_datetime_with_tz + } # 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 +1138,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 +1148,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 +1157,96 @@ 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 options given to ``execute()`` or the converter + decide how they are read: ``parse_dates`` (for ``timestamp`` and ``date``), a + ``dtype`` that names the column or is not a plain mapping, ``converters`` that + name the column (or any ``converters``, for ``timestamp with time zone``), or a + conversion function that the converter maps to the type, 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. + """ + dtype = self._kwargs.get("dtype") + converters = self._kwargs.get("converters") + 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 or "converters" in self._kwargs: + continue + elif type_ in mappings or "parse_dates" in self._kwargs: + continue + if "dtype" in self._kwargs and (type(dtype) is not dict or label in dtype): + continue + if converters and label in converters: + 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 +1479,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 +1501,24 @@ 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 are read as text, unless a dtype given to execute() + # replaces PyAthena's, with which pandas parses them first. + temporal = {} if "dtype" in self._kwargs else self._csv_temporal_converters read_csv_kwargs: dict[str, Any] = { "sep": sep, "header": header, "names": names, - "dtype": self.dtypes, + "dtype": {**self.dtypes, **dict.fromkeys(temporal, "str")}, "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 +1661,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 +1672,27 @@ 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) + # See _get_csv_read_options(). + temporal = {} if "dtype" in self._kwargs else 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 +2126,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..b4f7144a7 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() @@ -1035,7 +1028,7 @@ def test_complex_unload_as_pandas_pyarrow(self, pandas_cursor, parquet_engine): str, np.datetime64, np.object_, - 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..562f25aba 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -11,7 +11,7 @@ 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 import boto3 @@ -24,7 +24,7 @@ from botocore.response import StreamingBody from pandas.testing import assert_frame_equal -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 +35,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, @@ -888,6 +889,229 @@ def test_read_csv_json_custom_converters(self, tmp_path, converters, expected): ) assert df["j"].tolist() == expected + @pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"engine": "c"}, id="c"), + pytest.param({"engine": "python"}, id="python"), + pytest.param({"engine": "pyarrow"}, id="pyarrow"), + # An option that PyAthena's PyArrow reader does not take: pandas reads the file. + pytest.param({"engine": "pyarrow", "encoding": "utf-8"}, id="pandas-pyarrow"), + pytest.param({"chunksize": 1}, id="chunks"), + ], + ) + @pytest.mark.parametrize("infer_string", [True, False]) + def test_temporal_types(self, tmp_path, kwargs, 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. + """ + with pd.option_context("future.infer_string", infer_string): + self._check_temporal_types(tmp_path, kwargs) + + @staticmethod + def _check_temporal_types(tmp_path, kwargs): + """Check the dtypes, values, and fetched values of the temporal corpus. + + Args: + tmp_path: The directory of the result files. + kwargs: The result set's options. + """ + columns = temporal.columns() + 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() + assert row == temporal.fetched_row() + assert [type(v) for v in row] == [type(v) for v in temporal.fetched_row()] + + @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_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 +1273,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 +1315,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 +1324,48 @@ 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 text, as PyAthena reads + them with its default options before converting them. + """ missing = float("nan") return pd.DataFrame( { @@ -1120,12 +1381,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": ( + _string_series(["2024-02-29", missing], infer_string) + if temporal_text + else pd.Series([pd.Timestamp("2024-02-29"), pd.NaT], dtype="datetime64[us]") ), - "ts6": pd.Series( - [pd.Timestamp("2024-02-29 23:59:58.123456"), pd.NaT], dtype="datetime64[ns]" + "ts": ( + _string_series(["2024-02-29 23:59:58.123", missing], infer_string) + if temporal_text + else pd.Series( + [pd.Timestamp("2024-02-29 23:59:58.123"), pd.NaT], dtype="datetime64[ns]" + ) + ), + "ts6": ( + _string_series(["2024-02-29 23:59:58.123456", missing], infer_string) + 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 +1424,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 +1476,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]" - ), + _string_series(["2024-01-01", float("nan")], infer).rename("d"), ], axis=1, ), - [0, 1, 2], + [0, 1], id="duplicate_names", ), pytest.param( @@ -1303,13 +1575,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": _string_series(["2024-01-01", float("nan"), "2024-01-02"], infer), } ), - [1, 2], + [1], id="tab_separated", ), ], From b752a348a10c1b2615cf34e49eeffa780d7df721 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 12:03:38 +0900 Subject: [PATCH 02/13] Describe PandasCursor's temporal dtypes and fetched types Refs #1219 Co-Authored-By: Claude Opus 5.5 --- docs/pandas.md | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/docs/pandas.md b/docs/pandas.md index 88a16b404..1b642f2d5 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 `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. @@ -566,7 +568,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`, and `time with time zone`), 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 +577,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`, a `dtype` or `converters` entry that names a `date`, `timestamp`, or `timestamp with time zone` column, or a conversion function that the converter maps to one of these types, pandas reads the column as these options say, and the fetch methods return its values as pandas holds them. (pandas-s3-read-options)= @@ -604,7 +607,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 +772,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 `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 +786,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. From b765f4d810f369d8f99cc0063b99c8b192b67c2a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 12:15:10 +0900 Subject: [PATCH 03/13] Convert early local times, group zones once, and share the label dispatch A named zone's local times before its first transition, which pandas reports as nonexistent, are converted one at a time, and only a value without one instant raises DataError. The values of each zone are grouped once by a stable sort of the zone codes instead of a mask per zone, and any conversion failure raises DataError naming the column. The json and temporal converters find a column or an index level by the same helper. The PyArrow engine is chosen only when converters given to execute() and the converters of the columns not converted after reading would not be needed, so the results of a read do not depend on the engine. The fetch methods convert only NumPy-backed datetime columns, as pyarrow- backed ones keep their values. parse_dates given to execute() keeps PyAthena from converting any temporal column, and a dtype or converters entry keeps only the column it names. Refs #1219 Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 137 ++++++++++++++++++---------------- 1 file changed, 71 insertions(+), 66 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 276d82601..ca9030bbd 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -11,7 +11,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack from contextvars import copy_context -from datetime import timezone +from datetime import datetime, timezone from functools import partial from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper from itertools import chain, islice @@ -38,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 @@ -202,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. @@ -242,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 @@ -274,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 @@ -288,7 +299,7 @@ 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() @@ -328,27 +339,13 @@ def convert(self, df: DataFrame, name: Any) -> None: name: The label of the column, which can also be an index level. Raises: - DataError: If a ``timestamp with time zone`` value has an unknown zone or a - local time without one instant in its zone. The message names the - column. + 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. """ - import pandas as pd - try: - 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) - 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(pd.Series(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(pd.Series(index)), name=name) - except DataError as e: + _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: @@ -430,26 +427,31 @@ def _to_utc_timestamps(values: Series) -> Series: local = _to_timestamps(text.str.replace(r" [^ ]*$", "", regex=True)).to_numpy( dtype="datetime64[us]" ) - # Integer codes of the zones, -1 for NULL. + # 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): - mask = codes == code + positions = order[bounds[code + 1] : bounds[code + 2]] zone = _parse_time_zone(zone_name) if isinstance(zone, timezone): - offset = np.timedelta64(zone.utcoffset(None), "us") - utc[mask] = local[mask] - offset + utc[positions] = local[positions] - np.timedelta64(zone.utcoffset(None), "us") continue try: - localized = pd.Series(local[mask]).dt.tz_localize( + localized = pd.Series(local[positions]).dt.tz_localize( zone, ambiguous="raise", nonexistent="raise" ) except ValueError: - # Raise the DataError of the first value without one instant. - for value in text.to_numpy(dtype=object)[mask]: - _to_datetime_with_tz(value) - raise - utc[mask] = ( + # pandas reports local times before a zone's first transition as + # nonexistent, so the values are converted one at a time, which raises the + # DataError of a value without one instant. + texts = text.to_numpy(dtype=object)[positions] + offsets = [cast(datetime, _to_datetime_with_tz(value)).utcoffset() for value in texts] + utc[positions] = local[positions] - np.array(offsets, dtype="timedelta64[us]") + continue + utc[positions] = ( localized.dt.tz_convert("UTC").dt.tz_localize(None).to_numpy(dtype="datetime64[us]") ) return pd.Series(utc, index=values.index).dt.tz_localize("UTC") @@ -639,7 +641,7 @@ def _iter_row_blocks(self, df: DataFrame) -> Iterator[Iterator[tuple[Any, ...]]] series = df.iloc[:, i] array = series.array fetch_type = self._fetch_types.get(df.columns[i]) - if fetch_type is not None and series.dtype.kind == "M": + if fetch_type is not None and isinstance(array, pd.arrays.DatetimeArray): columns.append((array, fetch_type)) elif isinstance(array.dtype, pd.StringDtype): columns.append((array, "list")) @@ -997,12 +999,13 @@ def _get_csv_engine( is_compatible = ( effective_chunksize is None and self._quoting == 1 - # The default timestamp with time zone conversion runs after reading. - and not { - name - for name, converter in self.converters.items() - if converter is not _to_datetime_with_tz - } + # pandas' PyArrow engine takes no converters. Those of the temporal + # columns that are converted after reading are not used. + and "converters" not in self._kwargs + 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()). @@ -1215,11 +1218,11 @@ def _get_csv_temporal_converters( PyAthena converts the ``timestamp``, ``date``, and ``timestamp with time zone`` columns after reading, unless options given to ``execute()`` or the converter - decide how they are read: ``parse_dates`` (for ``timestamp`` and ``date``), a - ``dtype`` that names the column or is not a plain mapping, ``converters`` that - name the column (or any ``converters``, for ``timestamp with time zone``), or a - conversion function that the converter maps to the type, other than the - default ``timestamp with time zone`` one. + decide how they are read: ``parse_dates``, with which PyAthena converts none of + them, a ``dtype`` that is not a plain mapping, which applies to every column, a + ``dtype`` or ``converters`` entry that names the column, or a conversion + function that the converter maps to the type, other than the default + ``timestamp with time zone`` one. Args: columns: The label and the description of each column that is read. @@ -1227,8 +1230,10 @@ def _get_csv_temporal_converters( Returns: The converters keyed by column label. """ + if "parse_dates" in self._kwargs: + return {} dtype = self._kwargs.get("dtype") - converters = self._kwargs.get("converters") + converters = self._kwargs.get("converters") or {} mappings = self._converter.mappings temporal_converters = {} for label, d in columns: @@ -1236,13 +1241,13 @@ def _get_csv_temporal_converters( if type_ not in _CSV_TEMPORAL_TYPES: continue if type_ == "timestamp with time zone": - if mappings.get(type_) is not _to_datetime_with_tz or "converters" in self._kwargs: + if mappings.get(type_) is not _to_datetime_with_tz: continue - elif type_ in mappings or "parse_dates" in self._kwargs: + elif type_ in mappings: continue if "dtype" in self._kwargs and (type(dtype) is not dict or label in dtype): continue - if converters and label in converters: + if label in converters: continue temporal_converters[label] = _TemporalConverter(type_) return temporal_converters From d9a4ba08c67cda93409ce16e29da4dbcdfc6a1c7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 12:19:11 +0900 Subject: [PATCH 04/13] Test early local times, conversion errors, engine-independent time zone columns, and options The temporal tests cover a named zone's local times before its first transition, a value that fails to convert in each temporal type, time zone columns read with engine="pyarrow" as with the C engine under converters or dtype, and parse_dates and converters for time zone columns. The fetch methods again check the datetime64 dtype of a column, as every converted column is NumPy-backed. Refs #1219 Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 2 +- tests/pyathena/pandas/test_result_set.py | 79 ++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index ca9030bbd..2a0acaca2 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -641,7 +641,7 @@ def _iter_row_blocks(self, df: DataFrame) -> Iterator[Iterator[tuple[Any, ...]]] series = df.iloc[:, i] array = series.array fetch_type = self._fetch_types.get(df.columns[i]) - if fetch_type is not None and isinstance(array, pd.arrays.DatetimeArray): + 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")) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 562f25aba..0b59f6cb1 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1100,6 +1100,85 @@ def test_temporal_options_take_precedence(self, tmp_path, kwargs, converter_type 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. + """ + data = ( + '"tstz"\n' + '"1500-01-01 00: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 + ) + values = result_set.as_pandas()["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", "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.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( + ("kwargs", "expected"), + [ + # parse_dates keeps PyAthena from converting any temporal column: the + # converter's function gives the value in its zone. + pytest.param({"parse_dates": ["ts"]}, "+09:00", id="parse_dates"), + # converters that do not name the column keep its conversion, to UTC. + pytest.param({"converters": {}}, "+00:00", id="converters"), + pytest.param({"converters": {"ts": str}}, "+00:00", id="converters_of_other"), + ], + ) + def test_time_zone_options_take_precedence(self, tmp_path, kwargs, expected): + """Options apply to timestamp with time zone columns as to the other temporal ones.""" + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("ts", "timestamp"), ("tstz", "timestamp with time zone")], + '"ts","tstz"\n"2020-01-02 03:04:05.000","2020-01-02 03:04:05.000 Asia/Tokyo"\n', + **kwargs, + ) + assert result_set.fetchone()[1].isoformat().endswith(expected) + def test_custom_time_zone_converter(self, tmp_path): """A conversion function that the converter maps to timestamp with time zone wins.""" converter = DefaultPandasTypeConverter() From 244ae20444085df686ea7aa4f006e993b9d36a33 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 12:20:59 +0900 Subject: [PATCH 05/13] Describe the temporal conversion precedence, errors, and iterator rows Refs #1219 Co-Authored-By: Claude Opus 5.5 --- docs/pandas.md | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/docs/pandas.md b/docs/pandas.md index 1b642f2d5..4050d9701 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -222,7 +222,7 @@ for row in cursor: 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 `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. +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 @@ -392,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 @@ -568,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`, and `time with time zone`), 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. @@ -577,7 +578,14 @@ 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`, a `dtype` or `converters` entry that names a `date`, `timestamp`, or `timestamp with time zone` column, or a conversion function that the converter maps to one of these types, pandas reads the column as these options say, and the fetch methods return its values as pandas holds them. +PyAthena converts `date`, `timestamp`, and `timestamp with time zone` columns after reading, except for the columns that these options decide: + +- With `parse_dates`, PyAthena converts none of them. +- A `dtype` that is not a mapping applies to every column. +- A `dtype` or `converters` entry that names a column applies to that column only. +- A conversion function that the converter maps to one of these types, other than the default `timestamp with time zone` one, applies to the columns of that type. + +pandas reads these columns as the options say, and the fetch methods return their values as pandas holds them. (pandas-s3-read-options)= @@ -774,7 +782,7 @@ print(df.head()) 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 `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. +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 From 7afc37d1240bcc6c67e28bfc45919e4362f37a6e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 12:26:19 +0900 Subject: [PATCH 06/13] Ignore pandas' converter-and-dtype warning in the engine test Refs #1219 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_result_set.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 0b59f6cb1..ac4f15925 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1136,6 +1136,7 @@ def test_temporal_conversion_error(self, tmp_path, type_, infer_string): _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", [ From acf5baafac864d2f76137a12fb68d8ddd6bc7f9e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 12:43:15 +0900 Subject: [PATCH 07/13] Keep NULL in every engine, take positional and mapping options, and localize zones together The PyArrow engine reads the temporal text as pyarrow strings, as pandas' PyArrow engine turns NULL into "None" with the NumPy str dtype. dtype and converters entries name a column by its label or its position, as in pandas, any mapping other than a defaultdict names its columns, and a defaultdict applies to every column. The values of a zone are localized together, and only those that pandas does not localize, such as local times before the zone's first transition, are converted one at a time. The PyArrow cases of the temporal types test leave out the time with time zone columns, so they read with PyAthena's and pandas' PyArrow readers, which the test checks. Refs #1219 Co-Authored-By: Claude Opus 5.5 --- docs/pandas.md | 4 +- pyathena/pandas/result_set.py | 72 +++++++----- tests/pyathena/pandas/test_result_set.py | 142 ++++++++++++++++++----- 3 files changed, 160 insertions(+), 58 deletions(-) diff --git a/docs/pandas.md b/docs/pandas.md index 4050d9701..9ea2896bc 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -581,8 +581,8 @@ For example, `dtype` replaces the whole column type mapping, and `parse_dates` r PyAthena converts `date`, `timestamp`, and `timestamp with time zone` columns after reading, except for the columns that these options decide: - With `parse_dates`, PyAthena converts none of them. -- A `dtype` that is not a mapping applies to every column. -- A `dtype` or `converters` entry that names a column applies to that column only. +- A `dtype` that is not a mapping, or a `defaultdict`, applies to every column. +- A `dtype` or `converters` entry that names a column, by its label or its position, applies to that column only. - A conversion function that the converter maps to one of these types, other than the default `timestamp with time zone` one, applies to the columns of that type. pandas reads these columns as the options say, and the fetch methods return their values as pandas holds them. diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 2a0acaca2..9ea291143 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -6,7 +6,7 @@ import importlib import logging import sys -from collections import abc +from collections import abc, defaultdict from collections.abc import Callable, Iterable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack @@ -439,21 +439,22 @@ def _to_utc_timestamps(values: Series) -> Series: if isinstance(zone, timezone): utc[positions] = local[positions] - np.timedelta64(zone.utcoffset(None), "us") continue - try: - localized = pd.Series(local[positions]).dt.tz_localize( - zone, ambiguous="raise", nonexistent="raise" - ) - except ValueError: - # pandas reports local times before a zone's first transition as - # nonexistent, so the values are converted one at a time, which raises the - # DataError of a value without one instant. - texts = text.to_numpy(dtype=object)[positions] - offsets = [cast(datetime, _to_datetime_with_tz(value)).utcoffset() for value in texts] - utc[positions] = local[positions] - np.array(offsets, dtype="timedelta64[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") @@ -1004,7 +1005,7 @@ def _get_csv_engine( and "converters" not in self._kwargs and not self.converters.keys() - self._get_csv_temporal_converters( - [(d[0], d) for d in self.description or []] + [(i, d[0], d) for i, d in enumerate(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 @@ -1212,20 +1213,21 @@ def _get_fetch_types(self) -> dict[Any, str]: } def _get_csv_temporal_converters( - self, columns: list[tuple[Any, tuple[Any, ...]]] + self, columns: list[tuple[int, 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 options given to ``execute()`` or the converter decide how they are read: ``parse_dates``, with which PyAthena converts none of - them, a ``dtype`` that is not a plain mapping, which applies to every column, a - ``dtype`` or ``converters`` entry that names the column, or a conversion - function that the converter maps to the type, other than the default - ``timestamp with time zone`` one. + them, a ``dtype`` that is not a mapping or is a ``defaultdict``, which applies + to every column, a ``dtype`` or ``converters`` entry that names the column by + its label or its position, or a conversion function that the converter maps to + the type, other than the default ``timestamp with time zone`` one. Args: - columns: The label and the description of each column that is read. + columns: The position in the result file, the label, and the description + of each column that is read. Returns: The converters keyed by column label. @@ -1236,7 +1238,7 @@ def _get_csv_temporal_converters( converters = self._kwargs.get("converters") or {} mappings = self._converter.mappings temporal_converters = {} - for label, d in columns: + for position, label, d in columns: type_ = d[1] if type_ not in _CSV_TEMPORAL_TYPES: continue @@ -1245,9 +1247,16 @@ def _get_csv_temporal_converters( continue elif type_ in mappings: continue - if "dtype" in self._kwargs and (type(dtype) is not dict or label in dtype): + if "dtype" in self._kwargs and ( + # pandas applies the default of a defaultdict to the columns it does not + # name. + not isinstance(dtype, Mapping) + or isinstance(dtype, defaultdict) + or label in dtype + or position in dtype + ): continue - if label in converters: + if label in converters or position in converters: continue temporal_converters[label] = _TemporalConverter(type_) return temporal_converters @@ -1508,16 +1517,19 @@ def _get_csv_read_options(self, csv_engine: str, chunksize: int | None) -> dict[ description = self.description or [] self._csv_temporal_converters = self._get_csv_temporal_converters( - [(d[0], d) for d in description] + [(i, d[0], d) for i, d in enumerate(description)] ) # The temporal columns are read as text, unless a dtype given to execute() - # replaces PyAthena's, with which pandas parses them first. + # replaces PyAthena's, with which pandas parses them first. 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 = {} if "dtype" in self._kwargs else 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, **dict.fromkeys(temporal, "str")}, + "dtype": {**self.dtypes, **dict.fromkeys(temporal, text_dtype)}, "converters": { d[0]: self._get_csv_converter(d[1]) for d in description @@ -1677,7 +1689,13 @@ 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) + self._csv_temporal_converters = self._get_csv_temporal_converters( + [ + (i, label, d) + for i, (label, d) in enumerate(zip(labels, description, strict=True)) + if label is not None + ] + ) # See _get_csv_read_options(). temporal = {} if "dtype" in self._kwargs else self._csv_temporal_converters if "dtype" not in self._kwargs: diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index ac4f15925..d7a1f5333 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -10,7 +10,7 @@ import io import re import threading -from collections import defaultdict +from collections import OrderedDict, defaultdict from datetime import UTC, date, datetime from unittest.mock import MagicMock, PropertyMock, patch @@ -23,7 +23,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.converter import _to_datetime_with_tz from pyathena.error import DataError, OperationalError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution @@ -890,37 +892,52 @@ def test_read_csv_json_custom_converters(self, tmp_path, converters, expected): assert df["j"].tolist() == expected @pytest.mark.parametrize( - "kwargs", + ("kwargs", "reader"), [ - pytest.param({"engine": "c"}, id="c"), - pytest.param({"engine": "python"}, id="python"), - pytest.param({"engine": "pyarrow"}, id="pyarrow"), + 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"}, id="pandas-pyarrow"), - pytest.param({"chunksize": 1}, id="chunks"), + 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, infer_string): + 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. + without it, as objects. The PyArrow engine reads results without the ``time + with time zone`` columns, whose converter it does not take. """ - with pd.option_context("future.infer_string", infer_string): - self._check_temporal_types(tmp_path, kwargs) + 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, kwargs): - """Check the dtypes, values, and fetched values of the temporal corpus. + 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. """ - columns = temporal.columns() data = _temporal_csv(columns) assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES types = [(c.name, c.athena_type) for c in columns] @@ -943,8 +960,9 @@ def _check_temporal_types(tmp_path, kwargs): LocalBucketFileSystem(tmp_path / "rows"), types, data, **kwargs ) row = result_set.fetchone() - assert row == temporal.fetched_row() - assert [type(v) for v in row] == [type(v) for v in temporal.fetched_row()] + 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]) @@ -1100,25 +1118,90 @@ def test_temporal_options_take_precedence(self, tmp_path, kwargs, converter_type assert result_set.fetchone() == (expected,) assert type(result_set.as_pandas()["ts"].iloc[0]) is type(expected) + @pytest.mark.parametrize( + ("kwargs", "expected"), + [ + # pandas takes positions as converters and dtype keys. + pytest.param( + {"converters": {1: lambda value: "converted"}}, + (datetime(2020, 1, 2, 3, 4, 5, 123456), "converted"), + id="converters_position", + ), + # pandas parses the dates of PyAthena's parse_dates after the dtype. + pytest.param( + {"dtype": {1: "str"}}, + ( + datetime(2020, 1, 2, 3, 4, 5, 123456), + pd.Timestamp("2020-01-02 03:04:05.123456789"), + ), + id="dtype_position", + ), + pytest.param( + {"dtype": OrderedDict(n="Int64")}, + (datetime(2020, 1, 2, 3, 4, 5, 123456),) * 2, + id="dtype_mapping", + ), + # pandas applies the default to the columns that a defaultdict does not name. + pytest.param( + {"dtype": defaultdict(lambda: "str", n="Int64")}, + (pd.Timestamp("2020-01-02 03:04:05.123456789"),) * 2, + id="dtype_defaultdict", + ), + ], + ) + def test_temporal_options_by_position_and_mapping(self, tmp_path, kwargs, expected): + """dtype and converters entries name columns by label or position, as in pandas. + + A dtype mapping that names no temporal column keeps their conversion, and a + defaultdict applies to every column. The columns have the same name. + """ + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("ts", "timestamp"), ("ts", "timestamp")], + '"ts","ts"\n"2020-01-02 03:04:05.123456789","2020-01-02 03:04:05.123456789"\n', + **kwargs, + ) + assert result_set.fetchone() == 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. + instant. Only these values are converted one at a time, by the converter's + function, and the other values of their zones are localized together. """ 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 - ) - values = result_set.as_pandas()["tstz"].dt.tz_localize(None).to_numpy() + # The spy is the converter's function, which keeps the default conversion. + spy = MagicMock(wraps=_to_datetime_with_tz) + converter = DefaultPandasTypeConverter() + converter.set("timestamp with time zone", spy) + with patch("pyathena.pandas.result_set._to_datetime_with_tz", spy): + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("tstz", "timestamp with time zone")], + data, + converter=converter, + ) + df = result_set.as_pandas() + 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", "0000-12-31T14:41:01", "2020-01-01T18:04:05"] + 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]")) + assert [c.args for c in spy.call_args_list] == [ + ("1500-01-01 00:00:00.000 America/New_York",), + ("0001-01-01 00:00:00.000 Asia/Tokyo",), + ] @pytest.mark.parametrize("infer_string", [True, False]) @pytest.mark.parametrize("type_", ["timestamp", "date", "timestamp with time zone"]) @@ -1443,8 +1526,9 @@ def _string_series(values, infer_string): 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 text, as PyAthena reads - them with its default options before converting them. + 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( @@ -1462,19 +1546,19 @@ def _types_frame(infer_string, parse_time=True, temporal_text=False): "m": _string_series(["{k=1}", missing], infer_string), "rw": _string_series(["{a=1, b=x}", missing], infer_string), "dt": ( - _string_series(["2024-02-29", missing], infer_string) + 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": ( - _string_series(["2024-02-29 23:59:58.123", missing], infer_string) + 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": ( - _string_series(["2024-02-29 23:59:58.123456", missing], infer_string) + 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]" @@ -1556,7 +1640,7 @@ def _types_frame(infer_string, parse_time=True, temporal_text=False): [ pd.Series([1, None], name="x", dtype="Int64"), pd.Series([2, None], name="x", dtype="Int64"), - _string_series(["2024-01-01", float("nan")], infer).rename("d"), + pd.Series(["2024-01-01", None], name="d", dtype="string[pyarrow]"), ], axis=1, ), @@ -1655,7 +1739,7 @@ def _types_frame(infer_string, parse_time=True, temporal_text=False): { "a": _string_series(["x", float("nan"), "y y"], infer), "b": pd.Series([1, None, 3], dtype="Int64"), - "c": _string_series(["2024-01-01", float("nan"), "2024-01-02"], infer), + "c": pd.Series(["2024-01-01", None, "2024-01-02"], dtype="string[pyarrow]"), } ), [1], From 3da05c4e237f84ab7f7348f143c53334b8fb8a49 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 13:08:44 +0900 Subject: [PATCH 08/13] Resolve dtype and converters keys of temporal columns as pandas does A dtype or converters key names the column with that label, and a dtype key also the columns that pandas renamed from duplicates of that name. The C engine also takes an integer key as a position in the file, the Python engine, if no column has that label, as a position among the columns read, and the PyArrow engine takes no positions. The test of local times before a zone's first transition checks the values and the dtype only. Refs #1219 Co-Authored-By: Claude Opus 5.5 --- docs/pandas.md | 2 +- pyathena/pandas/result_set.py | 57 ++++++++-- tests/pyathena/pandas/test_result_set.py | 130 ++++++++++++++++------- 3 files changed, 138 insertions(+), 51 deletions(-) diff --git a/docs/pandas.md b/docs/pandas.md index 9ea2896bc..8ae759de8 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -582,7 +582,7 @@ PyAthena converts `date`, `timestamp`, and `timestamp with time zone` columns af - With `parse_dates`, PyAthena converts none of them. - A `dtype` that is not a mapping, or a `defaultdict`, applies to every column. -- A `dtype` or `converters` entry that names a column, by its label or its position, applies to that column only. +- A `dtype` or `converters` entry applies to the columns that pandas applies it to: the column with its label, and for `dtype`, the columns that pandas renames from duplicates of that name, such as `ts.1`. The C engine also takes an integer key as a position in the file, and the Python engine, if no column has that label, as a position among the columns read. - A conversion function that the converter maps to one of these types, other than the default `timestamp with time zone` one, applies to the columns of that type. pandas reads these columns as the options say, and the fetch methods return their values as pandas holds them. diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 9ea291143..88e1fa2f3 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -1005,7 +1005,7 @@ def _get_csv_engine( and "converters" not in self._kwargs and not self.converters.keys() - self._get_csv_temporal_converters( - [(i, d[0], d) for i, d in enumerate(self.description or [])] + [(i, d[0], d) for i, d in enumerate(self.description or [])], "pyarrow" ).keys() # The pyarrow engine does not rename columns with the same name, and # the column labels cannot be resolved for it (see @@ -1213,7 +1213,7 @@ def _get_fetch_types(self) -> dict[Any, str]: } def _get_csv_temporal_converters( - self, columns: list[tuple[int, Any, tuple[Any, ...]]] + self, columns: list[tuple[int, Any, tuple[Any, ...]]], csv_engine: str ) -> dict[Any, _TemporalConverter]: """Get the converters of the temporal columns that are converted after reading. @@ -1221,13 +1221,21 @@ def _get_csv_temporal_converters( columns after reading, unless options given to ``execute()`` or the converter decide how they are read: ``parse_dates``, with which PyAthena converts none of them, a ``dtype`` that is not a mapping or is a ``defaultdict``, which applies - to every column, a ``dtype`` or ``converters`` entry that names the column by - its label or its position, or a conversion function that the converter maps to - the type, other than the default ``timestamp with time zone`` one. + to every column, a ``dtype`` or ``converters`` entry that pandas applies to the + column, or a conversion function that the converter maps to the type, other + than the default ``timestamp with time zone`` one. + + The entries apply to columns as in ``pandas.read_csv()``: a key names the + column with that label, and a ``dtype`` key also the columns that pandas + renamed from a duplicate of that name, such as ``ts.1``. The C engine also + takes an integer key as the position of a column in the file, and the Python + engine, if no column has that label, as its position among the columns that + are read. The PyArrow engine takes no positions. Args: columns: The position in the result file, the label, and the description - of each column that is read. + of each column that is read, in the order of the file. + csv_engine: The CSV engine that reads the file. Returns: The converters keyed by column label. @@ -1237,6 +1245,33 @@ def _get_csv_temporal_converters( dtype = self._kwargs.get("dtype") converters = self._kwargs.get("converters") or {} mappings = self._converter.mappings + labels = {label for _, label, _ in columns} + # Without names given to execute(), pandas renames duplicates of a name, and + # applies the name's dtype to them. + renamed = "names" not in self._kwargs + read_positions = {position: i for i, (position, _, _) in enumerate(columns)} + + def applies(options: Mapping[Any, Any], position: int, label: Any, name: Any) -> bool: + """Whether pandas applies an entry of the options to a column. + + Args: + options: The ``dtype`` or ``converters`` mapping. + position: The position of the column in the file. + label: The label of the column. + name: The name of the column, which a ``dtype`` key also names. + + Returns: + True if an entry names the column. + """ + if label in options or (options is dtype and renamed and name in options): + return True + if csv_engine == "c": + return position in options + if csv_engine == "python": + output_position = read_positions[position] + return output_position not in labels and output_position in options + return False + temporal_converters = {} for position, label, d in columns: type_ = d[1] @@ -1252,11 +1287,10 @@ def _get_csv_temporal_converters( # name. not isinstance(dtype, Mapping) or isinstance(dtype, defaultdict) - or label in dtype - or position in dtype + or applies(dtype, position, label, d[0]) ): continue - if label in converters or position in converters: + if applies(converters, position, label, d[0]): continue temporal_converters[label] = _TemporalConverter(type_) return temporal_converters @@ -1517,7 +1551,7 @@ def _get_csv_read_options(self, csv_engine: str, chunksize: int | None) -> dict[ description = self.description or [] self._csv_temporal_converters = self._get_csv_temporal_converters( - [(i, d[0], d) for i, d in enumerate(description)] + [(i, d[0], d) for i, d in enumerate(description)], csv_engine ) # The temporal columns are read as text, unless a dtype given to execute() # replaces PyAthena's, with which pandas parses them first. pandas' PyArrow @@ -1694,7 +1728,8 @@ def _key_csv_columns_by_labels( (i, label, d) for i, (label, d) in enumerate(zip(labels, description, strict=True)) if label is not None - ] + ], + read_csv_kwargs["engine"], ) # See _get_csv_read_options(). temporal = {} if "dtype" in self._kwargs else self._csv_temporal_converters diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index d7a1f5333..19af7b788 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -25,7 +25,6 @@ from pandas.testing import assert_frame_equal from pyarrow import csv as pyarrow_csv -from pyathena.converter import _to_datetime_with_tz from pyathena.error import DataError, OperationalError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution @@ -51,6 +50,12 @@ _TEXT_CSV = '"v"\n"a"\n"b"\n' +# A timestamp with nanoseconds, as PyAthena converts it and as pandas parses it. +_TIMESTAMP_TEXT = "2020-01-02 03:04:05.123456789" +_CONVERTED = datetime(2020, 1, 2, 3, 4, 5, 123456) +_PARSED = pd.Timestamp(_TIMESTAMP_TEXT) + + 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] @@ -1119,56 +1124,114 @@ def test_temporal_options_take_precedence(self, tmp_path, kwargs, converter_type assert type(result_set.as_pandas()["ts"].iloc[0]) is type(expected) @pytest.mark.parametrize( - ("kwargs", "expected"), + ("engine", "kwargs", "expected"), [ - # pandas takes positions as converters and dtype keys. + # The C engine takes an integer key as a position in the file. pytest.param( + "c", {"converters": {1: lambda value: "converted"}}, - (datetime(2020, 1, 2, 3, 4, 5, 123456), "converted"), - id="converters_position", + (_CONVERTED, "converted"), + id="c-converters_position", ), - # pandas parses the dates of PyAthena's parse_dates after the dtype. + pytest.param("c", {"dtype": {1: "str"}}, (_CONVERTED, _PARSED), id="c-dtype_position"), + # The Python engine takes it as a position among the columns that are read, + # and parses the dates of parse_dates without the column's converter. pytest.param( - {"dtype": {1: "str"}}, - ( - datetime(2020, 1, 2, 3, 4, 5, 123456), - pd.Timestamp("2020-01-02 03:04:05.123456789"), - ), - id="dtype_position", + "python", + {"converters": {1: lambda value: "converted"}}, + (_CONVERTED, _PARSED), + id="python-converters_position", + ), + # The C engine takes an integer key as a label and a position, the Python + # engine as a label only, if a column has that label. + pytest.param( + "c", + {"names": [1, 0], "dtype": {0: "str"}}, + (_PARSED, _PARSED), + id="c-dtype_label_and_position", + ), + pytest.param( + "python", + {"names": [1, 0], "dtype": {0: "str"}}, + (_CONVERTED, _PARSED), + id="python-dtype_label", ), + # A dtype key also names the columns that pandas renamed from its duplicates, + # and a converters key does not. pytest.param( + "c", {"dtype": OrderedDict(ts="str")}, (_PARSED, _PARSED), id="c-dtype_duplicates" + ), + pytest.param( + "python", + {"dtype": OrderedDict(ts="str")}, + (_PARSED, _PARSED), + id="python-dtype_duplicates", + ), + pytest.param( + "c", + {"converters": {"ts": lambda value: "converted"}}, + ("converted", _CONVERTED), + id="c-converters_duplicates", + ), + # A dtype mapping that names no temporal column keeps their conversion. + pytest.param( + "c", {"dtype": OrderedDict(n="Int64")}, - (datetime(2020, 1, 2, 3, 4, 5, 123456),) * 2, - id="dtype_mapping", + (_CONVERTED, _CONVERTED), + id="c-dtype_mapping", ), # pandas applies the default to the columns that a defaultdict does not name. pytest.param( + "c", {"dtype": defaultdict(lambda: "str", n="Int64")}, - (pd.Timestamp("2020-01-02 03:04:05.123456789"),) * 2, - id="dtype_defaultdict", + (_PARSED, _PARSED), + id="c-dtype_defaultdict", ), ], ) - def test_temporal_options_by_position_and_mapping(self, tmp_path, kwargs, expected): - """dtype and converters entries name columns by label or position, as in pandas. + def test_temporal_options_by_key(self, tmp_path, engine, kwargs, expected): + """dtype and converters entries apply to temporal columns as pandas applies them. - A dtype mapping that names no temporal column keeps their conversion, and a - defaultdict applies to every column. The columns have the same name. + pandas parses the dates of PyAthena's parse_dates after the dtype, without + PyAthena's truncation. The columns have the same name. """ result_set = _local_result_set( LocalBucketFileSystem(tmp_path), [("ts", "timestamp"), ("ts", "timestamp")], - '"ts","ts"\n"2020-01-02 03:04:05.123456789","2020-01-02 03:04:05.123456789"\n', + f'"ts","ts"\n"{_TIMESTAMP_TEXT}","{_TIMESTAMP_TEXT}"\n', + engine=engine, **kwargs, ) assert result_set.fetchone() == expected + @pytest.mark.parametrize( + ("engine", "position", "expected"), + [ + # The C engine takes positions in the file, where ts is the second column. + pytest.param("c", 1, "converted", id="c-file_position"), + pytest.param("c", 0, _CONVERTED, id="c-unread_position"), + # The Python engine takes positions among the columns that are read, and + # parses the dates of parse_dates without the column's converter. + pytest.param("python", 0, _PARSED, id="python-read_position"), + ], + ) + def test_temporal_converters_position_with_usecols(self, tmp_path, engine, position, expected): + """A converters position names a column of the file or of the columns read, by engine.""" + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("n", "integer"), ("ts", "timestamp")], + f'"n","ts"\n"1","{_TIMESTAMP_TEXT}"\n', + engine=engine, + usecols=["ts"], + converters={position: lambda value: "converted"}, + ) + assert result_set.fetchone() == (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. Only these values are converted one at a time, by the converter's - function, and the other values of their zones are localized together. + instant. The other values of their zones keep the zones' later offsets. """ data = ( '"tstz"\n' @@ -1177,18 +1240,11 @@ def test_time_zone_before_first_transition(self, tmp_path): '"0001-01-01 00:00:00.000 Asia/Tokyo"\n' '"2020-01-02 03:04:05.000 Asia/Tokyo"\n' ) - # The spy is the converter's function, which keeps the default conversion. - spy = MagicMock(wraps=_to_datetime_with_tz) - converter = DefaultPandasTypeConverter() - converter.set("timestamp with time zone", spy) - with patch("pyathena.pandas.result_set._to_datetime_with_tz", spy): - result_set = _local_result_set( - LocalBucketFileSystem(tmp_path), - [("tstz", "timestamp with time zone")], - data, - converter=converter, - ) - df = result_set.as_pandas() + 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 = [ @@ -1198,10 +1254,6 @@ def test_time_zone_before_first_transition(self, tmp_path): "2020-01-01T18:04:05", ] np.testing.assert_array_equal(values, np.array(expected, dtype="datetime64[us]")) - assert [c.args for c in spy.call_args_list] == [ - ("1500-01-01 00:00:00.000 America/New_York",), - ("0001-01-01 00:00:00.000 Asia/Tokyo",), - ] @pytest.mark.parametrize("infer_string", [True, False]) @pytest.mark.parametrize("type_", ["timestamp", "date", "timestamp with time zone"]) From bab3840ac390c748b639e33d27c9e7ee3430953d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 13:09:22 +0900 Subject: [PATCH 09/13] Test that the PyArrow engine takes no dtype positions of temporal columns Refs #1219 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_result_set.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 19af7b788..4b860e042 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1204,6 +1204,19 @@ def test_temporal_options_by_key(self, tmp_path, engine, kwargs, expected): ) assert result_set.fetchone() == expected + def test_temporal_dtype_position_with_pyarrow(self, tmp_path): + """The PyArrow engine takes no dtype positions, so the columns are converted.""" + data = '"a","b"\n' + f'"{_TIMESTAMP_TEXT}","{_TIMESTAMP_TEXT}"\n' * 3 + assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("a", "timestamp"), ("b", "timestamp")], + data, + engine="pyarrow", + dtype={1: "str"}, + ) + assert result_set.fetchone() == (_CONVERTED, _CONVERTED) + @pytest.mark.parametrize( ("engine", "position", "expected"), [ From 36b23f0d26a9a4bb35e7eace66f4bccc4e5fde17 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 13:39:15 +0900 Subject: [PATCH 10/13] Leave temporal columns to pandas with parse_dates, dtype, or converters With any of these options given to execute(), PyAthena converts no temporal column after reading: pandas reads them with PyAthena's parse_dates and converter mapping unless the options replace them, and the fetch methods return pandas' values. The resolution of dtype and converters keys and its tests are removed. Refs #1219 Co-Authored-By: Claude Opus 5.5 --- docs/pandas.md | 9 +- pyathena/pandas/result_set.py | 94 +++-------- tests/pyathena/pandas/test_result_set.py | 192 +++++++---------------- 3 files changed, 80 insertions(+), 215 deletions(-) diff --git a/docs/pandas.md b/docs/pandas.md index 8ae759de8..92dfd5cc4 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -578,14 +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. -PyAthena converts `date`, `timestamp`, and `timestamp with time zone` columns after reading, except for the columns that these options decide: - -- With `parse_dates`, PyAthena converts none of them. -- A `dtype` that is not a mapping, or a `defaultdict`, applies to every column. -- A `dtype` or `converters` entry applies to the columns that pandas applies it to: the column with its label, and for `dtype`, the columns that pandas renames from duplicates of that name, such as `ts.1`. The C engine also takes an integer key as a position in the file, and the Python engine, if no column has that label, as a position among the columns read. -- A conversion function that the converter maps to one of these types, other than the default `timestamp with time zone` one, applies to the columns of that type. - -pandas reads these columns as the options say, and the fetch methods return their values as pandas holds them. +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)= diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 88e1fa2f3..0667cb97b 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -6,7 +6,7 @@ import importlib import logging import sys -from collections import abc, defaultdict +from collections import abc from collections.abc import Callable, Iterable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack @@ -1005,7 +1005,7 @@ def _get_csv_engine( and "converters" not in self._kwargs and not self.converters.keys() - self._get_csv_temporal_converters( - [(i, d[0], d) for i, d in enumerate(self.description or [])], "pyarrow" + [(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 @@ -1213,67 +1213,27 @@ def _get_fetch_types(self) -> dict[Any, str]: } def _get_csv_temporal_converters( - self, columns: list[tuple[int, Any, tuple[Any, ...]]], csv_engine: str + 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 options given to ``execute()`` or the converter - decide how they are read: ``parse_dates``, with which PyAthena converts none of - them, a ``dtype`` that is not a mapping or is a ``defaultdict``, which applies - to every column, a ``dtype`` or ``converters`` entry that pandas applies to the - column, or a conversion function that the converter maps to the type, other - than the default ``timestamp with time zone`` one. - - The entries apply to columns as in ``pandas.read_csv()``: a key names the - column with that label, and a ``dtype`` key also the columns that pandas - renamed from a duplicate of that name, such as ``ts.1``. The C engine also - takes an integer key as the position of a column in the file, and the Python - engine, if no column has that label, as its position among the columns that - are read. The PyArrow engine takes no positions. + 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 position in the result file, the label, and the description - of each column that is read, in the order of the file. - csv_engine: The CSV engine that reads the file. + columns: The label and the description of each column that is read. Returns: The converters keyed by column label. """ - if "parse_dates" in self._kwargs: + if self._kwargs.keys() & {"parse_dates", "dtype", "converters"}: return {} - dtype = self._kwargs.get("dtype") - converters = self._kwargs.get("converters") or {} mappings = self._converter.mappings - labels = {label for _, label, _ in columns} - # Without names given to execute(), pandas renames duplicates of a name, and - # applies the name's dtype to them. - renamed = "names" not in self._kwargs - read_positions = {position: i for i, (position, _, _) in enumerate(columns)} - - def applies(options: Mapping[Any, Any], position: int, label: Any, name: Any) -> bool: - """Whether pandas applies an entry of the options to a column. - - Args: - options: The ``dtype`` or ``converters`` mapping. - position: The position of the column in the file. - label: The label of the column. - name: The name of the column, which a ``dtype`` key also names. - - Returns: - True if an entry names the column. - """ - if label in options or (options is dtype and renamed and name in options): - return True - if csv_engine == "c": - return position in options - if csv_engine == "python": - output_position = read_positions[position] - return output_position not in labels and output_position in options - return False - temporal_converters = {} - for position, label, d in columns: + for label, d in columns: type_ = d[1] if type_ not in _CSV_TEMPORAL_TYPES: continue @@ -1282,16 +1242,6 @@ def applies(options: Mapping[Any, Any], position: int, label: Any, name: Any) -> continue elif type_ in mappings: continue - if "dtype" in self._kwargs and ( - # pandas applies the default of a defaultdict to the columns it does not - # name. - not isinstance(dtype, Mapping) - or isinstance(dtype, defaultdict) - or applies(dtype, position, label, d[0]) - ): - continue - if applies(converters, position, label, d[0]): - continue temporal_converters[label] = _TemporalConverter(type_) return temporal_converters @@ -1551,13 +1501,13 @@ def _get_csv_read_options(self, csv_engine: str, chunksize: int | None) -> dict[ description = self.description or [] self._csv_temporal_converters = self._get_csv_temporal_converters( - [(i, d[0], d) for i, d in enumerate(description)], csv_engine + [(d[0], d) for d in description] ) - # The temporal columns are read as text, unless a dtype given to execute() - # replaces PyAthena's, with which pandas parses them first. 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 = {} if "dtype" in self._kwargs else self._csv_temporal_converters + # 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, @@ -1723,16 +1673,8 @@ 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( - [ - (i, label, d) - for i, (label, d) in enumerate(zip(labels, description, strict=True)) - if label is not None - ], - read_csv_kwargs["engine"], - ) - # See _get_csv_read_options(). - temporal = {} if "dtype" in self._kwargs else self._csv_temporal_converters + 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"] = { **{ diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 4b860e042..58e67edd6 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -10,9 +10,10 @@ import io import re import threading -from collections import OrderedDict, defaultdict +from collections import defaultdict from datetime import UTC, date, datetime from unittest.mock import MagicMock, PropertyMock, patch +from zoneinfo import ZoneInfo import boto3 import numpy as np @@ -50,10 +51,12 @@ _TEXT_CSV = '"v"\n"a"\n"b"\n' -# A timestamp with nanoseconds, as PyAthena converts it and as pandas parses it. +# Temporal values as pandas reads them with PyAthena's parse_dates and converters. _TIMESTAMP_TEXT = "2020-01-02 03:04:05.123456789" -_CONVERTED = datetime(2020, 1, 2, 3, 4, 5, 123456) _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): @@ -1123,123 +1126,6 @@ def test_temporal_options_take_precedence(self, tmp_path, kwargs, converter_type assert result_set.fetchone() == (expected,) assert type(result_set.as_pandas()["ts"].iloc[0]) is type(expected) - @pytest.mark.parametrize( - ("engine", "kwargs", "expected"), - [ - # The C engine takes an integer key as a position in the file. - pytest.param( - "c", - {"converters": {1: lambda value: "converted"}}, - (_CONVERTED, "converted"), - id="c-converters_position", - ), - pytest.param("c", {"dtype": {1: "str"}}, (_CONVERTED, _PARSED), id="c-dtype_position"), - # The Python engine takes it as a position among the columns that are read, - # and parses the dates of parse_dates without the column's converter. - pytest.param( - "python", - {"converters": {1: lambda value: "converted"}}, - (_CONVERTED, _PARSED), - id="python-converters_position", - ), - # The C engine takes an integer key as a label and a position, the Python - # engine as a label only, if a column has that label. - pytest.param( - "c", - {"names": [1, 0], "dtype": {0: "str"}}, - (_PARSED, _PARSED), - id="c-dtype_label_and_position", - ), - pytest.param( - "python", - {"names": [1, 0], "dtype": {0: "str"}}, - (_CONVERTED, _PARSED), - id="python-dtype_label", - ), - # A dtype key also names the columns that pandas renamed from its duplicates, - # and a converters key does not. - pytest.param( - "c", {"dtype": OrderedDict(ts="str")}, (_PARSED, _PARSED), id="c-dtype_duplicates" - ), - pytest.param( - "python", - {"dtype": OrderedDict(ts="str")}, - (_PARSED, _PARSED), - id="python-dtype_duplicates", - ), - pytest.param( - "c", - {"converters": {"ts": lambda value: "converted"}}, - ("converted", _CONVERTED), - id="c-converters_duplicates", - ), - # A dtype mapping that names no temporal column keeps their conversion. - pytest.param( - "c", - {"dtype": OrderedDict(n="Int64")}, - (_CONVERTED, _CONVERTED), - id="c-dtype_mapping", - ), - # pandas applies the default to the columns that a defaultdict does not name. - pytest.param( - "c", - {"dtype": defaultdict(lambda: "str", n="Int64")}, - (_PARSED, _PARSED), - id="c-dtype_defaultdict", - ), - ], - ) - def test_temporal_options_by_key(self, tmp_path, engine, kwargs, expected): - """dtype and converters entries apply to temporal columns as pandas applies them. - - pandas parses the dates of PyAthena's parse_dates after the dtype, without - PyAthena's truncation. The columns have the same name. - """ - result_set = _local_result_set( - LocalBucketFileSystem(tmp_path), - [("ts", "timestamp"), ("ts", "timestamp")], - f'"ts","ts"\n"{_TIMESTAMP_TEXT}","{_TIMESTAMP_TEXT}"\n', - engine=engine, - **kwargs, - ) - assert result_set.fetchone() == expected - - def test_temporal_dtype_position_with_pyarrow(self, tmp_path): - """The PyArrow engine takes no dtype positions, so the columns are converted.""" - data = '"a","b"\n' + f'"{_TIMESTAMP_TEXT}","{_TIMESTAMP_TEXT}"\n' * 3 - assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES - result_set = _local_result_set( - LocalBucketFileSystem(tmp_path), - [("a", "timestamp"), ("b", "timestamp")], - data, - engine="pyarrow", - dtype={1: "str"}, - ) - assert result_set.fetchone() == (_CONVERTED, _CONVERTED) - - @pytest.mark.parametrize( - ("engine", "position", "expected"), - [ - # The C engine takes positions in the file, where ts is the second column. - pytest.param("c", 1, "converted", id="c-file_position"), - pytest.param("c", 0, _CONVERTED, id="c-unread_position"), - # The Python engine takes positions among the columns that are read, and - # parses the dates of parse_dates without the column's converter. - pytest.param("python", 0, _PARSED, id="python-read_position"), - ], - ) - def test_temporal_converters_position_with_usecols(self, tmp_path, engine, position, expected): - """A converters position names a column of the file or of the columns read, by engine.""" - result_set = _local_result_set( - LocalBucketFileSystem(tmp_path), - [("n", "integer"), ("ts", "timestamp")], - f'"n","ts"\n"1","{_TIMESTAMP_TEXT}"\n', - engine=engine, - usecols=["ts"], - converters={position: lambda value: "converted"}, - ) - assert result_set.fetchone() == (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. @@ -1307,26 +1193,70 @@ def test_time_zone_results_do_not_depend_on_engine(self, tmp_path, kwargs): assert_frame_equal(frames[1], frames[0]) assert rows[1] == rows[0] + @pytest.mark.parametrize("engine", ["c", "python", "pyarrow"]) @pytest.mark.parametrize( ("kwargs", "expected"), [ - # parse_dates keeps PyAthena from converting any temporal column: the - # converter's function gives the value in its zone. - pytest.param({"parse_dates": ["ts"]}, "+09:00", id="parse_dates"), - # converters that do not name the column keep its conversion, to UTC. - pytest.param({"converters": {}}, "+00:00", id="converters"), - pytest.param({"converters": {"ts": str}}, "+00:00", id="converters_of_other"), + # 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_time_zone_options_take_precedence(self, tmp_path, kwargs, expected): - """Options apply to timestamp with time zone columns as to the other temporal ones.""" + 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"), ("tstz", "timestamp with time zone")], - '"ts","tstz"\n"2020-01-02 03:04:05.000","2020-01-02 03:04:05.000 Asia/Tokyo"\n', + [ + ("ts", "timestamp"), + ("d", "date"), + ("tstz", "timestamp with time zone"), + ("n", "integer"), + ], + data, + engine=engine, **kwargs, ) - assert result_set.fetchone()[1].isoformat().endswith(expected) + assert result_set.fetchone() == expected + + @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"), + ], + ) + 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.""" From b59362ff142b48b78a9cf7efcdaed61162616dd0 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 13:51:36 +0900 Subject: [PATCH 11/13] Treat a None parse_dates, dtype, or converters as not given pandas' defaults for these options are None, so a None passed to execute() keeps PyAthena's temporal conversion and the PyArrow engine. The options test also checks that pandas keeps the zone of the time zone values. Refs #1219 Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 5 +++-- tests/pyathena/pandas/test_result_set.py | 23 ++++++++++++++++++++++- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 0667cb97b..9cfc0cd3a 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -1002,7 +1002,7 @@ def _get_csv_engine( and self._quoting == 1 # pandas' PyArrow engine takes no converters. Those of the temporal # columns that are converted after reading are not used. - and "converters" not in self._kwargs + 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 []] @@ -1229,7 +1229,8 @@ def _get_csv_temporal_converters( Returns: The converters keyed by column label. """ - if self._kwargs.keys() & {"parse_dates", "dtype", "converters"}: + # 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 = {} diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 58e67edd6..423e701e4 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1193,6 +1193,22 @@ def test_time_zone_results_do_not_depend_on_engine(self, tmp_path, kwargs): 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"), @@ -1231,7 +1247,12 @@ def test_temporal_options_read_by_pandas(self, tmp_path, engine, kwargs, expecte engine=engine, **kwargs, ) - assert result_set.fetchone() == expected + 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", From d39163c3172ae9105cd828e7c92892297a15134a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 13:52:13 +0900 Subject: [PATCH 12/13] Check that converters of None keep the PyArrow engine Refs #1219 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_result_set.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 423e701e4..793c3cf7a 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1261,6 +1261,8 @@ def test_temporal_options_read_by_pandas(self, tmp_path, engine, kwargs, expecte # 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): From 3b62b8d651e927e3da58bf69837e4237d23345d5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 14:03:34 +0900 Subject: [PATCH 13/13] Expect datetime64 for the UNLOAD date column of the complex pyarrow test The expectation had marked the binary column instead of the date column. Refs #1219 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_cursor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index b4f7144a7..ae2cd7d98 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -1027,12 +1027,12 @@ 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_, np.object_, + np.object_, ) rows = [ (