From 46b86f3541febbd42a510dcdaf954fec81d7f79e Mon Sep 17 00:00:00 2001 From: Fernando Herrera Date: Sat, 18 Dec 2021 17:45:09 +0000 Subject: [PATCH] Migration of series commands (#515) * corrected missing shellerror type * batch dataframe commands * removed option to find declaration with input * ordered dataframe folders * dataframe command name * series commands * date commands * series commands * series commands * clippy correction * rename commands --- crates/nu-command/Cargo.toml | 8 +- crates/nu-command/src/dataframe/append.rs | 13 +- crates/nu-command/src/dataframe/column.rs | 6 +- crates/nu-command/src/dataframe/command.rs | 2 +- crates/nu-command/src/dataframe/describe.rs | 6 +- crates/nu-command/src/dataframe/drop.rs | 6 +- crates/nu-command/src/dataframe/drop_nulls.rs | 120 ++++++++++++ crates/nu-command/src/dataframe/dtypes.rs | 6 +- crates/nu-command/src/dataframe/mod.rs | 57 +++++- crates/nu-command/src/dataframe/open.rs | 4 +- .../src/dataframe/series/all_false.rs | 102 ++++++++++ .../src/dataframe/series/all_true.rs | 102 ++++++++++ .../src/dataframe/series/arg_max.rs | 81 ++++++++ .../src/dataframe/series/arg_min.rs | 81 ++++++++ .../src/dataframe/series/cumulative.rs | 129 +++++++++++++ .../src/dataframe/series/date/get_day.rs | 87 +++++++++ .../src/dataframe/series/date/get_hour.rs | 87 +++++++++ .../src/dataframe/series/date/get_minute.rs | 87 +++++++++ .../src/dataframe/series/date/get_month.rs | 87 +++++++++ .../dataframe/series/date/get_nanosecond.rs | 87 +++++++++ .../src/dataframe/series/date/get_ordinal.rs | 87 +++++++++ .../src/dataframe/series/date/get_second.rs | 87 +++++++++ .../src/dataframe/series/date/get_week.rs | 87 +++++++++ .../src/dataframe/series/date/get_weekday.rs | 87 +++++++++ .../src/dataframe/series/date/get_year.rs | 87 +++++++++ .../src/dataframe/series/date/mod.rs | 21 ++ .../src/dataframe/series/indexes/arg_sort.rs | 95 ++++++++++ .../src/dataframe/series/indexes/arg_true.rs | 85 +++++++++ .../dataframe/series/indexes/arg_unique.rs | 86 +++++++++ .../src/dataframe/series/indexes/mod.rs | 9 + .../dataframe/series/indexes/set_with_idx.rs | 179 ++++++++++++++++++ .../dataframe/series/masks/is_duplicated.rs | 95 ++++++++++ .../src/dataframe/series/masks/is_in.rs | 104 ++++++++++ .../src/dataframe/series/masks/is_not_null.rs | 78 ++++++++ .../src/dataframe/series/masks/is_null.rs | 78 ++++++++ .../src/dataframe/series/masks/is_unique.rs | 90 +++++++++ .../src/dataframe/series/masks/mod.rs | 15 ++ .../src/dataframe/series/masks/not.rs | 82 ++++++++ .../src/dataframe/series/masks/set.rs | 163 ++++++++++++++++ crates/nu-command/src/dataframe/series/mod.rs | 36 ++++ .../nu-command/src/dataframe/series/n_null.rs | 79 ++++++++ .../src/dataframe/series/n_unique.rs | 85 +++++++++ .../nu-command/src/dataframe/series/rename.rs | 79 ++++++++ .../src/dataframe/series/rolling.rs | 163 ++++++++++++++++ .../nu-command/src/dataframe/series/shift.rs | 79 ++++++++ .../dataframe/series/string/concatenate.rs | 111 +++++++++++ .../src/dataframe/series/string/contains.rs | 98 ++++++++++ .../src/dataframe/series/string/mod.rs | 19 ++ .../src/dataframe/series/string/replace.rs | 116 ++++++++++++ .../dataframe/series/string/replace_all.rs | 116 ++++++++++++ .../dataframe/series/string/str_lengths.rs | 85 +++++++++ .../src/dataframe/series/string/str_slice.rs | 101 ++++++++++ .../src/dataframe/series/string/strftime.rs | 96 ++++++++++ .../dataframe/series/string/to_lowercase.rs | 90 +++++++++ .../dataframe/series/string/to_uppercase.rs | 90 +++++++++ .../nu-command/src/dataframe/series/unique.rs | 80 ++++++++ .../src/dataframe/series/value_counts.rs | 84 ++++++++ .../src/dataframe/test_dataframe.rs | 13 +- crates/nu-command/src/dataframe/to_df.rs | 12 +- .../nu-command/src/dataframe/with_column.rs | 100 ++++++++++ crates/nu-parser/src/parse_keywords.rs | 3 +- crates/nu-protocol/src/shell_error.rs | 4 + crates/nu-protocol/src/value/from_value.rs | 14 ++ 63 files changed, 4491 insertions(+), 35 deletions(-) create mode 100644 crates/nu-command/src/dataframe/drop_nulls.rs create mode 100644 crates/nu-command/src/dataframe/series/all_false.rs create mode 100644 crates/nu-command/src/dataframe/series/all_true.rs create mode 100644 crates/nu-command/src/dataframe/series/arg_max.rs create mode 100644 crates/nu-command/src/dataframe/series/arg_min.rs create mode 100644 crates/nu-command/src/dataframe/series/cumulative.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_day.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_hour.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_minute.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_month.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_nanosecond.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_ordinal.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_second.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_week.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_weekday.rs create mode 100644 crates/nu-command/src/dataframe/series/date/get_year.rs create mode 100644 crates/nu-command/src/dataframe/series/date/mod.rs create mode 100644 crates/nu-command/src/dataframe/series/indexes/arg_sort.rs create mode 100644 crates/nu-command/src/dataframe/series/indexes/arg_true.rs create mode 100644 crates/nu-command/src/dataframe/series/indexes/arg_unique.rs create mode 100644 crates/nu-command/src/dataframe/series/indexes/mod.rs create mode 100644 crates/nu-command/src/dataframe/series/indexes/set_with_idx.rs create mode 100644 crates/nu-command/src/dataframe/series/masks/is_duplicated.rs create mode 100644 crates/nu-command/src/dataframe/series/masks/is_in.rs create mode 100644 crates/nu-command/src/dataframe/series/masks/is_not_null.rs create mode 100644 crates/nu-command/src/dataframe/series/masks/is_null.rs create mode 100644 crates/nu-command/src/dataframe/series/masks/is_unique.rs create mode 100644 crates/nu-command/src/dataframe/series/masks/mod.rs create mode 100644 crates/nu-command/src/dataframe/series/masks/not.rs create mode 100644 crates/nu-command/src/dataframe/series/masks/set.rs create mode 100644 crates/nu-command/src/dataframe/series/n_null.rs create mode 100644 crates/nu-command/src/dataframe/series/n_unique.rs create mode 100644 crates/nu-command/src/dataframe/series/rename.rs create mode 100644 crates/nu-command/src/dataframe/series/rolling.rs create mode 100644 crates/nu-command/src/dataframe/series/shift.rs create mode 100644 crates/nu-command/src/dataframe/series/string/concatenate.rs create mode 100644 crates/nu-command/src/dataframe/series/string/contains.rs create mode 100644 crates/nu-command/src/dataframe/series/string/mod.rs create mode 100644 crates/nu-command/src/dataframe/series/string/replace.rs create mode 100644 crates/nu-command/src/dataframe/series/string/replace_all.rs create mode 100644 crates/nu-command/src/dataframe/series/string/str_lengths.rs create mode 100644 crates/nu-command/src/dataframe/series/string/str_slice.rs create mode 100644 crates/nu-command/src/dataframe/series/string/strftime.rs create mode 100644 crates/nu-command/src/dataframe/series/string/to_lowercase.rs create mode 100644 crates/nu-command/src/dataframe/series/string/to_uppercase.rs create mode 100644 crates/nu-command/src/dataframe/series/unique.rs create mode 100644 crates/nu-command/src/dataframe/series/value_counts.rs create mode 100644 crates/nu-command/src/dataframe/with_column.rs diff --git a/crates/nu-command/Cargo.toml b/crates/nu-command/Cargo.toml index 77e17f04ab..853b426a34 100644 --- a/crates/nu-command/Cargo.toml +++ b/crates/nu-command/Cargo.toml @@ -71,7 +71,11 @@ num = { version = "0.4.0", optional = true } [dependencies.polars] version = "0.18.0" optional = true -features = ["default", "parquet", "json", "serde", "object", "checked_arithmetic", "strings"] +features = [ + "default", "parquet", "json", "serde", "object", + "checked_arithmetic", "strings", "cum_agg", "is_in", + "rolling_window", "strings" +] [features] trash-support = ["trash"] @@ -79,4 +83,4 @@ plugin = ["nu-parser/plugin"] dataframe = ["polars", "num"] [build-dependencies] -shadow-rs = "0.8.1" \ No newline at end of file +shadow-rs = "0.8.1" diff --git a/crates/nu-command/src/dataframe/append.rs b/crates/nu-command/src/dataframe/append.rs index 914ff71891..8840ffd065 100644 --- a/crates/nu-command/src/dataframe/append.rs +++ b/crates/nu-command/src/dataframe/append.rs @@ -12,7 +12,7 @@ pub struct AppendDF; impl Command for AppendDF { fn name(&self) -> &str { - "dataframe append" + "dfr append" } fn usage(&self) -> &str { @@ -30,8 +30,8 @@ impl Command for AppendDF { vec![ Example { description: "Appends a dataframe as new columns", - example: r#"let a = ([[a b]; [1 2] [3 4]] | dataframe to-df); -$a | dataframe append $a"#, + example: r#"let a = ([[a b]; [1 2] [3 4]] | dfr to-df); + $a | dfr append $a"#, result: Some( NuDataFrame::try_from_columns(vec![ Column::new("a".to_string(), vec![1.into(), 3.into()]), @@ -45,9 +45,8 @@ $a | dataframe append $a"#, }, Example { description: "Appends a dataframe merging at the end of columns", - //example: r#"let a = ([[a b]; [1 2] [3 4]] | to df); $a | append-df $a -col"#, - example: r#"let a = ([[a b]; [1 2] [3 4]] | dataframe to-df); -$a | dataframe append $a --col"#, + example: r#"let a = ([[a b]; [1 2] [3 4]] | dfr to-df); + $a | dfr append $a --col"#, result: Some( NuDataFrame::try_from_columns(vec![ Column::new( @@ -104,6 +103,6 @@ mod test { #[test] fn test_examples() { - test_dataframe(AppendDF {}) + test_dataframe(vec![Box::new(AppendDF {})]) } } diff --git a/crates/nu-command/src/dataframe/column.rs b/crates/nu-command/src/dataframe/column.rs index d0d84c118d..2dc3efe7f4 100644 --- a/crates/nu-command/src/dataframe/column.rs +++ b/crates/nu-command/src/dataframe/column.rs @@ -12,7 +12,7 @@ pub struct ColumnDF; impl Command for ColumnDF { fn name(&self) -> &str { - "dataframe column" + "dfr column" } fn usage(&self) -> &str { @@ -28,7 +28,7 @@ impl Command for ColumnDF { fn examples(&self) -> Vec { vec![Example { description: "Returns the selected column as series", - example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe column a", + example: "[[a b]; [1 2] [3 4]] | dfr to-df | dfr column a", result: Some( NuDataFrame::try_from_columns(vec![Column::new( "a".to_string(), @@ -76,6 +76,6 @@ mod test { #[test] fn test_examples() { - test_dataframe(ColumnDF {}) + test_dataframe(vec![Box::new(ColumnDF {})]) } } diff --git a/crates/nu-command/src/dataframe/command.rs b/crates/nu-command/src/dataframe/command.rs index 5ea43a2ede..c51e67a640 100644 --- a/crates/nu-command/src/dataframe/command.rs +++ b/crates/nu-command/src/dataframe/command.rs @@ -10,7 +10,7 @@ pub struct Dataframe; impl Command for Dataframe { fn name(&self) -> &str { - "dataframe" + "dfr" } fn usage(&self) -> &str { diff --git a/crates/nu-command/src/dataframe/describe.rs b/crates/nu-command/src/dataframe/describe.rs index 93a4458bc0..9c3b953d4c 100644 --- a/crates/nu-command/src/dataframe/describe.rs +++ b/crates/nu-command/src/dataframe/describe.rs @@ -17,7 +17,7 @@ pub struct DescribeDF; impl Command for DescribeDF { fn name(&self) -> &str { - "dataframe describe" + "dfr describe" } fn usage(&self) -> &str { @@ -31,7 +31,7 @@ impl Command for DescribeDF { fn examples(&self) -> Vec { vec![Example { description: "dataframe description", - example: "[[a b]; [1 1] [1 1]] | dataframe to-df | dataframe describe", + example: "[[a b]; [1 1] [1 1]] | dfr to-df | dfr describe", result: Some( NuDataFrame::try_from_columns(vec![ Column::new( @@ -236,6 +236,6 @@ mod test { #[test] fn test_examples() { - test_dataframe(DescribeDF {}) + test_dataframe(vec![Box::new(DescribeDF {})]) } } diff --git a/crates/nu-command/src/dataframe/drop.rs b/crates/nu-command/src/dataframe/drop.rs index c632d0f0a7..fd0cc1aa80 100644 --- a/crates/nu-command/src/dataframe/drop.rs +++ b/crates/nu-command/src/dataframe/drop.rs @@ -13,7 +13,7 @@ pub struct DropDF; impl Command for DropDF { fn name(&self) -> &str { - "dataframe drop" + "dfr drop" } fn usage(&self) -> &str { @@ -29,7 +29,7 @@ impl Command for DropDF { fn examples(&self) -> Vec { vec![Example { description: "drop column a", - example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe drop a", + example: "[[a b]; [1 2] [3 4]] | dfr to-df | dfr drop a", result: Some( NuDataFrame::try_from_columns(vec![Column::new( "b".to_string(), @@ -106,6 +106,6 @@ mod test { #[test] fn test_examples() { - test_dataframe(DropDF {}) + test_dataframe(vec![Box::new(DropDF {})]) } } diff --git a/crates/nu-command/src/dataframe/drop_nulls.rs b/crates/nu-command/src/dataframe/drop_nulls.rs new file mode 100644 index 0000000000..f226b98f85 --- /dev/null +++ b/crates/nu-command/src/dataframe/drop_nulls.rs @@ -0,0 +1,120 @@ +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Value, +}; + +use super::values::utils::convert_columns; +use super::values::{Column, NuDataFrame}; + +#[derive(Clone)] +pub struct DropNulls; + +impl Command for DropNulls { + fn name(&self) -> &str { + "dfr drop-nulls" + } + + fn usage(&self) -> &str { + "Drops null values in dataframe" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .optional( + "subset", + SyntaxShape::Table, + "subset of columns to drop nulls", + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "drop null values in dataframe", + example: r#"let df = ([[a b]; [1 2] [3 0] [1 2]] | dfr to-df); + let res = ($df.b / $df.b); + let a = ($df | dfr with-column $res --name res); + $a | dfr drop-nulls"#, + result: Some( + NuDataFrame::try_from_columns(vec![ + Column::new("a".to_string(), vec![1.into(), 1.into()]), + Column::new("b".to_string(), vec![2.into(), 2.into()]), + Column::new("res".to_string(), vec![1.into(), 1.into()]), + ]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + Example { + description: "drop null values in dataframe", + example: r#"let s = ([1 2 0 0 3 4] | dfr to-df); + ($s / $s) | dfr drop-nulls"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "div_0_0".to_string(), + vec![1.into(), 1.into(), 1.into(), 1.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let columns: Option> = call.opt(engine_state, stack, 0)?; + + let (subset, col_span) = match columns { + Some(cols) => { + let (agg_string, col_span) = convert_columns(cols, call.head)?; + let agg_string = agg_string + .into_iter() + .map(|col| col.item) + .collect::>(); + (Some(agg_string), col_span) + } + None => (None, Span::unknown()), + }; + + let subset_slice = subset.as_ref().map(|cols| &cols[..]); + + df.as_ref() + .drop_nulls(subset_slice) + .map_err(|e| { + ShellError::SpannedLabeledError("Error dropping nulls".into(), e.to_string(), col_span) + }) + .map(|df| PipelineData::Value(NuDataFrame::dataframe_into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::test_dataframe::test_dataframe; + use super::super::WithColumn; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(DropNulls {}), Box::new(WithColumn {})]) + } +} diff --git a/crates/nu-command/src/dataframe/dtypes.rs b/crates/nu-command/src/dataframe/dtypes.rs index 56da61bf28..ae10f53ba2 100644 --- a/crates/nu-command/src/dataframe/dtypes.rs +++ b/crates/nu-command/src/dataframe/dtypes.rs @@ -10,7 +10,7 @@ pub struct DataTypes; impl Command for DataTypes { fn name(&self) -> &str { - "dataframe dtypes" + "dfr dtypes" } fn usage(&self) -> &str { @@ -24,7 +24,7 @@ impl Command for DataTypes { fn examples(&self) -> Vec { vec![Example { description: "Dataframe dtypes", - example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe dtypes", + example: "[[a b]; [1 2] [3 4]] | dfr to-df | dfr dtypes", result: Some( NuDataFrame::try_from_columns(vec![ Column::new( @@ -101,6 +101,6 @@ mod test { #[test] fn test_examples() { - test_dataframe(DataTypes {}) + test_dataframe(vec![Box::new(DataTypes {})]) } } diff --git a/crates/nu-command/src/dataframe/mod.rs b/crates/nu-command/src/dataframe/mod.rs index 20fe6bcc5d..8830987b1b 100644 --- a/crates/nu-command/src/dataframe/mod.rs +++ b/crates/nu-command/src/dataframe/mod.rs @@ -6,18 +6,24 @@ mod column; mod command; mod describe; mod drop; +mod drop_nulls; mod dtypes; mod open; mod to_df; +mod with_column; + +pub use series::*; pub use append::AppendDF; pub use column::ColumnDF; pub use command::Dataframe; pub use describe::DescribeDF; pub use drop::DropDF; +pub use drop_nulls::DropNulls; pub use dtypes::DataTypes; pub use open::OpenDataFrame; pub use to_df::ToDataFrame; +pub use with_column::WithColumn; use nu_protocol::engine::StateWorkingSet; @@ -31,6 +37,53 @@ pub fn add_dataframe_decls(working_set: &mut StateWorkingSet) { }; } + // Series commands + bind_command!( + AllFalse, + AllTrue, + ArgMax, + ArgMin, + ArgSort, + ArgTrue, + ArgUnique, + Concatenate, + Contains, + Cumulative, + GetDay, + GetHour, + GetMinute, + GetMonth, + GetNanosecond, + GetOrdinal, + GetSecond, + GetWeek, + GetWeekDay, + GetYear, + IsDuplicated, + IsIn, + IsNotNull, + IsNull, + IsUnique, + NNull, + NUnique, + NotSeries, + Rename, + Replace, + ReplaceAll, + Rolling, + SetSeries, + SetWithIndex, + Shift, + StrLengths, + StrSlice, + StrFTime, + ToLowerCase, + ToUpperCase, + Unique, + ValueCount + ); + + // Dataframe commands bind_command!( AppendDF, ColumnDF, @@ -38,8 +91,10 @@ pub fn add_dataframe_decls(working_set: &mut StateWorkingSet) { DataTypes, DescribeDF, DropDF, + DropNulls, OpenDataFrame, - ToDataFrame + ToDataFrame, + WithColumn ); } diff --git a/crates/nu-command/src/dataframe/open.rs b/crates/nu-command/src/dataframe/open.rs index 7c82f101c4..9f8ff47b59 100644 --- a/crates/nu-command/src/dataframe/open.rs +++ b/crates/nu-command/src/dataframe/open.rs @@ -14,7 +14,7 @@ pub struct OpenDataFrame; impl Command for OpenDataFrame { fn name(&self) -> &str { - "dataframe open" + "dfr open" } fn usage(&self) -> &str { @@ -63,7 +63,7 @@ impl Command for OpenDataFrame { fn examples(&self) -> Vec { vec![Example { description: "Takes a file name and creates a dataframe", - example: "dataframe open test.csv", + example: "dfr open test.csv", result: None, }] } diff --git a/crates/nu-command/src/dataframe/series/all_false.rs b/crates/nu-command/src/dataframe/series/all_false.rs new file mode 100644 index 0000000000..1320f97b52 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/all_false.rs @@ -0,0 +1,102 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, Value, +}; + +#[derive(Clone)] +pub struct AllFalse; + +impl Command for AllFalse { + fn name(&self) -> &str { + "dfr all-false" + } + + fn usage(&self) -> &str { + "Returns true if all values are false" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Returns true if all values are false", + example: "[$false $false $false] | dfr to-df | dfr all-false", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "all_false".to_string(), + vec![true.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + Example { + description: "Checks the result from a comparison", + example: r#"let s = ([5 6 2 10] | dfr to-df); + let res = ($s > 9); + $res | dfr all-false"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "all_false".to_string(), + vec![false.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let series = df.as_series(call.head)?; + let bool = series.bool().map_err(|_| { + ShellError::SpannedLabeledError( + "Error converting to bool".into(), + "all-false only works with series of type bool".into(), + call.head, + ) + })?; + + let value = Value::Bool { + val: bool.all_false(), + span: call.head, + }; + + NuDataFrame::try_from_columns(vec![Column::new("all_false".to_string(), vec![value])]) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(AllFalse {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/all_true.rs b/crates/nu-command/src/dataframe/series/all_true.rs new file mode 100644 index 0000000000..251e5a4384 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/all_true.rs @@ -0,0 +1,102 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, Value, +}; + +#[derive(Clone)] +pub struct AllTrue; + +impl Command for AllTrue { + fn name(&self) -> &str { + "dfr all-true" + } + + fn usage(&self) -> &str { + "Returns true if all values are true" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Returns true if all values are true", + example: "[$true $true $true] | dfr to-df | dfr all-true", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "all_true".to_string(), + vec![true.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + Example { + description: "Checks the result from a comparison", + example: r#"let s = ([5 6 2 8] | dfr to-df); + let res = ($s > 9); + $res | dfr all-true"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "all_true".to_string(), + vec![false.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let series = df.as_series(call.head)?; + let bool = series.bool().map_err(|_| { + ShellError::SpannedLabeledError( + "Error converting to bool".into(), + "all-false only works with series of type bool".into(), + call.head, + ) + })?; + + let value = Value::Bool { + val: bool.all_true(), + span: call.head, + }; + + NuDataFrame::try_from_columns(vec![Column::new("all_true".to_string(), vec![value])]) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(AllTrue {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/arg_max.rs b/crates/nu-command/src/dataframe/series/arg_max.rs new file mode 100644 index 0000000000..531e65c16c --- /dev/null +++ b/crates/nu-command/src/dataframe/series/arg_max.rs @@ -0,0 +1,81 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::{IntoSeries, NewChunkedArray, UInt32Chunked}; + +#[derive(Clone)] +pub struct ArgMax; + +impl Command for ArgMax { + fn name(&self) -> &str { + "dfr arg-max" + } + + fn usage(&self) -> &str { + "Return index for max value in series" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns index for max value", + example: "[1 3 2] | dfr to-df | dfr arg-max", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "arg_max".to_string(), + vec![1.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let res = series.arg_max(); + let chunked = match res { + Some(index) => UInt32Chunked::new_from_slice("arg_max", &[index as u32]), + None => UInt32Chunked::new_from_slice("arg_max", &[]), + }; + + let res = chunked.into_series(); + NuDataFrame::try_from_series(vec![res], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(ArgMax {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/arg_min.rs b/crates/nu-command/src/dataframe/series/arg_min.rs new file mode 100644 index 0000000000..4e49a8034c --- /dev/null +++ b/crates/nu-command/src/dataframe/series/arg_min.rs @@ -0,0 +1,81 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::{IntoSeries, NewChunkedArray, UInt32Chunked}; + +#[derive(Clone)] +pub struct ArgMin; + +impl Command for ArgMin { + fn name(&self) -> &str { + "dfr arg-min" + } + + fn usage(&self) -> &str { + "Return index for min value in series" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns index for min value", + example: "[1 3 2] | dfr to-df | dfr arg-min", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "arg_min".to_string(), + vec![0.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let res = series.arg_min(); + let chunked = match res { + Some(index) => UInt32Chunked::new_from_slice("arg_min", &[index as u32]), + None => UInt32Chunked::new_from_slice("arg_min", &[]), + }; + + let res = chunked.into_series(); + NuDataFrame::try_from_series(vec![res], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(ArgMin {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/cumulative.rs b/crates/nu-command/src/dataframe/series/cumulative.rs new file mode 100644 index 0000000000..b957b81e19 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/cumulative.rs @@ -0,0 +1,129 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, Spanned, SyntaxShape, +}; +use polars::prelude::{DataType, IntoSeries}; + +enum CumType { + Min, + Max, + Sum, +} + +impl CumType { + fn from_str(roll_type: &str, span: Span) -> Result { + match roll_type { + "min" => Ok(Self::Min), + "max" => Ok(Self::Max), + "sum" => Ok(Self::Sum), + _ => Err(ShellError::SpannedLabeledErrorHelp( + "Wrong operation".into(), + "Operation not valid for cumulative".into(), + span, + "Allowed values: max, min, sum".into(), + )), + } + } + + fn to_str(&self) -> &'static str { + match self { + CumType::Min => "cum_min", + CumType::Max => "cum_max", + CumType::Sum => "cum_sum", + } + } +} + +#[derive(Clone)] +pub struct Cumulative; + +impl Command for Cumulative { + fn name(&self) -> &str { + "dfr cum" + } + + fn usage(&self) -> &str { + "Cumulative calculation for a series" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("type", SyntaxShape::String, "rolling operation") + .switch("reverse", "Reverse cumulative calculation", Some('r')) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Cumulative sum for a series", + example: "[1 2 3 4 5] | dfr to-df | dfr cum sum", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0_cum_sum".to_string(), + vec![1.into(), 3.into(), 6.into(), 10.into(), 15.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let cum_type: Spanned = call.req(engine_state, stack, 0)?; + let reverse = call.has_flag("reverse"); + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + if let DataType::Object(_) = series.dtype() { + return Err(ShellError::SpannedLabeledError( + "Found object series".into(), + "Series of type object cannot be used for cumulative operation".into(), + call.head, + )); + } + + let cum_type = CumType::from_str(&cum_type.item, cum_type.span)?; + let mut res = match cum_type { + CumType::Max => series.cummax(reverse), + CumType::Min => series.cummin(reverse), + CumType::Sum => series.cumsum(reverse), + }; + + let name = format!("{}_{}", series.name(), cum_type.to_str()); + res.rename(&name); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(Cumulative {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_day.rs b/crates/nu-command/src/dataframe/series/date/get_day.rs new file mode 100644 index 0000000000..3d049edadc --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_day.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetDay; + +impl Command for GetDay { + fn name(&self) -> &str { + "dfr get-day" + } + + fn usage(&self) -> &str { + "Gets day from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns day from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-day"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![4.into(), 4.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.day().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetDay {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_hour.rs b/crates/nu-command/src/dataframe/series/date/get_hour.rs new file mode 100644 index 0000000000..0b9703c718 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_hour.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetHour; + +impl Command for GetHour { + fn name(&self) -> &str { + "dfr get-hour" + } + + fn usage(&self) -> &str { + "Gets hour from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns hour from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-hour"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![16.into(), 16.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.hour().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetHour {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_minute.rs b/crates/nu-command/src/dataframe/series/date/get_minute.rs new file mode 100644 index 0000000000..d5c44eba88 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_minute.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetMinute; + +impl Command for GetMinute { + fn name(&self) -> &str { + "dfr get-minute" + } + + fn usage(&self) -> &str { + "Gets minute from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns minute from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-minute"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![39.into(), 39.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.minute().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetMinute {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_month.rs b/crates/nu-command/src/dataframe/series/date/get_month.rs new file mode 100644 index 0000000000..6f4092e084 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_month.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetMonth; + +impl Command for GetMonth { + fn name(&self) -> &str { + "dfr get-month" + } + + fn usage(&self) -> &str { + "Gets month from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns month from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-month"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![8.into(), 8.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.month().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetMonth {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_nanosecond.rs b/crates/nu-command/src/dataframe/series/date/get_nanosecond.rs new file mode 100644 index 0000000000..a748ae096c --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_nanosecond.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetNanosecond; + +impl Command for GetNanosecond { + fn name(&self) -> &str { + "dfr get-nanosecond" + } + + fn usage(&self) -> &str { + "Gets nanosecond from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns nanosecond from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-nanosecond"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![0.into(), 0.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.nanosecond().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetNanosecond {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_ordinal.rs b/crates/nu-command/src/dataframe/series/date/get_ordinal.rs new file mode 100644 index 0000000000..08d22ab910 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_ordinal.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetOrdinal; + +impl Command for GetOrdinal { + fn name(&self) -> &str { + "dfr get-ordinal" + } + + fn usage(&self) -> &str { + "Gets ordinal from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns ordinal from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-ordinal"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![217.into(), 217.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.ordinal().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetOrdinal {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_second.rs b/crates/nu-command/src/dataframe/series/date/get_second.rs new file mode 100644 index 0000000000..d379488464 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_second.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetSecond; + +impl Command for GetSecond { + fn name(&self) -> &str { + "dfr get-second" + } + + fn usage(&self) -> &str { + "Gets second from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns second from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-second"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![18.into(), 18.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.second().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetSecond {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_week.rs b/crates/nu-command/src/dataframe/series/date/get_week.rs new file mode 100644 index 0000000000..117eb5bf9f --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_week.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetWeek; + +impl Command for GetWeek { + fn name(&self) -> &str { + "dfr get-week" + } + + fn usage(&self) -> &str { + "Gets week from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns week from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-week"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![32.into(), 32.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.week().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetWeek {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_weekday.rs b/crates/nu-command/src/dataframe/series/date/get_weekday.rs new file mode 100644 index 0000000000..9a517fed81 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_weekday.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetWeekDay; + +impl Command for GetWeekDay { + fn name(&self) -> &str { + "dfr get-weekday" + } + + fn usage(&self) -> &str { + "Gets weekday from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns weekday from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-weekday"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![1.into(), 1.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.weekday().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetWeekDay {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/get_year.rs b/crates/nu-command/src/dataframe/series/date/get_year.rs new file mode 100644 index 0000000000..0646a2ae73 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/get_year.rs @@ -0,0 +1,87 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct GetYear; + +impl Command for GetYear { + fn name(&self) -> &str { + "dfr get-year" + } + + fn usage(&self) -> &str { + "Gets year from date" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns year from a date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr get-year"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![2020.into(), 2020.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to datetime type".into(), + e.to_string(), + call.head, + ) + })?; + + let res = casted.year().into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(GetYear {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/date/mod.rs b/crates/nu-command/src/dataframe/series/date/mod.rs new file mode 100644 index 0000000000..fbbb75f5d3 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/date/mod.rs @@ -0,0 +1,21 @@ +mod get_day; +mod get_hour; +mod get_minute; +mod get_month; +mod get_nanosecond; +mod get_ordinal; +mod get_second; +mod get_week; +mod get_weekday; +mod get_year; + +pub use get_day::GetDay; +pub use get_hour::GetHour; +pub use get_minute::GetMinute; +pub use get_month::GetMonth; +pub use get_nanosecond::GetNanosecond; +pub use get_ordinal::GetOrdinal; +pub use get_second::GetSecond; +pub use get_week::GetWeek; +pub use get_weekday::GetWeekDay; +pub use get_year::GetYear; diff --git a/crates/nu-command/src/dataframe/series/indexes/arg_sort.rs b/crates/nu-command/src/dataframe/series/indexes/arg_sort.rs new file mode 100644 index 0000000000..6f155aa00c --- /dev/null +++ b/crates/nu-command/src/dataframe/series/indexes/arg_sort.rs @@ -0,0 +1,95 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct ArgSort; + +impl Command for ArgSort { + fn name(&self) -> &str { + "dfr arg-sort" + } + + fn usage(&self) -> &str { + "Returns indexes for a sorted series" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .switch("reverse", "reverse order", Some('r')) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Returns indexes for a sorted series", + example: "[1 2 2 3 3] | dfr to-df | dfr arg-sort", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "arg_sort".to_string(), + vec![0.into(), 1.into(), 2.into(), 3.into(), 4.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + Example { + description: "Returns indexes for a sorted series", + example: "[1 2 2 3 3] | dfr to-df | dfr arg-sort -r", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "arg_sort".to_string(), + vec![3.into(), 4.into(), 1.into(), 2.into(), 0.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let mut res = df + .as_series(call.head)? + .argsort(call.has_flag("reverse")) + .into_series(); + res.rename("arg_sort"); + + NuDataFrame::try_from_series(vec![res], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(ArgSort {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/indexes/arg_true.rs b/crates/nu-command/src/dataframe/series/indexes/arg_true.rs new file mode 100644 index 0000000000..345044d0d6 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/indexes/arg_true.rs @@ -0,0 +1,85 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct ArgTrue; + +impl Command for ArgTrue { + fn name(&self) -> &str { + "dfr arg-true" + } + + fn usage(&self) -> &str { + "Returns indexes where values are true" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns indexes where values are true", + example: "[$false $true $false] | dfr to-df | dfr arg-true", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "arg_true".to_string(), + vec![1.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let series = df.as_series(call.head)?; + let bool = series.bool().map_err(|_| { + ShellError::SpannedLabeledError( + "Error converting to bool".into(), + "all-false only works with series of type bool".into(), + call.head, + ) + })?; + + let mut res = bool.arg_true().into_series(); + res.rename("arg_true"); + + NuDataFrame::try_from_series(vec![res], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(ArgTrue {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/indexes/arg_unique.rs b/crates/nu-command/src/dataframe/series/indexes/arg_unique.rs new file mode 100644 index 0000000000..fe711cca2b --- /dev/null +++ b/crates/nu-command/src/dataframe/series/indexes/arg_unique.rs @@ -0,0 +1,86 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct ArgUnique; + +impl Command for ArgUnique { + fn name(&self) -> &str { + "dfr arg-unique" + } + + fn usage(&self) -> &str { + "Returns indexes for unique values" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns indexes for unique values", + example: "[1 2 2 3 3] | dfr to-df | dfr arg-unique", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "arg_unique".to_string(), + vec![0.into(), 1.into(), 3.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let mut res = df + .as_series(call.head)? + .arg_unique() + .map_err(|e| { + ShellError::SpannedLabeledError( + "Error extracting unique values".into(), + e.to_string(), + call.head, + ) + })? + .into_series(); + res.rename("arg_unique"); + + NuDataFrame::try_from_series(vec![res], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(ArgUnique {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/indexes/mod.rs b/crates/nu-command/src/dataframe/series/indexes/mod.rs new file mode 100644 index 0000000000..c0af8c8653 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/indexes/mod.rs @@ -0,0 +1,9 @@ +mod arg_sort; +mod arg_true; +mod arg_unique; +mod set_with_idx; + +pub use arg_sort::ArgSort; +pub use arg_true::ArgTrue; +pub use arg_unique::ArgUnique; +pub use set_with_idx::SetWithIndex; diff --git a/crates/nu-command/src/dataframe/series/indexes/set_with_idx.rs b/crates/nu-command/src/dataframe/series/indexes/set_with_idx.rs new file mode 100644 index 0000000000..60b542f0c8 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/indexes/set_with_idx.rs @@ -0,0 +1,179 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Value, +}; +use polars::prelude::{ChunkSet, DataType, IntoSeries}; + +#[derive(Clone)] +pub struct SetWithIndex; + +impl Command for SetWithIndex { + fn name(&self) -> &str { + "dfr set-with-idx" + } + + fn usage(&self) -> &str { + "Sets value in the given index" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("value", SyntaxShape::Any, "value to be inserted in series") + .required_named( + "indices", + SyntaxShape::Any, + "list of indices indicating where to set the value", + Some('i'), + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Set value in selected rows from series", + example: r#"let series = ([4 1 5 2 4 3] | dfr to-df); + let indices = ([0 2] | dfr to-df); + $series | dfr set-with-idx 6 -i $indices"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![6.into(), 1.into(), 6.into(), 2.into(), 4.into(), 3.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let value: Value = call.req(engine_state, stack, 0)?; + + let indices_value: Value = call + .get_flag(engine_state, stack, "indices")? + .expect("required named value"); + let indices_span = indices_value.span()?; + let indices = NuDataFrame::try_from_value(indices_value)?.as_series(indices_span)?; + + let casted = match indices.dtype() { + DataType::UInt32 | DataType::UInt64 | DataType::Int32 | DataType::Int64 => { + indices.as_ref().cast(&DataType::UInt32).map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting indices".into(), + e.to_string(), + indices_span, + ) + }) + } + _ => Err(ShellError::SpannedLabeledErrorHelp( + "Incorrect type".into(), + "Series with incorrect type".into(), + indices_span, + "Consider using a Series with type int type".into(), + )), + }?; + + let indices = casted + .u32() + .map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting indices".into(), + e.to_string(), + indices_span, + ) + })? + .into_iter() + .filter_map(|val| val.map(|v| v as usize)); + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let res = match value { + Value::Int { val, span } => { + let chunked = series.i64().map_err(|e| { + ShellError::SpannedLabeledError("Error casting to i64".into(), e.to_string(), span) + })?; + + let res = chunked.set_at_idx(indices, Some(val)).map_err(|e| { + ShellError::SpannedLabeledError("Error setting value".into(), e.to_string(), span) + })?; + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + } + Value::Float { val, span } => { + let chunked = series.as_ref().f64().map_err(|e| { + ShellError::SpannedLabeledError("Error casting to f64".into(), e.to_string(), span) + })?; + + let res = chunked.set_at_idx(indices, Some(val)).map_err(|e| { + ShellError::SpannedLabeledError("Error setting value".into(), e.to_string(), span) + })?; + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + } + Value::String { val, span } => { + let chunked = series.as_ref().utf8().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to string".into(), + e.to_string(), + span, + ) + })?; + + let res = chunked + .set_at_idx(indices, Some(val.as_ref())) + .map_err(|e| { + ShellError::SpannedLabeledError( + "Error setting value".into(), + e.to_string(), + span, + ) + })?; + + let mut res = res.into_series(); + res.rename("string"); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + } + _ => Err(ShellError::SpannedLabeledError( + "Incorrect value type".into(), + format!( + "this value cannot be set in a series of type '{}'", + series.dtype() + ), + value.span()?, + )), + }; + + res.map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(SetWithIndex {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/masks/is_duplicated.rs b/crates/nu-command/src/dataframe/series/masks/is_duplicated.rs new file mode 100644 index 0000000000..df111c297f --- /dev/null +++ b/crates/nu-command/src/dataframe/series/masks/is_duplicated.rs @@ -0,0 +1,95 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct IsDuplicated; + +impl Command for IsDuplicated { + fn name(&self) -> &str { + "dfr is-duplicated" + } + + fn usage(&self) -> &str { + "Creates mask indicating duplicated values" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Create mask indicating duplicated values", + example: "[5 6 6 6 8 8 8] | dfr to-df | dfr is-duplicated", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "is_duplicated".to_string(), + vec![ + false.into(), + true.into(), + true.into(), + true.into(), + true.into(), + true.into(), + true.into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let mut res = df + .as_series(call.head)? + .is_duplicated() + .map_err(|e| { + ShellError::SpannedLabeledError( + "Error finding duplicates".into(), + e.to_string(), + call.head, + ) + })? + .into_series(); + + res.rename("is_duplicated"); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(IsDuplicated {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/masks/is_in.rs b/crates/nu-command/src/dataframe/series/masks/is_in.rs new file mode 100644 index 0000000000..57478c72c8 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/masks/is_in.rs @@ -0,0 +1,104 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Value, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct IsIn; + +impl Command for IsIn { + fn name(&self) -> &str { + "dfr is-in" + } + + fn usage(&self) -> &str { + "Checks if elements from a series are contained in right series" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("other", SyntaxShape::Any, "right series") + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Checks if elements from a series are contained in right series", + example: r#"let other = ([1 3 6] | dfr to-df); + [5 6 6 6 8 8 8] | dfr to-df | dfr is-in $other"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "is_in".to_string(), + vec![ + false.into(), + true.into(), + true.into(), + true.into(), + false.into(), + false.into(), + false.into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let other_value: Value = call.req(engine_state, stack, 0)?; + let other_span = other_value.span()?; + let other_df = NuDataFrame::try_from_value(other_value)?; + let other = other_df.as_series(other_span)?; + + let mut res = df + .as_series(call.head)? + .is_in(&other) + .map_err(|e| { + ShellError::SpannedLabeledError( + "Error finding in other".into(), + e.to_string(), + call.head, + ) + })? + .into_series(); + + res.rename("is_in"); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(IsIn {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/masks/is_not_null.rs b/crates/nu-command/src/dataframe/series/masks/is_not_null.rs new file mode 100644 index 0000000000..42f6101c97 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/masks/is_not_null.rs @@ -0,0 +1,78 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct IsNotNull; + +impl Command for IsNotNull { + fn name(&self) -> &str { + "dfr is-not-null" + } + + fn usage(&self) -> &str { + "Creates mask where value is not null" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Create mask where values are not null", + example: r#"let s = ([5 6 0 8] | dfr to-df); + let res = ($s / $s); + $res | dfr is-not-null"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "is_not_null".to_string(), + vec![true.into(), true.into(), false.into(), true.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let mut res = df.as_series(call.head)?.is_not_null(); + res.rename("is_not_null"); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(IsNotNull {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/masks/is_null.rs b/crates/nu-command/src/dataframe/series/masks/is_null.rs new file mode 100644 index 0000000000..59f37c836e --- /dev/null +++ b/crates/nu-command/src/dataframe/series/masks/is_null.rs @@ -0,0 +1,78 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct IsNull; + +impl Command for IsNull { + fn name(&self) -> &str { + "dfr is-null" + } + + fn usage(&self) -> &str { + "Creates mask where value is null" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Create mask where values are null", + example: r#"let s = ([5 6 0 8] | dfr to-df); + let res = ($s / $s); + $res | dfr is-null"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "is_null".to_string(), + vec![false.into(), false.into(), true.into(), false.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let mut res = df.as_series(call.head)?.is_null(); + res.rename("is_null"); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(IsNull {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/masks/is_unique.rs b/crates/nu-command/src/dataframe/series/masks/is_unique.rs new file mode 100644 index 0000000000..971f501525 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/masks/is_unique.rs @@ -0,0 +1,90 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct IsUnique; + +impl Command for IsUnique { + fn name(&self) -> &str { + "dfr is-unique" + } + + fn usage(&self) -> &str { + "Creates mask indicating unique values" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Create mask indicating unique values", + example: "[5 6 6 6 8 8 8] | dfr to-df | dfr is-unique", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "is_unique".to_string(), + vec![ + true.into(), + false.into(), + false.into(), + false.into(), + false.into(), + false.into(), + false.into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let mut res = df.as_series(call.head)?.is_unique().map_err(|e| { + ShellError::SpannedLabeledError( + "Error finding unique values".into(), + e.to_string(), + call.head, + ) + })?; + res.rename("is_unique"); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(IsUnique {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/masks/mod.rs b/crates/nu-command/src/dataframe/series/masks/mod.rs new file mode 100644 index 0000000000..80c98b5ef0 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/masks/mod.rs @@ -0,0 +1,15 @@ +mod is_duplicated; +mod is_in; +mod is_not_null; +mod is_null; +mod is_unique; +mod not; +mod set; + +pub use is_duplicated::IsDuplicated; +pub use is_in::IsIn; +pub use is_not_null::IsNotNull; +pub use is_null::IsNull; +pub use is_unique::IsUnique; +pub use not::NotSeries; +pub use set::SetSeries; diff --git a/crates/nu-command/src/dataframe/series/masks/not.rs b/crates/nu-command/src/dataframe/series/masks/not.rs new file mode 100644 index 0000000000..cad0e4d388 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/masks/not.rs @@ -0,0 +1,82 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +use std::ops::Not; + +#[derive(Clone)] +pub struct NotSeries; + +impl Command for NotSeries { + fn name(&self) -> &str { + "dfr not" + } + + fn usage(&self) -> &str { + "Inverts boolean mask" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Inverts boolean mask", + example: "[$true $false $true] | dfr to-df | dfr not", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![false.into(), true.into(), false.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let bool = series.bool().map_err(|e| { + ShellError::SpannedLabeledError("Error inverting mask".into(), e.to_string(), call.head) + })?; + + let res = bool.not(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(NotSeries {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/masks/set.rs b/crates/nu-command/src/dataframe/series/masks/set.rs new file mode 100644 index 0000000000..610ed55ff2 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/masks/set.rs @@ -0,0 +1,163 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Value, +}; +use polars::prelude::{ChunkSet, DataType, IntoSeries}; + +#[derive(Clone)] +pub struct SetSeries; + +impl Command for SetSeries { + fn name(&self) -> &str { + "dfr set" + } + + fn usage(&self) -> &str { + "Sets value where given mask is true" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("value", SyntaxShape::Any, "value to be inserted in series") + .required_named( + "mask", + SyntaxShape::Any, + "mask indicating insertions", + Some('m'), + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Shifts the values by a given period", + example: r#"let s = ([1 2 2 3 3] | dfr to-df | dfr shift 2); + let mask = ($s | dfr is-null); + $s | dfr set 0 --mask $mask"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![0.into(), 0.into(), 1.into(), 2.into(), 2.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let value: Value = call.req(engine_state, stack, 0)?; + + let mask_value: Value = call + .get_flag(engine_state, stack, "mask")? + .expect("required named value"); + let mask_span = mask_value.span()?; + let mask = NuDataFrame::try_from_value(mask_value)?.as_series(mask_span)?; + + let bool_mask = match mask.dtype() { + DataType::Boolean => mask.bool().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to bool".into(), + e.to_string(), + mask_span, + ) + }), + _ => Err(ShellError::SpannedLabeledError( + "Incorrect type".into(), + "can only use bool series as mask".into(), + mask_span, + )), + }?; + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let res = match value { + Value::Int { val, span } => { + let chunked = series.i64().map_err(|e| { + ShellError::SpannedLabeledError("Error casting to i64".into(), e.to_string(), span) + })?; + + let res = chunked.set(bool_mask, Some(val)).map_err(|e| { + ShellError::SpannedLabeledError("Error setting value".into(), e.to_string(), span) + })?; + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + } + Value::Float { val, span } => { + let chunked = series.as_ref().f64().map_err(|e| { + ShellError::SpannedLabeledError("Error casting to f64".into(), e.to_string(), span) + })?; + + let res = chunked.set(bool_mask, Some(val)).map_err(|e| { + ShellError::SpannedLabeledError("Error setting value".into(), e.to_string(), span) + })?; + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + } + Value::String { val, span } => { + let chunked = series.as_ref().utf8().map_err(|e| { + ShellError::SpannedLabeledError( + "Error casting to string".into(), + e.to_string(), + span, + ) + })?; + + let res = chunked.set(bool_mask, Some(val.as_ref())).map_err(|e| { + ShellError::SpannedLabeledError("Error setting value".into(), e.to_string(), span) + })?; + + let mut res = res.into_series(); + res.rename("string"); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + } + _ => Err(ShellError::SpannedLabeledError( + "Incorrect value type".into(), + format!( + "this value cannot be set in a series of type '{}'", + series.dtype() + ), + value.span()?, + )), + }; + + res.map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::super::super::{IsNull, Shift}; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![ + Box::new(SetSeries {}), + Box::new(IsNull {}), + Box::new(Shift {}), + ]) + } +} diff --git a/crates/nu-command/src/dataframe/series/mod.rs b/crates/nu-command/src/dataframe/series/mod.rs index 8b13789179..9fa3b2221c 100644 --- a/crates/nu-command/src/dataframe/series/mod.rs +++ b/crates/nu-command/src/dataframe/series/mod.rs @@ -1 +1,37 @@ +mod date; +pub use date::*; +mod string; +pub use string::*; + +mod masks; +pub use masks::*; + +mod indexes; +pub use indexes::*; + +mod all_false; +mod all_true; +mod arg_max; +mod arg_min; +mod cumulative; +mod n_null; +mod n_unique; +mod rename; +mod rolling; +mod shift; +mod unique; +mod value_counts; + +pub use all_false::AllFalse; +pub use all_true::AllTrue; +pub use arg_max::ArgMax; +pub use arg_min::ArgMin; +pub use cumulative::Cumulative; +pub use n_null::NNull; +pub use n_unique::NUnique; +pub use rename::Rename; +pub use rolling::Rolling; +pub use shift::Shift; +pub use unique::Unique; +pub use value_counts::ValueCount; diff --git a/crates/nu-command/src/dataframe/series/n_null.rs b/crates/nu-command/src/dataframe/series/n_null.rs new file mode 100644 index 0000000000..4c21a120d0 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/n_null.rs @@ -0,0 +1,79 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, Value, +}; + +#[derive(Clone)] +pub struct NNull; + +impl Command for NNull { + fn name(&self) -> &str { + "dfr count-null" + } + + fn usage(&self) -> &str { + "Counts null values" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Counts null values", + example: r#"let s = ([1 1 0 0 3 3 4] | dfr to-df); + ($s / $s) | dfr count-null"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "count_null".to_string(), + vec![2.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let res = df.as_series(call.head)?.null_count(); + let value = Value::Int { + val: res as i64, + span: call.head, + }; + + NuDataFrame::try_from_columns(vec![Column::new("count_null".to_string(), vec![value])]) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(NNull {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/n_unique.rs b/crates/nu-command/src/dataframe/series/n_unique.rs new file mode 100644 index 0000000000..94910c63a2 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/n_unique.rs @@ -0,0 +1,85 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, Value, +}; + +#[derive(Clone)] +pub struct NUnique; + +impl Command for NUnique { + fn name(&self) -> &str { + "dfr count-unique" + } + + fn usage(&self) -> &str { + "Counts unique values" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Counts unique values", + example: "[1 1 2 2 3 3 4] | dfr to-df | dfr count-unique", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "count_unique".to_string(), + vec![4.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let res = df.as_series(call.head)?.n_unique().map_err(|e| { + ShellError::SpannedLabeledError( + "Error counting unique values".into(), + e.to_string(), + call.head, + ) + })?; + + let value = Value::Int { + val: res as i64, + span: call.head, + }; + + NuDataFrame::try_from_columns(vec![Column::new("count_unique".to_string(), vec![value])]) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(NUnique {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/rename.rs b/crates/nu-command/src/dataframe/series/rename.rs new file mode 100644 index 0000000000..8fb9ca2588 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/rename.rs @@ -0,0 +1,79 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, +}; + +#[derive(Clone)] +pub struct Rename; + +impl Command for Rename { + fn name(&self) -> &str { + "dfr rename" + } + + fn usage(&self) -> &str { + "Renames a series" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("name", SyntaxShape::String, "new series name") + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Renames a series", + example: "[5 6 7 8] | dfr to-df | dfr rename new_name", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "new_name".to_string(), + vec![5.into(), 6.into(), 7.into(), 8.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let name: String = call.req(engine_state, stack, 0)?; + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let mut series = df.as_series(call.head)?; + series.rename(&name); + + NuDataFrame::try_from_series(vec![series], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(Rename {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/rolling.rs b/crates/nu-command/src/dataframe/series/rolling.rs new file mode 100644 index 0000000000..4a5d7b96db --- /dev/null +++ b/crates/nu-command/src/dataframe/series/rolling.rs @@ -0,0 +1,163 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, Spanned, SyntaxShape, +}; +use polars::prelude::{DataType, IntoSeries, RollingOptions}; + +enum RollType { + Min, + Max, + Sum, + Mean, +} + +impl RollType { + fn from_str(roll_type: &str, span: Span) -> Result { + match roll_type { + "min" => Ok(Self::Min), + "max" => Ok(Self::Max), + "sum" => Ok(Self::Sum), + "mean" => Ok(Self::Mean), + _ => Err(ShellError::SpannedLabeledErrorHelp( + "Wrong operation".into(), + "Operation not valid for cumulative".into(), + span, + "Allowed values: min, max, sum, mean".into(), + )), + } + } + + fn to_str(&self) -> &'static str { + match self { + RollType::Min => "rolling_min", + RollType::Max => "rolling_max", + RollType::Sum => "rolling_sum", + RollType::Mean => "rolling_mean", + } + } +} + +#[derive(Clone)] +pub struct Rolling; + +impl Command for Rolling { + fn name(&self) -> &str { + "dfr rolling" + } + + fn usage(&self) -> &str { + "Rolling calculation for a series" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("type", SyntaxShape::String, "rolling operation") + .required("window", SyntaxShape::Int, "Window size for rolling") + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Rolling sum for a series", + example: "[1 2 3 4 5] | dfr to-df | dfr rolling sum 2 | dfr drop-nulls", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0_rolling_sum".to_string(), + vec![3.into(), 5.into(), 7.into(), 9.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + Example { + description: "Rolling max for a series", + example: "[1 2 3 4 5] | dfr to-df | dfr rolling max 2 | dfr drop-nulls", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0_rolling_max".to_string(), + vec![2.into(), 3.into(), 4.into(), 5.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let roll_type: Spanned = call.req(engine_state, stack, 0)?; + let window_size: usize = call.req(engine_state, stack, 1)?; + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + if let DataType::Object(_) = series.dtype() { + return Err(ShellError::SpannedLabeledError( + "Found object series".into(), + "Series of type object cannot be used for rolling operation".into(), + call.head, + )); + } + + let roll_type = RollType::from_str(&roll_type.item, roll_type.span)?; + + let rolling_opts = RollingOptions { + window_size, + min_periods: window_size, + weights: None, + center: false, + }; + let res = match roll_type { + RollType::Max => series.rolling_max(rolling_opts), + RollType::Min => series.rolling_min(rolling_opts), + RollType::Sum => series.rolling_sum(rolling_opts), + RollType::Mean => series.rolling_mean(rolling_opts), + }; + + let mut res = res.map_err(|e| { + ShellError::SpannedLabeledError( + "Error calculating rolling values".into(), + e.to_string(), + call.head, + ) + })?; + + let name = format!("{}_{}", series.name(), roll_type.to_str()); + res.rename(&name); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::super::super::DropNulls; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(Rolling {}), Box::new(DropNulls {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/shift.rs b/crates/nu-command/src/dataframe/series/shift.rs new file mode 100644 index 0000000000..f50c4ac6e1 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/shift.rs @@ -0,0 +1,79 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, +}; + +#[derive(Clone)] +pub struct Shift; + +impl Command for Shift { + fn name(&self) -> &str { + "dfr shift" + } + + fn usage(&self) -> &str { + "Shifts the values by a given period" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("period", SyntaxShape::Int, "shift period") + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Shifts the values by a given period", + example: "[1 2 2 3 3] | dfr to-df | dfr shift 2 | dfr drop-nulls", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![1.into(), 2.into(), 2.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let period: i64 = call.req(engine_state, stack, 0)?; + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?.shift(period); + + NuDataFrame::try_from_series(vec![series], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::super::super::DropNulls; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(Shift {}), Box::new(DropNulls {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/string/concatenate.rs b/crates/nu-command/src/dataframe/series/string/concatenate.rs new file mode 100644 index 0000000000..6abb80e363 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/concatenate.rs @@ -0,0 +1,111 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Value, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct Concatenate; + +impl Command for Concatenate { + fn name(&self) -> &str { + "dfr concatenate" + } + + fn usage(&self) -> &str { + "Concatenates strings with other array" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required( + "other", + SyntaxShape::Any, + "Other array with string to be concatenated", + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Concatenate string", + example: r#"let other = ([za xs cd] | dfr to-df); + [abc abc abc] | dfr to-df | dfr concatenate $other"#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![ + "abcza".to_string().into(), + "abcxs".to_string().into(), + "abccd".to_string().into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + + let other: Value = call.req(engine_state, stack, 0)?; + let other_span = other.span()?; + let other_df = NuDataFrame::try_from_value(other)?; + + let other_series = other_df.as_series(other_span)?; + let other_chunked = other_series.utf8().map_err(|e| { + ShellError::SpannedLabeledError( + "The concatenate only with string columns".into(), + e.to_string(), + other_span, + ) + })?; + + let series = df.as_series(call.head)?; + let chunked = series.utf8().map_err(|e| { + ShellError::SpannedLabeledError( + "The concatenate only with string columns".into(), + e.to_string(), + call.head, + ) + })?; + + let mut res = chunked.concat(other_chunked); + + res.rename(series.name()); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(Concatenate {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/string/contains.rs b/crates/nu-command/src/dataframe/series/string/contains.rs new file mode 100644 index 0000000000..605bd027d6 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/contains.rs @@ -0,0 +1,98 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct Contains; + +impl Command for Contains { + fn name(&self) -> &str { + "dfr contains" + } + + fn usage(&self) -> &str { + "Checks if a pattern is contained in a string" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required( + "pattern", + SyntaxShape::String, + "Regex pattern to be searched", + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns boolean indicating if pattern was found", + example: "[abc acb acb] | dfr to-df | dfr contains ab", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![true.into(), false.into(), false.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let pattern: String = call.req(engine_state, stack, 0)?; + + let series = df.as_series(call.head)?; + let chunked = series.utf8().map_err(|e| { + ShellError::SpannedLabeledError( + "The contains command only with string columns".into(), + e.to_string(), + call.head, + ) + })?; + + let res = chunked.contains(&pattern).map_err(|e| { + ShellError::SpannedLabeledError( + "Error searching in series".into(), + e.to_string(), + call.head, + ) + })?; + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(Contains {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/string/mod.rs b/crates/nu-command/src/dataframe/series/string/mod.rs new file mode 100644 index 0000000000..f2fa19cbaf --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/mod.rs @@ -0,0 +1,19 @@ +mod concatenate; +mod contains; +mod replace; +mod replace_all; +mod str_lengths; +mod str_slice; +mod strftime; +mod to_lowercase; +mod to_uppercase; + +pub use concatenate::Concatenate; +pub use contains::Contains; +pub use replace::Replace; +pub use replace_all::ReplaceAll; +pub use str_lengths::StrLengths; +pub use str_slice::StrSlice; +pub use strftime::StrFTime; +pub use to_lowercase::ToLowerCase; +pub use to_uppercase::ToUpperCase; diff --git a/crates/nu-command/src/dataframe/series/string/replace.rs b/crates/nu-command/src/dataframe/series/string/replace.rs new file mode 100644 index 0000000000..43437503a4 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/replace.rs @@ -0,0 +1,116 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct Replace; + +impl Command for Replace { + fn name(&self) -> &str { + "dfr replace" + } + + fn usage(&self) -> &str { + "Replace the leftmost (sub)string by a regex pattern" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required_named( + "pattern", + SyntaxShape::String, + "Regex pattern to be matched", + Some('p'), + ) + .required_named( + "replace", + SyntaxShape::String, + "replacing string", + Some('r'), + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Replaces string", + example: "[abc abc abc] | dfr to-df | dfr replace -p ab -r AB", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![ + "ABc".to_string().into(), + "ABc".to_string().into(), + "ABc".to_string().into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let pattern: String = call + .get_flag(engine_state, stack, "pattern")? + .expect("required value"); + let replace: String = call + .get_flag(engine_state, stack, "replace")? + .expect("required value"); + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + let chunked = series.utf8().map_err(|e| { + ShellError::SpannedLabeledError( + "Error convertion to string".into(), + e.to_string(), + call.head, + ) + })?; + + let mut res = chunked.replace(&pattern, &replace).map_err(|e| { + ShellError::SpannedLabeledError( + "Error finding pattern other".into(), + e.to_string(), + call.head, + ) + })?; + + res.rename(series.name()); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(Replace {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/string/replace_all.rs b/crates/nu-command/src/dataframe/series/string/replace_all.rs new file mode 100644 index 0000000000..63abefe934 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/replace_all.rs @@ -0,0 +1,116 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct ReplaceAll; + +impl Command for ReplaceAll { + fn name(&self) -> &str { + "dfr replace-all" + } + + fn usage(&self) -> &str { + "Replace all (sub)strings by a regex pattern" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required_named( + "pattern", + SyntaxShape::String, + "Regex pattern to be matched", + Some('p'), + ) + .required_named( + "replace", + SyntaxShape::String, + "replacing string", + Some('r'), + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Replaces string", + example: "[abac abac abac] | dfr to-df | dfr replace-all -p a -r A", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![ + "AbAc".to_string().into(), + "AbAc".to_string().into(), + "AbAc".to_string().into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let pattern: String = call + .get_flag(engine_state, stack, "pattern")? + .expect("required value"); + let replace: String = call + .get_flag(engine_state, stack, "replace")? + .expect("required value"); + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + let chunked = series.utf8().map_err(|e| { + ShellError::SpannedLabeledError( + "Error convertion to string".into(), + e.to_string(), + call.head, + ) + })?; + + let mut res = chunked.replace_all(&pattern, &replace).map_err(|e| { + ShellError::SpannedLabeledError( + "Error finding pattern other".into(), + e.to_string(), + call.head, + ) + })?; + + res.rename(series.name()); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(ReplaceAll {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/string/str_lengths.rs b/crates/nu-command/src/dataframe/series/string/str_lengths.rs new file mode 100644 index 0000000000..a13b03f4a6 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/str_lengths.rs @@ -0,0 +1,85 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct StrLengths; + +impl Command for StrLengths { + fn name(&self) -> &str { + "dfr str-lengths" + } + + fn usage(&self) -> &str { + "Get lengths of all strings" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns string lengths", + example: "[a ab abc] | dfr to-df | dfr str-lengths", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![1.into(), 2.into(), 3.into()], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let chunked = series.utf8().map_err(|e| { + ShellError::SpannedLabeledErrorHelp( + "Error casting to string".into(), + e.to_string(), + call.head, + "The str-lengths command can only be used with string columns".into(), + ) + })?; + + let res = chunked.as_ref().str_lengths().into_series(); + + NuDataFrame::try_from_series(vec![res], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(StrLengths {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/string/str_slice.rs b/crates/nu-command/src/dataframe/series/string/str_slice.rs new file mode 100644 index 0000000000..8240d24f71 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/str_slice.rs @@ -0,0 +1,101 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct StrSlice; + +impl Command for StrSlice { + fn name(&self) -> &str { + "dfr str-slice" + } + + fn usage(&self) -> &str { + "Slices the string from the start position until the selected length" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("start", SyntaxShape::Int, "start of slice") + .named("length", SyntaxShape::Int, "optional length", Some('l')) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Creates slices from the strings", + example: "[abcded abc321 abc123] | dfr to-df | dfr str-slice 1 -l 2", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![ + "bc".to_string().into(), + "bc".to_string().into(), + "bc".to_string().into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let start: i64 = call.req(engine_state, stack, 0)?; + + let length: Option = call.get_flag(engine_state, stack, "length")?; + let length = length.map(|v| v as u64); + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let chunked = series.utf8().map_err(|e| { + ShellError::SpannedLabeledErrorHelp( + "Error casting to string".into(), + e.to_string(), + call.head, + "The str-slice command can only be used with string columns".into(), + ) + })?; + + let mut res = chunked.str_slice(start, length).map_err(|e| { + ShellError::SpannedLabeledError("Error slicing series".into(), e.to_string(), call.head) + })?; + res.rename(series.name()); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(StrSlice {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/string/strftime.rs b/crates/nu-command/src/dataframe/series/string/strftime.rs new file mode 100644 index 0000000000..262d78101b --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/strftime.rs @@ -0,0 +1,96 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct StrFTime; + +impl Command for StrFTime { + fn name(&self) -> &str { + "dfr strftime" + } + + fn usage(&self) -> &str { + "Formats date based on string rule" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("fmt", SyntaxShape::String, "Format rule") + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Formats date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime -z 'UTC'); + let df = ([$dt $dt] | dfr to-df); + $df | dfr strftime "%Y/%m/%d""#, + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![ + "2020/08/04".to_string().into(), + "2020/08/04".to_string().into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let fmt: String = call.req(engine_state, stack, 0)?; + + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.datetime().map_err(|e| { + ShellError::SpannedLabeledErrorHelp( + "Error casting to date".into(), + e.to_string(), + call.head, + "The str-slice command can only be used with string columns".into(), + ) + })?; + + let res = casted.strftime(&fmt).into_series(); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::super::IntoDatetime; + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(StrFTime {}), Box::new(IntoDatetime {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/string/to_lowercase.rs b/crates/nu-command/src/dataframe/series/string/to_lowercase.rs new file mode 100644 index 0000000000..d396bd99a5 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/to_lowercase.rs @@ -0,0 +1,90 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct ToLowerCase; + +impl Command for ToLowerCase { + fn name(&self) -> &str { + "dfr to-lowercase" + } + + fn usage(&self) -> &str { + "Lowercase the strings in the column" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Modifies strings to lowercase", + example: "[Abc aBc abC] | dfr to-df | dfr to-lowercase", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![ + "abc".to_string().into(), + "abc".to_string().into(), + "abc".to_string().into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.utf8().map_err(|e| { + ShellError::SpannedLabeledErrorHelp( + "Error casting to string".into(), + e.to_string(), + call.head, + "The str-slice command can only be used with string columns".into(), + ) + })?; + + let mut res = casted.to_lowercase(); + res.rename(series.name()); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(ToLowerCase {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/string/to_uppercase.rs b/crates/nu-command/src/dataframe/series/string/to_uppercase.rs new file mode 100644 index 0000000000..a7bcbf72b8 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/string/to_uppercase.rs @@ -0,0 +1,90 @@ +use super::super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct ToUpperCase; + +impl Command for ToUpperCase { + fn name(&self) -> &str { + "dfr to-uppercase" + } + + fn usage(&self) -> &str { + "Uppercase the strings in the column" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Modifies strings to uppercase", + example: "[Abc aBc abC] | dfr to-df | dfr to-uppercase", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new( + "0".to_string(), + vec![ + "ABC".to_string().into(), + "ABC".to_string().into(), + "ABC".to_string().into(), + ], + )]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let casted = series.utf8().map_err(|e| { + ShellError::SpannedLabeledErrorHelp( + "Error casting to string".into(), + e.to_string(), + call.head, + "The str-slice command can only be used with string columns".into(), + ) + })?; + + let mut res = casted.to_uppercase(); + res.rename(series.name()); + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(ToUpperCase {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/unique.rs b/crates/nu-command/src/dataframe/series/unique.rs new file mode 100644 index 0000000000..1c23124c38 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/unique.rs @@ -0,0 +1,80 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; +use polars::prelude::IntoSeries; + +#[derive(Clone)] +pub struct Unique; + +impl Command for Unique { + fn name(&self) -> &str { + "dfr unique" + } + + fn usage(&self) -> &str { + "Returns unique values from a series" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Returns unique values from a series", + example: "[2 2 2 2 2] | dfr to-df | dfr unique", + result: Some( + NuDataFrame::try_from_columns(vec![Column::new("0".to_string(), vec![2.into()])]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let res = series.unique().map_err(|e| { + ShellError::SpannedLabeledErrorHelp( + "Error calculating unique values".into(), + e.to_string(), + call.head, + "The str-slice command can only be used with string columns".into(), + ) + })?; + + NuDataFrame::try_from_series(vec![res.into_series()], call.head) + .map(|df| PipelineData::Value(NuDataFrame::into_value(df, call.head), None)) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(Unique {})]) + } +} diff --git a/crates/nu-command/src/dataframe/series/value_counts.rs b/crates/nu-command/src/dataframe/series/value_counts.rs new file mode 100644 index 0000000000..42d5737751 --- /dev/null +++ b/crates/nu-command/src/dataframe/series/value_counts.rs @@ -0,0 +1,84 @@ +use super::super::values::{Column, NuDataFrame}; + +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, +}; + +#[derive(Clone)] +pub struct ValueCount; + +impl Command for ValueCount { + fn name(&self) -> &str { + "dfr value-counts" + } + + fn usage(&self) -> &str { + "Returns a dataframe with the counts for unique values in series" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()).category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Calculates value counts", + example: "[5 5 5 5 6 6] | dfr to-df | dfr value-counts", + result: Some( + NuDataFrame::try_from_columns(vec![ + Column::new("0".to_string(), vec![5.into(), 6.into()]), + Column::new("counts".to_string(), vec![4.into(), 2.into()]), + ]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + _engine_state: &EngineState, + _stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let df = NuDataFrame::try_from_pipeline(input, call.head)?; + let series = df.as_series(call.head)?; + + let res = series.value_counts().map_err(|e| { + ShellError::SpannedLabeledErrorHelp( + "Error calculating value counts values".into(), + e.to_string(), + call.head, + "The str-slice command can only be used with string columns".into(), + ) + })?; + + Ok(PipelineData::Value( + NuDataFrame::dataframe_into_value(res, call.head), + None, + )) +} + +#[cfg(test)] +mod test { + use super::super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(ValueCount {})]) + } +} diff --git a/crates/nu-command/src/dataframe/test_dataframe.rs b/crates/nu-command/src/dataframe/test_dataframe.rs index 9ce9ddd33a..afc8d7eb73 100644 --- a/crates/nu-command/src/dataframe/test_dataframe.rs +++ b/crates/nu-command/src/dataframe/test_dataframe.rs @@ -8,8 +8,13 @@ use nu_protocol::{ use super::ToDataFrame; use crate::Let; -pub fn test_dataframe(cmd: impl Command + 'static) { - let examples = cmd.examples(); +pub fn test_dataframe(cmds: Vec>) { + if cmds.is_empty() { + panic!("Empty commands vector") + } + + // The first element in the cmds vector must be the one tested + let examples = cmds[0].examples(); let mut engine_state = Box::new(EngineState::new()); let delta = { @@ -20,7 +25,9 @@ pub fn test_dataframe(cmd: impl Command + 'static) { working_set.add_decl(Box::new(ToDataFrame)); // Adding the command that is being tested to the working set - working_set.add_decl(Box::new(cmd)); + for cmd in cmds { + working_set.add_decl(cmd); + } working_set.render() }; diff --git a/crates/nu-command/src/dataframe/to_df.rs b/crates/nu-command/src/dataframe/to_df.rs index 308d61a72e..c8a1e683e8 100644 --- a/crates/nu-command/src/dataframe/to_df.rs +++ b/crates/nu-command/src/dataframe/to_df.rs @@ -11,7 +11,7 @@ pub struct ToDataFrame; impl Command for ToDataFrame { fn name(&self) -> &str { - "dataframe to-df" + "dfr to-df" } fn usage(&self) -> &str { @@ -26,7 +26,7 @@ impl Command for ToDataFrame { vec![ Example { description: "Takes a dictionary and creates a dataframe", - example: "[[a b];[1 2] [3 4]] | dataframe to-df", + example: "[[a b];[1 2] [3 4]] | dfr to-df", result: Some( NuDataFrame::try_from_columns(vec![ Column::new("a".to_string(), vec![1.into(), 3.into()]), @@ -38,7 +38,7 @@ impl Command for ToDataFrame { }, Example { description: "Takes a list of tables and creates a dataframe", - example: "[[1 2 a] [3 4 b] [5 6 c]] | dataframe to-df", + example: "[[1 2 a] [3 4 b] [5 6 c]] | dfr to-df", result: Some( NuDataFrame::try_from_columns(vec![ Column::new("0".to_string(), vec![1.into(), 3.into(), 5.into()]), @@ -58,7 +58,7 @@ impl Command for ToDataFrame { }, Example { description: "Takes a list and creates a dataframe", - example: "[a b c] | dataframe to-df", + example: "[a b c] | dfr to-df", result: Some( NuDataFrame::try_from_columns(vec![Column::new( "0".to_string(), @@ -74,7 +74,7 @@ impl Command for ToDataFrame { }, Example { description: "Takes a list of booleans and creates a dataframe", - example: "[$true $true $false] | dataframe to-df", + example: "[$true $true $false] | dfr to-df", result: Some( NuDataFrame::try_from_columns(vec![Column::new( "0".to_string(), @@ -106,6 +106,6 @@ mod test { #[test] fn test_examples() { - test_dataframe(ToDataFrame {}) + test_dataframe(vec![Box::new(ToDataFrame {})]) } } diff --git a/crates/nu-command/src/dataframe/with_column.rs b/crates/nu-command/src/dataframe/with_column.rs new file mode 100644 index 0000000000..1eabb9facf --- /dev/null +++ b/crates/nu-command/src/dataframe/with_column.rs @@ -0,0 +1,100 @@ +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, Spanned, SyntaxShape, Value, +}; + +use super::values::{Column, NuDataFrame}; + +#[derive(Clone)] +pub struct WithColumn; + +impl Command for WithColumn { + fn name(&self) -> &str { + "dfr with-column" + } + + fn usage(&self) -> &str { + "Adds a series to the dataframe" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("series", SyntaxShape::Any, "series to be added") + .required_named("name", SyntaxShape::String, "column name", Some('n')) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Adds a series to the dataframe", + example: + "[[a b]; [1 2] [3 4]] | dfr to-df | dfr with-column ([5 6] | dfr to-df) --name c", + result: Some( + NuDataFrame::try_from_columns(vec![ + Column::new("a".to_string(), vec![1.into(), 3.into()]), + Column::new("b".to_string(), vec![2.into(), 4.into()]), + Column::new("c".to_string(), vec![5.into(), 6.into()]), + ]) + .expect("simple df for test should not fail") + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + command(engine_state, stack, call, input) + } +} + +fn command( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let name: Spanned = call + .get_flag(engine_state, stack, "name")? + .expect("required named value"); + + let other_value: Value = call.req(engine_state, stack, 0)?; + let other_span = other_value.span()?; + let mut other = NuDataFrame::try_from_value(other_value)?.as_series(other_span)?; + let series = other.rename(&name.item).clone(); + + let mut df = NuDataFrame::try_from_pipeline(input, call.head)?; + + df.as_mut() + .with_column(series) + .map_err(|e| { + ShellError::SpannedLabeledError( + "Error adding column to dataframe".into(), + e.to_string(), + other_span, + ) + }) + .map(|df| { + PipelineData::Value( + NuDataFrame::dataframe_into_value(df.clone(), call.head), + None, + ) + }) +} + +#[cfg(test)] +mod test { + use super::super::test_dataframe::test_dataframe; + use super::*; + + #[test] + fn test_examples() { + test_dataframe(vec![Box::new(WithColumn {})]) + } +} diff --git a/crates/nu-parser/src/parse_keywords.rs b/crates/nu-parser/src/parse_keywords.rs index 73c15e08a3..5238fc94bc 100644 --- a/crates/nu-parser/src/parse_keywords.rs +++ b/crates/nu-parser/src/parse_keywords.rs @@ -12,7 +12,7 @@ use std::collections::{HashMap, HashSet}; use crate::{ lex, lite_parse, parser::{ - check_call, check_name, garbage, garbage_statement, parse, parse_block_expression, + check_name, garbage, garbage_statement, parse, parse_block_expression, parse_import_pattern, parse_internal_call, parse_multispan_value, parse_signature, parse_string, parse_var_with_opt_type, trim_quotes, }, @@ -1135,6 +1135,7 @@ pub fn parse_register( working_set: &mut StateWorkingSet, spans: &[Span], ) -> (Statement, Option) { + use crate::parser::check_call; use nu_plugin::{get_signature, EncodingType, PluginDeclaration}; use nu_protocol::Signature; diff --git a/crates/nu-protocol/src/shell_error.rs b/crates/nu-protocol/src/shell_error.rs index 3708c1ce8e..70faf4b8a1 100644 --- a/crates/nu-protocol/src/shell_error.rs +++ b/crates/nu-protocol/src/shell_error.rs @@ -246,6 +246,10 @@ pub enum ShellError { #[diagnostic()] SpannedLabeledError(String, String, #[label("{1}")] Span), + #[error("{0}")] + #[diagnostic(help("{3}"))] + SpannedLabeledErrorHelp(String, String, #[label("{1}")] Span, String), + #[error("{0}")] #[diagnostic()] LabeledError(String, String), diff --git a/crates/nu-protocol/src/value/from_value.rs b/crates/nu-protocol/src/value/from_value.rs index 2d07fc6270..2864054f4b 100644 --- a/crates/nu-protocol/src/value/from_value.rs +++ b/crates/nu-protocol/src/value/from_value.rs @@ -337,3 +337,17 @@ impl FromValue for Spanned { } } } + +impl FromValue for Vec { + fn from_value(v: &Value) -> Result { + // FIXME: we may want to fail a little nicer here + match v { + Value::List { vals, .. } => Ok(vals.clone()), + v => Err(ShellError::CantConvert( + "Vector of values".into(), + v.get_type().to_string(), + v.span()?, + )), + } + } +}