chore(rust): move replace_time_zone to polars-ops (#10078)

polars/polars-core/src/chunked_array/temporal/datetime.rs

@@ -8,8 +8,6 @@ use chrono::format::{DelayedFormat, StrftimeItems};
 use chrono::TimeZone as TimeZoneTrait;
 #[cfg(feature = "timezones")]
 use chrono_tz::Tz;
-#[cfg(feature = "timezones")]
-use polars_arrow::kernels::replace_time_zone;
 
 use super::conversion::{datetime_to_timestamp_ms, datetime_to_timestamp_ns};
 use super::*;
@@ -94,29 +92,6 @@ impl DatetimeChunked {
         }
     }
 
-    #[cfg(feature = "timezones")]
-    pub fn replace_time_zone(
-        &self,
-        time_zone: Option<&str>,
-        use_earliest: Option<bool>,
-    ) -> PolarsResult<DatetimeChunked> {
-        let out: PolarsResult<_> = {
-            let from = self.time_zone().as_deref().unwrap_or("UTC");
-            let to = time_zone.unwrap_or("UTC");
-            let chunks = self
-                .downcast_iter()
-                .map(|arr| {
-                    replace_time_zone(arr, self.time_unit().to_arrow(), from, to, use_earliest)
-                })
-                .collect::<PolarsResult<_>>()?;
-            let out = unsafe { ChunkedArray::from_chunks(self.name(), chunks) };
-            Ok(out.into_datetime(self.time_unit(), time_zone.map(|x| x.to_string())))
-        };
-        let mut out = out?;
-        out.set_sorted_flag(self.is_sorted_flag());
-        Ok(out)
-    }
-
     /// Convert from Datetime into Utf8 with the given format.
     /// See [chrono strftime/strptime](https://docs.rs/chrono/0.4.19/chrono/format/strftime/index.html).
     pub fn to_string(&self, format: &str) -> PolarsResult<Utf8Chunked> {

polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs

@@ -141,11 +141,10 @@ pub(super) fn ordinal_day(s: &Series) -> PolarsResult<Series> {
 pub(super) fn time(s: &Series) -> PolarsResult<Series> {
     match s.dtype() {
         #[cfg(feature = "timezones")]
-        DataType::Datetime(_, Some(_)) => s
-            .datetime()
-            .unwrap()
-            .replace_time_zone(None, None)?
-            .cast(&DataType::Time),
+        DataType::Datetime(_, Some(_)) => {
+            polars_ops::prelude::replace_time_zone(s.datetime().unwrap(), None, None)?
+                .cast(&DataType::Time)
+        }
         DataType::Datetime(_, _) => s.datetime().unwrap().cast(&DataType::Time),
         DataType::Date => s.datetime().unwrap().cast(&DataType::Time),
         DataType::Time => Ok(s.clone()),
@@ -156,11 +155,10 @@ pub(super) fn date(s: &Series) -> PolarsResult<Series> {
     match s.dtype() {
         #[cfg(feature = "timezones")]
         DataType::Datetime(_, Some(tz)) => {
-            let mut out = s
-                .datetime()
-                .unwrap()
-                .replace_time_zone(None, None)?
-                .cast(&DataType::Date)?;
+            let mut out = {
+                polars_ops::chunked_array::replace_time_zone(s.datetime().unwrap(), None, None)?
+                    .cast(&DataType::Date)?
+            };
             if tz != "UTC" {
                 // DST transitions may not preserve sortedness.
                 out.set_sorted_flag(IsSorted::Not);
@@ -176,11 +174,10 @@ pub(super) fn datetime(s: &Series) -> PolarsResult<Series> {
     match s.dtype() {
         #[cfg(feature = "timezones")]
         DataType::Datetime(tu, Some(tz)) => {
-            let mut out = s
-                .datetime()
-                .unwrap()
-                .replace_time_zone(None, None)?
-                .cast(&DataType::Datetime(*tu, None))?;
+            let mut out = {
+                polars_ops::chunked_array::replace_time_zone(s.datetime().unwrap(), None, None)?
+                    .cast(&DataType::Datetime(*tu, None))?
+            };
             if tz != "UTC" {
                 // DST transitions may not preserve sortedness.
                 out.set_sorted_flag(IsSorted::Not);
@@ -319,14 +316,3 @@ pub(super) fn round(s: &Series, every: &str, offset: &str) -> PolarsResult<Serie
         dt => polars_bail!(opq = round, got = dt, expected = "date/datetime"),
     })
 }
-
-#[cfg(feature = "timezones")]
-pub(super) fn replace_time_zone(
-    s: &Series,
-    time_zone: Option<&str>,
-    use_earliest: Option<bool>,
-) -> PolarsResult<Series> {
-    let ca = s.datetime()?;
-    ca.replace_time_zone(time_zone, use_earliest)
-        .map(|ca| ca.into_series())
-}

polars/polars-lazy/polars-plan/src/dsl/function_expr/dispatch.rs

@@ -32,3 +32,13 @@ pub(super) fn set_sorted_flag(s: &Series, sorted: IsSorted) -> PolarsResult<Seri
     s.set_sorted_flag(sorted);
     Ok(s)
 }
+
+#[cfg(feature = "timezones")]
+pub(super) fn replace_time_zone(
+    s: &Series,
+    time_zone: Option<&str>,
+    use_earliest: Option<bool>,
+) -> PolarsResult<Series> {
+    let ca = s.datetime().unwrap();
+    Ok(polars_ops::prelude::replace_time_zone(ca, time_zone, use_earliest)?.into_series())
+}

polars/polars-lazy/polars-plan/src/dsl/function_expr/mod.rs

@@ -746,7 +746,7 @@ impl From<TemporalFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
             Round(every, offset) => map!(datetime::round, &every, &offset),
             #[cfg(feature = "timezones")]
             ReplaceTimeZone(tz, use_earliest) => {
-                map!(datetime::replace_time_zone, tz.as_deref(), use_earliest)
+                map!(dispatch::replace_time_zone, tz.as_deref(), use_earliest)
             }
             Combine(tu) => map_as_slice!(temporal::combine, tu),
             DateRange {

polars/polars-lazy/polars-plan/src/dsl/function_expr/temporal.rs

@@ -79,11 +79,12 @@ pub(super) fn combine(s: &[Series], tu: TimeUnit) -> PolarsResult<Series> {
     let result_naive = datetime + duration;
     match tz {
         #[cfg(feature = "timezones")]
-        Some(tz) => Ok(result_naive
-            .datetime()
-            .unwrap()
-            .replace_time_zone(Some(tz), None)?
-            .into()),
+        Some(tz) => Ok(polars_ops::prelude::replace_time_zone(
+            result_naive.datetime().unwrap(),
+            Some(tz),
+            None,
+        )?
+        .into()),
         _ => Ok(result_naive),
     }
 }
@@ -132,21 +133,22 @@ pub(super) fn temporal_range_dispatch(
     let (mut start, mut stop) = match dtype {
         #[cfg(feature = "timezones")]
         DataType::Datetime(_, Some(_)) => (
-            start
-                .cast(&dtype)?
-                .datetime()
-                .unwrap()
-                .replace_time_zone(None, None)?
-                .into_series()
-                .to_physical_repr()
-                .cast(&DataType::Int64)?,
-            stop.cast(&dtype)?
-                .datetime()
-                .unwrap()
-                .replace_time_zone(None, None)?
-                .into_series()
-                .to_physical_repr()
-                .cast(&DataType::Int64)?,
+            polars_ops::prelude::replace_time_zone(
+                start.cast(&dtype)?.datetime().unwrap(),
+                None,
+                None,
+            )?
+            .into_series()
+            .to_physical_repr()
+            .cast(&DataType::Int64)?,
+            polars_ops::prelude::replace_time_zone(
+                stop.cast(&dtype)?.datetime().unwrap(),
+                None,
+                None,
+            )?
+            .into_series()
+            .to_physical_repr()
+            .cast(&DataType::Int64)?,
         ),
         _ => (
             start

polars/polars-ops/Cargo.toml

@@ -13,6 +13,8 @@ description = "More operations on polars data structures"
 argminmax = { version = "0.6.1", default-features = false, features = ["float"] }
 arrow.workspace = true
 base64 = { version = "0.21", optional = true }
+chrono = { version = "0.4", default-features = false, features = ["std"], optional = true }
+chrono-tz = { version = "0.8", optional = true }
 either.workspace = true
 hex = { version = "0.4", optional = true }
 indexmap.workspace = true
@@ -52,6 +54,7 @@ approx_unique = []
 fused = []
 cutqcut = ["dtype-categorical", "dtype-struct"]
 rle = ["dtype-struct"]
+timezones = ["chrono-tz", "chrono"]
 
 # extra utilities for BinaryChunked
 binary_encoding = ["base64", "hex"]

polars/polars-ops/src/chunked_array/datetime/mod.rs

@@ -0,0 +1,4 @@
+#[cfg(feature = "timezones")]
+mod replace_time_zone;
+#[cfg(feature = "timezones")]
+pub use replace_time_zone::*;

polars/polars-ops/src/chunked_array/datetime/replace_time_zone.rs

@@ -0,0 +1,24 @@
+use polars_arrow::kernels::replace_time_zone as replace_time_zone_kernel;
+use polars_core::prelude::*;
+
+pub fn replace_time_zone(
+    ca: &DatetimeChunked,
+    time_zone: Option<&str>,
+    use_earliest: Option<bool>,
+) -> PolarsResult<DatetimeChunked> {
+    let out: PolarsResult<_> = {
+        let from = ca.time_zone().as_deref().unwrap_or("UTC");
+        let to = time_zone.unwrap_or("UTC");
+        let chunks = ca
+            .downcast_iter()
+            .map(|arr| {
+                replace_time_zone_kernel(arr, ca.time_unit().to_arrow(), from, to, use_earliest)
+            })
+            .collect::<PolarsResult<_>>()?;
+        let out = unsafe { ChunkedArray::from_chunks(ca.name(), chunks) };
+        Ok(out.into_datetime(ca.time_unit(), time_zone.map(|x| x.to_string())))
+    };
+    let mut out = out?;
+    out.set_sorted_flag(ca.is_sorted_flag());
+    Ok(out)
+}

polars/polars-ops/src/chunked_array/mod.rs

@@ -1,6 +1,8 @@
 #[cfg(feature = "dtype-array")]
 pub mod array;
 mod binary;
+#[cfg(feature = "timezones")]
+pub mod datetime;
 #[cfg(feature = "interpolate")]
 mod interpolate;
 pub mod list;
@@ -13,6 +15,8 @@ mod sum;
 mod top_k;
 
 pub use binary::*;
+#[cfg(feature = "timezones")]
+pub use datetime::*;
 #[cfg(feature = "interpolate")]
 pub use interpolate::*;
 pub use list::*;

polars/polars-time/Cargo.toml

@@ -30,7 +30,7 @@ dtype-time = ["polars-core/dtype-time", "polars-core/temporal"]
 dtype-duration = ["polars-core/dtype-duration", "polars-core/temporal"]
 rolling_window = ["polars-core/rolling_window", "dtype-duration"]
 fmt = ["polars-core/fmt"]
-timezones = ["chrono-tz", "dtype-datetime", "polars-core/timezones", "polars-arrow/timezones"]
+timezones = ["chrono-tz", "dtype-datetime", "polars-core/timezones", "polars-arrow/timezones", "polars-ops/timezones"]
 
 test = ["dtype-date", "dtype-datetime", "polars-core/fmt"]
 

polars/polars-time/src/chunkedarray/utf8/infer.rs

@@ -523,14 +523,14 @@ pub(crate) fn to_datetime(
                 Pattern::DatetimeYMDZ => infer.coerce_utf8(ca).datetime().map(|ca| {
                     let mut ca = ca.clone();
                     ca.set_time_unit(tu);
-                    ca.replace_time_zone(Some("UTC"), None)
+                    polars_ops::prelude::replace_time_zone(&ca, Some("UTC"), None)
                 })?,
                 _ => infer.coerce_utf8(ca).datetime().map(|ca| {
                     let mut ca = ca.clone();
                     ca.set_time_unit(tu);
                     match tz {
                         #[cfg(feature = "timezones")]
-                        Some(tz) => ca.replace_time_zone(Some(tz), None),
+                        Some(tz) => polars_ops::prelude::replace_time_zone(&ca, Some(tz), None),
                         _ => Ok(ca),
                     }
                 })?,

polars/polars-time/src/chunkedarray/utf8/mod.rs

@@ -303,7 +303,9 @@ pub trait Utf8Methods: AsUtf8 {
         ca.rename(utf8_ca.name());
         match (tz_aware, tz) {
             #[cfg(feature = "timezones")]
-            (false, Some(tz)) => ca.into_datetime(tu, None).replace_time_zone(Some(tz), None),
+            (false, Some(tz)) => {
+                polars_ops::prelude::replace_time_zone(&ca.into_datetime(tu, None), Some(tz), None)
+            }
             #[cfg(feature = "timezones")]
             (true, _) => Ok(ca.into_datetime(tu, Some("UTC".to_string()))),
             _ => Ok(ca.into_datetime(tu, None)),
@@ -516,7 +518,11 @@ pub trait Utf8Methods: AsUtf8 {
             ca.rename(utf8_ca.name());
             match tz {
                 #[cfg(feature = "timezones")]
-                Some(tz) => ca.into_datetime(tu, None).replace_time_zone(Some(tz), None),
+                Some(tz) => polars_ops::prelude::replace_time_zone(
+                    &ca.into_datetime(tu, None),
+                    Some(tz),
+                    None,
+                ),
                 _ => Ok(ca.into_datetime(tu, None)),
             }
         }