diff --git a/crates/polars-arrow/src/kernels/time.rs b/crates/polars-arrow/src/kernels/time.rs index b9cc64d612cf..b9774d932020 100644 --- a/crates/polars-arrow/src/kernels/time.rs +++ b/crates/polars-arrow/src/kernels/time.rs @@ -7,7 +7,6 @@ use arrow::temporal_conversions::{ }; use chrono::{LocalResult, NaiveDateTime, TimeZone}; use chrono_tz::Tz; -use polars_error::polars_bail; use crate::error::PolarsResult; @@ -37,11 +36,11 @@ fn convert_to_naive_local( } } -fn convert_to_timestamp( - from_tz: Tz, - to_tz: Tz, +pub fn replace_time_zone( arr: &PrimitiveArray, tu: TimeUnit, + from_tz: &Tz, + to_tz: &Tz, use_earliest: Option, ) -> PolarsResult> { let res = match tu { @@ -49,7 +48,7 @@ fn convert_to_timestamp( arr, |value| { let ndt = timestamp_ms_to_datetime(value); - Ok(convert_to_naive_local(&from_tz, &to_tz, ndt, use_earliest)?.timestamp_millis()) + Ok(convert_to_naive_local(from_tz, to_tz, ndt, use_earliest)?.timestamp_millis()) }, ArrowDataType::Int64, ), @@ -57,7 +56,7 @@ fn convert_to_timestamp( arr, |value| { let ndt = timestamp_us_to_datetime(value); - Ok(convert_to_naive_local(&from_tz, &to_tz, ndt, use_earliest)?.timestamp_micros()) + Ok(convert_to_naive_local(from_tz, to_tz, ndt, use_earliest)?.timestamp_micros()) }, ArrowDataType::Int64, ), @@ -65,7 +64,7 @@ fn convert_to_timestamp( arr, |value| { let ndt = timestamp_ns_to_datetime(value); - Ok(convert_to_naive_local(&from_tz, &to_tz, ndt, use_earliest)?.timestamp_nanos()) + Ok(convert_to_naive_local(from_tz, to_tz, ndt, use_earliest)?.timestamp_nanos()) }, ArrowDataType::Int64, ), @@ -73,19 +72,3 @@ fn convert_to_timestamp( }; Ok(res?) } - -pub fn replace_time_zone( - arr: &PrimitiveArray, - tu: TimeUnit, - from: &str, - to: &str, - use_earliest: Option, -) -> PolarsResult> { - match from.parse::() { - Ok(from_tz) => match to.parse::() { - Ok(to_tz) => convert_to_timestamp(from_tz, to_tz, arr, tu, use_earliest), - Err(_) => polars_bail!(ComputeError: "unable to parse time zone: '{}'", to), - }, - Err(_) => polars_bail!(ComputeError: "unable to parse time zone: '{}'", from), - } -} diff --git a/crates/polars-ops/src/chunked_array/datetime/replace_time_zone.rs b/crates/polars-ops/src/chunked_array/datetime/replace_time_zone.rs index 77954a64e065..795da6be9285 100644 --- a/crates/polars-ops/src/chunked_array/datetime/replace_time_zone.rs +++ b/crates/polars-ops/src/chunked_array/datetime/replace_time_zone.rs @@ -9,8 +9,21 @@ pub fn replace_time_zone( let out: PolarsResult<_> = { let from = ca.time_zone().as_deref().unwrap_or("UTC"); let to = time_zone.unwrap_or("UTC"); + let (from_tz, to_tz) = match from.parse::() { + Ok(from_tz) => match to.parse::() { + Ok(to_tz) => (from_tz, to_tz), + Err(_) => polars_bail!(ComputeError: "unable to parse time zone: '{}'", to), + }, + Err(_) => polars_bail!(ComputeError: "unable to parse time zone: '{}'", from), + }; let chunks = ca.downcast_iter().map(|arr| { - replace_time_zone_kernel(arr, ca.time_unit().to_arrow(), from, to, use_earliest) + replace_time_zone_kernel( + arr, + ca.time_unit().to_arrow(), + &from_tz, + &to_tz, + use_earliest, + ) }); let out = ChunkedArray::try_from_chunk_iter(ca.name(), chunks)?; Ok(out.into_datetime(ca.time_unit(), time_zone.map(|x| x.to_string())))