add more tests to the polars plugin (#12719)

# Description

I added some more tests to our mighty `polars` ~~, yet I don't know how
to add expected results in some of them. I would like to ask for help.~~

~~My experiments are in the last commit: [polars:
experiments](f7e5e72019).
Without those experiments `cargo test` goes well.~~
 
UPD. I moved out my unsuccessful test experiments into a separate
[branch](https://github.com/maxim-uvarov/nushell/blob/polars-tests-broken2/).
So, this branch seems ready for a merge.

@ayax79, maybe you'll find time for me please? It's not urgent for sure.

P.S. I'm very new to git. Please feel free to give me any suggestions on
how I should use it better
This commit is contained in:
Maxim Uvarov 2024-05-04 09:14:55 +08:00 committed by GitHub
parent 406df7f208
commit a1287f7b3f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 169 additions and 46 deletions

View File

@ -1,14 +1,14 @@
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Spanned, SyntaxShape,
Type,
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Spanned,
SyntaxShape, Type, Value,
};
use polars::prelude::NamedFrom;
use polars::series::Series;
use crate::{values::CustomValueSupport, PolarsPlugin};
use super::super::values::NuDataFrame;
use super::super::values::{Column, NuDataFrame};
#[derive(Clone)]
pub struct SampleDF;
@ -66,6 +66,23 @@ impl PluginCommand for SampleDF {
"[[a b]; [1 2] [3 4] [5 6]] | polars into-df | polars sample --fraction 0.5 --replace",
result: None, // No expected value because sampling is random
},
Example {
description: "Shows sample row using using predefined seed 1",
example:
"[[a b]; [1 2] [3 4] [5 6]] | polars into-df | polars sample --seed 1 --n-rows 1",
result: Some(
NuDataFrame::try_from_columns(
vec![
Column::new("a".to_string(), vec![Value::test_int(5)]),
Column::new("b".to_string(), vec![Value::test_int(6)]),
],
None,
)
.expect("should not fail")
.into_value(Span::test_data()),
)
},
]
}
@ -133,3 +150,14 @@ fn command(
let df = NuDataFrame::new(false, df?);
df.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(&SampleDF)
}
}

View File

@ -33,10 +33,16 @@ impl PluginCommand for ToLazyFrame {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Takes a dictionary and creates a lazy dataframe",
description: "Takes a table and creates a lazyframe",
example: "[[a b];[1 2] [3 4]] | polars into-lazy",
result: None,
}]
},
Example {
description: "Takes a table, creates a lazyframe, assigns column 'b' type str, displays the schema",
example: "[[a b];[1 2] [3 4]] | polars into-lazy --schema {b: str} | polars schema",
result: None
},
]
}
fn run(

View File

@ -67,27 +67,50 @@ impl PluginCommand for Cumulative {
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Cumulative sum for a series",
example: "[1 2 3 4 5] | polars into-df | polars cumulative sum",
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"0_cumulative_sum".to_string(),
vec![
Value::test_int(1),
Value::test_int(3),
Value::test_int(6),
Value::test_int(10),
Value::test_int(15),
],
)],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
vec![
Example {
description: "Cumulative sum for a series",
example: "[1 2 3 4 5] | polars into-df | polars cumulative sum",
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"0_cumulative_sum".to_string(),
vec![
Value::test_int(1),
Value::test_int(3),
Value::test_int(6),
Value::test_int(10),
Value::test_int(15),
],
)],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
Example {
description: "Cumulative sum for a series in reverse order",
example: "[1 2 3 4 5] | polars into-df | polars cumulative sum --reverse",
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"0_cumulative_sum".to_string(),
vec![
Value::test_int(15),
Value::test_int(14),
Value::test_int(12),
Value::test_int(9),
Value::test_int(5),
],
)],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
]
}
fn run(

View File

@ -41,11 +41,18 @@ impl PluginCommand for AsDate {
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Converts string to date",
example: r#"["2021-12-30" "2021-12-31"] | polars into-df | polars as-date "%Y-%m-%d""#,
result: None,
}]
vec![
Example {
description: "Converts string to date",
example: r#"["2021-12-30" "2021-12-31"] | polars into-df | polars as-date "%Y-%m-%d""#,
result: None, // help is needed on how to provide results
},
Example {
description: "Converts string to date",
example: r#"["2021-12-30" "2021-12-31 21:00:00"] | polars into-df | polars as-date "%Y-%m-%d" --not-exact"#,
result: None,
},
]
}
fn run(
@ -99,3 +106,14 @@ fn command(
let df = NuDataFrame::try_from_series_vec(vec![res], call.head)?;
df.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(&AsDate)
}
}

View File

@ -116,6 +116,30 @@ impl PluginCommand for AsDateTime {
.into_value(Span::test_data()),
),
},
Example {
description: "Converts string to datetime using the `--not-exact` flag even with excessive symbols",
example: r#"["2021-12-30 00:00:00 GMT+4"] | polars into-df | polars as-datetime "%Y-%m-%d %H:%M:%S" --not-exact"#,
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"datetime".to_string(),
vec![
Value::date(
DateTime::parse_from_str(
"2021-12-30 00:00:00 +0000",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
],
)],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
]
}

View File

@ -45,21 +45,45 @@ impl PluginCommand for Shift {
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Shifts the values by a given period",
example: "[1 2 2 3 3] | polars into-df | polars shift 2 | polars drop-nulls",
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![Value::test_int(1), Value::test_int(2), Value::test_int(2)],
)],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
vec![
Example {
description: "Shifts the values by a given period",
example: "[1 2 2 3 3] | polars into-df | polars shift 2 | polars drop-nulls",
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![Value::test_int(1), Value::test_int(2), Value::test_int(2)],
)],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
Example {
description: "Shifts the values by a given period, fill absent values with 0",
example:
"[1 2 2 3 3] | polars into-lazy | polars shift 2 --fill 0 | polars collect",
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"0".to_string(),
vec![
Value::test_int(0),
Value::test_int(0),
Value::test_int(1),
Value::test_int(2),
Value::test_int(2),
],
)],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
]
}
fn run(