Skip to content

Test the temporal types of each synchronous cursor #1226

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
laughingman7743 merged 3 commits into from
Oct 17, 2026
Merged

Test the temporal types of each synchronous cursor #1226

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5065d62
Test the temporal types of each synchronous cursor
laughingman7743 Oct 17, 2026
25569f1
Compare PandasCursor's aware values by instant
laughingman7743 Oct 17, 2026
66bfa6a
Run each temporal test once and cover more aware formats
laughingman7743 Oct 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions tests/pyathena/arrow/test_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.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(
"""
Expand Down
53 changes: 53 additions & 0 deletions tests/pyathena/pandas/test_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -42,6 +43,17 @@ def _pandas_converter_without_bigint_dtype():
return converter


def _pandas_value(column: temporal.Column):

@laughingman7743 laughingman7743 Oct 17, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round one (implementation, reworked design). Base 18c1eea1 (master), head 25569f1a065ff9c322eade38529dc1b297fe01f2. Result: FINDINGS. Two findings, repaired in 25569f1.

Inventory: tests/pyathena/temporal.py (columns, UNLOAD subset, query, expected row, zones), and test_temporal_types in the five sync cursor test classes.

  1. The pandas test checked the zone of fetched aware values. PandasCursor holds them in UTC (Define and enforce the contract for Athena temporal types across cursors #1212, Make PandasCursor's temporal dtypes and fetched types depend only on the Athena type #1219), so the check would fail once Make PandasCursor's temporal dtypes and fetched types depend only on the Athena type #1219 lands, although that behavior is correct. Repair: pandas compares aware values by instant only.
  2. The expected DataFrame values came from a nested conditional expression. Repair: the _pandas_value() helper.

Verified, no change needed:

  • Each xfail fails at its dtype check with --runxfail.
  • The UNLOAD form selects only TIMESTAMP(3) and date.
  • zones() compares ZoneInfo and timezone objects by equality, so timezone.utc and ZoneInfo('UTC') differ.
  • The NULL columns compare as None in fetch and as the container's null in frames.
  • pd.Timestamp == date is False, so a fetched Timestamp for date fails as intended.

"""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)]
Expand Down Expand Up @@ -509,6 +521,47 @@ 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.

Aware values are compared by instant, as the DataFrame holds them in UTC.
"""
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]] == [
_pandas_value(c) for c in columns
]
# The fetch methods read the same DataFrame, whose aware values are in UTC.
row = pandas_cursor.fetchone()

@laughingman7743 laughingman7743 Oct 17, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

3. [P2] The DataFrame tests ran the query twice (relayed), for the frame and again for the fetch methods.

Repair (66bfa6a): the DataFrame, Arrow, and Polars tests fetch from the execution they read. A probe confirmed that fetchone() after as_pandas() or as_arrow() returns the row, both with and without UNLOAD. Each test parameter now runs one query (8 per run).

Validation at 66bfa6a: just lint passed. Live Athena on Python 3.15 and 3.11: 3 passed and 5 xfailed. With --runxfail, each xfail fails at its dtype check.

assert row == temporal.fetched_row(unload)

@pytest.mark.parametrize(
("pandas_cursor", "chunksize"),
[
Expand Down
30 changes: 30 additions & 0 deletions tests/pyathena/polars/test_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.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(
"""
Expand Down
7 changes: 7 additions & 0 deletions tests/pyathena/s3fs/test_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
"""
Expand Down
192 changes: 192 additions & 0 deletions tests/pyathena/temporal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,192 @@
# 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_min",

@laughingman7743 laughingman7743 Oct 17, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2. [P2] Aware timestamps lacked years 1 and 9999 (relayed). A conversion through nanosecond storage would pass with 2020 values only.

Repair (66bfa6a): added tstz_min and tstz_max in UTC.

"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
),
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_p12",

@laughingman7743 laughingman7743 Oct 17, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent review (relayed). Reviewer: Codex CLI 0.162.1, codex exec -s read-only, static only. Snapshot: detached worktree at 25569f1a065ff9c322eade38529dc1b297fe01f2, base 18c1eea1. Prompt: the diff, the contract, and the repository conventions, without the PR text or the earlier findings. Coverage, as the reviewer reported it: all six files; SQL and expected values; precision and year boundaries; NULLs; pandas equality (Timestamp, date, NaT); container dtypes; fetched zones; xfail placement; and conventions. It found no contradiction between the SQL and the expected values. Result: FINDINGS (3).

1. [P2] time with time zone with 12 fractional digits was missing. A path that rejects or rounds TIME(12) WITH TIME ZONE would pass.

Repair (66bfa6a): added time_tz_p12. The fetch methods expect .123456+05:30, and Arrow and Polars expect Athena's text 03:04:05.123456789012+05:30.

@laughingman7743 laughingman7743 Oct 17, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up on the repair (relayed): CLEAN. Reviewer: Codex CLI 0.162.1, codex exec -s read-only, static only. Snapshot: 66bfa6a47079c7a34d25e8ac7d2df0b1deb53e29. Scope: the diff 25569f1a..66bfa6a4. All three findings are resolved. The new expected values match their SQL, and the Arrow and Polars text matches Athena's. The reviewer confirmed in the cursor code that unchunked pandas, Arrow, and Polars reads return cached frames without advancing their fetch iterators, with and without UNLOAD, so fetching from the same execution is valid.

"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),
]

# UNLOAD writes only TIMESTAMP(3) and date of these types.
UNLOAD_COLUMNS = [

@laughingman7743 laughingman7743 Oct 17, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round two (claims and operations). Head 25569f1a065ff9c322eade38529dc1b297fe01f2. Result: CLEAN after the round-one repairs.

Claims checked against evidence:

  • The gap table in the PR: measured per value on live Athena before the tests were written, and each xfail reproduces it with --runxfail.
  • "UNLOAD writes only TIMESTAMP(3) and date": the Define and enforce the contract for Athena temporal types across cursors #1212 facts. Other precisions failed with NOT_SUPPORTED in the earlier audit.
  • "The async cursors and managed results return the same values": the earlier measurement for Test the temporal types of each synchronous cursor #1226's first design found no difference between sync, futures-based, and asyncio forms, and managed GetQueryResults text equals the CSV text.
  • Python 3.11 with Polars 1.44.2 and pyarrow 25.0.1 gives the same 3 passed and 5 xfailed.

AWS operator: 14 queries per run, all on literals (8 SELECTs and 6 UNLOADs). No new fixtures or session resources.

*(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]
Loading
Loading