Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
73 changes: 47 additions & 26 deletions benchmarks/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,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.
Neither workload converts every library's output into a common Python object representation.
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.
No workload converts every library's output into a common Python object representation.

## Dependencies and local checks

Expand Down Expand Up @@ -239,12 +240,18 @@ When the temporary directory is a tmpfs, as `/tmp` is on Amazon Linux 2023, file
| --- | --- |
| Cursor, DictCursor, S3FS | Row batches with arraysize 100 and 1000 |
| Pandas CSV | Full DataFrame or streamed chunks |
| Pandas UNLOAD | Full Parquet read; requested chunked cases are marked unsupported as bounded-memory measurements |
| Pandas UNLOAD | Full DataFrame or chunks, each also with `dtype_backend="pyarrow"` |
| Arrow CSV and UNLOAD | Eager table loading; arraysize changes row extraction, not S3 streaming |
| Polars CSV and UNLOAD | Full result or chunk iterator |
| Wrangler CSV | Scalar projections; nested cases are marked unsupported |
| 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 CTAS and UNLOAD | Full DataFrame or chunks, including nested workloads |

The SQLAlchemy families `sqla_rest`, `sqla_pandas`, `sqla_arrow`, `sqla_polars`, and `sqla_s3fs` read with the synchronous API in the single-query suite only, and the result-file drivers also read UNLOAD results.
Their engines connect with the same workgroup, output location, polling interval, and disabled result reuse as the cursor cases, their statements run without the client cache (`cache_size` defaults to 0), and their cursors take the same cursor options; with `chunksize`, the pandas and Polars drivers' cursors also read the result file in chunks.
As a cursor is, the engine's connection is opened during setup, outside `total_seconds`.
`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.
Unsupported combinations are included in the plan and report with reasons.
Expand Down Expand Up @@ -291,6 +298,7 @@ While an API-heavy job waits for slots, workers do not start later API-heavy job
It orders API-heavy jobs first and longer jobs earlier.
With `--split-pages`, a job whose single trial needs at least that many pages becomes one job per measured repetition without a warmup, so those repetitions run on different hosts.
`run` accepts the matching `--arraysize`, `--warmups`, and `--repetitions` options.
`--chunksize` keeps only the given chunk sizes, with 0 for reads without chunks, in `plan`, `run`, and `jobs`; `jobs` passes it to each job and adds it to the job ID.
`jobs` accepts `--warmups` and `--repetitions` too; with either option, every job that is not split carries both values as arguments, so one queue can combine jobs with different trial counts under one stored configuration.
`queue --jobs` accepts several job files, rejects duplicate job IDs, and orders the combined jobs as `jobs` orders one selection.

Expand Down Expand Up @@ -330,41 +338,53 @@ uv run --no-sync python -m pyathena_bench --config results/fleet.toml queue \
```

Several `jobs` selections can make up one bounded matrix.
For example, the following selection, with `chunksizes = [10000, 100000, 1000000]`, `concurrency = [1, 10]`, and `timeout_seconds = 5400.0` in `results/fleet.toml`, measures:
For example, the following selection, with `chunksizes = [10000, 100000, 1000000]` and `timeout_seconds = 5400.0` in `results/fleet.toml`, measures the synchronous API only:

- medium in full, with five repetitions;
- large with arraysize 100 paging limited to Cursor with the synchronous API, concurrent row paging limited to flat Cursor jobs, and initialization limited to the S3FS and DataFrame families, with five repetitions, except three for those two Cursor selections;
- xlarge with the synchronous API only, without DictCursor and without arraysize 100 for the DataFrame families' row output, with three repetitions.
- medium as a small regression subset: Cursor and S3FS rows with arraysize 1000, and the native reads of the DataFrame families and Wrangler without chunks, flat, three repetitions;
- large, flat and nested: Cursor and S3FS rows with arraysize 1000 (five repetitions), Cursor with arraysize 100 (three, split into single-repetition jobs), DictCursor flat (three), every native read of the DataFrame families and Wrangler (five), the json workload's CSV reads without chunks and with 100,000-row chunks (five), and `read_sql_query()` through the pandas, Arrow, Polars, and S3FS drivers without chunks and with 100,000-row chunks (five) and through the REST driver without chunks (three, split); `sqla_s3fs` covers the chunked `fetchmany()` path without GetQueryResults paging;
- xlarge: S3FS rows flat, every native read of the DataFrame families and Wrangler, the json workload without chunks and with 100,000-row chunks, and `read_sql_query()` through the pandas, Arrow, Polars, and S3FS drivers flat without chunks and with chunks, and through the pandas, Arrow, and Polars drivers nested with chunks only, all with three repetitions.

It leaves out the ThreadPool and native asyncio APIs, whose single-query timings matched the synchronous API, the initialization suite, the concurrency suite, arraysize 100 for every family but Cursor, and DictCursor beyond one scale.
Every selection runs without warmups because each trial already runs in a fresh process.

```bash
jobs() { uv run --no-sync python -m pyathena_bench --config results/fleet.toml jobs --warmups 0 "$@"; }
jobs --suite single concurrent init --scale medium --shape flat nested --repetitions 5 > results/jobs-medium.json
jobs --suite single --scale large --shape flat nested --arraysize 1000 --repetitions 5 > results/jobs-large-1000.json
jobs --suite single --scale large --shape flat nested --arraysize 100 --family s3fs pandas arrow polars \
--repetitions 5 > results/jobs-large-100.json
jobs --suite single --scale large --shape flat nested --arraysize 100 --family cursor --api sync \
--repetitions 3 --split-pages 5000 > results/jobs-large-cursor-100.json
jobs --suite concurrent --scale large --shape flat nested --family s3fs pandas arrow polars \
--repetitions 5 > results/jobs-large-concurrent.json
jobs --suite concurrent --scale large --shape flat --family cursor \
--repetitions 3 --split-pages 5000 > results/jobs-large-concurrent-cursor.json
jobs --suite init --scale large --shape flat nested --family s3fs pandas arrow polars \
--repetitions 5 > results/jobs-large-init.json
jobs --suite single --scale xlarge --shape flat nested --api sync --family cursor --arraysize 1000 \
--repetitions 3 --split-pages 5000 > results/jobs-xlarge-cursor.json
jobs --suite single --scale xlarge --shape flat nested --api sync --family s3fs \
--repetitions 3 > results/jobs-xlarge-s3fs.json
jobs --suite single --scale xlarge --shape flat nested --api sync --family pandas arrow polars wrangler \
--arraysize 1000 --repetitions 3 > results/jobs-xlarge-dataframes.json
jobs --suite single --scale medium --shape flat --api sync --arraysize 1000 --chunksize 0 \
--family cursor s3fs --output-kind rows --repetitions 3 > results/jobs-medium-rows.json
jobs --suite single --scale medium --shape flat --api sync --chunksize 0 \
--family pandas arrow polars wrangler --output-kind native --repetitions 3 > results/jobs-medium-native.json
jobs --suite single --scale large --shape flat nested --api sync --arraysize 1000 \
--family cursor s3fs --output-kind rows --repetitions 5 > results/jobs-large-rows.json
jobs --suite single --scale large --shape flat nested --api sync --arraysize 100 \
--family cursor --repetitions 3 --split-pages 5000 > results/jobs-large-cursor-100.json
jobs --suite single --scale large --shape flat --api sync --arraysize 1000 \
--family dict --repetitions 3 > results/jobs-large-dict.json
jobs --suite single --scale large --shape flat nested --api sync \
--family pandas arrow polars wrangler --output-kind native --repetitions 5 > results/jobs-large-native.json
jobs --suite single --scale large --shape json --api sync --transport csv --chunksize 0 100000 \
--family pandas wrangler --repetitions 5 > results/jobs-large-json.json
jobs --suite single --scale large --shape flat nested --chunksize 0 100000 \
--family sqla_pandas sqla_arrow sqla_polars sqla_s3fs --repetitions 5 > results/jobs-large-sqla.json
jobs --suite single --scale large --shape flat nested --chunksize 0 \
--family sqla_rest --repetitions 3 --split-pages 1000 > results/jobs-large-sqla-rest.json
jobs --suite single --scale xlarge --shape flat --api sync --arraysize 1000 \
--family s3fs --output-kind rows --repetitions 3 > results/jobs-xlarge-s3fs.json
jobs --suite single --scale xlarge --shape flat nested --api sync \
--family pandas arrow polars wrangler --output-kind native --repetitions 3 > results/jobs-xlarge-native.json
jobs --suite single --scale xlarge --shape json --api sync --transport csv --chunksize 0 100000 \
--family pandas wrangler --repetitions 3 > results/jobs-xlarge-json.json
jobs --suite single --scale xlarge --shape flat --chunksize 0 100000 \
--family sqla_pandas sqla_arrow sqla_polars sqla_s3fs --repetitions 3 > results/jobs-xlarge-sqla-flat.json
jobs --suite single --scale xlarge --shape nested --chunksize 100000 \
--family sqla_pandas sqla_arrow sqla_polars --repetitions 3 > results/jobs-xlarge-sqla-nested.json
uv run --no-sync python -m pyathena_bench --config results/fleet.toml queue \
--stack "$BENCHMARK_STACK_ID" --name large-1 --jobs results/jobs-medium.json results/jobs-large-*.json \
--stack "$BENCHMARK_STACK_ID" --name large-1 --jobs results/jobs-medium-*.json results/jobs-large-*.json \
results/jobs-xlarge-*.json --manifest results/input-bounded.json
```

The manifest for this selection must be prepared with `--scale medium large xlarge`.
It is published as `large-1` in place of the previous example, so the worker and status commands below apply unchanged.
The selection has 96 jobs; from the job durations of the 2026-10-11 run, they take about 14 hours of job time on `r7g.2xlarge` hosts, and the longest, xlarge nested pandas CSV, about 75 minutes.

`queue` reserves the name with a conditional write, stores the configuration, manifest, and jobs under `fleet/<name>/` in the scratch bucket, and writes the `queue.json` marker last.
Workers refuse a queue without the marker.
Expand Down Expand Up @@ -405,6 +425,7 @@ Every output directory contains `environment.json`, a copy of the input manifest
Environment metadata includes Git SHA, dirty state, Python and dependency versions, lockfile hash, CPU count, RAM, and settings.
SDK events preserve query IDs, SQL, output locations, Athena statistics, and failed operations without storing result rows.
Successful measurements generate `summary.csv` and `summary.md`; failed trials and unsupported cases remain visible.
After each successful measured UNLOAD or CTAS trial, `run` reads the query manifests (`<query_id>-manifest.csv`) below the trial's output location and records the number and sizes of the data files they list in `output_files`; the summary shows the most files and the largest file of the measured trials, and notes a measured trial whose files could not be summarized.
Regenerate reports after interruption with:

```bash
Expand Down
36 changes: 18 additions & 18 deletions benchmarks/pyathena_bench/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,10 @@
stack_resources,
validate_manifest,
)
from pyathena_bench.cases import FAMILIES, matrix
from pyathena_bench.cases import ALL_FAMILIES, SHAPES, matrix
from pyathena_bench.config import Settings, read_json
from pyathena_bench.fleet import (
API_FAMILIES,
Filters,
Queue,
expand_jobs,
Expand Down Expand Up @@ -62,7 +63,7 @@ def parser() -> argparse.ArgumentParser:
suites = ("single", "concurrent", "init")
if name == "jobs":
command.add_argument("--suite", choices=suites, nargs="+", required=True)
command.add_argument("--shape", choices=("flat", "nested"), nargs="+", default=["flat"])
command.add_argument("--shape", choices=SHAPES, nargs="+", default=["flat"])
command.add_argument(
"--split-pages",
type=int,
Expand All @@ -71,17 +72,23 @@ def parser() -> argparse.ArgumentParser:
)
else:
command.add_argument("--suite", choices=suites, required=True)
command.add_argument("--shape", choices=("flat", "nested"), default="flat")
command.add_argument("--shape", choices=SHAPES, default="flat")
command.add_argument("--warmups", type=int, help="Override the configured warmups")
command.add_argument("--repetitions", type=int, help="Override the configured repetitions")
command.add_argument("--scale", nargs="+", required=True)
command.add_argument("--family", choices=(*FAMILIES, "wrangler"), nargs="+")
command.add_argument("--family", choices=ALL_FAMILIES, nargs="+")
command.add_argument(
"--api", choices=("sync", "thread", "aio", "direct", "to_thread"), nargs="+"
)
command.add_argument("--transport", choices=("csv", "unload", "ctas"), nargs="+")
command.add_argument("--output-kind", choices=("rows", "native"))
command.add_argument("--arraysize", type=int, nargs="+")
command.add_argument(
"--chunksize",
type=int,
nargs="+",
help="Chunk sizes to keep; 0 selects unchunked reads",
)
if name == "run":
command.add_argument("--manifest", type=Path, required=True)
command.add_argument("--out", type=Path, required=True)
Expand Down Expand Up @@ -186,13 +193,16 @@ def main(argv: list[str] | None = None) -> int:
retire_host(queue, client(session_, "autoscaling"), ran)
else:
sys.stdout.write(json.dumps(status(queue), indent=2) + "\n")
elif args.command == "jobs":
filters = Filters(args.family, args.api, args.transport, args.output_kind, args.arraysize)
else:
filters = Filters(
args.family, args.api, args.transport, args.output_kind, args.arraysize, args.chunksize
)
overrides = {
k: v for k in ("warmups", "repetitions") if (v := getattr(args, k)) is not None
}
if overrides:
settings = replace(settings, **overrides)
if args.command == "jobs":
jobs = expand_jobs(
settings,
args.suite,
Expand All @@ -203,23 +213,13 @@ def main(argv: list[str] | None = None) -> int:
explicit_trials=bool(overrides),
)
sys.stdout.write(json.dumps(jobs, indent=1) + "\n")
else:
overrides = {
k: v for k in ("warmups", "repetitions") if (v := getattr(args, k)) is not None
}
if overrides:
settings = replace(settings, **overrides)
filters = Filters(args.family, args.api, args.transport, args.output_kind, args.arraysize)
elif args.command in {"plan", "run"}:
cases = filters.select(matrix(settings, args.suite, args.shape))
if not cases:
raise ValueError("No cases match the selected filters")
if args.command == "plan":
page_size = min(
(
c.arraysize
for c in cases
if c.family in {"cursor", "dict"} and c.suite != "init"
),
(c.arraysize for c in cases if c.family in API_FAMILIES and c.suite != "init"),
default=0,
)
sys.stdout.write(
Expand Down
Loading
Loading