diff --git a/crates/nu-command/src/commands/dataframe/pivot.rs b/crates/nu-command/src/commands/dataframe/pivot.rs index b466cda97f..84ae324cf9 100644 --- a/crates/nu-command/src/commands/dataframe/pivot.rs +++ b/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, }] } diff --git a/crates/nu-command/src/commands/dataframe/with_column.rs b/crates/nu-command/src/commands/dataframe/with_column.rs index 65842be82d..394ff80c2d 100644 --- a/crates/nu-command/src/commands/dataframe/with_column.rs +++ b/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 { @@ -33,7 +33,8 @@ impl WholeStreamCommand for DataFrame { fn examples(&self) -> Vec { 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 { let tag = args.call_info.name_tag.clone(); let mut args = args.evaluate_once()?; let value: Value = args.req(0)?; + let name: Tagged = 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 { )), }?; + 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( diff --git a/crates/nu-value-ext/src/lib.rs b/crates/nu-value-ext/src/lib.rs index ce7939aa07..4a561c94fd 100644 --- a/crates/nu-value-ext/src/lib.rs +++ b/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 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),