Skip to content

Benchmark read_sql_query and trim the next benchmark selection - #1146

Merged
laughingman7743 merged 5 commits into
masterfrom
feat/1133-benchmark-selection
Oct 12, 2026
Merged

laughingman7743 merged 5 commits into
masterfrom
feat/1133-benchmark-selection

Conversation

@laughingman7743

@laughingman7743 laughingman7743 commented Oct 12, 2026 •

Copy link
Copy Markdown
Member

WHAT

Benchmark harness changes for the next run (#644), per #1133:

  • SQLAlchemy pandas.read_sql_query(): new families sqla_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 configured chunksize. 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, outside total_seconds. Its connections pass the same cursor options through cursor_kwargs, including the pandas/Polars cursor chunksize, so a chunked read also streams the result file. Rows are validated as for the other cases; sqla_rest counts as GetQueryResults paging for the API slots.
  • json workload (--shape json, single-query suite only): the flat projection plus CAST(details AS JSON) AS details_json (NULL where details is 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 hold json (Wrangler UNLOAD fails with NOT_SUPPORTED: Unsupported Hive type: json).
  • Pandas UNLOAD: chunked reads are measured again. Since Read chunked UNLOAD results in bounded memory #1135, PandasCursor reads UNLOAD results in chunks by default, so the stale "unsupported" mark is removed. Every pandas UNLOAD read also runs with dtype_backend="pyarrow" (new Case.dtype_backend, passed to execute(), and added to the case ID only when set).
  • UNLOAD/CTAS output files (Start chunked UNLOAD reads sooner and read large Parquet files with fewer bytes #1138): after each successful measured UNLOAD or CTAS trial, run reads 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 in output_files. summary.csv and summary.md show the most files and the largest file. A listing error is recorded without failing the trial and noted in the summary.
  • --chunksize filter for plan/run/jobs (0 selects unchunked reads). jobs passes it to the job and names it in the job ID.
  • README: the capability table, the json workload, and a new job-selection example with the trimmed matrix. It covers the synchronous API only, a medium regression subset, arraysize 100 only for Cursor, DictCursor at one scale, read_sql_query() through the REST driver without chunks only, json CSV reads, and no initialization or concurrency suite. The benchmark depends on PyAthena[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:

    • the matrix: SQLAlchemy families, the json workload, pandas UNLOAD chunk and backend variants;
    • the json projection;
    • engines for every SQLAlchemy case: connected without AWS requests, checking connection options and cursor type/options;
    • read_sql_query consumption of frames and chunks, and the connection opened during setup;
    • dtype_backend as an execute option;
    • the --chunksize filter, job args, and IDs;
    • sqla_rest paging;
    • output_files filtering on paginated listings, recording only for UNLOAD/CTAS trials;
    • the report columns.
  • Live smoke test against Athena: run_adapter with LIMIT 1000 queries over the benchmark source, using the repository test workgroup instead of a stack. Each case returned 1,000 rows and status ok:

    • sqla_rest flat, and nested with chunks;
    • sqla_s3fs;
    • sqla_pandas CSV, and nested UNLOAD with chunks;
    • sqla_arrow nested UNLOAD;
    • sqla_polars CSV with chunks, and nested UNLOAD;
    • pandas json without and with chunks;
    • pandas nested UNLOAD with chunks and dtype_backend="pyarrow";
    • Wrangler json.

    output_files matched the manifests to the listings: Wrangler UNLOAD 18 of 18 data files, Wrangler CTAS 15 of 15 (manifest under tables/).

  • Not run: any EC2 stack or fleet job; the selection and timings are untested at scale until the pilot.

🤖 Generated with Claude Code

laughingman7743 and others added 2 commits October 12, 2026 15:28
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>
Comment thread benchmarks/pyathena_bench/adapters.py Outdated
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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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): a read_sql_query() case opened its SQLAlchemy connection inside the measured query, so boto3 client creation and the dialect's first-connect setup counted in total_seconds, while cursor cases create connection and cursor in setup. The connection is now opened in run_adapter setup and closed before engine.dispose(); test_read_sql_query_connects_during_setup checks the order. Live smoke rerun: all 12 new case types ok, 1,000 rows.
  • Checked, no change: Connection.cursor() passes every cursor_kwargs entry, and unknown names raise (BREAKING: Reject unknown cursor keyword arguments #1102), so cursor_kwargs(family) builds only the options the driver's cursor accepts (engine test creates a cursor for every SQLAlchemy case without AWS requests); pandas read_sql_query() runs a string with exec_driver_sql(), so the SQL is not parsed for bind parameters; DataFrames have no close(), and the chunk generator closes in finally; dtype_backend cases are excluded from the concurrent and init filters; a --chunksize selection that leaves a group empty creates no job; job IDs stay unique across chunk selections (merge_jobs test).
  • Deferred: sqla_arrow with chunksize still loads the whole Arrow table (ArrowCursor has no chunked read); kept as a measurement of that behavior.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Round one, continued: relayed /simplify (reuse, simplification, efficiency, altitude) by the coordinating session on 78db697f, and repairs in 610edf3f

Applied, behavior otherwise unchanged:

  1. One SQLALCHEMY_FAMILIES mapping (family → PyAthena driver), RESULT_FILE_FAMILIES shared by the cursor matrix, the SQLAlchemy UNLOAD transports, and the concurrent/init filters, and ALL_FAMILIES for the CLI; Adapter.engine() reads the mapping (a driver's cursor takes the options of the cursor family of the same name; rest gets the plain Cursor options); tests use the constants.
  2. The dtype_backend="pyarrow" pandas UNLOAD variants are generated for the single-query suite only, in one loop over backends; the dtype_backend is None filters in concurrent/init are gone.
  3. measure_query(): one branch for Wrangler and read_sql_query(), differing only in the call that returns the DataFrame or chunk iterator.
  4. output_files() parses both locations with pyathena.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.
  5. CLI: Filters (and the trial overrides) built once for jobs/plan/run.
  6. jobs under --chunksize names and passes each group's actual chunk sizes (an S3FS rows group under 0 100000 becomes -c0).
  7. Tests: RUN_MANIFEST and a run_stubs fixture shared by the run() tests.
  8. Selection: sqla_rest without chunks only (sqla_s3fs covers the chunked fetchmany() 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 and plan with 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).

Comment thread benchmarks/pyathena_bench/runner.py Outdated
return result


def output_files(settings: Settings, location: str) -> dict[str, int]:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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() defaults dtype_backend to numpy_nullable and the harness passes engine as a cursor option, not a read option; the live smoke test read nested UNLOAD chunks with dtype_backend="pyarrow".
  • "Parquet cannot hold json": Wrangler UNLOAD of the projection failed with NOT_SUPPORTED: Unsupported Hive type: json (local query, test workgroup); CAST(details AS JSON) returns type json with NULLs where details is NULL.
  • Operator: one ListObjectsV2 pagination per successful UNLOAD/CTAS trial; the instance role already grants s3:ListBucket on the scratch bucket (cloudformation/benchmark.yaml); a listing error is recorded without failing the trial. sqla_rest takes 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, --family gain choices; old trial records without output_files report N/A.
  • README: the selection commands are the ones the estimate ran (jobs output: 96 jobs); just docs lint clean.

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":

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent review (relayed Codex result): FINDINGS, two items

  • Reviewer: Codex CLI 0.162.1 (codex exec --sandbox read-only --ephemeral), thread 01a12ed0-b59b-7921-bfb5-5f1af389a524; static review, no edits, builds, tests, network, or GitHub access. Base d3d4c4df55cf5e64e143aaaf16c794e702bf87c6, head 3f7fa39a5110d2eac36376839f44563d08209c96, 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).
  1. P2 (benchmarks/pyathena_bench/cases.py:122): the json filter leaves pandas cases for init, so run --suite init --shape json is accepted, but fixture preparation (aws.py:243) also prepares a Parquet UNLOAD fixture, which rejects the json column, aborting before measurement.
  2. P2 (benchmarks/pyathena_bench/report.py:125): a trial whose S3 listing failed keeps its successful timing with output_files_error, but the report drops it from the file statistics without noting the error, so maxima can come from fewer trials silently.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Repairs and independent follow-up

  1. Finding 1 → c350c739: matrix() returns no json cases for the init and concurrent suites, so run/plan --suite init --shape json reject the empty selection before fixture preparation, and jobs omits those combinations. test_json_shape_covers_pandas_and_wrangler_csv_reads asserts both empty matrices (fails on 3f7fa39a). README: the json workload covers the single-query suite.
  2. Finding 2 → c350c739: a successful trial with output_files_error keeps its timing, and the summary notes Output files not listed: <error>; the report test asserts the note and that the maxima come from the listed trial (fails on 3f7fa39a). README notes it.
  3. Self-review of the repair: behavior — mixed jobs selections keep their valid json/single combinations (README selection regenerated: still 96 jobs); claims — README wording re-read against the code. Validation on c350c739: just benchmark lint and just docs lint clean, just benchmark test 128 passed, 1 skipped.
  4. Codex follow-up on c350c739 (thread 01a12ed9-2f0c-7850-af5a-fec744377ff3, frozen snapshot unchanged): CLEAN. Covered all matrix() callers, fixture preparation, listing-error reporting in CSV/Markdown, README wording, and test strength.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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")

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round one, continued (relayed /code-review high by the coordinating session on c350c7398a3e22bdf1165bc18aa24ca82ea78c75): FINDINGS, eight items, all verified

  1. cases.py:125: the json shape dropped UNLOAD/CTAS cases instead of marking them unsupported, contrary to matrix()'s contract and the README ("Unsupported combinations are included in the plan and report with reasons").
  2. README.md:345: the xlarge bullet said read_sql_query() covers "the result-file drivers" flat and nested, but the flat job includes sqla_s3fs and the nested one does not; RESULT_FILE_SQLALCHEMY excludes sqla_s3fs.
  3. runner.py:331: output_files() counted every object not ending in .csv/.metadata/.txt/.json as Parquet; prefer the query manifest.
  4. report.py:75: the "Output files not listed" note included warmup trials while the columns use measured trials only.
  5. fleet.py:180: the chunk-size job-ID suffix depended on argument order and repeats (-c0_100000 vs -c100000_0), bypassing merge_jobs' duplicate guard.
  6. adapters.py:372: engine.connect() ran before the try/finally, so a connect failure skipped engine.dispose().
  7. 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.
  8. runner.py:327: output_files() created a boto3 session and S3 client per trial.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Repairs (78db697f) and independent follow-up

  1. 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/concurrent still have no json cases.
  2. README xlarge bullet: pandas, Arrow, Polars, and S3FS drivers flat; pandas, Arrow, and Polars nested with chunks only.
  3. output_files(s3, location) reads every <query_id>-manifest.csv below the trial prefix and sizes the listed files from the listing (HeadObject outside it); no manifest raises, which is recorded as output_files_error. Live check on the test bucket: Wrangler UNLOAD 18/18 files, Wrangler CTAS 15/15 (manifest under tables/), PyAthena UNLOAD manifests at the prefix root.
  4. The listing-failure note covers measured trials only.
  5. Chunk selections are sorted and deduplicated in job IDs and arguments.
  6. Setup (SQLAlchemy connect, cursor creation) runs inside the try; the finally disposes the engine and closes what was created.
  7. Adapter.execute_kwargs() returns cache_size=0, result_reuse_enable=False, and dtype_backend, used only by the cursor branches; Adapter.engine() documents that result reuse is off through the connection and the client cache through the default cache_size (ExecuteOptions.cache_size = 0; the dialect's do_execute() passes no cache_size). README says the same.
  8. 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>
@laughingman7743
laughingman7743 marked this pull request as ready for review October 12, 2026 08:34
@laughingman7743
laughingman7743 merged commit 7f5eddd into master Oct 12, 2026
15 checks passed
@laughingman7743
laughingman7743 deleted the feat/1133-benchmark-selection branch October 12, 2026 09:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant