From 3bdd65d9e0d61628da7e3341c970c2238c481f90 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 20 Oct 2026 00:17:32 +0900 Subject: [PATCH 1/3] Parse temporal text with fromisoformat and convert PandasCursor time columns after reading (GH-1308) The timestamp, date, time, and time zone converters parsed with datetime.strptime(), at about 3.5 microseconds a value. fromisoformat() parses Athena's text, truncating digits beyond microseconds, about 50 times faster. This speeds up the row output of every cursor with temporal columns. PandasCursor left time columns to parse_dates, for which pandas cannot infer a format and falls back to dateutil for each value. They are now read as text and converted after reading, as the other temporal columns are. This also fixes a parse_dates that does not name a time column, which raised AttributeError, and index_col naming a time column, which raised KeyError. Co-Authored-By: Claude Opus 5.5 --- docs/pandas.md | 2 +- pyathena/converter.py | 63 ++++++-------------- pyathena/pandas/result_set.py | 73 ++++++++++++++++++------ tests/pyathena/pandas/test_result_set.py | 49 ++++++++++++---- 4 files changed, 110 insertions(+), 77 deletions(-) diff --git a/docs/pandas.md b/docs/pandas.md index c8e81e78d..1ecaecb7e 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -584,7 +584,7 @@ It also changes the values of a column with a string dtype whose values all look Apart from PandasCursor's own options such as `engine` and `chunksize`, an option passed here replaces the value PyAthena sets for the same pandas.read_csv() argument. For example, `dtype` replaces the whole column type mapping, and `parse_dates` replaces the list of date, time, and timestamp columns. -With `parse_dates`, `dtype`, or `converters` passed here, PyAthena does not convert the `date`, `timestamp`, and `timestamp with time zone` columns after reading. pandas reads them with these options and the values PyAthena sets for the others: `parse_dates` for date, time, and timestamp columns, and the converter mapping, which converts `timestamp with time zone` values in their zones. The fetch methods return their values as pandas holds them. A conversion function that the converter maps to one of these types also replaces PyAthena's conversion of the columns of that type. +With `parse_dates`, `dtype`, or `converters` passed here, PyAthena does not convert the `date`, `time`, `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. `time` columns that pandas parses become `time` objects, and the fetch methods return the values of the other columns 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/converter.py b/pyathena/converter.py index 1d8c6df29..757360096 100644 --- a/pyathena/converter.py +++ b/pyathena/converter.py @@ -35,38 +35,22 @@ def _to_date(value: str | datetime | date | None) -> date | None: return value.date() if isinstance(value, date): return value - return datetime.strptime(value, "%Y-%m-%d").date() - - -def _parse_datetime(value: str) -> datetime: - """Parse an Athena TIMESTAMP value of any precision. - - Args: - value: The value as ``YYYY-MM-DD HH:MM:SS`` followed by an optional - fraction of up to 12 digits. - - Returns: - The datetime. Digits beyond microseconds are truncated. - """ - seconds, _, fraction = value.partition(".") - parsed = datetime.strptime(seconds, "%Y-%m-%d %H:%M:%S") - if fraction: - parsed = parsed.replace(microsecond=int(fraction[:6].ljust(6, "0"))) - return parsed + return date.fromisoformat(value) def _to_datetime(varchar_value: str | None) -> datetime | None: - """Convert an Athena TIMESTAMP value to a datetime. + """Convert an Athena TIMESTAMP value of any precision to a datetime. Args: - varchar_value: The value as text, or None. + varchar_value: The value as ``YYYY-MM-DD HH:MM:SS`` followed by an optional + fraction of up to 12 digits, or None. Returns: - The datetime, or None. + The datetime, or None. Digits beyond microseconds are truncated. """ if varchar_value is None: return None - return _parse_datetime(varchar_value) + return datetime.fromisoformat(varchar_value) # The length of Athena timestamp text, such as "2020-01-02 03:04:05.123456", that holds @@ -133,7 +117,7 @@ def _to_datetime_with_tz(varchar_value: str | None) -> datetime | None: for NULL, is None. Returns: - The aware datetime, or None. + The aware datetime, or None. Digits beyond microseconds are truncated. Raises: DataError: If the zone name is unknown, or the local time is ambiguous or @@ -143,7 +127,7 @@ def _to_datetime_with_tz(varchar_value: str | None) -> datetime | None: return None datetime_, _, tz = varchar_value.rpartition(" ") zone = _parse_time_zone(tz) - value = _parse_datetime(datetime_).replace(tzinfo=zone) + value = datetime.fromisoformat(datetime_).replace(tzinfo=zone) if isinstance(zone, timezone): return value # PEP 495: fold=0 and fold=1 give different offsets only at a transition. At a @@ -161,35 +145,19 @@ def _to_datetime_with_tz(varchar_value: str | None) -> datetime | None: return value -def _parse_time(value: str) -> time: - """Parse an Athena TIME value of any precision. - - Args: - value: The value as ``HH:MM:SS`` followed by an optional fraction of up to - 12 digits. - - Returns: - The time. Digits beyond microseconds are truncated. - """ - seconds, _, fraction = value.partition(".") - parsed = datetime.strptime(seconds, "%H:%M:%S").time() - if fraction: - parsed = parsed.replace(microsecond=int(fraction[:6].ljust(6, "0"))) - return parsed - - def _to_time(varchar_value: str | None) -> time | None: - """Convert an Athena TIME value to a time. + """Convert an Athena TIME value of any precision to a time. Args: - varchar_value: The value as text, or None. An empty string is None. + varchar_value: The value as ``HH:MM:SS`` followed by an optional fraction of + up to 12 digits, or None. An empty string is None. Returns: - The time, or None. + The time, or None. Digits beyond microseconds are truncated. """ if not varchar_value: return None - return _parse_time(varchar_value) + return time.fromisoformat(varchar_value) def _to_time_with_tz(varchar_value: str | None) -> time | None: @@ -201,12 +169,13 @@ def _to_time_with_tz(varchar_value: str | None) -> time | None: None. Returns: - The time with a fixed-offset ``tzinfo``, or None. + The time with a fixed-offset ``tzinfo``, or None. Digits beyond microseconds + are truncated. """ if not varchar_value: return None index = max(varchar_value.rfind("+"), varchar_value.rfind("-")) - return _parse_time(varchar_value[:index]).replace( + return time.fromisoformat(varchar_value[:index]).replace( tzinfo=_parse_utc_offset(varchar_value[index:]) ) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 81ebf702b..6de643732 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -30,7 +30,7 @@ from pyathena import OperationalError from pyathena.arrow.reader import ParquetChunkReader from pyathena.arrow.util import to_column_info -from pyathena.converter import Converter, _parse_time_zone, _to_datetime_with_tz +from pyathena.converter import Converter, _parse_time_zone, _to_datetime_with_tz, _to_time from pyathena.error import DataError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution @@ -57,7 +57,7 @@ # 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") +_CSV_TEMPORAL_TYPES = ("timestamp", "date", "timestamp with time zone", "time") # 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") @@ -314,9 +314,10 @@ def _convert_values(self, values: Series) -> Series: 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]``. + ``timestamp``, ``date``, ``timestamp with time zone``, and ``time`` 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]``, + ``datetime64[us, UTC]``, and ``object`` with ``time`` objects and None for NULL. 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. @@ -366,6 +367,8 @@ def _convert_values(self, values: Series) -> Series: """ if self._type == "timestamp with time zone": return _to_utc_timestamps(values) + if self._type == "time": + return _to_times(values) return _to_timestamps(values, date=self._type == "date") @@ -448,6 +451,27 @@ def _to_timestamps(values: Series, date: bool = False) -> Series: return pd.to_datetime(text, format="ISO8601").dt.as_unit("us") +def _to_times(values: Series) -> Series: + """Convert time text to ``time`` objects. + + pandas cannot infer the format of time text, which it would parse one value at a + time with ``dateutil``. + + Args: + values: The text values, with missing values for NULL. + + Returns: + The values as ``time`` objects, with None for NULL. + """ + import pandas as pd + + return pd.Series( + [_to_time(v) if isinstance(v, str) else None for v in values.to_numpy(dtype=object)], + index=values.index, + dtype=object, + ) + + def _to_utc_timestamps(values: Series) -> Series: """Convert ``timestamp with time zone`` text to ``datetime64[us, UTC]``. @@ -1515,6 +1539,7 @@ def _get_fetch_types(self) -> dict[Any, str]: return { label: "date" if converter.type == "date" else "datetime" for label, converter in self._csv_temporal_converters.items() + if converter.type != "time" } if self._kwargs.keys() & {"dtype_backend", "coerce_int96_timestamp_unit"}: return {} @@ -1529,11 +1554,11 @@ def _get_csv_temporal_converters( ) -> dict[Any, _TemporalConverter]: """Get the converters of the temporal columns that are converted after reading. - PyAthena converts the ``timestamp``, ``date``, and ``timestamp with time zone`` - columns after reading, unless ``parse_dates``, ``dtype``, or ``converters`` is - given to ``execute()``, with which pandas reads them, or the converter maps the - type to a conversion function other than the default ``timestamp with time - zone`` one. + PyAthena converts the ``timestamp``, ``date``, ``timestamp with time zone``, and + ``time`` columns after reading, unless ``parse_dates``, ``dtype``, or + ``converters`` is given to ``execute()``, with which pandas reads them, or the + converter maps the type to a conversion function other than the default + ``timestamp with time zone`` one. Args: columns: The label and the description of each column that is read. @@ -1584,13 +1609,27 @@ def _get_csv_converter( return converter def _trunc_date(self, df: DataFrame) -> DataFrame: - if self._time_columns: - # A NULL is None, as with the GetQueryResults fallback and the other types. - truncated = df.loc[:, self._time_columns].apply( - lambda r: r.dt.time.astype(object).where(r.notna(), None) - ) - for time_col in self._time_columns: - df.isetitem(df.columns.get_loc(time_col), truncated[time_col]) + """Convert the time columns that pandas parsed as datetimes to ``time`` objects. + + pandas parses the time columns with ``parse_dates`` when ``parse_dates``, + ``dtype``, or ``converters`` given to ``execute()`` keeps PyAthena from + converting them after reading. Time columns that pandas did not parse, or that + are not columns of the DataFrame, are left as they are. + + Args: + df: The DataFrame or chunk, updated in place. + + Returns: + The same DataFrame. + """ + for time_col in self._time_columns: + if time_col not in df.columns: + continue + i = df.columns.get_loc(time_col) + values = df.iloc[:, i] + if values.dtype.kind == "M": + # A NULL is None, as with the GetQueryResults fallback and the other types. + df.isetitem(i, values.dt.time.astype(object).where(values.notna(), None)) return df @override diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index c28116a33..2c18ead19 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, date, datetime +from datetime import UTC, date, datetime, time from unittest.mock import MagicMock, PropertyMock, patch from zoneinfo import ZoneInfo @@ -58,6 +58,8 @@ _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")) +_TIME_TEXT = "03:04:05.123" +_TIME = time(3, 4, 5, 123000) def _typed_values(rows): @@ -1073,6 +1075,20 @@ def test_temporal_values_in_one_column(self, tmp_path, chunksize, infer_string): assert pd.isna(df["ts"].iloc[2]) assert pd.isna(df["tstz"].iloc[2]) + @pytest.mark.parametrize("chunksize", [None, 1]) + def test_temporal_index(self, tmp_path, chunksize): + """A temporal column that index_col reads as the index is converted as a column is.""" + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("tm", "time"), ("ts", "timestamp")], + f'"tm","ts"\n"{_TIME_TEXT}","{_TIMESTAMP_TEXT}"\n', + chunksize=chunksize, + index_col="tm", + ) + df = pd.concat(list(result_set.iter_chunks())) + assert df.index.tolist() == [_TIME] + assert df["ts"].tolist() == [pd.Timestamp("2020-01-02 03:04:05.123456")] + @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.""" @@ -1257,13 +1273,14 @@ def test_time_zone_before_first_transition(self, tmp_path): 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"]) + @pytest.mark.parametrize("type_", ["timestamp", "date", "timestamp with time zone", "time"]) 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", + "time": "25:00:00.000", }[type_] with ( pd.option_context("future.infer_string", infer_string), @@ -1311,18 +1328,22 @@ def test_temporal_options_none(self, tmp_path, engine, key): assert str(result_set.as_pandas()["d"].dtype) == "datetime64[us]" assert result_set.fetchone() == (date(2020, 1, 2), 1) + # pandas parses the time column without a format. + @pytest.mark.filterwarnings("ignore:Could not infer format") @pytest.mark.parametrize("engine", ["c", "python", "pyarrow"]) @pytest.mark.parametrize( ("kwargs", "expected"), [ # PyAthena's parse_dates and converters read the temporal columns. - pytest.param({"dtype": {"n": "Int64"}}, (_PARSED, _DATE, _TOKYO, 1), id="dtype"), + pytest.param({"dtype": {"n": "Int64"}}, (_PARSED, _DATE, _TOKYO, _TIME, 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( + {"converters": {}}, (_PARSED, _DATE, _TOKYO_TEXT, _TIME, 1), id="converters" + ), + # parse_dates replaces PyAthena's, and the columns it does not name stay text. pytest.param( {"parse_dates": ["ts"]}, - (_PARSED, "2020-01-02", _TOKYO, 1), + (_PARSED, "2020-01-02", _TOKYO, _TIME_TEXT, 1), id="parse_dates", ), ], @@ -1333,8 +1354,8 @@ def test_temporal_options_read_by_pandas(self, tmp_path, engine, kwargs, expecte 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 + data = '"ts","d","tstz","tm","n"\n' + ( + f'"{_TIMESTAMP_TEXT}","2020-01-02","{_TOKYO_TEXT}","{_TIME_TEXT}","1"\n' * 2 ) assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES result_set = _local_result_set( @@ -1343,6 +1364,7 @@ def test_temporal_options_read_by_pandas(self, tmp_path, engine, kwargs, expecte ("ts", "timestamp"), ("d", "date"), ("tstz", "timestamp with time zone"), + ("tm", "time"), ("n", "integer"), ], data, @@ -1667,8 +1689,8 @@ 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 pyarrow strings, as - PyAthena reads them for the PyArrow engine with its default options before + With ``temporal_text``, the date, timestamp, and time columns are pyarrow strings, + as PyAthena reads them for the PyArrow engine with its default options before converting them. """ missing = float("nan") @@ -1707,7 +1729,9 @@ def _types_frame(infer_string, parse_time=True, temporal_text=False): ), # pandas supplies today's date when it parses a time without a date. "tm": ( - pd.Series(pd.to_datetime(["12:34:56.789", None])) + pd.Series(["12:34:56.789", None], dtype="string[pyarrow]") + if temporal_text + else pd.Series(pd.to_datetime(["12:34:56.789", None])) if parse_time else _string_series(["12:34:56.789", None], infer_string) ), @@ -1729,7 +1753,8 @@ def _types_frame(infer_string, parse_time=True, temporal_text=False): _TYPES_CSV, _TYPES, {}, - # The date and timestamp columns are read as text and converted afterwards. + # The date, timestamp, and time 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", From 9ebd8ff35c115d431af91572e9972729444690a1 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 20 Oct 2026 00:21:45 +0900 Subject: [PATCH 2/3] Convert a parsed time index level to time objects (GH-1308) With parse_dates, dtype, or converters given to execute(), pandas parses a time column that index_col names into the index, which _trunc_date() skipped. It now converts the column or index level through _replace_by_label(), as the after-read conversion does. Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 22 ++++++++++++---------- tests/pyathena/pandas/test_result_set.py | 20 +++++++++++++++----- 2 files changed, 27 insertions(+), 15 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 6de643732..04ca1355a 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -458,13 +458,16 @@ def _to_times(values: Series) -> Series: time with ``dateutil``. Args: - values: The text values, with missing values for NULL. + values: The text values, with missing values for NULL, or datetime values that + pandas parsed. Returns: The values as ``time`` objects, with None for NULL. """ import pandas as pd + if values.dtype.kind == "M": + return values.dt.time.astype(object).where(values.notna(), None) return pd.Series( [_to_time(v) if isinstance(v, str) else None for v in values.to_numpy(dtype=object)], index=values.index, @@ -1613,8 +1616,8 @@ def _trunc_date(self, df: DataFrame) -> DataFrame: pandas parses the time columns with ``parse_dates`` when ``parse_dates``, ``dtype``, or ``converters`` given to ``execute()`` keeps PyAthena from - converting them after reading. Time columns that pandas did not parse, or that - are not columns of the DataFrame, are left as they are. + converting them after reading. A time column can also be an index level. Time + columns that pandas did not parse are left as they are. Args: df: The DataFrame or chunk, updated in place. @@ -1623,13 +1626,12 @@ def _trunc_date(self, df: DataFrame) -> DataFrame: The same DataFrame. """ for time_col in self._time_columns: - if time_col not in df.columns: - continue - i = df.columns.get_loc(time_col) - values = df.iloc[:, i] - if values.dtype.kind == "M": - # A NULL is None, as with the GetQueryResults fallback and the other types. - df.isetitem(i, values.dt.time.astype(object).where(values.notna(), None)) + # A NULL is None, as with the GetQueryResults fallback and the other types. + _replace_by_label( + df, + time_col, + lambda values: _to_times(values) if values.dtype.kind == "M" else values, + ) return df @override diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 2c18ead19..00e9075b5 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1075,19 +1075,29 @@ def test_temporal_values_in_one_column(self, tmp_path, chunksize, infer_string): assert pd.isna(df["ts"].iloc[2]) assert pd.isna(df["tstz"].iloc[2]) + # pandas parses the time column without a format. + @pytest.mark.filterwarnings("ignore:Could not infer format") + @pytest.mark.parametrize( + "kwargs", + [ + pytest.param({}, id="default"), + # pandas parses the time column, which then becomes time objects. + pytest.param({"dtype": {}}, id="pandas_parses"), + ], + ) @pytest.mark.parametrize("chunksize", [None, 1]) - def test_temporal_index(self, tmp_path, chunksize): - """A temporal column that index_col reads as the index is converted as a column is.""" + def test_time_index(self, tmp_path, kwargs, chunksize): + """A time column that index_col reads as the index holds time objects, as a column does.""" result_set = _local_result_set( LocalBucketFileSystem(tmp_path), - [("tm", "time"), ("ts", "timestamp")], - f'"tm","ts"\n"{_TIME_TEXT}","{_TIMESTAMP_TEXT}"\n', + [("tm", "time"), ("n", "integer")], + f'"tm","n"\n"{_TIME_TEXT}","1"\n', chunksize=chunksize, index_col="tm", + **kwargs, ) df = pd.concat(list(result_set.iter_chunks())) assert df.index.tolist() == [_TIME] - assert df["ts"].tolist() == [pd.Timestamp("2020-01-02 03:04:05.123456")] @pytest.mark.parametrize("chunksize", [None, 1]) def test_unload_temporal_types(self, tmp_path, chunksize): From f5029cfefee44133cfddda1ccc9b9940c48c9191 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 20 Oct 2026 00:50:27 +0900 Subject: [PATCH 3/3] Keep time values that pyarrow parsed, and skip converted time columns (GH-1308) With dtype=None given to execute(), PyAthena's dtype mapping is replaced, and pyarrow parses time text without a fraction into time objects, which _to_times() turned into None. It now keeps them. _trunc_date() also skips the time columns that PyAthena converted after reading, which it rebuilt again when they were index levels. Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 22 ++++++++++++++-------- tests/pyathena/pandas/test_result_set.py | 11 +++++++---- 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 04ca1355a..34c183acb 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -12,7 +12,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack from contextvars import copy_context -from datetime import datetime, timezone +from datetime import datetime, time, timezone from functools import partial from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper from itertools import chain, islice @@ -318,9 +318,9 @@ class _TemporalConverter: read as text and converted as whole columns after reading, so that their dtypes depend only on the Athena type: ``datetime64[us]``, ``datetime64[us]``, ``datetime64[us, UTC]``, and ``object`` with ``time`` objects and None for NULL. - 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. + Fractional digits beyond microseconds are truncated. A column that pandas or + pyarrow already parsed, as when ``dtype=None`` given to ``execute()`` replaces + PyAthena's dtype mapping, gets the same dtype. """ __slots__ = ("_type",) @@ -458,8 +458,8 @@ def _to_times(values: Series) -> Series: time with ``dateutil``. Args: - values: The text values, with missing values for NULL, or datetime values that - pandas parsed. + values: The text values, with missing values for NULL, datetime values that + pandas parsed, or ``time`` values that pyarrow parsed. Returns: The values as ``time`` objects, with None for NULL. @@ -469,7 +469,10 @@ def _to_times(values: Series) -> Series: if values.dtype.kind == "M": return values.dt.time.astype(object).where(values.notna(), None) return pd.Series( - [_to_time(v) if isinstance(v, str) else None for v in values.to_numpy(dtype=object)], + [ + _to_time(v) if isinstance(v, str) else v if isinstance(v, time) else None + for v in values.to_numpy(dtype=object) + ], index=values.index, dtype=object, ) @@ -1617,7 +1620,8 @@ def _trunc_date(self, df: DataFrame) -> DataFrame: pandas parses the time columns with ``parse_dates`` when ``parse_dates``, ``dtype``, or ``converters`` given to ``execute()`` keeps PyAthena from converting them after reading. A time column can also be an index level. Time - columns that pandas did not parse are left as they are. + columns that PyAthena converted after reading or that pandas did not parse are + left as they are. Args: df: The DataFrame or chunk, updated in place. @@ -1626,6 +1630,8 @@ def _trunc_date(self, df: DataFrame) -> DataFrame: The same DataFrame. """ for time_col in self._time_columns: + if time_col in self._csv_temporal_converters: + continue # A NULL is None, as with the GetQueryResults fallback and the other types. _replace_by_label( df, diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 00e9075b5..23a33cb18 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1325,18 +1325,21 @@ def test_time_zone_results_do_not_depend_on_engine(self, tmp_path, kwargs): @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 + """parse_dates, dtype, or converters of None, their default, is as if not given. + + With dtype of None, pyarrow parses the dates and the times without a fraction. + """ + data = '"d","tm","n"\n' + '"2020-01-02","03:04:05","1"\n' * 10 assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES result_set = _local_result_set( LocalBucketFileSystem(tmp_path), - [("d", "date"), ("n", "integer")], + [("d", "date"), ("tm", "time"), ("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) + assert result_set.fetchone() == (date(2020, 1, 2), time(3, 4, 5), 1) # pandas parses the time column without a format. @pytest.mark.filterwarnings("ignore:Could not infer format")