Enable interrupt handling

Author: kosiew

src/dataframe.rs

@@ -391,7 +391,7 @@ impl PyDataFrame {
     /// Unless some order is specified in the plan, there is no
     /// guarantee of the order of the result.
     fn collect(&self, py: Python) -> PyResult<Vec<PyObject>> {
-        let batches = wait_for_future(py, self.df.as_ref().clone().collect())
+        let batches = wait_for_future(py, self.df.as_ref().clone().collect())?
             .map_err(PyDataFusionError::from)?;
         // cannot use PyResult<Vec<RecordBatch>> return type due to
         // https://github.com/PyO3/pyo3/issues/1813
@@ -407,8 +407,9 @@ impl PyDataFrame {
     /// Executes this DataFrame and collects all results into a vector of vector of RecordBatch
     /// maintaining the input partitioning.
     fn collect_partitioned(&self, py: Python) -> PyResult<Vec<Vec<PyObject>>> {
-        let batches = wait_for_future(py, self.df.as_ref().clone().collect_partitioned())
-            .map_err(PyDataFusionError::from)?;
+        let batches =
+            wait_for_future(py, self.df.as_ref().clone().collect_partitioned())?
+                .map_err(PyDataFusionError::from)?;
 
         batches
             .into_iter()

src/record_batch.rs

@@ -63,7 +63,7 @@ impl PyRecordBatchStream {
 impl PyRecordBatchStream {
     fn next(&mut self, py: Python) -> PyResult<PyRecordBatch> {
         let stream = self.stream.clone();
-        wait_for_future(py, next_stream(stream, true))
+        wait_for_future(py, next_stream(stream, true))?
     }
 
     fn __next__(&mut self, py: Python) -> PyResult<PyRecordBatch> {

src/utils.rs

@@ -24,6 +24,7 @@ use datafusion::logical_expr::Volatility;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::PyCapsule;
+use pyo3::PyErr;
 use std::future::Future;
 use std::sync::OnceLock;
 use tokio::runtime::Runtime;
@@ -47,14 +48,44 @@ pub(crate) fn get_global_ctx() -> &'static SessionContext {
     CTX.get_or_init(SessionContext::new)
 }
 
-/// Utility to collect rust futures with GIL released
-pub fn wait_for_future<F>(py: Python, f: F) -> F::Output
+/// Utility to collect rust futures with GIL released and respond to
+/// Python interrupts such as ``KeyboardInterrupt``. If a signal is
+/// received while the future is running, the future is aborted and the
+/// corresponding Python exception is raised.
+pub fn wait_for_future<F>(py: Python, f: F) -> PyResult<F::Output>
 where
-    F: Future + Send,
-    F::Output: Send,
+    F: Future + Send + 'static,
+    F::Output: Send + 'static,
 {
+    use std::{thread, time::Duration};
+    use tokio::task::JoinHandle;
+
     let runtime: &Runtime = &get_tokio_runtime().0;
-    py.allow_threads(|| runtime.block_on(f))
+
+    // Spawn the future so it can be aborted if a signal is received
+    let handle: JoinHandle<F::Output> = runtime.spawn(f);
+
+    let mut interrupt: Option<PyErr> = None;
+    py.allow_threads(|| {
+        while !handle.is_finished() {
+            thread::sleep(Duration::from_millis(10));
+            Python::with_gil(|py| {
+                if let Err(err) = py.check_signals() {
+                    handle.abort();
+                    interrupt = Some(err);
+                }
+            });
+            if interrupt.is_some() {
+                break;
+            }
+        }
+    });
+
+    if let Some(err) = interrupt {
+        return Err(err);
+    }
+
+    Ok(runtime.block_on(handle).expect("Tokio task panicked"))
 }
 
 pub(crate) fn parse_volatility(value: &str) -> PyDataFusionResult<Volatility> {