Convert time zone values to aware times and empty TIME/JSON text to NULL

docs/pandas.md

@@ -554,7 +554,7 @@ Common performance options:
 - `dtype`: Explicit column data types
 - `parse_dates`: Columns to parse as dates
 
-With `engine="pyarrow"`, PandasCursor uses the PyArrow engine only when 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`, and `json` 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 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.
 Otherwise, it falls back to the C engine.
 
 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.

docs/s3fs.md

@@ -125,6 +125,7 @@ The following type mappings are used:
 | timestamp | datetime.datetime |
 | timestamp with time zone | datetime.datetime (timezone-aware) |
 | time | datetime.time |
+| time with time zone | datetime.time (timezone-aware) |
 | varbinary | bytes |
 | array, map, row (struct) | Parsed into Python list/dict (see {ref}`usage-type-hints` for the types of nested values); values too complex to parse are returned as the original string |
 | json | Parsed JSON value (dict, list, or scalar) |

pyathena/arrow/converter.py

@@ -9,12 +9,14 @@
 
 from pyathena.converter import (
     Converter,
-    _csv_to_json,
     _to_binary,
     _to_date,
+    _to_datetime_with_tz,
     _to_decimal,
     _to_default,
+    _to_json,
     _to_time,
+    _to_time_with_tz,
 )
 from pyathena.util import override
 
@@ -24,9 +26,11 @@
 _DEFAULT_ARROW_CONVERTERS: dict[str, Callable[[str | None], Any | None]] = {
     "date": _to_date,
     "time": _to_time,
+    "time with time zone": _to_time_with_tz,
+    "timestamp with time zone": _to_datetime_with_tz,
     "decimal": _to_decimal,
     "varbinary": _to_binary,
-    "json": _csv_to_json,
+    "json": _to_json,
 }
 
 
@@ -85,6 +89,8 @@ def _dtypes(self) -> dict[str, type[Any]]:
                 "timestamp": pa.timestamp("ms"),
                 "date": pa.timestamp("ms"),
                 "time": pa.string(),
+                "time with time zone": pa.string(),
+                "timestamp with time zone": pa.string(),
                 "varbinary": pa.string(),
                 "array": pa.string(),
                 "map": pa.string(),

pyathena/arrow/result_set.py

@@ -12,7 +12,7 @@
 
 from pyathena import OperationalError
 from pyathena.arrow.util import to_column_info
-from pyathena.converter import Converter, _json_text_converter, _to_default
+from pyathena.converter import Converter, _text_value_converter, _to_default
 from pyathena.error import ProgrammingError
 from pyathena.model import AthenaQueryExecution
 from pyathena.result_set import AthenaResultSet
@@ -147,7 +147,8 @@ def __init__(
 
             self._table = pa.Table.from_pydict({})
         # The fetch methods convert the values read from a result file. GetQueryResults
-        # values are already converted, except json values, which stay text.
+        # values are already converted, except json and time with time zone values,
+        # which stay text.
         self._convert_rows = bool(self.output_location)
         self._batches = iter(self._table.to_batches(arraysize))
 
@@ -255,7 +256,9 @@ def _fetch(self) -> None:
         else:
             dict_rows = rows.to_pydict()
             converters = (
-                self.converters if self._convert_rows else self._json_converters(self.converters)
+                self.converters
+                if self._convert_rows
+                else self._text_value_converters(self.converters)
             )
             if converters:
                 column_names = dict_rows.keys()
@@ -387,12 +390,13 @@ def _as_arrow_from_api(self, converter: Converter | None = None) -> Table:
 
         Args:
             converter: Type converter for result values. Defaults to
-                ``DefaultTypeConverter`` with json values kept as text, as in
-                the CSV result file.
+                ``DefaultTypeConverter`` with json and time with time zone values kept as
+                text, as in the CSV result file. Arrow has no type for JSON values or for
+                times with a time zone.
         """
         import pyarrow as pa
 
-        rows = self._fetch_all_rows(converter or _json_text_converter())
+        rows = self._fetch_all_rows(converter or _text_value_converter())
         if not rows:
             return pa.Table.from_pydict({})
         description = self.description if self.description else []

pyathena/converter.py

@@ -9,7 +9,7 @@
 from abc import ABCMeta, abstractmethod
 from collections.abc import Callable
 from copy import deepcopy
-from datetime import date, datetime, time
+from datetime import date, datetime, time, timedelta, timezone
 from decimal import Decimal
 from typing import Any, ClassVar
 
@@ -67,25 +67,92 @@ def _to_datetime(varchar_value: str | None) -> datetime | None:
     return _parse_datetime(varchar_value)
 
 
+_UTC_OFFSET_PATTERN: re.Pattern[str] = re.compile(r"([+-])(\d{2}):(\d{2})")
+
+
+def _parse_utc_offset(value: str) -> timezone | None:
+    """Parse a ``+HH:MM`` or ``-HH:MM`` UTC offset.
+
+    Args:
+        value: The text to parse.
+
+    Returns:
+        The fixed-offset time zone, or None if the text is not an offset.
+    """
+    match = _UTC_OFFSET_PATTERN.fullmatch(value)
+    if not match:
+        return None
+    sign, hours, minutes = match.groups()
+    offset = timedelta(hours=int(hours), minutes=int(minutes))
+    return timezone(-offset if sign == "-" else offset)
+
+
 def _to_datetime_with_tz(varchar_value: str | None) -> datetime | None:
     """Convert an Athena TIMESTAMP WITH TIME ZONE value to an aware datetime.
 
     Args:
-        varchar_value: The value as text with a trailing zone name, or None.
+        varchar_value: The value as text with a trailing zone name or ``+HH:MM`` /
+            ``-HH:MM`` UTC offset, or None. An empty string, which pandas and the
+            Arrow CSV reader return for NULL, is None.
 
     Returns:
         The aware datetime, or None.
     """
-    if varchar_value is None:
+    if not varchar_value:
         return None
     datetime_, _, tz = varchar_value.rpartition(" ")
-    return _parse_datetime(datetime_).replace(tzinfo=gettz(tz))
+    return _parse_datetime(datetime_).replace(tzinfo=_parse_utc_offset(tz) or gettz(tz))
+
+
+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:
-    if varchar_value is None:
+    """Convert an Athena TIME value to a time.
+
+    Args:
+        varchar_value: The value as text, or None. An empty string, which the Arrow
+            CSV reader returns for NULL, is None.
+
+    Returns:
+        The time, or None.
+    """
+    if not varchar_value:
         return None
-    return datetime.strptime(varchar_value, "%H:%M:%S.%f").time()
+    return _parse_time(varchar_value)
+
+
+def _to_time_with_tz(varchar_value: str | None) -> time | None:
+    """Convert an Athena TIME WITH TIME ZONE value to an aware time.
+
+    Args:
+        varchar_value: The value as text with a trailing ``+HH:MM`` or ``-HH:MM``
+            offset, or None. An empty string, which pandas and the Arrow CSV reader
+            return for NULL, is None.
+
+    Returns:
+        The time with a fixed-offset ``tzinfo``, or None.
+    """
+    if not varchar_value:
+        return None
+    index = max(varchar_value.rfind("+"), varchar_value.rfind("-"))
+    return _parse_time(varchar_value[:index]).replace(
+        tzinfo=_parse_utc_offset(varchar_value[index:])
+    )
 
 
 def _to_float(varchar_value: str | None) -> float | None:
@@ -119,17 +186,12 @@ def _to_binary(varchar_value: str | None) -> bytes | None:
 
 
 def _to_json(varchar_value: str | None) -> Any | None:
-    if varchar_value is None:
-        return None
-    return json.loads(varchar_value)
-
-
-def _csv_to_json(varchar_value: str | None) -> Any | None:
-    """Convert an Athena JSON value read from a CSV result file.
+    """Convert an Athena JSON value to a Python value.
 
     Args:
-        varchar_value: The value as JSON text, or None. An empty string, which
-            CSV results use for SQL NULL, is also treated as NULL.
+        varchar_value: The JSON text, or None. An empty string, which pandas and the
+            Arrow CSV reader return for NULL, is None; Athena never returns empty JSON
+            text.
 
     Returns:
         The decoded value, or None for SQL NULL.
@@ -494,6 +556,7 @@ def _to_default(varchar_value: str | None) -> str | None:
     "timestamp with time zone": _to_datetime_with_tz,
     "date": _to_date,
     "time": _to_time,
+    "time with time zone": _to_time_with_tz,
     "varbinary": _to_binary,
     "array": _to_array,
     "map": _to_map,
@@ -744,12 +807,31 @@ def _parse_type_hint(self, type_hint: str) -> TypeNode:
         return self._parsed_hints[normalized]
 
 
-def _json_text_converter() -> DefaultTypeConverter:
-    """Return a ``DefaultTypeConverter`` that keeps json values as text.
+# The types whose values the Arrow and Polars GetQueryResults fallbacks keep as text,
+# as in a CSV result file, and convert when the rows are fetched.
+_TEXT_VALUE_TYPES: tuple[str, ...] = ("json", "time with time zone", "timestamp with time zone")
+
+
+def _text_value_converter() -> DefaultTypeConverter:
+    """Return a ``DefaultTypeConverter`` that keeps ``_TEXT_VALUE_TYPES`` values as text.
+
+    Values nested in typed complex values keep only the time zone types as text,
+    because Arrow and Polars time and timestamp types hold one time zone per column;
+    nested JSON values are decoded as before.
 
     Returns:
         The converter.
     """
     converter = DefaultTypeConverter()
-    converter.set("json", _to_default)
+    for type_ in _TEXT_VALUE_TYPES:
+        converter.set(type_, _to_default)
+    converter._typed_converter = TypedValueConverter(
+        converters={
+            **_DEFAULT_CONVERTERS,
+            "time with time zone": _to_default,
+            "timestamp with time zone": _to_default,
+        },
+        default_converter=_to_default,
+        struct_parser=_to_struct,
+    )
     return converter

pyathena/pandas/converter.py

@@ -9,11 +9,13 @@
 
 from pyathena.converter import (
     Converter,
-    _csv_to_json,
     _to_binary,
     _to_boolean,
+    _to_datetime_with_tz,
     _to_decimal,
     _to_default,
+    _to_json,
+    _to_time_with_tz,
 )
 from pyathena.util import override
 
@@ -24,7 +26,9 @@
     "boolean": _to_boolean,
     "decimal": _to_decimal,
     "varbinary": _to_binary,
-    "json": _csv_to_json,
+    "json": _to_json,
+    "time with time zone": _to_time_with_tz,
+    "timestamp with time zone": _to_datetime_with_tz,
 }
 
 

pyathena/pandas/result_set.py

@@ -257,9 +257,7 @@ class AthenaPandasResultSet(AthenaResultSet):
     _PARSE_DATES: ClassVar[list[str]] = [
         "date",
         "time",
-        "time with time zone",
         "timestamp",
-        "timestamp with time zone",
     ]
 
     def __init__(
@@ -343,9 +341,7 @@ def __init__(
 
         # Cache time column names for efficient _trunc_date processing
         description = self.description if self.description else []
-        self._time_columns: list[str] = [
-            d[0] for d in description if d[1] in ("time", "time with time zone")
-        ]
+        self._time_columns: list[str] = [d[0] for d in description if d[1] == "time"]
 
         import pandas as pd
 

pyathena/parser.py

@@ -141,7 +141,11 @@ def parse(self, type_str: str) -> TypeNode:
                 return TypeNode(type_name=type_name, children=[key_type, value_type])
             return TypeNode(type_name=type_name)
 
-        # Types with parameters like decimal(10, 2), varchar(255)
+        # Types with parameters like decimal(10, 2), varchar(255), or
+        # time(3) with time zone, whose suffix is part of the type name.
+        # Other trailing text is ignored.
+        if " ".join(type_str[close_idx + 1 :].lower().split()) == "with time zone":
+            type_name = f"{type_name} with time zone"
         return TypeNode(type_name=type_name)
 
     def _split_type_args(self, s: str) -> list[str]:
@@ -268,19 +272,21 @@ def convert(self, value: str, type_node: TypeNode) -> Any:
         return converter_fn(value)
 
     @staticmethod
-    def _to_json_str(value: Any) -> str:
+    def _to_json_str(value: Any, type_node: TypeNode) -> str:
         """Convert a JSON-parsed value back to a string for further conversion.
 
-        Uses json.dumps for dict/list to produce valid JSON, and str() for
-        scalar types to produce converter-compatible strings.
+        Uses json.dumps for dict/list values and for every value of a JSON type,
+        so that the JSON converter decodes the original JSON text, and str() for
+        the other scalar types to produce converter-compatible strings.
 
         Args:
             value: A value from json.loads output.
+            type_node: The type of the value.
 
         Returns:
             String representation suitable for type conversion.
         """
-        if isinstance(value, (dict, list)):
+        if isinstance(value, (dict, list)) or type_node.type_name == "json":
             return json.dumps(value)
         return str(value)
 
@@ -324,7 +330,7 @@ def _convert_typed_array(self, value: str, type_node: TypeNode) -> list[Any] | N
                     return [
                         None
                         if elem is None
-                        else self.convert(self._to_json_str(elem), element_type)
+                        else self.convert(self._to_json_str(elem, element_type), element_type)
                         for elem in parsed
                     ]
             except json.JSONDecodeError:
@@ -379,8 +385,12 @@ def _convert_typed_map(self, value: str, type_node: TypeNode) -> dict[str, Any]
                 parsed = json.loads(value)
                 if isinstance(parsed, dict):
                     return {
-                        str(self.convert(self._to_json_str(k), key_type) if k is not None else k): (
-                            self.convert(self._to_json_str(v), value_type)
+                        str(
+                            self.convert(self._to_json_str(k, key_type), key_type)
+                            if k is not None
+                            else k
+                        ): (
+                            self.convert(self._to_json_str(v, value_type), value_type)
                             if v is not None
                             else None
                         )
@@ -447,7 +457,7 @@ def _convert_typed_struct(self, value: str, type_node: TypeNode) -> dict[str, An
                     for i, (k, v) in enumerate(parsed.items()):
                         ft = self._get_field_type(k, type_node, i)
                         result[k] = (
-                            self.convert(self._to_json_str(v), ft) if v is not None else None
+                            self.convert(self._to_json_str(v, ft), ft) if v is not None else None
                         )
                     return result
             except json.JSONDecodeError: