Skip to content

Commit

Permalink
catch unwrap on panics with polars collect (nushell#13850)
Browse files Browse the repository at this point in the history
# Description
This resurrects the work from nushell#12866 and fixes nushell#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.
  • Loading branch information
ayax79 authored Sep 15, 2024
1 parent c9cb620 commit c535c24
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 11 deletions.
25 changes: 15 additions & 10 deletions crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
21 changes: 20 additions & 1 deletion crates/nu_plugin_polars/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -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::*;
Expand Down

0 comments on commit c535c24

Please sign in to comment.