From 5065d6212ab7a5efd47f04c7185e5f782ef34f83 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 11:56:03 +0900 Subject: [PATCH 1/3] Test the temporal types of each synchronous cursor Each synchronous cursor runs one SELECT whose columns cover the temporal formats Athena writes: 0 to 12 fractional digits, years 1, 1600, 2300, and 9999, NULL, UTC, a named zone, an offset, dates, and times. Cursors that read UNLOAD results also run its UNLOAD form. The values, dtypes, and time zones are checked against the contract of #1212, and the current gaps are strict xfails that name the issue that fixes them. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_cursor.py | 36 ++++++ tests/pyathena/pandas/test_cursor.py | 44 +++++++ tests/pyathena/polars/test_cursor.py | 30 +++++ tests/pyathena/s3fs/test_cursor.py | 7 ++ tests/pyathena/temporal.py | 171 +++++++++++++++++++++++++++ tests/pyathena/test_cursor.py | 7 ++ 6 files changed, 295 insertions(+) create mode 100644 tests/pyathena/temporal.py diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index 180a1c281..c45c3b207 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -26,6 +26,7 @@ from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig from tests import ENV +from tests.pyathena import temporal from tests.pyathena.conftest import connect from tests.pyathena.util import ( AMBIGUOUS_TIME_ZONE_QUERY, @@ -150,6 +151,41 @@ def test_invalid_arraysize(self, arrow_cursor): with pytest.raises(ProgrammingError): arrow_cursor.arraysize = -1 + @pytest.mark.parametrize( + ("arrow_cursor", "unload"), + [ + pytest.param( + {"cursor_kwargs": {"unload": False}}, + False, + marks=pytest.mark.xfail(strict=True, reason="#1216"), + ), + pytest.param( + {"cursor_kwargs": {"unload": True}}, + True, + marks=pytest.mark.xfail(strict=True, reason="#1214"), + ), + ], + indirect=["arrow_cursor"], + ids=["csv", "unload"], + ) + def test_temporal_types(self, arrow_cursor, unload): + """Temporal values convert as the contract of #1212 defines. + + Types that Arrow cannot represent stay Athena's text. + """ + types = {"timestamp": pa.timestamp("us"), "date": pa.date32()} + columns = temporal.columns(unload) + table = arrow_cursor.execute(temporal.query(unload)).as_arrow() + assert {field.name: field.type for field in table.schema} == { + c.name: types.get(c.athena_type, pa.string()) for c in columns + } + assert [table.column(c.name)[0].as_py() for c in columns] == [ + c.value if c.athena_type in types else c.text for c in columns + ] + row = arrow_cursor.execute(temporal.query(unload)).fetchone() + assert row == temporal.fetched_row(unload) + assert temporal.zones(row) == temporal.zones(temporal.fetched_row(unload)) + def test_complex(self, arrow_cursor): arrow_cursor.execute( """ diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 0dbd79e72..8e3678344 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -26,6 +26,7 @@ ) from pyathena.util import RetryConfig from tests import ENV +from tests.pyathena import temporal from tests.pyathena.conftest import connect from tests.pyathena.util import ( AMBIGUOUS_TIME_ZONE_QUERY, @@ -509,6 +510,49 @@ def test_invalid_arraysize(self, pandas_cursor): with pytest.raises(ProgrammingError): pandas_cursor.arraysize = -1 + @pytest.mark.parametrize( + ("pandas_cursor", "unload"), + [ + pytest.param( + {"cursor_kwargs": {"unload": False}}, + False, + marks=pytest.mark.xfail(strict=True, reason="#1219"), + ), + pytest.param( + {"cursor_kwargs": {"unload": True}}, + True, + marks=pytest.mark.xfail(strict=True, reason="#1214, #1219"), + ), + ], + indirect=["pandas_cursor"], + ids=["csv", "unload"], + ) + def test_temporal_types(self, pandas_cursor, unload): + """Temporal values convert as the contract of #1212 defines.""" + dtypes = { + "timestamp": "datetime64[us]", + "date": "datetime64[us]", + "timestamp with time zone": "datetime64[us, UTC]", + "time": "object", + "time with time zone": "object", + } + columns = temporal.columns(unload) + df = pandas_cursor.execute(temporal.query(unload)).as_pandas() + assert {name: str(dtype) for name, dtype in df.dtypes.items()} == { + c.name: dtypes[c.athena_type] for c in columns + } + assert [None if pd.isna(v) else v for v in df.iloc[0]] == [ + pd.Timestamp(c.value).tz_convert("UTC") + if c.athena_type == "timestamp with time zone" and c.value is not None + else pd.Timestamp(c.value) + if c.athena_type in ("timestamp", "date") and c.value is not None + else c.value + for c in columns + ] + row = pandas_cursor.execute(temporal.query(unload)).fetchone() + assert row == temporal.fetched_row(unload) + assert temporal.zones(row) == temporal.zones(temporal.fetched_row(unload)) + @pytest.mark.parametrize( ("pandas_cursor", "chunksize"), [ diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index e84be4990..9e9759427 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -24,6 +24,7 @@ from pyathena.polars.result_set import AthenaPolarsResultSet from pyathena.util import RetryConfig from tests import ENV +from tests.pyathena import temporal from tests.pyathena.conftest import connect from tests.pyathena.util import ( AMBIGUOUS_TIME_ZONE_QUERY, @@ -148,6 +149,35 @@ def test_many_as_polars(self, polars_cursor): assert df.height == 10000 assert df.width == 1 + @pytest.mark.parametrize( + ("polars_cursor", "unload"), + [ + ({"cursor_kwargs": {"unload": False}}, False), + pytest.param( + {"cursor_kwargs": {"unload": True}}, + True, + marks=pytest.mark.xfail(strict=True, reason="#1214"), + ), + ], + indirect=["polars_cursor"], + ids=["csv", "unload"], + ) + def test_temporal_types(self, polars_cursor, unload): + """Temporal values convert as the contract of #1212 defines. + + Types that Polars cannot represent stay Athena's text. + """ + types = {"timestamp": pl.Datetime("us"), "date": pl.Date()} + columns = temporal.columns(unload) + df = polars_cursor.execute(temporal.query(unload)).as_polars() + assert dict(df.schema) == {c.name: types.get(c.athena_type, pl.String()) for c in columns} + assert [df[c.name][0] for c in columns] == [ + c.value if c.athena_type in types else c.text for c in columns + ] + row = polars_cursor.execute(temporal.query(unload)).fetchone() + assert row == temporal.fetched_row(unload) + assert temporal.zones(row) == temporal.zones(temporal.fetched_row(unload)) + def test_complex(self, polars_cursor): polars_cursor.execute( """ diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index fbc335864..09abbb4c5 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -18,6 +18,7 @@ from pyathena.s3fs.result_set import AthenaS3FSResultSet from pyathena.util import RetryConfig from tests import ENV +from tests.pyathena import temporal from tests.pyathena.conftest import connect from tests.pyathena.util import ( AMBIGUOUS_TIME_ZONE_QUERY, @@ -78,6 +79,12 @@ def test_invalid_arraysize(self, s3fs_cursor): with pytest.raises(ProgrammingError): s3fs_cursor.arraysize = -1 + def test_temporal_types(self, s3fs_cursor): + """Temporal values convert as the contract of #1212 defines.""" + row = s3fs_cursor.execute(temporal.query()).fetchone() + assert row == temporal.fetched_row() + assert temporal.zones(row) == temporal.zones(temporal.fetched_row()) + def test_complex(self, s3fs_cursor): s3fs_cursor.execute( """ diff --git a/tests/pyathena/temporal.py b/tests/pyathena/temporal.py new file mode 100644 index 000000000..f954ff943 --- /dev/null +++ b/tests/pyathena/temporal.py @@ -0,0 +1,171 @@ +# Copyright 2026 The PyAthena authors +# +# Licensed under the MIT License. +# See LICENSE or https://opensource.org/licenses/MIT. +# +# SPDX-License-Identifier: MIT +"""A query of Athena temporal values and what each cursor returns for them (#1212). + +Each column is one format. The fetch methods return ``value``, fractional digits +beyond six truncated. ArrowCursor and PolarsCursor keep the types that they cannot +represent as Athena's text, ``text``. +""" + +from datetime import date, datetime, time, timedelta, timezone +from typing import Any, NamedTuple +from zoneinfo import ZoneInfo + + +class Column(NamedTuple): + name: str + sql: str + athena_type: str + value: Any + text: str | None = None + + +_UTC = ZoneInfo("UTC") +_PLUS_0530 = timezone(timedelta(hours=5, minutes=30)) + +COLUMNS = [ + Column( + "ts_p0", + "CAST('2020-01-02 03:04:05' AS TIMESTAMP(0))", + "timestamp", + datetime(2020, 1, 2, 3, 4, 5), + ), + Column( + "ts_p3", + "TIMESTAMP '2020-01-02 03:04:05.123'", + "timestamp", + datetime(2020, 1, 2, 3, 4, 5, 123000), + ), + Column( + "ts_p6", + "CAST('2020-01-02 03:04:05.123456' AS TIMESTAMP(6))", + "timestamp", + datetime(2020, 1, 2, 3, 4, 5, 123456), + ), + Column( + "ts_p9", + "CAST('2020-01-02 03:04:05.123456789' AS TIMESTAMP(9))", + "timestamp", + datetime(2020, 1, 2, 3, 4, 5, 123456), + ), + Column( + "ts_p12", + "CAST('2020-01-02 03:04:05.123456789012' AS TIMESTAMP(12))", + "timestamp", + datetime(2020, 1, 2, 3, 4, 5, 123456), + ), + Column( + "ts_before_epoch", + "CAST('1969-12-31 23:59:59.987654321' AS TIMESTAMP(9))", + "timestamp", + datetime(1969, 12, 31, 23, 59, 59, 987654), + ), + Column("ts_min", "TIMESTAMP '0001-01-01 00:00:00.000'", "timestamp", datetime(1, 1, 1)), + Column( + "ts_1600", + "TIMESTAMP '1600-01-02 03:04:05.123'", + "timestamp", + datetime(1600, 1, 2, 3, 4, 5, 123000), + ), + Column( + "ts_2300", + "TIMESTAMP '2300-01-02 03:04:05.123'", + "timestamp", + datetime(2300, 1, 2, 3, 4, 5, 123000), + ), + Column( + "ts_max", + "CAST('9999-12-31 23:59:59.999999999999' AS TIMESTAMP(12))", + "timestamp", + datetime(9999, 12, 31, 23, 59, 59, 999999), + ), + Column("ts_null", "CAST(NULL AS TIMESTAMP(3))", "timestamp", None), + Column("date", "DATE '2020-01-02'", "date", date(2020, 1, 2)), + Column("date_min", "DATE '0001-01-01'", "date", date(1, 1, 1)), + Column("date_max", "DATE '9999-12-31'", "date", date(9999, 12, 31)), + Column("date_null", "CAST(NULL AS DATE)", "date", None), + Column( + "tstz_utc", + "TIMESTAMP '2020-01-02 03:04:05.123 UTC'", + "timestamp with time zone", + datetime(2020, 1, 2, 3, 4, 5, 123000, tzinfo=_UTC), + "2020-01-02 03:04:05.123 UTC", + ), + Column( + "tstz_zone", + "TIMESTAMP '2020-01-02 03:04:05.123 Asia/Tokyo'", + "timestamp with time zone", + datetime(2020, 1, 2, 3, 4, 5, 123000, tzinfo=ZoneInfo("Asia/Tokyo")), + "2020-01-02 03:04:05.123 Asia/Tokyo", + ), + Column( + "tstz_offset", + "TIMESTAMP '2020-01-02 03:04:05.123 +05:30'", + "timestamp with time zone", + datetime(2020, 1, 2, 3, 4, 5, 123000, tzinfo=_PLUS_0530), + "2020-01-02 03:04:05.123 +05:30", + ), + Column( + "tstz_p12", + "CAST('2020-01-02 03:04:05.123456789012 UTC' AS TIMESTAMP(12) WITH TIME ZONE)", + "timestamp with time zone", + datetime(2020, 1, 2, 3, 4, 5, 123456, tzinfo=_UTC), + "2020-01-02 03:04:05.123456789012 UTC", + ), + Column( + "tstz_null", "CAST(NULL AS TIMESTAMP(3) WITH TIME ZONE)", "timestamp with time zone", None + ), + Column("time_p3", "TIME '03:04:05.123'", "time", time(3, 4, 5, 123000), "03:04:05.123"), + Column( + "time_p12", + "CAST('03:04:05.123456789012' AS TIME(12))", + "time", + time(3, 4, 5, 123456), + "03:04:05.123456789012", + ), + Column("time_null", "CAST(NULL AS TIME(3))", "time", None), + Column( + "time_tz", + "TIME '03:04:05.123 +05:30'", + "time with time zone", + time(3, 4, 5, 123000, tzinfo=_PLUS_0530), + "03:04:05.123+05:30", + ), + Column("time_tz_null", "CAST(NULL AS TIME(3) WITH TIME ZONE)", "time with time zone", None), +] + +# UNLOAD writes only TIMESTAMP(3) and date of these types. +UNLOAD_COLUMNS = [ + *(c for c in COLUMNS if c.name in ("ts_p3", "ts_min", "ts_1600", "ts_2300", "ts_null")), + Column( + "ts_max", + "TIMESTAMP '9999-12-31 23:59:59.999'", + "timestamp", + datetime(9999, 12, 31, 23, 59, 59, 999000), + ), + *(c for c in COLUMNS if c.athena_type == "date"), +] + + +def columns(unload: bool = False) -> list[Column]: + """The columns of the query, or of its UNLOAD form.""" + return UNLOAD_COLUMNS if unload else COLUMNS + + +def query(unload: bool = False) -> str: + """The SELECT query of the columns.""" + return "SELECT " + ", ".join(f"{c.sql} AS {c.name}" for c in columns(unload)) + + +def fetched_row(unload: bool = False) -> tuple[Any, ...]: + """The row that the fetch methods return.""" + return tuple(c.value for c in columns(unload)) + + +def zones(row: tuple[Any, ...]) -> list[Any]: + """The time zone of each aware value: a fixed-offset timezone or a ZoneInfo.""" + return [v.tzinfo for v in row if isinstance(v, datetime | time) and v.tzinfo is not None] diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 50faeebf2..b5604e854 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -44,6 +44,7 @@ from pyathena.result_set import AthenaDictResultSet from pyathena.util import RetryConfig, stop_retries_on from tests import ENV +from tests.pyathena import temporal from tests.pyathena.conftest import connect from tests.pyathena.tables import TABLES, VIEWS from tests.pyathena.util import ( @@ -771,6 +772,12 @@ def test_query_execution_initial(self, cursor): assert cursor.selected_engine_version is None assert cursor.effective_engine_version is None + def test_temporal_types(self, cursor): + """Temporal values convert as the contract of #1212 defines.""" + row = cursor.execute(temporal.query()).fetchone() + assert row == temporal.fetched_row() + assert temporal.zones(row) == temporal.zones(temporal.fetched_row()) + def test_complex(self, cursor): cursor.execute( """ From 25569f1a065ff9c322eade38529dc1b297fe01f2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 11:57:45 +0900 Subject: [PATCH 2/3] Compare PandasCursor's aware values by instant PandasCursor holds aware values in UTC, so the zone of the fetched values is not checked, and the expected DataFrame values come from a helper instead of a nested conditional. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_cursor.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 8e3678344..a6d08deb0 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -43,6 +43,17 @@ def _pandas_converter_without_bigint_dtype(): return converter +def _pandas_value(column: temporal.Column): + """The value that a DataFrame holds for a temporal column.""" + if column.value is None: + return None + if column.athena_type == "timestamp with time zone": + return pd.Timestamp(column.value).tz_convert("UTC") + if column.athena_type in ("timestamp", "date"): + return pd.Timestamp(column.value) + return column.value + + class TestPandasCursor: @pytest.mark.parametrize( ("engine", "chunksize"), [("auto", None), ("c", 2), ("python", 2), ("pyarrow", None)] @@ -528,7 +539,10 @@ def test_invalid_arraysize(self, pandas_cursor): ids=["csv", "unload"], ) def test_temporal_types(self, pandas_cursor, unload): - """Temporal values convert as the contract of #1212 defines.""" + """Temporal values convert as the contract of #1212 defines. + + Aware values are compared by instant, as the DataFrame holds them in UTC. + """ dtypes = { "timestamp": "datetime64[us]", "date": "datetime64[us]", @@ -542,16 +556,11 @@ def test_temporal_types(self, pandas_cursor, unload): c.name: dtypes[c.athena_type] for c in columns } assert [None if pd.isna(v) else v for v in df.iloc[0]] == [ - pd.Timestamp(c.value).tz_convert("UTC") - if c.athena_type == "timestamp with time zone" and c.value is not None - else pd.Timestamp(c.value) - if c.athena_type in ("timestamp", "date") and c.value is not None - else c.value - for c in columns + _pandas_value(c) for c in columns ] + # The fetch methods read the DataFrame, whose aware values are in UTC. row = pandas_cursor.execute(temporal.query(unload)).fetchone() assert row == temporal.fetched_row(unload) - assert temporal.zones(row) == temporal.zones(temporal.fetched_row(unload)) @pytest.mark.parametrize( ("pandas_cursor", "chunksize"), From 66bfa6a47079c7a34d25e8ac7d2df0b1deb53e29 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 12:07:22 +0900 Subject: [PATCH 3/3] Run each temporal test once and cover more aware formats - Fetch from the execution whose DataFrame or table was checked, so each test parameter runs one query. - Add timestamp with time zone values in years 1 and 9999, and a time with time zone value with 12 fractional digits. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_cursor.py | 2 +- tests/pyathena/pandas/test_cursor.py | 4 ++-- tests/pyathena/polars/test_cursor.py | 2 +- tests/pyathena/temporal.py | 21 +++++++++++++++++++++ 4 files changed, 25 insertions(+), 4 deletions(-) diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index c45c3b207..a8d48d554 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -182,7 +182,7 @@ def test_temporal_types(self, arrow_cursor, unload): assert [table.column(c.name)[0].as_py() for c in columns] == [ c.value if c.athena_type in types else c.text for c in columns ] - row = arrow_cursor.execute(temporal.query(unload)).fetchone() + row = arrow_cursor.fetchone() assert row == temporal.fetched_row(unload) assert temporal.zones(row) == temporal.zones(temporal.fetched_row(unload)) diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index a6d08deb0..e4fd7fc9f 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -558,8 +558,8 @@ def test_temporal_types(self, pandas_cursor, unload): assert [None if pd.isna(v) else v for v in df.iloc[0]] == [ _pandas_value(c) for c in columns ] - # The fetch methods read the DataFrame, whose aware values are in UTC. - row = pandas_cursor.execute(temporal.query(unload)).fetchone() + # The fetch methods read the same DataFrame, whose aware values are in UTC. + row = pandas_cursor.fetchone() assert row == temporal.fetched_row(unload) @pytest.mark.parametrize( diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 9e9759427..2d0cb4631 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -174,7 +174,7 @@ def test_temporal_types(self, polars_cursor, unload): assert [df[c.name][0] for c in columns] == [ c.value if c.athena_type in types else c.text for c in columns ] - row = polars_cursor.execute(temporal.query(unload)).fetchone() + row = polars_cursor.fetchone() assert row == temporal.fetched_row(unload) assert temporal.zones(row) == temporal.zones(temporal.fetched_row(unload)) diff --git a/tests/pyathena/temporal.py b/tests/pyathena/temporal.py index f954ff943..bc7eb6aa2 100644 --- a/tests/pyathena/temporal.py +++ b/tests/pyathena/temporal.py @@ -116,6 +116,20 @@ class Column(NamedTuple): datetime(2020, 1, 2, 3, 4, 5, 123456, tzinfo=_UTC), "2020-01-02 03:04:05.123456789012 UTC", ), + Column( + "tstz_min", + "CAST('0001-01-01 00:00:00.000 UTC' AS TIMESTAMP(3) WITH TIME ZONE)", + "timestamp with time zone", + datetime(1, 1, 1, tzinfo=_UTC), + "0001-01-01 00:00:00.000 UTC", + ), + Column( + "tstz_max", + "CAST('9999-12-31 23:59:59.999 UTC' AS TIMESTAMP(3) WITH TIME ZONE)", + "timestamp with time zone", + datetime(9999, 12, 31, 23, 59, 59, 999000, tzinfo=_UTC), + "9999-12-31 23:59:59.999 UTC", + ), Column( "tstz_null", "CAST(NULL AS TIMESTAMP(3) WITH TIME ZONE)", "timestamp with time zone", None ), @@ -135,6 +149,13 @@ class Column(NamedTuple): time(3, 4, 5, 123000, tzinfo=_PLUS_0530), "03:04:05.123+05:30", ), + Column( + "time_tz_p12", + "CAST('03:04:05.123456789012 +05:30' AS TIME(12) WITH TIME ZONE)", + "time with time zone", + time(3, 4, 5, 123456, tzinfo=_PLUS_0530), + "03:04:05.123456789012+05:30", + ), Column("time_tz_null", "CAST(NULL AS TIME(3) WITH TIME ZONE)", "time with time zone", None), ]