catch unwrap on panics with polars collect (#13850)

# Description
This resurrects the work from #12866 and fixes #12732. 

Polars panics for a plethora or reasons. While handling panics is
generally frowned upon, in cases like with `polars collect` a panic
cause a lot of work to be lost. Often you might have multiple dataframes
in memory and you are trying one operation and lose all state.

While it possible the panic can leave things a strange state, it is
pretty unlikely as part of a polars pipeline. Most of the time polars
objects are not manipulating dataframes in memory mutability, but rather
creating a new dataframe the operations being applied. This is always
the case with a lazy pipeline. After the collect call, the original
dataframes are intact still and I haven't observed any side effects.
This commit is contained in:
Jack Wright 2024-09-15 05:21:02 -07:00 committed by GitHub
parent c9cb62067c
commit c535c24d03
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 35 additions and 11 deletions

View File

@ -55,16 +55,21 @@ impl NuLazyFrame {
}
pub fn collect(self, span: Span) -> Result<NuDataFrame, ShellError> {
self.to_polars()
.collect()
.map_err(|e| ShellError::GenericError {
error: "Error collecting lazy frame".into(),
msg: e.to_string(),
span: Some(span),
help: None,
inner: vec![],
})
.map(|df| NuDataFrame::new(true, df))
crate::handle_panic(
|| {
self.to_polars()
.collect()
.map_err(|e| ShellError::GenericError {
error: "Error collecting lazy frame".into(),
msg: e.to_string(),
span: Some(span),
help: None,
inner: vec![],
})
.map(|df| NuDataFrame::new(true, df))
},
span,
)
}
pub fn apply_with_expr<F>(self, expr: NuExpression, f: F) -> Self

View File

@ -1,4 +1,7 @@
use std::cmp::Ordering;
use std::{
cmp::Ordering,
panic::{catch_unwind, AssertUnwindSafe},
};
use cache::cache_commands;
pub use cache::{Cache, Cacheable};
@ -209,6 +212,22 @@ impl Plugin for PolarsPlugin {
}
}
pub(crate) fn handle_panic<F, R>(f: F, span: Span) -> Result<R, ShellError>
where
F: FnOnce() -> Result<R, ShellError>,
{
match catch_unwind(AssertUnwindSafe(f)) {
Ok(inner_result) => inner_result,
Err(_) => Err(ShellError::GenericError {
error: "Panic occurred".into(),
msg: "".into(),
span: Some(span),
help: None,
inner: vec![],
}),
}
}
#[cfg(test)]
pub mod test {
use super::*;