diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs new file mode 100644 index 0000000000..ae12239800 --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs @@ -0,0 +1,92 @@ +use crate::dataframe::values::NuExpression; +use crate::values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, PolarsPluginObject, PolarsPluginType, +}; +use crate::PolarsPlugin; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, +}; +use polars::df; +use polars::series::Series; + +pub struct ExprAggGroups; + +impl PluginCommand for ExprAggGroups { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars agg-groups" + } + + fn description(&self) -> &str { + "Creates an agg_groups expression." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_types(vec![( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + )]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Get the groiup index of the group by operations.", + example: r#"[[group value]; [one 94] [one 95] [one 96] [two 97] [two 98] [two 99]] + | polars into-df + | polars group-by group + | polars agg (polars col value | polars agg-groups) + | polars collect + | polars sort-by group"#, + result: Some( + NuDataFrame::from( + df!( + "group"=> ["one", "two"], + "values" => [[0i64, 1, 2].iter().collect::(), [3i64, 4, 5].iter().collect::()], + ) + .expect("should not fail"), + ) + .into_value(Span::test_data()), + ), + }] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + NuExpression::from(expr.into_polars().agg_groups()).to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprAggGroups) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs index fdce9ea311..73ff561a28 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs @@ -40,81 +40,44 @@ impl PluginCommand for LazyAggregate { } fn examples(&self) -> Vec { - vec![ - Example { - description: "Group by and perform an aggregation", - example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] - | polars into-df - | polars group-by a - | polars agg [ - (polars col b | polars min | polars as "b_min") - (polars col b | polars max | polars as "b_max") - (polars col b | polars sum | polars as "b_sum") - ]"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_int(1), Value::test_int(2)], - ), - Column::new( - "b_min".to_string(), - vec![Value::test_int(2), Value::test_int(4)], - ), - Column::new( - "b_max".to_string(), - vec![Value::test_int(4), Value::test_int(6)], - ), - Column::new( - "b_sum".to_string(), - vec![Value::test_int(6), Value::test_int(10)], - ), - ], - None, - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - Example { - description: "Group by and perform an aggregation", - example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] - | polars into-lazy - | polars group-by a - | polars agg [ - (polars col b | polars min | polars as "b_min") - (polars col b | polars max | polars as "b_max") - (polars col b | polars sum | polars as "b_sum") - ] - | polars collect"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_int(1), Value::test_int(2)], - ), - Column::new( - "b_min".to_string(), - vec![Value::test_int(2), Value::test_int(4)], - ), - Column::new( - "b_max".to_string(), - vec![Value::test_int(4), Value::test_int(6)], - ), - Column::new( - "b_sum".to_string(), - vec![Value::test_int(6), Value::test_int(10)], - ), - ], - None, - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - ] + vec![Example { + description: "Group by and perform an aggregation", + example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] + | polars into-lazy + | polars group-by a + | polars agg [ + (polars col b | polars min | polars as "b_min") + (polars col b | polars max | polars as "b_max") + (polars col b | polars sum | polars as "b_sum") + ] + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_int(1), Value::test_int(2)], + ), + Column::new( + "b_min".to_string(), + vec![Value::test_int(2), Value::test_int(4)], + ), + Column::new( + "b_max".to_string(), + vec![Value::test_int(4), Value::test_int(6)], + ), + Column::new( + "b_sum".to_string(), + vec![Value::test_int(6), Value::test_int(10)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }] } fn run( diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregation_commands.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregation_commands.rs deleted file mode 100644 index 4f2a4875da..0000000000 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregation_commands.rs +++ /dev/null @@ -1,360 +0,0 @@ -use crate::dataframe::values::{Column, NuDataFrame, NuExpression, NuLazyFrame}; -use crate::values::CustomValueSupport; -use crate::PolarsPlugin; -use crate::{expr_command, lazy_expr_command}; -use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; -use nu_protocol::{Category, Example, LabeledError, PipelineData, Signature, Span, Type, Value}; - -// ExprList command -// Expands to a command definition for a list expression -expr_command!( - ExprList, - "polars implode", - "Aggregates a group to a Series.", - vec![Example { - description: "", - example: "", - result: None, - }], - implode, - test_implode -); - -// ExprAggGroups command -// Expands to a command definition for a agg groups expression -expr_command!( - ExprAggGroups, - "polars agg-groups", - "Creates an agg_groups expression.", - vec![Example { - description: "", - example: "", - result: None, - }], - agg_groups, - test_groups -); - -// ExprCount command -// Expands to a command definition for a count expression -expr_command!( - ExprCount, - "polars count", - "Creates a count expression.", - vec![Example { - description: "", - example: "", - result: None, - }], - count, - test_count -); - -// ExprMax command -// Expands to a command definition for max aggregation -lazy_expr_command!( - ExprMax, - "polars max", - "Creates a max expression or aggregates columns to their max value.", - vec![ - Example { - description: "Max value from columns in a dataframe", - example: "[[a b]; [6 2] [1 4] [4 1]] | polars into-df | polars max", - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new("a".to_string(), vec![Value::test_int(6)],), - Column::new("b".to_string(), vec![Value::test_int(4)],), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - Example { - description: "Max aggregation for a group-by", - example: r#"[[a b]; [one 2] [one 4] [two 1]] - | polars into-df - | polars group-by a - | polars agg (polars col b | polars max)"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_string("one"), Value::test_string("two")], - ), - Column::new( - "b".to_string(), - vec![Value::test_int(4), Value::test_int(1)], - ), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - ], - max, - test_max -); - -// ExprMin command -// Expands to a command definition for min aggregation -lazy_expr_command!( - ExprMin, - "polars min", - "Creates a min expression or aggregates columns to their min value.", - vec![ - Example { - description: "Min value from columns in a dataframe", - example: "[[a b]; [6 2] [1 4] [4 1]] | polars into-df | polars min", - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new("a".to_string(), vec![Value::test_int(1)],), - Column::new("b".to_string(), vec![Value::test_int(1)],), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - Example { - description: "Min aggregation for a group-by", - example: r#"[[a b]; [one 2] [one 4] [two 1]] - | polars into-df - | polars group-by a - | polars agg (polars col b | polars min)"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_string("one"), Value::test_string("two")], - ), - Column::new( - "b".to_string(), - vec![Value::test_int(2), Value::test_int(1)], - ), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - ], - min, - test_min -); - -// ExprSum command -// Expands to a command definition for sum aggregation -lazy_expr_command!( - ExprSum, - "polars sum", - "Creates a sum expression for an aggregation or aggregates columns to their sum value.", - vec![ - Example { - description: "Sums all columns in a dataframe", - example: "[[a b]; [6 2] [1 4] [4 1]] | polars into-df | polars sum", - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new("a".to_string(), vec![Value::test_int(11)],), - Column::new("b".to_string(), vec![Value::test_int(7)],), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - Example { - description: "Sum aggregation for a group-by", - example: r#"[[a b]; [one 2] [one 4] [two 1]] - | polars into-df - | polars group-by a - | polars agg (polars col b | polars sum)"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_string("one"), Value::test_string("two")], - ), - Column::new( - "b".to_string(), - vec![Value::test_int(6), Value::test_int(1)], - ), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - ], - sum, - test_sum -); - -// ExprMean command -// Expands to a command definition for mean aggregation -lazy_expr_command!( - ExprMean, - "polars mean", - "Creates a mean expression for an aggregation or aggregates columns to their mean value.", - vec![ - Example { - description: "Mean value from columns in a dataframe", - example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars mean", - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new("a".to_string(), vec![Value::test_float(4.0)],), - Column::new("b".to_string(), vec![Value::test_float(2.0)],), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - Example { - description: "Mean aggregation for a group-by", - example: r#"[[a b]; [one 2] [one 4] [two 1]] - | polars into-df - | polars group-by a - | polars agg (polars col b | polars mean)"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_string("one"), Value::test_string("two")], - ), - Column::new( - "b".to_string(), - vec![Value::test_float(3.0), Value::test_float(1.0)], - ), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - ], - mean, - test_mean -); - -// ExprStd command -// Expands to a command definition for std aggregation -lazy_expr_command!( - ExprStd, - "polars std", - "Creates a std expression for an aggregation of std value from columns in a dataframe.", - vec![ - Example { - description: "Std value from columns in a dataframe", - example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars std", - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new("a".to_string(), vec![Value::test_float(2.0)],), - Column::new("b".to_string(), vec![Value::test_float(0.0)],), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - Example { - description: "Std aggregation for a group-by", - example: r#"[[a b]; [one 2] [one 2] [two 1] [two 1]] - | polars into-df - | polars group-by a - | polars agg (polars col b | polars std)"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_string("one"), Value::test_string("two")], - ), - Column::new( - "b".to_string(), - vec![Value::test_float(0.0), Value::test_float(0.0)], - ), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - ], - std, - test_std, - 1 -); - -// ExprVar command -// Expands to a command definition for var aggregation -lazy_expr_command!( - ExprVar, - "polars var", - "Create a var expression for an aggregation.", - vec![ - Example { - description: - "Var value from columns in a dataframe or aggregates columns to their var value", - example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars var", - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new("a".to_string(), vec![Value::test_float(4.0)],), - Column::new("b".to_string(), vec![Value::test_float(0.0)],), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - Example { - description: "Var aggregation for a group-by", - example: r#"[[a b]; [one 2] [one 2] [two 1] [two 1]] - | polars into-df - | polars group-by a - | polars agg (polars col b | polars var)"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_string("one"), Value::test_string("two")], - ), - Column::new( - "b".to_string(), - vec![Value::test_float(0.0), Value::test_float(0.0)], - ), - ], - None - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - ], - var, - test_var, - 1 -); diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs new file mode 100644 index 0000000000..8bf8163413 --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs @@ -0,0 +1,90 @@ +use crate::dataframe::values::NuExpression; +use crate::values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, PolarsPluginObject, PolarsPluginType, +}; +use crate::PolarsPlugin; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, +}; +use polars::df; + +pub struct ExprCount; + +impl PluginCommand for ExprCount { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars count" + } + + fn description(&self) -> &str { + "Returns the number of non-null values in the column." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_types(vec![( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + )]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + // to add an example with a result that contains a null we will need to be able to + // allow null values to be entered into the dataframe from nushell + // and retain the correct dtype. Right now null values cause the dtype to be object + vec![Example { + description: "Count the number of non-null values in a column", + example: r#"[[a]; ["foo"] ["bar"]] | polars into-df + | polars select (polars col a | polars count) + | polars collect"#, + result: Some( + NuDataFrame::from( + df!( + "a" => [2] + ) + .expect("should not fail"), + ) + .into_value(Span::unknown()), + ), + }] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + NuExpression::from(expr.into_polars().count()).to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprCount) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs index db3affbd25..7d68391dfb 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs @@ -82,7 +82,10 @@ impl PluginCommand for Cumulative { vec![ Example { description: "Cumulative sum for a column", - example: "[[a]; [1] [2] [3] [4] [5]] | polars into-df | polars select (polars col a | polars cumulative sum | polars as cum_a) | polars collect", + example: "[[a]; [1] [2] [3] [4] [5]] + | polars into-df + | polars select (polars col a | polars cumulative sum | polars as cum_a) + | polars collect", result: Some( NuDataFrame::try_from_columns( vec![Column::new( diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs index 060bdcd63a..24372e47aa 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs @@ -1,5 +1,5 @@ use crate::{ - dataframe::values::{Column, NuDataFrame, NuExpression, NuLazyFrame, NuLazyGroupBy}, + dataframe::values::{NuDataFrame, NuExpression, NuLazyFrame, NuLazyGroupBy}, values::CustomValueSupport, PolarsPlugin, }; @@ -8,7 +8,7 @@ use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Value, }; -use polars::prelude::Expr; +use polars::{df, prelude::Expr}; #[derive(Clone)] pub struct ToLazyGroupBy; @@ -39,46 +39,9 @@ impl PluginCommand for ToLazyGroupBy { } fn examples(&self) -> Vec { - vec![ - Example { - description: "Group by and perform an aggregation", - example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] - | polars into-df - | polars group-by a - | polars agg [ - (polars col b | polars min | polars as "b_min") - (polars col b | polars max | polars as "b_max") - (polars col b | polars sum | polars as "b_sum") - ]"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_int(1), Value::test_int(2)], - ), - Column::new( - "b_min".to_string(), - vec![Value::test_int(2), Value::test_int(4)], - ), - Column::new( - "b_max".to_string(), - vec![Value::test_int(4), Value::test_int(6)], - ), - Column::new( - "b_sum".to_string(), - vec![Value::test_int(6), Value::test_int(10)], - ), - ], - None, - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - Example { - description: "Group by and perform an aggregation", - example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] + vec![Example { + description: "Group by and perform an aggregation", + example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] | polars into-lazy | polars group-by a | polars agg [ @@ -86,34 +49,21 @@ impl PluginCommand for ToLazyGroupBy { (polars col b | polars max | polars as "b_max") (polars col b | polars sum | polars as "b_sum") ] - | polars collect"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_int(1), Value::test_int(2)], - ), - Column::new( - "b_min".to_string(), - vec![Value::test_int(2), Value::test_int(4)], - ), - Column::new( - "b_max".to_string(), - vec![Value::test_int(4), Value::test_int(6)], - ), - Column::new( - "b_sum".to_string(), - vec![Value::test_int(6), Value::test_int(10)], - ), - ], - None, + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::from( + df!( + "a" => &[1i64, 2], + "b_min" => &[2i64, 4], + "b_max" => &[4i64, 6], + "b_sum" => &[6i64, 10], ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }, - ] + .expect("should not fail"), + ) + .into_value(Span::test_data()), + ), + }] } fn run( diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs new file mode 100644 index 0000000000..c46dd6f7a8 --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs @@ -0,0 +1,85 @@ +use crate::dataframe::values::NuExpression; +use crate::values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, PolarsPluginObject, PolarsPluginType, +}; +use crate::PolarsPlugin; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, +}; +use polars::df; +use polars::series::Series; + +pub struct ExprImplode; + +impl PluginCommand for ExprImplode { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars implode" + } + + fn description(&self) -> &str { + "Aggregates values into a list." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_type( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Create two lists for columns a and b with all the rows as values.", + example: "[[a b]; [1 4] [2 5] [3 6]] | polars into-df | polars select (polars col '*' | polars implode) | polars collect", + result: Some( + NuDataFrame::from(df!( + "a"=> [[1i64, 2, 3].iter().collect::()], + "b"=> [[4i64, 5, 6].iter().collect::()], + ).expect("should not fail")) + .into_value(Span::unknown()) + ), + }] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + let res: NuExpression = expr.into_polars().implode().into(); + res.to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprImplode) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/max.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/max.rs new file mode 100644 index 0000000000..d75ea65ede --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/max.rs @@ -0,0 +1,141 @@ +use crate::dataframe::values::NuExpression; +use crate::values::{ + cant_convert_err, Column, CustomValueSupport, NuDataFrame, NuLazyFrame, PolarsPluginObject, + PolarsPluginType, +}; +use crate::PolarsPlugin; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, Value, +}; + +pub struct ExprMax; + +impl PluginCommand for ExprMax { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars max" + } + + fn description(&self) -> &str { + "Creates a max expression or aggregates columns to their max value." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_types(vec![ + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Max value from columns in a dataframe", + example: "[[a b]; [6 2] [1 4] [4 1]] | polars into-df | polars max", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new("a".to_string(), vec![Value::test_int(6)]), + Column::new("b".to_string(), vec![Value::test_int(4)]), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Max aggregation for a group-by", + example: r#"[[a b]; [one 2] [one 4] [two 1]] + | polars into-df + | polars group-by a + | polars agg (polars col b | polars max) + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_string("one"), Value::test_string("two")], + ), + Column::new( + "b".to_string(), + vec![Value::test_int(4), Value::test_int(1)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + NuExpression::from(expr.into_polars().max()).to_pipeline_data(plugin, engine, call.head) +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + let res: NuLazyFrame = lazy.to_polars().max().into(); + + res.to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprMax) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/mean.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/mean.rs new file mode 100644 index 0000000000..1dab512020 --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/mean.rs @@ -0,0 +1,141 @@ +use crate::dataframe::values::NuExpression; +use crate::values::{ + cant_convert_err, Column, CustomValueSupport, NuDataFrame, NuLazyFrame, PolarsPluginObject, + PolarsPluginType, +}; +use crate::PolarsPlugin; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, Value, +}; + +pub struct ExprMean; + +impl PluginCommand for ExprMean { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars mean" + } + + fn description(&self) -> &str { + "Creates a mean expression for an aggregation or aggregates columns to their mean value." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_types(vec![ + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Mean value from columns in a dataframe", + example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars mean", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new("a".to_string(), vec![Value::test_float(4.0)]), + Column::new("b".to_string(), vec![Value::test_float(2.0)]), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Mean aggregation for a group-by", + example: r#"[[a b]; [one 2] [one 4] [two 1]] + | polars into-df + | polars group-by a + | polars agg (polars col b | polars mean) + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_string("one"), Value::test_string("two")], + ), + Column::new( + "b".to_string(), + vec![Value::test_float(3.0), Value::test_float(1.0)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + NuExpression::from(expr.into_polars().mean()).to_pipeline_data(plugin, engine, call.head) +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + let res: NuLazyFrame = lazy.to_polars().mean().into(); + + res.to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprMean) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/median.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/median.rs index 92578da7e0..e4beadecca 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/median.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/median.rs @@ -43,9 +43,11 @@ impl PluginCommand for LazyMedian { Example { description: "Median aggregation for a group-by", example: r#"[[a b]; [one 2] [one 4] [two 1]] - | polars into-df - | polars group-by a - | polars agg (polars col b | polars median)"#, + | polars into-df + | polars group-by a + | polars agg (polars col b | polars median) + | polars collect + | polars sort-by a"#, result: Some( NuDataFrame::try_from_columns( vec![ @@ -66,7 +68,8 @@ impl PluginCommand for LazyMedian { }, Example { description: "Median value from columns in a dataframe", - example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars median", + example: + "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars median | polars collect", result: Some( NuDataFrame::try_from_columns( vec![ diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/min.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/min.rs new file mode 100644 index 0000000000..750099d32f --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/min.rs @@ -0,0 +1,141 @@ +use crate::dataframe::values::NuExpression; +use crate::values::{ + cant_convert_err, Column, CustomValueSupport, NuDataFrame, NuLazyFrame, PolarsPluginObject, + PolarsPluginType, +}; +use crate::PolarsPlugin; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, Value, +}; + +pub struct ExprMin; + +impl PluginCommand for ExprMin { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars min" + } + + fn description(&self) -> &str { + "Creates a min expression or aggregates columns to their min value." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_types(vec![ + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Min value from columns in a dataframe", + example: "[[a b]; [6 2] [1 4] [4 1]] | polars into-df | polars min", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new("a".to_string(), vec![Value::test_int(1)]), + Column::new("b".to_string(), vec![Value::test_int(1)]), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Min aggregation for a group-by", + example: r#"[[a b]; [one 2] [one 4] [two 1]] + | polars into-df + | polars group-by a + | polars agg (polars col b | polars min) + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_string("one"), Value::test_string("two")], + ), + Column::new( + "b".to_string(), + vec![Value::test_int(2), Value::test_int(1)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + NuExpression::from(expr.into_polars().min()).to_pipeline_data(plugin, engine, call.head) +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + let res: NuLazyFrame = lazy.to_polars().min().into(); + + res.to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprMin) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/mod.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/mod.rs index ea14500075..f86b267df1 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/mod.rs @@ -1,31 +1,47 @@ +mod agg_groups; mod aggregate; -mod aggregation_commands; +mod count; mod cumulative; pub mod groupby; +mod implode; +mod max; +mod mean; mod median; +mod min; mod n_null; mod n_unique; mod quantile; mod rolling; +mod std; +mod sum; mod value_counts; +mod var; use crate::PolarsPlugin; +use agg_groups::ExprAggGroups; use nu_plugin::PluginCommand; pub use aggregate::LazyAggregate; -pub use aggregation_commands::*; +use count::ExprCount; pub use cumulative::Cumulative; +use implode::ExprImplode; +use max::ExprMax; +use mean::ExprMean; +use min::ExprMin; pub use n_null::NNull; pub use n_unique::NUnique; pub use rolling::Rolling; +use std::ExprStd; +pub use sum::ExprSum; pub use value_counts::ValueCount; +use var::ExprVar; pub(crate) fn aggregation_commands() -> Vec>> { vec![ Box::new(Cumulative), Box::new(ExprAggGroups), Box::new(ExprCount), - Box::new(ExprList), + Box::new(ExprImplode), Box::new(ExprMax), Box::new(ExprMin), Box::new(ExprSum), diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/quantile.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/quantile.rs index 79804b2937..def85a3e01 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/quantile.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/quantile.rs @@ -66,9 +66,11 @@ impl PluginCommand for LazyQuantile { Example { description: "Quantile aggregation for a group-by", example: r#"[[a b]; [one 2] [one 4] [two 1]] - | polars into-df - | polars group-by a - | polars agg (polars col b | polars quantile 0.5)"#, + | polars into-df + | polars group-by a + | polars agg (polars col b | polars quantile 0.5) + | polars collect + | polars sort-by a"#, result: Some( NuDataFrame::try_from_columns( vec![ diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/std.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/std.rs new file mode 100644 index 0000000000..46411d1b0d --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/std.rs @@ -0,0 +1,134 @@ +use crate::values::{ + cant_convert_err, Column, CustomValueSupport, NuLazyFrame, PolarsPluginObject, PolarsPluginType, +}; +use crate::PolarsPlugin; +use crate::{dataframe::values::NuExpression, values::NuDataFrame}; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, Value, +}; +use polars::df; + +pub struct ExprStd; + +impl PluginCommand for ExprStd { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars std" + } + + fn description(&self) -> &str { + "Creates a std expression for an aggregation of std value from columns in a dataframe." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_types(vec![ + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Std value from columns in a dataframe", + example: + "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars std | polars collect", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new("a".to_string(), vec![Value::test_float(2.0)]), + Column::new("b".to_string(), vec![Value::test_float(0.0)]), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Std aggregation for a group-by", + example: r#"[[a b]; [one 2] [one 2] [two 1] [two 1]] + | polars into-df + | polars group-by a + | polars agg (polars col b | polars std) + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::from( + df!( + "a" => &["one", "two"], + "b" => &[0.0f64, 0.0], + ) + .expect("should not fail"), + ) + .into_value(Span::test_data()), + ), + }, + ] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + NuExpression::from(expr.into_polars().std(1)).to_pipeline_data(plugin, engine, call.head) +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + let res: NuLazyFrame = lazy.to_polars().std(1).into(); + + res.to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprStd) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/sum.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/sum.rs new file mode 100644 index 0000000000..6aaf3320fe --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/sum.rs @@ -0,0 +1,142 @@ +use crate::dataframe::values::NuExpression; +use crate::values::{ + cant_convert_err, Column, CustomValueSupport, NuDataFrame, NuLazyFrame, PolarsPluginObject, + PolarsPluginType, +}; +use crate::PolarsPlugin; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, Value, +}; + +pub struct ExprSum; + +impl PluginCommand for ExprSum { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars sum" + } + + fn description(&self) -> &str { + "Creates a sum expression for an aggregation or aggregates columns to their sum value." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_types(vec![ + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Sums all columns in a dataframe", + example: + "[[a b]; [6 2] [1 4] [4 1]] | polars into-df | polars sum | polars collect", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new("a".to_string(), vec![Value::test_int(11)]), + Column::new("b".to_string(), vec![Value::test_int(7)]), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Sum aggregation for a group-by", + example: r#"[[a b]; [one 2] [one 4] [two 1]] + | polars into-df + | polars group-by a + | polars agg (polars col b | polars sum) + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_string("one"), Value::test_string("two")], + ), + Column::new( + "b".to_string(), + vec![Value::test_int(6), Value::test_int(1)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + NuExpression::from(expr.into_polars().sum()).to_pipeline_data(plugin, engine, call.head) +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + let res: NuLazyFrame = lazy.to_polars().sum().into(); + + res.to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprSum) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs index e05a1c401c..dd3fc57261 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs @@ -1,12 +1,12 @@ -use crate::values::{Column, CustomValueSupport, NuDataFrame}; +use crate::values::{CustomValueSupport, NuDataFrame}; use crate::PolarsPlugin; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, - Value, }; +use polars::df; use polars::prelude::SeriesMethods; #[derive(Clone)] @@ -53,22 +53,15 @@ impl PluginCommand for ValueCount { fn examples(&self) -> Vec { vec![Example { description: "Calculates value counts", - example: "[5 5 5 5 6 6] | polars into-df | polars value-counts", + example: "[5 5 5 5 6 6] | polars into-df | polars value-counts | polars sort-by count", result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "0".to_string(), - vec![Value::test_int(5), Value::test_int(6)], - ), - Column::new( - "count".to_string(), - vec![Value::test_int(4), Value::test_int(2)], - ), - ], - None, + NuDataFrame::from( + df!( + "0" => &[6i64, 5], + "count" => &[2i64, 4], + ) + .expect("should not fail"), ) - .expect("simple df for test should not fail") .into_value(Span::test_data()), ), }] diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/var.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/var.rs new file mode 100644 index 0000000000..bc1ee8ed04 --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/var.rs @@ -0,0 +1,138 @@ +use crate::dataframe::values::NuExpression; +use crate::values::{ + cant_convert_err, Column, CustomValueSupport, NuDataFrame, NuLazyFrame, PolarsPluginObject, + PolarsPluginType, +}; +use crate::PolarsPlugin; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, Value, +}; +use polars::df; + +pub struct ExprVar; + +impl PluginCommand for ExprVar { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars var" + } + + fn description(&self) -> &str { + "Create a var expression for an aggregation." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_types(vec![ + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: + "Var value from columns in a dataframe or aggregates columns to their var value", + example: + "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars var | polars collect", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new("a".to_string(), vec![Value::test_float(4.0)]), + Column::new("b".to_string(), vec![Value::test_float(0.0)]), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Var aggregation for a group-by", + example: r#"[[a b]; [one 2] [one 2] [two 1] [two 1]] + | polars into-df + | polars group-by a + | polars agg (polars col b | polars var) + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::from( + df!( + "a" => &["one", "two"], + "b" => &[0.0, 0.0], + ) + .expect("should not fail"), + ) + .into_value(Span::test_data()), + ), + }, + ] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + NuExpression::from(expr.into_polars().var(1)).to_pipeline_data(plugin, engine, call.head) +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + let res: NuLazyFrame = lazy.to_polars().var(1).into(); + + res.to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use crate::test::test_polars_plugin_command; + + use super::*; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprVar) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/expr_not.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/expr_not.rs index 886d89eeab..c70d921e32 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/expr_not.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/expr_not.rs @@ -1,21 +1,98 @@ -use crate::expr_command; -use crate::values::CustomValueSupport; -use crate::values::NuExpression; +use crate::dataframe::values::NuExpression; +use crate::values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, PolarsPluginObject, PolarsPluginType, +}; use crate::PolarsPlugin; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; -use nu_protocol::{Category, Example, LabeledError, PipelineData, Signature, Type}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, +}; +use polars::df; -// ExprNot command -// Expands to a command definition for a not expression -expr_command!( - ExprNot, - "polars expr-not", - "Creates a not expression.", - vec![Example { - description: "Creates a not expression", - example: "(polars col a) > 2) | polars expr-not", - result: None, - },], - not, - test_not -); +pub struct ExprNot; + +impl PluginCommand for ExprNot { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars expr-not" + } + + fn description(&self) -> &str { + "Creates a not expression." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_types(vec![( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + )]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Creates a not expression", + example: "(polars col a) > 2) | polars expr-not", + result: None, + }, + Example { + description: "Adds a column showing which values of col a are not greater than 2", + example: "[[a]; [1] [2] [3] [4] [5]] | polars into-df + | polars with-column [(((polars col a) > 2) + | polars expr-not + | polars as a_expr_not)] + | polars collect + | polars sort-by a", + result: Some( + NuDataFrame::from( + df!( + "a" => [1, 2, 3, 4, 5], + "b" => [true, true, false, false, false] + ) + .expect("should not fail"), + ) + .into_value(Span::test_data()), + ), + }, + ] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), + } + .map_err(LabeledError::from) + } +} + +fn command_expr( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + expr: NuExpression, +) -> Result { + NuExpression::from(expr.into_polars().not()).to_pipeline_data(plugin, engine, call.head) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&ExprNot) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/cache.rs b/crates/nu_plugin_polars/src/dataframe/command/core/cache.rs index ce68f7712b..23c972654e 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/cache.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/cache.rs @@ -1,16 +1,77 @@ -use crate::lazy_command; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{Category, Example, LabeledError, PipelineData, Signature, Span, Type}; +use polars::df; -// LazyCache command -// Expands to a command definition for cache -lazy_command!( - LazyCache, - "polars cache", - "Caches operations in a new LazyFrame.", - vec![Example { - description: "Caches the result into a new LazyFrame", - example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars reverse | polars cache", - result: None, - }], - cache, - test_cache -); +use crate::{ + values::{CustomValueSupport, NuDataFrame, NuLazyFrame}, + PolarsPlugin, +}; + +pub struct LazyCache; + +impl PluginCommand for LazyCache { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars cache" + } + + fn description(&self) -> &str { + "Caches operations in a new LazyFrame." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_type( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Caches the result into a new LazyFrame", + example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df + | polars reverse + | polars cache + | polars sort-by a", + result: Some( + NuDataFrame::from( + df!( + "a" => [2i64, 4, 6], + "b" => [2i64, 2, 2], + ) + .expect("should not fail"), + ) + .into_value(Span::test_data()), + ), + }] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let lazy = NuLazyFrame::try_from_pipeline_coerce(plugin, input, call.head) + .map_err(LabeledError::from)?; + let lazy = NuLazyFrame::new(lazy.from_eager, lazy.to_polars().cache()); + lazy.to_pipeline_data(plugin, engine, call.head) + .map_err(LabeledError::from) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&LazyCache) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs b/crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs index 5f82ce13f6..4ea26a78b5 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs @@ -3,13 +3,14 @@ use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Value, }; +use polars::df; use polars::prelude::UniqueKeepStrategy; use crate::values::CustomValueSupport; use crate::PolarsPlugin; use crate::values::utils::convert_columns_string; -use crate::values::{Column, NuDataFrame}; +use crate::values::NuDataFrame; #[derive(Clone)] pub struct DropDuplicates; @@ -48,22 +49,17 @@ impl PluginCommand for DropDuplicates { fn examples(&self) -> Vec { vec![Example { description: "drop duplicates", - example: "[[a b]; [1 2] [3 4] [1 2]] | polars into-df | polars drop-duplicates", + example: "[[a b]; [1 2] [3 4] [1 2]] | polars into-df + | polars drop-duplicates + | polars sort-by a", result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_int(3), Value::test_int(1)], - ), - Column::new( - "b".to_string(), - vec![Value::test_int(4), Value::test_int(2)], - ), - ], - None, + NuDataFrame::from( + df!( + "a" => &[1i64, 3], + "b" => &[2i64, 4], + ) + .expect("should not fail"), ) - .expect("simple df for test should not fail") .into_value(Span::test_data()), ), }] diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/explode.rs b/crates/nu_plugin_polars/src/dataframe/command/data/explode.rs index 3266640d34..e2bc6b4e49 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/explode.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/explode.rs @@ -46,7 +46,11 @@ impl PluginCommand for LazyExplode { vec![ Example { description: "Explode the specified dataframe", - example: "[[id name hobbies]; [1 Mercy [Cycling Knitting]] [2 Bob [Skiing Football]]] | polars into-df | polars explode hobbies | polars collect", + example: "[[id name hobbies]; [1 Mercy [Cycling Knitting]] [2 Bob [Skiing Football]]] + | polars into-df + | polars explode hobbies + | polars collect + | polars sort-by [id, name]", result: Some( NuDataFrame::try_from_columns(vec![ Column::new( diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/reverse.rs b/crates/nu_plugin_polars/src/dataframe/command/data/reverse.rs index cb1e6af35f..9665b7aa2b 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/reverse.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/reverse.rs @@ -1,37 +1,80 @@ -use nu_protocol::{Span, Value}; +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{Category, Example, LabeledError, PipelineData, Signature, Span, Type, Value}; use crate::{ - lazy_command, - values::{Column, NuDataFrame}, + values::{Column, CustomValueSupport, NuDataFrame, NuLazyFrame}, + PolarsPlugin, }; -// LazyReverse command -// Expands to a command definition for reverse -lazy_command!( - LazyReverse, - "polars reverse", - "Reverses the LazyFrame", - vec![Example { - description: "Reverses the dataframe.", - example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars reverse", - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_int(2), Value::test_int(4), Value::test_int(6),], - ), - Column::new( - "b".to_string(), - vec![Value::test_int(2), Value::test_int(2), Value::test_int(2),], - ), - ], - None +pub struct LazyReverse; + +impl PluginCommand for LazyReverse { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars reverse" + } + + fn description(&self) -> &str { + "Reverses the LazyFrame" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_type( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - },], - reverse, - test_reverse -); + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Reverses the dataframe.", + example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars reverse", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_int(2), Value::test_int(4), Value::test_int(6)], + ), + Column::new( + "b".to_string(), + vec![Value::test_int(2), Value::test_int(2), Value::test_int(2)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + input: PipelineData, + ) -> Result { + let lazy = NuLazyFrame::try_from_pipeline_coerce(plugin, input, call.head) + .map_err(LabeledError::from)?; + let lazy = NuLazyFrame::new(lazy.from_eager, lazy.to_polars().reverse()); + lazy.to_pipeline_data(plugin, engine, call.head) + .map_err(LabeledError::from) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command; + use nu_protocol::ShellError; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command(&LazyReverse) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/macro_commands.rs b/crates/nu_plugin_polars/src/dataframe/command/macro_commands.rs deleted file mode 100644 index 3986a53037..0000000000 --- a/crates/nu_plugin_polars/src/dataframe/command/macro_commands.rs +++ /dev/null @@ -1,440 +0,0 @@ -/// Definition of multiple Expression commands using a macro rule -/// All of these expressions have an identical body and only require -/// to have a change in the name, description and expression function - -#[macro_export] -macro_rules! lazy_command { - ($command: ident, $name: expr, $desc: expr, $examples: expr, $func: ident, $test: ident) => { - use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; - use nu_protocol::{Category, Example, LabeledError, PipelineData, Signature, Type}; - /// Definition of multiple lazyframe commands using a macro rule - /// All of these commands have an identical body and only require - /// to have a change in the name, description and function - use $crate::dataframe::values::NuLazyFrame; - use $crate::values::CustomValueSupport; - use $crate::PolarsPlugin; - - #[derive(Clone)] - pub struct $command; - - impl PluginCommand for $command { - type Plugin = PolarsPlugin; - - fn name(&self) -> &str { - $name - } - - fn description(&self) -> &str { - $desc - } - - fn signature(&self) -> Signature { - Signature::build(self.name()) - .description($desc) - .input_output_type( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ) - .category(Category::Custom("lazyframe".into())) - } - - fn examples(&self) -> Vec { - $examples - } - - fn run( - &self, - plugin: &Self::Plugin, - engine: &EngineInterface, - call: &EvaluatedCall, - input: PipelineData, - ) -> Result { - let lazy = NuLazyFrame::try_from_pipeline_coerce(plugin, input, call.head) - .map_err(LabeledError::from)?; - let lazy = NuLazyFrame::new(lazy.from_eager, lazy.to_polars().$func()); - lazy.to_pipeline_data(plugin, engine, call.head) - .map_err(LabeledError::from) - } - } - - #[cfg(test)] - mod $test { - use super::*; - use nu_protocol::ShellError; - use $crate::test::test_polars_plugin_command; - - #[test] - fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&$command) - } - } - }; - - ($command: ident, $name: expr, $desc: expr, $examples: expr, $func: ident, $test: ident, $ddot: expr) => { - #[derive(Clone)] - pub struct $command; - - impl PluginCommand for $command { - type Plugin = PolarsPlugin; - - fn signature(&self) -> Signature { - Signature::build($name) - .description($desc) - .input_output_type( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ) - .category(Category::Custom("lazyframe".into())) - .plugin_examples($examples) - } - - fn run( - &self, - _plugin: &Self::Plugin, - engine: &EngineInterface, - call: &EvaluatedCall, - input: PipelineData, - ) -> Result { - let lazy = NuLazyFrame::try_from_pipeline_coerce(plugin, input, call.head) - .map_err(LabeledError::from)?; - let lazy = NuLazyFrame::new(lazy.from_eager, lazy.into_polars().$func($ddot)); - lazy.to_pipeline_data(plugin, engine, call.head) - .map_err(LabeledError::from) - } - } - - #[cfg(test)] - mod $test { - use super::*; - use nu_protocol::ShellError; - use $crate::test::test_polars_plugin_command; - - #[test] - fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&$command) - } - } - }; - - ($command: ident, $name: expr, $desc: expr, $examples: expr, $func: ident?, $test: ident) => { - #[derive(Clone)] - pub struct $command; - - impl PluginCommand for $command { - type Plugin = PolarsPlugin; - - fn name(&self) -> &str { - $name - } - - fn description(&self) -> &str { - $desc - } - fn signature(&self) -> Signature { - Signature::build(self.name()) - .input_output_type( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ) - .category(Category::Custom("lazyframe".into())) - } - - fn examples(&self) -> Vec { - $examples - } - - fn run( - &self, - plugin: &Self::Plugin, - engine: &EngineInterface, - call: &EvaluatedCall, - input: PipelineData, - ) -> Result { - let lazy = NuLazyFrame::try_from_pipeline_coerce(plugin, input, call.head) - .map_err(LabeledError::from)?; - - let lazy = NuLazyFrame::new( - lazy.from_eager, - lazy.to_polars() - .$func() - .map_err(|e| ShellError::GenericError { - error: "Dataframe Error".into(), - msg: e.to_string(), - help: None, - span: None, - inner: vec![], - }) - .map_err(LabeledError::from)?, - ); - - lazy.to_pipeline_data(plugin, engine, call.head) - .map_err(LabeledError::from) - } - } - - #[cfg(test)] - mod $test { - use super::*; - use nu_protocol::ShellError; - use $crate::test::test_polars_plugin_command; - - #[test] - fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&$command) - } - } - }; -} - -// The structs defined in this file are structs that form part of other commands -// since they share a similar name -#[macro_export] -macro_rules! expr_command { - ($command: ident, $name: expr, $desc: expr, $examples: expr, $func: ident, $test: ident) => { - #[derive(Clone)] - pub struct $command; - - impl PluginCommand for $command { - type Plugin = PolarsPlugin; - - fn name(&self) -> &str { - $name - } - - fn description(&self) -> &str { - $desc - } - - fn signature(&self) -> Signature { - Signature::build(self.name()) - .description($desc) - .input_output_type( - Type::Custom("expression".into()), - Type::Custom("expression".into()), - ) - .category(Category::Custom("expression".into())) - } - - fn examples(&self) -> Vec { - $examples - } - - fn run( - &self, - plugin: &Self::Plugin, - engine: &EngineInterface, - call: &EvaluatedCall, - input: PipelineData, - ) -> Result { - let expr = NuExpression::try_from_pipeline(plugin, input, call.head) - .map_err(LabeledError::from)?; - let expr: NuExpression = expr.into_polars().$func().into(); - expr.to_pipeline_data(plugin, engine, call.head) - .map_err(LabeledError::from) - } - } - - #[cfg(test)] - mod $test { - use super::*; - use nu_protocol::ShellError; - use $crate::test::test_polars_plugin_command; - - #[test] - fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&$command) - } - } - }; - - ($command: ident, $name: expr, $desc: expr, $examples: expr, $func: ident, $test: ident, $ddof: expr) => { - #[derive(Clone)] - pub struct $command; - - impl PluginCommand for $command { - type Plugin = PolarsPlugin; - - fn signature(&self) -> Signature { - Signature::build(self.name()) - .description($desc) - .input_output_type( - Type::Custom("expression".into()), - Type::Custom("expression".into()), - ) - .category(Category::Custom("expression".into())) - .plugin_examples($examples) - } - - fn run( - &self, - _plugin: &Self::Plugin, - engine: &EngineInterface, - call: &EvaluatedCall, - input: PipelineData, - ) -> Result { - let expr = NuExpression::try_from_pipeline(input, call.head) - .map_err(LabeledError::from)?; - let expr: NuExpression = expr.into_polars().$func($ddof).into(); - expr.to_pipeline_data(plugin, engine, call.head) - .map_err(LabeledError::from) - } - } - - #[cfg(test)] - mod $test { - use super::*; - use $crate::test::test_polars_plugin_command; - - #[test] - fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&$command) - } - } - }; -} - -// The structs defined in this file are structs that form part of other commands -// since they share a similar name -#[macro_export] -macro_rules! lazy_expr_command { - ($command: ident, $name: expr, $desc: expr, $examples: expr, $func: ident, $test: ident) => { - #[derive(Clone)] - pub struct $command; - - impl PluginCommand for $command { - type Plugin = PolarsPlugin; - - fn name(&self) -> &str { - $name - } - - fn description(&self) -> &str { - $desc - } - - fn signature(&self) -> Signature { - Signature::build(self.name()) - .description($desc) - .input_output_types(vec![ - ( - Type::Custom("expression".into()), - Type::Custom("expression".into()), - ), - ( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ), - ]) - .category(Category::Custom("expression".into())) - } - - fn examples(&self) -> Vec { - $examples - } - - fn run( - &self, - plugin: &Self::Plugin, - engine: &EngineInterface, - call: &EvaluatedCall, - input: PipelineData, - ) -> Result { - let value = input.into_value(call.head)?; - if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) { - let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value) - .map_err(LabeledError::from)?; - let lazy = NuLazyFrame::new(lazy.from_eager, lazy.to_polars().$func()); - lazy.to_pipeline_data(plugin, engine, call.head) - .map_err(LabeledError::from) - } else { - let expr = - NuExpression::try_from_value(plugin, &value).map_err(LabeledError::from)?; - let expr: NuExpression = expr.into_polars().$func().into(); - expr.to_pipeline_data(plugin, engine, call.head) - .map_err(LabeledError::from) - } - } - } - - #[cfg(test)] - mod $test { - use super::*; - use nu_protocol::ShellError; - use $crate::test::test_polars_plugin_command; - - #[test] - fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&$command) - } - } - }; - - ($command: ident, $name: expr, $desc: expr, $examples: expr, $func: ident, $test: ident, $ddof: expr) => { - #[derive(Clone)] - pub struct $command; - - impl PluginCommand for $command { - type Plugin = PolarsPlugin; - - fn name(&self) -> &str { - $name - } - - fn description(&self) -> &str { - $desc - } - fn signature(&self) -> Signature { - Signature::build(self.name()) - .input_output_types(vec![ - ( - Type::Custom("expression".into()), - Type::Custom("expression".into()), - ), - ( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ), - ]) - .category(Category::Custom("expression".into())) - } - - fn examples(&self) -> Vec { - $examples - } - - fn run( - &self, - plugin: &Self::Plugin, - engine: &EngineInterface, - call: &EvaluatedCall, - input: PipelineData, - ) -> Result { - let value = input.into_value(call.head)?; - if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) { - let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value) - .map_err(LabeledError::from)?; - let lazy = NuLazyFrame::new(lazy.from_eager, lazy.to_polars().$func($ddof)); - lazy.to_pipeline_data(plugin, engine, call.head) - .map_err(LabeledError::from) - } else { - let expr = NuExpression::try_from_value(plugin, &value)?; - let expr: NuExpression = expr.into_polars().$func($ddof).into(); - expr.to_pipeline_data(plugin, engine, call.head) - .map_err(LabeledError::from) - } - } - } - - #[cfg(test)] - mod $test { - use super::*; - use nu_protocol::ShellError; - use $crate::test::test_polars_plugin_command; - - #[test] - fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&$command) - } - } - }; -} diff --git a/crates/nu_plugin_polars/src/dataframe/command/mod.rs b/crates/nu_plugin_polars/src/dataframe/command/mod.rs index bab0f3c9f2..50c0d522a1 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/mod.rs @@ -5,6 +5,5 @@ pub mod data; pub mod datetime; pub mod index; pub mod integer; -pub mod macro_commands; pub mod string; pub mod stub; diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs index 89a8c862e7..e782cb9265 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs @@ -8,10 +8,7 @@ pub use operations::Axis; use indexmap::map::IndexMap; use nu_protocol::{did_you_mean, PipelineData, Record, ShellError, Span, Value}; -use polars::{ - chunked_array::ops::SortMultipleOptions, - prelude::{DataFrame, DataType, IntoLazy, PolarsObject, Series}, -}; +use polars::prelude::{DataFrame, DataType, IntoLazy, PolarsObject, Series}; use polars_plan::prelude::{lit, Expr, Null}; use polars_utils::total_ord::{TotalEq, TotalHash}; use std::{ @@ -429,80 +426,14 @@ impl NuDataFrame { // Dataframes are considered equal if they have the same shape, column name and values pub fn is_equal(&self, other: &Self) -> Option { - if self.as_ref().width() == 0 { - // checking for empty dataframe - return None; + let polars_self = self.to_polars(); + let polars_other = other.to_polars(); + + if polars_self == polars_other { + Some(Ordering::Equal) + } else { + None } - - if self.as_ref().get_column_names() != other.as_ref().get_column_names() { - // checking both dataframes share the same names - return None; - } - - if self.as_ref().height() != other.as_ref().height() { - // checking both dataframes have the same row size - return None; - } - - // sorting dataframe by the first column - let column_names = self.as_ref().get_column_names(); - let first_col = column_names - .first() - .expect("already checked that dataframe is different than 0"); - - // if unable to sort, then unable to compare - let lhs = match self - .as_ref() - .sort(vec![*first_col], SortMultipleOptions::default()) - { - Ok(df) => df, - Err(_) => return None, - }; - - let rhs = match other - .as_ref() - .sort(vec![*first_col], SortMultipleOptions::default()) - { - Ok(df) => df, - Err(_) => return None, - }; - - for name in self.as_ref().get_column_names() { - let self_series = lhs.column(name).expect("name from dataframe names"); - - let other_series = rhs - .column(name) - .expect("already checked that name in other"); - - // Casting needed to compare other numeric types with nushell numeric type. - // In nushell we only have i64 integer numeric types and any array created - // with nushell untagged primitives will be of type i64 - let self_series = match self_series.dtype() { - DataType::UInt32 | DataType::Int32 if *other_series.dtype() == DataType::Int64 => { - match self_series.cast(&DataType::Int64) { - Ok(series) => series, - Err(_) => return None, - } - } - _ => self_series.clone(), - }; - - let other_series = match other_series.dtype() { - DataType::UInt32 | DataType::Int32 if *self_series.dtype() == DataType::Int64 => { - match other_series.cast(&DataType::Int64) { - Ok(series) => series, - Err(_) => return None, - } - } - _ => other_series.clone(), - }; - - if !self_series.equals(&other_series) { - return None; - } - } - - Some(Ordering::Equal) } pub fn schema(&self) -> NuSchema { diff --git a/crates/nu_plugin_polars/src/lib.rs b/crates/nu_plugin_polars/src/lib.rs index c537ef53fe..f70ab2b6ef 100644 --- a/crates/nu_plugin_polars/src/lib.rs +++ b/crates/nu_plugin_polars/src/lib.rs @@ -267,13 +267,11 @@ pub mod test { } } - let mut plugin_test = PluginTest::new("polars", plugin.into())?; + let mut plugin_test = PluginTest::new(command.name(), plugin.into())?; for decl in decls { let _ = plugin_test.add_decl(decl)?; } - plugin_test.test_examples(&examples)?; - - Ok(()) + plugin_test.test_examples(&examples) } }