holm/nushell
1
1
Fork 0
forked from extern/nushell
Code Issues Pull Requests Packages Projects Releases Wiki Activity

All is a DataFrame (#3812)

Browse Source 
* nuframe in its own type in UntaggedValue

* Removed eager dataframe from enum

* Dataframe created from list of values

* Corrected order in dataframe columns

* Returned tag from stream collection

* Removed series from dataframe commands

* Arithmetic operators

* forced push

* forced push

* Replace all command

* String commands

* appending operations with dfs

* Testing suite for dataframes

* Unit test for dataframe commands

* improved equality for dataframes
This commit is contained in:
Fernando Herrera
2021-07-25 11:01:54 +01:00
committed by GitHub
parent 9120a64cfb
commit f1ee9113ac
83 changed files with 3293 additions and 1422 deletions
Download Patch File Download Diff File Expand all files Collapse all files

212
crates/nu-command/src/commands/dataframe/aggregate.rs
View File

@ -2,14 +2,11 @@ use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, PolarsData},
Signature, SyntaxShape, TaggedDictBuilder, UntaggedValue, Value,
dataframe::{Column, FrameStruct, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::Tagged;
use polars::{
frame::groupby::GroupBy,
prelude::{DataType, PolarsError, Series},
};
use polars::{frame::groupby::GroupBy, prelude::PolarsError};
enum Operation {
Mean,
@ -111,17 +108,40 @@ impl WholeStreamCommand for DataFrame {
description: "Aggregate sum by grouping by column a and summing on col b",
example:
"[[a b]; [one 1] [one 2]] | dataframe to-df | dataframe group-by a | dataframe aggregate sum",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("a".to_string(), vec![UntaggedValue::string("one").into()]),
Column::new("b".to_string(), vec![UntaggedValue::int(3).into()]),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "Aggregate sum in dataframe columns",
example: "[[a b]; [4 1] [5 2]] | dataframe to-df | dataframe aggregate sum",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("a".to_string(), vec![UntaggedValue::int(9).into()]),
Column::new("b".to_string(), vec![UntaggedValue::int(3).into()]),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "Aggregate sum in series",
example: "[4 1 5 6] | dataframe to-series | dataframe aggregate sum",
result: None,
example: "[4 1 5 6] | dataframe to-df | dataframe aggregate sum",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("0".to_string(), vec![UntaggedValue::int(16).into()]),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
]
}
@ -139,7 +159,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
})?;
match value.value {
UntaggedValue::DataFrame(PolarsData::GroupBy(nu_groupby)) => {
UntaggedValue::FrameStruct(FrameStruct::GroupBy(nu_groupby)) => {
let groupby = nu_groupby.to_groupby()?;
let res = perform_groupby_aggregation(
@ -152,18 +172,13 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, tag)))
}
UntaggedValue::DataFrame(PolarsData::EagerDataFrame(df)) => {
UntaggedValue::DataFrame(df) => {
let df = df.as_ref();
let res = perform_dataframe_aggregation(&df, op, &operation.tag)?;
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, tag)))
}
UntaggedValue::DataFrame(PolarsData::Series(series)) => {
let value = perform_series_aggregation(series.as_ref(), op, &operation.tag)?;
Ok(OutputStream::one(value))
}
_ => Err(ShellError::labeled_error(
"No groupby, dataframe or series in stream",
"no groupby, dataframe or series found in input stream",
@ -264,162 +279,15 @@ fn perform_dataframe_aggregation(
}
}
fn perform_series_aggregation(
series: &Series,
operation: Operation,
operation_tag: &Tag,
) -> Result<Value, ShellError> {
match operation {
Operation::Mean => {
let res = match series.mean() {
Some(val) => UntaggedValue::Primitive(val.into()),
None => UntaggedValue::Primitive(0.into()),
};
#[cfg(test)]
mod tests {
use super::DataFrame;
use super::ShellError;
let value = Value {
value: res,
tag: operation_tag.clone(),
};
#[test]
fn examples_work_as_expected() -> Result<(), ShellError> {
use crate::examples::test_dataframe as test_examples;
let mut data = TaggedDictBuilder::new(operation_tag.clone());
data.insert_value(series.name(), value);
Ok(data.into_value())
}
Operation::Median => {
let res = match series.median() {
Some(val) => UntaggedValue::Primitive(val.into()),
None => UntaggedValue::Primitive(0.into()),
};
let value = Value {
value: res,
tag: operation_tag.clone(),
};
let mut data = TaggedDictBuilder::new(operation_tag.clone());
data.insert_value(series.name(), value);
Ok(data.into_value())
}
Operation::Sum => {
let untagged = match series.dtype() {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => {
let res: i64 = series.sum().unwrap_or(0);
Ok(UntaggedValue::Primitive(res.into()))
}
DataType::Float32 | DataType::Float64 => {
let res: f64 = series.sum().unwrap_or(0.0);
Ok(UntaggedValue::Primitive(res.into()))
}
_ => Err(ShellError::labeled_error(
"Not valid type",
format!(
"this operation can not be performed with series of type {}",
series.dtype()
),
&operation_tag.span,
)),
}?;
let value = Value {
value: untagged,
tag: operation_tag.clone(),
};
let mut data = TaggedDictBuilder::new(operation_tag.clone());
data.insert_value(series.name(), value);
Ok(data.into_value())
}
Operation::Max => {
let untagged = match series.dtype() {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => {
let res: i64 = series.max().unwrap_or(0);
Ok(UntaggedValue::Primitive(res.into()))
}
DataType::Float32 | DataType::Float64 => {
let res: f64 = series.max().unwrap_or(0.0);
Ok(UntaggedValue::Primitive(res.into()))
}
_ => Err(ShellError::labeled_error(
"Not valid type",
format!(
"this operation can not be performed with series of type {}",
series.dtype()
),
&operation_tag.span,
)),
}?;
let value = Value {
value: untagged,
tag: operation_tag.clone(),
};
let mut data = TaggedDictBuilder::new(operation_tag.clone());
data.insert_value(series.name(), value);
Ok(data.into_value())
}
Operation::Min => {
let untagged = match series.dtype() {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => {
let res: i64 = series.min().unwrap_or(0);
Ok(UntaggedValue::Primitive(res.into()))
}
DataType::Float32 | DataType::Float64 => {
let res: f64 = series.min().unwrap_or(0.0);
Ok(UntaggedValue::Primitive(res.into()))
}
_ => Err(ShellError::labeled_error(
"Not valid type",
format!(
"this operation can not be performed with series of type {}",
series.dtype()
),
&operation_tag.span,
)),
}?;
let value = Value {
value: untagged,
tag: operation_tag.clone(),
};
let mut data = TaggedDictBuilder::new(operation_tag.clone());
data.insert_value(series.name(), value);
Ok(data.into_value())
}
_ => Err(ShellError::labeled_error_with_secondary(
"Not valid operation",
"operation not valid for series",
&operation_tag.span,
"Perhaps you want: mean, median, sum, max, min",
&operation_tag.span,
)),
test_examples(DataFrame {})
}
}

35
crates/nu-command/src/commands/dataframe/column.rs
View File

@ -2,8 +2,8 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, NuSeries},
Signature, SyntaxShape,
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::Tagged;
@ -32,7 +32,15 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "Returns the selected column as series",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe column a",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"a".to_string(),
vec![UntaggedValue::int(1).into(), UntaggedValue::int(3).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -41,15 +49,26 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let column: Tagged<String> = args.req(0)?;
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = df
.as_ref()
.column(column.item.as_ref())
.map_err(|e| parse_polars_error::<&str>(&e, &column.tag.span, None))?;
Ok(OutputStream::one(NuSeries::series_to_value(
res.clone(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.clone()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

30
crates/nu-command/src/commands/dataframe/drop.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, SyntaxShape, Value};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
use super::utils::{convert_columns, parse_polars_error};
@ -28,7 +31,15 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "drop column a",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe drop a",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![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())]),
}]
}
}
@ -39,7 +50,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let columns: Vec<Value> = args.rest(0)?;
let (col_string, col_span) = convert_columns(&columns, &tag)?;
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let new_df = match col_string.get(0) {
Some(col) => df
@ -63,3 +74,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, 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 {})
}
}

36
crates/nu-command/src/commands/dataframe/drop_duplicates.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, SyntaxShape, Value};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
use super::utils::{convert_columns, parse_polars_error};
@ -34,7 +37,21 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "drop duplicates",
example: "[[a b]; [1 2] [3 4] [1 2]] | dataframe to-df | dataframe drop-duplicates",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"a".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())]),
}]
}
}
@ -52,7 +69,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
None => (None, Span::unknown()),
};
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let subset_slice = subset.as_ref().map(|cols| &cols[..]);
@ -63,3 +80,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, 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 {})
}
}

61
crates/nu-command/src/commands/dataframe/drop_nulls.rs
View File

@ -2,7 +2,7 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, NuSeries, PolarsData},
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
@ -38,15 +38,45 @@ impl WholeStreamCommand for DataFrame {
example: r#"let df = ([[a b]; [1 2] [3 0] [1 2]] | dataframe to-df);
let res = ($df.b / $df.b);
let df = ($df | dataframe with-column $res --name res);
$df | dataframe drop-nulls
"#,
result: None,
$df | dataframe drop-nulls"#,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"a".to_string(),
vec![UntaggedValue::int(1).into(), UntaggedValue::int(1).into()],
),
Column::new(
"b".to_string(),
vec![UntaggedValue::int(2).into(), UntaggedValue::int(2).into()],
),
Column::new(
"res".to_string(),
vec![UntaggedValue::int(1).into(), UntaggedValue::int(1).into()],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "drop null values in dataframe",
example: r#"let s = ([1 2 0 0 3 4] | dataframe to-series);
example: r#"let s = ([1 2 0 0 3 4] | dataframe to-df);
($s / $s) | dataframe drop-nulls"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"div_0_0".to_string(),
vec![
UntaggedValue::int(1).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(1).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
]
}
@ -60,7 +90,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
})?;
match value.value {
UntaggedValue::DataFrame(PolarsData::EagerDataFrame(df)) => {
UntaggedValue::DataFrame(df) => {
// Extracting the selection columns of the columns to perform the aggregation
let columns: Option<Vec<Value>> = args.opt(0)?;
let (subset, col_span) = match columns {
@ -80,10 +110,6 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, tag)))
}
UntaggedValue::DataFrame(PolarsData::Series(series)) => {
let res = series.as_ref().drop_nulls();
Ok(OutputStream::one(NuSeries::series_to_value(res, tag)))
}
_ => Err(ShellError::labeled_error(
"Incorrect type",
"drop nulls cannot be done with this value",
@ -91,3 +117,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}
}
#[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 {})
}
}

82
crates/nu-command/src/commands/dataframe/dtypes.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, TaggedDictBuilder};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue, Value,
};
pub struct DataFrame;
@ -26,7 +29,27 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "drop column a",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe dtypes",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"column".to_string(),
vec![
UntaggedValue::string("a").into(),
UntaggedValue::string("b").into(),
],
),
Column::new(
"dtype".to_string(),
vec![
UntaggedValue::string("i64").into(),
UntaggedValue::string("i64").into(),
],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -35,26 +58,49 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let col_names = df
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let mut dtypes: Vec<Value> = Vec::new();
let names: Vec<Value> = df
.as_ref()
.get_column_names()
.iter()
.map(|v| v.to_string())
.collect::<Vec<String>>();
.map(|v| {
let dtype = df
.as_ref()
.column(v)
.expect("using name from list of names from dataframe")
.dtype();
let values = df
.as_ref()
.dtypes()
.into_iter()
.zip(col_names.into_iter())
.map(move |(dtype, name)| {
let mut data = TaggedDictBuilder::new(tag.clone());
data.insert_value("column", name.as_ref());
data.insert_value("dtype", format!("{}", dtype));
let dtype_str = format!("{}", dtype);
dtypes.push(Value {
value: dtype_str.into(),
tag: Tag::default(),
});
data.into_value()
});
Value {
value: v.to_string().into(),
tag: Tag::default(),
}
})
.collect();
Ok(OutputStream::from_stream(values))
let names_col = Column::new("column".to_string(), names);
let dtypes_col = Column::new("dtype".to_string(), dtypes);
let df = NuDataFrame::try_from_columns(vec![names_col, dtypes_col], &tag.span)?;
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 {})
}
}

92
crates/nu-command/src/commands/dataframe/dummies.rs
View File

@ -2,7 +2,7 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, PolarsData},
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
@ -32,12 +32,70 @@ impl WholeStreamCommand for DataFrame {
Example {
description: "Create new dataframe with dummy variables from a dataframe",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe to-dummies",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"a_1".to_string(),
vec![UntaggedValue::int(1).into(), UntaggedValue::int(0).into()],
),
Column::new(
"a_3".to_string(),
vec![UntaggedValue::int(0).into(), UntaggedValue::int(1).into()],
),
Column::new(
"b_2".to_string(),
vec![UntaggedValue::int(1).into(), UntaggedValue::int(0).into()],
),
Column::new(
"b_4".to_string(),
vec![UntaggedValue::int(0).into(), UntaggedValue::int(1).into()],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "Create new dataframe with dummy variables from a series",
example: "[1 2 2 3 3] | dataframe to-series | dataframe to-dummies",
result: None,
example: "[1 2 2 3 3] | dataframe to-df | dataframe to-dummies",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"0_1".to_string(),
vec![
UntaggedValue::int(1).into(),
UntaggedValue::int(0).into(),
UntaggedValue::int(0).into(),
UntaggedValue::int(0).into(),
UntaggedValue::int(0).into(),
],
),
Column::new(
"0_2".to_string(),
vec![
UntaggedValue::int(0).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(0).into(),
UntaggedValue::int(0).into(),
],
),
Column::new(
"0_3".to_string(),
vec![
UntaggedValue::int(0).into(),
UntaggedValue::int(0).into(),
UntaggedValue::int(0).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(1).into(),
],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
]
}
@ -51,7 +109,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
})?;
match value.value {
UntaggedValue::DataFrame(PolarsData::EagerDataFrame(df)) => {
UntaggedValue::DataFrame(df) => {
let res = df.as_ref().to_dummies().map_err(|e| {
parse_polars_error(
&e,
@ -62,17 +120,6 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, tag)))
}
UntaggedValue::DataFrame(PolarsData::Series(series)) => {
let res = series.as_ref().to_dummies().map_err(|e| {
parse_polars_error(
&e,
&tag.span,
Some("The only allowed column types for dummies are String or Int"),
)
})?;
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, tag)))
}
_ => Err(ShellError::labeled_error(
"Incorrect type",
"dummies cannot be done with this value",
@ -80,3 +127,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}
}
#[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 {})
}
}

41
crates/nu-command/src/commands/dataframe/filter.rs
View File

@ -2,7 +2,7 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, PolarsData},
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
@ -34,13 +34,21 @@ impl WholeStreamCommand for DataFrame {
vec![
Example {
description: "Filter dataframe using a bool mask",
example: r#"let mask = ([$true $false] | dataframe to-series);
example: r#"let mask = ([$true $false] | dataframe to-df);
[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe filter-with $mask"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("a".to_string(), vec![UntaggedValue::int(1).into()]),
Column::new("b".to_string(), vec![UntaggedValue::int(2).into()]),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "Filter dataframe by creating a mask from operation",
example: r#"let mask = (([5 6] | dataframe to-series) > 5);
example: r#"let mask = (([5 6] | dataframe to-df) > 5);
[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe filter-with $mask"#,
result: None,
},
@ -53,16 +61,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let value: Value = args.req(0)?;
let series_span = value.tag.span;
let series = match value.value {
UntaggedValue::DataFrame(PolarsData::Series(series)) => Ok(series),
let df = match value.value {
UntaggedValue::DataFrame(df) => Ok(df),
_ => Err(ShellError::labeled_error(
"Incorrect type",
"can only add a series to a dataframe",
value.tag.span,
)),
}?;
let casted = series.as_ref().bool().map_err(|e| {
let series = df.as_series(&series_span)?;
let casted = series.bool().map_err(|e| {
parse_polars_error(
&e,
&&series_span,
@ -70,12 +78,25 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)
})?;
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = df
.as_ref()
.filter(&casted)
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
.map_err(|e| parse_polars_error::<&str>(&e, &df_tag.span, None))?;
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, 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 {})
}
}

32
crates/nu-command/src/commands/dataframe/first.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, SyntaxShape};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::Tagged;
@ -31,8 +34,16 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Create new dataframe with head rows",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe first",
result: None,
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe first 1",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("a".to_string(), vec![UntaggedValue::int(1).into()]),
Column::new("b".to_string(), vec![UntaggedValue::int(2).into()]),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -46,8 +57,21 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
None => 5,
};
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = df.as_ref().head(Some(rows));
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, 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 {})
}
}

30
crates/nu-command/src/commands/dataframe/get.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, SyntaxShape, Value};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
use super::utils::{convert_columns, parse_polars_error};
pub struct DataFrame;
@ -27,7 +30,15 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "Creates dataframe with selected columns",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe get a",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"a".to_string(),
vec![UntaggedValue::int(1).into(), UntaggedValue::int(3).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -38,7 +49,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let (col_string, col_span) = convert_columns(&columns, &tag)?;
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = df
.as_ref()
@ -47,3 +58,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, 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 {})
}
}

6
crates/nu-command/src/commands/dataframe/groupby.rs
View File

@ -2,7 +2,7 @@ use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, NuGroupBy, PolarsData},
dataframe::{FrameStruct, NuDataFrame, NuGroupBy},
Signature, SyntaxShape, UntaggedValue, Value,
};
@ -43,7 +43,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let by_columns: Vec<Value> = args.rest(0)?;
let (columns_string, col_span) = convert_columns(&by_columns, &tag)?;
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
// This is the expensive part of the groupby; to create the
// groups that will be used for grouping the data in the
@ -57,7 +57,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let groups = groupby.get_groups().to_vec();
let groupby = Value {
tag,
value: UntaggedValue::DataFrame(PolarsData::GroupBy(NuGroupBy::new(
value: UntaggedValue::FrameStruct(FrameStruct::GroupBy(NuGroupBy::new(
NuDataFrame::new(df.as_ref().clone()),
columns_string,
groups,

75
crates/nu-command/src/commands/dataframe/join.rs
View File

@ -2,7 +2,7 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, PolarsData},
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
@ -51,20 +51,50 @@ impl WholeStreamCommand for DataFrame {
}
fn examples(&self) -> Vec<Example> {
vec![
Example {
description: "inner join dataframe",
example: r#"let right = ([[a b c]; [1 2 5] [3 4 5] [5 6 6]] | dataframe to-df);
vec![Example {
description: "inner join dataframe",
example: r#"let right = ([[a b c]; [1 2 5] [3 4 5] [5 6 6]] | dataframe to-df);
$right | dataframe join $right -l [a b] -r [a b]"#,
result: None,
},
Example {
description: "right join dataframe",
example: r#"let right = ([[a b c]; [1 2 3] [3 4 5] [5 6 7]] | dataframe to-df);
$right | dataframe join $right -l [a c] -r [a c] -t inner"#,
result: None,
},
]
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"a".to_string(),
vec![
UntaggedValue::int(1).into(),
UntaggedValue::int(3).into(),
UntaggedValue::int(5).into(),
],
),
Column::new(
"b".to_string(),
vec![
UntaggedValue::int(2).into(),
UntaggedValue::int(4).into(),
UntaggedValue::int(6).into(),
],
),
Column::new(
"c".to_string(),
vec![
UntaggedValue::int(5).into(),
UntaggedValue::int(5).into(),
UntaggedValue::int(6).into(),
],
),
Column::new(
"c_right".to_string(),
vec![
UntaggedValue::int(5).into(),
UntaggedValue::int(5).into(),
UntaggedValue::int(6).into(),
],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -97,10 +127,10 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let (l_col_string, l_col_span) = convert_columns(&l_col, &tag)?;
let (r_col_string, r_col_span) = convert_columns(&r_col, &tag)?;
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = match r_df.value {
UntaggedValue::DataFrame(PolarsData::EagerDataFrame(r_df)) => {
UntaggedValue::DataFrame(r_df) => {
// Checking the column types before performing the join
check_column_datatypes(
df.as_ref(),
@ -173,3 +203,16 @@ fn check_column_datatypes<T: AsRef<str>>(
Ok(())
}
#[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 {})
}
}

32
crates/nu-command/src/commands/dataframe/last.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, SyntaxShape};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::Tagged;
pub struct DataFrame;
@ -30,8 +33,16 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Create new dataframe with last rows",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe last",
result: None,
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe last 1",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("a".to_string(), vec![UntaggedValue::int(3).into()]),
Column::new("b".to_string(), vec![UntaggedValue::int(4).into()]),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -45,9 +56,22 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
None => 5,
};
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = df.as_ref().tail(Some(rows));
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, 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 {})
}
}

95
crates/nu-command/src/commands/dataframe/list.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::PolarsData, Signature, TaggedDictBuilder, UntaggedValue};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue, Value,
};
pub struct DataFrame;
@ -19,46 +22,94 @@ impl WholeStreamCommand for DataFrame {
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
let values = args
let data = args
.context
.scope
.get_vars()
.into_iter()
.filter_map(|(name, value)| {
if let UntaggedValue::DataFrame(PolarsData::EagerDataFrame(df)) = &value.value {
let mut data = TaggedDictBuilder::new(value.tag.clone());
if let UntaggedValue::DataFrame(df) = &value.value {
let rows = Value {
value: (df.as_ref().height() as i64).into(),
tag: Tag::default(),
};
let rows = df.as_ref().height();
let cols = df.as_ref().width();
let cols = Value {
value: (df.as_ref().width() as i64).into(),
tag: Tag::default(),
};
data.insert_value("name", name.as_ref());
data.insert_value("rows", format!("{}", rows));
data.insert_value("columns", format!("{}", cols));
let location = match value.tag.anchor {
Some(AnchorLocation::File(name)) => name,
Some(AnchorLocation::Url(name)) => name,
Some(AnchorLocation::Source(text)) => text.slice(0..text.end).text,
None => "stream".to_string(),
};
match value.tag.anchor {
Some(AnchorLocation::File(name)) => data.insert_value("location", name),
Some(AnchorLocation::Url(name)) => data.insert_value("location", name),
Some(AnchorLocation::Source(text)) => {
let loc_name = text.slice(0..text.end);
data.insert_value("location", loc_name.text)
}
None => data.insert_value("location", "stream"),
}
let location = Value {
value: location.into(),
tag: Tag::default(),
};
Some(data.into_value())
let name = Value {
value: name.into(),
tag: Tag::default(),
};
Some((name, rows, cols, location))
} else {
None
}
});
Ok(OutputStream::from_stream(values))
let mut name = Column::new_empty("name".to_string());
let mut rows = Column::new_empty("rows".to_string());
let mut cols = Column::new_empty("columns".to_string());
let mut location = Column::new_empty("location".to_string());
for tuple in data {
name.push(tuple.0);
rows.push(tuple.1);
cols.push(tuple.2);
location.push(tuple.3);
}
let tag = args.call_info.name_tag;
let df = NuDataFrame::try_from_columns(vec![name, rows, cols, location], &tag.span)?;
Ok(OutputStream::one(df.into_value(tag)))
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Lists loaded dataframes in current scope",
example: "dataframe list",
result: None,
example: "let a = ([[a b];[1 2] [3 4]] | dataframe to-df); dataframe list",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("name".to_string(), vec![UntaggedValue::string("$a").into()]),
Column::new("rows".to_string(), vec![UntaggedValue::int(2).into()]),
Column::new("columns".to_string(), vec![UntaggedValue::int(2).into()]),
Column::new(
"location".to_string(),
vec![UntaggedValue::string("stream").into()],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
#[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 {})
}
}

93
crates/nu-command/src/commands/dataframe/melt.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, SyntaxShape, Value};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
use super::utils::convert_columns;
@ -18,8 +21,18 @@ impl WholeStreamCommand for DataFrame {
fn signature(&self) -> Signature {
Signature::build("dataframe melt")
.required("id_columns", SyntaxShape::Table, "Id columns for melting")
.rest(SyntaxShape::Any, "columns used as value columns")
.required_named(
"columns",
SyntaxShape::Table,
"column names for melting",
Some('c'),
)
.required_named(
"values",
SyntaxShape::Table,
"column names used as value columns",
Some('v'),
)
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
@ -29,8 +42,59 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "melt dataframe",
example: "[[a b]; [a 2] [b 4] [a 6]] | dataframe to-df | dataframe melt a b",
result: None,
example:
"[[a b c d]; [x 1 4 a] [y 2 5 b] [z 3 6 c]] | dataframe to-df | dataframe melt -c [b c] -v [a d]",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"b".to_string(),
vec![
UntaggedValue::int(1).into(),
UntaggedValue::int(2).into(),
UntaggedValue::int(3).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(2).into(),
UntaggedValue::int(3).into(),
],
),
Column::new(
"c".to_string(),
vec![
UntaggedValue::int(4).into(),
UntaggedValue::int(5).into(),
UntaggedValue::int(6).into(),
UntaggedValue::int(4).into(),
UntaggedValue::int(5).into(),
UntaggedValue::int(6).into(),
],
),
Column::new(
"variable".to_string(),
vec![
UntaggedValue::string("a").into(),
UntaggedValue::string("a").into(),
UntaggedValue::string("a").into(),
UntaggedValue::string("d").into(),
UntaggedValue::string("d").into(),
UntaggedValue::string("d").into(),
],
),
Column::new(
"value".to_string(),
vec![
UntaggedValue::string("x").into(),
UntaggedValue::string("y").into(),
UntaggedValue::string("z").into(),
UntaggedValue::string("a").into(),
UntaggedValue::string("b").into(),
UntaggedValue::string("c").into(),
],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -38,13 +102,13 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let id_col: Vec<Value> = args.req(0)?;
let val_col: Vec<Value> = args.rest(1)?;
let id_col: Vec<Value> = args.req_named("columns")?;
let val_col: Vec<Value> = args.req_named("values")?;
let (id_col_string, id_col_span) = convert_columns(&id_col, &tag)?;
let (val_col_string, val_col_span) = convert_columns(&val_col, &tag)?;
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
check_column_datatypes(df.as_ref(), &id_col_string, &id_col_span)?;
check_column_datatypes(df.as_ref(), &val_col_string, &val_col_span)?;
@ -99,3 +163,16 @@ fn check_column_datatypes<T: AsRef<str>>(
Ok(())
}
#[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 {})
}
}

9
crates/nu-command/src/commands/dataframe/mod.rs
View File

@ -26,7 +26,6 @@ pub mod take;
pub mod to_csv;
pub mod to_df;
pub mod to_parquet;
pub mod to_series;
pub(crate) mod utils;
pub mod where_;
pub mod with_column;
@ -59,7 +58,6 @@ pub use take::DataFrame as DataFrameTake;
pub use to_csv::DataFrame as DataFrameToCsv;
pub use to_df::DataFrame as DataFrameToDF;
pub use to_parquet::DataFrame as DataFrameToParquet;
pub use to_series::DataFrame as DataFrameToSeries;
pub use where_::DataFrame as DataFrameWhere;
pub use with_column::DataFrame as DataFrameWithColumn;
@ -71,6 +69,8 @@ pub use series::DataFrameArgMin;
pub use series::DataFrameArgSort;
pub use series::DataFrameArgTrue;
pub use series::DataFrameArgUnique;
pub use series::DataFrameConcatenate;
pub use series::DataFrameContains;
pub use series::DataFrameIsDuplicated;
pub use series::DataFrameIsIn;
pub use series::DataFrameIsNotNull;
@ -80,9 +80,14 @@ pub use series::DataFrameNNull;
pub use series::DataFrameNUnique;
pub use series::DataFrameNot;
pub use series::DataFrameReplace;
pub use series::DataFrameReplaceAll;
pub use series::DataFrameSeriesRename;
pub use series::DataFrameSet;
pub use series::DataFrameSetWithIdx;
pub use series::DataFrameShift;
pub use series::DataFrameStringLengths;
pub use series::DataFrameStringSlice;
pub use series::DataFrameToLowercase;
pub use series::DataFrameToUppercase;
pub use series::DataFrameUnique;
pub use series::DataFrameValueCounts;

2
crates/nu-command/src/commands/dataframe/pivot.rs
View File

@ -73,7 +73,7 @@ impl WholeStreamCommand for DataFrame {
description: "Pivot a dataframe on b and aggregation on col c",
example:
"[[a b c]; [one x 1] [two y 2]] | dataframe to-df | dataframe group-by a | dataframe pivot b c sum",
result: None,
result: None, // No sample because there are nulls in the result dataframe
}]
}
}

6
crates/nu-command/src/commands/dataframe/sample.rs
View File

@ -42,13 +42,13 @@ impl WholeStreamCommand for DataFrame {
Example {
description: "Sample rows from dataframe",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe sample -r 1",
result: None,
result: None, // No expected value because sampling is random
},
Example {
description: "Shows sample row using fraction and replace",
example:
"[[a b]; [1 2] [3 4] [5 6]] | dataframe to-df | dataframe sample -f 0.5 -e",
result: None,
result: None, // No expected value because sampling is random
},
]
}
@ -61,7 +61,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let fraction: Option<Tagged<f64>> = args.get_flag("fraction")?;
let replace: bool = args.has_flag("replace");
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = match (rows, fraction) {
(Some(rows), None) => df

30
crates/nu-command/src/commands/dataframe/select.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, SyntaxShape, Value};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
use super::utils::{convert_columns, parse_polars_error};
@ -28,7 +31,15 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "Create new dataframe with column a",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe select a",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"a".to_string(),
vec![UntaggedValue::int(1).into(), UntaggedValue::int(3).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -40,7 +51,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let (col_string, col_span) = convert_columns(&columns, &tag)?;
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = df
.as_ref()
@ -49,3 +60,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, 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 {})
}
}

55
crates/nu-command/src/commands/dataframe/series/all_false.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature, TaggedDictBuilder, UntaggedValue, Value};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue, Value,
};
pub struct DataFrame;
@ -26,15 +29,31 @@ impl WholeStreamCommand for DataFrame {
vec![
Example {
description: "Returns true if all values are false",
example: "[$false $false $false] | dataframe to-series | dataframe all-false",
result: None,
example: "[$false $false $false] | dataframe to-df | dataframe all-false",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"all_false".to_string(),
vec![UntaggedValue::boolean(true).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "Checks the result from a comparison",
example: r#"let s = ([5 6 2 8] | dataframe to-series);
example: r#"let s = ([5 6 2 10] | dataframe to-df);
let res = ($s > 9);
$res | dataframe all-false"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"all_false".to_string(),
vec![UntaggedValue::boolean(false).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
]
}
@ -43,9 +62,10 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let bool = series.as_ref().bool().map_err(|e| {
let series = df.as_series(&df_tag.span)?;
let bool = series.bool().map_err(|e| {
parse_polars_error::<&str>(
&e,
&tag.span,
@ -60,8 +80,23 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
tag: tag.clone(),
};
let mut data = TaggedDictBuilder::new(tag);
data.insert_value("all_false", value);
let df = NuDataFrame::try_from_columns(
vec![Column::new("all_false".to_string(), vec![value])],
&tag.span,
)?;
Ok(OutputStream::one(data.into_value()))
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 {})
}
}

55
crates/nu-command/src/commands/dataframe/series/all_true.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature, TaggedDictBuilder, UntaggedValue, Value};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue, Value,
};
pub struct DataFrame;
@ -26,15 +29,31 @@ impl WholeStreamCommand for DataFrame {
vec![
Example {
description: "Returns true if all values are true",
example: "[$true $true $true] | dataframe to-series | dataframe all-true",
result: None,
example: "[$true $true $true] | dataframe to-df | dataframe all-true",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"all_true".to_string(),
vec![UntaggedValue::boolean(true).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "Checks the result from a comparison",
example: r#"let s = ([5 6 2 8] | dataframe to-series);
example: r#"let s = ([5 6 2 8] | dataframe to-df);
let res = ($s > 9);
$res | dataframe all-true"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"all_true".to_string(),
vec![UntaggedValue::boolean(false).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
]
}
@ -43,9 +62,10 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let bool = series.as_ref().bool().map_err(|e| {
let series = df.as_series(&df_tag.span)?;
let bool = series.bool().map_err(|e| {
parse_polars_error::<&str>(
&e,
&tag.span,
@ -60,8 +80,23 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
tag: tag.clone(),
};
let mut data = TaggedDictBuilder::new(tag);
data.insert_value("all_true", value);
let df = NuDataFrame::try_from_columns(
vec![Column::new("all_true".to_string(), vec![value])],
&tag.span,
)?;
Ok(OutputStream::one(data.into_value()))
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 {})
}
}

38
crates/nu-command/src/commands/dataframe/series/arg_max.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::{IntoSeries, NewChunkedArray, UInt32Chunked};
@ -27,8 +30,16 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Returns index for max value",
example: "[1 3 2] | dataframe to-series | dataframe arg-max",
result: None,
example: "[1 3 2] | dataframe to-df | dataframe arg-max",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"arg_max".to_string(),
vec![UntaggedValue::int(1).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -36,9 +47,10 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
let res = series.as_ref().arg_max();
let res = series.arg_max();
let chunked = match res {
Some(index) => UInt32Chunked::new_from_slice("arg_max", &[index as u32]),
@ -46,6 +58,20 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
};
let res = chunked.into_series();
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(NuSeries::series_to_value(res, tag)))
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

37
crates/nu-command/src/commands/dataframe/series/arg_min.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::{IntoSeries, NewChunkedArray, UInt32Chunked};
@ -27,8 +30,16 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Returns index for min value",
example: "[1 3 2] | dataframe to-series | dataframe arg-min",
result: None,
example: "[1 3 2] | dataframe to-df | dataframe arg-min",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"arg_min".to_string(),
vec![UntaggedValue::int(0).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -36,9 +47,9 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series.as_ref().arg_min();
let res = df.as_series(&df_tag.span)?.arg_min();
let chunked = match res {
Some(index) => UInt32Chunked::new_from_slice("arg_min", &[index as u32]),
@ -46,6 +57,20 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
};
let res = chunked.into_series();
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(NuSeries::series_to_value(res, tag)))
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

47
crates/nu-command/src/commands/dataframe/series/arg_sort.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
@ -26,8 +29,22 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Returns indexes for a sorted series",
example: "[1 2 2 3 3] | dataframe to-series | dataframe arg-sort",
result: None,
example: "[1 2 2 3 3] | dataframe to-df | dataframe arg-sort",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"arg_sort".to_string(),
vec![
UntaggedValue::int(0).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(2).into(),
UntaggedValue::int(3).into(),
UntaggedValue::int(4).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -36,12 +53,24 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let reverse = args.has_flag("reverse");
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series.as_ref().argsort(reverse);
let mut res = df.as_series(&df_tag.span)?.argsort(reverse).into_series();
res.rename("arg_sort");
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

40
crates/nu-command/src/commands/dataframe/series/arg_true.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
@ -26,8 +29,16 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Returns indexes where values are true",
example: "[$false $true $false] | dataframe to-series | dataframe arg-true",
result: None,
example: "[$false $true $false] | dataframe to-df | dataframe arg-true",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"arg_true".to_string(),
vec![UntaggedValue::int(1).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -35,9 +46,10 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let bool = series.as_ref().bool().map_err(|e| {
let series = df.as_series(&df_tag.span)?;
let bool = series.bool().map_err(|e| {
parse_polars_error::<&str>(
&e,
&tag.span,
@ -46,7 +58,21 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
})?;
let mut res = bool.arg_true().into_series();
res.rename("int");
res.rename("arg_true");
Ok(OutputStream::one(NuSeries::series_to_value(res, tag)))
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

51
crates/nu-command/src/commands/dataframe/series/arg_unique.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
@ -26,8 +29,20 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Returns indexes for unique values",
example: "[1 2 2 3 3] | dataframe to-series | dataframe arg-unique",
result: None,
example: "[1 2 2 3 3] | dataframe to-df | dataframe arg-unique",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"arg_unique".to_string(),
vec![
UntaggedValue::int(0).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(3).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -35,15 +50,29 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series
.as_ref()
let mut res = df
.as_series(&df_tag.span)?
.arg_unique()
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?
.into_series();
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
res.rename("arg_unique");
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

107
crates/nu-command/src/commands/dataframe/series/concatenate.rs Normal file
View File

@ -0,0 +1,107 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
impl WholeStreamCommand for DataFrame {
fn name(&self) -> &str {
"dataframe concatenate"
}
fn usage(&self) -> &str {
"[Series] Concatenates strings with other array"
}
fn signature(&self) -> Signature {
Signature::build("dataframe concatenate").required(
"other",
SyntaxShape::Any,
"Other array with string to be concatenated",
)
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
command(args)
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Concatenate string",
example: r#"let other = ([za xs cd] | dataframe to-df);
[abc abc abc] | dataframe to-df | dataframe concatenate $other"#,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::string("abcza").into(),
UntaggedValue::string("abcxs").into(),
UntaggedValue::string("abccd").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 other: Value = args.req(0)?;
let other_df = match &other.value {
UntaggedValue::DataFrame(df) => Ok(df),
_ => Err(ShellError::labeled_error(
"Incorrect type",
"can only concatenate another series",
other.tag.span,
)),
}?;
let other_series = other_df.as_series(&other.tag.span)?;
let other_chunked = other_series.utf8().map_err(|e| {
parse_polars_error::<&str>(
&e,
&other.tag.span,
Some("The concatenate command can only be used with string columns"),
)
})?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
let chunked = series.utf8().map_err(|e| {
parse_polars_error::<&str>(
&e,
&df_tag.span,
Some("The concatenate command can only be used with string columns"),
)
})?;
let mut res = chunked.concat(&other_chunked);
res.rename(series.name());
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

90
crates/nu-command/src/commands/dataframe/series/contains.rs Normal file
View File

@ -0,0 +1,90 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::Tagged;
use polars::prelude::IntoSeries;
pub struct DataFrame;
impl WholeStreamCommand for DataFrame {
fn name(&self) -> &str {
"dataframe contains"
}
fn usage(&self) -> &str {
"[Series] Checks if a patter is contained in a string"
}
fn signature(&self) -> Signature {
Signature::build("dataframe contains").required_named(
"pattern",
SyntaxShape::String,
"Regex pattern to be searched",
Some('p'),
)
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
command(args)
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Returns boolean indicating if patter was found",
example: "[abc acb acb] | dataframe to-df | dataframe contains -p ab",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(false).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 pattern: Tagged<String> = args.req_named("pattern")?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
let chunked = series.utf8().map_err(|e| {
parse_polars_error::<&str>(
&e,
&df_tag.span,
Some("The contains command can only be used with string columns"),
)
})?;
let res = chunked
.contains(pattern.as_str())
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

55
crates/nu-command/src/commands/dataframe/series/is_duplicated.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
@ -26,8 +29,24 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Create mask indicating duplicated values",
example: "[5 6 6 6 8 8 8] | dataframe to-series | dataframe is-duplicated",
result: None,
example: "[5 6 6 6 8 8 8] | dataframe to-df | dataframe is-duplicated",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"is_duplicated".to_string(),
vec![
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(true).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -35,15 +54,29 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series
.as_ref()
let mut res = df
.as_series(&df_tag.span)?
.is_duplicated()
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?
.into_series();
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
res.rename("is_duplicated");
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

62
crates/nu-command/src/commands/dataframe/series/is_in.rs
View File

@ -2,7 +2,7 @@ use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuSeries, PolarsData},
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
use polars::prelude::IntoSeries;
@ -29,9 +29,25 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Checks if elements from a series are contained in right series",
example: r#"let other = ([1 3 6] | dataframe to-series);
[5 6 6 6 8 8 8] | dataframe to-series | dataframe is-in $other"#,
result: None,
example: r#"let other = ([1 3 6] | dataframe to-df);
[5 6 6 6 8 8 8] | dataframe to-df | dataframe is-in $other"#,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"is_in".to_string(),
vec![
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(false).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -40,8 +56,8 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let value: Value = args.req(0)?;
let other = match value.value {
UntaggedValue::DataFrame(PolarsData::Series(series)) => Ok(series),
let other_df = match value.value {
UntaggedValue::DataFrame(df) => Ok(df),
_ => Err(ShellError::labeled_error(
"Incorrect type",
"can only search in a series",
@ -49,15 +65,31 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}?;
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let other = other_df.as_series(&value.tag.span)?;
let res = series
.as_ref()
.is_in(other.as_ref())
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let mut res = df
.as_series(&df_tag.span)?
.is_in(&other)
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?
.into_series();
res.rename("is_in");
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

45
crates/nu-command/src/commands/dataframe/series/is_not_null.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
@ -26,10 +29,23 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Create mask where values are not null",
example: r#"let s = ([5 6 0 8] | dataframe to-series);
example: r#"let s = ([5 6 0 8] | dataframe to-df);
let res = ($s / $s);
$res | dataframe is-not-null"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"is_not_null".to_string(),
vec![
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(true).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -37,12 +53,23 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series.as_ref().is_not_null();
let res = df.as_series(&df_tag.span)?.is_not_null();
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

45
crates/nu-command/src/commands/dataframe/series/is_null.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
@ -26,10 +29,23 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Create mask where values are null",
example: r#"let s = ([5 6 0 8] | dataframe to-series);
example: r#"let s = ([5 6 0 8] | dataframe to-df);
let res = ($s / $s);
$res | dataframe is-null"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"is_null".to_string(),
vec![
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(false).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -37,12 +53,23 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series.as_ref().is_null();
let res = df.as_series(&df_tag.span)?.is_null();
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

55
crates/nu-command/src/commands/dataframe/series/is_unique.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
@ -26,8 +29,24 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Create mask indicating unique values",
example: "[5 6 6 6 8 8 8] | dataframe to-series | dataframe is-unique",
result: None,
example: "[5 6 6 6 8 8 8] | dataframe to-df | dataframe is-unique",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"is_unique".to_string(),
vec![
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(false).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -35,15 +54,29 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series
.as_ref()
let mut res = df
.as_series(&df_tag.span)?
.is_unique()
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?
.into_series();
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
res.rename("is_unique");
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

14
crates/nu-command/src/commands/dataframe/series/mod.rs
View File

@ -5,6 +5,8 @@ pub mod arg_min;
pub mod arg_sort;
pub mod arg_true;
pub mod arg_unique;
pub mod concatenate;
pub mod contains;
pub mod is_duplicated;
pub mod is_in;
pub mod is_not_null;
@ -15,9 +17,14 @@ pub mod n_unique;
pub mod not;
pub mod rename;
pub mod replace;
pub mod replace_all;
pub mod set;
pub mod set_with_idx;
pub mod shift;
pub mod str_lengths;
pub mod str_slice;
pub mod to_lowercase;
pub mod to_uppercase;
pub mod unique;
pub mod value_counts;
@ -28,6 +35,8 @@ pub use arg_min::DataFrame as DataFrameArgMin;
pub use arg_sort::DataFrame as DataFrameArgSort;
pub use arg_true::DataFrame as DataFrameArgTrue;
pub use arg_unique::DataFrame as DataFrameArgUnique;
pub use concatenate::DataFrame as DataFrameConcatenate;
pub use contains::DataFrame as DataFrameContains;
pub use is_duplicated::DataFrame as DataFrameIsDuplicated;
pub use is_in::DataFrame as DataFrameIsIn;
pub use is_not_null::DataFrame as DataFrameIsNotNull;
@ -38,8 +47,13 @@ pub use n_unique::DataFrame as DataFrameNUnique;
pub use not::DataFrame as DataFrameNot;
pub use rename::DataFrame as DataFrameSeriesRename;
pub use replace::DataFrame as DataFrameReplace;
pub use replace_all::DataFrame as DataFrameReplaceAll;
pub use set::DataFrame as DataFrameSet;
pub use set_with_idx::DataFrame as DataFrameSetWithIdx;
pub use shift::DataFrame as DataFrameShift;
pub use str_lengths::DataFrame as DataFrameStringLengths;
pub use str_slice::DataFrame as DataFrameStringSlice;
pub use to_lowercase::DataFrame as DataFrameToLowercase;
pub use to_uppercase::DataFrame as DataFrameToUppercase;
pub use unique::DataFrame as DataFrameUnique;
pub use value_counts::DataFrame as DataFrameValueCounts;

42
crates/nu-command/src/commands/dataframe/series/n_null.rs
View File

@ -2,7 +2,8 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::NuSeries, Primitive, Signature, TaggedDictBuilder, UntaggedValue, Value,
dataframe::{Column, NuDataFrame},
Primitive, Signature, UntaggedValue, Value,
};
pub struct DataFrame;
@ -27,9 +28,17 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Counts null values",
example: r#"let s = ([1 1 0 0 3 3 4] | dataframe to-series);
($s / ss) | dataframe count-null"#,
result: None,
example: r#"let s = ([1 1 0 0 3 3 4] | dataframe to-df);
($s / $s) | dataframe count-null"#,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"count_null".to_string(),
vec![UntaggedValue::int(2).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -37,17 +46,32 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series.as_ref().null_count();
let res = df.as_series(&df_tag.span)?.null_count();
let value = Value {
value: UntaggedValue::Primitive(Primitive::Int(res as i64)),
tag: tag.clone(),
};
let mut data = TaggedDictBuilder::new(tag);
data.insert_value("count-null", value);
let df = NuDataFrame::try_from_columns(
vec![Column::new("count_null".to_string(), vec![value])],
&tag.span,
)?;
Ok(OutputStream::one(data.into_value()))
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 {})
}
}

42
crates/nu-command/src/commands/dataframe/series/n_unique.rs
View File

@ -2,7 +2,8 @@ use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::NuSeries, Primitive, Signature, TaggedDictBuilder, UntaggedValue, Value,
dataframe::{Column, NuDataFrame},
Primitive, Signature, UntaggedValue, Value,
};
pub struct DataFrame;
@ -27,8 +28,16 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Counts unique values",
example: "[1 1 2 2 3 3 4] | dataframe to-series | dataframe count-unique",
result: None,
example: "[1 1 2 2 3 3 4] | dataframe to-df | dataframe count-unique",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"count_unique".to_string(),
vec![UntaggedValue::int(4).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -36,10 +45,10 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series
.as_ref()
let res = df
.as_series(&df_tag.span)?
.n_unique()
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
@ -48,8 +57,23 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
tag: tag.clone(),
};
let mut data = TaggedDictBuilder::new(tag);
data.insert_value("count-unique", value);
let df = NuDataFrame::try_from_columns(
vec![Column::new("count_unique".to_string(), vec![value])],
&tag.span,
)?;
Ok(OutputStream::one(data.into_value()))
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 {})
}
}

45
crates/nu-command/src/commands/dataframe/series/not.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
use std::ops::Not;
@ -27,8 +30,20 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Inverts boolean mask",
example: "[$true $false $true] | dataframe to-series | dataframe not",
result: None,
example: "[$true $false $true] | dataframe to-df | dataframe not",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::boolean(false).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(false).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -36,9 +51,10 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
let bool = series.as_ref().bool().map_err(|e| {
let bool = series.bool().map_err(|e| {
parse_polars_error::<&str>(
&e,
&tag.span,
@ -48,8 +64,19 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let res = bool.not();
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

44
crates/nu-command/src/commands/dataframe/series/rename.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature, SyntaxShape};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::Tagged;
pub struct DataFrame;
@ -30,8 +33,21 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Renames a series",
example: "[5 6 7 8] | dataframe to-series | dataframe rename-series new_name",
result: None,
example: "[5 6 7 8] | dataframe to-df | dataframe rename new_name",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"new_name".to_string(),
vec![
UntaggedValue::int(5).into(),
UntaggedValue::int(6).into(),
UntaggedValue::int(7).into(),
UntaggedValue::int(8).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -40,9 +56,25 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let name: Tagged<String> = args.req(0)?;
let mut series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
series.as_mut().rename(name.item.as_ref());
let mut series = df.as_series(&df_tag.span)?;
Ok(OutputStream::one(series.into_value(tag)))
series.rename(name.item.as_ref());
let df = NuDataFrame::try_from_series(vec![series], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

56
crates/nu-command/src/commands/dataframe/series/replace.rs
View File

@ -1,8 +1,11 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature, SyntaxShape};
use nu_source::Tagged;
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::{Span, Tagged};
use polars::prelude::IntoSeries;
pub struct DataFrame;
@ -39,8 +42,20 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Replaces string",
example: "[abc abc abc] | dataframe to-series | dataframe replace -p ab -r AB",
result: None,
example: "[abc abc abc] | dataframe to-df | dataframe replace -p ab -r AB",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::string("ABc").into(),
UntaggedValue::string("ABc").into(),
UntaggedValue::string("ABc").into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -50,23 +65,36 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let pattern: Tagged<String> = args.req_named("pattern")?;
let replace: Tagged<String> = args.req_named("replace")?;
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let chunked = series.as_ref().utf8().map_err(|e| {
let series = df.as_series(&df_tag.span)?;
let chunked = series.utf8().map_err(|e| {
parse_polars_error::<&str>(
&e,
&tag.span,
Some("The replace command can only be used with string columns"),
&df_tag.span,
Some("The replace-all command can only be used with string columns"),
)
})?;
let res = chunked
.as_ref()
let mut res = chunked
.replace(pattern.as_str(), replace.as_str())
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
res.rename(series.name());
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

100
crates/nu-command/src/commands/dataframe/series/replace_all.rs Normal file
View File

@ -0,0 +1,100 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::Tagged;
use polars::prelude::IntoSeries;
pub struct DataFrame;
impl WholeStreamCommand for DataFrame {
fn name(&self) -> &str {
"dataframe replace-all"
}
fn usage(&self) -> &str {
"[Series] Replace all (sub)strings by a regex pattern"
}
fn signature(&self) -> Signature {
Signature::build("dataframe replace")
.required_named(
"pattern",
SyntaxShape::String,
"Regex pattern to be matched",
Some('p'),
)
.required_named(
"replace",
SyntaxShape::String,
"replacing string",
Some('r'),
)
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
command(args)
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Replaces string",
example: "[abac abac abac] | dataframe to-df | dataframe replace-all -p a -r A",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::string("AbAc").into(),
UntaggedValue::string("AbAc").into(),
UntaggedValue::string("AbAc").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 pattern: Tagged<String> = args.req_named("pattern")?;
let replace: Tagged<String> = args.req_named("replace")?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
let chunked = series.utf8().map_err(|e| {
parse_polars_error::<&str>(
&e,
&df_tag.span,
Some("The replace command can only be used with string columns"),
)
})?;
let mut res = chunked
.replace_all(pattern.as_str(), replace.as_str())
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
res.rename(series.name());
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

90
crates/nu-command/src/commands/dataframe/series/set.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Primitive, Signature, SyntaxShape, UntaggedValue, Value};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Primitive, Signature, SyntaxShape, UntaggedValue, Value,
};
use polars::prelude::{ChunkSet, DataType, IntoSeries};
pub struct DataFrame;
@ -33,10 +36,24 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Shifts the values by a given period",
example: r#"let s = ([1 2 2 3 3] | dataframe to-series | dataframe shift 2);
example: r#"let s = ([1 2 2 3 3] | dataframe to-df | dataframe shift 2);
let mask = ($s | dataframe is-null);
$s | dataframe set 0 --mask $mask"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::int(0).into(),
UntaggedValue::int(0).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(2).into(),
UntaggedValue::int(2).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -46,20 +63,21 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let value: Value = args.req(0)?;
let mask: Value = args.req_named("mask")?;
let bool_mask = match &mask.value {
UntaggedValue::DataFrame(nu_protocol::dataframe::PolarsData::Series(series)) => {
match series.as_ref().dtype() {
DataType::Boolean => series
.as_ref()
.bool()
.map_err(|e| parse_polars_error::<&str>(&e, &mask.tag.span, None)),
_ => Err(ShellError::labeled_error(
"Incorrect type",
"can only use bool series as mask",
value.tag.span,
)),
}
}
let mask_df = match &mask.value {
UntaggedValue::DataFrame(df) => Ok(df),
_ => Err(ShellError::labeled_error(
"Incorrect type",
"can only use a series as mask",
value.tag.span,
)),
}?;
let mask_series = mask_df.as_series(&mask.tag.span)?;
let bool_mask = match mask_series.dtype() {
DataType::Boolean => mask_series
.bool()
.map_err(|e| parse_polars_error::<&str>(&e, &mask.tag.span, None)),
_ => Err(ShellError::labeled_error(
"Incorrect type",
"can only use bool series as mask",
@ -67,11 +85,12 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}?;
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
match &value.value {
UntaggedValue::Primitive(Primitive::Int(val)) => {
let chunked = series.as_ref().i64().map_err(|e| {
let chunked = series.i64().map_err(|e| {
parse_polars_error::<&str>(
&e,
&value.tag.span,
@ -83,10 +102,8 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
.set(bool_mask, Some(*val))
.map_err(|e| parse_polars_error::<&str>(&e, &value.tag.span, None))?;
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_tag)))
}
UntaggedValue::Primitive(Primitive::Decimal(val)) => {
let chunked = series.as_ref().f64().map_err(|e| {
@ -107,10 +124,8 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)
.map_err(|e| parse_polars_error::<&str>(&e, &value.tag.span, None))?;
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_tag)))
}
UntaggedValue::Primitive(Primitive::String(val)) => {
let chunked = series.as_ref().utf8().map_err(|e| {
@ -128,18 +143,29 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let mut res = res.into_series();
res.rename("string");
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_tag)))
}
_ => Err(ShellError::labeled_error(
"Incorrect type",
format!(
"this value cannot be set in a series of type '{}'",
series.as_ref().dtype()
series.dtype()
),
value.tag.span,
)),
}
}
#[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 {})
}
}

68
crates/nu-command/src/commands/dataframe/series/set_with_idx.rs
View File

@ -1,7 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Primitive, Signature, SyntaxShape, UntaggedValue, Value};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Primitive, Signature, SyntaxShape, UntaggedValue, Value,
};
use polars::prelude::{ChunkSet, DataType, IntoSeries};
pub struct DataFrame;
@ -33,10 +36,25 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Set value in selected rows from series",
example: r#"let series = ([4 1 5 2 4 3] | dataframe to-series);
let indices = ([0 2] | dataframe to-series);
example: r#"let series = ([4 1 5 2 4 3] | dataframe to-df);
let indices = ([0 2] | dataframe to-df);
$series | dataframe set-with-idx 6 -i $indices"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::int(6).into(),
UntaggedValue::int(1).into(),
UntaggedValue::int(6).into(),
UntaggedValue::int(2).into(),
UntaggedValue::int(4).into(),
UntaggedValue::int(3).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -47,7 +65,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let indices: Value = args.req_named("indices")?;
let indices = match &indices.value {
UntaggedValue::DataFrame(nu_protocol::dataframe::PolarsData::Series(series)) => Ok(series),
UntaggedValue::DataFrame(df) => Ok(df),
_ => Err(ShellError::labeled_error(
"Incorrect type",
"can only use a series for set command",
@ -55,7 +73,9 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}?;
let casted = match indices.as_ref().dtype() {
let indices = indices.as_series(&value.tag.span)?;
let casted = match indices.dtype() {
DataType::UInt32 | DataType::UInt64 | DataType::Int32 | DataType::Int64 => indices
.as_ref()
.cast_with_dtype(&DataType::UInt32)
@ -75,11 +95,12 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
.into_iter()
.filter_map(|val| val.map(|v| v as usize));
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
match &value.value {
UntaggedValue::Primitive(Primitive::Int(val)) => {
let chunked = series.as_ref().i64().map_err(|e| {
let chunked = series.i64().map_err(|e| {
parse_polars_error::<&str>(
&e,
&value.tag.span,
@ -91,10 +112,8 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
.set_at_idx(indices, Some(*val))
.map_err(|e| parse_polars_error::<&str>(&e, &value.tag.span, None))?;
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_tag)))
}
UntaggedValue::Primitive(Primitive::Decimal(val)) => {
let chunked = series.as_ref().f64().map_err(|e| {
@ -115,10 +134,8 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)
.map_err(|e| parse_polars_error::<&str>(&e, &value.tag.span, None))?;
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_tag)))
}
UntaggedValue::Primitive(Primitive::String(val)) => {
let chunked = series.as_ref().utf8().map_err(|e| {
@ -136,10 +153,8 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let mut res = res.into_series();
res.rename("string");
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_tag)))
}
_ => Err(ShellError::labeled_error(
"Incorrect type",
@ -151,3 +166,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}
}
#[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 {})
}
}

15
crates/nu-command/src/commands/dataframe/series/shift.rs
View File

@ -1,9 +1,8 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature, SyntaxShape};
use nu_protocol::{dataframe::NuDataFrame, Signature, SyntaxShape};
use nu_source::Tagged;
use polars::prelude::IntoSeries;
pub struct DataFrame;
@ -27,7 +26,7 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Shifts the values by a given period",
example: "[1 2 2 3 3] | dataframe to-series | dataframe shift 2",
example: "[1 2 2 3 3] | dataframe to-df | dataframe shift 2",
result: None,
}]
}
@ -37,12 +36,10 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let period: Tagged<i64> = args.req(0)?;
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series.as_ref().shift(period.item);
let res = df.as_series(&df_tag.span)?.shift(period.item);
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_tag)))
}

81
crates/nu-command/src/commands/dataframe/series/str_lengths.rs Normal file
View File

@ -0,0 +1,81 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
impl WholeStreamCommand for DataFrame {
fn name(&self) -> &str {
"dataframe str-lengths"
}
fn usage(&self) -> &str {
"[Series] Get lengths of all strings"
}
fn signature(&self) -> Signature {
Signature::build("dataframe str-lengths")
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
command(args)
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Returns string lengths",
example: "[a ab abc] | dataframe to-df | dataframe str-lengths",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::int(1).into(),
UntaggedValue::int(2).into(),
UntaggedValue::int(3).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 (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
let chunked = series.utf8().map_err(|e| {
parse_polars_error::<&str>(
&e,
&df_tag.span,
Some("The str-lengths command can only be used with string columns"),
)
})?;
let res = chunked.as_ref().str_lengths();
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

92
crates/nu-command/src/commands/dataframe/series/str_slice.rs Normal file
View File

@ -0,0 +1,92 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::Tagged;
use polars::prelude::IntoSeries;
pub struct DataFrame;
impl WholeStreamCommand for DataFrame {
fn name(&self) -> &str {
"dataframe str-slice"
}
fn usage(&self) -> &str {
"[Series] Slices the string from the start position until the selected length"
}
fn signature(&self) -> Signature {
Signature::build("dataframe replace")
.required_named("start", SyntaxShape::Int, "start of slice", Some('s'))
.named("length", SyntaxShape::Int, "optional length", Some('l'))
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
command(args)
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Creates slices from the strings",
example: "[abcded abc321 abc123] | dataframe to-df | dataframe str-slice -s 1 -l 2",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::string("bc").into(),
UntaggedValue::string("bc").into(),
UntaggedValue::string("bc").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 start: Tagged<i64> = args.req_named("start")?;
let length: Option<Tagged<i64>> = args.get_flag("length")?;
let length = length.map(|v| v.item as u64);
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
let chunked = series.utf8().map_err(|e| {
parse_polars_error::<&str>(
&e,
&df_tag.span,
Some("The str-slice command can only be used with string columns"),
)
})?;
let mut res = chunked
.str_slice(start.item, length)
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
res.rename(series.name());
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

82
crates/nu-command/src/commands/dataframe/series/to_lowercase.rs Normal file
View File

@ -0,0 +1,82 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
impl WholeStreamCommand for DataFrame {
fn name(&self) -> &str {
"dataframe to-lowercase"
}
fn usage(&self) -> &str {
"[Series] Lowercase the strings in the column"
}
fn signature(&self) -> Signature {
Signature::build("dataframe to-lowercase")
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
command(args)
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Modifies strings to lowercase",
example: "[Abc aBc abC] | dataframe to-df | dataframe to-lowercase",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::string("abc").into(),
UntaggedValue::string("abc").into(),
UntaggedValue::string("abc").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 (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
let chunked = series.utf8().map_err(|e| {
parse_polars_error::<&str>(
&e,
&df_tag.span,
Some("The to-lowercase command can only be used with string columns"),
)
})?;
let mut res = chunked.to_lowercase();
res.rename(series.name());
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

82
crates/nu-command/src/commands/dataframe/series/to_uppercase.rs Normal file
View File

@ -0,0 +1,82 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use polars::prelude::IntoSeries;
pub struct DataFrame;
impl WholeStreamCommand for DataFrame {
fn name(&self) -> &str {
"dataframe to-uppercase"
}
fn usage(&self) -> &str {
"[Series] Uppercase the strings in the column"
}
fn signature(&self) -> Signature {
Signature::build("dataframe to-uppercase")
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
command(args)
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Modifies strings to uppercase",
example: "[Abc aBc abC] | dataframe to-df | dataframe to-uppercase",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::string("ABC").into(),
UntaggedValue::string("ABC").into(),
UntaggedValue::string("ABC").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 (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = df.as_series(&df_tag.span)?;
let chunked = series.utf8().map_err(|e| {
parse_polars_error::<&str>(
&e,
&df_tag.span,
Some("The to-uppercase command can only be used with string columns"),
)
})?;
let mut res = chunked.to_uppercase();
res.rename(series.name());
let df = NuDataFrame::try_from_series(vec![res.into_series()], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

43
crates/nu-command/src/commands/dataframe/series/unique.rs
View File

@ -1,8 +1,10 @@
use crate::{commands::dataframe::utils::parse_polars_error, prelude::*};
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature};
use polars::prelude::IntoSeries;
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
pub struct DataFrame;
@ -26,8 +28,16 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Returns unique values from a series",
example: "[1 2 2 3 3] | dataframe to-series | dataframe unique",
result: None,
example: "[2 2 2 2 2] | dataframe to-df | dataframe unique",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![UntaggedValue::int(2).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -35,15 +45,26 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = series
.as_ref()
let res = df
.as_series(&df_tag.span)?
.unique()
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
Ok(OutputStream::one(NuSeries::series_to_value(
res.into_series(),
tag,
)))
let df = NuDataFrame::try_from_series(vec![res], &tag.span)?;
Ok(OutputStream::one(df.into_value(df_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 {})
}
}

45
crates/nu-command/src/commands/dataframe/series/value_counts.rs
View File

@ -2,8 +2,8 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, NuSeries},
Signature,
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
use crate::commands::dataframe::utils::parse_polars_error;
@ -30,8 +30,22 @@ impl WholeStreamCommand for DataFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Calculates value counts",
example: "[5 5 6 6] | dataframe to-series | dataframe value-counts",
result: None,
example: "[5 5 5 5 6 6] | dataframe to-df | dataframe value-counts",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"0".to_string(),
vec![UntaggedValue::int(5).into(), UntaggedValue::int(6).into()],
),
Column::new(
"counts".to_string(),
vec![UntaggedValue::int(4).into(), UntaggedValue::int(2).into()],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -39,12 +53,27 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let series = NuSeries::try_from_stream(&mut args.input, &tag.span)?;
let (df, df_tag) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let df = series
.as_ref()
let df_new = df
.as_series(&df_tag.span)?
.value_counts()
.map_err(|e| parse_polars_error::<&str>(&e, &tag.span, None))?;
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(df, tag)))
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(
df_new, 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 {})
}
}

49
crates/nu-command/src/commands/dataframe/shape.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, TaggedDictBuilder};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue, Value,
};
pub struct DataFrame;
@ -26,7 +29,15 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "Shows row and column shape",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe shape",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("rows".to_string(), vec![UntaggedValue::int(2).into()]),
Column::new("columns".to_string(), vec![UntaggedValue::int(2).into()]),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -34,14 +45,34 @@ impl WholeStreamCommand for DataFrame {
fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let rows = df.as_ref().height();
let cols = df.as_ref().width();
let rows = Value {
value: (df.as_ref().height() as i64).into(),
tag: Tag::default(),
};
let mut data = TaggedDictBuilder::new(&tag);
data.insert_value("rows", format!("{}", rows));
data.insert_value("columns", format!("{}", cols));
let cols = Value {
value: (df.as_ref().width() as i64).into(),
tag: Tag::default(),
};
Ok(OutputStream::one(data.into_value()))
let rows_col = Column::new("rows".to_string(), vec![rows]);
let cols_col = Column::new("columns".to_string(), vec![cols]);
let df = NuDataFrame::try_from_columns(vec![rows_col, cols_col], &tag.span)?;
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 {})
}
}

2
crates/nu-command/src/commands/dataframe/show.rs
View File

@ -53,7 +53,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let rows: Option<Tagged<usize>> = args.get_flag("n_rows")?;
let tail: bool = args.has_flag("tail");
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let rows = rows.map(|v| v.item);
let values = if tail { df.tail(rows)? } else { df.head(rows)? };

30
crates/nu-command/src/commands/dataframe/slice.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature, SyntaxShape};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue,
};
use nu_source::Tagged;
pub struct DataFrame;
@ -29,7 +32,15 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "Create new dataframe from a slice of the rows",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe slice 0 1",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("a".to_string(), vec![UntaggedValue::int(1).into()]),
Column::new("b".to_string(), vec![UntaggedValue::int(2).into()]),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -40,8 +51,21 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let offset: Tagged<usize> = args.req(0)?;
let size: Tagged<usize> = args.req(1)?;
let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let res = df.as_ref().slice(offset.item as i64, size.item);
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, 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 {})
}
}

89
crates/nu-command/src/commands/dataframe/sort.rs
View File

@ -2,7 +2,7 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, NuSeries, PolarsData},
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
@ -33,12 +33,39 @@ impl WholeStreamCommand for DataFrame {
Example {
description: "Create new sorted dataframe",
example: "[[a b]; [3 4] [1 2]] | dataframe to-df | dataframe sort a",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"a".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())]),
},
Example {
description: "Create new sorted series",
example: "[3 4 1 2] | dataframe to-series | dataframe sort",
result: None,
example: "[3 4 1 2] | dataframe to-df | dataframe sort",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::int(1).into(),
UntaggedValue::int(2).into(),
UntaggedValue::int(3).into(),
UntaggedValue::int(4).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
]
}
@ -53,31 +80,38 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let reverse = args.has_flag("reverse");
match value.value {
UntaggedValue::DataFrame(PolarsData::EagerDataFrame(df)) => {
let columns: Vec<Value> = args.rest(0)?;
if !columns.is_empty() {
let (col_string, col_span) = convert_columns(&columns, &tag)?;
match &value.value {
UntaggedValue::DataFrame(df) => {
if df.is_series() {
let columns = df.as_ref().get_column_names();
let res = df
.as_ref()
.sort(&col_string, reverse)
.map_err(|e| parse_polars_error::<&str>(&e, &col_span, None))?;
.sort(columns, reverse)
.map_err(|e| parse_polars_error::<&str>(&e, &value.tag.span, None))?;
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, tag)))
} else {
Err(ShellError::labeled_error(
"Missing columns",
"missing column name to perform sort",
&tag.span,
))
let columns: Vec<Value> = args.rest(0)?;
if !columns.is_empty() {
let (col_string, col_span) = convert_columns(&columns, &tag)?;
let res = df
.as_ref()
.sort(&col_string, reverse)
.map_err(|e| parse_polars_error::<&str>(&e, &col_span, None))?;
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, tag)))
} else {
Err(ShellError::labeled_error(
"Missing columns",
"missing column name to perform sort",
&tag.span,
))
}
}
}
UntaggedValue::DataFrame(PolarsData::Series(series)) => {
let res = series.as_ref().sort(reverse);
Ok(OutputStream::one(NuSeries::series_to_value(res, tag)))
}
_ => Err(ShellError::labeled_error(
"Incorrect type",
"sort cannot be done with this value",
@ -85,3 +119,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}
}
#[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 {})
}
}

62
crates/nu-command/src/commands/dataframe/take.rs
View File

@ -2,7 +2,7 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, NuSeries, PolarsData},
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
use polars::prelude::DataType;
@ -37,16 +37,38 @@ impl WholeStreamCommand for DataFrame {
Example {
description: "Takes selected rows from dataframe",
example: r#"let df = ([[a b]; [4 1] [5 2] [4 3]] | dataframe to-df);
let indices = ([0 2] | dataframe to-series);
let indices = ([0 2] | dataframe to-df);
$df | dataframe take $indices"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"a".to_string(),
vec![UntaggedValue::int(4).into(), UntaggedValue::int(4).into()],
),
Column::new(
"b".to_string(),
vec![UntaggedValue::int(1).into(), UntaggedValue::int(3).into()],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "Takes selected rows from series",
example: r#"let series = ([4 1 5 2 4 3] | dataframe to-series);
let indices = ([0 2] | dataframe to-series);
example: r#"let series = ([4 1 5 2 4 3] | dataframe to-df);
let indices = ([0 2] | dataframe to-df);
$series | dataframe take $indices"#,
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![UntaggedValue::int(4).into(), UntaggedValue::int(5).into()],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
]
}
@ -56,8 +78,8 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let value: Value = args.req(0)?;
let series = match &value.value {
UntaggedValue::DataFrame(PolarsData::Series(series)) => Ok(series),
let df = match &value.value {
UntaggedValue::DataFrame(df) => Ok(df),
_ => Err(ShellError::labeled_error(
"Incorrect type",
"can only use a series for take command",
@ -65,7 +87,9 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}?;
let casted = match series.as_ref().dtype() {
let series = df.as_series(&value.tag.span)?;
let casted = match series.dtype() {
DataType::UInt32 | DataType::UInt64 | DataType::Int32 | DataType::Int64 => series
.as_ref()
.cast_with_dtype(&DataType::UInt32)
@ -88,16 +112,11 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
})?;
match value.value {
UntaggedValue::DataFrame(PolarsData::EagerDataFrame(df)) => {
UntaggedValue::DataFrame(df) => {
let res = df.as_ref().take(indices);
Ok(OutputStream::one(NuDataFrame::dataframe_to_value(res, tag)))
}
UntaggedValue::DataFrame(PolarsData::Series(series)) => {
let res = series.as_ref().take(indices);
Ok(OutputStream::one(NuSeries::series_to_value(res, tag)))
}
_ => Err(ShellError::labeled_error(
"No dataframe or series in stream",
"no dataframe or series found in input stream",
@ -105,3 +124,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}
}
#[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 {})
}
}

2
crates/nu-command/src/commands/dataframe/to_csv.rs
View File

@ -64,7 +64,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let delimiter: Option<Tagged<String>> = args.get_flag("delimiter")?;
let no_header: bool = args.has_flag("no_header");
let mut df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (mut df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let mut file = File::create(&file_name.item).map_err(|e| {
ShellError::labeled_error(

115
crates/nu-command/src/commands/dataframe/to_df.rs
View File

@ -1,7 +1,10 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuDataFrame, Signature};
use nu_protocol::{
dataframe::{Column, NuDataFrame},
Signature, UntaggedValue,
};
pub struct DataFrame;
@ -11,7 +14,7 @@ impl WholeStreamCommand for DataFrame {
}
fn usage(&self) -> &str {
"Converts a pipelined Table or List into a polars dataframe"
"Converts a List, Table or Dictionary into a polars dataframe"
}
fn signature(&self) -> Signature {
@ -27,10 +30,108 @@ impl WholeStreamCommand for DataFrame {
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Takes an input stream and converts it to a polars dataframe",
example: "[[a b];[1 2] [3 4]] | dataframe to-df",
result: None,
}]
vec![
Example {
description: "Takes a dictionary and creates a dataframe",
example: "[[a b];[1 2] [3 4]] | dataframe to-df",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"a".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())]),
},
Example {
description: "Takes a list of tables and creates a dataframe",
example: "[[1 2 a] [3 4 b] [5 6 c]] | dataframe to-df",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"0".to_string(),
vec![
UntaggedValue::int(1).into(),
UntaggedValue::int(3).into(),
UntaggedValue::int(5).into(),
],
),
Column::new(
"1".to_string(),
vec![
UntaggedValue::int(2).into(),
UntaggedValue::int(4).into(),
UntaggedValue::int(6).into(),
],
),
Column::new(
"2".to_string(),
vec![
UntaggedValue::string("a").into(),
UntaggedValue::string("b").into(),
UntaggedValue::string("c").into(),
],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "Takes a list and creates a dataframe",
example: "[a b c] | dataframe to-df",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::string("a").into(),
UntaggedValue::string("b").into(),
UntaggedValue::string("c").into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
Example {
description: "Takes a list of booleans and creates a dataframe",
example: "[$true $true $false] | dataframe to-df",
result: Some(vec![NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(true).into(),
UntaggedValue::boolean(false).into(),
],
)],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
},
]
}
}
#[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 {})
}
}

2
crates/nu-command/src/commands/dataframe/to_parquet.rs
View File

@ -48,7 +48,7 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let file_name: Tagged<PathBuf> = args.req(0)?;
let mut df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let (mut df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let file = File::create(&file_name.item).map_err(|e| {
ShellError::labeled_error(

44
crates/nu-command/src/commands/dataframe/to_series.rs
View File

@ -1,44 +0,0 @@
use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{dataframe::NuSeries, Signature, SyntaxShape};
use nu_source::Tagged;
pub struct DataFrame;
impl WholeStreamCommand for DataFrame {
fn name(&self) -> &str {
"dataframe to-series"
}
fn usage(&self) -> &str {
"Converts a pipelined List into a polars series"
}
fn signature(&self) -> Signature {
Signature::build("dataframe to-series").optional(
"name",
SyntaxShape::String,
"Optional series name",
)
}
fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
let tag = args.call_info.name_tag.clone();
let name: Option<Tagged<String>> = args.opt(0)?;
let name = name.map(|v| v.item);
let series = NuSeries::try_from_iter(args.input, name)?;
Ok(InputStream::one(series.into_value(tag)))
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Takes an input stream and converts it to a polars series",
example: "[1 2 3 4] | dataframe to-series my-col",
result: None,
}]
}
}

27
crates/nu-command/src/commands/dataframe/where_.rs
View File

@ -2,7 +2,7 @@ use crate::prelude::*;
use nu_engine::{evaluate_baseline_expr, WholeStreamCommand};
use nu_errors::ShellError;
use nu_protocol::{
dataframe::NuDataFrame,
dataframe::{Column, NuDataFrame},
hir::{CapturedBlock, ClassifiedCommand, Expression, Literal, Operator, SpannedExpression},
Primitive, Signature, SyntaxShape, UnspannedPathMember, UntaggedValue, Value,
};
@ -37,7 +37,15 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "Filter dataframe based on column a",
example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe where a == 1",
result: None,
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new("a".to_string(), vec![UntaggedValue::int(1).into()]),
Column::new("b".to_string(), vec![UntaggedValue::int(2).into()]),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -143,7 +151,7 @@ fn filter_dataframe(
}?;
let span = args.call_info.name_tag.span;
let df = NuDataFrame::try_from_stream(&mut args.input, &span)?;
let (df, _) = NuDataFrame::try_from_stream(&mut args.input, &span)?;
let col = df
.as_ref()
@ -214,3 +222,16 @@ fn filter_dataframe(
args.call_info.name_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 {})
}
}

56
crates/nu-command/src/commands/dataframe/with_column.rs
View File

@ -2,7 +2,7 @@ use crate::prelude::*;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{
dataframe::{NuDataFrame, PolarsData},
dataframe::{Column, NuDataFrame},
Signature, SyntaxShape, UntaggedValue, Value,
};
use nu_source::Tagged;
@ -33,8 +33,35 @@ impl WholeStreamCommand for DataFrame {
vec![Example {
description: "Adds a series to the dataframe",
example:
"[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe with-column ([5 6] | dataframe to-series) --name c",
result: None,
"[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe with-column ([5 6] | dataframe to-df) --name c",
result: Some(vec![NuDataFrame::try_from_columns(
vec![
Column::new(
"a".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(),
],
),
Column::new(
"c".to_string(),
vec![
UntaggedValue::int(5).into(),
UntaggedValue::int(6).into(),
],
),
],
&Span::default(),
)
.expect("simple df for test should not fail")
.into_value(Tag::default())]),
}]
}
}
@ -44,8 +71,8 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
let value: Value = args.req(0)?;
let name: Tagged<String> = args.req_named("name")?;
let mut series = match value.value {
UntaggedValue::DataFrame(PolarsData::Series(series)) => Ok(series),
let df = match value.value {
UntaggedValue::DataFrame(df) => Ok(df),
_ => Err(ShellError::labeled_error(
"Incorrect type",
"can only add a series to a dataframe",
@ -53,9 +80,11 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
)),
}?;
let series = series.as_mut().rename(name.item.as_ref()).clone();
let mut series = df.as_series(&value.tag.span)?;
let mut df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
let series = series.rename(name.item.as_ref()).clone();
let (mut df, _) = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
df.as_mut()
.with_column(series)
@ -63,3 +92,16 @@ fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
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 {})
}
}

12
crates/nu-command/src/commands/filters/first.rs
View File

@ -126,11 +126,13 @@ fn first(args: CommandArgs) -> Result<OutputStream, ShellError> {
tag,
)),
#[cfg(all(not(target_arch = "wasm32"), feature = "dataframe"))]
UntaggedValue::DataFrame(_) => Err(ShellError::labeled_error(
"unsure how to handled UntaggedValue::DataFrame",
"found dataframe",
tag,
)),
UntaggedValue::DataFrame(_) | UntaggedValue::FrameStruct(_) => {
Err(ShellError::labeled_error(
"unsure how to handled dataframe struct",
"found dataframe",
tag,
))
}
},
None => Ok(input_peek.take(rows_desired).into_output_stream()),
}

2
crates/nu-command/src/commands/filters/uniq.rs
View File

@ -166,7 +166,7 @@ fn uniq(args: CommandArgs) -> Result<ActionStream, ShellError> {
))
}
#[cfg(feature = "dataframe")]
UntaggedValue::DataFrame(_) => {
UntaggedValue::DataFrame(_) | UntaggedValue::FrameStruct(_) => {
return Err(ShellError::labeled_error(
"uniq -c cannot operate on data structs",
"source",

2
crates/nu-command/src/commands/formats/to/json.rs
View File

@ -115,7 +115,7 @@ pub fn value_to_json_value(v: &Value) -> Result<serde_json::Value, ShellError> {
serde_json::Value::Null
}
#[cfg(feature = "dataframe")]
UntaggedValue::DataFrame(_) => serde_json::Value::Null,
UntaggedValue::DataFrame(_) | UntaggedValue::FrameStruct(_) => serde_json::Value::Null,
UntaggedValue::Primitive(Primitive::Binary(b)) => serde_json::Value::Array(
b.iter()
.map(|x| {

4
crates/nu-command/src/commands/formats/to/toml.rs
View File

@ -74,7 +74,9 @@ fn helper(v: &Value) -> Result<toml::Value, ShellError> {
UntaggedValue::Error(e) => return Err(e.clone()),
UntaggedValue::Block(_) => toml::Value::String("<Block>".to_string()),
#[cfg(feature = "dataframe")]
UntaggedValue::DataFrame(_) => toml::Value::String("<Data>".to_string()),
UntaggedValue::DataFrame(_) | UntaggedValue::FrameStruct(_) => {
toml::Value::String("<Data>".to_string())
}
UntaggedValue::Primitive(Primitive::Range(_)) => toml::Value::String("<Range>".to_string()),
UntaggedValue::Primitive(Primitive::Binary(b)) => {
toml::Value::Array(b.iter().map(|x| toml::Value::Integer(*x as i64)).collect())

2
crates/nu-command/src/commands/formats/to/yaml.rs
View File

@ -96,7 +96,7 @@ pub fn value_to_yaml_value(v: &Value) -> Result<serde_yaml::Value, ShellError> {
serde_yaml::Value::Null
}
#[cfg(feature = "dataframe")]
UntaggedValue::DataFrame(_) => serde_yaml::Value::Null,
UntaggedValue::DataFrame(_) | UntaggedValue::FrameStruct(_) => serde_yaml::Value::Null,
UntaggedValue::Primitive(Primitive::Binary(b)) => serde_yaml::Value::Sequence(
b.iter()
.map(|x| serde_yaml::Value::Number(serde_yaml::Number::from(*x)))

19
crates/nu-command/src/commands/mod.rs
View File

@ -27,15 +27,16 @@ pub use core_commands::*;
pub use dataframe::{
DataFrame, DataFrameAggregate, DataFrameAllFalse, DataFrameAllTrue, DataFrameArgMax,
DataFrameArgMin, DataFrameArgSort, DataFrameArgTrue, DataFrameArgUnique, DataFrameColumn,
DataFrameDTypes, DataFrameDrop, DataFrameDropDuplicates, DataFrameDropNulls, DataFrameDummies,
DataFrameFilter, DataFrameFirst, DataFrameGet, DataFrameGroupBy, DataFrameIsDuplicated,
DataFrameIsIn, DataFrameIsNotNull, DataFrameIsNull, DataFrameIsUnique, DataFrameJoin,
DataFrameLast, DataFrameList, DataFrameMelt, DataFrameNNull, DataFrameNUnique, DataFrameNot,
DataFrameOpen, DataFramePivot, DataFrameReplace, DataFrameSample, DataFrameSelect,
DataFrameSeriesRename, DataFrameSet, DataFrameSetWithIdx, DataFrameShape, DataFrameShift,
DataFrameShow, DataFrameSlice, DataFrameSort, DataFrameTake, DataFrameToCsv, DataFrameToDF,
DataFrameToParquet, DataFrameToSeries, DataFrameUnique, DataFrameValueCounts, DataFrameWhere,
DataFrameWithColumn,
DataFrameConcatenate, DataFrameContains, DataFrameDTypes, DataFrameDrop,
DataFrameDropDuplicates, DataFrameDropNulls, DataFrameDummies, DataFrameFilter, DataFrameFirst,
DataFrameGet, DataFrameGroupBy, DataFrameIsDuplicated, DataFrameIsIn, DataFrameIsNotNull,
DataFrameIsNull, DataFrameIsUnique, DataFrameJoin, DataFrameLast, DataFrameList, DataFrameMelt,
DataFrameNNull, DataFrameNUnique, DataFrameNot, DataFrameOpen, DataFramePivot,
DataFrameReplace, DataFrameReplaceAll, DataFrameSample, DataFrameSelect, DataFrameSeriesRename,
DataFrameSet, DataFrameSetWithIdx, DataFrameShape, DataFrameShift, DataFrameShow,
DataFrameSlice, DataFrameSort, DataFrameStringLengths, DataFrameStringSlice, DataFrameTake,
DataFrameToCsv, DataFrameToDF, DataFrameToLowercase, DataFrameToParquet, DataFrameToUppercase,
DataFrameUnique, DataFrameValueCounts, DataFrameWhere, DataFrameWithColumn,
};
pub use env::*;
pub use filesystem::*;

20
crates/nu-command/src/commands/viewers/autoview/command.rs
View File

@ -9,7 +9,7 @@ use nu_protocol::{Primitive, Signature, UntaggedValue, Value};
use nu_table::TextStyle;
#[cfg(feature = "dataframe")]
use nu_protocol::dataframe::PolarsData;
use nu_protocol::dataframe::FrameStruct;
pub struct Command;
@ -239,7 +239,7 @@ pub fn autoview(args: CommandArgs) -> Result<OutputStream, ShellError> {
}
#[cfg(feature = "dataframe")]
Value {
value: UntaggedValue::DataFrame(PolarsData::EagerDataFrame(df)),
value: UntaggedValue::DataFrame(df),
tag,
} => {
if let Some(table) = table {
@ -253,7 +253,7 @@ pub fn autoview(args: CommandArgs) -> Result<OutputStream, ShellError> {
}
#[cfg(feature = "dataframe")]
Value {
value: UntaggedValue::DataFrame(PolarsData::GroupBy(groupby)),
value: UntaggedValue::FrameStruct(FrameStruct::GroupBy(groupby)),
tag,
} => {
if let Some(table) = table {
@ -265,20 +265,6 @@ pub fn autoview(args: CommandArgs) -> Result<OutputStream, ShellError> {
let _ = result.collect::<Vec<_>>();
}
}
#[cfg(feature = "dataframe")]
Value {
value: UntaggedValue::DataFrame(PolarsData::Series(series)),
tag,
} => {
if let Some(table) = table {
// TODO. Configure the parameter rows from file. It can be
// adjusted to see a certain amount of values in the head
let command_args =
create_default_command_args(&context, series.print()?.into(), tag);
let result = table.run(command_args)?;
let _ = result.collect::<Vec<_>>();
}
}
Value {
value: UntaggedValue::Primitive(Primitive::Nothing),
..

8
crates/nu-command/src/default_context.rs
View File

@ -287,7 +287,6 @@ pub fn create_default_context(interactive: bool) -> Result<EvaluationContext, Bo
whole_stream_command(DataFramePivot),
whole_stream_command(DataFrameWhere),
whole_stream_command(DataFrameToDF),
whole_stream_command(DataFrameToSeries),
whole_stream_command(DataFrameToParquet),
whole_stream_command(DataFrameToCsv),
whole_stream_command(DataFrameSort),
@ -321,6 +320,13 @@ pub fn create_default_context(interactive: bool) -> Result<EvaluationContext, Bo
whole_stream_command(DataFrameSetWithIdx),
whole_stream_command(DataFrameShape),
whole_stream_command(DataFrameReplace),
whole_stream_command(DataFrameReplaceAll),
whole_stream_command(DataFrameStringLengths),
whole_stream_command(DataFrameContains),
whole_stream_command(DataFrameToLowercase),
whole_stream_command(DataFrameToUppercase),
whole_stream_command(DataFrameStringSlice),
whole_stream_command(DataFrameConcatenate),
]);
#[cfg(feature = "clipboard-cli")]

84
crates/nu-command/src/examples.rs
View File

@ -14,6 +14,11 @@ use nu_protocol::hir::{ClassifiedBlock, ExternalRedirection};
use nu_protocol::{ShellTypeName, Value};
use nu_source::AnchorLocation;
#[cfg(feature = "dataframe")]
use crate::commands::{
DataFrameGroupBy, DataFrameIsNull, DataFrameShift, DataFrameToDF, DataFrameWithColumn,
};
use crate::commands::{
Append, BuildString, Each, Echo, First, Get, Keep, Last, Let, Math, MathMode, Nth, Select,
StrCollect, Wrap,
@ -149,6 +154,85 @@ pub fn test(cmd: impl WholeStreamCommand + 'static) -> Result<(), ShellError> {
Ok(())
}
#[cfg(feature = "dataframe")]
pub fn test_dataframe(cmd: impl WholeStreamCommand + 'static) -> Result<(), ShellError> {
use nu_protocol::UntaggedValue;
let examples = cmd.examples();
let base_context = EvaluationContext::basic();
base_context.add_commands(vec![
whole_stream_command(cmd),
// Commands used with dataframe
whole_stream_command(DataFrameToDF),
whole_stream_command(DataFrameShift),
whole_stream_command(DataFrameIsNull),
whole_stream_command(DataFrameGroupBy),
whole_stream_command(DataFrameWithColumn),
// Base commands for context
whole_stream_command(Math),
whole_stream_command(MathMode {}),
whole_stream_command(Echo {}),
whole_stream_command(BuildString {}),
whole_stream_command(Get {}),
whole_stream_command(Keep {}),
whole_stream_command(Each {}),
whole_stream_command(Let {}),
whole_stream_command(Select),
whole_stream_command(StrCollect),
whole_stream_command(Wrap),
]);
for sample_pipeline in examples {
let mut ctx = base_context.clone();
println!("{:?}", &sample_pipeline.example);
let block = parse_line(sample_pipeline.example, &ctx)?;
if let Some(expected) = &sample_pipeline.result {
let start = std::time::Instant::now();
let result = evaluate_block(block, &mut ctx)?;
println!("input: {}", sample_pipeline.example);
println!("result: {:?}", result);
println!("done: {:?}", start.elapsed());
let value = match result.get(0) {
Some(v) => v,
None => panic!(
"Unable to extract a value after parsing example: {}",
sample_pipeline.example
),
};
let df = match &value.value {
UntaggedValue::DataFrame(df) => df,
_ => panic!(
"Unable to extract dataframe from parsed example: {}",
sample_pipeline.example
),
};
let expected = match expected.get(0) {
Some(v) => v,
None => panic!("Empty vector in result example"),
};
let df_expected = match &expected.value {
UntaggedValue::DataFrame(df) => df,
_ => panic!("Unable to extract dataframe from example result"),
};
println!("expected: {:?}", df_expected);
assert_eq!(df, df_expected)
}
}
Ok(())
}
pub fn test_anchors(cmd: Command) -> Result<(), ShellError> {
let examples = cmd.examples();