Skip to content

Replace the cross-thread cancel example in the Spark documentation #1270

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

Merged
laughingman7743 merged 1 commit into from
Oct 18, 2026
Merged

Replace the cross-thread cancel example in the Spark documentation #1270

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 25 additions & 14 deletions docs/spark.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,22 +277,13 @@ Athena cancels the calculation on a best-effort basis:
- Canceling a calculation that has already finished does not raise an error or change its state.

The `execute()` method raises `OperationalError` when the calculation ends in the `CANCELED` state.
The following example cancels the calculation from another thread after 60 seconds:

```python
import threading
from pyathena import connect
from pyathena.spark.cursor import SparkCursor
As `pyathena.threadsafety = 2` declares, calling `cancel()` from another thread while `execute()` runs on the same cursor is not supported.

@laughingman7743 laughingman7743 Oct 18, 2026

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 1 (implementation behavior): CLEAN

Base e086eac7bb2120e434c08c2eababccd9d0167fb0, head 32dbdd45bf46285c1fc5c77d3a55bc02795e9c19. Scope: docs/spark.md (Cancellation and AioSparkCursor sections).

  • The removed Timer example was the only cross-thread cancel() usage in the docs (grep for Timer and "another thread" across docs/ and pyathena/spark, pyathena/aio/spark). docs/sqlalchemy.md cancels from another thread through a new cursor on the shared connection, which threadsafety = 2 allows.
  • Each listed alternative is checked against the code: SparkCursor._calculate/_poll (KeyboardInterrupt), AsyncSparkCursor.execute() returning (calculation_id, future) and cancel(query_id) (pyathena/spark/async_cursor.py:205-253), and AioSparkCursor._calculate/_poll (pyathena/aio/spark/cursor.py:98-257), which cancel the calculation on task cancellation with kill_on_interrupt.
  • The existing AioSparkCursor create_task + await cursor.cancel() example runs on one event loop thread, so it is not affected and is kept.
  • The tests that call SparkCursor.cancel() from another thread (tests/pyathena/spark/test_spark_cursor.py:162) are kept, as decided in Break result sets after failed or cancelled fetches and drop the cross-thread guards #1238.

A long calculation can be stopped in the following ways:

conn = connect(work_group="YOUR_SPARK_WORKGROUP", cursor_class=SparkCursor)
with conn.cursor() as cursor:
timer = threading.Timer(60, cursor.cancel)
timer.start()
try:
cursor.execute("""print(spark.read.parquet("s3://YOUR_S3_BUCKET/large_dataset/").count())""")
finally:
timer.cancel()
```
- A `KeyboardInterrupt` while `execute()` runs, with `kill_on_interrupt` enabled, as described below.
- {ref}`async-spark-cursor`, whose `execute()` returns the calculation ID with a future, and whose `cancel()` takes that ID.
- {ref}`aio-spark-cursor` with `asyncio.wait_for()`, whose timeout cancels the calculation.

With `kill_on_interrupt` enabled, which is the default, a `KeyboardInterrupt` while `execute()` waits for the calculation
requests cancellation, waits until the calculation reaches a terminal state, and then propagates.
Expand Down Expand Up @@ -535,3 +526,23 @@ Cancelling the task while `execute()` is still starting the calculation first wa
and then cancels the calculation it started in the same way.
If the task is cancelled before `execute()` begins the request, the request is never sent.
Cancelling the task again during this wait raises `asyncio.CancelledError` at once without cancelling the calculation.

With `kill_on_interrupt` enabled, a timeout from `asyncio.wait_for()` while `execute()` starts or waits for the calculation therefore cancels it and raises `asyncio.TimeoutError`.

@laughingman7743 laughingman7743 Oct 18, 2026

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 2 (claims, callers, documentation reader): CLEAN

Base e086eac7bb2120e434c08c2eababccd9d0167fb0, head 32dbdd45bf46285c1fc5c77d3a55bc02795e9c19.

Claims checked:

  • "a timeout from asyncio.wait_for() ... cancels it and raises asyncio.TimeoutError": verified offline with an AioSparkCursor whose first status request blocks (mocked _calculate, _poll_until_terminal, _cancel). On Python 3.11 (where wait_for runs the coroutine in a separate task) and 3.14 (same task via asyncio.timeout), the result was TimeoutError, _cancel awaited once with the calculation ID, calculation_id == "calculation_id", and state == "CANCELED". The example's cursor.calculation_id and cursor.state therefore print the cancelled calculation. requires-python = ">=3.11".
  • The start-phase case (timeout while StartCalculationExecution runs) follows the existing paragraph above, so it is not repeated.
  • {ref} targets async-spark-cursor and aio-spark-cursor resolve in a worktree sphinx-build (spark.html).
  • The example mirrors the SQL one in docs/aio.md (Task cancellation) and uses the same aio_connect / asyncio.to_thread(conn.cursor) pattern as the other AioSparkCursor examples.

Limits: no live Spark run, since this is a documentation-only change. The live task-cancellation path is covered by tests/pyathena/aio/spark/test_cursor.py::test_cancel.

@laughingman7743 laughingman7743 Oct 18, 2026

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 CLI 0.163.0 (codex exec -s read-only, session 01a14cf6-50bd-78c2-a423-58dafce47c29), static review on a detached snapshot of 32dbdd45bf46285c1fc5c77d3a55bc02795e9c19, base e086eac7bb2120e434c08c2eababccd9d0167fb0. The prompt contained only the diff, the threadsafety contract, and the review request (no PR text or prior findings). The snapshot and the PR worktree were unchanged afterwards.

Result: CLEAN.

Covered: both changed hunks; SparkCursor and the shared start, poll, interrupt, cancel, and state handling; AsyncSparkCursor's calculation ID/future return and ID-based cancellation; AioSparkCursor's timeout handling, defaults, cursor creation, context managers, and reported state; related usage, asyncio, SQLAlchemy, cursor/API docs, cross-references, and relevant tests.

Non-actionable notes: the example follows the page's top-level-await convention (a script needs an async function and asyncio.run()), and since cancellation is best effort, the wait can exceed 60 seconds while the cancellation finishes. Both match the rest of the page and the existing "Cancellation" rules, so there is no change.


```python
import asyncio
from pyathena import aio_connect
from pyathena.aio.spark.cursor import AioSparkCursor

async with await aio_connect(work_group="YOUR_SPARK_WORKGROUP",
cursor_class=AioSparkCursor) as conn:
cursor = await asyncio.to_thread(conn.cursor)
async with cursor:
try:
await asyncio.wait_for(
cursor.execute("""print(spark.read.parquet("s3://YOUR_S3_BUCKET/large_dataset/").count())"""),
timeout=60,
)
except asyncio.TimeoutError:
print(f"Calculation timed out: {cursor.calculation_id} ({cursor.state})")
```
Loading