Add Interruptible Query Execution in Jupyter via KeyboardInterrupt Support #1141

kosiew · 2025-06-04T09:25:37Z

Which issue does this PR close?

Rationale for this change

Currently, when executing long-running queries in Jupyter notebooks using DataFusion, the interrupt (Ctrl+C) mechanism does not function as expected. Users are forced to restart the kernel, causing workflow disruption and data/context loss. This PR introduces a graceful interrupt mechanism allowing rapid development iterations and better notebook UX.

What changes are included in this PR?

Introduces test_collect_interrupted to validate interruptibility via KeyboardInterrupt on long-running queries.
Adds polling and signaling logic to safely trigger interrupts in main thread.
Replaces all uses of wait_for_future(py, fut)? with wait_for_future(py, fut)?? to properly propagate nested Result types.
Updates wait_for_future utility to periodically check for Python signals using py.check_signals() to allow interrupt handling.
Enhances resilience and correctness of wait_for_future and associated async/await error propagation.
Ensures consistent behavior across contexts like SQL queries, file reads, and execution plans.

Note that in Jupyter, you click the Stop button to interrupt a query:

Jun-04-2025.17-23-59.mp4

Are these changes tested?

✅ Yes

A dedicated test test_collect_interrupted ensures interrupt behavior works reliably.
Additional tests continue to pass after modifications to async handling logic.

Are there any user-facing changes?

✅ Yes

Users can now interrupt long-running DataFusion queries in Jupyter notebooks without terminating the entire kernel session.
Graceful recovery is possible — queries can be modified and re-executed post-interrupt.
Improved error messages and behavior for missing tables (e.g., PyKeyError for "No table named").

…yDataFrame

…ecution plan

…atch iteration

…it_for_future to use it

This reverts commit 52a5efe.

…xt methods

…or_future result

…Fusion

This reverts commit ca2d892.

This reverts commit b8ce3e4.

…errupted

…Stream

…for_future

- Simplified async handling by removing unnecessary cloning of strings and context in various methods. - Streamlined the use of `wait_for_future` to directly handle futures without intermediate variables. - Improved error handling by directly propagating results from futures. - Enhanced readability by reducing boilerplate code in methods related to reading and writing data. - Updated the `wait_for_future` function to improve signal checking and future handling.

timsaucer

This is amazing!

One suggestion in two places to avoid having a panic.

src/dataframe.rs

kosiew · 2025-06-13T15:43:16Z

@timsaucer
Thanks for the review and nipping the panic.

… PyDataFrame

timsaucer

Thank you!

kosiew added 30 commits June 3, 2025 11:30

fix: enhance error handling in async wait_for_future function

60ff7f2

feat: implement async execution for execution plans in PySessionContext

b168afb

fix: improve error message for execution failures in PySessionContext

cca1cf8

fix: enhance error handling and improve execution plan retrieval in P…

81e4204

…yDataFrame

fix: ensure 'static lifetime for futures in wait_for_future function

cfc3c40

fix: handle potential errors when caching DataFrame and retrieving ex…

1ee861a

…ecution plan

fix: flatten batches in PyDataFrame to ensure proper schema conversion

569dc26

fix: correct error handling in batch processing for schema conversion

45c9fa3

fix: flatten nested structure in PyDataFrame to ensure proper RecordB…

5336cf4

…atch iteration

fix: improve error handling in PyDataFrame stream execution

fb814cc

fix: add utility to get Tokio Runtime with time enabled and update wa…

a24e280

…it_for_future to use it

fix: store result of converting RecordBatches to PyArrow for debugging

6f0ef29

fix: handle error from wait_for_future in PyDataFrame collect method

ec2abf1

fix: propagate error from wait_for_future in collect_partitioned method

aa71328

fix: enable IO in Tokio runtime with time support

d9bfed5

main register_listing_table

52a5efe

Revert "main register_listing_table"

b778911

This reverts commit 52a5efe.

fix: propagate error correctly from wait_for_future in PySessionConte…

f10652c

…xt methods

fix: simplify error handling in PySessionContext by unwrapping wait_f…

97f86dc

…or_future result

test: add interruption handling test for long-running queries in Data…

3f17c9c

…Fusion

test: move test_collect_interrupted to test_dataframe.py

dafa3a5

fix: add const for interval in wait_for_future utility

973e690

fix: use get_tokio_runtime instead of the custom get_runtime

ca2d892

Revert "fix: use get_tokio_runtime instead of the custom get_runtime"

ada6faa

This reverts commit ca2d892.

fix: use get_tokio_runtime instead of the custom get_runtime

f14b51f

.

b8ce3e4

Revert "."

5073070

This reverts commit b8ce3e4.

fix: improve query interruption handling in test_collect_interrupted

63a1ab7

fix: ensure proper handling of query interruption in test_collect_int…

c6e5205

…errupted

fix: improve error handling in database table retrieval

d0f2d37

kosiew added 15 commits June 4, 2025 14:04

create_csv_read_options -> PyDataFusionResult

40b4345

revert to before create_csv_read_options

bcc4f81

refactor: simplify file compression type parsing in PySessionContext

8b0e2e1

fix: parse_compression_type once only

b1e67df

add create_ndjson_read_options

393e7ca

refactor comment for clarity in wait_for_future function

268a855

refactor wait_for_future to avoid spawn

bf91cf8

remove unused py_err_to_datafusion_err function

c3d808c

add comment to clarify error handling in next method of PyRecordBatch…

ccb4a5c

…Stream

handle error from wait_for_future in PySubstraitSerializer

ea3673b

clarify comment on future pinning in wait_for_future function

8c50e7f

refactor wait_for_future to use Duration for signal check interval

aead2a0

handle error from wait_for_future in count method of PyDataFrame

84f31bb

fix ruff errors

e2f1e1c

fix clippy errors

d5cf16d

kosiew force-pushed the interrupt-1136 branch from dffc0ed to d5cf16d Compare June 4, 2025 10:09

kosiew added 3 commits June 4, 2025 18:18

remove unused get_and_enter_tokio_runtime function and simplify wait_…

48e4f7c

…for_future

Organize imports in utils.rs for improved readability

0179619

kosiew changed the title ~~Improve async Python integration with better error handling, signal checks, and KeyboardInterrupt support in query execution~~ Add Interruptible Query Execution in Jupyter via KeyboardInterrupt Support Jun 5, 2025

timsaucer reviewed Jun 13, 2025

View reviewed changes

src/dataframe.rs Outdated Show resolved Hide resolved

src/dataframe.rs Outdated Show resolved Hide resolved

kosiew marked this pull request as draft June 13, 2025 16:33

kosiew added 2 commits June 14, 2025 10:14

map_err instead of panic

bc6f5a6

Fix error handling in async stream execution for PySessionContext and…

42ab687

… PyDataFrame

kosiew force-pushed the interrupt-1136 branch from 00e180f to 42ab687 Compare June 14, 2025 02:14

kosiew marked this pull request as ready for review June 14, 2025 06:40

Merge branch 'main' into interrupt-1136

79ebe1e

timsaucer approved these changes Jun 16, 2025

View reviewed changes

timsaucer merged commit dc0d35a into apache:main Jun 16, 2025
17 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add Interruptible Query Execution in Jupyter via KeyboardInterrupt Support #1141

Add Interruptible Query Execution in Jupyter via KeyboardInterrupt Support #1141

Uh oh!

kosiew commented Jun 4, 2025 •

edited

Loading

Uh oh!

timsaucer left a comment

Uh oh!

Uh oh!

Uh oh!

kosiew commented Jun 13, 2025

Uh oh!

timsaucer left a comment

Uh oh!

Uh oh!

Uh oh!

Add Interruptible Query Execution in Jupyter via KeyboardInterrupt Support #1141

Add Interruptible Query Execution in Jupyter via KeyboardInterrupt Support #1141

Uh oh!

Conversation

kosiew commented Jun 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

timsaucer left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

kosiew commented Jun 13, 2025

Uh oh!

timsaucer left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

kosiew commented Jun 4, 2025 •

edited

Loading