diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 869a185b4..69cd7063c 100644 --- a/pyathena/arrow/result_set.py +++ b/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) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index db8d7f926..1c2690ac6 100644 --- a/pyathena/pandas/result_set.py +++ b/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 } diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index d4017a402..c3b31dfba 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,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) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 97184e012..9c4d870d6 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,79 @@ 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. + """ + # 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 = ( + 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] ) -> tuple[list[dict[str, Any]], str | None]: @@ -956,48 +999,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..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="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/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index b29123c48..6aac375bf 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()) @@ -1087,6 +1098,36 @@ 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") + @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. + + 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") + + 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), + [("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( ("kwargs", "converter_types", "expected"), [ diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 1b1985818..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): + 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..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): + 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..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="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.