diff --git a/benchmarks/README.md b/benchmarks/README.md index 60383d6b8..9d1995573 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 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 @@ -239,12 +240,18 @@ When the temporary directory is a tmpfs, as `/tmp` is on Amazon Linux 2023, file | --- | --- | | Cursor, DictCursor, S3FS | Row batches with arraysize 100 and 1000 | | Pandas CSV | Full DataFrame or streamed chunks | -| Pandas UNLOAD | Full Parquet read; requested chunked cases are marked unsupported as bounded-memory measurements | +| Pandas UNLOAD | Full DataFrame or chunks, each also with `dtype_backend="pyarrow"` | | Arrow CSV and UNLOAD | Eager table loading; arraysize changes row extraction, not S3 streaming | | Polars CSV and UNLOAD | Full result or chunk iterator | -| Wrangler CSV | Scalar projections; nested cases are marked unsupported | +| SQLAlchemy (`sqla_*`) | `pandas.read_sql_query()` through an engine of each PyAthena driver, whole or with `chunksize` | +| Wrangler CSV | Scalar projections and the json workload; nested cases are marked unsupported | | Wrangler CTAS and UNLOAD | Full DataFrame or chunks, including nested workloads | +The SQLAlchemy families `sqla_rest`, `sqla_pandas`, `sqla_arrow`, `sqla_polars`, and `sqla_s3fs` read with the synchronous API in the single-query suite only, and the result-file drivers also read UNLOAD results. +Their engines connect with the same workgroup, output location, polling interval, and disabled result reuse as the cursor cases, their statements run without the client cache (`cache_size` defaults to 0), and their cursors take the same cursor options; with `chunksize`, the pandas and Polars drivers' cursors also read the result file in chunks. +As a cursor is, the engine's connection is opened during setup, outside `total_seconds`. +`pandas.read_sql_query()` builds DataFrames from the rows of `fetchall()` or `fetchmany()`, so these cases include the driver's row output; `sqla_rest` pages through GetQueryResults like Cursor. + The [Wrangler API documentation](https://aws-sdk-pandas.readthedocs.io/en/stable/stubs/awswrangler.athena.read_sql_query.html) describes its CSV nested-type restriction. CTAS and UNLOAD have separate limitations, including timestamp-with-time-zone columns; the supplied timestamp is a bigint. Unsupported combinations are included in the plan and report with reasons. @@ -291,6 +298,7 @@ While an API-heavy job waits for slots, workers do not start later API-heavy job It orders API-heavy jobs first and longer jobs earlier. With `--split-pages`, a job whose single trial needs at least that many pages becomes one job per measured repetition without a warmup, so those repetitions run on different hosts. `run` accepts the matching `--arraysize`, `--warmups`, and `--repetitions` options. +`--chunksize` keeps only the given chunk sizes, with 0 for reads without chunks, in `plan`, `run`, and `jobs`; `jobs` passes it to each job and adds it to the job ID. `jobs` accepts `--warmups` and `--repetitions` too; with either option, every job that is not split carries both values as arguments, so one queue can combine jobs with different trial counts under one stored configuration. `queue --jobs` accepts several job files, rejects duplicate job IDs, and orders the combined jobs as `jobs` orders one selection. @@ -330,41 +338,53 @@ uv run --no-sync python -m pyathena_bench --config results/fleet.toml queue \ ``` Several `jobs` selections can make up one bounded matrix. -For example, the following selection, with `chunksizes = [10000, 100000, 1000000]`, `concurrency = [1, 10]`, and `timeout_seconds = 5400.0` in `results/fleet.toml`, measures: +For example, the following selection, with `chunksizes = [10000, 100000, 1000000]` and `timeout_seconds = 5400.0` in `results/fleet.toml`, measures the synchronous API only: -- medium in full, with five repetitions; -- large with arraysize 100 paging limited to Cursor with the synchronous API, concurrent row paging limited to flat Cursor jobs, and initialization limited to the S3FS and DataFrame families, with five repetitions, except three for those two Cursor selections; -- xlarge with the synchronous API only, without DictCursor and without arraysize 100 for the DataFrame families' row output, with three repetitions. +- medium as a small regression subset: Cursor and S3FS rows with arraysize 1000, and the native reads of the DataFrame families and Wrangler without chunks, flat, three repetitions; +- large, flat and nested: Cursor and S3FS rows with arraysize 1000 (five repetitions), Cursor with arraysize 100 (three, split into single-repetition jobs), DictCursor flat (three), every native read of the DataFrame families and Wrangler (five), the json workload's CSV reads without chunks and with 100,000-row chunks (five), and `read_sql_query()` through the pandas, Arrow, Polars, and S3FS drivers without chunks and with 100,000-row chunks (five) and through the REST driver without chunks (three, split); `sqla_s3fs` covers the chunked `fetchmany()` path without GetQueryResults paging; +- xlarge: S3FS rows flat, every native read of the DataFrame families and Wrangler, the json workload without chunks and with 100,000-row chunks, and `read_sql_query()` through the pandas, Arrow, Polars, and S3FS drivers flat without chunks and with chunks, and through the pandas, Arrow, and Polars drivers nested with chunks only, all with three repetitions. +It leaves out the ThreadPool and native asyncio APIs, whose single-query timings matched the synchronous API, the initialization suite, the concurrency suite, arraysize 100 for every family but Cursor, and DictCursor beyond one scale. Every selection runs without warmups because each trial already runs in a fresh process. ```bash jobs() { uv run --no-sync python -m pyathena_bench --config results/fleet.toml jobs --warmups 0 "$@"; } -jobs --suite single concurrent init --scale medium --shape flat nested --repetitions 5 > results/jobs-medium.json -jobs --suite single --scale large --shape flat nested --arraysize 1000 --repetitions 5 > results/jobs-large-1000.json -jobs --suite single --scale large --shape flat nested --arraysize 100 --family s3fs pandas arrow polars \ - --repetitions 5 > results/jobs-large-100.json -jobs --suite single --scale large --shape flat nested --arraysize 100 --family cursor --api sync \ - --repetitions 3 --split-pages 5000 > results/jobs-large-cursor-100.json -jobs --suite concurrent --scale large --shape flat nested --family s3fs pandas arrow polars \ - --repetitions 5 > results/jobs-large-concurrent.json -jobs --suite concurrent --scale large --shape flat --family cursor \ - --repetitions 3 --split-pages 5000 > results/jobs-large-concurrent-cursor.json -jobs --suite init --scale large --shape flat nested --family s3fs pandas arrow polars \ - --repetitions 5 > results/jobs-large-init.json -jobs --suite single --scale xlarge --shape flat nested --api sync --family cursor --arraysize 1000 \ - --repetitions 3 --split-pages 5000 > results/jobs-xlarge-cursor.json -jobs --suite single --scale xlarge --shape flat nested --api sync --family s3fs \ - --repetitions 3 > results/jobs-xlarge-s3fs.json -jobs --suite single --scale xlarge --shape flat nested --api sync --family pandas arrow polars wrangler \ - --arraysize 1000 --repetitions 3 > results/jobs-xlarge-dataframes.json +jobs --suite single --scale medium --shape flat --api sync --arraysize 1000 --chunksize 0 \ + --family cursor s3fs --output-kind rows --repetitions 3 > results/jobs-medium-rows.json +jobs --suite single --scale medium --shape flat --api sync --chunksize 0 \ + --family pandas arrow polars wrangler --output-kind native --repetitions 3 > results/jobs-medium-native.json +jobs --suite single --scale large --shape flat nested --api sync --arraysize 1000 \ + --family cursor s3fs --output-kind rows --repetitions 5 > results/jobs-large-rows.json +jobs --suite single --scale large --shape flat nested --api sync --arraysize 100 \ + --family cursor --repetitions 3 --split-pages 5000 > results/jobs-large-cursor-100.json +jobs --suite single --scale large --shape flat --api sync --arraysize 1000 \ + --family dict --repetitions 3 > results/jobs-large-dict.json +jobs --suite single --scale large --shape flat nested --api sync \ + --family pandas arrow polars wrangler --output-kind native --repetitions 5 > results/jobs-large-native.json +jobs --suite single --scale large --shape json --api sync --transport csv --chunksize 0 100000 \ + --family pandas wrangler --repetitions 5 > results/jobs-large-json.json +jobs --suite single --scale large --shape flat nested --chunksize 0 100000 \ + --family sqla_pandas sqla_arrow sqla_polars sqla_s3fs --repetitions 5 > results/jobs-large-sqla.json +jobs --suite single --scale large --shape flat nested --chunksize 0 \ + --family sqla_rest --repetitions 3 --split-pages 1000 > results/jobs-large-sqla-rest.json +jobs --suite single --scale xlarge --shape flat --api sync --arraysize 1000 \ + --family s3fs --output-kind rows --repetitions 3 > results/jobs-xlarge-s3fs.json +jobs --suite single --scale xlarge --shape flat nested --api sync \ + --family pandas arrow polars wrangler --output-kind native --repetitions 3 > results/jobs-xlarge-native.json +jobs --suite single --scale xlarge --shape json --api sync --transport csv --chunksize 0 100000 \ + --family pandas wrangler --repetitions 3 > results/jobs-xlarge-json.json +jobs --suite single --scale xlarge --shape flat --chunksize 0 100000 \ + --family sqla_pandas sqla_arrow sqla_polars sqla_s3fs --repetitions 3 > results/jobs-xlarge-sqla-flat.json +jobs --suite single --scale xlarge --shape nested --chunksize 100000 \ + --family sqla_pandas sqla_arrow sqla_polars --repetitions 3 > results/jobs-xlarge-sqla-nested.json uv run --no-sync python -m pyathena_bench --config results/fleet.toml queue \ - --stack "$BENCHMARK_STACK_ID" --name large-1 --jobs results/jobs-medium.json results/jobs-large-*.json \ + --stack "$BENCHMARK_STACK_ID" --name large-1 --jobs results/jobs-medium-*.json results/jobs-large-*.json \ results/jobs-xlarge-*.json --manifest results/input-bounded.json ``` The manifest for this selection must be prepared with `--scale medium large xlarge`. It is published as `large-1` in place of the previous example, so the worker and status commands below apply unchanged. +The selection has 96 jobs; from the job durations of the 2026-10-11 run, they take about 14 hours of job time on `r7g.2xlarge` hosts, and the longest, xlarge nested pandas CSV, about 75 minutes. `queue` reserves the name with a conditional write, stores the configuration, manifest, and jobs under `fleet//` in the scratch bucket, and writes the `queue.json` marker last. Workers refuse a queue without the marker. @@ -405,6 +425,7 @@ Every output directory contains `environment.json`, a copy of the input manifest Environment metadata includes Git SHA, dirty state, Python and dependency versions, lockfile hash, CPU count, RAM, and settings. SDK events preserve query IDs, SQL, output locations, Athena statistics, and failed operations without storing result rows. Successful measurements generate `summary.csv` and `summary.md`; failed trials and unsupported cases remain visible. +After each successful measured UNLOAD or CTAS trial, `run` reads the query manifests (`-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 40e1883c1..f4909acbe 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 ALL_FAMILIES, SHAPES, 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,23 @@ def parser() -> argparse.ArgumentParser: ) else: command.add_argument("--suite", choices=suites, required=True) - command.add_argument("--shape", choices=("flat", "nested"), default="flat") + command.add_argument("--shape", choices=SHAPES, default="flat") command.add_argument("--warmups", type=int, help="Override the configured warmups") command.add_argument("--repetitions", type=int, help="Override the configured repetitions") command.add_argument("--scale", nargs="+", required=True) - command.add_argument("--family", choices=(*FAMILIES, "wrangler"), nargs="+") + command.add_argument("--family", choices=ALL_FAMILIES, nargs="+") command.add_argument( "--api", choices=("sync", "thread", "aio", "direct", "to_thread"), nargs="+" ) command.add_argument("--transport", choices=("csv", "unload", "ctas"), nargs="+") command.add_argument("--output-kind", choices=("rows", "native")) command.add_argument("--arraysize", type=int, nargs="+") + command.add_argument( + "--chunksize", + type=int, + nargs="+", + help="Chunk sizes to keep; 0 selects unchunked reads", + ) if name == "run": command.add_argument("--manifest", type=Path, required=True) command.add_argument("--out", type=Path, required=True) @@ -186,13 +193,16 @@ def main(argv: list[str] | None = None) -> int: retire_host(queue, client(session_, "autoscaling"), ran) else: sys.stdout.write(json.dumps(status(queue), indent=2) + "\n") - elif args.command == "jobs": - filters = Filters(args.family, args.api, args.transport, args.output_kind, args.arraysize) + else: + filters = Filters( + args.family, args.api, args.transport, args.output_kind, args.arraysize, args.chunksize + ) overrides = { k: v for k in ("warmups", "repetitions") if (v := getattr(args, k)) is not None } if overrides: settings = replace(settings, **overrides) + if args.command == "jobs": jobs = expand_jobs( settings, args.suite, @@ -203,23 +213,13 @@ def main(argv: list[str] | None = None) -> int: explicit_trials=bool(overrides), ) sys.stdout.write(json.dumps(jobs, indent=1) + "\n") - else: - overrides = { - k: v for k in ("warmups", "repetitions") if (v := getattr(args, k)) is not None - } - if overrides: - settings = replace(settings, **overrides) - filters = Filters(args.family, args.api, args.transport, args.output_kind, args.arraysize) + elif args.command in {"plan", "run"}: cases = filters.select(matrix(settings, args.suite, args.shape)) if not cases: raise ValueError("No cases match the selected filters") if args.command == "plan": page_size = min( - ( - c.arraysize - for c in cases - if c.family in {"cursor", "dict"} and c.suite != "init" - ), + (c.arraysize for c in cases if c.family in API_FAMILIES and c.suite != "init"), default=0, ) sys.stdout.write( diff --git a/benchmarks/pyathena_bench/adapters.py b/benchmarks/pyathena_bench/adapters.py index e452722d0..b0b888ee7 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,104 @@ 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, 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. + """ 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 + 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 + dialect executes statements without ``cache_size``, which defaults to 0. + + Returns: + The engine. + """ + 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(driver), + }, + ) + + def execute_kwargs(self) -> dict[str, Any]: + """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: + The keyword arguments for the cursor's ``execute()``. + """ + 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] = { "boto3_session": self.session, @@ -247,25 +311,30 @@ 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) + if case.family == "wrangler" or case.family in SQLALCHEMY_FAMILIES: + # The "cursor" of a SQLAlchemy case is its connection, opened during setup. + 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 # The botocore observer captures IDs even for iterators. elif case.api == "thread": query_id, future = await asyncio.to_thread( - cursor.execute, sql, cache_size=0, result_reuse_enable=False + 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, cache_size=0, result_reuse_enable=False) + 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, cache_size=0, result_reuse_enable=False) + cursor.execute(sql, **adapter.execute_kwargs()) ready = time.perf_counter() query_id = cursor.query_id count = consume(cursor.result_set, case) @@ -305,20 +374,27 @@ 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() + sqlalchemy = case.family in SQLALCHEMY_FAMILIES + connection = None if case.family == "wrangler" or sqlalchemy else adapter.connection() + engine = adapter.engine() if sqlalchemy else None cursors: list[Any] = [] - 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() + 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") @@ -365,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) @@ -374,3 +451,5 @@ async def run_adapter( cursor.close() if connection is not None: connection.close() + if engine is not None: + engine.dispose() diff --git a/benchmarks/pyathena_bench/cases.py b/benchmarks/pyathena_bench/cases.py index 3cd6125d1..386e1b5e8 100644 --- a/benchmarks/pyathena_bench/cases.py +++ b/benchmarks/pyathena_bench/cases.py @@ -14,6 +14,20 @@ from pyathena_bench.config import Settings FAMILIES = ("cursor", "dict", "s3fs", "pandas", "arrow", "polars") +# 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") @dataclass(frozen=True) @@ -28,45 +42,71 @@ 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 + 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. + 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") + if shape == "json" and suite != "single": + return [] cases: list[Case] = [] arraysizes = settings.arraysizes if suite == "single" else [1000] 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: - 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, - ) + # 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, ) + for chunk in chunks + for backend in backends + ) cases.extend( Case(family, api, "unload", arraysize=a, shape=shape) for a in arraysizes ) @@ -87,6 +127,22 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: unsupported=reason, ) ) + if suite == "single": + 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 + for chunk in (None, *settings.chunksizes) + ) + 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.output == "native" + ] if suite == "concurrent": # Keep the scaling suite bounded; chunk and arraysize sweeps belong to single. cases = [ @@ -98,7 +154,7 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: 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 ] @@ -112,7 +168,7 @@ def matrix(settings: Settings, suite: str, shape: str) -> list[Case]: 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/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..77b730429 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. 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 @@ -170,6 +175,12 @@ def expand_jobs( if arraysize is not None: base += f"-a{arraysize}" args += ["--arraysize", str(arraysize)] + if 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] 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..255b2ed10 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 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 prefix = "Warmup: " if trial.get("warmup") else "" failures = [q for q in trial.get("queries", []) if q["status"] != "ok"] @@ -122,6 +125,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 +180,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 +204,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..bfd774737 100644 --- a/benchmarks/pyathena_bench/runner.py +++ b/benchmarks/pyathena_bench/runner.py @@ -31,8 +31,15 @@ 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, 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 +309,49 @@ def supervise( return result +def output_files(s3: Any, location: str) -> dict[str, int]: + """Summarize the data files that a trial's UNLOAD or CTAS queries wrote. + + 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: + 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 = 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", [])) + 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 = 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"] + sizes.append(size) + 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( @@ -334,6 +384,7 @@ def run( write_json(output / "environment.json", environment(settings)) write_json(output / "manifest.json", manifest) success = True + s3 = client(session(settings), "s3") for scale in scales: for case in cases: if case.unsupported: @@ -370,6 +421,17 @@ def run( } append_json(output / "events.jsonl", {"event": "trial_start", **payload}) 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: + 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}" 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..92fc59275 100644 --- a/benchmarks/tests/test_adapters.py +++ b/benchmarks/tests/test_adapters.py @@ -27,9 +27,21 @@ 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 +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 +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.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: @@ -143,6 +155,131 @@ 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 in SQLALCHEMY_FAMILIES] + 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 + driver = SQLALCHEMY_FAMILIES[case.family] + if driver in RESULT_FILE_FAMILIES: + assert cursor._unload is (case.transport == "unload") + if driver in {"pandas", "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)) + + 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, "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"] + + 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(), + Case("pandas", transport="unload", output="native", dtype_backend="pyarrow"), + None, + "scratch", + "s3://out/", + "temp", + ) + 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): session = boto3.Session( diff --git a/benchmarks/tests/test_config.py b/benchmarks/tests/test_config.py index fbcd01328..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 @@ -43,14 +43,47 @@ 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 in SQLALCHEMY_FAMILIES] + 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 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(): + 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"), + } + # 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") == [] + + 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 +101,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..93fd583a5 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,68 @@ 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 + # 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]) + ) + 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)] == [ + 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..4a64b0ea1 100644 --- a/benchmarks/tests/test_report.py +++ b/benchmarks/tests/test_report.py @@ -109,3 +109,31 @@ 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 + 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" + # 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 07edbf1b9..09e015d79 100644 --- a/benchmarks/tests/test_runner.py +++ b/benchmarks/tests/test_runner.py @@ -9,12 +9,13 @@ import os import time from dataclasses import replace +from io import BytesIO from pathlib import Path 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): @@ -41,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 = [] @@ -106,29 +129,117 @@ 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"] result = json.loads((tmp_path / "run/trials.jsonl").read_text()) assert result["status"] == status + + +class FakeS3: + """An S3 client with a fixed paginated listing and object bodies.""" + + 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): + assert name == "list_objects_v2" + return self + + def paginate(self, **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 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": 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}, + ] + }, + {}, + ], + 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}, + ) + assert output_files(s3, f"s3://out/{root}") == { + "count": 3, + "bytes": 450, + "max_bytes": 300, + "min_bytes": 50, + } + 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( + ("case", "listed"), + [(Case("pandas", transport="unload", output="native"), True), (Case("pandas"), False)], +) +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"}) + calls = [] + + def summarize(s3, location): + calls.append((s3, location)) + return {"count": 1, "bytes": 5, "max_bytes": 5, "min_bytes": 5} + + monkeypatch.setattr("pyathena_bench.runner.output_files", summarize) + 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()] + # 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 + ) 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]