diff --git a/docs/pandas.md b/docs/pandas.md index 3edc0dda..20387867 100644 --- a/docs/pandas.md +++ b/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. diff --git a/docs/s3fs.md b/docs/s3fs.md index eb1f5370..31e44d73 100644 --- a/docs/s3fs.md +++ b/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) | diff --git a/pyathena/arrow/converter.py b/pyathena/arrow/converter.py index 88943efb..7da7bcef 100644 --- a/pyathena/arrow/converter.py +++ b/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(), diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 1acd654b..64f65040 100644 --- a/pyathena/arrow/result_set.py +++ b/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 [] diff --git a/pyathena/converter.py b/pyathena/converter.py index eff98742..e573a365 100644 --- a/pyathena/converter.py +++ b/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 diff --git a/pyathena/pandas/converter.py b/pyathena/pandas/converter.py index a3b207e9..1db9b2d7 100644 --- a/pyathena/pandas/converter.py +++ b/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, } diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 63864272..2a78aeef 100644 --- a/pyathena/pandas/result_set.py +++ b/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 diff --git a/pyathena/parser.py b/pyathena/parser.py index 1694cebc..120e2998 100644 --- a/pyathena/parser.py +++ b/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: diff --git a/pyathena/polars/converter.py b/pyathena/polars/converter.py index d0b2caf6..4d46c702 100644 --- a/pyathena/polars/converter.py +++ b/pyathena/polars/converter.py @@ -18,9 +18,11 @@ Converter, _to_binary, _to_date, + _to_datetime_with_tz, _to_default, _to_json, _to_time, + _to_time_with_tz, ) from pyathena.util import override @@ -30,6 +32,8 @@ _DEFAULT_POLARS_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, "varbinary": _to_binary, "json": _to_json, } @@ -89,6 +93,8 @@ def _dtypes(self) -> dict[str, Any]: "timestamp": pl.Datetime, "date": pl.Date, "time": pl.String, + "time with time zone": pl.String, + "timestamp with time zone": pl.String, "varbinary": pl.String, "array": pl.String, "map": pl.String, diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 4641614a..7f9bea74 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -20,7 +20,7 @@ ) from pyathena import OperationalError -from pyathena.converter import Converter, _json_text_converter +from pyathena.converter import Converter, _text_value_converter from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.polars.util import to_column_info @@ -277,8 +277,9 @@ def __init__( self._df_iter = self._create_dataframe_iterator() elif self.state == AthenaQueryExecution.STATE_SUCCEEDED: self._df = self._as_polars_from_api() - # GetQueryResults values are already converted, except json values kept as text. - self._df_converters = self._json_converters(self.converters) + # GetQueryResults values are already converted, except json and time with + # time zone values kept as text. + self._df_converters = self._text_value_converters(self.converters) else: self._df = pl.DataFrame() if self._df is not None: @@ -582,12 +583,12 @@ def _as_polars_from_api(self, converter: Converter | None = None) -> pl.DataFram 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. A Polars ``Time`` has no time zone. """ import polars as pl - rows = self._fetch_all_rows(converter or _json_text_converter()) + rows = self._fetch_all_rows(converter or _text_value_converter()) if not rows: return pl.DataFrame() description = self.description if self.description else [] diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 65842761..9ecc9187 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -12,7 +12,7 @@ ) from pyathena.common import BaseCursor, CursorIterator -from pyathena.converter import Converter, DefaultTypeConverter +from pyathena.converter import _TEXT_VALUE_TYPES, Converter, DefaultTypeConverter from pyathena.error import DataError, OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig, override, parse_output_location, retry_api_call @@ -681,19 +681,19 @@ def _is_first_row_column_labels(self, rows: list[dict[str, Any]]) -> bool: return False return True - def _json_converters( + def _text_value_converters( self, converters: dict[str, Callable[[str | None], Any | None]] ) -> dict[str, Callable[[str | None], Any | None]]: - """Select the converters of the json columns. + """Select the converters of the columns that the fallbacks keep as text. Args: converters: The converters keyed by column name. Returns: - The converters of the columns whose Athena type is json. + The converters of the columns whose Athena type is in ``_TEXT_VALUE_TYPES``. """ description = self.description if self.description else [] - return {d[0]: converters[d[0]] for d in description if d[1] == "json"} + return {d[0]: converters[d[0]] for d in description if d[1] in _TEXT_VALUE_TYPES} def _fetch_all_rows( self, diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index fd3f04b4..055d47b8 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -19,6 +19,7 @@ import pyarrow as pa import pytest +from pyathena.arrow.converter import DefaultArrowTypeConverter from pyathena.arrow.cursor import ArrowCursor from pyathena.arrow.result_set import AthenaArrowResultSet from pyathena.error import DatabaseError, ProgrammingError @@ -26,6 +27,7 @@ from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect +from tests.pyathena.util import CONVERTED_VALUES_QUERY, CONVERTED_VALUES_ROW class TestArrowCursor: @@ -972,6 +974,19 @@ def test_null_vs_empty_string(self, arrow_cursor): assert values[3] == "N/A" assert values[4] == "NULL" + def test_fetch_converts_each_value_once(self): + """The fetch methods call the converter once per value, after a set() too.""" + calls = [] + converter = DefaultArrowTypeConverter() + with ( + contextlib.closing(connect()) as conn, + conn.cursor(ArrowCursor, converter=converter) as cursor, + ): + cursor.execute("SELECT * FROM (VALUES 'a', 'b') AS t(v) ORDER BY v") + converter.set("varchar", lambda value: calls.append(value) or value) + assert cursor.fetchall() == [("a",), ("b",)] + assert calls == ["a", "b"] + @pytest.mark.parametrize( "arrow_cursor", [ @@ -997,18 +1012,36 @@ def test_fetch_all_rows(self, arrow_cursor): ,json_parse('{"a": 1}') AS col_json ,CAST('{"a": 1}' AS JSON) AS col_json_string ,CAST(NULL AS JSON) AS col_json_null + ,CAST(NULL AS TIME) AS col_time_null UNION ALL SELECT - 2, CAST('12:34:56' AS TIME), X'0102', json_parse('[1, "x"]'), json_parse('"s"'), NULL + 2, CAST('12:34:56' AS TIME), X'0102', json_parse('[1, "x"]'), json_parse('"s"'), NULL, + NULL ORDER BY col """ ) assert arrow_cursor.fetchall() == [ - (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x01\x02", {"a": 1}, '{"a": 1}', None), - (2, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x01\x02", [1, "x"], "s", None), + ( + 1, + datetime(2017, 1, 1, 12, 34, 56).time(), + b"\x01\x02", + {"a": 1}, + '{"a": 1}', + None, + None, + ), + (2, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x01\x02", [1, "x"], "s", None, None), ] assert arrow_cursor.as_arrow().schema.field("col_json").type == pa.string() + arrow_cursor.execute(CONVERTED_VALUES_QUERY) + assert arrow_cursor.fetchall() == [CONVERTED_VALUES_ROW] + # An Arrow time type has no time zone, so the table keeps the text. + assert arrow_cursor.as_arrow().column("col_time_tz").to_pylist() == ["12:34:56.789+09:00"] + assert arrow_cursor.as_arrow().column("col_timestamp_tz").to_pylist() == [ + "2024-02-29 23:59:58.123 +05:30" + ] + @pytest.mark.parametrize( "execute_kwargs", [{}, {"connect_timeout": 3.0, "request_timeout": 4.0}] ) diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 7c646ded..a6916536 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -23,7 +23,7 @@ from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect -from tests.pyathena.util import cached_file_systems +from tests.pyathena.util import CONVERTED_VALUES_QUERY, CONVERTED_VALUES_ROW, cached_file_systems class TestPandasCursor: @@ -1617,6 +1617,9 @@ def test_fetch_all_rows(self, pandas_cursor): (2, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x00\x01", [1, "x"], "s", None), ] + pandas_cursor.execute(CONVERTED_VALUES_QUERY) + assert pandas_cursor.fetchall() == [CONVERTED_VALUES_ROW] + @pytest.mark.parametrize( "execute_kwargs", [ diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 7988a29c..a943d268 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -24,7 +24,7 @@ from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect -from tests.pyathena.util import cached_file_systems +from tests.pyathena.util import CONVERTED_VALUES_QUERY, CONVERTED_VALUES_ROW, cached_file_systems class TestPolarsCursor: @@ -779,6 +779,14 @@ def test_fetch_all_rows(self, polars_cursor): ] assert polars_cursor.as_polars()["col_json"].dtype == pl.String + polars_cursor.execute(CONVERTED_VALUES_QUERY) + assert polars_cursor.fetchall() == [CONVERTED_VALUES_ROW] + # A Polars Time has no time zone, so the DataFrame keeps the text. + assert polars_cursor.as_polars()["col_time_tz"].to_list() == ["12:34:56.789+09:00"] + assert polars_cursor.as_polars()["col_timestamp_tz"].to_list() == [ + "2024-02-29 23:59:58.123 +05:30" + ] + @pytest.mark.parametrize( "execute_kwargs", [{}, {"block_size": 2048, "cache_type": "none", "max_workers": 3, "chunksize": 20}], diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index 4ebaa413..4b643d5b 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -19,7 +19,7 @@ from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect -from tests.pyathena.util import cached_file_systems +from tests.pyathena.util import CONVERTED_VALUES_QUERY, CONVERTED_VALUES_ROW, cached_file_systems def _s3fs_converter_with_json_text(): @@ -537,8 +537,8 @@ def test_quoted_string_with_comma(self, csv_reader_class): indirect=["s3fs_cursor"], ) def test_fetch_all_rows(self, s3fs_cursor): - s3fs_cursor.execute("SELECT 1 AS col") - assert s3fs_cursor.fetchall() == [(1,)] + s3fs_cursor.execute(CONVERTED_VALUES_QUERY) + assert s3fs_cursor.fetchall() == [CONVERTED_VALUES_ROW] @pytest.mark.parametrize( "s3fs_cursor", diff --git a/tests/pyathena/test_converter.py b/tests/pyathena/test_converter.py index 5d7ad401..08092519 100644 --- a/tests/pyathena/test_converter.py +++ b/tests/pyathena/test_converter.py @@ -1,16 +1,19 @@ -from datetime import datetime +from datetime import datetime, time, timedelta, timezone import pytest from dateutil.tz import gettz from pyathena.converter import ( DefaultTypeConverter, - _csv_to_json, + _text_value_converter, _to_array, _to_datetime, _to_datetime_with_tz, + _to_json, _to_map, _to_struct, + _to_time, + _to_time_with_tz, ) @@ -326,22 +329,6 @@ def test_to_array_invalid_formats(input_value): assert _to_array(input_value) is None -@pytest.mark.parametrize( - ("input_value", "expected"), - [ - (None, None), - ("", None), - ('{"a":1}', {"a": 1}), - ("[1,2]", [1, 2]), - ('""', ""), - ('"[1, 2]"', "[1, 2]"), - ("null", None), - ], -) -def test_csv_to_json(input_value, expected): - assert _csv_to_json(input_value) == expected - - @pytest.mark.parametrize( ("value", "type_hint", "expected"), [ @@ -590,3 +577,152 @@ def test_normalize_hive_syntax_mixed(self): type_hint="array", ) assert result == [{"a": 1, "b": "hello"}] + + +@pytest.mark.parametrize( + ("input_value", "expected"), + [ + (None, None), + ("", None), + ("12:34:56", time(12, 34, 56)), + ("12:34:56.1", time(12, 34, 56, 100000)), + ("12:34:56.123", time(12, 34, 56, 123000)), + ("12:34:56.123456", time(12, 34, 56, 123456)), + ("12:34:56.123456789012", time(12, 34, 56, 123456)), + ], +) +def test_to_time_any_precision(input_value, expected): + assert _to_time(input_value) == expected + + +@pytest.mark.parametrize( + ("input_value", "expected"), + [ + (None, None), + ("", None), + ("12:34:56+09:00", time(12, 34, 56, tzinfo=timezone(timedelta(hours=9)))), + ( + "12:34:56.789-05:30", + time(12, 34, 56, 789000, tzinfo=timezone(-timedelta(hours=5, minutes=30))), + ), + ("00:00:00.123456789012+00:00", time(0, 0, 0, 123456, tzinfo=timezone(timedelta(0)))), + ("23:59:59.9-14:00", time(23, 59, 59, 900000, tzinfo=timezone(-timedelta(hours=14)))), + ], +) +def test_to_time_with_tz(input_value, expected): + result = _to_time_with_tz(input_value) + assert result == expected + if expected is not None: + assert result.utcoffset() == expected.utcoffset() + + +@pytest.mark.parametrize( + ("input_value", "expected"), + [ + (None, None), + ("", None), + ('""', ""), + ('"[1, 2]"', "[1, 2]"), + ('{"a": 1}', {"a": 1}), + ("[1, 2]", [1, 2]), + ("null", None), + ], +) +def test_to_json(input_value, expected): + assert _to_json(input_value) == expected + + +@pytest.mark.parametrize( + ("type_hint", "value", "expected"), + [ + ("array(json)", '[""]', [""]), + ( + "array(json)", + '[{"a":1}, "x", 1, true, null, ""]', + [{"a": 1}, "x", 1, True, None, ""], + ), + # JSON string scalars whose text looks like JSON stay strings. + ( + "array(json)", + '["{\\"a\\": 1}", "123", "true", "null"]', + ['{"a": 1}', "123", "true", "null"], + ), + ( + "map(varchar,json)", + '{"k": "", "n": 1, "b": true, "z": null}', + {"k": "", "n": 1, "b": True, "z": None}, + ), + ("row(a json, b json)", '{"a": "x", "b": {"c": 1}}', {"a": "x", "b": {"c": 1}}), + ("row(a json, b json)", '{"a": "", "b": true}', {"a": "", "b": True}), + ("array(varchar)", '["a", "123"]', ["a", "123"]), + ], +) +def test_typed_json_elements(type_hint, value, expected): + """JSON elements of typed complex values decode their original JSON text.""" + type_ = type_hint.split("(", 1)[0] + assert DefaultTypeConverter().convert(type_, value, type_hint=type_hint) == expected + + +def test_typed_time_with_tz_elements(): + """Parameterized time zone types in type hints keep their time zone.""" + converter = DefaultTypeConverter() + jst = timezone(timedelta(hours=9)) + assert converter.convert( + "array", "[12:34:56.789+09:00, null]", type_hint="array(time(3) with time zone)" + ) == [time(12, 34, 56, 789000, tzinfo=jst), None] + assert converter.convert( + "map", "{a=12:34:56+09:00}", type_hint="map(varchar, time(0) with time zone)" + ) == {"a": time(12, 34, 56, tzinfo=jst)} + + +def test_text_value_converter(): + """The fallback converter keeps text values and nested time zones as text.""" + converter = _text_value_converter() + assert converter.convert("json", '{"a": 1}') == '{"a": 1}' + assert converter.convert("time with time zone", "12:34:56+09:00") == "12:34:56+09:00" + assert converter.convert( + "array", "[12:34:56.789+09:00]", type_hint="array(time with time zone)" + ) == ["12:34:56.789+09:00"] + assert converter.convert("array", '[{"a": 1}]', type_hint="array(json)") == [{"a": 1}] + # Other converters keep the default mappings. + assert DefaultTypeConverter().convert( + "array", "[12:34:56.789+09:00]", type_hint="array(time with time zone)" + ) == [time(12, 34, 56, 789000, tzinfo=timezone(timedelta(hours=9)))] + + +@pytest.mark.parametrize( + ("input_value", "expected"), + [ + (None, None), + ("", None), + ( + "2024-02-29 23:59:58.123 +05:30", + datetime( + 2024, 2, 29, 23, 59, 58, 123000, tzinfo=timezone(timedelta(hours=5, minutes=30)) + ), + ), + ( + "2024-02-29 23:59:58.123456 -08:00", + datetime(2024, 2, 29, 23, 59, 58, 123456, tzinfo=timezone(-timedelta(hours=8))), + ), + ( + "2024-02-29 23:59:58 +00:00", + datetime(2024, 2, 29, 23, 59, 58, tzinfo=timezone(timedelta(0))), + ), + ( + "2024-02-29 23:59:58.123 UTC", + datetime(2024, 2, 29, 23, 59, 58, 123000, tzinfo=gettz("UTC")), + ), + ( + "2024-02-29 23:59:58.123 America/New_York", + datetime(2024, 2, 29, 23, 59, 58, 123000, tzinfo=gettz("America/New_York")), + ), + ], +) +def test_to_datetime_with_tz_offsets_and_zone_names(input_value, expected): + """Numeric UTC offsets give fixed-offset time zones; zone names keep their zone.""" + result = _to_datetime_with_tz(input_value) + assert result == expected + if expected is not None: + assert result.utcoffset() == expected.utcoffset() + assert result.tzinfo is not None diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index f4a6b552..6b012f7f 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -41,6 +41,8 @@ from tests.pyathena.conftest import connect from tests.pyathena.tables import TABLES, VIEWS from tests.pyathena.util import ( + CONVERTED_VALUES_QUERY, + CONVERTED_VALUES_ROW, EVENT_TIMEOUT, interrupt_start_waits, succeeded_query_execution, @@ -1688,8 +1690,8 @@ def test_null_vs_empty_string(self, cursor): indirect=["cursor"], ) def test_fetch_all_rows(self, cursor): - cursor.execute("SELECT 1 AS col") - assert cursor.fetchall() == [(1,)] + cursor.execute(CONVERTED_VALUES_QUERY) + assert cursor.fetchall() == [CONVERTED_VALUES_ROW] @staticmethod def _metadata_view(metadata): diff --git a/tests/pyathena/test_parser.py b/tests/pyathena/test_parser.py index 00b03cca..dd3b827f 100644 --- a/tests/pyathena/test_parser.py +++ b/tests/pyathena/test_parser.py @@ -23,6 +23,18 @@ def test_simple_type(self): assert node.children == [] assert node.field_names is None + @pytest.mark.parametrize( + ("type_str", "expected"), + [ + ("time(3) with time zone", "time with time zone"), + ("TIMESTAMP(6) WITH TIME ZONE", "timestamp with time zone"), + ("decimal(10, 2)", "decimal"), + ("varchar(255)", "varchar"), + ], + ) + def test_parameterized_type_keeps_suffix(self, type_str, expected): + assert TypeSignatureParser().parse(type_str).type_name == expected + def test_simple_type_case_insensitive(self): parser = TypeSignatureParser() node = parser.parse("VARCHAR") diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index b2e0cf40..b8479d36 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -7,11 +7,13 @@ import time from concurrent.futures import wait +from datetime import datetime, timedelta, timezone from pathlib import Path from unittest.mock import patch from botocore.config import Config from botocore.exceptions import ClientError +from dateutil.tz import gettz from jinja2 import Environment, FileSystemLoader from sqlalchemy import types @@ -107,6 +109,38 @@ def unreachable_glue(connection): ) +# TIME values of several precisions, with and without a time zone, TIMESTAMP WITH TIME +# ZONE values with UTC offsets and a zone name, a NULL JSON value, and the row that +# every cursor should fetch for them. +CONVERTED_VALUES_QUERY = """ +SELECT + 1 AS col + ,CAST('12:34:56' AS TIME(0)) AS col_time_0 + ,CAST('12:34:56.123456789' AS TIME(9)) AS col_time_9 + ,CAST('12:34:56.789 +09:00' AS TIME WITH TIME ZONE) AS col_time_tz + ,CAST('12:34:56 -05:30' AS TIME(0) WITH TIME ZONE) AS col_time_tz_0 + ,CAST(NULL AS TIME WITH TIME ZONE) AS col_time_tz_null + ,CAST(NULL AS JSON) AS col_json_null + ,TIMESTAMP '2024-02-29 23:59:58.123 +05:30' AS col_timestamp_tz + ,TIMESTAMP '2024-02-29 23:59:58.123 -08:00' AS col_timestamp_tz_negative + ,TIMESTAMP '2024-02-29 23:59:58.123 America/New_York' AS col_timestamp_tz_name + ,CAST(NULL AS TIMESTAMP WITH TIME ZONE) AS col_timestamp_tz_null +""" +CONVERTED_VALUES_ROW = ( + 1, + datetime(2000, 1, 1, 12, 34, 56).time(), + datetime(2000, 1, 1, 12, 34, 56, 123456).time(), + datetime(2000, 1, 1, 12, 34, 56, 789000, tzinfo=timezone(timedelta(hours=9))).timetz(), + datetime(2000, 1, 1, 12, 34, 56, tzinfo=timezone(-timedelta(hours=5, minutes=30))).timetz(), + None, + None, + datetime(2024, 2, 29, 23, 59, 58, 123000, tzinfo=timezone(timedelta(hours=5, minutes=30))), + datetime(2024, 2, 29, 23, 59, 58, 123000, tzinfo=timezone(-timedelta(hours=8))), + datetime(2024, 2, 29, 23, 59, 58, 123000, tzinfo=gettz("America/New_York")), + None, +) + + # A Spark job whose executor tasks sleep, so that StopCalculationExecution can cancel it. CANCELABLE_SPARK_JOB = """ import time