Public/nushell
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

formatting

Browse Source
This commit is contained in:
Jack Wright 2024-07-01 14:33:58 -07:00
parent 074ec52754
commit a467939f67
1 changed files with 22 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
crates/nu_plugin_polars/src/dataframe/eager/pivot.rs
View File

@ -1,7 +1,7 @@
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{ use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
SyntaxShape, Type, Value, Value,
}; };
use polars_ops::pivot::pivot; use polars_ops::pivot::pivot;
@ -73,7 +73,9 @@ impl PluginCommand for PivotDF {
) -> Result<PipelineData, LabeledError> { ) -> Result<PipelineData, LabeledError> {
match PolarsPluginObject::try_from_pipeline(plugin, input, call.head)? { match PolarsPluginObject::try_from_pipeline(plugin, input, call.head)? {
PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df),
PolarsPluginObject::NuLazyFrame(lazy) => command_eager(plugin, engine, call, lazy.collect(call.head)?), PolarsPluginObject::NuLazyFrame(lazy) => {
command_eager(plugin, engine, call, lazy.collect(call.head)?)
}
_ => Err(ShellError::GenericError { _ => Err(ShellError::GenericError {
error: "Must be a dataframe or lazy dataframe".into(), error: "Must be a dataframe or lazy dataframe".into(),
msg: "".into(), msg: "".into(),
@ -103,17 +105,25 @@ fn command_eager(
check_column_datatypes(df.as_ref(), &on_col_string, id_col_span)?; check_column_datatypes(df.as_ref(), &on_col_string, id_col_span)?;
check_column_datatypes(df.as_ref(), &index_col_string, index_col_span)?; check_column_datatypes(df.as_ref(), &index_col_string, index_col_span)?;
check_column_datatypes(df.as_ref(), &val_col_string, val_col_span)?; check_column_datatypes(df.as_ref(), &val_col_string, val_col_span)?;
let polars_df = df.to_polars(); let polars_df = df.to_polars();
// todo add other args // todo add other args
let pivoted = pivot(&polars_df, &on_col_string, Some(&index_col_string), Some(&val_col_string), false, None, None) let pivoted = pivot(
.map_err(|e| ShellError::GenericError { &polars_df,
error: format!("Pivot error: {e}"), &on_col_string,
msg: "".into(), Some(&index_col_string),
span: Some(call.head), Some(&val_col_string),
help: None, false,
inner: vec![], None,
})?; None,
)
.map_err(|e| ShellError::GenericError {
error: format!("Pivot error: {e}"),
msg: "".into(),
span: Some(call.head),
help: None,
inner: vec![],
})?;
let res = NuDataFrame::new(false, pivoted); let res = NuDataFrame::new(false, pivoted);
res.to_pipeline_data(plugin, engine, call.head) res.to_pipeline_data(plugin, engine, call.head)