From bc27a086e5e69ddc5706be5947db09943fa02652 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 15:24:20 +0900 Subject: [PATCH 01/10] Name the column in DataError from the fetch methods of every cursor (GH-1217) A value that fails to convert in Cursor, DictCursor, S3FSCursor, ArrowCursor, or PolarsCursor, such as a timestamp with time zone value with an unknown zone or a local time that its zone repeats or skips, now raises DataError naming the column, as PandasCursor does since #1273. ArrowCursor converts each column of a record batch at once instead of each row. Co-Authored-By: Claude Opus 5.5 --- pyathena/arrow/result_set.py | 25 +++-- pyathena/polars/result_set.py | 32 ++++-- pyathena/result_set.py | 143 +++++++++++++++------------ pyathena/s3fs/result_set.py | 23 ++--- tests/pyathena/arrow/test_cursor.py | 2 +- tests/pyathena/polars/test_cursor.py | 2 +- tests/pyathena/s3fs/test_cursor.py | 2 +- tests/pyathena/test_cursor.py | 7 +- 8 files changed, 133 insertions(+), 103 deletions(-) diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 869a185b4..f13428b9e 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -10,7 +10,7 @@ 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,15 +214,15 @@ 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 [] - 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) - ] - else: - processed_rows = list(zip(*columns, strict=False)) - self._rows.extend(processed_rows) + for i, d in enumerate(description): + convert = self._converter.get(d[1]) + if convert is _to_default: + continue + try: + columns[i] = [convert(v) for v in columns[i]] + except DataError as e: + raise DataError(f"Column {d[0]!r}: {e}") from e + self._rows.extend(zip(*columns, strict=False)) def _open_csv(self, location: str) -> Any: """Open the CSV result file for ``pyarrow.csv.read_csv()``. diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index d4017a402..5fc465ab5 100644 --- a/pyathena/polars/result_set.py +++ b/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,22 @@ 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: + yield converter(value) + except DataError as e: + raise DataError(f"Column {name!r}: {e}") from e 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 +246,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) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 97184e012..57d54f5b5 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -6,7 +6,7 @@ import logging import operator from datetime import datetime -from itertools import islice +from itertools import islice, repeat from typing import ( TYPE_CHECKING, Any, @@ -20,7 +20,7 @@ from pyathena.util import RetryConfig, override, parse_output_location, retry_api_call if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterable, Iterator from pyathena.connection import Connection @@ -687,36 +687,76 @@ def _get_rows( converter: Converter | None = None, ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: conv = converter or self._converter - col_types = self._column_types - col_hints = self._column_type_hints - if col_hints and col_types: - return [ - tuple( - conv.convert(col_type, row.get("VarCharValue"), type_hint=hint) - if hint - else conv.convert(col_type, row.get("VarCharValue")) - for col_type, row, hint in zip( - col_types, rows[i].get("Data", []), col_hints, strict=False - ) - ) - for i in range(offset, len(rows)) - ] - if col_types: - return [ - tuple( - conv.convert(col_type, row.get("VarCharValue")) - for col_type, row in zip(col_types, rows[i].get("Data", []), strict=False) - ) - for i in range(offset, len(rows)) - ] + names, types = self._get_column_names_and_types(metadata) + hints = self._column_type_hints or repeat(None) return [ tuple( - conv.convert(meta.get("Type"), row.get("VarCharValue")) - for meta, row in zip(metadata, rows[i].get("Data", []), strict=False) + self._convert_values( + conv, + (data.get("VarCharValue") for data in rows[i].get("Data", [])), + names, + types, + hints, + ) ) for i in range(offset, len(rows)) ] + def _get_column_names_and_types( + self, metadata: tuple[Any, ...] + ) -> tuple[tuple[Any, ...], tuple[Any, ...]]: + """Get the names and types of the result columns. + + Args: + metadata: The ``ColumnInfo`` of the result, for a result set whose column + types are not set. + + Returns: + The names and the types, in column order. + """ + if self._column_names and self._column_types: + return self._column_names, self._column_types + return ( + tuple(m.get("Name") for m in metadata), + tuple(m.get("Type") for m in metadata), + ) + + def _convert_values( + self, + converter: Converter, + values: Iterable[str | None], + names: Iterable[Any], + types: Iterable[Any], + hints: Iterable[str | None], + ) -> Iterator[Any | None]: + """Convert the values of a row. + + Args: + converter: The converter of the values. + values: The values as text, in column order. + names: The names of the columns. + types: The Athena types of the columns. + hints: The type hints of the columns, None for a column without one. + + Yields: + The converted values. A row with fewer values than columns, or more, gives + as many values as the shorter of the two. + + 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. + """ + for name, type_, hint, value in zip(names, types, hints, values, strict=False): + try: + yield ( + converter.convert(type_, value, type_hint=hint) + if hint + else converter.convert(type_, value) + ) + except DataError as e: + raise DataError(f"Column {name!r}: {e}") from e + def _parse_result_rows( self, response: dict[str, Any] ) -> tuple[list[dict[str, Any]], str | None]: @@ -956,48 +996,21 @@ def _get_rows( converter: Converter | None = None, ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: conv = converter or self._converter - col_types = self._column_types - col_names = self._column_names - col_hints = self._column_type_hints - if col_hints and col_types and col_names: - return [ - self.dict_type( - ( - name, - conv.convert(col_type, row.get("VarCharValue"), type_hint=hint) - if hint - else conv.convert(col_type, row.get("VarCharValue")), - ) - for name, col_type, row, hint in zip( - col_names, - col_types, - rows[i].get("Data", []), - col_hints, - strict=False, - ) - ) - for i in range(offset, len(rows)) - ] - if col_types and col_names: - return [ - self.dict_type( - ( - name, - conv.convert(col_type, row.get("VarCharValue")), - ) - for name, col_type, row in zip( - col_names, col_types, rows[i].get("Data", []), strict=False - ) - ) - for i in range(offset, len(rows)) - ] + names, types = self._get_column_names_and_types(metadata) + hints = self._column_type_hints or repeat(None) return [ self.dict_type( - ( - meta.get("Name"), - conv.convert(meta.get("Type"), row.get("VarCharValue")), + zip( + names, + self._convert_values( + conv, + (data.get("VarCharValue") for data in rows[i].get("Data", [])), + names, + types, + hints, + ), + strict=False, ) - for meta, row in zip(metadata, rows[i].get("Data", []), strict=False) ) for i in range(offset, len(rows)) ] diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py index 7569a92f3..35915fd99 100644 --- a/pyathena/s3fs/result_set.py +++ b/pyathena/s3fs/result_set.py @@ -11,6 +11,7 @@ import logging from io import TextIOWrapper +from itertools import repeat from typing import TYPE_CHECKING, Any, ClassVar from fsspec import AbstractFileSystem @@ -222,11 +223,13 @@ def _fetch(self) -> None: if not self._csv_reader: return + col_names = self._column_names col_types = self._column_types - if not col_types: + if not col_names or not col_types: description = self.description if self.description else [] + col_names = tuple(d[0] for d in description) col_types = tuple(d[1] for d in description) - col_hints = self._column_type_hints + col_hints = self._column_type_hints or repeat(None) rows_fetched = 0 while rows_fetched < self._arraysize: @@ -238,19 +241,9 @@ def _fetch(self) -> None: self._csv_reader = None break - if col_hints: - converted_row = tuple( - self._converter.convert(col_type, value, type_hint=hint) - if hint - else self._converter.convert(col_type, value) - for col_type, value, hint in zip(col_types, row, col_hints, strict=False) - ) - else: - converted_row = tuple( - self._converter.convert(col_type, value) - for col_type, value in zip(col_types, row, strict=False) - ) - self._rows.append(converted_row) + self._rows.append( + tuple(self._convert_values(self._converter, row, col_names, col_types, col_hints)) + ) rows_fetched += 1 @override diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index 4d33f4ee1..222ebf844 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -907,7 +907,7 @@ def test_ambiguous_time_zone_value(self, arrow_cursor): The table keeps the text, as for every TIMESTAMP WITH TIME ZONE value. """ arrow_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) - with pytest.raises(DataError, match="is ambiguous"): + with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): arrow_cursor.fetchall() @pytest.mark.skipif(not ENV.managed_work_group, reason="AWS_ATHENA_MANAGED_WORKGROUP not set") diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 1b1985818..f786bc103 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -892,7 +892,7 @@ def test_ambiguous_time_zone_value(self, polars_cursor): The table keeps the text, as for every TIMESTAMP WITH TIME ZONE value. """ polars_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) - with pytest.raises(DataError): + with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): polars_cursor.fetchall() @pytest.mark.skipif(not ENV.managed_work_group, reason="AWS_ATHENA_MANAGED_WORKGROUP not set") diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index e83bae4d8..ce551c884 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -323,7 +323,7 @@ def test_fetch_all_rows(self, s3fs_cursor): def test_ambiguous_time_zone_value(self, s3fs_cursor): """A local time that its zone repeats raises DataError instead of choosing one.""" s3fs_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) - with pytest.raises(DataError): + with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): s3fs_cursor.fetchall() @pytest.mark.parametrize( diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 6623f3963..57584d9c4 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -1682,7 +1682,7 @@ def test_ambiguous_time_zone_value(self, cursor): The first page of rows is converted by execute(). """ - with pytest.raises(DataError, match="is ambiguous"): + with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) @pytest.mark.parametrize( @@ -1950,6 +1950,11 @@ def __init__(self, **kwargs): assert isinstance(cursor.result_set, CustomResultSet) assert type(cursor.fetchone()) is OrderedDict + def test_ambiguous_time_zone_value(self, dict_cursor): + """A local time that its zone repeats raises DataError naming the column.""" + with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): + dict_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) + def test_null_vs_empty_string(self, dict_cursor): """ DictCursor should properly distinguish NULL from empty string. From 914200e121dcbecb5e20bbf91d10b25d52c362f9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 15:26:19 +0900 Subject: [PATCH 02/10] Keep the yield outside the conversion's try (GH-1217) Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 3 ++- pyathena/result_set.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 5fc465ab5..c3b31dfba 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -228,9 +228,10 @@ def convert( # RuntimeError instead of ending the rows early. for value in values: try: - yield converter(value) + 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): diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 57d54f5b5..6971c7068 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -749,13 +749,14 @@ def _convert_values( """ for name, type_, hint, value in zip(names, types, hints, values, strict=False): try: - yield ( + converted = ( converter.convert(type_, value, type_hint=hint) if hint else converter.convert(type_, value) ) except DataError as e: raise DataError(f"Column {name!r}: {e}") from e + yield converted def _parse_result_rows( self, response: dict[str, Any] From 80b8f91d96c6c5401665f1fe49dcccb1d9ce1442 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 15:30:41 +0900 Subject: [PATCH 03/10] Keep StopIteration from a converter a RuntimeError in the Arrow fetch (GH-1217) The column-wise conversion uses a generator, as the row-wise one did, so a StopIteration from a converter does not end an iteration of the rows. The GetQueryResults rows build their values in a list. Co-Authored-By: Claude Opus 5.5 --- pyathena/arrow/result_set.py | 4 +++- pyathena/result_set.py | 6 ++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index f13428b9e..0cdd412a4 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -219,7 +219,9 @@ def _fetch(self) -> None: if convert is _to_default: continue try: - columns[i] = [convert(v) for v in columns[i]] + # Unlike a list comprehension, a generator turns a StopIteration from the + # converter into a RuntimeError instead of ending an iteration of the rows. + columns[i] = list(convert(v) for v in columns[i]) # noqa: C400 except DataError as e: raise DataError(f"Column {d[0]!r}: {e}") from e self._rows.extend(zip(*columns, strict=False)) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 6971c7068..9c4d870d6 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -693,7 +693,7 @@ def _get_rows( tuple( self._convert_values( conv, - (data.get("VarCharValue") for data in rows[i].get("Data", [])), + [data.get("VarCharValue") for data in rows[i].get("Data", [])], names, types, hints, @@ -747,6 +747,8 @@ def _convert_values( zone`` value with an unknown zone or a local time without one instant in its zone. The message names the column. """ + # 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, type_, hint, value in zip(names, types, hints, values, strict=False): try: converted = ( @@ -1005,7 +1007,7 @@ def _get_rows( names, self._convert_values( conv, - (data.get("VarCharValue") for data in rows[i].get("Data", [])), + [data.get("VarCharValue") for data in rows[i].get("Data", [])], names, types, hints, From f81e5afa60d63e2180a6054830078ea93fb1c086 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 15:40:12 +0900 Subject: [PATCH 04/10] Convert the Arrow record batches row by row again (GH-1217) Converting a column at a time changed the order in which converters see values, and so the results of a stateful custom converter and the value whose failure is reported first. Each row is converted in column order again, naming the column of a failure. The tests anchor the column prefix at the start of the message. Co-Authored-By: Claude Opus 5.5 --- pyathena/arrow/result_set.py | 32 +++++++++++++++++----------- tests/pyathena/arrow/test_cursor.py | 2 +- tests/pyathena/polars/test_cursor.py | 2 +- tests/pyathena/s3fs/test_cursor.py | 2 +- tests/pyathena/test_cursor.py | 4 ++-- 5 files changed, 25 insertions(+), 17 deletions(-) diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 0cdd412a4..69cd7063c 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging -from collections.abc import Callable +from collections.abc import Callable, Iterator from typing import ( TYPE_CHECKING, Any, @@ -214,17 +214,25 @@ 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 [] - for i, d in enumerate(description): - convert = self._converter.get(d[1]) - if convert is _to_default: - continue - try: - # Unlike a list comprehension, a generator turns a StopIteration from the - # converter into a RuntimeError instead of ending an iteration of the rows. - columns[i] = list(convert(v) for v in columns[i]) # noqa: C400 - except DataError as e: - raise DataError(f"Column {d[0]!r}: {e}") from e - self._rows.extend(zip(*columns, strict=False)) + names = [d[0] for d in description] + converters = [self._converter.get(d[1]) for d in description] + + 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) def _open_csv(self, location: str) -> Any: """Open the CSV result file for ``pyarrow.csv.read_csv()``. diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index 222ebf844..7123c4a89 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -907,7 +907,7 @@ def test_ambiguous_time_zone_value(self, arrow_cursor): The table keeps the text, as for every TIMESTAMP WITH TIME ZONE value. """ arrow_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) - with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): + with pytest.raises(DataError, match=r"^Column 'col_timestamp_tz': .* is ambiguous"): arrow_cursor.fetchall() @pytest.mark.skipif(not ENV.managed_work_group, reason="AWS_ATHENA_MANAGED_WORKGROUP not set") diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index f786bc103..0c18ac4ce 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -892,7 +892,7 @@ def test_ambiguous_time_zone_value(self, polars_cursor): The table keeps the text, as for every TIMESTAMP WITH TIME ZONE value. """ polars_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) - with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): + with pytest.raises(DataError, match=r"^Column 'col_timestamp_tz': .* is ambiguous"): polars_cursor.fetchall() @pytest.mark.skipif(not ENV.managed_work_group, reason="AWS_ATHENA_MANAGED_WORKGROUP not set") diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index ce551c884..7e563f60f 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -323,7 +323,7 @@ def test_fetch_all_rows(self, s3fs_cursor): def test_ambiguous_time_zone_value(self, s3fs_cursor): """A local time that its zone repeats raises DataError instead of choosing one.""" s3fs_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) - with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): + with pytest.raises(DataError, match=r"^Column 'col_timestamp_tz': .* is ambiguous"): s3fs_cursor.fetchall() @pytest.mark.parametrize( diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 57584d9c4..3f76a15d2 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -1682,7 +1682,7 @@ def test_ambiguous_time_zone_value(self, cursor): The first page of rows is converted by execute(). """ - with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): + with pytest.raises(DataError, match=r"^Column 'col_timestamp_tz': .* is ambiguous"): cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) @pytest.mark.parametrize( @@ -1952,7 +1952,7 @@ def __init__(self, **kwargs): def test_ambiguous_time_zone_value(self, dict_cursor): """A local time that its zone repeats raises DataError naming the column.""" - with pytest.raises(DataError, match=r"Column 'col_timestamp_tz': .* is ambiguous"): + with pytest.raises(DataError, match=r"^Column 'col_timestamp_tz': .* is ambiguous"): dict_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) def test_null_vs_empty_string(self, dict_cursor): From 067a1eff81f3aa719fb1e25dbf1e6d4e5f25920b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 16:47:15 +0900 Subject: [PATCH 05/10] Name the column when pandas converts time zone values as it reads (GH-1217) With parse_dates or dtype given to PandasCursor.execute(), or with a custom timestamp with time zone conversion function, pandas applies the column's conversion function to each value as it reads. A value that failed to convert raised OperationalError without the column name in a whole read, and DataError without it in a chunked read. The conversion function now names the column, and the read lets its DataError through. Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 63 ++++++++++++++++++++++-- tests/pyathena/pandas/test_result_set.py | 15 +++++- 2 files changed, 71 insertions(+), 7 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index db8d7f926..43f6fbb2f 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -363,6 +363,48 @@ 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") + + 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 + + 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: + raise DataError(f"Column {self._label!r}: {e}") from e + + def _to_timestamps(values: Series, date: bool = False) -> Series: """Convert timestamp or date text to ``datetime64[us]``. @@ -1246,15 +1288,19 @@ 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. @@ -1262,6 +1308,8 @@ def _get_csv_converter(self, type_: str) -> Callable[[str | None], Any] | _JSONC 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 +1365,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 @@ -1413,6 +1463,9 @@ def _read_csv(self) -> TextFileReader | DataFrame: return result + except DataError: + # A value that fails to convert, whose converter names the column. + raise except Exception as e: _logger.exception(f"Failed to read {location}.") raise OperationalError(*e.args) from e @@ -1519,7 +1572,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 +1743,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 } diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index b29123c48..15fa26f2b 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1055,6 +1055,16 @@ def test_unload_temporal_options(self, tmp_path, kwargs): result_set = _local_result_set(LocalBucketFileSystem(tmp_path), tables=tables, **kwargs) assert result_set.as_pandas()["d"].dtype != "datetime64[us]" + @pytest.mark.parametrize( + "kwargs", + [ + pytest.param({}, id="default"), + # With dtype given, pandas converts each value as it reads. + pytest.param({"dtype": {}}, id="pandas_converts"), + # skiprows leaves the column labels to pandas. + pytest.param({"dtype": {}, "skiprows": 0}, id="pandas_converts_unlabeled"), + ], + ) @pytest.mark.parametrize("chunksize", [None, 1]) @pytest.mark.parametrize( ("value", "reason"), @@ -1065,11 +1075,11 @@ def test_unload_temporal_options(self, tmp_path, kwargs): ], ids=["ambiguous", "nonexistent", "unknown_zone"], ) - def test_time_zone_errors(self, tmp_path, chunksize, value, reason): + def test_time_zone_errors(self, tmp_path, kwargs, 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. + chunk with the value, also when pandas converts the values as it reads. """ data = f'"tstz"\n"2020-01-02 03:04:05.000 UTC"\n"{value}"\n' @@ -1079,6 +1089,7 @@ def read(): [("tstz", "timestamp with time zone")], data, chunksize=chunksize, + **kwargs, ) return list(result_set.iter_chunks()) From d555a2c29e617881f0e95c854942e080f2ed9afb Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 16:48:05 +0900 Subject: [PATCH 06/10] Describe the _TimeZoneConverter that _get_csv_converter() returns (GH-1217) Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 43f6fbb2f..ae10265e4 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -1303,7 +1303,8 @@ def _get_csv_converter( 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": From 4dd1dac1c2be3e895890ca0b399e1e7379a32ab3 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 16:56:38 +0900 Subject: [PATCH 07/10] Raise only the time zone conversion's DataError as it is (GH-1217) A DataError from a conversion function given in converters is still raised as OperationalError by a whole read. _TimeZoneConverter raises a private DataError subclass, which _read_csv() lets through. Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 20 ++++++++++++++------ tests/pyathena/pandas/test_result_set.py | 18 ++++++++++++++++++ 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index ae10265e4..298ef39d2 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -363,6 +363,15 @@ def _convert_values(self, values: Series) -> Series: return _to_timestamps(values, date=self._type == "date") +class _ColumnDataError(DataError): + """The ``DataError`` of a value that ``_TimeZoneConverter`` failed to convert. + + Its message names the column. Unlike an error from a conversion function given in + ``converters``, which a whole read raises as ``OperationalError``, a whole read + raises it as it is. + """ + + class _TimeZoneConverter: """A converter that ``pandas.read_csv()`` applies to each ``timestamp with time zone`` value. @@ -395,14 +404,14 @@ def __call__(self, value: str | None) -> Any: 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. + _ColumnDataError: 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: - raise DataError(f"Column {self._label!r}: {e}") from e + raise _ColumnDataError(f"Column {self._label!r}: {e}") from e def _to_timestamps(values: Series, date: bool = False) -> Series: @@ -1464,8 +1473,7 @@ def _read_csv(self) -> TextFileReader | DataFrame: return result - except DataError: - # A value that fails to convert, whose converter names the column. + except _ColumnDataError: raise except Exception as e: _logger.exception(f"Failed to read {location}.") diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 15fa26f2b..c5e9aa329 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1098,6 +1098,24 @@ def read(): assert str(error.value).startswith("Column 'tstz': ") assert reason in str(error.value) + @pytest.mark.filterwarnings("ignore:Both a converter and dtype:pandas.errors.ParserWarning") + def test_converter_data_error(self, tmp_path): + """A DataError from a conversion function given in converters is OperationalError. + + Only PyAthena's time zone conversion raises its DataError as it is. + """ + + def reject(value): + raise DataError("invalid number") + + with pytest.raises(OperationalError, match="invalid number"): + _local_result_set( + LocalBucketFileSystem(tmp_path), + [("n", "varchar")], + '"n"\n"bad"\n', + converters={"n": reject}, + ) + @pytest.mark.parametrize( ("kwargs", "converter_types", "expected"), [ From 8f6101f1a556c3340e790ad06231a4224fb20abf Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 16:56:52 +0900 Subject: [PATCH 08/10] Say that the converter error test reads the whole result (GH-1217) Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_result_set.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index c5e9aa329..915f2d89f 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1100,7 +1100,7 @@ def read(): @pytest.mark.filterwarnings("ignore:Both a converter and dtype:pandas.errors.ParserWarning") def test_converter_data_error(self, tmp_path): - """A DataError from a conversion function given in converters is OperationalError. + """A whole read raises a converters function's DataError as OperationalError. Only PyAthena's time zone conversion raises its DataError as it is. """ From 37e97783d41a57b8b9f75583b5249c9cf588b1e9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 17:37:39 +0900 Subject: [PATCH 09/10] Raise the time zone conversion's DataError as DataError (GH-1217) Instead of a private DataError subclass, _TimeZoneConverter keeps the DataError it raised, and _read_csv() raises the error as it is when one of its converters raised it. The error's class is DataError, as from the other paths. Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 30 ++++++++++++++---------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 298ef39d2..9f5e3fda3 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -363,15 +363,6 @@ def _convert_values(self, values: Series) -> Series: return _to_timestamps(values, date=self._type == "date") -class _ColumnDataError(DataError): - """The ``DataError`` of a value that ``_TimeZoneConverter`` failed to convert. - - Its message names the column. Unlike an error from a conversion function given in - ``converters``, which a whole read raises as ``OperationalError``, a whole read - raises it as it is. - """ - - class _TimeZoneConverter: """A converter that ``pandas.read_csv()`` applies to each ``timestamp with time zone`` value. @@ -382,7 +373,7 @@ class _TimeZoneConverter: ``DataError`` of a value that fails to convert, as ``_TemporalConverter`` does. """ - __slots__ = ("_converter", "_label") + __slots__ = ("_converter", "_label", "error") def __init__(self, converter: Callable[[str | None], Any], label: Any) -> None: """Wrap the conversion function of a column. @@ -393,6 +384,8 @@ def __init__(self, converter: Callable[[str | None], Any], label: Any) -> None: """ 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. @@ -404,14 +397,15 @@ def __call__(self, value: str | None) -> Any: The converted value. Raises: - _ColumnDataError: 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. + 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: - raise _ColumnDataError(f"Column {self._label!r}: {e}") from e + self.error = DataError(f"Column {self._label!r}: {e}") + raise self.error from e def _to_timestamps(values: Series, date: bool = False) -> Series: @@ -1473,9 +1467,13 @@ def _read_csv(self) -> TextFileReader | DataFrame: return result - except _ColumnDataError: - raise except Exception as e: + converters = read_csv_kwargs.get("converters") or {} + if 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 From df712324f396a0127c1c6ee815c9b2c0056da2b5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 17:45:02 +0900 Subject: [PATCH 10/10] Keep OperationalError for invalid converters and other DataErrors (GH-1217) _read_csv() checks the converters for the time zone conversion's error only when they are a mapping, so converters of another type given to execute() still raise OperationalError. The read error test covers a DataError from a mapped conversion function read with a time zone column, and invalid converters. Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 5 +++-- tests/pyathena/pandas/test_result_set.py | 26 +++++++++++++++++------- 2 files changed, 22 insertions(+), 9 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 9f5e3fda3..1c2690ac6 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -1468,8 +1468,9 @@ def _read_csv(self) -> TextFileReader | DataFrame: return result except Exception as e: - converters = read_csv_kwargs.get("converters") or {} - if any( + # 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() ): diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 915f2d89f..6aac375bf 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -1099,21 +1099,33 @@ def read(): assert reason in str(error.value) @pytest.mark.filterwarnings("ignore:Both a converter and dtype:pandas.errors.ParserWarning") - def test_converter_data_error(self, tmp_path): - """A whole read raises a converters function's DataError as OperationalError. + @pytest.mark.parametrize("case", ["mapping", "converters", "invalid_converters"]) + def test_read_error(self, tmp_path, case): + """A whole read raises an error other than the time zone conversion's as OperationalError. - Only PyAthena's time zone conversion raises its DataError as it is. + This includes a DataError from another conversion function, such as one that the + converter maps, which pandas applies with the time zone conversion when dtype + is given. """ def reject(value): raise DataError("invalid number") - with pytest.raises(OperationalError, match="invalid number"): + converter = DefaultPandasTypeConverter() + if case == "mapping": + converter.set("varchar", reject) + kwargs = { + "mapping": {"dtype": {}}, + "converters": {"converters": {"n": reject}}, + "invalid_converters": {"converters": [reject]}, + }[case] + with pytest.raises(OperationalError): _local_result_set( LocalBucketFileSystem(tmp_path), - [("n", "varchar")], - '"n"\n"bad"\n', - converters={"n": reject}, + [("tstz", "timestamp with time zone"), ("n", "varchar")], + '"tstz","n"\n"2020-01-02 03:04:05.000 UTC","bad"\n', + converter=converter, + **kwargs, ) @pytest.mark.parametrize(