Repository navigation
Benchmark read_sql_query and trim the next benchmark selection - #1146
Conversation
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
| 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. |
There was a problem hiding this comment.
Self-review round one (implementation behavior): FINDINGS, repaired
Base d3d4c4df55cf5e64e143aaaf16c794e702bf87c6, head 3f7fa39a5110d2eac36376839f44563d08209c96. Covered: cases.matrix (SQLAlchemy families, json filter, pandas UNLOAD chunk/backend variants, concurrent/init filters), config.select_sql, Adapter (connection_kwargs, cursor_kwargs, engine, execute_kwargs), measure_query/run_adapter, runner.output_files and its call, report, fleet.Filters/expand_jobs/API_FAMILIES, CLI choices, README, lock. Done by hand from the /code-review and /simplify perspectives (this session could not launch their subagents).
- Finding (fixed in
3f7fa39a): aread_sql_query()case opened its SQLAlchemy connection inside the measured query, so boto3 client creation and the dialect's first-connect setup counted intotal_seconds, while cursor cases create connection and cursor in setup. The connection is now opened inrun_adaptersetup and closed beforeengine.dispose();test_read_sql_query_connects_during_setupchecks the order. Live smoke rerun: all 12 new case typesok, 1,000 rows. - Checked, no change:
Connection.cursor()passes everycursor_kwargsentry, and unknown names raise (BREAKING: Reject unknown cursor keyword arguments #1102), socursor_kwargs(family)builds only the options the driver's cursor accepts (engine test creates a cursor for every SQLAlchemy case without AWS requests); pandasread_sql_query()runs a string withexec_driver_sql(), so the SQL is not parsed for bind parameters; DataFrames have noclose(), and the chunk generator closes infinally;dtype_backendcases are excluded from the concurrent and init filters; a--chunksizeselection that leaves a group empty creates no job; job IDs stay unique across chunk selections (merge_jobstest). - Deferred:
sqla_arrowwithchunksizestill loads the whole Arrow table (ArrowCursor has no chunked read); kept as a measurement of that behavior.
There was a problem hiding this comment.
Round one, continued: relayed /simplify (reuse, simplification, efficiency, altitude) by the coordinating session on 78db697f, and repairs in 610edf3f
Applied, behavior otherwise unchanged:
- One
SQLALCHEMY_FAMILIESmapping (family → PyAthena driver),RESULT_FILE_FAMILIESshared by the cursor matrix, the SQLAlchemy UNLOAD transports, and the concurrent/init filters, andALL_FAMILIESfor the CLI;Adapter.engine()reads the mapping (a driver's cursor takes the options of the cursor family of the same name;restgets the plain Cursor options); tests use the constants. - The
dtype_backend="pyarrow"pandas UNLOAD variants are generated for the single-query suite only, in one loop over backends; thedtype_backend is Nonefilters in concurrent/init are gone. measure_query(): one branch for Wrangler andread_sql_query(), differing only in the call that returns the DataFrame or chunk iterator.output_files()parses both locations withpyathena.util.parse_output_location(); warmup trials are not summarized;run()creates the S3 client once before the trials; the test checks the measured trials and the client identity.- CLI:
Filters(and the trial overrides) built once forjobs/plan/run. jobsunder--chunksizenames and passes each group's actual chunk sizes (an S3FS rows group under0 100000becomes-c0).- Tests:
RUN_MANIFESTand arun_stubsfixture shared by therun()tests. - Selection:
sqla_restwithout chunks only (sqla_s3fscovers the chunkedfetchmany()path without GetQueryResults paging); the large json workload on its own line with--transport csv --chunksize 0 100000(the unsupported json UNLOAD/CTAS rows stay in the matrix andplanwith their reason, without jobs). Regenerated: 96 jobs, 13.7 h of job time (×44.1/37.2 ≈ 16.3 host-hours, about $2.6–2.8 Spot); longest job still about 75 min. The README commands are identical to the ones the estimate ran.
Skipped, as agreed: keeping bytes/min_bytes in output_files (raw evidence); moving run_adapter setup into an ExitStack method (class restructuring needs maintainer approval).
Self-review of the repair — behavior: the CLI dispatch reaches the filters only for jobs/plan/run; the client is created per run() (which already calls AWS through validate_manifest); chunk-size job IDs follow the group's cases, so equal selections still share IDs; the engine smoke (12 case types against Athena, LIMIT 1000) is still ok with 1,000 rows each. Claims: README selection bullets and counts re-read against the regenerated jobs. Validation on 610edf3f: just benchmark lint, just lint, just docs lint clean; just benchmark test 130 passed, 1 skipped; the warmup-skip and per-group chunk-ID assertions fail on 78db697f's source (the other changes are refactors that keep behavior).
| return result | ||
|
|
||
|
|
||
| def output_files(settings: Settings, location: str) -> dict[str, int]: |
There was a problem hiding this comment.
Self-review round two (claims, callers, operations): CLEAN
Base d3d4c4df55cf5e64e143aaaf16c794e702bf87c6, head 3f7fa39a5110d2eac36376839f44563d08209c96.
- "PandasCursor has read UNLOAD results in chunks since Read chunked UNLOAD results in bounded memory #1135":
_reads_parquet_in_chunks()defaultsdtype_backendtonumpy_nullableand the harness passesengineas a cursor option, not a read option; the live smoke test read nested UNLOAD chunks withdtype_backend="pyarrow". - "Parquet cannot hold
json": Wrangler UNLOAD of the projection failed withNOT_SUPPORTED: Unsupported Hive type: json(local query, test workgroup);CAST(details AS JSON)returns typejsonwith NULLs wheredetailsis NULL. - Operator: one
ListObjectsV2pagination per successful UNLOAD/CTAS trial; the instance role already grantss3:ListBucketon the scratch bucket (cloudformation/benchmark.yaml); a listing error is recorded without failing the trial.sqla_resttakes an API slot like Cursor. - Estimate (96 jobs, 14.3 job-hours, about 17 host-hours, about $3 Spot, longest about 75 min): computed from the 444 done markers of the 2026-10-11 queue with new cases mapped to measured equivalents (factors in the PR body); it is an estimate, not a measurement of the new cases.
- Existing callers: case IDs are unchanged for
dtype_backend=None;--shape,--familygain choices; old trial records withoutoutput_filesreportN/A. - README: the selection commands are the ones the estimate ran (
jobsoutput: 96 jobs);just docs lintclean.
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 <noreply@anthropic.com>
| for transport in transports | ||
| for chunk in (None, *settings.chunksizes) | ||
| ) | ||
| if shape == "json": |
There was a problem hiding this comment.
Independent review (relayed Codex result): FINDINGS, two items
- Reviewer: Codex CLI 0.162.1 (
codex exec --sandbox read-only --ephemeral), thread01a12ed0-b59b-7921-bfb5-5f1af389a524; static review, no edits, builds, tests, network, or GitHub access. Based3d4c4df55cf5e64e143aaaf16c794e702bf87c6, head3f7fa39a5110d2eac36376839f44563d08209c96, frozen detached worktree (unchanged afterwards); prompt without the PR number, description, commit messages, or prior findings. - Covered: all 15 changed files; SQLAlchemy execution and driver options; timing, caching, result reuse, and row validation; json/UNLOAD/dtype cases; chunk filters, job IDs, and API slots; file statistics, reports, tests, and README (the README selection expands to 96 jobs).
- P2 (
benchmarks/pyathena_bench/cases.py:122): the json filter leaves pandas cases forinit, sorun --suite init --shape jsonis accepted, but fixture preparation (aws.py:243) also prepares a Parquet UNLOAD fixture, which rejects thejsoncolumn, aborting before measurement. - P2 (
benchmarks/pyathena_bench/report.py:125): a trial whose S3 listing failed keeps its successful timing withoutput_files_error, but the report drops it from the file statistics without noting the error, so maxima can come from fewer trials silently.
There was a problem hiding this comment.
Repairs and independent follow-up
- Finding 1 →
c350c739:matrix()returns no json cases for theinitandconcurrentsuites, sorun/plan --suite init --shape jsonreject the empty selection before fixture preparation, andjobsomits those combinations.test_json_shape_covers_pandas_and_wrangler_csv_readsasserts both empty matrices (fails on3f7fa39a). README: the json workload covers the single-query suite. - Finding 2 →
c350c739: a successful trial withoutput_files_errorkeeps its timing, and the summary notesOutput files not listed: <error>; the report test asserts the note and that the maxima come from the listed trial (fails on3f7fa39a). README notes it. - Self-review of the repair: behavior — mixed
jobsselections keep their valid json/single combinations (README selection regenerated: still 96 jobs); claims — README wording re-read against the code. Validation onc350c739:just benchmark lintandjust docs lintclean,just benchmark test128 passed, 1 skipped. - Codex follow-up on
c350c739(thread01a12ed9-2f0c-7850-af5a-fec744377ff3, frozen snapshot unchanged): CLEAN. Covered allmatrix()callers, fixture preparation, listing-error reporting in CSV/Markdown, README wording, and test strength.
There was a problem hiding this comment.
Codex follow-up on the /simplify repairs (610edf3f): CLEAN
Reviewer: Codex CLI 0.162.1 (codex exec --sandbox read-only --ephemeral), thread 01a12f14-8251-7900-a266-55ed67620d8e, frozen detached snapshot at 610edf3f6bfa474f537c77b62f55d44aabb88974 (unchanged afterwards); patch 78db697f..610edf3f plus the full diff from d3d4c4df. Covered all ten CLI subcommands (filters, overrides, dispatch), all nine suite/shape combinations (contents, order, case IDs, dtype variants, unsupported cases), the five SQLAlchemy drivers (engine/cursor options, merged measurement, cleanup), run() (eager client, manifest parsing, warmup/init exclusions, summary errors, fail-fast), job grouping with effective chunk IDs and duplicate rejection, tests, and docs; the README selection totals 96 jobs (11 medium + 51 large + 34 xlarge). No unintended behavioral regression.
Noted limitation, accepted: the shared run_stubs client stub returns the same object on every call, so the tests no longer count client creations; the single creation before the trial loop is confirmed statically.
- 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 <noreply@anthropic.com>
| cases = [ | ||
| c | ||
| if c.transport == "csv" | ||
| else replace(c, unsupported="Athena cannot write the json type to Parquet") |
There was a problem hiding this comment.
Self-review round one, continued (relayed /code-review high by the coordinating session on c350c7398a3e22bdf1165bc18aa24ca82ea78c75): FINDINGS, eight items, all verified
cases.py:125: the json shape dropped UNLOAD/CTAS cases instead of marking them unsupported, contrary tomatrix()'s contract and the README ("Unsupported combinations are included in the plan and report with reasons").README.md:345: the xlarge bullet saidread_sql_query()covers "the result-file drivers" flat and nested, but the flat job includessqla_s3fsand the nested one does not;RESULT_FILE_SQLALCHEMYexcludessqla_s3fs.runner.py:331:output_files()counted every object not ending in.csv/.metadata/.txt/.jsonas Parquet; prefer the query manifest.report.py:75: the "Output files not listed" note included warmup trials while the columns use measured trials only.fleet.py:180: the chunk-size job-ID suffix depended on argument order and repeats (-c0_100000vs-c100000_0), bypassingmerge_jobs' duplicate guard.adapters.py:372:engine.connect()ran before thetry/finally, so a connect failure skippedengine.dispose().adapters.py:305: the execute options dict was built for the Wrangler/SQLAlchemy branches that never use it; the SQLAlchemy path's result reuse and client cache settings were not explicit.runner.py:327:output_files()created a boto3 session and S3 client per trial.
There was a problem hiding this comment.
Repairs (78db697f) and independent follow-up
- json UNLOAD/CTAS reads stay in the matrix, marked "Athena cannot write the json type to Parquet" (Wrangler UNLOAD of the projection fails with
NOT_SUPPORTED: Unsupported Hive type: json);init/concurrentstill have no json cases. - README xlarge bullet: pandas, Arrow, Polars, and S3FS drivers flat; pandas, Arrow, and Polars nested with chunks only.
output_files(s3, location)reads every<query_id>-manifest.csvbelow the trial prefix and sizes the listed files from the listing (HeadObjectoutside it); no manifest raises, which is recorded asoutput_files_error. Live check on the test bucket: Wrangler UNLOAD 18/18 files, Wrangler CTAS 15/15 (manifest undertables/), PyAthena UNLOAD manifests at the prefix root.- The listing-failure note covers measured trials only.
- Chunk selections are sorted and deduplicated in job IDs and arguments.
- Setup (SQLAlchemy connect, cursor creation) runs inside the
try; thefinallydisposes the engine and closes what was created. Adapter.execute_kwargs()returnscache_size=0,result_reuse_enable=False, anddtype_backend, used only by the cursor branches;Adapter.engine()documents that result reuse is off through the connection and the client cache through the defaultcache_size(ExecuteOptions.cache_size = 0; the dialect'sdo_execute()passes nocache_size). README says the same.- One S3 client per
run(), created on first use.
README selection: xlarge json reads CSV only; 99 jobs (three only record the unsupported json UNLOAD/CTAS cases), estimate unchanged at 14.3 h of job time.
Self-review of the repair — behavior: setup failures now also close a partly created cursor list and the PyAthena connection; manifests of several queries in one trial are summed; a deleted listed file surfaces as an error. Claims: README, docstrings, and the commit message re-checked against the code and the live listings. Validation on 78db697f: just benchmark lint, just lint, just docs lint clean; just benchmark test 130 passed, 1 skipped; 8 new/changed tests fail with the source reverted to c350c739.
Codex follow-up on 78db697f (thread 01a12f00-1fcf-7bb3-a96f-aabdd0bad969, frozen snapshot unchanged): CLEAN. Covered json capability reporting, the 99-job README selection, manifest discovery/sizing/fallbacks and the shared client, measured-only notes, canonical chunk IDs, run_adapter setup/finally, cursor-only execute options, cache/reuse documentation, and test strength.
- 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 <noreply@anthropic.com>
WHAT
Benchmark harness changes for the next run (#644), per #1133:
pandas.read_sql_query(): new familiessqla_rest,sqla_pandas,sqla_arrow,sqla_polars,sqla_s3fs(single-query suite, synchronous API, flat and nested). The result-file drivers also read UNLOAD results. Each case reads without chunks and with every configuredchunksize. The engine (NullPool) connects with the same arguments as the cursor cases: workgroup, output location, polling interval, result reuse disabled, session, and botocore config. As with a cursor, its connection is opened during trial setup, outsidetotal_seconds. Its connections pass the same cursor options throughcursor_kwargs, including the pandas/Polars cursorchunksize, so a chunked read also streams the result file. Rows are validated as for the other cases;sqla_restcounts as GetQueryResults paging for the API slots.--shape json, single-query suite only): the flat projection plusCAST(details AS JSON) AS details_json(NULL wheredetailsis NULL), for native pandas and Wrangler reads (Parse CSV results with json columns without forcing low_memory=False #1140); their UNLOAD and CTAS reads are marked unsupported because Parquet cannot holdjson(Wrangler UNLOAD fails withNOT_SUPPORTED: Unsupported Hive type: json).dtype_backend="pyarrow"(newCase.dtype_backend, passed toexecute(), and added to the case ID only when set).runreads the query manifests (<query_id>-manifest.csv) below the trial's output prefix, which Athena writes for PyAthena and Wrangler UNLOAD and for CTAS, and records the number, total, largest, and smallest sizes of the data files they list inoutput_files.summary.csvandsummary.mdshow the most files and the largest file. A listing error is recorded without failing the trial and noted in the summary.--chunksizefilter forplan/run/jobs(0 selects unchunked reads).jobspasses it to the job and names it in the job ID.read_sql_query()through the REST driver without chunks only, json CSV reads, and no initialization or concurrency suite. The benchmark depends onPyAthena[sqlalchemy](lock: extras only).WHY
Refs #1133.
#254 reports that
read_sql_query()is much slower than AWS SDK for pandas, and the 4.0.0 run did not measure that path.The next run should confirm that PyAthena is not clearly slower than Wrangler for the same output and that chunking bounds memory. It should also measure the follow-ups (#1137, #1138, #1140) and drop cases that carried no information.
Selection estimate: 96 jobs (11 medium, 51 large, 34 xlarge), about 13.7 hours of job time. The estimate scales the per-case trial durations of the 2026-10-11 run, mapping new cases to a measured equivalent:
sqla_*from the driver family's row output ×1.2, json from flat ×1.3, pandas UNLOAD variants from the whole read. With that run's ratio of host time to job time (44.1/37.2), that is about 16 host-hours, roughly $2.6–2.8 at r7g.2xlarge Spot prices. The longest job, xlarge nested pandas CSV, takes about 75 minutes.TEST
Tested commit: 610edf3.
just benchmark lint,just lint,just docs lint: passed.just benchmark test: 130 passed, 1 skipped. New tests cover:read_sql_queryconsumption of frames and chunks, and the connection opened during setup;dtype_backendas an execute option;--chunksizefilter, job args, and IDs;sqla_restpaging;output_filesfiltering on paginated listings, recording only for UNLOAD/CTAS trials;Live smoke test against Athena:
run_adapterwithLIMIT 1000queries over the benchmark source, using the repository test workgroup instead of a stack. Each case returned 1,000 rows and statusok:sqla_restflat, and nested with chunks;sqla_s3fs;sqla_pandasCSV, and nested UNLOAD with chunks;sqla_arrownested UNLOAD;sqla_polarsCSV with chunks, and nested UNLOAD;dtype_backend="pyarrow";output_filesmatched the manifests to the listings: Wrangler UNLOAD 18 of 18 data files, Wrangler CTAS 15 of 15 (manifest undertables/).Not run: any EC2 stack or fleet job; the selection and timings are untested at scale until the pilot.
🤖 Generated with Claude Code