Column selector using FullColumnPath (#3572)

* Column selector using FullColumnPath

* column name with as

* standar group by name
This commit is contained in:
Fernando Herrera
2021-06-08 03:34:37 +01:00
committed by GitHub
parent 7eadbd938d
commit 4e6c2c0fa1
3 changed files with 34 additions and 9 deletions

View File

@ -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,
}]
}

View File

@ -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(