Skip to content

feat(athena): add start_query_executions for parallel query execution #3190

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 7 commits into from
Closed

feat(athena): add start_query_executions for parallel query execution #3190

wants to merge 7 commits into from

Conversation

@ggiallo28
Copy link

@ggiallo28 ggiallo28 commented Aug 28, 2025 •
edited
Loading

Feature or Bugfix

  • Feature

Detail

  • Added start_query_executions to submit multiple Athena queries in one call.
  • Enabled parallel query submission and wait, significantly reducing end-to-end execution time.
  • Introduced configurable concurrency to adapt performance to available system resources.

Relates

  • Improves efficiency and responsiveness for workflows requiring multiple Athena queries.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

@ggiallo28
feat(athena): add start_query_executions for async multi-query execution
17cd753 
Introduce `wr.athena.start_query_executions` as a parallelized variant of
`start_query_execution`. It allows submitting multiple queries in one call,
with support for:

- Sequential or threaded submission (`use_threads`)
- Lazy or eager consumption of results (`as_iterator`)
- Per-query `client_request_token` (string or list)
- Optional workgroup checks (`check_workgroup`, `enforce_workgroup`)
- Full Athena cache integration

This improves performance when dispatching batches of queries by reducing
workgroup lookups and enabling concurrent execution.
@ggiallo28
feat(athena): improve start_query_executions with simplified tokens a…
b6e4d88 
…nd parallel wait

- Simplified client_request_token handling:
  - Removed manual padding/truncation.
  - Let Athena enforce length constraints.
  - Tokens generated as `<base_token>-<index>` or provided as list.
- Improved wait logic:
  - Added optional wait handling directly inside _submit.
  - Queries can now be waited in parallel with submission (reduced overhead).
- Configurable default threads:
  - Replaced hardcoded defaults with os.cpu_count().
  - Added support for AWSWRANGLER_THREADS_DEFAULT env var override.
@ggiallo28
chore: cleanup and CI adjustments
5e1c7f0 
- Removed unused `reduce` import from Athena module.
- Applied ruff formatting to `start_query_executions`.
- Fixed static check issues to pass CI.
- Added ruff check on Athena tests file.
kukushking
kukushking requested changes Sep 3, 2025
View reviewed changes
Copy link
Contributor

@kukushking kukushking left a comment

Choose a reason for hiding this comment

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

Hi, thank you for opening this PR!

What is the reason to build this into the SDK? Is this a common use case to submit x queries in parallel? Feels like it is specific to your concrete application logic

@ggiallo28
Copy link
Author

ggiallo28 commented Sep 3, 2025

Hi @kukushking, thanks for the review!

I see that this pattern shows up in data workflows where many short Athena queries must run together. A few common examples:

  • Dashboard refresh/precompute: populate multiple queries that feed BI tiles concurrently. Running them one by one slows down the process and forces repeated checks for each query even when I make small changes in the queries, like checking the workgroup every time, and so on.
  • Data quality checks: run the same validation across dozens of tables/prefixes in parallel.

API symmetry & ergonomics:
Wrangler already provides athena.get_query_executions for fetching many query details in parallel. start_query_executions is the natural counterpart for submitting many queries at once.

Parallel submission and coordinated wait help improve performance while respecting quotas via a configurable concurrency. The implementation is opt-in, non-breaking, and uses the same guardrails already present for single-query flows.

Why I contributed this:
I took inspiration from awswrangler.athena.get_query_executions(...), found myself repeatedly re-implementing the batch submit and wait pattern for data use cases, and decided to contribute a reusable, documented version back to the community.

Happy to adjust naming, docs, or move it behind a helper/recipe if you prefer, but I believe the symmetry and the prevalence of these data workflows justify having this in the SDK. In the future, it would be great to support retrieving results in the same parallel manner.

@jaidisido
Copy link
Contributor

jaidisido commented Sep 5, 2025

AWS CodeBuild CI Report

  • CodeBuild project: GitHubCodeBuild8756EF16-4rfo0GHQ0u9a
  • Commit ID: b8a607f
  • Result: SUCCEEDED
  • Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

@jaidisido
Copy link
Contributor

jaidisido commented Sep 5, 2025

AWS CodeBuild CI Report

  • CodeBuild project: GitHubCodeBuild8756EF16-4rfo0GHQ0u9a
  • Commit ID: 89449cb
  • Result: SUCCEEDED
  • Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

@jaidisido
Copy link
Contributor

jaidisido commented Sep 5, 2025

AWS CodeBuild CI Report

  • CodeBuild project: GitHubCodeBuild8756EF16-4rfo0GHQ0u9a
  • Commit ID: f20d694
  • Result: SUCCEEDED
  • Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

@jaidisido
Copy link
Contributor

jaidisido commented Sep 6, 2025

AWS CodeBuild CI Report

  • CodeBuild project: GitHubDistributedCodeBuild6-jWcl5DLmvupS
  • Commit ID: b8a607f
  • Result: FAILED
  • Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

@jaidisido
Copy link
Contributor

jaidisido commented Sep 6, 2025

AWS CodeBuild CI Report

  • CodeBuild project: GitHubDistributedCodeBuild6-jWcl5DLmvupS
  • Commit ID: 89449cb
  • Result: FAILED
  • Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

@jaidisido
Copy link
Contributor

jaidisido commented Sep 6, 2025

AWS CodeBuild CI Report

  • CodeBuild project: GitHubDistributedCodeBuild6-jWcl5DLmvupS
  • Commit ID: f20d694
  • Result: FAILED
  • Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

ggiallo28

This comment was marked as duplicate.

Sign in to view

@jaidisido
Copy link
Contributor

jaidisido commented Sep 9, 2025

AWS CodeBuild CI Report

  • CodeBuild project: GitHubCodeBuild8756EF16-4rfo0GHQ0u9a
  • Commit ID: 281f5c4
  • Result: SUCCEEDED
  • Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

@jaidisido
Copy link
Contributor

jaidisido commented Sep 9, 2025

AWS CodeBuild CI Report

  • CodeBuild project: GitHubDistributedCodeBuild6-jWcl5DLmvupS
  • Commit ID: 281f5c4
  • Result: FAILED
  • Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

@kukushking
Copy link
Contributor

kukushking commented Oct 21, 2025

Thanks for the comment @ggiallo28 .

  • Dashboard refresh/precompute: populate multiple queries that feed BI tiles concurrently. Running them one by one slows down the process and forces repeated checks for each query even when I make small changes in the queries, like checking the workgroup every time, and so on.
  • Data quality checks: run the same validation across dozens of tables/prefixes in parallel.

I agree running queries one by one slows down the process, however that does not explain why this should be handled in the SDK, and not your BI application logic. I am a bit cautious to add this to SDK as the same reasoning may be extended to any other call the library provides, adding overhead and negatively impacting maintainability.

API symmetry & ergonomics: Wrangler already provides athena.get_query_executions for fetching many query details in parallel. start_query_executions is the natural counterpart for submitting many queries at once.

athena.get_query_executions is a batch API, it uses BatchGetQueryExecution under the hood. We maintain API symmetry between SDK and AWS API.

@kukushking kukushking closed this Oct 21, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kukushking kukushking kukushking requested changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ggiallo28 @jaidisido @kukushking