Improve exception handling in wait_for_future

Updated wait_for_future to surface pending Python exceptions by
executing bytecode during signal checks, ensuring that asynchronous
interrupts are processed promptly. Enhanced PartitionedDataFrameStreamReader
to cancel remaining partition streams on projection errors or Python
interrupts, allowing for clean iteration stops. Added regression tests
to validate interrupted Arrow C stream reads and improve timing
for RecordBatchReader.read_all cancellations.

Author: kosiew

python/tests/test_dataframe.py

@@ -3349,17 +3349,92 @@ def read_stream():
         pytest.fail("Stream read operation timed out after 10 seconds")
 
     # Check if we got the expected KeyboardInterrupt
-    if read_exception and isinstance(read_exception[0], type) and read_exception[0] == KeyboardInterrupt:
-        interrupted = True
-    elif read_exception:
-        interrupt_error = read_exception[0]
+    if read_exception:
+        if isinstance(read_exception[0], type) and read_exception[0] == KeyboardInterrupt:
+            interrupted = True
+        elif "KeyboardInterrupt" in str(read_exception[0]):
+            interrupted = True
+        else:
+            interrupt_error = read_exception[0]
 
     if not interrupted:
         pytest.fail(f"Stream was not interrupted; got error: {interrupt_error}")
 
     interrupt_thread.join(timeout=1.0)
 
 
+def test_record_batch_reader_interrupt_exits_quickly(ctx):
+    df = ctx.sql(
+        """
+        SELECT t1.value AS a, t2.value AS a2
+        FROM range(0, 1000000, 1) AS t1
+        JOIN range(0, 1000000, 1) AS t2 ON t1.value = t2.value
+        """
+    )
+
+    reader = pa.RecordBatchReader.from_stream(df)
+
+    query_started = threading.Event()
+    read_thread_id = None
+    interrupt_time = None
+    completion_time = None
+    read_exception = []
+
+    def trigger_interrupt():
+        nonlocal interrupt_time
+        if not query_started.wait(timeout=5.0):
+            pytest.fail("Query did not start in time")
+
+        time.sleep(0.1)
+        interrupt_time = time.time()
+
+        if read_thread_id is None:
+            pytest.fail("Read thread did not record an identifier")
+
+        res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
+            ctypes.c_long(read_thread_id), ctypes.py_object(KeyboardInterrupt)
+        )
+        if res != 1:
+            ctypes.pythonapi.PyThreadState_SetAsyncExc(
+                ctypes.c_long(read_thread_id), ctypes.py_object(0)
+            )
+            pytest.fail("Failed to raise KeyboardInterrupt in read thread")
+
+    def read_stream():
+        nonlocal read_thread_id, completion_time
+        read_thread_id = threading.get_ident()
+        try:
+            query_started.set()
+            reader.read_all()
+        except KeyboardInterrupt:
+            completion_time = time.time()
+        except Exception as exc:  # pragma: no cover - unexpected failure path
+            completion_time = time.time()
+            read_exception.append(exc)
+
+    read_thread = threading.Thread(target=read_stream, daemon=True)
+    interrupt_thread = threading.Thread(target=trigger_interrupt, daemon=True)
+
+    read_thread.start()
+    interrupt_thread.start()
+
+    read_thread.join(timeout=10.0)
+    if read_thread.is_alive():
+        pytest.fail("Stream read operation timed out after 10 seconds")
+
+    interrupt_thread.join(timeout=1.0)
+
+    if read_exception and "KeyboardInterrupt" not in str(read_exception[0]):
+        pytest.fail(f"Read thread raised unexpected exception: {read_exception[0]}")
+
+    assert completion_time is not None, "Read thread did not finish"
+    assert interrupt_time is not None, "Interrupt was not sent"
+
+    elapsed = completion_time - interrupt_time
+    assert elapsed >= 0, "Completion recorded before interrupt was sent"
+    assert elapsed < 1.5, f"Cancellation took too long: {elapsed}s"
+
+
 def test_show_select_where_no_rows(capsys) -> None:
     ctx = SessionContext()
     df = ctx.sql("SELECT 1 WHERE 1=0")

src/dataframe.rs

@@ -375,6 +375,13 @@ struct PartitionedDataFrameStreamReader {
     current: usize,
 }
 
+impl PartitionedDataFrameStreamReader {
+    fn cancel_streams(&mut self) {
+        self.streams.drain(self.current..);
+        self.current = self.streams.len();
+    }
+}
+
 impl Iterator for PartitionedDataFrameStreamReader {
     type Item = Result<RecordBatch, ArrowError>;
 
@@ -389,7 +396,10 @@ impl Iterator for PartitionedDataFrameStreamReader {
                     let batch = if let Some(ref schema) = self.projection {
                         match record_batch_into_schema(batch, schema.as_ref()) {
                             Ok(b) => b,
-                            Err(e) => return Some(Err(e)),
+                            Err(e) => {
+                                self.cancel_streams();
+                                return Some(Err(e));
+                            }
                         }
                     } else {
                         batch
@@ -401,9 +411,11 @@ impl Iterator for PartitionedDataFrameStreamReader {
                     continue;
                 }
                 Ok(Err(e)) => {
+                    self.cancel_streams();
                     return Some(Err(ArrowError::ExternalError(Box::new(e))));
                 }
                 Err(e) => {
+                    self.cancel_streams();
                     return Some(Err(ArrowError::ExternalError(Box::new(e))));
                 }
             }

src/utils.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::ffi::CString;
 use std::future::Future;
 use std::sync::{Arc, OnceLock};
 use std::time::Duration;
@@ -27,6 +28,7 @@ use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::PyCapsule;
+use pyo3::PyErr;
 use tokio::runtime::Runtime;
 use tokio::task::JoinHandle;
 use tokio::time::sleep;
@@ -84,7 +86,15 @@ where
                 tokio::select! {
                     res = &mut fut => break Ok(res),
                     _ = sleep(INTERVAL_CHECK_SIGNALS) => {
-                        Python::attach(|py| py.check_signals())?;
+                        Python::attach(|py| {
+                            if let Some(err) = PyErr::take(py) {
+                                Err(err)
+                            } else {
+                                let code = CString::new("pass").unwrap();
+                                py.run(code.as_c_str(), None, None)?;
+                                py.check_signals()
+                            }
+                        })?;
                     }
                 }
             }