Add a temporal benchmark shape derived from the snapshot timestamp (GH-1287)

benchmarks/DATASET.md

@@ -111,7 +111,7 @@ The main groups are:
 
 The exported probe file stores `timestamp` as physical Parquet `INT64` with a microsecond timestamp annotation (`isAdjustedToUTC=false`).
 The registered Athena table deliberately exposes its integer values as `bigint`; the original probe validation checked these values against the exported file.
-The benchmark projections preserve this integer type, so this workload does not benchmark native timestamp conversion.
+The benchmark projections preserve this integer type; only the temporal projection, described in [Tables used for measurements](#tables-used-for-measurements), adds temporal columns derived from it.
 For a human-readable UTC display, use `from_unixtime(CAST("timestamp" AS double) / 1000000, 'UTC')`; retain the integer column when exact microsecond values matter.
 Nested fields can be null and remain nested in the exported files.
 
@@ -233,6 +233,9 @@ Reuse a prepared manifest and its tables when comparing revisions.
 
 The flat projection selects 15 scalar columns, including selected fields from `file`, `details`, and `http`.
 The nested projection selects nine top-level columns and preserves all three structures.
-Both omit `download_date` from the measured result and read the same prepared table without a further `LIMIT`; their definitions are in [config.py](pyathena_bench/config.py).
+The json projection adds `details` as a column of Athena's `json` type to the flat projection.
+The temporal projection adds six columns derived from `from_unixtime("timestamp" / 1e6)` to the flat projection, at millisecond precision: a `timestamp(3)`, a `date`, `timestamp(3) with time zone` columns in `UTC`, `Asia/Tokyo`, and `+09:00`, and a `time(3)`.
+Its UNLOAD and CTAS reads add only the `timestamp(3)` and `date` columns, because UNLOAD and CTAS cannot write the other types.
+All projections omit `download_date` from the measured result and read the same prepared table without a further `LIMIT`; their definitions are in [config.py](pyathena_bench/config.py).
 Full cleanup removes these scratch tables and their data, while `cleanup --trials-only` preserves them for another run.
 Neither mode deletes the original source table or its S3 partition.

benchmarks/README.md

@@ -31,6 +31,8 @@ The Athena timings describe queries against these smaller snapshots, not a scan
 The flat workload projects scalar fields from `file`, `details`, and `http` alongside the top-level scalar columns.
 The nested workload preserves those structures.
 The json workload adds `CAST(details AS JSON) AS details_json` to the flat projection; it covers native pandas and Wrangler reads in the single-query suite, with their UNLOAD and CTAS reads marked unsupported, because Parquet cannot hold the `json` type.
+The temporal workload adds columns derived from `"timestamp"` with `from_unixtime()` to the flat projection: `event_time` (`timestamp(3)`), `event_date` (`date`), `event_time_utc`, `event_time_zone`, and `event_time_offset` (`timestamp(3) with time zone` in `UTC`, `Asia/Tokyo`, and `+09:00`), and `event_clock` (`time(3)`).
+It covers every family in the single-query suite; its UNLOAD and CTAS reads add only `event_time` and `event_date`, because UNLOAD and CTAS cannot write `time` or `with time zone` types.
 No workload converts every library's output into a common Python object representation.
 
 ## Dependencies and local checks
@@ -176,6 +178,9 @@ uv run --no-sync python -m pyathena_bench run \
 uv run --no-sync python -m pyathena_bench run \
   --manifest results/input.json --out results/nested-small \
   --suite single --scale small --shape nested --output-kind native
+uv run --no-sync python -m pyathena_bench run \
+  --manifest results/input.json --out results/temporal-small \
+  --suite single --scale small --shape temporal
 ```
 
 The concurrency suite compares ThreadPool and native asyncio APIs at 1, 10, 50, and 100 simultaneous queries by default.
@@ -300,7 +305,7 @@ When the temporary directory is a tmpfs, as `/tmp` is on Amazon Linux 2023, file
 | Arrow CSV and UNLOAD | Eager table loading; arraysize changes row extraction, not S3 streaming |
 | Polars CSV and UNLOAD | Full result or chunk iterator |
 | SQLAlchemy (`sqla_*`) | `pandas.read_sql_query()` through an engine of each PyAthena driver, whole or with `chunksize` |
-| Wrangler CSV | Scalar projections and the json workload; nested cases are marked unsupported |
+| Wrangler CSV | Scalar projections, including the json and temporal workloads; nested cases are marked unsupported |
 | Wrangler CTAS and UNLOAD | Full DataFrame or chunks, including nested workloads |
 | S3 filesystem (`s3`) | Read, write, and multipart copy of one object with `S3FileSystem` and `AioS3FileSystem` |
 
@@ -310,7 +315,7 @@ As a cursor is, the engine's connection is opened during setup, outside `total_s
 `pandas.read_sql_query()` builds DataFrames from the rows of `fetchall()` or `fetchmany()`, so these cases include the driver's row output; `sqla_rest` pages through GetQueryResults like Cursor.
 
 The [Wrangler API documentation](https://aws-sdk-pandas.readthedocs.io/en/stable/stubs/awswrangler.athena.read_sql_query.html) describes its CSV nested-type restriction.
-CTAS and UNLOAD have separate limitations, including timestamp-with-time-zone columns; the supplied timestamp is a bigint.
+CTAS and UNLOAD have separate limitations, including `time` and timestamp-with-time-zone columns, which the temporal workload leaves out of its CTAS and UNLOAD reads.
 Unsupported combinations are included in the plan and report with reasons.
 An attempted operation that fails is recorded as a failure, not silently reclassified as unsupported.
 Row-count validation detects incomplete consumption; it is not a proof that different libraries produce identical dtypes or nested Python objects.

benchmarks/pyathena_bench/aws.py

@@ -305,7 +305,7 @@ def prepare_fixtures(
             if key in manifest["fixtures"]:
                 continue
             sql = select_sql(
-                resources["ScratchDatabase"], manifest["scales"][scale]["table"], shape
+                resources["ScratchDatabase"], manifest["scales"][scale]["table"], shape, transport
             )
             root = (
                 f"s3://{resources['Bucket']}/runs/{manifest['run_id']}/fixtures/{uuid.uuid4().hex}/"

benchmarks/pyathena_bench/cases.py

@@ -33,7 +33,7 @@
 FILESYSTEM_OPERATIONS = ("read", "write", "copy")
 ALL_FAMILIES = (*FAMILIES, "wrangler", *SQLALCHEMY_FAMILIES, FILESYSTEM_FAMILY)
 SUITES = ("single", "concurrent", "init", "filesystem")
-SHAPES = ("flat", "nested", "json")
+SHAPES = ("flat", "nested", "json", "temporal")
 
 
 @dataclass(frozen=True)
@@ -100,14 +100,19 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]:
     ``json`` type; the initialization suite, which prepares UNLOAD fixtures too,
     has no json cases.
 
+    The ``temporal`` shape adds temporal columns to the flat projection and
+    covers every family in the single-query suite, UNLOAD and CTAS reads
+    included (see ``select_sql()``). Like json, it has no cases in the other
+    suites.
+
     The ``filesystem`` suite does not query Athena and has no shape; its cases
     belong to the ``flat`` shape.
 
     Args:
         settings: Configuration with the arraysizes, chunk sizes, and concurrency,
             and the filesystem suite's max_workers and connection pools.
         suite: ``single``, ``concurrent``, ``init``, or ``filesystem``.
-        shape: ``flat``, ``nested``, or ``json``.
+        shape: ``flat``, ``nested``, ``json``, or ``temporal``.
 
     Returns:
         The cases in a stable order.
@@ -135,7 +140,7 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]:
             for workers in settings.filesystem_max_workers
             for pool in settings.filesystem_pool_connections
         ]
-    if shape == "json" and suite != "single":
+    if shape in {"json", "temporal"} and suite != "single":
         return []
     cases: list[Case] = []
     arraysizes = settings.arraysizes if suite == "single" else [1000]

benchmarks/pyathena_bench/config.py

@@ -30,6 +30,23 @@
 )
 # The flat projection with one column of Athena's json type, NULL where details is NULL.
 JSON_COLUMNS = f"{FLAT_COLUMNS}, CAST(details AS JSON) AS details_json"
+# The flat projection with temporal columns derived from the epoch microseconds of
+# "timestamp"; from_unixtime() returns timestamp(3). UNLOAD and CTAS write Parquet
+# through Hive, which has no time or with time zone types, so their reads project
+# the timestamp and date columns only. The zone has no daylight saving time, whose
+# ambiguous local times PyAthena rejects.
+_SECONDS = '"timestamp" / 1e6'
+PARQUET_TEMPORAL_COLUMNS = (
+    f"{FLAT_COLUMNS}, from_unixtime({_SECONDS}) AS event_time, "
+    f"CAST(from_unixtime({_SECONDS}) AS date) AS event_date"
+)
+TEMPORAL_COLUMNS = (
+    f"{PARQUET_TEMPORAL_COLUMNS}, "
+    f"from_unixtime({_SECONDS}, 'UTC') AS event_time_utc, "
+    f"from_unixtime({_SECONDS}, 'Asia/Tokyo') AS event_time_zone, "
+    f"from_unixtime({_SECONDS}, '+09:00') AS event_time_offset, "
+    f"CAST(from_unixtime({_SECONDS}) AS time(3)) AS event_clock"
+)
 
 
 def identifier(value: str) -> str:
@@ -139,23 +156,35 @@ def load(cls, path: Path) -> Settings:
         )
 
 
-def select_sql(database: str, table: str, shape: str) -> str:
+def select_sql(database: str, table: str, shape: str, transport: str) -> str:
     """Return the measured query of a shape over a snapshot table.
 
     Args:
         database: Database of the snapshot.
         table: Snapshot table.
-        shape: ``flat``, ``nested``, or ``json``.
+        shape: ``flat``, ``nested``, ``json``, or ``temporal``.
+        transport: ``csv``, ``unload``, or ``ctas``. The ``temporal`` shape
+            projects only the columns that Parquet can hold for ``unload`` and
+            ``ctas``.
 
     Returns:
         The SELECT statement without a ``LIMIT``.
 
     Raises:
-        ValueError: If the shape is unknown.
+        ValueError: If the shape or transport is unknown.
     """
-    columns = {"flat": FLAT_COLUMNS, "nested": NESTED_COLUMNS, "json": JSON_COLUMNS}.get(shape)
+    columns = {
+        "flat": FLAT_COLUMNS,
+        "nested": NESTED_COLUMNS,
+        "json": JSON_COLUMNS,
+        "temporal": TEMPORAL_COLUMNS,
+    }.get(shape)
     if columns is None:
         raise ValueError(f"Unknown shape: {shape}")
+    if transport not in {"csv", "unload", "ctas"}:
+        raise ValueError(f"Unknown transport: {transport}")
+    if shape == "temporal" and transport != "csv":
+        columns = PARQUET_TEMPORAL_COLUMNS
     return f'SELECT {columns} FROM "{identifier(database)}"."{identifier(table)}"'
 
 

benchmarks/pyathena_bench/runner.py

@@ -532,6 +532,7 @@ def run(
                             resources["ScratchDatabase"],
                             manifest["scales"][scale]["table"],
                             case.shape,
+                            case.transport,
                         ),
                         expected_rows=manifest["scales"][scale]["rows"],
                         fixture=manifest["fixtures"].get(f"{scale}-{case.shape}-{case.transport}"),

benchmarks/tests/test_config.py

@@ -90,6 +90,15 @@ def test_json_shape_covers_pandas_and_wrangler_reads():
     assert matrix(Settings(), "concurrent", "json") == []
 
 
+def test_temporal_shape_covers_every_flat_case_in_the_single_query_suite():
+    # UNLOAD and CTAS read a Parquet projection instead of being unsupported.
+    assert matrix(Settings(), "single", "temporal") == [
+        replace(c, shape="temporal") for c in matrix(Settings(), "single", "flat")
+    ]
+    assert matrix(Settings(), "init", "temporal") == []
+    assert matrix(Settings(), "concurrent", "temporal") == []
+
+
 def test_concurrency_and_init_matrices_are_bounded():
     cases = matrix(Settings(), "concurrent", "flat")
     assert {c.concurrency for c in cases} == {1, 10, 50, 100}
@@ -127,17 +136,46 @@ def test_filesystem_section_is_loaded_with_its_own_scales(tmp_path):
 
 
 def test_flat_projection_and_nested_projection_share_the_snapshot():
-    flat = select_sql("scratch", "b_run_small", "flat")
-    nested = select_sql("scratch", "b_run_small", "nested")
+    flat = select_sql("scratch", "b_run_small", "flat", "csv")
+    nested = select_sql("scratch", "b_run_small", "nested", "csv")
     assert flat.endswith('FROM "scratch"."b_run_small"')
     assert nested.endswith('FROM "scratch"."b_run_small"')
     assert "file.filename AS filename" in flat
     assert "project, file, details, http" in nested
     assert "LIMIT" not in flat
-    json = select_sql("scratch", "b_run_small", "json")
+    json = select_sql("scratch", "b_run_small", "json", "csv")
     assert json == flat.replace(
         ' FROM "scratch"', ', CAST(details AS JSON) AS details_json FROM "scratch"'
     )
+    # The other shapes read the same projection with every transport.
+    for shape in ("flat", "nested", "json"):
+        for transport in ("unload", "ctas"):
+            assert select_sql("scratch", "b_run_small", shape, transport) == select_sql(
+                "scratch", "b_run_small", shape, "csv"
+            )
+
+
+def test_temporal_projection_of_parquet_reads_has_no_time_or_time_zone_types():
+    flat = select_sql("scratch", "b_run_small", "flat", "csv")
+    parquet = flat.replace(
+        ' FROM "scratch"',
+        ', from_unixtime("timestamp" / 1e6) AS event_time, '
+        'CAST(from_unixtime("timestamp" / 1e6) AS date) AS event_date FROM "scratch"',
+    )
+    # Athena's UNLOAD and CTAS fail with "Unsupported Hive type" for time and with
+    # time zone columns.
+    assert select_sql("scratch", "b_run_small", "temporal", "unload") == parquet
+    assert select_sql("scratch", "b_run_small", "temporal", "ctas") == parquet
+    assert select_sql("scratch", "b_run_small", "temporal", "csv") == parquet.replace(
+        ' FROM "scratch"',
+        ", from_unixtime(\"timestamp\" / 1e6, 'UTC') AS event_time_utc, "
+        "from_unixtime(\"timestamp\" / 1e6, 'Asia/Tokyo') AS event_time_zone, "
+        "from_unixtime(\"timestamp\" / 1e6, '+09:00') AS event_time_offset, "
+        'CAST(from_unixtime("timestamp" / 1e6) AS time(3)) AS event_clock FROM "scratch"',
+    )
+    # A misspelled transport fails instead of reading the projection of CSV results.
+    with pytest.raises(ValueError, match="Unknown transport"):
+        select_sql("scratch", "b_run_small", "temporal", "UNLOAD")
 
 
 @pytest.mark.parametrize("family", ["cursor", "dict", "arrow"])

benchmarks/tests/test_runner.py

@@ -246,6 +246,25 @@ def summarize(s3, location):
     )
 
 
+def test_temporal_trials_query_the_projection_of_their_transport(monkeypatch, tmp_path, run_stubs):
+    payloads = []
+    monkeypatch.setattr(
+        "pyathena_bench.runner.supervise",
+        lambda payload, events, settings: payloads.append(payload) or {"status": "ok"},
+    )
+    monkeypatch.setattr("pyathena_bench.runner.output_files", lambda *args: {})
+    cases = [
+        Case("arrow", transport=transport, output="native", shape="temporal")
+        for transport in ("csv", "unload")
+    ]
+    settings = replace(Settings(), warmups=0, repetitions=1)
+    assert run(settings, RUN_MANIFEST, cases, ["small"], tmp_path / "run")
+    csv, unload = (p["sql"] for p in payloads)
+    assert "AS event_time_utc" in csv
+    assert "AS event_date" in unload
+    assert "AS event_time_utc" not in unload
+
+
 def test_filesystem_trials_transfer_the_prepared_object_within_the_scratch_bucket(
     monkeypatch, tmp_path, run_stubs
 ):