diff --git a/crates/nu-command/src/commands/dataframe/melt.rs b/crates/nu-command/src/commands/dataframe/melt.rs index 88f1f76fb..a1e509f70 100644 --- a/crates/nu-command/src/commands/dataframe/melt.rs +++ b/crates/nu-command/src/commands/dataframe/melt.rs @@ -5,6 +5,7 @@ use nu_protocol::{ dataframe::{Column, NuDataFrame}, Signature, SyntaxShape, UntaggedValue, Value, }; +use nu_source::Tagged; use super::utils::convert_columns; @@ -33,6 +34,18 @@ impl WholeStreamCommand for DataFrame { "column names used as value columns", Some('v'), ) + .named( + "variable_name", + SyntaxShape::String, + "optional name for variable column", + Some('r'), + ) + .named( + "value_name", + SyntaxShape::String, + "optional name for value column", + Some('l'), + ) } fn run(&self, args: CommandArgs) -> Result { @@ -105,6 +118,9 @@ fn command(mut args: CommandArgs) -> Result { let id_col: Vec = args.req_named("columns")?; let val_col: Vec = args.req_named("values")?; + let value_name: Option> = args.get_flag("value_name")?; + let variable_name: Option> = args.get_flag("variable_name")?; + let (id_col_string, id_col_span) = convert_columns(&id_col, &tag)?; let (val_col_string, val_col_span) = convert_columns(&val_col, &tag)?; @@ -113,11 +129,21 @@ fn command(mut args: CommandArgs) -> Result { check_column_datatypes(df.as_ref(), &id_col_string, &id_col_span)?; check_column_datatypes(df.as_ref(), &val_col_string, &val_col_span)?; - let res = df + let mut res = df .as_ref() .melt(&id_col_string, &val_col_string) .map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?; + if let Some(name) = &variable_name { + res.rename("variable", name.item.as_str()) + .map_err(|e| parse_polars_error::<&str>(&e, &name.tag.span, None))?; + } + + if let Some(name) = &value_name { + res.rename("value", name.item.as_str()) + .map_err(|e| parse_polars_error::<&str>(&e, &name.tag.span, None))?; + } + Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, tag))) } diff --git a/crates/nu-command/src/commands/dataframe/mod.rs b/crates/nu-command/src/commands/dataframe/mod.rs index 227f22bb0..4df4a935f 100644 --- a/crates/nu-command/src/commands/dataframe/mod.rs +++ b/crates/nu-command/src/commands/dataframe/mod.rs @@ -18,6 +18,7 @@ pub mod list; pub mod melt; pub mod open; pub mod pivot; +pub mod rename; pub mod sample; pub mod select; pub mod shape; @@ -52,6 +53,7 @@ pub use list::DataFrame as DataFrameList; pub use melt::DataFrame as DataFrameMelt; pub use open::DataFrame as DataFrameOpen; pub use pivot::DataFrame as DataFramePivot; +pub use rename::DataFrame as DataFrameRename; pub use sample::DataFrame as DataFrameSample; pub use select::DataFrame as DataFrameSelect; pub use shape::DataFrame as DataFrameShape; diff --git a/crates/nu-command/src/commands/dataframe/rename.rs b/crates/nu-command/src/commands/dataframe/rename.rs new file mode 100644 index 000000000..d29e50cf7 --- /dev/null +++ b/crates/nu-command/src/commands/dataframe/rename.rs @@ -0,0 +1,81 @@ +use crate::prelude::*; +use nu_engine::WholeStreamCommand; +use nu_errors::ShellError; +use nu_protocol::{ + dataframe::{Column, NuDataFrame}, + Signature, SyntaxShape, UntaggedValue, +}; + +use nu_source::Tagged; + +use super::utils::parse_polars_error; +pub struct DataFrame; + +impl WholeStreamCommand for DataFrame { + fn name(&self) -> &str { + "dataframe rename-col" + } + + fn usage(&self) -> &str { + "[DataFrame] rename a dataframe column" + } + + fn signature(&self) -> Signature { + Signature::build("dataframe rename-col") + .required("from", SyntaxShape::String, "column name to be renamed") + .required("to", SyntaxShape::String, "new column name") + } + + fn run(&self, args: CommandArgs) -> Result { + command(args) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Renames a dataframe column", + example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe rename-col a ab", + result: Some(vec![NuDataFrame::try_from_columns( + vec![ + Column::new( + "ab".to_string(), + vec![UntaggedValue::int(1).into(), UntaggedValue::int(3).into()], + ), + Column::new( + "b".to_string(), + vec![UntaggedValue::int(2).into(), UntaggedValue::int(4).into()], + ), + ], + &Span::default(), + ) + .expect("simple df for test should not fail") + .into_value(Tag::default())]), + }] + } +} + +fn command(mut args: CommandArgs) -> Result { + let tag = args.call_info.name_tag.clone(); + let from: Tagged = args.req(0)?; + let to: Tagged = args.req(1)?; + + let (mut df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?; + + df.as_mut() + .rename(from.item.as_str(), to.item.as_str()) + .map_err(|e| parse_polars_error::<&str>(&e, &df_tag.span, None))?; + + Ok(OutputStream::one(df.into_value(tag))) +} + +#[cfg(test)] +mod tests { + use super::DataFrame; + use super::ShellError; + + #[test] + fn examples_work_as_expected() -> Result<(), ShellError> { + use crate::examples::test_dataframe as test_examples; + + test_examples(DataFrame {}) + } +} diff --git a/crates/nu-command/src/commands/mod.rs b/crates/nu-command/src/commands/mod.rs index b4d326158..7b598735c 100644 --- a/crates/nu-command/src/commands/mod.rs +++ b/crates/nu-command/src/commands/mod.rs @@ -34,7 +34,7 @@ pub use dataframe::{ DataFrameGetOrdinal, DataFrameGetSecond, DataFrameGetWeek, DataFrameGetWeekDay, DataFrameGetYear, DataFrameGroupBy, DataFrameIsDuplicated, DataFrameIsIn, DataFrameIsNotNull, DataFrameIsNull, DataFrameIsUnique, DataFrameJoin, DataFrameLast, DataFrameList, DataFrameMelt, - DataFrameNNull, DataFrameNUnique, DataFrameNot, DataFrameOpen, DataFramePivot, + DataFrameNNull, DataFrameNUnique, DataFrameNot, DataFrameOpen, DataFramePivot, DataFrameRename, DataFrameReplace, DataFrameReplaceAll, DataFrameRolling, DataFrameSample, DataFrameSelect, DataFrameSeriesRename, DataFrameSet, DataFrameSetWithIdx, DataFrameShape, DataFrameShift, DataFrameShow, DataFrameSlice, DataFrameSort, DataFrameStrFTime, DataFrameStringLengths, diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index dcf1c5398..46f2c5d72 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -347,6 +347,7 @@ pub fn create_default_context(interactive: bool) -> Result