From 885b87a8423f788b8bbabf8cd5e29c2e4637c971 Mon Sep 17 00:00:00 2001 From: pyz4 <42039243+pyz4@users.noreply.github.com> Date: Fri, 11 Apr 2025 17:08:40 -0400 Subject: [PATCH] `polars`: add new command `polars convert-time-zone` (#15550) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Description This is a direct port of the python polars command `convert_time_zone` (https://docs.pola.rs/api/python/stable/reference/series/api/polars.Series.dt.convert_time_zone.html). Consistent with the rust/python implementation, naive datetimes are treated as if they are in UTC time. ```nushell # Convert timezone for timezone-aware datetime > ["2025-04-10 09:30:00 -0400" "2025-04-10 10:30:00 -0400"] | polars into-df | polars as-datetime "%Y-%m-%d %H:%M:%S %z" | polars select (polars col datetime | polars convert-time-zone "Europe/Lisbon") ╭───┬───────────────────────╮ │ # │ datetime │ ├───┼───────────────────────┤ │ 0 │ 04/10/2025 02:30:00PM │ │ 1 │ 04/10/2025 03:30:00PM │ ╰───┴───────────────────────╯ # Timezone conversions for timezone-naive datetime will assume the original timezone is UTC > ["2025-04-10 09:30:00" "2025-04-10 10:30:00"] | polars into-df | polars as-datetime "%Y-%m-%d %H:%M:%S" --naive | polars select (polars col datetime | polars convert-time-zone "America/New_York") ╭───┬───────────────────────╮ │ # │ datetime │ ├───┼───────────────────────┤ │ 0 │ 04/10/2025 05:30:00AM │ │ 1 │ 04/10/2025 06:30:00AM │ ╰───┴───────────────────────╯ ``` # User-Facing Changes No breaking changes. Users have access to a new command `polars convert-time-zone` # Tests + Formatting Example tests have been added. # After Submitting --- .../command/datetime/convert_time_zone.rs | 169 ++++++++++++++++++ .../src/dataframe/command/datetime/mod.rs | 3 + 2 files changed, 172 insertions(+) create mode 100644 crates/nu_plugin_polars/src/dataframe/command/datetime/convert_time_zone.rs diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/convert_time_zone.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/convert_time_zone.rs new file mode 100644 index 0000000000..3572578b42 --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/convert_time_zone.rs @@ -0,0 +1,169 @@ +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, Signature, Span, SyntaxShape, Type, Value, +}; + +use chrono::DateTime; +use polars::prelude::*; + +#[derive(Clone)] +pub struct ConvertTimeZone; + +impl PluginCommand for ConvertTimeZone { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars convert-time-zone" + } + + fn description(&self) -> &str { + "Convert datetime to target timezone." + } + + 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.", + ) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Convert timezone for timezone-aware datetime", + example: r#"["2025-04-10 09:30:00 -0400" "2025-04-10 10:30:00 -0400"] | polars into-df + | polars as-datetime "%Y-%m-%d %H:%M:%S %z" + | polars select (polars col datetime | polars convert-time-zone "Europe/Lisbon")"#, + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "datetime".to_string(), + vec![ + Value::date( + DateTime::parse_from_str( + "2025-04-10 14:30:00 +0100", + "%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-04-10 15:30:00 +0100", + "%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("Europe/Lisbon")), + ), + ), + ])))), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Timezone conversions for timezone-naive datetime will assume the original timezone is UTC", + example: r#"["2025-04-10 09:30:00" "2025-04-10 10:30:00"] | polars into-df + | polars as-datetime "%Y-%m-%d %H:%M:%S" --naive + | polars select (polars col datetime | polars convert-time-zone "America/New_York")"#, + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "datetime".to_string(), + vec![ + Value::date( + DateTime::parse_from_str( + "2025-04-10 05:30: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-04-10 06:30: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 { + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuExpression(expr) => { + let time_zone: String = call.req(0)?; + let expr: NuExpression = expr + .into_polars() + .dt() + .convert_time_zone(PlSmallStr::from_str(&time_zone)) + .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(&ConvertTimeZone) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/mod.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/mod.rs index 7802d9a764..80c997f6ff 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/mod.rs @@ -1,5 +1,6 @@ mod as_date; mod as_datetime; +mod convert_time_zone; mod datepart; mod get_day; mod get_hour; @@ -19,6 +20,7 @@ use nu_plugin::PluginCommand; pub use as_date::AsDate; pub use as_datetime::AsDateTime; +pub use convert_time_zone::ConvertTimeZone; pub use datepart::ExprDatePart; pub use get_day::GetDay; pub use get_hour::GetHour; @@ -50,5 +52,6 @@ pub(crate) fn datetime_commands() -> Vec