extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2024-11-22 08:23:24 +01:00
Code Issues Packages Projects Releases Wiki Activity

Migrating polars commands away from macros, removed custom DataFrame comparison. (#13829)

Browse Source 
# Description
This PR:
- Removes the lazy_command, expr_command macros and migrates the
commands that were utilizing them.
- Removes the custom logic in DataFrameValues::is_equals to use the
polars DataFrame version of PartialEq
- Adds examples to commands that previously did not have examples or had
inadequate ones.

NOTE: A lot of examples now have a `polars sort` at the end. This is
needed due to the comparison in the result. The new polars version of
equals cares about the ordering. I removed the custom equals logic as it
causes comparisons to lock up when comparing dataframes that contain a
row that contains a list. I discovered this issue when adding examples
to `polars implode`
This commit is contained in:
Jack Wright 2024-09-11 10:33:05 -07:00 committed by GitHub
parent 0c139c7411
commit 8d60c0d35d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
26 changed files with 1477 additions and 1134 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

92
crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs Normal file
View File

@ -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<Example> {
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::<Series>(), [3i64, 4, 5].iter().collect::<Series>()],
)
.expect("should not fail"),
)
.into_value(Span::test_data()),
),
}]
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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)
}
}

113
crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs
View File

@ -40,81 +40,44 @@ impl PluginCommand for LazyAggregate {
} }
fn examples(&self) -> Vec<Example> { fn examples(&self) -> Vec<Example> {
vec![ vec![Example {
Example { description: "Group by and perform an aggregation",
description: "Group by and perform an aggregation", example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]]
example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] | polars into-lazy
| polars into-df | polars group-by a
| polars group-by a | polars agg [
| polars agg [ (polars col b | polars min | polars as "b_min")
(polars col b | polars min | polars as "b_min") (polars col b | polars max | polars as "b_max")
(polars col b | polars max | polars as "b_max") (polars col b | polars sum | polars as "b_sum")
(polars col b | polars sum | polars as "b_sum") ]
]"#, | polars collect
result: Some( | polars sort-by a"#,
NuDataFrame::try_from_columns( result: Some(
vec![ NuDataFrame::try_from_columns(
Column::new( vec![
"a".to_string(), Column::new(
vec![Value::test_int(1), Value::test_int(2)], "a".to_string(),
), vec![Value::test_int(1), Value::test_int(2)],
Column::new( ),
"b_min".to_string(), Column::new(
vec![Value::test_int(2), Value::test_int(4)], "b_min".to_string(),
), vec![Value::test_int(2), Value::test_int(4)],
Column::new( ),
"b_max".to_string(), Column::new(
vec![Value::test_int(4), Value::test_int(6)], "b_max".to_string(),
), vec![Value::test_int(4), Value::test_int(6)],
Column::new( ),
"b_sum".to_string(), Column::new(
vec![Value::test_int(6), Value::test_int(10)], "b_sum".to_string(),
), vec![Value::test_int(6), Value::test_int(10)],
], ),
None, ],
) None,
.expect("simple df for test should not fail") )
.into_value(Span::test_data()), .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()),
),
},
]
} }
fn run( fn run(

360
crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregation_commands.rs
View File

@ -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
);

90
crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs Normal file
View File

@ -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<Example> {
// 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<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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)
}
}

5
crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs
View File

@ -82,7 +82,10 @@ impl PluginCommand for Cumulative {
vec![ vec![
Example { Example {
description: "Cumulative sum for a column", 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( result: Some(
NuDataFrame::try_from_columns( NuDataFrame::try_from_columns(
vec![Column::new( vec![Column::new(

88
crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs
View File

@ -1,5 +1,5 @@
use crate::{ use crate::{
dataframe::values::{Column, NuDataFrame, NuExpression, NuLazyFrame, NuLazyGroupBy}, dataframe::values::{NuDataFrame, NuExpression, NuLazyFrame, NuLazyGroupBy},
values::CustomValueSupport, values::CustomValueSupport,
PolarsPlugin, PolarsPlugin,
}; };
@ -8,7 +8,7 @@ use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value, Value,
}; };
use polars::prelude::Expr; use polars::{df, prelude::Expr};
#[derive(Clone)] #[derive(Clone)]
pub struct ToLazyGroupBy; pub struct ToLazyGroupBy;
@ -39,46 +39,9 @@ impl PluginCommand for ToLazyGroupBy {
} }
fn examples(&self) -> Vec<Example> { fn examples(&self) -> Vec<Example> {
vec![ vec![Example {
Example { description: "Group by and perform an aggregation",
description: "Group by and perform an aggregation", example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]]
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 into-lazy
| polars group-by a | polars group-by a
| polars agg [ | polars agg [
@ -86,34 +49,21 @@ impl PluginCommand for ToLazyGroupBy {
(polars col b | polars max | polars as "b_max") (polars col b | polars max | polars as "b_max")
(polars col b | polars sum | polars as "b_sum") (polars col b | polars sum | polars as "b_sum")
] ]
| polars collect"#, | polars collect
result: Some( | polars sort-by a"#,
NuDataFrame::try_from_columns( result: Some(
vec![ NuDataFrame::from(
Column::new( df!(
"a".to_string(), "a" => &[1i64, 2],
vec![Value::test_int(1), Value::test_int(2)], "b_min" => &[2i64, 4],
), "b_max" => &[4i64, 6],
Column::new( "b_sum" => &[6i64, 10],
"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") .expect("should not fail"),
.into_value(Span::test_data()), )
), .into_value(Span::test_data()),
}, ),
] }]
} }
fn run( fn run(

85
crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs Normal file
View File

@ -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<Example> {
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::<Series>()],
"b"=> [[4i64, 5, 6].iter().collect::<Series>()],
).expect("should not fail"))
.into_value(Span::unknown())
),
}]
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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)
}
}

141
crates/nu_plugin_polars/src/dataframe/command/aggregation/max.rs Normal file
View File

@ -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<Example> {
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<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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<PipelineData, ShellError> {
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)
}
}

141
crates/nu_plugin_polars/src/dataframe/command/aggregation/mean.rs Normal file
View File

@ -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<Example> {
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<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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<PipelineData, ShellError> {
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)
}
}

11
crates/nu_plugin_polars/src/dataframe/command/aggregation/median.rs
View File

@ -43,9 +43,11 @@ impl PluginCommand for LazyMedian {
Example { Example {
description: "Median aggregation for a group-by", description: "Median aggregation for a group-by",
example: r#"[[a b]; [one 2] [one 4] [two 1]] example: r#"[[a b]; [one 2] [one 4] [two 1]]
| polars into-df | polars into-df
| polars group-by a | polars group-by a
| polars agg (polars col b | polars median)"#, | polars agg (polars col b | polars median)
| polars collect
| polars sort-by a"#,
result: Some( result: Some(
NuDataFrame::try_from_columns( NuDataFrame::try_from_columns(
vec![ vec![
@ -66,7 +68,8 @@ impl PluginCommand for LazyMedian {
}, },
Example { Example {
description: "Median value from columns in a dataframe", 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( result: Some(
NuDataFrame::try_from_columns( NuDataFrame::try_from_columns(
vec![ vec![

141
crates/nu_plugin_polars/src/dataframe/command/aggregation/min.rs Normal file
View File

@ -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<Example> {
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<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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<PipelineData, ShellError> {
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)
}
}

22
crates/nu_plugin_polars/src/dataframe/command/aggregation/mod.rs
View File

@ -1,31 +1,47 @@
mod agg_groups;
mod aggregate; mod aggregate;
mod aggregation_commands; mod count;
mod cumulative; mod cumulative;
pub mod groupby; pub mod groupby;
mod implode;
mod max;
mod mean;
mod median; mod median;
mod min;
mod n_null; mod n_null;
mod n_unique; mod n_unique;
mod quantile; mod quantile;
mod rolling; mod rolling;
mod std;
mod sum;
mod value_counts; mod value_counts;
mod var;
use crate::PolarsPlugin; use crate::PolarsPlugin;
use agg_groups::ExprAggGroups;
use nu_plugin::PluginCommand; use nu_plugin::PluginCommand;
pub use aggregate::LazyAggregate; pub use aggregate::LazyAggregate;
pub use aggregation_commands::*; use count::ExprCount;
pub use cumulative::Cumulative; pub use cumulative::Cumulative;
use implode::ExprImplode;
use max::ExprMax;
use mean::ExprMean;
use min::ExprMin;
pub use n_null::NNull; pub use n_null::NNull;
pub use n_unique::NUnique; pub use n_unique::NUnique;
pub use rolling::Rolling; pub use rolling::Rolling;
use std::ExprStd;
pub use sum::ExprSum;
pub use value_counts::ValueCount; pub use value_counts::ValueCount;
use var::ExprVar;
pub(crate) fn aggregation_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin>>> { pub(crate) fn aggregation_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin>>> {
vec![ vec![
Box::new(Cumulative), Box::new(Cumulative),
Box::new(ExprAggGroups), Box::new(ExprAggGroups),
Box::new(ExprCount), Box::new(ExprCount),
Box::new(ExprList), Box::new(ExprImplode),
Box::new(ExprMax), Box::new(ExprMax),
Box::new(ExprMin), Box::new(ExprMin),
Box::new(ExprSum), Box::new(ExprSum),

8
crates/nu_plugin_polars/src/dataframe/command/aggregation/quantile.rs
View File

@ -66,9 +66,11 @@ impl PluginCommand for LazyQuantile {
Example { Example {
description: "Quantile aggregation for a group-by", description: "Quantile aggregation for a group-by",
example: r#"[[a b]; [one 2] [one 4] [two 1]] example: r#"[[a b]; [one 2] [one 4] [two 1]]
| polars into-df | polars into-df
| polars group-by a | polars group-by a
| polars agg (polars col b | polars quantile 0.5)"#, | polars agg (polars col b | polars quantile 0.5)
| polars collect
| polars sort-by a"#,
result: Some( result: Some(
NuDataFrame::try_from_columns( NuDataFrame::try_from_columns(
vec![ vec![

134
crates/nu_plugin_polars/src/dataframe/command/aggregation/std.rs Normal file
View File

@ -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<Example> {
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<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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<PipelineData, ShellError> {
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)
}
}

142
crates/nu_plugin_polars/src/dataframe/command/aggregation/sum.rs Normal file
View File

@ -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<Example> {
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<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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<PipelineData, ShellError> {
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)
}
}

25
crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs
View File

@ -1,12 +1,12 @@
use crate::values::{Column, CustomValueSupport, NuDataFrame}; use crate::values::{CustomValueSupport, NuDataFrame};
use crate::PolarsPlugin; use crate::PolarsPlugin;
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{ use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value,
}; };
use polars::df;
use polars::prelude::SeriesMethods; use polars::prelude::SeriesMethods;
#[derive(Clone)] #[derive(Clone)]
@ -53,22 +53,15 @@ impl PluginCommand for ValueCount {
fn examples(&self) -> Vec<Example> { fn examples(&self) -> Vec<Example> {
vec![Example { vec![Example {
description: "Calculates value counts", 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( result: Some(
NuDataFrame::try_from_columns( NuDataFrame::from(
vec![ df!(
Column::new( "0" => &[6i64, 5],
"0".to_string(), "count" => &[2i64, 4],
vec![Value::test_int(5), Value::test_int(6)], )
), .expect("should not fail"),
Column::new(
"count".to_string(),
vec![Value::test_int(4), Value::test_int(2)],
),
],
None,
) )
.expect("simple df for test should not fail")
.into_value(Span::test_data()), .into_value(Span::test_data()),
), ),
}] }]

138
crates/nu_plugin_polars/src/dataframe/command/aggregation/var.rs Normal file
View File

@ -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<Example> {
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<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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<PipelineData, ShellError> {
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)
}
}

113
crates/nu_plugin_polars/src/dataframe/command/boolean/expr_not.rs
View File

@ -1,21 +1,98 @@
use crate::expr_command; use crate::dataframe::values::NuExpression;
use crate::values::CustomValueSupport; use crate::values::{
use crate::values::NuExpression; cant_convert_err, CustomValueSupport, NuDataFrame, PolarsPluginObject, PolarsPluginType,
};
use crate::PolarsPlugin; use crate::PolarsPlugin;
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; 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 pub struct ExprNot;
// Expands to a command definition for a not expression
expr_command!( impl PluginCommand for ExprNot {
ExprNot, type Plugin = PolarsPlugin;
"polars expr-not",
"Creates a not expression.", fn name(&self) -> &str {
vec![Example { "polars expr-not"
description: "Creates a not expression", }
example: "(polars col a) > 2) | polars expr-not",
result: None, fn description(&self) -> &str {
},], "Creates a not expression."
not, }
test_not
); 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<Example> {
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<PipelineData, LabeledError> {
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<PipelineData, ShellError> {
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)
}
}

91
crates/nu_plugin_polars/src/dataframe/command/core/cache.rs
View File

@ -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 use crate::{
// Expands to a command definition for cache values::{CustomValueSupport, NuDataFrame, NuLazyFrame},
lazy_command!( PolarsPlugin,
LazyCache, };
"polars cache",
"Caches operations in a new LazyFrame.", pub struct LazyCache;
vec![Example {
description: "Caches the result into a new LazyFrame", impl PluginCommand for LazyCache {
example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars reverse | polars cache", type Plugin = PolarsPlugin;
result: None,
}], fn name(&self) -> &str {
cache, "polars cache"
test_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<Example> {
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<PipelineData, LabeledError> {
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)
}
}

26
crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs
View File

@ -3,13 +3,14 @@ use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value, Value,
}; };
use polars::df;
use polars::prelude::UniqueKeepStrategy; use polars::prelude::UniqueKeepStrategy;
use crate::values::CustomValueSupport; use crate::values::CustomValueSupport;
use crate::PolarsPlugin; use crate::PolarsPlugin;
use crate::values::utils::convert_columns_string; use crate::values::utils::convert_columns_string;
use crate::values::{Column, NuDataFrame}; use crate::values::NuDataFrame;
#[derive(Clone)] #[derive(Clone)]
pub struct DropDuplicates; pub struct DropDuplicates;
@ -48,22 +49,17 @@ impl PluginCommand for DropDuplicates {
fn examples(&self) -> Vec<Example> { fn examples(&self) -> Vec<Example> {
vec![Example { vec![Example {
description: "drop duplicates", 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( result: Some(
NuDataFrame::try_from_columns( NuDataFrame::from(
vec![ df!(
Column::new( "a" => &[1i64, 3],
"a".to_string(), "b" => &[2i64, 4],
vec![Value::test_int(3), Value::test_int(1)], )
), .expect("should not fail"),
Column::new(
"b".to_string(),
vec![Value::test_int(4), Value::test_int(2)],
),
],
None,
) )
.expect("simple df for test should not fail")
.into_value(Span::test_data()), .into_value(Span::test_data()),
), ),
}] }]

6
crates/nu_plugin_polars/src/dataframe/command/data/explode.rs
View File

@ -46,7 +46,11 @@ impl PluginCommand for LazyExplode {
vec![ vec![
Example { Example {
description: "Explode the specified dataframe", 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( result: Some(
NuDataFrame::try_from_columns(vec![ NuDataFrame::try_from_columns(vec![
Column::new( Column::new(

107
crates/nu_plugin_polars/src/dataframe/command/data/reverse.rs
View File

@ -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::{ use crate::{
lazy_command, values::{Column, CustomValueSupport, NuDataFrame, NuLazyFrame},
values::{Column, NuDataFrame}, PolarsPlugin,
}; };
// LazyReverse command pub struct LazyReverse;
// Expands to a command definition for reverse
lazy_command!( impl PluginCommand for LazyReverse {
LazyReverse, type Plugin = PolarsPlugin;
"polars reverse",
"Reverses the LazyFrame", fn name(&self) -> &str {
vec![Example { "polars reverse"
description: "Reverses the dataframe.", }
example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars reverse",
result: Some( fn description(&self) -> &str {
NuDataFrame::try_from_columns( "Reverses the LazyFrame"
vec![ }
Column::new(
"a".to_string(), fn signature(&self) -> Signature {
vec![Value::test_int(2), Value::test_int(4), Value::test_int(6),], Signature::build(self.name())
), .input_output_type(
Column::new( Type::Custom("dataframe".into()),
"b".to_string(), Type::Custom("dataframe".into()),
vec![Value::test_int(2), Value::test_int(2), Value::test_int(2),],
),
],
None
) )
.expect("simple df for test should not fail") .category(Category::Custom("dataframe".into()))
.into_value(Span::test_data()), }
),
},], fn examples(&self) -> Vec<Example> {
reverse, vec![Example {
test_reverse 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<PipelineData, LabeledError> {
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)
}
}

440
crates/nu_plugin_polars/src/dataframe/command/macro_commands.rs
View File

@ -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<Example> {
$examples
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
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<PipelineData, LabeledError> {
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<Example> {
$examples
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
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<Example> {
$examples
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
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<PipelineData, LabeledError> {
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<Example> {
$examples
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
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<Example> {
$examples
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
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)
}
}
};
}

1
crates/nu_plugin_polars/src/dataframe/command/mod.rs
View File

@ -5,6 +5,5 @@ pub mod data;
pub mod datetime; pub mod datetime;
pub mod index; pub mod index;
pub mod integer; pub mod integer;
pub mod macro_commands;
pub mod string; pub mod string;
pub mod stub; pub mod stub;

85
crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs
View File

@ -8,10 +8,7 @@ pub use operations::Axis;
use indexmap::map::IndexMap; use indexmap::map::IndexMap;
use nu_protocol::{did_you_mean, PipelineData, Record, ShellError, Span, Value}; use nu_protocol::{did_you_mean, PipelineData, Record, ShellError, Span, Value};
use polars::{ use polars::prelude::{DataFrame, DataType, IntoLazy, PolarsObject, Series};
chunked_array::ops::SortMultipleOptions,
prelude::{DataFrame, DataType, IntoLazy, PolarsObject, Series},
};
use polars_plan::prelude::{lit, Expr, Null}; use polars_plan::prelude::{lit, Expr, Null};
use polars_utils::total_ord::{TotalEq, TotalHash}; use polars_utils::total_ord::{TotalEq, TotalHash};
use std::{ use std::{
@ -429,80 +426,14 @@ impl NuDataFrame {
// Dataframes are considered equal if they have the same shape, column name and values // Dataframes are considered equal if they have the same shape, column name and values
pub fn is_equal(&self, other: &Self) -> Option<Ordering> { pub fn is_equal(&self, other: &Self) -> Option<Ordering> {
if self.as_ref().width() == 0 { let polars_self = self.to_polars();
// checking for empty dataframe let polars_other = other.to_polars();
return None;
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 { pub fn schema(&self) -> NuSchema {

6
crates/nu_plugin_polars/src/lib.rs
View File

@ -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 { for decl in decls {
let _ = plugin_test.add_decl(decl)?; let _ = plugin_test.add_decl(decl)?;
} }
plugin_test.test_examples(&examples)?; plugin_test.test_examples(&examples)
Ok(())
} }
} }