feat: Improved support for KeyboardInterrupts (#20961)

Author: orlp

Cargo.lock

@@ -123,12 +123,6 @@ pub fn register_object_builder(
     })
 }
 
-pub fn is_object_builder_registered() -> bool {
-    let reg = GLOBAL_OBJECT_REGISTRY.deref();
-    let reg = reg.read().unwrap();
-    reg.is_some()
-}
-
 #[cold]
 pub fn get_object_physical_type() -> ArrowDataType {
     let reg = GLOBAL_OBJECT_REGISTRY.read().unwrap();

crates/polars-core/src/chunked_array/object/registry.rs

@@ -1,5 +1,5 @@
 use arrow::legacy::kernels::sort_partition::{create_clean_partitions, partition_to_groups};
-use polars_error::check_signals;
+use polars_error::signals::try_raise_keyboard_interrupt;
 use polars_utils::total_ord::{ToTotalOrd, TotalHash};
 
 use super::*;
@@ -235,7 +235,7 @@ where
                 num_groups_proxy(ca, multithreaded, sorted)
             },
         };
-        check_signals()?;
+        try_raise_keyboard_interrupt();
         Ok(out)
     }
 }
@@ -287,7 +287,7 @@ impl IntoGroupsType for BinaryChunked {
         } else {
             group_by(bh[0].iter(), sorted)
         };
-        check_signals()?;
+        try_raise_keyboard_interrupt();
         Ok(out)
     }
 }

crates/polars-core/src/frame/group_by/into_groups.rs

@@ -37,6 +37,7 @@ pub use crate::chunked_array::StructChunked;
 #[cfg(feature = "dtype-categorical")]
 pub use crate::datatypes::string_cache::StringCacheHolder;
 pub use crate::datatypes::{ArrayCollectIterExt, *};
+pub use crate::error::signals::try_raise_keyboard_interrupt;
 pub use crate::error::{
     polars_bail, polars_ensure, polars_err, polars_warn, PolarsError, PolarsResult,
 };

crates/polars-core/src/prelude.rs

@@ -16,5 +16,8 @@ regex = { workspace = true, optional = true }
 simdutf8 = { workspace = true }
 thiserror = { workspace = true }
 
+[target.'cfg(not(target_family = "wasm"))'.dependencies]
+signal-hook = "0.3"
+
 [features]
 python = []

crates/polars-error/Cargo.toml

@@ -9,9 +9,8 @@ use std::fmt::{self, Display, Formatter, Write};
 use std::ops::Deref;
 use std::sync::{Arc, LazyLock};
 use std::{env, io};
-mod signals;
+pub mod signals;
 
-pub use signals::{check_signals, set_signals_function};
 pub use warning::*;
 
 enum ErrorStrategy {

crates/polars-error/src/lib.rs

@@ -1,23 +1,95 @@
-use crate::PolarsResult;
-
-type SignalsFunction = fn() -> PolarsResult<()>;
-static mut SIGNALS_FUNCTION: Option<SignalsFunction> = None;
-
-/// Set the function that will be called check_signals.
-/// This can be set on startup to enable stopping a query when user input like `ctrl-c` is called.
-///
-/// # Safety
-/// The caller must ensure there is no other thread accessing this function
-/// or calling `check_signals`.
-pub unsafe fn set_signals_function(function: SignalsFunction) {
-    SIGNALS_FUNCTION = Some(function)
+use std::panic::{catch_unwind, UnwindSafe};
+use std::sync::atomic::{AtomicU64, Ordering};
+
+/// Python hooks SIGINT to instead generate a KeyboardInterrupt exception.
+/// So we do the same to try and abort long-running computations and return to
+/// Python so that the Python exception can be generated.
+pub struct KeyboardInterrupt;
+
+// Bottom bit: interrupt flag.
+// Top 63 bits: number of alive interrupt catchers.
+static INTERRUPT_STATE: AtomicU64 = AtomicU64::new(0);
+
+pub fn register_polars_keyboard_interrupt_hook() {
+    let default_hook = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |p| {
+        // Suppress panic output on KeyboardInterrupt.
+        if p.payload().downcast_ref::<KeyboardInterrupt>().is_none() {
+            default_hook(p);
+        }
+    }));
+
+    // WASM doesn't support signals, so we just skip installing the hook there.
+    #[cfg(not(target_family = "wasm"))]
+    unsafe {
+        // SAFETY: we only do an atomic op in the signal handler, which is allowed.
+        signal_hook::low_level::register(signal_hook::consts::signal::SIGINT, move || {
+            // Set the interrupt flag, but only if there are active catchers.
+            INTERRUPT_STATE
+                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |state| {
+                    let num_catchers = state >> 1;
+                    if num_catchers > 0 {
+                        Some(state | 1)
+                    } else {
+                        None
+                    }
+                })
+                .ok();
+        })
+        .unwrap();
+    }
+}
+
+/// Checks if the keyboard interrupt flag is set, and if yes panics with a
+/// KeyboardInterrupt. This function is very cheap.
+#[inline(always)]
+pub fn try_raise_keyboard_interrupt() {
+    if INTERRUPT_STATE.load(Ordering::Relaxed) & 1 != 0 {
+        try_raise_keyboard_interrupt_slow()
+    }
+}
+
+#[inline(never)]
+#[cold]
+fn try_raise_keyboard_interrupt_slow() {
+    std::panic::panic_any(KeyboardInterrupt);
+}
+
+/// Runs the passed function, catching any KeyboardInterrupts if they occur
+/// while running the function.
+pub fn catch_keyboard_interrupt<R, F: FnOnce() -> R + UnwindSafe>(
+    try_fn: F,
+) -> Result<R, KeyboardInterrupt> {
+    // Try to register this catcher (or immediately return if there is an
+    // uncaught interrupt).
+    try_register_catcher()?;
+    let ret = catch_unwind(try_fn);
+    unregister_catcher();
+    ret.map_err(|p| match p.downcast::<KeyboardInterrupt>() {
+        Ok(_) => KeyboardInterrupt,
+        Err(p) => std::panic::resume_unwind(p),
+    })
 }
 
-fn default() -> PolarsResult<()> {
+fn try_register_catcher() -> Result<(), KeyboardInterrupt> {
+    let old_state = INTERRUPT_STATE.fetch_add(2, Ordering::Relaxed);
+    if old_state & 1 != 0 {
+        unregister_catcher();
+        return Err(KeyboardInterrupt);
+    }
     Ok(())
 }
 
-pub fn check_signals() -> PolarsResult<()> {
-    let f = unsafe { SIGNALS_FUNCTION.unwrap_or(default) };
-    f()
+fn unregister_catcher() {
+    INTERRUPT_STATE
+        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |state| {
+            let num_catchers = state >> 1;
+            if num_catchers > 1 {
+                Some(state - 2)
+            } else {
+                // Last catcher, clear interrupt flag.
+                Some(0)
+            }
+        })
+        .ok();
 }

crates/polars-error/src/signals.rs

@@ -5,7 +5,6 @@ use std::sync::{Mutex, RwLock};
 use bitflags::bitflags;
 use once_cell::sync::OnceCell;
 use polars_core::config::verbose;
-use polars_core::error::check_signals;
 use polars_core::prelude::*;
 use polars_ops::prelude::ChunkJoinOptIds;
 
@@ -150,7 +149,7 @@ impl ExecutionState {
 
     // This is wrong when the U64 overflows which will never happen.
     pub fn should_stop(&self) -> PolarsResult<()> {
-        check_signals()?;
+        try_raise_keyboard_interrupt();
         polars_ensure!(!self.stop.load(Ordering::Relaxed), ComputeError: "query interrupted");
         Ok(())
     }

crates/polars-expr/src/state/execution_state.rs

@@ -6,7 +6,6 @@ use std::cmp::Ordering;
 use default::*;
 pub use groups::AsofJoinBy;
 use polars_core::prelude::*;
-use polars_error::check_signals;
 use polars_utils::pl_str::PlSmallStr;
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
@@ -334,7 +333,7 @@ pub trait AsofJoin: IntoDf {
                 join_asof_numeric(ca, &right_key, strategy, tolerance, allow_eq)
             },
         }?;
-        check_signals()?;
+        try_raise_keyboard_interrupt();
 
         // Drop right join column.
         let other = if coalesce && left_key.name() == right_key.name() {

crates/polars-ops/src/frame/join/asof/mod.rs

@@ -133,7 +133,7 @@ fn cross_join_dfs(
         }
     };
     let (l_df, r_df) = if parallel {
-        check_signals()?;
+        try_raise_keyboard_interrupt();
         POOL.install(|| rayon::join(create_left_df, create_right_df))
     } else {
         (create_left_df(), create_right_df())

crates/polars-ops/src/frame/join/cross_join.rs

@@ -87,7 +87,7 @@ pub fn materialize_left_join_from_series(
     } else {
         right.drop(s_right.name()).unwrap()
     };
-    check_signals()?;
+    try_raise_keyboard_interrupt();
 
     #[cfg(feature = "chunked_ids")]
     match (left_idx, right_idx) {

crates/polars-ops/src/frame/join/dispatch_left_right.rs

@@ -155,7 +155,7 @@ pub trait JoinDispatch: IntoDf {
         let (mut join_idx_l, mut join_idx_r) =
             s_left.hash_join_outer(s_right, args.validation, args.join_nulls)?;
 
-        check_signals()?;
+        try_raise_keyboard_interrupt();
         if let Some((offset, len)) = args.slice {
             let (offset, len) = slice_offsets(offset, len, join_idx_l.len());
             join_idx_l.slice(offset, len);

crates/polars-ops/src/frame/join/hash_join/mod.rs

@@ -69,7 +69,7 @@ where
     } else {
         build_tables(build, join_nulls)
     };
-    check_signals()?;
+    try_raise_keyboard_interrupt();
 
     let n_tables = hash_tbls.len();
     let offsets = probe_to_offsets(&probe);

crates/polars-ops/src/frame/join/hash_join/single_keys_inner.rs

@@ -136,7 +136,7 @@ where
     } else {
         build_tables(build, join_nulls)
     };
-    check_signals()?;
+    try_raise_keyboard_interrupt();
     let n_tables = hash_tbls.len();
 
     // we determine the offset so that we later know which index to store in the join tuples

crates/polars-ops/src/frame/join/hash_join/single_keys_left.rs

@@ -231,7 +231,7 @@ where
     let (probe_hashes, _) = create_hash_and_keys_threaded_vectorized(probe, Some(random_state));
 
     let n_tables = hash_tbls.len();
-    check_signals()?;
+    try_raise_keyboard_interrupt();
 
     // probe the hash table.
     // Note: indexes from b that are not matched will be None, Some(idx_b)

crates/polars-ops/src/frame/join/hash_join/single_keys_outer.rs

@@ -12,7 +12,7 @@ use polars_core::prelude::*;
 use polars_core::series::IsSorted;
 use polars_core::utils::{_set_partition_size, split};
 use polars_core::{with_match_physical_numeric_polars_type, POOL};
-use polars_error::{check_signals, polars_err, PolarsResult};
+use polars_error::{polars_err, PolarsResult};
 use polars_utils::binary_search::ExponentialSearch;
 use polars_utils::itertools::Itertools;
 use polars_utils::total_ord::{TotalEq, TotalOrd};
@@ -362,7 +362,7 @@ unsafe fn materialize_join(
     right_row_idx: &IdxCa,
     suffix: Option<PlSmallStr>,
 ) -> PolarsResult<DataFrame> {
-    check_signals()?;
+    try_raise_keyboard_interrupt();
     let (join_left, join_right) = {
         POOL.join(
             || left.take_unchecked(left_row_idx),

crates/polars-ops/src/frame/join/iejoin/mod.rs

@@ -45,7 +45,6 @@ use polars_core::utils::slice_offsets;
 #[allow(unused_imports)]
 use polars_core::utils::slice_slice;
 use polars_core::POOL;
-use polars_error::check_signals;
 use polars_utils::hashing::BytesHash;
 use rayon::prelude::*;
 
@@ -565,7 +564,7 @@ trait DataFrameJoinOpsPrivate: IntoDf {
                 args.maintain_order,
                 MaintainOrderJoin::Left | MaintainOrderJoin::LeftRight
             );
-        check_signals()?;
+        try_raise_keyboard_interrupt();
         let (df_left, df_right) =
             if args.maintain_order != MaintainOrderJoin::None && !already_left_sorted {
                 let mut df =

crates/polars-ops/src/frame/join/mod.rs

@@ -10,6 +10,7 @@ use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
 
 use crate::error::PyPolarsErr;
+use crate::utils::EnterPolarsExt;
 use crate::{PyDataFrame, Wrap};
 
 #[pyclass]
@@ -136,13 +137,7 @@ impl PyBatchedCsv {
 
     fn next_batches(&self, py: Python, n: usize) -> PyResult<Option<Vec<PyDataFrame>>> {
         let reader = &self.reader;
-        let batches = py.allow_threads(move || {
-            reader
-                .lock()
-                .map_err(|e| PyPolarsErr::Other(e.to_string()))?
-                .next_batches(n)
-                .map_err(PyPolarsErr::from)
-        })?;
+        let batches = py.enter_polars(move || reader.lock().unwrap().next_batches(n))?;
 
         // SAFETY: same memory layout
         let batches = unsafe {