Name the column in DataError from the fetch methods of every cursor

pyathena/arrow/result_set.py

@@ -3,14 +3,14 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from typing import (
     TYPE_CHECKING,
     Any,
     ClassVar,
 )
 
-from pyathena import OperationalError
+from pyathena import DataError, OperationalError
 from pyathena.arrow.converter import _to_null, _to_timestamp
 from pyathena.arrow.util import to_column_info
 from pyathena.converter import Converter, _to_default
@@ -201,6 +201,11 @@ def _fetch(self) -> None:
 
         A table can hold empty record batches between others, such as one of
         an empty result file, which would otherwise end the fetches early.
+
+        Raises:
+            DataError: If a value fails to convert, such as a ``timestamp with time
+                zone`` value with an unknown zone or a local time without one instant in
+                its zone. The message names the column.
         """
         rows = next((batch for batch in self._batches if batch.num_rows), None)
         if rows is None:
@@ -209,12 +214,22 @@ def _fetch(self) -> None:
         # converters property keep one column per name.
         columns = [column.to_pylist() for column in rows.columns]
         description = self.description if self.description else []
+        names = [d[0] for d in description]
         converters = [self._converter.get(d[1]) for d in description]
-        if any(convert is not _to_default for convert in converters):
-            processed_rows = [
-                tuple(convert(v) for convert, v in zip(converters, row, strict=False))
-                for row in zip(*columns, strict=False)
-            ]
+
+        def convert(row: tuple[Any, ...]) -> Iterator[Any]:
+            # Unlike a loop that returns the values, a generator turns a StopIteration
+            # from the converter into a RuntimeError instead of ending an iteration of
+            # the rows.
+            for name, converter, value in zip(names, converters, row, strict=False):
+                try:
+                    converted = converter(value)
+                except DataError as e:
+                    raise DataError(f"Column {name!r}: {e}") from e
+                yield converted
+
+        if any(converter is not _to_default for converter in converters):
+            processed_rows = [tuple(convert(row)) for row in zip(*columns, strict=False)]
         else:
             processed_rows = list(zip(*columns, strict=False))
         self._rows.extend(processed_rows)

pyathena/pandas/result_set.py

@@ -363,6 +363,51 @@ def _convert_values(self, values: Series) -> Series:
         return _to_timestamps(values, date=self._type == "date")
 
 
+class _TimeZoneConverter:
+    """A converter that ``pandas.read_csv()`` applies to each ``timestamp with time zone`` value.
+
+    With ``parse_dates`` or ``dtype`` given to ``execute()``, or with a conversion
+    function other than the default one, pandas applies the column's conversion
+    function to each value as it reads, unless ``converters`` given to ``execute()``
+    replaces the conversion functions. This names the column in the
+    ``DataError`` of a value that fails to convert, as ``_TemporalConverter`` does.
+    """
+
+    __slots__ = ("_converter", "_label", "error")
+
+    def __init__(self, converter: Callable[[str | None], Any], label: Any) -> None:
+        """Wrap the conversion function of a column.
+
+        Args:
+            converter: The conversion function.
+            label: The label of the column.
+        """
+        self._converter = converter
+        self._label = label
+        # The DataError that the converter raised, which a whole read raises as it is.
+        self.error: DataError | None = None
+
+    def __call__(self, value: str | None) -> Any:
+        """Convert a value.
+
+        Args:
+            value: The text of the value.
+
+        Returns:
+            The converted value.
+
+        Raises:
+            DataError: If the value fails to convert, such as a value with an unknown
+                zone or a local time without one instant in its zone. The message names
+                the column.
+        """
+        try:
+            return self._converter(value)
+        except DataError as e:
+            self.error = DataError(f"Column {self._label!r}: {e}")
+            raise self.error from e
+
+
 def _to_timestamps(values: Series, date: bool = False) -> Series:
     """Convert timestamp or date text to ``datetime64[us]``.
 
@@ -1246,22 +1291,29 @@ def _get_csv_temporal_converters(
             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.
+    def _get_csv_converter(
+        self, type_: str, label: Any
+    ) -> Callable[[str | None], Any] | _JSONConverter:
+        """Get the converter for a column of the CSV result file.
 
         json columns get a ``_JSONConverter``, which ``_read_json_columns_as_text()``
         replaces in the converters for ``pandas.read_csv()``, and which converts the
-        columns after reading.
+        columns after reading. ``timestamp with time zone`` columns get a
+        ``_TimeZoneConverter``, which names the column in a ``DataError``.
 
         Args:
             type_: The Athena type of the column.
+            label: The label of the column.
 
         Returns:
-            The conversion function, or the ``_JSONConverter`` of a json column.
+            The conversion function, the ``_JSONConverter`` of a json column, or the
+            ``_TimeZoneConverter`` of a ``timestamp with time zone`` column.
         """
         converter = self._converter.get(type_)
         if type_ == "json":
             return _JSONConverter(converter)
+        if type_ == "timestamp with time zone":
+            return _TimeZoneConverter(converter, label)
         return converter
 
     def _trunc_date(self, df: DataFrame) -> DataFrame:
@@ -1317,6 +1369,8 @@ def _read_csv(self) -> TextFileReader | DataFrame:
             The DataFrame of the results, or a reader of their chunks.
 
         Raises:
+            DataError: If a ``timestamp with time zone`` value that pandas converts as it
+                reads fails to convert, naming the column.
             OperationalError: If reading the results fails.
         """
         import pandas as pd
@@ -1414,6 +1468,13 @@ def _read_csv(self) -> TextFileReader | DataFrame:
             return result
 
         except Exception as e:
+            # converters given to execute() can be of any type.
+            converters = read_csv_kwargs.get("converters")
+            if isinstance(converters, Mapping) and any(
+                isinstance(converter, _TimeZoneConverter) and converter.error is e
+                for converter in converters.values()
+            ):
+                raise
             _logger.exception(f"Failed to read {location}.")
             raise OperationalError(*e.args) from e
 
@@ -1519,7 +1580,7 @@ def _get_csv_read_options(self, csv_engine: str, chunksize: int | None) -> dict[
             "names": names,
             "dtype": {**self.dtypes, **dict.fromkeys(temporal, text_dtype)},
             "converters": {
-                d[0]: self._get_csv_converter(d[1])
+                d[0]: self._get_csv_converter(d[1], d[0])
                 for d in description
                 if d[1] in self._converter.mappings and d[0] not in self._csv_temporal_converters
             },
@@ -1690,7 +1751,7 @@ def _key_csv_columns_by_labels(
             }
         if "converters" not in self._kwargs:
             read_csv_kwargs["converters"] = {
-                label: self._get_csv_converter(d[1])
+                label: self._get_csv_converter(d[1], label)
                 for label, d in columns
                 if d[1] in self._converter.mappings and label not in self._csv_temporal_converters
             }

pyathena/polars/result_set.py

@@ -31,7 +31,7 @@
 from pyathena import OperationalError
 from pyathena.arrow.reader import ParquetChunkReader
 from pyathena.converter import Converter, _to_default
-from pyathena.error import ProgrammingError
+from pyathena.error import DataError, ProgrammingError
 from pyathena.filesystem.s3 import S3FileSystem
 from pyathena.model import AthenaQueryExecution
 from pyathena.polars.converter import _to_datetimes
@@ -192,16 +192,25 @@ def _iter_row_blocks(self, df: pl.DataFrame) -> Iterator[Iterator[tuple[Any, ...
 
         Yields:
             An iterator of the rows of each block.
+
+        Raises:
+            DataError: While the rows are iterated, 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 polars as pl
 
-        # Each column's Series, or None for a name the chunk does not have, and its
-        # converter, or None for the identity.
-        columns: list[tuple[pl.Series | None, Callable[[Any], Any] | None]] = []
+        # Each column's name, its Series, or None for a name the chunk does not have,
+        # and its converter, or None for the identity.
+        columns: list[tuple[str, pl.Series | None, Callable[[Any], Any] | None]] = []
         for name in self._column_names:
             converter = self._converters.get(name)
             columns.append(
-                (df.get_column(name, default=None), None if converter is _to_default else converter)
+                (
+                    name,
+                    df.get_column(name, default=None),
+                    None if converter is _to_default else converter,
+                )
             )
 
         def take_values(block: pl.Series) -> Iterator[Any]:
@@ -212,17 +221,23 @@ def take_values(block: pl.Series) -> Iterator[Any]:
             for i in range(block.len()):
                 yield block.slice(i, 1).to_list()[0]
 
-        def convert(converter: Callable[[Any], Any], values: Iterable[Any]) -> Iterator[Any]:
+        def convert(
+            name: str, converter: Callable[[Any], Any], values: Iterable[Any]
+        ) -> Iterator[Any]:
             # Unlike map(), a generator turns a StopIteration from the converter into a
             # RuntimeError instead of ending the rows early.
             for value in values:
-                yield converter(value)
+                try:
+                    converted = converter(value)
+                except DataError as e:
+                    raise DataError(f"Column {name!r}: {e}") from e
+                yield converted
 
         size = self._ROW_BLOCK_SIZE
         for start in range(0, df.height, size):
             length = min(size, df.height - start)
             blocks: list[Iterable[Any]] = []
-            for series, converter in columns:
+            for name, series, converter in columns:
                 values: Iterable[Any]
                 if series is None:
                     values = [None] * length
@@ -232,7 +247,7 @@ def convert(converter: Callable[[Any], Any], values: Iterable[Any]) -> Iterator[
                         values = block.to_list()
                     except (Exception, pl.exceptions.PanicException):
                         values = take_values(block)
-                blocks.append(values if converter is None else convert(converter, values))
+                blocks.append(values if converter is None else convert(name, converter, values))
             # Without column names, each row is empty, as iter_rows() gave it.
             yield zip(*blocks, strict=True) if blocks else repeat((), length)