perf(rust, python): parse time zones outside of downcast_iter() in re…

…place_time_zone
pola-rs · Aug 24, 2023 · 38e3e7d · 38e3e7d
1 parent 9bfa5b6
commit 38e3e7d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 24 deletions.
diff --git 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,55 +36,39 @@ fn convert_to_naive_local(
     }
 }
 
-fn convert_to_timestamp(
-    from_tz: Tz,
-    to_tz: Tz,
+pub fn replace_time_zone(
     arr: &PrimitiveArray<i64>,
     tu: TimeUnit,
+    from_tz: &Tz,
+    to_tz: &Tz,
     use_earliest: Option<bool>,
 ) -> PolarsResult<PrimitiveArray<i64>> {
     let res = match tu {
         TimeUnit::Millisecond => try_unary(
             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,
         ),
         TimeUnit::Microsecond => try_unary(
             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,
         ),
         TimeUnit::Nanosecond => try_unary(
             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,
         ),
         _ => unreachable!(),
     };
     Ok(res?)
 }
-
-pub fn replace_time_zone(
-    arr: &PrimitiveArray<i64>,
-    tu: TimeUnit,
-    from: &str,
-    to: &str,
-    use_earliest: Option<bool>,
-) -> PolarsResult<PrimitiveArray<i64>> {
-    match from.parse::<chrono_tz::Tz>() {
-        Ok(from_tz) => match to.parse::<chrono_tz::Tz>() {
-            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
@@ -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::<chrono_tz::Tz>() {
+            Ok(from_tz) => match to.parse::<chrono_tz::Tz>() {
+                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())))