Parse temporal text with fromisoformat and convert PandasCursor time columns after reading (GH-1308)

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

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

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
@@ -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,12 +314,13 @@ 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]``.
-    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.
+    ``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 or
+    pyarrow already parsed, as when ``dtype=None`` given to ``execute()`` replaces
+    PyAthena's dtype mapping, gets the same dtype.
     """
 
     __slots__ = ("_type",)
@@ -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,33 @@ 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, datetime values that
+            pandas parsed, or ``time`` values that pyarrow 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 v if isinstance(v, time) 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 +1545,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 +1560,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 +1615,29 @@ def _get_csv_converter(
         return converter
 
     def _trunc_date(self, df: DataFrame) -> DataFrame:
-        if self._time_columns:
+        """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. A time column can also be an index level. Time
+        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.
+
+        Returns:
+            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.
-            truncated = df.loc[:, self._time_columns].apply(
-                lambda r: r.dt.time.astype(object).where(r.notna(), None)
+            _replace_by_label(
+                df,
+                time_col,
+                lambda values: _to_times(values) if values.dtype.kind == "M" else values,
             )
-            for time_col in self._time_columns:
-                df.isetitem(df.columns.get_loc(time_col), truncated[time_col])
         return df
 
     @override