From 6f45a088054321fd5b7aa71bfdf61782273f91a9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 15:28:56 +0900 Subject: [PATCH 1/5] Benchmark read_sql_query and trim the next benchmark selection Add pandas.read_sql_query() through a SQLAlchemy engine of each PyAthena driver (sqla_rest, sqla_pandas, sqla_arrow, sqla_polars, sqla_s3fs), whole and with chunksize, with the same connection and cursor options as the cursor cases. Add a json workload (the flat projection with CAST(details AS JSON)) for pandas and Wrangler CSV reads, pandas UNLOAD reads with dtype_backend="pyarrow", and a record of the Parquet files each UNLOAD or CTAS trial wrote. Pandas UNLOAD chunks are measured again: PandasCursor has read them in chunks since #1135. A --chunksize filter selects chunk sizes and is carried by jobs. The README selection measures the synchronous API only and leaves out the cases the 4.0.0 run showed to carry no information. Refs #1133. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 70 +++++++++------ benchmarks/pyathena_bench/__main__.py | 31 ++++--- benchmarks/pyathena_bench/adapters.py | 117 ++++++++++++++++++++------ benchmarks/pyathena_bench/cases.py | 75 +++++++++++++---- benchmarks/pyathena_bench/config.py | 19 ++++- benchmarks/pyathena_bench/fleet.py | 15 +++- benchmarks/pyathena_bench/report.py | 16 +++- benchmarks/pyathena_bench/runner.py | 48 ++++++++++- benchmarks/pyproject.toml | 2 +- benchmarks/tests/test_adapters.py | 89 +++++++++++++++++++- benchmarks/tests/test_config.py | 38 +++++++-- benchmarks/tests/test_fleet.py | 45 ++++++++++ benchmarks/tests/test_report.py | 21 +++++ benchmarks/tests/test_runner.py | 77 ++++++++++++++++- uv.lock | 4 +- 15 files changed, 570 insertions(+), 97 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 60383d6b8..e5dcd76b4 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -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 only native pandas and Wrangler reads of CSV results, because Parquet cannot hold the `json` type. +No workload converts every library's output into a common Python object representation. ## Dependencies and local checks @@ -239,12 +240,17 @@ 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, and their cursors take the same cursor options; with `chunksize`, the pandas and Polars drivers' cursors also read the result file in chunks. +`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. @@ -291,6 +297,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. @@ -330,41 +337,51 @@ 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 including the json workload (five), and `read_sql_query()` for every driver without chunks and with 100,000-row chunks (five; three for `sqla_rest`, split); +- 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()` for the result-file drivers flat without chunks and with chunks, and 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 json --api sync \ + --family pandas arrow polars wrangler --output-kind native --repetitions 5 > results/jobs-large-native.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 100000 \ + --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 --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 host-hours 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//` in the scratch bucket, and writes the `queue.json` marker last. Workers refuse a queue without the marker. @@ -405,6 +422,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 UNLOAD or CTAS trial, `run` lists the trial's output location and records the number and sizes of the Parquet files in `output_files`; the summary shows the most files and the largest file. Regenerate reports after interruption with: ```bash diff --git a/benchmarks/pyathena_bench/__main__.py b/benchmarks/pyathena_bench/__main__.py index 40e1883c1..082f2ee64 100644 --- a/benchmarks/pyathena_bench/__main__.py +++ b/benchmarks/pyathena_bench/__main__.py @@ -26,9 +26,10 @@ stack_resources, validate_manifest, ) -from pyathena_bench.cases import FAMILIES, matrix +from pyathena_bench.cases import FAMILIES, SHAPES, SQLALCHEMY_FAMILIES, matrix from pyathena_bench.config import Settings, read_json from pyathena_bench.fleet import ( + API_FAMILIES, Filters, Queue, expand_jobs, @@ -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, @@ -71,17 +72,25 @@ 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=(*FAMILIES, "wrangler", *SQLALCHEMY_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) @@ -187,7 +196,9 @@ def main(argv: list[str] | None = None) -> int: 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) + 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 } @@ -209,17 +220,15 @@ def main(argv: list[str] | None = None) -> int: } if overrides: settings = replace(settings, **overrides) - filters = Filters(args.family, args.api, args.transport, args.output_kind, args.arraysize) + filters = Filters( + args.family, args.api, args.transport, args.output_kind, args.arraysize, args.chunksize + ) 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( diff --git a/benchmarks/pyathena_bench/adapters.py b/benchmarks/pyathena_bench/adapters.py index e452722d0..c3c6d8ac9 100644 --- a/benchmarks/pyathena_bench/adapters.py +++ b/benchmarks/pyathena_bench/adapters.py @@ -20,6 +20,8 @@ import polars as pl import pyarrow as pa from botocore.config import Config +from sqlalchemy import create_engine +from sqlalchemy.pool import NullPool from pyathena import connect from pyathena.aio.arrow.cursor import AioArrowCursor @@ -43,7 +45,7 @@ from pyathena.s3fs.async_cursor import AsyncS3FSCursor from pyathena.s3fs.cursor import S3FSCursor from pyathena.s3fs.result_set import AthenaS3FSResultSet -from pyathena_bench.cases import Case +from pyathena_bench.cases import SQLALCHEMY_FAMILIES, Case from pyathena_bench.config import Settings BACKEND_VERSIONS = { @@ -158,42 +160,95 @@ class Adapter: output: str temp_table: str - def connection(self) -> Any: - return connect( - session=self.session, - region_name=self.settings.region, - schema_name=self.database, - work_group=self.settings.workgroup, - s3_staging_dir=self.output, - poll_interval=self.settings.poll_interval, - result_reuse_enable=False, - config=Config( + def connection_kwargs(self) -> dict[str, Any]: + """Return the ``pyathena.connect()`` arguments shared by every PyAthena case. + + Returns: + The connection arguments, with Athena result reuse disabled. + """ + return { + "session": self.session, + "region_name": self.settings.region, + "schema_name": self.database, + "work_group": self.settings.workgroup, + "s3_staging_dir": self.output, + "poll_interval": self.settings.poll_interval, + "result_reuse_enable": False, + "config": Config( max_pool_connections=max( 10, 2 * self.settings.executor_workers, self.case.concurrency ), connect_timeout=10, read_timeout=60, ), - ) + } - def cursor(self, connection: Any) -> Any: + def connection(self) -> Any: + return connect(**self.connection_kwargs()) + + def cursor_kwargs(self, family: str) -> dict[str, Any]: + """Return the cursor options of the case for a cursor family. + + Args: + family: The cursor family, which a SQLAlchemy case maps from its driver. + + Returns: + The keyword arguments that the family's cursor accepts. + """ case = self.case kwargs: dict[str, Any] = {"arraysize": case.arraysize} if case.api == "thread": kwargs["max_workers"] = self.settings.executor_workers - if case.family in {"pandas", "polars"}: + if family in {"pandas", "polars"}: kwargs["s3_max_workers"] = self.settings.executor_workers - if case.family in {"pandas", "polars", "arrow"}: + if family in {"pandas", "polars", "arrow"}: kwargs["unload"] = case.transport == "unload" - if case.family in {"pandas", "polars"}: + if family in {"pandas", "polars"}: kwargs["chunksize"] = case.chunksize - if case.family == "pandas" and case.transport == "unload": + if family == "pandas" and case.transport == "unload": kwargs["engine"] = "pyarrow" - cursor = connection.cursor(CURSORS[case.api][case.family], **kwargs) + return kwargs + + def cursor(self, connection: Any) -> Any: + case = self.case + cursor = connection.cursor( + CURSORS[case.api][case.family], **self.cursor_kwargs(case.family) + ) # Some cursor bases reset arraysize during construction. cursor.arraysize = case.arraysize return cursor + def engine(self) -> Any: + """Create the SQLAlchemy engine of a ``read_sql_query()`` case. + + The engine connects with the same arguments as the cursor cases, and its + connections create cursors with the cursor options of the driver's + family, including ``chunksize`` for the pandas and Polars drivers, so a + chunked read also streams the result file. ``NullPool`` opens one + connection per query. + + Returns: + The engine. + """ + driver = self.case.family.removeprefix("sqla_") + family = "cursor" if driver == "rest" else driver + return create_engine( + f"awsathena+{driver}://athena.{self.settings.region}.amazonaws.com:443/{self.database}", + poolclass=NullPool, + connect_args={ + **self.connection_kwargs(), + "cursor_kwargs": self.cursor_kwargs(family), + }, + ) + + def execute_kwargs(self) -> dict[str, Any]: + """Return the ``execute()`` options of the case besides cache and reuse settings. + + Returns: + ``dtype_backend`` for the pandas cases that set it, or no options. + """ + return {"dtype_backend": self.case.dtype_backend} if self.case.dtype_backend else {} + def wrangler(self, sql: str) -> Any: options: dict[str, Any] = { "boto3_session": self.session, @@ -247,25 +302,32 @@ async def measure_query(adapter: Adapter, cursor: Any, sql: str, expected: int) case = adapter.case result = None try: + options = {"cache_size": 0, "result_reuse_enable": False, **adapter.execute_kwargs()} if case.family == "wrangler": result = adapter.wrangler(sql) ready = time.perf_counter() count = consume(result, case) query_id = None # The botocore observer captures IDs even for iterators. + elif case.family in SQLALCHEMY_FAMILIES: + # The "cursor" of a SQLAlchemy case is its engine; the connection stays + # open while chunks are read. + with cursor.connect() as connection: + result = pd.read_sql_query(sql, connection, chunksize=case.chunksize) + ready = time.perf_counter() + count = consume(result, case) + query_id = None elif case.api == "thread": - query_id, future = await asyncio.to_thread( - cursor.execute, sql, cache_size=0, result_reuse_enable=False - ) + query_id, future = await asyncio.to_thread(cursor.execute, sql, **options) result = await asyncio.wrap_future(future) ready = time.perf_counter() count = await asyncio.to_thread(consume, result, case) elif case.api == "aio": - await cursor.execute(sql, cache_size=0, result_reuse_enable=False) + await cursor.execute(sql, **options) ready = time.perf_counter() query_id = cursor.query_id count = await consume_aio(cursor, case) else: - cursor.execute(sql, cache_size=0, result_reuse_enable=False) + cursor.execute(sql, **options) ready = time.perf_counter() query_id = cursor.query_id count = consume(cursor.result_set, case) @@ -305,8 +367,9 @@ async def run_adapter( loop.set_default_executor(ThreadPoolExecutor(max_workers=adapter.settings.executor_workers)) setup_start = time.perf_counter() case = adapter.case - connection = None if case.family == "wrangler" else adapter.connection() - cursors: list[Any] = [] + sqlalchemy = case.family in SQLALCHEMY_FAMILIES + connection = None if case.family == "wrangler" or sqlalchemy else adapter.connection() + cursors: list[Any] = [adapter.engine()] if sqlalchemy else [] if case.suite != "init" and connection is not None: # AsyncCursor owns one shared pool; Aio cursors carry per-query state. cursors = [ @@ -368,7 +431,9 @@ async def run_adapter( stop.set() await monitor for cursor in cursors: - if case.api == "thread": + if sqlalchemy: + cursor.dispose() + elif case.api == "thread": cursor.close(wait=True) else: cursor.close() diff --git a/benchmarks/pyathena_bench/cases.py b/benchmarks/pyathena_bench/cases.py index 3cd6125d1..8421d37cd 100644 --- a/benchmarks/pyathena_bench/cases.py +++ b/benchmarks/pyathena_bench/cases.py @@ -14,6 +14,11 @@ from pyathena_bench.config import Settings FAMILIES = ("cursor", "dict", "s3fs", "pandas", "arrow", "polars") +# pandas.read_sql_query() through a SQLAlchemy engine of each PyAthena driver. +SQLALCHEMY_FAMILIES = ("sqla_rest", "sqla_pandas", "sqla_arrow", "sqla_polars", "sqla_s3fs") +SHAPES = ("flat", "nested", "json") +# The SQLAlchemy drivers that read result files and can read UNLOAD results. +RESULT_FILE_SQLALCHEMY = {"sqla_pandas", "sqla_arrow", "sqla_polars"} @dataclass(frozen=True) @@ -28,17 +33,37 @@ class Case: concurrency: int = 1 suite: str = "single" unsupported: str | None = None + # The pandas.read_csv()/read_parquet() dtype_backend given to execute(). + dtype_backend: str | None = None @property def id(self) -> str: + backend = f"-{self.dtype_backend}" if self.dtype_backend else "" return ( f"{self.suite}-{self.family}-{self.api}-{self.transport}-{self.output}-" - f"a{self.arraysize}-c{self.chunksize or 0}-{self.shape}-n{self.concurrency}" + f"a{self.arraysize}-c{self.chunksize or 0}{backend}-{self.shape}-n{self.concurrency}" ) def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: - if suite not in {"single", "concurrent", "init"} or shape not in {"flat", "nested"}: + """Return every case of a suite and result shape, including unsupported ones. + + The ``json`` shape adds a ``json`` column to the flat projection and covers + only the native pandas and Wrangler reads of CSV results: Parquet cannot + hold the ``json`` type. + + Args: + settings: Configuration with the arraysizes, chunk sizes, and concurrency. + suite: ``single``, ``concurrent``, or ``init``. + shape: ``flat``, ``nested``, or ``json``. + + Returns: + The cases in a stable order. + + Raises: + ValueError: If the suite or shape is unknown. + """ + if suite not in {"single", "concurrent", "init"} or shape not in SHAPES: raise ValueError("Unknown suite or shape") cases: list[Case] = [] arraysizes = settings.arraysizes if suite == "single" else [1000] @@ -51,22 +76,21 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: [None] if family == "arrow" else [None, *settings.chunksizes] ) for chunk in chunks: - reason = None - if family == "pandas" and transport == "unload" and chunk: - reason = ( - "Pandas UNLOAD reads all Parquet; chunksize does not bound memory" - ) cases.append( - Case( - family, - api, - transport, - "native", - chunksize=chunk, - shape=shape, - unsupported=reason, - ) + Case(family, api, transport, "native", chunksize=chunk, shape=shape) ) + if family == "pandas" and transport == "unload": + cases.append( + Case( + family, + api, + transport, + "native", + chunksize=chunk, + shape=shape, + dtype_backend="pyarrow", + ) + ) cases.extend( Case(family, api, "unload", arraysize=a, shape=shape) for a in arraysizes ) @@ -87,6 +111,23 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: unsupported=reason, ) ) + if suite == "single": + for family in SQLALCHEMY_FAMILIES: + transports = ["csv", "unload"] if family in RESULT_FILE_SQLALCHEMY else ["csv"] + cases.extend( + Case(family, transport=transport, output="native", chunksize=chunk, shape=shape) + for transport in transports + for chunk in (None, *settings.chunksizes) + ) + if shape == "json": + cases = [ + c + for c in cases + if c.family in {"pandas", "wrangler"} + and c.transport == "csv" + and c.output == "native" + and c.dtype_backend is None + ] if suite == "concurrent": # Keep the scaling suite bounded; chunk and arraysize sweeps belong to single. cases = [ @@ -95,6 +136,7 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: if c.api in {"thread", "aio"} and c.arraysize == 1000 and c.chunksize is None + and c.dtype_backend is None and c.transport == "csv" and ( (c.family in {"cursor", "dict", "s3fs"} and c.output == "rows") @@ -109,6 +151,7 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: if c.api == "sync" and c.arraysize == 1000 and c.chunksize is None + and c.dtype_backend is None and c.family != "wrangler" and ( (c.family in {"cursor", "dict", "s3fs"} and c.output == "rows") diff --git a/benchmarks/pyathena_bench/config.py b/benchmarks/pyathena_bench/config.py index a71793d5e..d4d12d137 100644 --- a/benchmarks/pyathena_bench/config.py +++ b/benchmarks/pyathena_bench/config.py @@ -28,6 +28,8 @@ NESTED_COLUMNS = ( '"timestamp", country_code, url, project, file, details, http, tls_protocol, tls_cipher' ) +# 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" def identifier(value: str) -> str: @@ -100,9 +102,22 @@ def load(cls, path: Path) -> Settings: def select_sql(database: str, table: str, shape: str) -> str: - if shape not in {"flat", "nested"}: + """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``. + + Returns: + The SELECT statement without a ``LIMIT``. + + Raises: + ValueError: If the shape is unknown. + """ + columns = {"flat": FLAT_COLUMNS, "nested": NESTED_COLUMNS, "json": JSON_COLUMNS}.get(shape) + if columns is None: raise ValueError(f"Unknown shape: {shape}") - columns = FLAT_COLUMNS if shape == "flat" else NESTED_COLUMNS return f'SELECT {columns} FROM "{identifier(database)}"."{identifier(table)}"' diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index b14fafad1..fc13b97cd 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -27,8 +27,9 @@ from pyathena_bench.cases import Case, matrix from pyathena_bench.config import Settings -# Row cursors page through GetQueryResults, which has an account-wide request-rate quota. -API_FAMILIES = {"cursor", "dict"} +# Row cursors page through GetQueryResults, which has an account-wide request-rate quota; +# so does read_sql_query() through the REST driver. +API_FAMILIES = {"cursor", "dict", "sqla_rest"} QUEUE_ROOT = "fleet" # A host with this file keeps running after its worker exits, for inspection. KEEP_HOST = Path("results/keep-host") @@ -44,6 +45,8 @@ class Filters: transport: list[str] | None = None output_kind: str | None = None arraysize: list[int] | None = None + # Chunk sizes to keep, with 0 for reads without chunks. + chunksize: list[int] | None = None def select(self, cases: Iterable[Case]) -> list[Case]: """Return the cases that match every configured filter. @@ -62,6 +65,7 @@ def select(self, cases: Iterable[Case]) -> list[Case]: and (not self.transport or c.transport in self.transport) and (not self.output_kind or c.output == self.output_kind) and (not self.arraysize or c.arraysize in self.arraysize) + and (not self.chunksize or (c.chunksize or 0) in self.chunksize) ] @@ -127,7 +131,8 @@ def expand_jobs( """Split a selection into independent ``run`` invocations. A job covers one suite, scale, shape, family, API, transport, and output - kind, and additionally one arraysize for row output. Jobs that page through + kind, and additionally one arraysize for row output. A chunk size filter + is passed on to each job and named in its ID. Jobs that page through GetQueryResults are marked with the number of simultaneous paging queries so workers can bound the fleet's request rate. When one trial of such a job needs at least ``split_pages`` pages, each measured repetition becomes a @@ -170,6 +175,10 @@ def expand_jobs( if arraysize is not None: base += f"-a{arraysize}" args += ["--arraysize", str(arraysize)] + if filters.chunksize: + # The job runs the same chunk selection; the ID keeps it unique. + base += "-c" + "_".join(map(str, filters.chunksize)) + args += ["--chunksize", *map(str, filters.chunksize)] runnable = [c for c in cases if c.unsupported is None] pages = max((api_pages(c, rows) for c in runnable), default=0) api_weight = max( diff --git a/benchmarks/pyathena_bench/report.py b/benchmarks/pyathena_bench/report.py index 311488587..41776876b 100644 --- a/benchmarks/pyathena_bench/report.py +++ b/benchmarks/pyathena_bench/report.py @@ -122,6 +122,14 @@ def summarize(trials: list[dict[str, Any]]) -> list[dict[str, Any]]: "loop_lag_p95_seconds": percentile(lag, 0.95), "loop_lag_max_seconds": max(lag, default=None), "peak_threads": max((t["max_threads"] for t in good), default=None), + "max_output_files": max( + (t["output_files"]["count"] for t in good if "output_files" in t), + default=None, + ), + "max_output_file_bytes": max( + (t["output_files"]["max_bytes"] for t in good if "output_files" in t), + default=None, + ), "notes": "; ".join(sorted(reasons)), } ) @@ -169,11 +177,13 @@ def report(directory: Path) -> None: "Total time includes Athena execution and polling; init time excludes validation.", "Thread counts include SDK and native library threads, not just executor workers.", "Peak temp dir is the per-trial Polars temporary directory, which is outside the RSS.", + "Output files: the most Parquet files an UNLOAD or CTAS trial wrote, and the largest.", "", "| Scale | Case | OK | Fail | Warmup fail | Unsupported | Total (s) | " "Client result (s) | Init (s) | Peak RSS (B) | RSS increase (B) | " - "Peak temp dir (B) | Notes |", - "| --- | --- | ---: | ---: | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | --- |", + "Peak temp dir (B) | Output files | Largest file (B) | Notes |", + "| --- | --- | ---: | ---: | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: " + "| ---: | ---: | --- |", ] for row in rows: values = [ @@ -191,6 +201,8 @@ def report(directory: Path) -> None: "peak_rss_bytes", "median_rss_increase_bytes", "peak_temp_dir_bytes", + "max_output_files", + "max_output_file_bytes", "notes", ) ] diff --git a/benchmarks/pyathena_bench/runner.py b/benchmarks/pyathena_bench/runner.py index a55ac8721..02f8331e0 100644 --- a/benchmarks/pyathena_bench/runner.py +++ b/benchmarks/pyathena_bench/runner.py @@ -32,7 +32,13 @@ import psutil from pyathena_bench.adapters import BACKEND_VERSIONS, Adapter, run_adapter -from pyathena_bench.aws import cancel_queries, session, validate_inputs, validate_manifest +from pyathena_bench.aws import ( + cancel_queries, + client, + session, + validate_inputs, + validate_manifest, +) from pyathena_bench.cases import Case from pyathena_bench.config import Settings, select_sql, write_json @@ -302,6 +308,36 @@ def supervise( return result +def output_files(settings: Settings, location: str) -> dict[str, int]: + """Summarize the data files that a trial wrote below its output location. + + UNLOAD and CTAS write Parquet files there; query result CSV files, their + metadata, and manifests are excluded. + + Args: + settings: Configuration with the AWS region and profile. + location: The trial's ``s3://`` output prefix. + + Returns: + The number of files and their total, largest, and smallest sizes in bytes. + """ + bucket, _, prefix = location.removeprefix("s3://").partition("/") + sizes = [ + item["Size"] + for page in client(session(settings), "s3") + .get_paginator("list_objects_v2") + .paginate(Bucket=bucket, Prefix=prefix) + for item in page.get("Contents", []) + if not item["Key"].endswith((".csv", ".metadata", ".txt", ".json")) + ] + return { + "count": len(sizes), + "bytes": sum(sizes), + "max_bytes": max(sizes, default=0), + "min_bytes": min(sizes, default=0), + } + + def environment(settings: Settings) -> dict[str, Any]: root = Path(__file__).resolve().parents[2] git = subprocess.run( @@ -370,6 +406,16 @@ def run( } append_json(output / "events.jsonl", {"event": "trial_start", **payload}) result = supervise(payload, output / "events.jsonl", settings) + if ( + result["status"] == "ok" + and case.transport in {"unload", "ctas"} + and case.suite != "init" + ): + try: + result["output_files"] = output_files(settings, payload["output"]) + except Exception as exc: + # The measurement stands; only the file summary is missing. + result["output_files_error"] = f"{type(exc).__name__}: {exc}" append_json( output / "trials.jsonl", { diff --git a/benchmarks/pyproject.toml b/benchmarks/pyproject.toml index 17b6633a3..33afecb01 100644 --- a/benchmarks/pyproject.toml +++ b/benchmarks/pyproject.toml @@ -13,7 +13,7 @@ description = "Reproducible Athena cursor measurements" requires-python = ">=3.11" classifiers = ["Private :: Do Not Upload"] dependencies = [ - "PyAthena[pandas,arrow,polars]", + "PyAthena[pandas,arrow,polars,sqlalchemy]", "awswrangler>=3.17,<4", "polars", "psutil>=7,<8", diff --git a/benchmarks/tests/test_adapters.py b/benchmarks/tests/test_adapters.py index 5eb67dd7a..eba371d7f 100644 --- a/benchmarks/tests/test_adapters.py +++ b/benchmarks/tests/test_adapters.py @@ -8,6 +8,7 @@ import asyncio import weakref from concurrent.futures import Future +from contextlib import contextmanager from io import BytesIO from types import SimpleNamespace @@ -30,6 +31,12 @@ from pyathena_bench.cases import Case, matrix from pyathena_bench.config import Settings +from pyathena.arrow.cursor import ArrowCursor +from pyathena.cursor import Cursor as PyAthenaCursor +from pyathena.pandas.cursor import PandasCursor +from pyathena.polars.cursor import PolarsCursor +from pyathena.s3fs.cursor import S3FSCursor + class Rows: def __init__(self): @@ -131,7 +138,7 @@ def unexpected_request(*args, **kwargs): ) settings = Settings(executor_workers=3) for case in matrix(settings, "single", "flat"): - if case.family == "wrangler" or case.unsupported: + if case.family == "wrangler" or case.family.startswith("sqla_") or case.unsupported: continue adapter = Adapter(settings, case, session, "scratch", "s3://out/", "temp") with adapter.connection() as connection: @@ -143,6 +150,86 @@ def unexpected_request(*args, **kwargs): assert cursor._s3_max_workers == 3 cursor.close() + def test_sqlalchemy_engines_create_cursors_with_the_case_options(self, monkeypatch): + def unexpected_request(*args, **kwargs): + pytest.fail("Connecting the engine made an AWS request") + + monkeypatch.setattr(BaseClient, "_make_api_call", unexpected_request) + session = boto3.Session( + aws_access_key_id="testing", aws_secret_access_key="testing", region_name="us-west-2" + ) + settings = Settings(executor_workers=3, poll_interval=0.5) + expected = { + "sqla_rest": PyAthenaCursor, + "sqla_s3fs": S3FSCursor, + "sqla_pandas": PandasCursor, + "sqla_arrow": ArrowCursor, + "sqla_polars": PolarsCursor, + } + cases = [c for c in matrix(settings, "single", "flat") if c.family.startswith("sqla_")] + for case in cases: + engine = Adapter(settings, case, session, "scratch", "s3://out/", "temp").engine() + try: + with engine.connect() as connection: + raw = connection.connection.dbapi_connection + assert raw.work_group == settings.workgroup + assert raw.s3_staging_dir == "s3://out/" + assert raw.schema_name == "scratch" + assert raw.poll_interval == 0.5 + assert raw.result_reuse_enable is False + cursor = raw.cursor() + assert type(cursor) is expected[case.family] + assert cursor.arraysize == 1000 + if case.family in {"sqla_pandas", "sqla_arrow", "sqla_polars"}: + assert cursor._unload is (case.transport == "unload") + if case.family in {"sqla_pandas", "sqla_polars"}: + assert cursor._chunksize == case.chunksize + assert cursor._s3_max_workers == 3 + cursor.close() + finally: + engine.dispose() + + async def test_read_sql_query_consumes_frames_and_chunks(self, monkeypatch): + calls = [] + + def read_sql_query(sql, connection, chunksize=None): + calls.append((sql, chunksize)) + if chunksize is None: + return pd.DataFrame({"a": range(5)}) + return (pd.DataFrame({"a": range(n)}) for n in (2, 2, 1)) + + class Engine: + @contextmanager + def connect(self): + yield "connection" + + monkeypatch.setattr(pd, "read_sql_query", read_sql_query) + for chunksize in (None, 2): + adapter = Adapter( + Settings(), + Case("sqla_pandas", output="native", chunksize=chunksize), + None, + "scratch", + "s3://out/", + "temp", + ) + result = await measure_query(adapter, Engine(), "SELECT a", 5) + assert result["status"] == "ok", result + assert result["rows"] == 5 + assert calls == [("SELECT a", None), ("SELECT a", 2)] + + def test_pandas_dtype_backend_is_an_execute_option(self): + adapter = Adapter( + Settings(), + Case("pandas", transport="unload", output="native", dtype_backend="pyarrow"), + None, + "scratch", + "s3://out/", + "temp", + ) + assert adapter.execute_kwargs() == {"dtype_backend": "pyarrow"} + assert Adapter(Settings(), Case("pandas"), None, "s", "s3://o/", "t").execute_kwargs() == {} + @pytest.mark.parametrize("api", ["sync", "thread", "aio"]) async def test_real_api_cursors_consume_stubbed_athena_results(self, api): session = boto3.Session( diff --git a/benchmarks/tests/test_config.py b/benchmarks/tests/test_config.py index fbcd01328..d63be0c86 100644 --- a/benchmarks/tests/test_config.py +++ b/benchmarks/tests/test_config.py @@ -43,14 +43,38 @@ def test_matrix_preserves_unsupported_capabilities(): wrangler = [c for c in cases if c.family == "wrangler"] assert all(c.unsupported for c in wrangler if c.transport == "csv") assert all(c.unsupported is None for c in wrangler if c.transport != "csv") - assert all( - c.unsupported - for c in cases - if c.family == "pandas" and c.transport == "unload" and c.chunksize - ) + pandas_unload = [c for c in cases if c.family == "pandas" and c.transport == "unload"] + # PandasCursor reads UNLOAD results in chunks, with or without dtype_backend. + assert all(c.unsupported is None for c in pandas_unload) + assert {(c.chunksize, c.dtype_backend) for c in pandas_unload if c.output == "native"} == { + (chunk, backend) for chunk in (None, 10000, 100000) for backend in (None, "pyarrow") + } assert {c.api for c in cases if c.family == "cursor"} == {"sync", "thread", "aio"} +def test_sqlalchemy_cases_are_synchronous_single_query_reads(): + cases = [c for c in matrix(Settings(), "single", "flat") if c.family.startswith("sqla_")] + assert {(c.family, c.transport) for c in cases} == { + ("sqla_rest", "csv"), + ("sqla_s3fs", "csv"), + *((f"sqla_{d}", t) for d in ("pandas", "arrow", "polars") for t in ("csv", "unload")), + } + assert {c.chunksize for c in cases} == {None, 10000, 100000} + assert {(c.api, c.output, c.unsupported) for c in cases} == {("sync", "native", None)} + for suite in ("concurrent", "init"): + assert not [c for c in matrix(Settings(), suite, "flat") if c.family.startswith("sqla_")] + + +def test_json_shape_covers_pandas_and_wrangler_csv_reads(): + cases = matrix(Settings(), "single", "json") + assert {(c.family, c.transport, c.output) for c in cases} == { + ("pandas", "csv", "native"), + ("wrangler", "csv", "native"), + } + assert all(c.unsupported is None for c in cases) + assert {c.chunksize for c in cases} == {None, 10000, 100000} + + def test_concurrency_and_init_matrices_are_bounded(): cases = matrix(Settings(), "concurrent", "flat") assert {c.concurrency for c in cases} == {1, 10, 50, 100} @@ -68,6 +92,10 @@ def test_flat_projection_and_nested_projection_share_the_snapshot(): 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") + assert json == flat.replace( + ' FROM "scratch"', ', CAST(details AS JSON) AS details_json FROM "scratch"' + ) @pytest.mark.parametrize("family", ["cursor", "dict", "arrow"]) diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 6078ac0c3..f37dd25b8 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -14,6 +14,7 @@ import pytest from botocore.exceptions import ClientError from pyathena_bench.__main__ import main +from pyathena_bench.cases import Case from pyathena_bench.config import Settings from pyathena_bench.fleet import ( Filters, @@ -194,6 +195,50 @@ def test_concurrent_row_jobs_weigh_simultaneous_paging_queries(self): assert [(j["api_heavy"], j["api_weight"]) for j in jobs] == [(True, 10)] +class TestExpandJobsSelection: + def test_chunksize_filter_is_passed_to_the_job_and_named_in_its_id(self): + selection = Filters( + family=["pandas"], api=["sync"], transport=["csv"], output_kind="native" + ) + jobs = [ + job + for chunks in ([0], [100000]) + for job in expand_jobs( + SETTINGS, + ["single"], + ["small"], + ["flat"], + replace(selection, chunksize=chunks), + ) + ] + assert [j["id"] for j in jobs] == [ + "single-small-flat-pandas-sync-csv-native-c0", + "single-small-flat-pandas-sync-csv-native-c100000", + ] + assert [j["args"][-2:] for j in jobs] == [["--chunksize", "0"], ["--chunksize", "100000"]] + assert merge_jobs([jobs[:1], jobs[1:]]) == jobs + cases = [Case("pandas", output="native", chunksize=c) for c in (None, 10000, 100000)] + assert [c.chunksize for c in replace(selection, chunksize=[0]).select(cases)] == [None] + assert [c.chunksize for c in replace(selection, chunksize=[100000]).select(cases)] == [ + 100000 + ] + + def test_read_sql_query_through_the_rest_driver_pages_through_the_api(self): + jobs = expand_jobs( + SETTINGS, + ["single"], + ["small"], + ["flat"], + Filters(family=["sqla_rest", "sqla_pandas"], chunksize=[0]), + ) + heavy = {j["id"]: (j["api_heavy"], j["api_weight"]) for j in jobs} + assert heavy == { + "single-small-flat-sqla_rest-sync-csv-native-c0": (True, 1), + "single-small-flat-sqla_pandas-sync-csv-native-c0": (False, 0), + "single-small-flat-sqla_pandas-sync-unload-native-c0": (False, 0), + } + + class TestQueue: def test_submit_refuses_an_existing_queue_and_duplicate_ids(self, tmp_path): s3, queue = queue_with([job("a")], tmp_path) diff --git a/benchmarks/tests/test_report.py b/benchmarks/tests/test_report.py index 7e59e598d..4ab2067ce 100644 --- a/benchmarks/tests/test_report.py +++ b/benchmarks/tests/test_report.py @@ -109,3 +109,24 @@ def test_report_exposes_client_init_and_memory_increase(self, tmp_path): assert "Client result (s)" in markdown assert "Init (s)" in markdown assert "RSS increase (B)" in markdown + + def test_report_shows_the_most_output_files_and_the_largest_file(self, tmp_path): + base = { + "case_id": "case", + "scale": "small", + "status": "ok", + "queries": [{"total_seconds": 1}], + "successful_queries_per_second": 1, + "rss_peak_bytes": 1, + "max_threads": 1, + } + trials = [ + {**base, "output_files": {"count": 30, "bytes": 60, "max_bytes": 4, "min_bytes": 1}}, + {**base, "output_files": {"count": 31, "bytes": 60, "max_bytes": 3, "min_bytes": 1}}, + ] + (tmp_path / "trials.jsonl").write_text("".join(json.dumps(t) + "\n" for t in trials)) + report(tmp_path) + row = summarize(trials)[0] + assert (row["max_output_files"], row["max_output_file_bytes"]) == (31, 4) + assert "| 31 | 4 |" in (tmp_path / "summary.md").read_text() + assert summarize([base])[0]["max_output_files"] is None diff --git a/benchmarks/tests/test_runner.py b/benchmarks/tests/test_runner.py index 07edbf1b9..08279c1b5 100644 --- a/benchmarks/tests/test_runner.py +++ b/benchmarks/tests/test_runner.py @@ -14,7 +14,7 @@ import pytest from pyathena_bench.cases import Case from pyathena_bench.config import Settings -from pyathena_bench.runner import Observer, run, supervise +from pyathena_bench.runner import Observer, output_files, run, supervise def successful_worker(pipe, payload): @@ -132,3 +132,78 @@ def fail(payload, events, settings): assert calls[0]["warmup"] result = json.loads((tmp_path / "run/trials.jsonl").read_text()) assert result["status"] == status + + +class FakeS3Listing: + """An S3 client whose paginated listing returns fixed pages.""" + + def __init__(self, pages): + self.pages = pages + self.requests = [] + + def get_paginator(self, name): + assert name == "list_objects_v2" + return self + + def paginate(self, **kwargs): + self.requests.append(kwargs) + return iter(self.pages) + + +def test_output_files_count_parquet_files_below_the_trial_prefix(monkeypatch): + s3 = FakeS3Listing( + [ + { + "Contents": [ + {"Key": "runs/r/trials/t/q.csv", "Size": 10}, + {"Key": "runs/r/trials/t/q.csv.metadata", "Size": 1}, + {"Key": "runs/r/trials/t/q-manifest.csv", "Size": 2}, + {"Key": "runs/r/trials/t/unload/20261012_0001", "Size": 300}, + ] + }, + {"Contents": [{"Key": "runs/r/trials/t/unload/20261012_0002", "Size": 100}]}, + {}, + ] + ) + monkeypatch.setattr("pyathena_bench.runner.session", lambda settings: None) + monkeypatch.setattr("pyathena_bench.runner.client", lambda session_, service: s3) + assert output_files(Settings(), "s3://out/runs/r/trials/t/") == { + "count": 2, + "bytes": 400, + "max_bytes": 300, + "min_bytes": 100, + } + assert s3.requests == [{"Bucket": "out", "Prefix": "runs/r/trials/t/"}] + + +@pytest.mark.parametrize( + ("case", "listed"), + [(Case("pandas", transport="unload", output="native"), True), (Case("pandas"), False)], +) +def test_unload_trials_record_their_output_files(monkeypatch, tmp_path, case, listed): + monkeypatch.setattr( + "pyathena_bench.runner.validate_manifest", + lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out"}, + ) + monkeypatch.setattr("pyathena_bench.runner.validate_inputs", lambda *args: None) + monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) + monkeypatch.setattr("pyathena_bench.runner.supervise", lambda *args: {"status": "ok"}) + locations = [] + + def summarize(settings, location): + locations.append(location) + return {"count": 1, "bytes": 5, "max_bytes": 5, "min_bytes": 5} + + monkeypatch.setattr("pyathena_bench.runner.output_files", summarize) + manifest = { + "run_id": "a" * 32, + "scales": {"small": {"table": "input", "rows": 10}}, + "fixtures": {}, + } + settings = replace(Settings(), warmups=0, repetitions=1) + assert run(settings, manifest, [case], ["small"], tmp_path / "run") + trial = json.loads((tmp_path / "run/trials.jsonl").read_text()) + assert ("output_files" in trial) is listed + assert [loc.startswith(f"s3://out/runs/{'a' * 32}/trials/") for loc in locations] == ( + [True] if listed else [] + ) diff --git a/uv.lock b/uv.lock index b1874e3ef..d172a5dbe 100644 --- a/uv.lock +++ b/uv.lock @@ -1019,7 +1019,7 @@ dependencies = [ { name = "awswrangler" }, { name = "polars" }, { name = "psutil" }, - { name = "pyathena", extra = ["arrow", "pandas", "polars"] }, + { name = "pyathena", extra = ["arrow", "pandas", "polars", "sqlalchemy"] }, ] [package.dev-dependencies] @@ -1036,7 +1036,7 @@ requires-dist = [ { name = "awswrangler", specifier = ">=3.17,<4" }, { name = "polars" }, { name = "psutil", specifier = ">=7,<8" }, - { name = "pyathena", extras = ["pandas", "arrow", "polars"], editable = "." }, + { name = "pyathena", extras = ["pandas", "arrow", "polars", "sqlalchemy"], editable = "." }, ] [package.metadata.requires-dev] From 3f7fa39a5110d2eac36376839f44563d08209c96 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 15:32:01 +0900 Subject: [PATCH 2/5] Open the SQLAlchemy connection during trial setup The connection, with its boto3 clients and the dialect's first-connect initialization, was opened inside the measured query, unlike the cursor cases, which create their connection and cursor in setup. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 1 + benchmarks/pyathena_bench/adapters.py | 20 +++++++-------- benchmarks/tests/test_adapters.py | 37 ++++++++++++++++++++++----- 3 files changed, 41 insertions(+), 17 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index e5dcd76b4..8f02a31ce 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -249,6 +249,7 @@ When the temporary directory is a tmpfs, as `/tmp` is on Amazon Linux 2023, file 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, 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. diff --git a/benchmarks/pyathena_bench/adapters.py b/benchmarks/pyathena_bench/adapters.py index c3c6d8ac9..b91341e48 100644 --- a/benchmarks/pyathena_bench/adapters.py +++ b/benchmarks/pyathena_bench/adapters.py @@ -309,12 +309,10 @@ async def measure_query(adapter: Adapter, cursor: Any, sql: str, expected: int) count = consume(result, case) query_id = None # The botocore observer captures IDs even for iterators. elif case.family in SQLALCHEMY_FAMILIES: - # The "cursor" of a SQLAlchemy case is its engine; the connection stays - # open while chunks are read. - with cursor.connect() as connection: - result = pd.read_sql_query(sql, connection, chunksize=case.chunksize) - ready = time.perf_counter() - count = consume(result, case) + # The "cursor" of a SQLAlchemy case is its connection, opened during setup. + result = pd.read_sql_query(sql, cursor, chunksize=case.chunksize) + ready = time.perf_counter() + count = consume(result, case) query_id = None elif case.api == "thread": query_id, future = await asyncio.to_thread(cursor.execute, sql, **options) @@ -369,7 +367,9 @@ async def run_adapter( case = adapter.case sqlalchemy = case.family in SQLALCHEMY_FAMILIES connection = None if case.family == "wrangler" or sqlalchemy else adapter.connection() - cursors: list[Any] = [adapter.engine()] if sqlalchemy else [] + engine = adapter.engine() if sqlalchemy else None + # Like a cursor, the SQLAlchemy connection is set up outside the measurement. + cursors: list[Any] = [engine.connect()] if engine is not None else [] if case.suite != "init" and connection is not None: # AsyncCursor owns one shared pool; Aio cursors carry per-query state. cursors = [ @@ -431,11 +431,11 @@ async def run_adapter( stop.set() await monitor for cursor in cursors: - if sqlalchemy: - cursor.dispose() - elif case.api == "thread": + if case.api == "thread": cursor.close(wait=True) else: cursor.close() if connection is not None: connection.close() + if engine is not None: + engine.dispose() diff --git a/benchmarks/tests/test_adapters.py b/benchmarks/tests/test_adapters.py index eba371d7f..dff80d359 100644 --- a/benchmarks/tests/test_adapters.py +++ b/benchmarks/tests/test_adapters.py @@ -8,7 +8,6 @@ import asyncio import weakref from concurrent.futures import Future -from contextlib import contextmanager from io import BytesIO from types import SimpleNamespace @@ -198,11 +197,6 @@ def read_sql_query(sql, connection, chunksize=None): return pd.DataFrame({"a": range(5)}) return (pd.DataFrame({"a": range(n)}) for n in (2, 2, 1)) - class Engine: - @contextmanager - def connect(self): - yield "connection" - monkeypatch.setattr(pd, "read_sql_query", read_sql_query) for chunksize in (None, 2): adapter = Adapter( @@ -213,11 +207,40 @@ def connect(self): "s3://out/", "temp", ) - result = await measure_query(adapter, Engine(), "SELECT a", 5) + result = await measure_query(adapter, "connection", "SELECT a", 5) assert result["status"] == "ok", result assert result["rows"] == 5 assert calls == [("SELECT a", None), ("SELECT a", 2)] + async def test_read_sql_query_connects_during_setup(self, monkeypatch): + events = [] + + class Connection: + def close(self): + events.append("close") + + class Engine: + def connect(self): + events.append("connect") + return Connection() + + def dispose(self): + events.append("dispose") + + async def measure(adapter, cursor, sql, expected): + assert isinstance(cursor, Connection) + events.append("measure") + return {"status": "ok", "rows": expected} + + monkeypatch.setattr(Adapter, "engine", lambda self: Engine()) + monkeypatch.setattr("pyathena_bench.adapters.measure_query", measure) + adapter = Adapter( + Settings(), Case("sqla_rest", output="native"), None, "scratch", "s3://out/", "temp" + ) + result = await run_adapter(adapter, "SELECT a", 5) + assert result["status"] == "ok" + assert events == ["connect", "measure", "close", "dispose"] + def test_pandas_dtype_backend_is_an_execute_option(self): adapter = Adapter( Settings(), From c350c7398a3e22bdf1165bc18aa24ca82ea78c75 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 15:41:48 +0900 Subject: [PATCH 3/5] Keep the json workload to single queries and note failed file listings The initialization suite prepares UNLOAD fixtures too, which the json type cannot use, so the json workload has no initialization or concurrency cases. A trial whose output listing failed keeps its timing, and the report now notes the error instead of silently taking the file maxima from the other trials. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 4 ++-- benchmarks/pyathena_bench/cases.py | 7 +++++-- benchmarks/pyathena_bench/report.py | 3 +++ benchmarks/tests/test_config.py | 3 +++ benchmarks/tests/test_report.py | 4 ++++ 5 files changed, 17 insertions(+), 4 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 8f02a31ce..63fd9041e 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -30,7 +30,7 @@ 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 only native pandas and Wrangler reads of CSV results, because Parquet cannot hold the `json` type. +The json workload adds `CAST(details AS JSON) AS details_json` to the flat projection; it covers only native pandas and Wrangler reads of CSV results in the single-query suite, because Parquet cannot hold the `json` type. No workload converts every library's output into a common Python object representation. ## Dependencies and local checks @@ -423,7 +423,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 UNLOAD or CTAS trial, `run` lists the trial's output location and records the number and sizes of the Parquet files in `output_files`; the summary shows the most files and the largest file. +After each successful UNLOAD or CTAS trial, `run` lists the trial's output location and records the number and sizes of the Parquet files in `output_files`; the summary shows the most files and the largest file, and notes a trial whose listing failed. Regenerate reports after interruption with: ```bash diff --git a/benchmarks/pyathena_bench/cases.py b/benchmarks/pyathena_bench/cases.py index 8421d37cd..6237eecf6 100644 --- a/benchmarks/pyathena_bench/cases.py +++ b/benchmarks/pyathena_bench/cases.py @@ -49,8 +49,9 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: """Return every case of a suite and result shape, including unsupported ones. The ``json`` shape adds a ``json`` column to the flat projection and covers - only the native pandas and Wrangler reads of CSV results: Parquet cannot - hold the ``json`` type. + only the native pandas and Wrangler reads of CSV results in the single-query + suite: Parquet cannot hold the ``json`` type, and the initialization suite + prepares UNLOAD fixtures too. Args: settings: Configuration with the arraysizes, chunk sizes, and concurrency. @@ -65,6 +66,8 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: """ if suite not in {"single", "concurrent", "init"} or shape not in SHAPES: raise ValueError("Unknown suite or shape") + if shape == "json" and suite != "single": + return [] cases: list[Case] = [] arraysizes = settings.arraysizes if suite == "single" else [1000] for family in FAMILIES: diff --git a/benchmarks/pyathena_bench/report.py b/benchmarks/pyathena_bench/report.py index 41776876b..993962746 100644 --- a/benchmarks/pyathena_bench/report.py +++ b/benchmarks/pyathena_bench/report.py @@ -70,6 +70,9 @@ def summarize(trials: list[dict[str, Any]]) -> list[dict[str, Any]]: reasons = set() for trial in group: if trial["status"] == "ok": + if "output_files_error" in trial: + # The timing stands, but the file columns miss this trial. + reasons.add(f"Output files not listed: {trial['output_files_error']}") continue prefix = "Warmup: " if trial.get("warmup") else "" failures = [q for q in trial.get("queries", []) if q["status"] != "ok"] diff --git a/benchmarks/tests/test_config.py b/benchmarks/tests/test_config.py index d63be0c86..e9e86f78c 100644 --- a/benchmarks/tests/test_config.py +++ b/benchmarks/tests/test_config.py @@ -73,6 +73,9 @@ def test_json_shape_covers_pandas_and_wrangler_csv_reads(): } assert all(c.unsupported is None for c in cases) assert {c.chunksize for c in cases} == {None, 10000, 100000} + # The initialization suite would also prepare an UNLOAD fixture, which json rejects. + assert matrix(Settings(), "init", "json") == [] + assert matrix(Settings(), "concurrent", "json") == [] def test_concurrency_and_init_matrices_are_bounded(): diff --git a/benchmarks/tests/test_report.py b/benchmarks/tests/test_report.py index 4ab2067ce..35108c04a 100644 --- a/benchmarks/tests/test_report.py +++ b/benchmarks/tests/test_report.py @@ -130,3 +130,7 @@ def test_report_shows_the_most_output_files_and_the_largest_file(self, tmp_path) assert (row["max_output_files"], row["max_output_file_bytes"]) == (31, 4) assert "| 31 | 4 |" in (tmp_path / "summary.md").read_text() assert summarize([base])[0]["max_output_files"] is None + failed_listing = {**base, "output_files_error": "ClientError: denied"} + row = summarize([trials[0], failed_listing])[0] + assert (row["successful_trials"], row["max_output_files"]) == (2, 30) + assert row["notes"] == "Output files not listed: ClientError: denied" From 78db697f5b8322aa3f8fff9e387a44caea4dfe99 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 16:24:06 +0900 Subject: [PATCH 4/5] Repair the code-review findings on the benchmark selection - Mark json UNLOAD/CTAS reads unsupported instead of dropping them. - Summarize UNLOAD/CTAS output files from the query manifests that Athena writes for PyAthena and Wrangler, with one S3 client per run; a trial without a manifest records an error instead of a guess. - Note listing failures for measured trials only, as the columns do. - Name chunk selections in job IDs and arguments sorted and without repeats, so equal selections cannot bypass the duplicate check. - Dispose of the engine when the SQLAlchemy connection fails, and keep the cursor execute options in Adapter.execute_kwargs(); document how the SQLAlchemy path disables result reuse and the client cache. - README: describe which drivers the xlarge selection reads, and read only CSV for the xlarge json workload. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 12 ++-- benchmarks/pyathena_bench/adapters.py | 60 ++++++++++------- benchmarks/pyathena_bench/cases.py | 14 ++-- benchmarks/pyathena_bench/fleet.py | 8 ++- benchmarks/pyathena_bench/report.py | 2 +- benchmarks/pyathena_bench/runner.py | 45 +++++++++---- benchmarks/tests/test_adapters.py | 24 ++++++- benchmarks/tests/test_config.py | 10 ++- benchmarks/tests/test_fleet.py | 7 ++ benchmarks/tests/test_report.py | 3 + benchmarks/tests/test_runner.py | 93 ++++++++++++++++++++------- 11 files changed, 195 insertions(+), 83 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 63fd9041e..eabf75564 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -30,7 +30,7 @@ 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 only native pandas and Wrangler reads of CSV results in the single-query suite, because Parquet cannot hold the `json` type. +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 @@ -248,7 +248,7 @@ When the temporary directory is a tmpfs, as `/tmp` is on Amazon Linux 2023, file | 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, and their cursors take the same cursor options; with `chunksize`, the pandas and Polars drivers' cursors also read the result file in chunks. +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. @@ -342,7 +342,7 @@ For example, the following selection, with `chunksizes = [10000, 100000, 1000000 - 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 including the json workload (five), and `read_sql_query()` for every driver without chunks and with 100,000-row chunks (five; three for `sqla_rest`, split); -- 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()` for the result-file drivers flat without chunks and with chunks, and nested with chunks only, all with three repetitions. +- 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. @@ -369,7 +369,7 @@ 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 --chunksize 0 100000 \ +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 @@ -382,7 +382,7 @@ uv run --no-sync python -m pyathena_bench --config results/fleet.toml queue \ 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 host-hours on `r7g.2xlarge` hosts, and the longest, xlarge nested pandas CSV, about 75 minutes. +The selection has 99 jobs, three of which only record the unsupported json UNLOAD and CTAS cases; 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//` in the scratch bucket, and writes the `queue.json` marker last. Workers refuse a queue without the marker. @@ -423,7 +423,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 UNLOAD or CTAS trial, `run` lists the trial's output location and records the number and sizes of the Parquet files in `output_files`; the summary shows the most files and the largest file, and notes a trial whose listing failed. +After each successful UNLOAD or CTAS trial, `run` reads the query manifests (`-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 diff --git a/benchmarks/pyathena_bench/adapters.py b/benchmarks/pyathena_bench/adapters.py index b91341e48..8d699abba 100644 --- a/benchmarks/pyathena_bench/adapters.py +++ b/benchmarks/pyathena_bench/adapters.py @@ -225,7 +225,9 @@ def engine(self) -> Any: connections create cursors with the cursor options of the driver's family, including ``chunksize`` for the pandas and Polars drivers, so a chunked read also streams the result file. ``NullPool`` opens one - connection per query. + connection per query. Result reuse is disabled through the connection's + ``result_reuse_enable``, and the client cache stays off because the + dialect executes statements without ``cache_size``, which defaults to 0. Returns: The engine. @@ -242,12 +244,18 @@ def engine(self) -> Any: ) def execute_kwargs(self) -> dict[str, Any]: - """Return the ``execute()`` options of the case besides cache and reuse settings. + """Return the ``execute()`` options of a cursor case. + + The client cache and Athena result reuse are disabled, and pandas cases + that set ``dtype_backend`` pass it. Returns: - ``dtype_backend`` for the pandas cases that set it, or no options. + The keyword arguments for the cursor's ``execute()``. """ - return {"dtype_backend": self.case.dtype_backend} if self.case.dtype_backend else {} + options: dict[str, Any] = {"cache_size": 0, "result_reuse_enable": False} + if self.case.dtype_backend: + options["dtype_backend"] = self.case.dtype_backend + return options def wrangler(self, sql: str) -> Any: options: dict[str, Any] = { @@ -302,7 +310,6 @@ async def measure_query(adapter: Adapter, cursor: Any, sql: str, expected: int) case = adapter.case result = None try: - options = {"cache_size": 0, "result_reuse_enable": False, **adapter.execute_kwargs()} if case.family == "wrangler": result = adapter.wrangler(sql) ready = time.perf_counter() @@ -315,17 +322,19 @@ async def measure_query(adapter: Adapter, cursor: Any, sql: str, expected: int) count = consume(result, case) query_id = None elif case.api == "thread": - query_id, future = await asyncio.to_thread(cursor.execute, sql, **options) + query_id, future = await asyncio.to_thread( + cursor.execute, sql, **adapter.execute_kwargs() + ) result = await asyncio.wrap_future(future) ready = time.perf_counter() count = await asyncio.to_thread(consume, result, case) elif case.api == "aio": - await cursor.execute(sql, **options) + await cursor.execute(sql, **adapter.execute_kwargs()) ready = time.perf_counter() query_id = cursor.query_id count = await consume_aio(cursor, case) else: - cursor.execute(sql, **options) + cursor.execute(sql, **adapter.execute_kwargs()) ready = time.perf_counter() query_id = cursor.query_id count = consume(cursor.result_set, case) @@ -368,20 +377,24 @@ async def run_adapter( sqlalchemy = case.family in SQLALCHEMY_FAMILIES connection = None if case.family == "wrangler" or sqlalchemy else adapter.connection() engine = adapter.engine() if sqlalchemy else None - # Like a cursor, the SQLAlchemy connection is set up outside the measurement. - cursors: list[Any] = [engine.connect()] if engine is not None else [] - if case.suite != "init" and connection is not None: - # AsyncCursor owns one shared pool; Aio cursors carry per-query state. - cursors = [ - adapter.cursor(connection) for _ in range(case.concurrency if case.api == "aio" else 1) - ] - setup_seconds = time.perf_counter() - setup_start - samples: list[float] = [] - stop = asyncio.Event() - monitor = asyncio.create_task(heartbeat(adapter.settings.heartbeat_interval, stop, samples)) - await asyncio.sleep(0) - start = time.perf_counter() + cursors: list[Any] = [] + monitor = None try: + if engine is not None: + # Like a cursor, the SQLAlchemy connection is set up outside the measurement. + cursors = [engine.connect()] + if case.suite != "init" and connection is not None: + # AsyncCursor owns one shared pool; Aio cursors carry per-query state. + cursors = [ + adapter.cursor(connection) + for _ in range(case.concurrency if case.api == "aio" else 1) + ] + setup_seconds = time.perf_counter() - setup_start + samples: list[float] = [] + stop = asyncio.Event() + monitor = asyncio.create_task(heartbeat(adapter.settings.heartbeat_interval, stop, samples)) + await asyncio.sleep(0) + start = time.perf_counter() if case.suite == "init": if fixture is None: raise ValueError("Missing completed query fixture") @@ -428,8 +441,9 @@ async def run_adapter( else None, } finally: - stop.set() - await monitor + if monitor is not None: + stop.set() + await monitor for cursor in cursors: if case.api == "thread": cursor.close(wait=True) diff --git a/benchmarks/pyathena_bench/cases.py b/benchmarks/pyathena_bench/cases.py index 6237eecf6..5ffbc56b4 100644 --- a/benchmarks/pyathena_bench/cases.py +++ b/benchmarks/pyathena_bench/cases.py @@ -49,9 +49,10 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: """Return every case of a suite and result shape, including unsupported ones. The ``json`` shape adds a ``json`` column to the flat projection and covers - only the native pandas and Wrangler reads of CSV results in the single-query - suite: Parquet cannot hold the ``json`` type, and the initialization suite - prepares UNLOAD fixtures too. + the native pandas and Wrangler reads in the single-query suite. Their UNLOAD + and CTAS reads are marked unsupported, because Parquet cannot hold the + ``json`` type; the initialization suite, which prepares UNLOAD fixtures too, + has no json cases. Args: settings: Configuration with the arraysizes, chunk sizes, and concurrency. @@ -125,11 +126,10 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: if shape == "json": cases = [ c + if c.transport == "csv" + else replace(c, unsupported="Athena cannot write the json type to Parquet") for c in cases - if c.family in {"pandas", "wrangler"} - and c.transport == "csv" - and c.output == "native" - and c.dtype_backend is None + if c.family in {"pandas", "wrangler"} and c.output == "native" ] if suite == "concurrent": # Keep the scaling suite bounded; chunk and arraysize sweeps belong to single. diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index fc13b97cd..146715275 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -176,9 +176,11 @@ def expand_jobs( base += f"-a{arraysize}" args += ["--arraysize", str(arraysize)] if filters.chunksize: - # The job runs the same chunk selection; the ID keeps it unique. - base += "-c" + "_".join(map(str, filters.chunksize)) - args += ["--chunksize", *map(str, filters.chunksize)] + # The job runs the same chunk selection, which its ID names + # in one order so that equal selections share an ID. + chunks = [str(c) for c in sorted(set(filters.chunksize))] + base += "-c" + "_".join(chunks) + args += ["--chunksize", *chunks] runnable = [c for c in cases if c.unsupported is None] pages = max((api_pages(c, rows) for c in runnable), default=0) api_weight = max( diff --git a/benchmarks/pyathena_bench/report.py b/benchmarks/pyathena_bench/report.py index 993962746..255b2ed10 100644 --- a/benchmarks/pyathena_bench/report.py +++ b/benchmarks/pyathena_bench/report.py @@ -70,7 +70,7 @@ def summarize(trials: list[dict[str, Any]]) -> list[dict[str, Any]]: reasons = set() for trial in group: if trial["status"] == "ok": - if "output_files_error" in trial: + if "output_files_error" in trial and not trial.get("warmup"): # The timing stands, but the file columns miss this trial. reasons.add(f"Output files not listed: {trial['output_files_error']}") continue diff --git a/benchmarks/pyathena_bench/runner.py b/benchmarks/pyathena_bench/runner.py index 02f8331e0..e6686e66d 100644 --- a/benchmarks/pyathena_bench/runner.py +++ b/benchmarks/pyathena_bench/runner.py @@ -308,28 +308,41 @@ def supervise( return result -def output_files(settings: Settings, location: str) -> dict[str, int]: - """Summarize the data files that a trial wrote below its output location. +def output_files(s3: Any, location: str) -> dict[str, int]: + """Summarize the data files that a trial's UNLOAD or CTAS queries wrote. - UNLOAD and CTAS write Parquet files there; query result CSV files, their - metadata, and manifests are excluded. + Athena writes a ``-manifest.csv`` that lists the data files of + each UNLOAD or CTAS query below the query's output location, for PyAthena + and Wrangler alike, so the manifests below the trial's location name the + files. Sizes come from the listing of that location, or from + ``HeadObject`` for a file outside it. Args: - settings: Configuration with the AWS region and profile. + s3: S3 client. location: The trial's ``s3://`` output prefix. Returns: The number of files and their total, largest, and smallest sizes in bytes. + + Raises: + ValueError: If no manifest is below the location. """ bucket, _, prefix = location.removeprefix("s3://").partition("/") - sizes = [ - item["Size"] - for page in client(session(settings), "s3") - .get_paginator("list_objects_v2") - .paginate(Bucket=bucket, Prefix=prefix) - for item in page.get("Contents", []) - if not item["Key"].endswith((".csv", ".metadata", ".txt", ".json")) - ] + listed: dict[str, int] = {} + for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix): + listed.update((item["Key"], item["Size"]) for item in page.get("Contents", [])) + manifests = [key for key in listed if key.endswith("-manifest.csv")] + if not manifests: + raise ValueError(f"No query manifest below {location}") + sizes = [] + for manifest in manifests: + body: str = s3.get_object(Bucket=bucket, Key=manifest)["Body"].read().decode() + for line in filter(None, body.splitlines()): + file_bucket, _, key = line.removeprefix("s3://").partition("/") + size = listed.get(key) if file_bucket == bucket else None + if size is None: + size = s3.head_object(Bucket=file_bucket, Key=key)["ContentLength"] + sizes.append(size) return { "count": len(sizes), "bytes": sum(sizes), @@ -370,6 +383,8 @@ def run( write_json(output / "environment.json", environment(settings)) write_json(output / "manifest.json", manifest) success = True + # Created on first use and shared by the run's output file summaries. + s3: Any = None for scale in scales: for case in cases: if case.unsupported: @@ -412,7 +427,9 @@ def run( and case.suite != "init" ): try: - result["output_files"] = output_files(settings, payload["output"]) + if s3 is None: + s3 = client(session(settings), "s3") + result["output_files"] = output_files(s3, payload["output"]) except Exception as exc: # The measurement stands; only the file summary is missing. result["output_files_error"] = f"{type(exc).__name__}: {exc}" diff --git a/benchmarks/tests/test_adapters.py b/benchmarks/tests/test_adapters.py index dff80d359..0d41523f9 100644 --- a/benchmarks/tests/test_adapters.py +++ b/benchmarks/tests/test_adapters.py @@ -241,6 +241,24 @@ async def measure(adapter, cursor, sql, expected): assert result["status"] == "ok" assert events == ["connect", "measure", "close", "dispose"] + async def test_failed_sqlalchemy_connect_disposes_the_engine(self, monkeypatch): + events = [] + + class Engine: + def connect(self): + raise OSError("connect failed") + + def dispose(self): + events.append("dispose") + + monkeypatch.setattr(Adapter, "engine", lambda self: Engine()) + adapter = Adapter( + Settings(), Case("sqla_rest", output="native"), None, "scratch", "s3://out/", "temp" + ) + with pytest.raises(OSError, match="connect failed"): + await run_adapter(adapter, "SELECT a", 5) + assert events == ["dispose"] + def test_pandas_dtype_backend_is_an_execute_option(self): adapter = Adapter( Settings(), @@ -250,8 +268,10 @@ def test_pandas_dtype_backend_is_an_execute_option(self): "s3://out/", "temp", ) - assert adapter.execute_kwargs() == {"dtype_backend": "pyarrow"} - assert Adapter(Settings(), Case("pandas"), None, "s", "s3://o/", "t").execute_kwargs() == {} + disabled = {"cache_size": 0, "result_reuse_enable": False} + assert adapter.execute_kwargs() == {**disabled, "dtype_backend": "pyarrow"} + plain = Adapter(Settings(), Case("pandas"), None, "s", "s3://o/", "t") + assert plain.execute_kwargs() == disabled @pytest.mark.parametrize("api", ["sync", "thread", "aio"]) async def test_real_api_cursors_consume_stubbed_athena_results(self, api): diff --git a/benchmarks/tests/test_config.py b/benchmarks/tests/test_config.py index e9e86f78c..ae56c1c8e 100644 --- a/benchmarks/tests/test_config.py +++ b/benchmarks/tests/test_config.py @@ -65,14 +65,18 @@ def test_sqlalchemy_cases_are_synchronous_single_query_reads(): assert not [c for c in matrix(Settings(), suite, "flat") if c.family.startswith("sqla_")] -def test_json_shape_covers_pandas_and_wrangler_csv_reads(): +def test_json_shape_covers_pandas_and_wrangler_reads(): cases = matrix(Settings(), "single", "json") assert {(c.family, c.transport, c.output) for c in cases} == { ("pandas", "csv", "native"), + ("pandas", "unload", "native"), ("wrangler", "csv", "native"), + ("wrangler", "ctas", "native"), + ("wrangler", "unload", "native"), } - assert all(c.unsupported is None for c in cases) - assert {c.chunksize for c in cases} == {None, 10000, 100000} + # Parquet cannot hold json, so UNLOAD and CTAS stay visible as unsupported. + assert all((c.unsupported is None) == (c.transport == "csv") for c in cases) + assert {c.chunksize for c in cases if c.transport == "csv"} == {None, 10000, 100000} # The initialization suite would also prepare an UNLOAD fixture, which json rejects. assert matrix(Settings(), "init", "json") == [] assert matrix(Settings(), "concurrent", "json") == [] diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index f37dd25b8..e66b1d3af 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -217,6 +217,13 @@ def test_chunksize_filter_is_passed_to_the_job_and_named_in_its_id(self): ] assert [j["args"][-2:] for j in jobs] == [["--chunksize", "0"], ["--chunksize", "100000"]] assert merge_jobs([jobs[:1], jobs[1:]]) == jobs + # Equal selections in another order or with repeats get the same job. + again = expand_jobs( + SETTINGS, ["single"], ["small"], ["flat"], replace(selection, chunksize=[100000, 0, 0]) + ) + assert [(j["id"], j["args"][-3:]) for j in again] == [ + ("single-small-flat-pandas-sync-csv-native-c0_100000", ["--chunksize", "0", "100000"]) + ] cases = [Case("pandas", output="native", chunksize=c) for c in (None, 10000, 100000)] assert [c.chunksize for c in replace(selection, chunksize=[0]).select(cases)] == [None] assert [c.chunksize for c in replace(selection, chunksize=[100000]).select(cases)] == [ diff --git a/benchmarks/tests/test_report.py b/benchmarks/tests/test_report.py index 35108c04a..4a64b0ea1 100644 --- a/benchmarks/tests/test_report.py +++ b/benchmarks/tests/test_report.py @@ -134,3 +134,6 @@ def test_report_shows_the_most_output_files_and_the_largest_file(self, tmp_path) row = summarize([trials[0], failed_listing])[0] assert (row["successful_trials"], row["max_output_files"]) == (2, 30) assert row["notes"] == "Output files not listed: ClientError: denied" + # A warmup is not in the columns, so its listing failure is not noted either. + row = summarize([trials[0], {**failed_listing, "warmup": True}])[0] + assert row["notes"] == "" diff --git a/benchmarks/tests/test_runner.py b/benchmarks/tests/test_runner.py index 08279c1b5..54393c7a1 100644 --- a/benchmarks/tests/test_runner.py +++ b/benchmarks/tests/test_runner.py @@ -9,6 +9,7 @@ import os import time from dataclasses import replace +from io import BytesIO from pathlib import Path import pytest @@ -134,11 +135,13 @@ def fail(payload, events, settings): assert result["status"] == status -class FakeS3Listing: - """An S3 client whose paginated listing returns fixed pages.""" +class FakeS3: + """An S3 client with a fixed paginated listing and object bodies.""" - def __init__(self, pages): + def __init__(self, pages, bodies=None, heads=None): self.pages = pages + self.bodies = bodies or {} + self.heads = heads or {} self.requests = [] def get_paginator(self, name): @@ -146,34 +149,62 @@ def get_paginator(self, name): return self def paginate(self, **kwargs): - self.requests.append(kwargs) + self.requests.append(("list", kwargs)) return iter(self.pages) + def get_object(self, **kwargs): + self.requests.append(("get", kwargs["Key"])) + return {"Body": BytesIO(self.bodies[kwargs["Key"]].encode())} -def test_output_files_count_parquet_files_below_the_trial_prefix(monkeypatch): - s3 = FakeS3Listing( + def head_object(self, **kwargs): + self.requests.append(("head", kwargs["Bucket"], kwargs["Key"])) + return {"ContentLength": self.heads[(kwargs["Bucket"], kwargs["Key"])]} + + +def test_output_files_follow_the_query_manifests(): + root = "runs/r/trials/t/" + s3 = FakeS3( [ { "Contents": [ - {"Key": "runs/r/trials/t/q.csv", "Size": 10}, - {"Key": "runs/r/trials/t/q.csv.metadata", "Size": 1}, - {"Key": "runs/r/trials/t/q-manifest.csv", "Size": 2}, - {"Key": "runs/r/trials/t/unload/20261012_0001", "Size": 300}, + {"Key": f"{root}q.csv", "Size": 10}, + {"Key": f"{root}u-manifest.csv", "Size": 2}, + {"Key": f"{root}u.metadata", "Size": 1}, + {"Key": f"{root}unload/20261012_0001", "Size": 300}, + # Not in a manifest, so not a data file of the query. + {"Key": f"{root}unload/leftover", "Size": 999}, + ] + }, + { + "Contents": [ + {"Key": f"{root}unload/20261012_0002", "Size": 100}, + {"Key": f"{root}tables/c-manifest.csv", "Size": 2}, ] }, - {"Contents": [{"Key": "runs/r/trials/t/unload/20261012_0002", "Size": 100}]}, {}, - ] + ], + bodies={ + f"{root}u-manifest.csv": ( + f"s3://out/{root}unload/20261012_0001\ns3://out/{root}unload/20261012_0002\n" + ), + f"{root}tables/c-manifest.csv": "s3://elsewhere/ctas/f1\n", + }, + heads={("elsewhere", "ctas/f1"): 50}, ) - monkeypatch.setattr("pyathena_bench.runner.session", lambda settings: None) - monkeypatch.setattr("pyathena_bench.runner.client", lambda session_, service: s3) - assert output_files(Settings(), "s3://out/runs/r/trials/t/") == { - "count": 2, - "bytes": 400, + assert output_files(s3, f"s3://out/{root}") == { + "count": 3, + "bytes": 450, "max_bytes": 300, - "min_bytes": 100, + "min_bytes": 50, } - assert s3.requests == [{"Bucket": "out", "Prefix": "runs/r/trials/t/"}] + assert s3.requests[0] == ("list", {"Bucket": "out", "Prefix": root}) + assert ("head", "elsewhere", "ctas/f1") in s3.requests + + +def test_output_files_without_a_manifest_is_an_error(): + s3 = FakeS3([{"Contents": [{"Key": "runs/r/trials/t/q.csv", "Size": 10}]}]) + with pytest.raises(ValueError, match="No query manifest"): + output_files(s3, "s3://out/runs/r/trials/t/") @pytest.mark.parametrize( @@ -189,21 +220,35 @@ def test_unload_trials_record_their_output_files(monkeypatch, tmp_path, case, li monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) monkeypatch.setattr("pyathena_bench.runner.supervise", lambda *args: {"status": "ok"}) locations = [] + clients = [] - def summarize(settings, location): + def summarize(s3, location): locations.append(location) + clients.append(s3) return {"count": 1, "bytes": 5, "max_bytes": 5, "min_bytes": 5} + created = [] + + def make_client(session_, service): + assert service == "s3" + created.append(object()) + return created[-1] + monkeypatch.setattr("pyathena_bench.runner.output_files", summarize) + monkeypatch.setattr("pyathena_bench.runner.session", lambda settings: None) + monkeypatch.setattr("pyathena_bench.runner.client", make_client) manifest = { "run_id": "a" * 32, "scales": {"small": {"table": "input", "rows": 10}}, "fixtures": {}, } - settings = replace(Settings(), warmups=0, repetitions=1) + settings = replace(Settings(), warmups=0, repetitions=2) assert run(settings, manifest, [case], ["small"], tmp_path / "run") - trial = json.loads((tmp_path / "run/trials.jsonl").read_text()) - assert ("output_files" in trial) is listed + trials = [json.loads(line) for line in (tmp_path / "run/trials.jsonl").read_text().splitlines()] + assert all(("output_files" in trial) is listed for trial in trials) assert [loc.startswith(f"s3://out/runs/{'a' * 32}/trials/") for loc in locations] == ( - [True] if listed else [] + [True, True] if listed else [] ) + # One client serves every trial of the run. + assert len(created) == (1 if listed else 0) + assert clients == created * 2 From 610edf3f6bfa474f537c77b62f55d44aabb88974 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 16:46:38 +0900 Subject: [PATCH 5/5] Apply the simplify cleanups to the benchmark selection - One SQLAlchemy family mapping to drivers, RESULT_FILE_FAMILIES shared with the cursor matrix, and ALL_FAMILIES for the CLI; the engine reads the mapping. - Generate the pyarrow dtype variants in the single-query suite only. - One branch for the DataFrame reads of Wrangler and read_sql_query(). - output_files(): parse locations with parse_output_location(), skip warmups, and use one S3 client created before the trials. - Build the case filters once in the CLI. - Name and pass each job's actual chunk sizes under a chunk filter. - Share the run() stubs in tests. - Selection: read_sql_query() through the REST driver without chunks only, and the large json workload as CSV reads on its own line. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 12 ++-- benchmarks/pyathena_bench/__main__.py | 19 ++---- benchmarks/pyathena_bench/adapters.py | 26 ++++----- benchmarks/pyathena_bench/cases.py | 62 +++++++++++--------- benchmarks/pyathena_bench/fleet.py | 10 ++-- benchmarks/pyathena_bench/runner.py | 11 ++-- benchmarks/tests/test_adapters.py | 17 ++++-- benchmarks/tests/test_config.py | 8 ++- benchmarks/tests/test_fleet.py | 11 ++++ benchmarks/tests/test_runner.py | 83 ++++++++++++--------------- 10 files changed, 136 insertions(+), 123 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index eabf75564..9d1995573 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -341,7 +341,7 @@ Several `jobs` selections can make up one bounded matrix. 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 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 including the json workload (five), and `read_sql_query()` for every driver without chunks and with 100,000-row chunks (five; three for `sqla_rest`, split); +- 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. @@ -359,11 +359,13 @@ 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 json --api sync \ +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 100000 \ +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 @@ -382,7 +384,7 @@ uv run --no-sync python -m pyathena_bench --config results/fleet.toml queue \ 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 99 jobs, three of which only record the unsupported json UNLOAD and CTAS cases; 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. +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//` in the scratch bucket, and writes the `queue.json` marker last. Workers refuse a queue without the marker. @@ -423,7 +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 UNLOAD or CTAS trial, `run` reads the query manifests (`-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. +After each successful measured UNLOAD or CTAS trial, `run` reads the query manifests (`-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 diff --git a/benchmarks/pyathena_bench/__main__.py b/benchmarks/pyathena_bench/__main__.py index 082f2ee64..f4909acbe 100644 --- a/benchmarks/pyathena_bench/__main__.py +++ b/benchmarks/pyathena_bench/__main__.py @@ -26,7 +26,7 @@ stack_resources, validate_manifest, ) -from pyathena_bench.cases import FAMILIES, SHAPES, SQLALCHEMY_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, @@ -76,9 +76,7 @@ def parser() -> argparse.ArgumentParser: 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", *SQLALCHEMY_FAMILIES), nargs="+" - ) + command.add_argument("--family", choices=ALL_FAMILIES, nargs="+") command.add_argument( "--api", choices=("sync", "thread", "aio", "direct", "to_thread"), nargs="+" ) @@ -195,7 +193,7 @@ 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": + else: filters = Filters( args.family, args.api, args.transport, args.output_kind, args.arraysize, args.chunksize ) @@ -204,6 +202,7 @@ def main(argv: list[str] | None = None) -> int: } if overrides: settings = replace(settings, **overrides) + if args.command == "jobs": jobs = expand_jobs( settings, args.suite, @@ -214,15 +213,7 @@ 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, args.chunksize - ) + 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") diff --git a/benchmarks/pyathena_bench/adapters.py b/benchmarks/pyathena_bench/adapters.py index 8d699abba..b0b888ee7 100644 --- a/benchmarks/pyathena_bench/adapters.py +++ b/benchmarks/pyathena_bench/adapters.py @@ -190,7 +190,9 @@ def cursor_kwargs(self, family: str) -> dict[str, Any]: """Return the cursor options of the case for a cursor family. Args: - family: The cursor family, which a SQLAlchemy case maps from its driver. + family: The cursor family, or the SQLAlchemy driver of a + ``read_sql_query()`` case, whose cursor takes the options of the + family of the same name. Returns: The keyword arguments that the family's cursor accepts. @@ -223,7 +225,7 @@ def engine(self) -> Any: The engine connects with the same arguments as the cursor cases, and its connections create cursors with the cursor options of the driver's - family, including ``chunksize`` for the pandas and Polars drivers, so a + name, including ``chunksize`` for the pandas and Polars drivers, so a chunked read also streams the result file. ``NullPool`` opens one connection per query. Result reuse is disabled through the connection's ``result_reuse_enable``, and the client cache stays off because the @@ -232,14 +234,13 @@ def engine(self) -> Any: Returns: The engine. """ - driver = self.case.family.removeprefix("sqla_") - family = "cursor" if driver == "rest" else driver + driver = SQLALCHEMY_FAMILIES[self.case.family] return create_engine( f"awsathena+{driver}://athena.{self.settings.region}.amazonaws.com:443/{self.database}", poolclass=NullPool, connect_args={ **self.connection_kwargs(), - "cursor_kwargs": self.cursor_kwargs(family), + "cursor_kwargs": self.cursor_kwargs(driver), }, ) @@ -310,17 +311,16 @@ async def measure_query(adapter: Adapter, cursor: Any, sql: str, expected: int) case = adapter.case result = None try: - if case.family == "wrangler": - result = adapter.wrangler(sql) - ready = time.perf_counter() - count = consume(result, case) - query_id = None # The botocore observer captures IDs even for iterators. - elif case.family in SQLALCHEMY_FAMILIES: + if case.family == "wrangler" or case.family in SQLALCHEMY_FAMILIES: # The "cursor" of a SQLAlchemy case is its connection, opened during setup. - result = pd.read_sql_query(sql, cursor, chunksize=case.chunksize) + result = ( + adapter.wrangler(sql) + if case.family == "wrangler" + else pd.read_sql_query(sql, cursor, chunksize=case.chunksize) + ) ready = time.perf_counter() count = consume(result, case) - query_id = None + query_id = None # The botocore observer captures IDs even for iterators. elif case.api == "thread": query_id, future = await asyncio.to_thread( cursor.execute, sql, **adapter.execute_kwargs() diff --git a/benchmarks/pyathena_bench/cases.py b/benchmarks/pyathena_bench/cases.py index 5ffbc56b4..386e1b5e8 100644 --- a/benchmarks/pyathena_bench/cases.py +++ b/benchmarks/pyathena_bench/cases.py @@ -14,11 +14,20 @@ from pyathena_bench.config import Settings FAMILIES = ("cursor", "dict", "s3fs", "pandas", "arrow", "polars") -# pandas.read_sql_query() through a SQLAlchemy engine of each PyAthena driver. -SQLALCHEMY_FAMILIES = ("sqla_rest", "sqla_pandas", "sqla_arrow", "sqla_polars", "sqla_s3fs") +# The cursor families that read result files as DataFrames or tables, also from UNLOAD. +RESULT_FILE_FAMILIES = ("pandas", "arrow", "polars") +# pandas.read_sql_query() through a SQLAlchemy engine, by the PyAthena driver it uses. +# A driver's cursor takes the options of the cursor family of the same name; the rest +# driver's Cursor takes those of no family in particular. +SQLALCHEMY_FAMILIES = { + "sqla_rest": "rest", + "sqla_pandas": "pandas", + "sqla_arrow": "arrow", + "sqla_polars": "polars", + "sqla_s3fs": "s3fs", +} +ALL_FAMILIES = (*FAMILIES, "wrangler", *SQLALCHEMY_FAMILIES) SHAPES = ("flat", "nested", "json") -# The SQLAlchemy drivers that read result files and can read UNLOAD results. -RESULT_FILE_SQLALCHEMY = {"sqla_pandas", "sqla_arrow", "sqla_polars"} @dataclass(frozen=True) @@ -74,27 +83,30 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: for family in FAMILIES: for api in ("sync", "thread", "aio"): cases.extend(Case(family, api, arraysize=a, shape=shape) for a in arraysizes) - if family in {"pandas", "arrow", "polars"}: + if family in RESULT_FILE_FAMILIES: for transport in ("csv", "unload"): chunks: list[int | None] = ( [None] if family == "arrow" else [None, *settings.chunksizes] ) - for chunk in chunks: - cases.append( - Case(family, api, transport, "native", chunksize=chunk, shape=shape) + # Single queries also measure pandas UNLOAD reads with pyarrow dtypes. + backends: tuple[str | None, ...] = ( + (None, "pyarrow") + if suite == "single" and family == "pandas" and transport == "unload" + else (None,) + ) + cases.extend( + Case( + family, + api, + transport, + "native", + chunksize=chunk, + shape=shape, + dtype_backend=backend, ) - if family == "pandas" and transport == "unload": - cases.append( - Case( - family, - api, - transport, - "native", - chunksize=chunk, - shape=shape, - dtype_backend="pyarrow", - ) - ) + for chunk in chunks + for backend in backends + ) cases.extend( Case(family, api, "unload", arraysize=a, shape=shape) for a in arraysizes ) @@ -116,8 +128,8 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: ) ) if suite == "single": - for family in SQLALCHEMY_FAMILIES: - transports = ["csv", "unload"] if family in RESULT_FILE_SQLALCHEMY else ["csv"] + for family, driver in SQLALCHEMY_FAMILIES.items(): + transports = ["csv", "unload"] if driver in RESULT_FILE_FAMILIES else ["csv"] cases.extend( Case(family, transport=transport, output="native", chunksize=chunk, shape=shape) for transport in transports @@ -139,11 +151,10 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: if c.api in {"thread", "aio"} and c.arraysize == 1000 and c.chunksize is None - and c.dtype_backend is None and c.transport == "csv" and ( (c.family in {"cursor", "dict", "s3fs"} and c.output == "rows") - or (c.family in {"pandas", "arrow", "polars"} and c.output == "native") + or (c.family in RESULT_FILE_FAMILIES and c.output == "native") ) for n in settings.concurrency ] @@ -154,11 +165,10 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: if c.api == "sync" and c.arraysize == 1000 and c.chunksize is None - and c.dtype_backend is None and c.family != "wrangler" and ( (c.family in {"cursor", "dict", "s3fs"} and c.output == "rows") - or (c.family in {"pandas", "arrow", "polars"} and c.output == "native") + or (c.family in RESULT_FILE_FAMILIES and c.output == "native") ) for api in ("direct", "to_thread") ] diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index 146715275..77b730429 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -131,8 +131,8 @@ def expand_jobs( """Split a selection into independent ``run`` invocations. A job covers one suite, scale, shape, family, API, transport, and output - kind, and additionally one arraysize for row output. A chunk size filter - is passed on to each job and named in its ID. Jobs that page through + kind, and additionally one arraysize for row output. With a chunk size + filter, each job runs and names the chunk sizes of its group's cases. Jobs that page through GetQueryResults are marked with the number of simultaneous paging queries so workers can bound the fleet's request rate. When one trial of such a job needs at least ``split_pages`` pages, each measured repetition becomes a @@ -176,9 +176,9 @@ def expand_jobs( base += f"-a{arraysize}" args += ["--arraysize", str(arraysize)] if filters.chunksize: - # The job runs the same chunk selection, which its ID names - # in one order so that equal selections share an ID. - chunks = [str(c) for c in sorted(set(filters.chunksize))] + # The job runs its group's chunk sizes, which its ID names in + # order, so equal selections share an ID. + chunks = [str(c) for c in sorted({c.chunksize or 0 for c in cases})] base += "-c" + "_".join(chunks) args += ["--chunksize", *chunks] runnable = [c for c in cases if c.unsupported is None] diff --git a/benchmarks/pyathena_bench/runner.py b/benchmarks/pyathena_bench/runner.py index e6686e66d..bfd774737 100644 --- a/benchmarks/pyathena_bench/runner.py +++ b/benchmarks/pyathena_bench/runner.py @@ -31,6 +31,7 @@ import psutil +from pyathena.util import parse_output_location from pyathena_bench.adapters import BACKEND_VERSIONS, Adapter, run_adapter from pyathena_bench.aws import ( cancel_queries, @@ -327,7 +328,7 @@ def output_files(s3: Any, location: str) -> dict[str, int]: Raises: ValueError: If no manifest is below the location. """ - bucket, _, prefix = location.removeprefix("s3://").partition("/") + bucket, prefix = parse_output_location(location) listed: dict[str, int] = {} for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix): listed.update((item["Key"], item["Size"]) for item in page.get("Contents", [])) @@ -338,7 +339,7 @@ def output_files(s3: Any, location: str) -> dict[str, int]: for manifest in manifests: body: str = s3.get_object(Bucket=bucket, Key=manifest)["Body"].read().decode() for line in filter(None, body.splitlines()): - file_bucket, _, key = line.removeprefix("s3://").partition("/") + file_bucket, key = parse_output_location(line) size = listed.get(key) if file_bucket == bucket else None if size is None: size = s3.head_object(Bucket=file_bucket, Key=key)["ContentLength"] @@ -383,8 +384,7 @@ def run( write_json(output / "environment.json", environment(settings)) write_json(output / "manifest.json", manifest) success = True - # Created on first use and shared by the run's output file summaries. - s3: Any = None + s3 = client(session(settings), "s3") for scale in scales: for case in cases: if case.unsupported: @@ -423,12 +423,11 @@ def run( result = supervise(payload, output / "events.jsonl", settings) if ( result["status"] == "ok" + and repetition >= 0 and case.transport in {"unload", "ctas"} and case.suite != "init" ): try: - if s3 is None: - s3 = client(session(settings), "s3") result["output_files"] = output_files(s3, payload["output"]) except Exception as exc: # The measurement stands; only the file summary is missing. diff --git a/benchmarks/tests/test_adapters.py b/benchmarks/tests/test_adapters.py index 0d41523f9..92fc59275 100644 --- a/benchmarks/tests/test_adapters.py +++ b/benchmarks/tests/test_adapters.py @@ -27,7 +27,13 @@ measure_query, run_adapter, ) -from pyathena_bench.cases import Case, matrix +from pyathena_bench.cases import ( + FAMILIES, + RESULT_FILE_FAMILIES, + SQLALCHEMY_FAMILIES, + Case, + matrix, +) from pyathena_bench.config import Settings from pyathena.arrow.cursor import ArrowCursor @@ -137,7 +143,7 @@ def unexpected_request(*args, **kwargs): ) settings = Settings(executor_workers=3) for case in matrix(settings, "single", "flat"): - if case.family == "wrangler" or case.family.startswith("sqla_") or case.unsupported: + if case.family not in FAMILIES or case.unsupported: continue adapter = Adapter(settings, case, session, "scratch", "s3://out/", "temp") with adapter.connection() as connection: @@ -165,7 +171,7 @@ def unexpected_request(*args, **kwargs): "sqla_arrow": ArrowCursor, "sqla_polars": PolarsCursor, } - cases = [c for c in matrix(settings, "single", "flat") if c.family.startswith("sqla_")] + cases = [c for c in matrix(settings, "single", "flat") if c.family in SQLALCHEMY_FAMILIES] for case in cases: engine = Adapter(settings, case, session, "scratch", "s3://out/", "temp").engine() try: @@ -179,9 +185,10 @@ def unexpected_request(*args, **kwargs): cursor = raw.cursor() assert type(cursor) is expected[case.family] assert cursor.arraysize == 1000 - if case.family in {"sqla_pandas", "sqla_arrow", "sqla_polars"}: + driver = SQLALCHEMY_FAMILIES[case.family] + if driver in RESULT_FILE_FAMILIES: assert cursor._unload is (case.transport == "unload") - if case.family in {"sqla_pandas", "sqla_polars"}: + if driver in {"pandas", "polars"}: assert cursor._chunksize == case.chunksize assert cursor._s3_max_workers == 3 cursor.close() diff --git a/benchmarks/tests/test_config.py b/benchmarks/tests/test_config.py index ae56c1c8e..95a215f34 100644 --- a/benchmarks/tests/test_config.py +++ b/benchmarks/tests/test_config.py @@ -10,7 +10,7 @@ import pytest from pyathena_bench.__main__ import main -from pyathena_bench.cases import matrix +from pyathena_bench.cases import SQLALCHEMY_FAMILIES, matrix from pyathena_bench.config import Settings, identifier, select_sql @@ -53,7 +53,7 @@ def test_matrix_preserves_unsupported_capabilities(): def test_sqlalchemy_cases_are_synchronous_single_query_reads(): - cases = [c for c in matrix(Settings(), "single", "flat") if c.family.startswith("sqla_")] + cases = [c for c in matrix(Settings(), "single", "flat") if c.family in SQLALCHEMY_FAMILIES] assert {(c.family, c.transport) for c in cases} == { ("sqla_rest", "csv"), ("sqla_s3fs", "csv"), @@ -62,7 +62,9 @@ def test_sqlalchemy_cases_are_synchronous_single_query_reads(): assert {c.chunksize for c in cases} == {None, 10000, 100000} assert {(c.api, c.output, c.unsupported) for c in cases} == {("sync", "native", None)} for suite in ("concurrent", "init"): - assert not [c for c in matrix(Settings(), suite, "flat") if c.family.startswith("sqla_")] + assert not [c for c in matrix(Settings(), suite, "flat") if c.family in SQLALCHEMY_FAMILIES] + # The pyarrow dtype variants belong to the single-query suite too. + assert not [c for c in matrix(Settings(), suite, "flat") if c.dtype_backend] def test_json_shape_covers_pandas_and_wrangler_reads(): diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index e66b1d3af..93fd583a5 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -217,6 +217,17 @@ def test_chunksize_filter_is_passed_to_the_job_and_named_in_its_id(self): ] assert [j["args"][-2:] for j in jobs] == [["--chunksize", "0"], ["--chunksize", "100000"]] assert merge_jobs([jobs[:1], jobs[1:]]) == jobs + # A group without chunked cases names only what it runs. + rows = expand_jobs( + SETTINGS, + ["single"], + ["small"], + ["flat"], + Filters(family=["s3fs"], api=["sync"], arraysize=[1000], chunksize=[0, 100000]), + ) + assert [(j["id"], j["args"][-2:]) for j in rows] == [ + ("single-small-flat-s3fs-sync-csv-rows-a1000-c0", ["--chunksize", "0"]) + ] # Equal selections in another order or with repeats get the same job. again = expand_jobs( SETTINGS, ["single"], ["small"], ["flat"], replace(selection, chunksize=[100000, 0, 0]) diff --git a/benchmarks/tests/test_runner.py b/benchmarks/tests/test_runner.py index 54393c7a1..09e015d79 100644 --- a/benchmarks/tests/test_runner.py +++ b/benchmarks/tests/test_runner.py @@ -42,6 +42,28 @@ def cache_writing_worker(pipe, payload): pipe.close() +RUN_MANIFEST = { + "run_id": "a" * 32, + "scales": {"small": {"table": "input", "rows": 10}}, + "fixtures": {}, +} + + +@pytest.fixture +def run_stubs(monkeypatch): + """Replace the AWS calls of ``run()`` and return the S3 client it creates.""" + s3 = object() + monkeypatch.setattr( + "pyathena_bench.runner.validate_manifest", + lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out"}, + ) + monkeypatch.setattr("pyathena_bench.runner.validate_inputs", lambda *args: None) + monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) + monkeypatch.setattr("pyathena_bench.runner.session", lambda settings: None) + monkeypatch.setattr("pyathena_bench.runner.client", lambda session_, service: s3) + return s3 + + def test_observer_keeps_first_completion_and_query_id(): events = [] @@ -107,27 +129,18 @@ def test_timeout_does_not_become_a_successful_timing(self, tmp_path): assert result["status"] == "timeout" @pytest.mark.parametrize("status", ["timeout", "memory_limit", "worker_exit", "error"]) - def test_failed_trial_stops_suite_before_any_more_queries(self, monkeypatch, tmp_path, status): + def test_failed_trial_stops_suite_before_any_more_queries( + self, monkeypatch, tmp_path, run_stubs, status + ): calls = [] - monkeypatch.setattr( - "pyathena_bench.runner.validate_manifest", - lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out"}, - ) - monkeypatch.setattr("pyathena_bench.runner.validate_inputs", lambda *args: None) - monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) def fail(payload, events, settings): calls.append(payload) return {"status": status} monkeypatch.setattr("pyathena_bench.runner.supervise", fail) - manifest = { - "run_id": "a" * 32, - "scales": {"small": {"table": "input", "rows": 10}}, - "fixtures": {}, - } assert not run( - Settings(), manifest, [Case("cursor"), Case("dict")], ["small"], tmp_path / "run" + Settings(), RUN_MANIFEST, [Case("cursor"), Case("dict")], ["small"], tmp_path / "run" ) assert len(calls) == 1 assert calls[0]["warmup"] @@ -211,44 +224,22 @@ def test_output_files_without_a_manifest_is_an_error(): ("case", "listed"), [(Case("pandas", transport="unload", output="native"), True), (Case("pandas"), False)], ) -def test_unload_trials_record_their_output_files(monkeypatch, tmp_path, case, listed): - monkeypatch.setattr( - "pyathena_bench.runner.validate_manifest", - lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out"}, - ) - monkeypatch.setattr("pyathena_bench.runner.validate_inputs", lambda *args: None) - monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) +def test_unload_trials_record_their_output_files(monkeypatch, tmp_path, run_stubs, case, listed): monkeypatch.setattr("pyathena_bench.runner.supervise", lambda *args: {"status": "ok"}) - locations = [] - clients = [] + calls = [] def summarize(s3, location): - locations.append(location) - clients.append(s3) + calls.append((s3, location)) return {"count": 1, "bytes": 5, "max_bytes": 5, "min_bytes": 5} - created = [] - - def make_client(session_, service): - assert service == "s3" - created.append(object()) - return created[-1] - monkeypatch.setattr("pyathena_bench.runner.output_files", summarize) - monkeypatch.setattr("pyathena_bench.runner.session", lambda settings: None) - monkeypatch.setattr("pyathena_bench.runner.client", make_client) - manifest = { - "run_id": "a" * 32, - "scales": {"small": {"table": "input", "rows": 10}}, - "fixtures": {}, - } - settings = replace(Settings(), warmups=0, repetitions=2) - assert run(settings, manifest, [case], ["small"], tmp_path / "run") + settings = replace(Settings(), warmups=1, repetitions=2) + assert run(settings, RUN_MANIFEST, [case], ["small"], tmp_path / "run") trials = [json.loads(line) for line in (tmp_path / "run/trials.jsonl").read_text().splitlines()] - assert all(("output_files" in trial) is listed for trial in trials) - assert [loc.startswith(f"s3://out/runs/{'a' * 32}/trials/") for loc in locations] == ( - [True, True] if listed else [] + # Measured trials only, with the run's one client. + assert [("output_files" in t) for t in trials] == [False, listed, listed] + assert len(calls) == (2 if listed else 0) + assert all( + s3 is run_stubs and location.startswith(f"s3://out/runs/{'a' * 32}/trials/") + for s3, location in calls ) - # One client serves every trial of the run. - assert len(created) == (1 if listed else 0) - assert clients == created * 2