Column selector using FullColumnPath (#3572)

* Column selector using FullColumnPath

* column name with as

* standar group by name

crates/nu-command/src/commands/dataframe/pivot.rs

@@ -72,7 +72,7 @@ impl WholeStreamCommand for DataFrame {
         vec![Example {
             description: "Pivot a dataframe on b and aggregation on col c",
             example:
-                "[[a b c]; [one x 1] [two y 2]] | pls to-df | pls groupby [a] | pls pivot b c sum",
+                "[[a b c]; [one x 1] [two y 2]] | pls to-df | pls group-by [a] | pls pivot b c sum",
             result: None,
         }]
     }

crates/nu-command/src/commands/dataframe/with_column.rs

@@ -5,6 +5,7 @@ use nu_protocol::{
     dataframe::{NuDataFrame, PolarsData},
     Signature, SyntaxShape, UntaggedValue, Value,
 };
+use nu_source::Tagged;
 
 use super::utils::parse_polars_error;
 pub struct DataFrame;
@@ -19,11 +20,10 @@ impl WholeStreamCommand for DataFrame {
     }
 
     fn signature(&self) -> Signature {
-        Signature::build("pls with-column").required(
-            "series",
-            SyntaxShape::Any,
-            "series to be added",
-        )
+        Signature::build("pls with-column")
+            .required("series", SyntaxShape::Any, "series to be added")
+            .required("as", SyntaxShape::String, "the word 'as'")
+            .required("name", SyntaxShape::String, "column name")
     }
 
     fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
@@ -33,7 +33,8 @@ impl WholeStreamCommand for DataFrame {
     fn examples(&self) -> Vec<Example> {
         vec![Example {
             description: "Adds a series to the dataframe",
-            example: "[[a b]; [1 2] [3 4]] | pls to-df | pls with-column ([5 6] | pls to-series)",
+            example:
+                "[[a b]; [1 2] [3 4]] | pls to-df | pls with-column ([5 6] | pls to-series) as c",
             result: None,
         }]
     }
@@ -43,8 +44,9 @@ fn command(args: CommandArgs) -> Result<OutputStream, ShellError> {
     let tag = args.call_info.name_tag.clone();
     let mut args = args.evaluate_once()?;
     let value: Value = args.req(0)?;
+    let name: Tagged<String> = args.req(2)?;
 
-    let series = match value.value {
+    let mut series = match value.value {
         UntaggedValue::DataFrame(PolarsData::Series(series)) => Ok(series),
         _ => Err(ShellError::labeled_error(
             "Incorrect type",
@@ -53,11 +55,13 @@ fn command(args: CommandArgs) -> Result<OutputStream, ShellError> {
         )),
     }?;
 
+    let series = series.as_mut().rename(name.item.as_ref()).clone();
+
     let mut df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
 
     let res = df
         .as_mut()
-        .with_column(series.series())
+        .with_column(series)
         .map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
 
     Ok(OutputStream::one(NuDataFrame::dataframe_to_value(

crates/nu-value-ext/src/lib.rs

@@ -11,6 +11,9 @@ use nu_source::{
 };
 use num_traits::cast::ToPrimitive;
 
+#[cfg(feature = "dataframe")]
+use nu_protocol::dataframe::{NuSeries, PolarsData};
+
 pub trait ValueExt {
     fn into_parts(self) -> (UntaggedValue, Tag);
     fn get_data(&self, desc: &str) -> MaybeOwned<'_, Value>;
@@ -199,6 +202,24 @@ pub fn get_data_by_member(value: &Value, name: &PathMember) -> Result<Value, She
                 }
             }
         }
+        #[cfg(feature = "dataframe")]
+        UntaggedValue::DataFrame(PolarsData::EagerDataFrame(df)) => match &name.unspanned {
+            UnspannedPathMember::String(string) => {
+                let column = df.as_ref().column(string.as_ref()).map_err(|e| {
+                    ShellError::labeled_error("Dataframe error", format!("{}", e), &name.span)
+                })?;
+
+                Ok(NuSeries::series_to_value(
+                    column.clone(),
+                    Tag::new(value.anchor(), name.span),
+                ))
+            }
+            _ => Err(ShellError::labeled_error(
+                "Integer as column",
+                "Only string as column name",
+                &name.span,
+            )),
+        },
         other => Err(ShellError::type_error(
             "row or table",
             other.type_name().spanned(value.tag.span),