extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2025-06-20 01:48:09 +02:00
Code Issues Packages Projects Releases Wiki Activity

polars: add new command polars replace-time-zone (#15538)

Browse Source 
<!--
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 add a direct port of the python polars
`replace_time_zone` command in the `dt` namespace
(https://docs.pola.rs/api/python/stable/reference/series/api/polars.Series.dt.replace_time_zone.html).

Please note: I opted for two keywords "dt" and "replace-time-zone" to
map directly with the implementation in both the rust and python
packages, but I'm open to simplifying it to just one keyword, or `polars
replace-time-zone`

```nushell
#  Apply timezone to a naive datetime
  > ["2021-12-30 00:00:00" "2021-12-31 00:00:00"] | polars into-df
                    | polars as-datetime "%Y-%m-%d %H:%M:%S" --naive
                    | polars select (polars col datetime | polars dt replace-time-zone "America/New_York")
  ╭───┬─────────────────────╮
  │ # │      datetime       │
  ├───┼─────────────────────┤
  │ 0 │ 12/30/21 12:00:00AM │
  │ 1 │ 12/31/21 12:00:00AM │
  ╰───┴─────────────────────╯

#  Apply timezone with ambiguous datetime
  > ["2025-11-02 00:00:00", "2025-11-02 01:00:00", "2025-11-02 02:00:00", "2025-11-02 03:00:00"]
                    | polars into-df
                    | polars as-datetime "%Y-%m-%d %H:%M:%S" --naive
                    | polars select (polars col datetime | polars dt replace-time-zone "America/New_York" --ambiguous null)
  ╭───┬─────────────────────╮
  │ # │      datetime       │
  ├───┼─────────────────────┤
  │ 0 │ 11/02/25 12:00:00AM │
  │ 1 │                     │
  │ 2 │ 11/02/25 02:00:00AM │
  │ 3 │ 11/02/25 03:00:00AM │
  ╰───┴─────────────────────╯

#  Apply timezone with nonexistent datetime
  > ["2025-03-09 01:00:00", "2025-03-09 02:00:00", "2025-03-09 03:00:00", "2025-03-09 04:00:00"]
                    | polars into-df
                    | polars as-datetime "%Y-%m-%d %H:%M:%S" --naive
                    | polars select (polars col datetime | polars dt replace-time-zone "America/New_York" --nonexistent null)
  ╭───┬─────────────────────╮
  │ # │      datetime       │
  ├───┼─────────────────────┤
  │ 0 │ 03/09/25 01:00:00AM │
  │ 1 │                     │
  │ 2 │ 03/09/25 03:00:00AM │
  │ 3 │ 03/09/25 04:00:00AM │
  ╰───┴─────────────────────╯
```

# User-Facing Changes
No breaking changes. The user will be able to access the new command.

# Tests + Formatting
See example tests.

# After Submitting
This commit is contained in:
pyz4 2025-04-11 12:09:37 -04:00 committed by GitHub
parent d75aa7ed1b
commit 1a0778d77e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 290 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
crates/nu_plugin_polars/src/dataframe/command/datetime/mod.rs
View File

@ -11,6 +11,8 @@ mod get_second;
mod get_week; mod get_week;
mod get_weekday; mod get_weekday;
mod get_year; mod get_year;
mod replace_time_zone;
mod strftime;
use crate::PolarsPlugin; use crate::PolarsPlugin;
use nu_plugin::PluginCommand; use nu_plugin::PluginCommand;
@ -28,8 +30,7 @@ pub use get_second::GetSecond;
pub use get_week::GetWeek; pub use get_week::GetWeek;
pub use get_weekday::GetWeekDay; pub use get_weekday::GetWeekDay;
pub use get_year::GetYear; pub use get_year::GetYear;
mod strftime; pub use replace_time_zone::ReplaceTimeZone;
pub use strftime::StrFTime; pub use strftime::StrFTime;
pub(crate) fn datetime_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin>>> { pub(crate) fn datetime_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin>>> {
@ -48,5 +49,6 @@ pub(crate) fn datetime_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPl
Box::new(GetWeekDay), Box::new(GetWeekDay),
Box::new(GetYear), Box::new(GetYear),
Box::new(StrFTime), Box::new(StrFTime),
Box::new(ReplaceTimeZone),
] ]
} }

286
crates/nu_plugin_polars/src/dataframe/command/datetime/replace_time_zone.rs Normal file
View File

@ -0,0 +1,286 @@
use crate::values::{Column, NuDataFrame, NuSchema};
use crate::{
dataframe::values::NuExpression,
values::{cant_convert_err, CustomValueSupport, PolarsPluginObject, PolarsPluginType},
PolarsPlugin,
};
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value,
};
use chrono::DateTime;
use polars::prelude::*;
#[derive(Clone)]
pub struct ReplaceTimeZone;
impl PluginCommand for ReplaceTimeZone {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars replace-time-zone"
}
fn description(&self) -> &str {
"Replace the timezone information in a datetime column."
}
fn signature(&self) -> Signature {
Signature::build(self.name())
.input_output_types(vec![(
Type::Custom("expression".into()),
Type::Custom("expression".into()),
)])
.required(
"time_zone",
SyntaxShape::String,
"Timezone for the Datetime Series. Pass `null` to unset time zone.",
)
.named(
"ambiguous",
SyntaxShape::OneOf(vec![SyntaxShape::String, SyntaxShape::Nothing]),
r#"Determine how to deal with ambiguous datetimes:
`raise` (default): raise error
`earliest`: use the earliest datetime
`latest`: use the latest datetime
`null`: set to null"#,
Some('a'),
)
.named(
"nonexistent",
SyntaxShape::OneOf(vec![SyntaxShape::String, SyntaxShape::Nothing]),
r#"Determine how to deal with non-existent datetimes: raise (default) or null."#,
Some('n'),
)
.category(Category::Custom("dataframe".into()))
}
fn examples(&self) -> Vec<Example> {
vec![
Example {
description: "Apply timezone to a naive datetime",
example: r#"["2021-12-30 00:00:00" "2021-12-31 00:00:00"] | polars into-df
| polars as-datetime "%Y-%m-%d %H:%M:%S" --naive
| polars select (polars col datetime | polars replace-time-zone "America/New_York")"#,
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 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::date(
DateTime::parse_from_str(
"2021-12-31 00:00:00 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
],
)],
Some(NuSchema::new(Arc::new(Schema::from_iter(vec![
Field::new(
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
Some(PlSmallStr::from_static("America/New_York")),
),
),
])))),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
Example {
description: "Apply timezone with ambiguous datetime",
example: r#"["2025-11-02 00:00:00", "2025-11-02 01:00:00", "2025-11-02 02:00:00", "2025-11-02 03:00:00"]
| polars into-df
| polars as-datetime "%Y-%m-%d %H:%M:%S" --naive
| polars select (polars col datetime | polars replace-time-zone "America/New_York" --ambiguous null)"#,
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"datetime".to_string(),
vec![
Value::date(
DateTime::parse_from_str(
"2025-11-02 00:00:00 -0400",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::nothing(Span::test_data()),
Value::date(
DateTime::parse_from_str(
"2025-11-02 02:00:00 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::date(
DateTime::parse_from_str(
"2025-11-02 03:00:00 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
],
)],
Some(NuSchema::new(Arc::new(Schema::from_iter(vec![
Field::new(
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
Some(PlSmallStr::from_static("America/New_York")),
),
),
])))),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
Example {
description: "Apply timezone with nonexistent datetime",
example: r#"["2025-03-09 01:00:00", "2025-03-09 02:00:00", "2025-03-09 03:00:00", "2025-03-09 04:00:00"]
| polars into-df
| polars as-datetime "%Y-%m-%d %H:%M:%S" --naive
| polars select (polars col datetime | polars replace-time-zone "America/New_York" --nonexistent null)"#,
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"datetime".to_string(),
vec![
Value::date(
DateTime::parse_from_str(
"2025-03-09 01:00:00 -0500",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::nothing(Span::test_data()),
Value::date(
DateTime::parse_from_str(
"2025-03-09 03:00:00 -0400",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
Value::date(
DateTime::parse_from_str(
"2025-03-09 04:00:00 -0400",
"%Y-%m-%d %H:%M:%S %z",
)
.expect("date calculation should not fail in test"),
Span::test_data(),
),
],
)],
Some(NuSchema::new(Arc::new(Schema::from_iter(vec![
Field::new(
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
Some(PlSmallStr::from_static("America/New_York")),
),
),
])))),
)
.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)?;
let ambiguous = match call.get_flag::<Value>("ambiguous")? {
Some(Value::String { val, internal_span }) => match val.as_str() {
"raise" | "earliest" | "latest" => Ok(val),
_ => Err(ShellError::GenericError {
error: "Invalid argument value".into(),
msg: "`ambiguous` must be one of raise, earliest, latest, or null".into(),
span: Some(internal_span),
help: None,
inner: vec![],
}),
},
Some(Value::Nothing { .. }) => Ok("null".into()),
Some(_) => unreachable!("Argument only accepts string or null."),
None => Ok("raise".into()),
}
.map_err(LabeledError::from)?;
let nonexistent = match call.get_flag::<Value>("nonexistent")? {
Some(Value::String { val, internal_span }) => match val.as_str() {
"raise" => Ok(NonExistent::Raise),
_ => Err(ShellError::GenericError {
error: "Invalid argument value".into(),
msg: "`nonexistent` must be one of raise or null".into(),
span: Some(internal_span),
help: None,
inner: vec![],
}),
},
Some(Value::Nothing { .. }) => Ok(NonExistent::Null),
Some(_) => unreachable!("Argument only accepts string or null."),
None => Ok(NonExistent::Raise),
}
.map_err(LabeledError::from)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuExpression(expr) => {
let time_zone: String = call.req(0)?;
let expr: NuExpression = expr
.into_polars()
.dt()
.replace_time_zone(
Some(PlSmallStr::from_str(&time_zone)),
Expr::Literal(LiteralValue::String(PlSmallStr::from_string(ambiguous))),
nonexistent,
)
.into();
expr.to_pipeline_data(plugin, engine, call.head)
}
_ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])),
}
.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(&ReplaceTimeZone)
}
}