polars: update polars lit to handle nushell Value::Duration and Value::Date types (#15564)

<!--
if this PR closes one or more issues, you can automatically link the PR
with
them by using one of the [*linking
keywords*](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword),
e.g.
- this PR should close #xxxx
- fixes #xxxx

you can also mention related issues, PRs or discussions!
-->

# Description
<!--
Thank you for improving Nushell. Please, check our [contributing
guide](../CONTRIBUTING.md) and talk to the core team before making major
changes.

Description of your pull request goes here. **Provide examples and/or
screenshots** if your changes affect the user experience.
-->
This PR seeks to expand `polars lit` to handle additional nushell types:
Value::Date and Value::Duration. This change is especially relevant to
the `polars filter` command, where expressions would then directly
incorporate Value::Date and Value::Duration types as literals. See one
such example below.

```nushell
#  Filter dataframe for rows where dt is within the last 2 days of the maximum dt value
  > [[dt val]; [2025-04-01 1] [2025-04-02 2] [2025-04-03 3] [2025-04-04 4]] | polars into-df | polars filter ((polars col dt) > ((polars col dt | polars max | $in - 2day)))
  ╭───┬─────────────────────┬─────╮
  │ # │          dt         │ val │
  ├───┼─────────────────────┼─────┤
  │ 0 │ 04/03/25 12:00:00AM │   3 │
  │ 1 │ 04/04/25 12:00:00AM │   4 │
  ╰───┴─────────────────────┴─────╯
```

# User-Facing Changes
<!-- List of all changes that impact the user experience here. This
helps us keep track of breaking changes. -->
No breaking changes. Users now can directly access Value::Date and
Value::Duration types as literals in polars expressions.

# Tests + Formatting
<!--
Don't forget to add tests that cover your changes.

Make sure you've run and fixed any issues with these commands:

- `cargo fmt --all -- --check` to check standard code formatting (`cargo
fmt --all` applies these changes)
- `cargo clippy --workspace -- -D warnings -D clippy::unwrap_used` to
check that you're using the standard code style
- `cargo test --workspace` to check that all tests pass (on Windows make
sure to [enable developer
mode](https://learn.microsoft.com/en-us/windows/apps/get-started/developer-mode-features-and-debugging))
- `cargo run -- -c "use toolkit.nu; toolkit test stdlib"` to run the
tests for the standard library

> **Note**
> from `nushell` you can also use the `toolkit` as follows
> ```bash
> use toolkit.nu # or use an `env_change` hook to activate it
automatically
> toolkit check pr
> ```
-->
Several additional examples added to `polars lit` and `polars filter`

# After Submitting
<!-- If your PR had any user-facing changes, update [the
documentation](https://github.com/nushell/nushell.github.io) after the
PR is merged, if necessary. This will help us keep the docs up to date.
-->
This commit is contained in:
pyz4 2025-04-14 11:58:07 -04:00 committed by GitHub
parent eff9305eb3
commit 9dd30d7756
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 111 additions and 18 deletions

View File

@ -39,7 +39,8 @@ impl PluginCommand for LazyFilter {
}
fn examples(&self) -> Vec<Example> {
vec![Example {
vec![
Example {
description: "Filter dataframe using an expression",
example:
"[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars filter ((polars col a) >= 4)",
@ -60,7 +61,47 @@ impl PluginCommand for LazyFilter {
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
},
Example {
description: "Filter dataframe for rows where dt is within the last 2 days of the maximum dt value",
example:
"[[dt val]; [2025-04-01 1] [2025-04-02 2] [2025-04-03 3] [2025-04-04 4]] | polars into-df | polars filter ((polars col dt) > ((polars col dt | polars max | $in - 2day)))",
result: Some(
NuDataFrame::try_from_columns(
vec![
Column::new(
"dt".to_string(),
vec![
Value::date(
chrono::DateTime::parse_from_str(
"2025-04-03 00:00:00 +0000",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::date(
chrono::DateTime::parse_from_str(
"2025-04-04 00:00:00 +0000",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
]
),
Column::new(
"val".to_string(),
vec![Value::test_int(3), Value::test_int(4)],
),
],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
]
}
fn run(

View File

@ -30,14 +30,40 @@ impl PluginCommand for ExprLit {
}
fn examples(&self) -> Vec<Example> {
vec![Example {
vec![
Example {
description: "Created a literal expression and converts it to a nu object",
example: "polars lit 2 | polars into-nu",
result: Some(Value::test_record(record! {
"expr" => Value::test_string("literal"),
"value" => Value::test_string("dyn int: 2"),
})),
}]
},
Example {
description: "Create a literal expression from date",
example: "polars lit 2025-04-13 | polars into-nu",
result: Some(Value::test_record(record! {
"expr" => Value::test_record(record! {
"expr" => Value::test_string("literal"),
"value" => Value::test_string("dyn int: 1744502400000000000"),
}),
"dtype" => Value::test_string("Datetime(Nanoseconds, None)"),
"cast_options" => Value::test_string("STRICT")
})),
},
Example {
description: "Create a literal expression from duration",
example: "polars lit 3hr | polars into-nu",
result: Some(Value::test_record(record! {
"expr" => Value::test_record(record! {
"expr" => Value::test_string("literal"),
"value" => Value::test_string("dyn int: 10800000000000"),
}),
"dtype" => Value::test_string("Duration(Nanoseconds)"),
"cast_options" => Value::test_string("STRICT")
})),
},
]
}
fn search_terms(&self) -> Vec<&str> {

View File

@ -523,6 +523,32 @@ impl CustomValueSupport for NuExpression {
Value::Int { val, .. } => Ok(val.to_owned().lit().into()),
Value::Bool { val, .. } => Ok(val.to_owned().lit().into()),
Value::Float { val, .. } => Ok(val.to_owned().lit().into()),
Value::Date { val, .. } => val
.to_owned()
.timestamp_nanos_opt()
.ok_or_else(|| ShellError::GenericError {
error: "Integer overflow".into(),
msg: "Provided datetime in nanoseconds is too large for i64".into(),
span: Some(value.span()),
help: None,
inner: vec![],
})
.map(|nanos| -> NuExpression {
nanos
.lit()
.strict_cast(polars::prelude::DataType::Datetime(
polars::prelude::TimeUnit::Nanoseconds,
None,
))
.into()
}),
Value::Duration { val, .. } => Ok(val
.to_owned()
.lit()
.strict_cast(polars::prelude::DataType::Duration(
polars::prelude::TimeUnit::Nanoseconds,
))
.into()),
x => Err(ShellError::CantConvert {
to_type: "lazy expression".into(),
from_type: x.get_type().to_string(),