rename command and rename for melt (#3968)

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<OutputStream, ShellError> {
@@ -105,6 +118,9 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
     let id_col: Vec<Value> = args.req_named("columns")?;
     let val_col: Vec<Value> = args.req_named("values")?;
 
+    let value_name: Option<Tagged<String>> = args.get_flag("value_name")?;
+    let variable_name: Option<Tagged<String>> = 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<OutputStream, ShellError> {
     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)))
 }
 

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;

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<OutputStream, ShellError> {
+        command(args)
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        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<OutputStream, ShellError> {
+    let tag = args.call_info.name_tag.clone();
+    let from: Tagged<String> = args.req(0)?;
+    let to: Tagged<String> = 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 {})
+    }
+}

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,

crates/nu-command/src/default_context.rs

@@ -347,6 +347,7 @@ pub fn create_default_context(interactive: bool) -> Result<EvaluationContext, Bo
             whole_stream_command(DataFrameDescribe),
             whole_stream_command(DataFrameRolling),
             whole_stream_command(DataFrameCumulative),
+            whole_stream_command(DataFrameRename),
         ]);
 
         #[cfg(feature = "clipboard-cli")]