feat: add non_existent arg to replace_time_zone

crates/polars-arrow/src/legacy/kernels/mod.rs

@@ -24,9 +24,9 @@ pub mod string;
 pub mod take_agg;
 mod time;
 
-pub use time::Ambiguous;
 #[cfg(feature = "timezones")]
 pub use time::{convert_to_naive_local, convert_to_naive_local_opt};
+pub use time::{Ambiguous, NonExistent};
 
 /// Internal state of [SlicesIterator]
 #[derive(Debug, PartialEq)]

crates/polars-arrow/src/legacy/kernels/time.rs

@@ -7,6 +7,8 @@ use chrono_tz::Tz;
 #[cfg(feature = "timezones")]
 use polars_error::PolarsResult;
 use polars_error::{polars_bail, PolarsError};
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
 
 pub enum Ambiguous {
     Earliest,
@@ -30,12 +32,20 @@ impl FromStr for Ambiguous {
     }
 }
 
+#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum NonExistent {
+    Null,
+    Raise,
+}
+
 #[cfg(feature = "timezones")]
 pub fn convert_to_naive_local(
     from_tz: &Tz,
     to_tz: &Tz,
     ndt: NaiveDateTime,
     ambiguous: Ambiguous,
+    non_existent: NonExistent,
 ) -> PolarsResult<Option<NaiveDateTime>> {
     let ndt = from_tz.from_utc_datetime(&ndt).naive_local();
     match to_tz.from_local_datetime(&ndt) {
@@ -48,10 +58,13 @@ pub fn convert_to_naive_local(
                 polars_bail!(ComputeError: "datetime '{}' is ambiguous in time zone '{}'. Please use `ambiguous` to tell how it should be localized.", ndt, to_tz)
             },
         },
-        LocalResult::None => polars_bail!(ComputeError:
-                "datetime '{}' is non-existent in time zone '{}'. Non-existent datetimes are not yet supported",
+        LocalResult::None => match non_existent {
+            NonExistent::Raise => polars_bail!(ComputeError:
+                "datetime '{}' is non-existent in time zone '{}'. You can use `non_existent='null'` to return `null` in this case.",
                 ndt, to_tz
-        ),
+            ),
+            NonExistent::Null => Ok(None),
+        },
     }
 }
 

crates/polars-arrow/src/legacy/prelude.rs

@@ -5,6 +5,7 @@ pub use crate::legacy::bitmap::mutable::MutableBitmapExtension;
 pub use crate::legacy::index::*;
 pub use crate::legacy::kernels::rolling::no_nulls::QuantileInterpolOptions;
 pub use crate::legacy::kernels::rolling::{DynArgs, RollingQuantileParams, RollingVarParams};
+pub use crate::legacy::kernels::{Ambiguous, NonExistent};
 
 pub type LargeStringArray = Utf8Array<i64>;
 pub type LargeBinaryArray = BinaryArray<i64>;

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

@@ -1,6 +1,6 @@
 use std::str::FromStr;
 
-use arrow::legacy::kernels::{convert_to_naive_local, Ambiguous};
+use arrow::legacy::kernels::convert_to_naive_local;
 use arrow::temporal_conversions::{
     timestamp_ms_to_datetime, timestamp_ns_to_datetime, timestamp_us_to_datetime,
 };
@@ -14,6 +14,7 @@ pub fn replace_time_zone(
     datetime: &Logical<DatetimeType, Int64Type>,
     time_zone: Option<&str>,
     ambiguous: &StringChunked,
+    non_existent: NonExistent,
 ) -> PolarsResult<DatetimeChunked> {
     let from_time_zone = datetime.time_zone().as_deref().unwrap_or("UTC");
     let from_tz = parse_time_zone(from_time_zone)?;
@@ -39,7 +40,10 @@ pub fn replace_time_zone(
         TimeUnit::Nanoseconds => datetime_to_timestamp_ns,
     };
 
-    let out = if ambiguous.len() == 1 && ambiguous.get(0) != Some("null") {
+    let out = if ambiguous.len() == 1
+        && ambiguous.get(0) != Some("null")
+        && non_existent == NonExistent::Raise
+    {
         impl_replace_time_zone_fast(
             datetime,
             ambiguous.get(0),
@@ -52,6 +56,7 @@ pub fn replace_time_zone(
         impl_replace_time_zone(
             datetime,
             ambiguous,
+            non_existent,
             timestamp_to_datetime,
             datetime_to_timestamp,
             &from_tz,
@@ -84,8 +89,14 @@ pub fn impl_replace_time_zone_fast(
         Some(ambiguous) => datetime.0.try_apply(|timestamp| {
             let ndt = timestamp_to_datetime(timestamp);
             Ok(datetime_to_timestamp(
-                convert_to_naive_local(from_tz, to_tz, ndt, Ambiguous::from_str(ambiguous)?)?
-                    .expect("we didn't use Ambiguous::Null"),
+                convert_to_naive_local(
+                    from_tz,
+                    to_tz,
+                    ndt,
+                    Ambiguous::from_str(ambiguous)?,
+                    NonExistent::Raise,
+                )?
+                .expect("we didn't use Ambiguous::Null or NonExistent::Null"),
             ))
         }),
         _ => Ok(datetime.0.apply(|_| None)),
@@ -95,14 +106,14 @@ pub fn impl_replace_time_zone_fast(
 pub fn impl_replace_time_zone(
     datetime: &Logical<DatetimeType, Int64Type>,
     ambiguous: &StringChunked,
+    non_existent: NonExistent,
     timestamp_to_datetime: fn(i64) -> NaiveDateTime,
     datetime_to_timestamp: fn(NaiveDateTime) -> i64,
     from_tz: &chrono_tz::Tz,
     to_tz: &chrono_tz::Tz,
 ) -> PolarsResult<Int64Chunked> {
     match ambiguous.len() {
         1 => {
-            debug_assert!(ambiguous.get(0) == Some("null"));
             let iter = datetime.0.downcast_iter().map(|arr| {
                 let element_iter = arr.iter().map(|timestamp_opt| match timestamp_opt {
                     Some(timestamp) => {
@@ -111,7 +122,8 @@ pub fn impl_replace_time_zone(
                             from_tz,
                             to_tz,
                             ndt,
-                            Ambiguous::from_str("null")?,
+                            Ambiguous::from_str(ambiguous.get(0).unwrap())?,
+                            non_existent,
                         )?;
                         Ok::<_, PolarsError>(res.map(datetime_to_timestamp))
                     },
@@ -130,6 +142,7 @@ pub fn impl_replace_time_zone(
                         to_tz,
                         ndt,
                         Ambiguous::from_str(ambiguous)?,
+                        non_existent,
                     )?
                     .map(datetime_to_timestamp))
                 },

crates/polars-plan/src/dsl/dt.rs

@@ -238,9 +238,14 @@ impl DateLikeNameSpace {
     }
 
     #[cfg(feature = "timezones")]
-    pub fn replace_time_zone(self, time_zone: Option<TimeZone>, ambiguous: Expr) -> Expr {
+    pub fn replace_time_zone(
+        self,
+        time_zone: Option<TimeZone>,
+        ambiguous: Expr,
+        non_existent: NonExistent,
+    ) -> Expr {
         self.0.map_many_private(
-            FunctionExpr::TemporalExpr(TemporalFunction::ReplaceTimeZone(time_zone)),
+            FunctionExpr::TemporalExpr(TemporalFunction::ReplaceTimeZone(time_zone, non_existent)),
             &[ambiguous],
             false,
             false,

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

@@ -58,7 +58,7 @@ pub enum TemporalFunction {
     DSTOffset,
     Round(String, String),
     #[cfg(feature = "timezones")]
-    ReplaceTimeZone(Option<TimeZone>),
+    ReplaceTimeZone(Option<TimeZone>, NonExistent),
     Combine(TimeUnit),
     DatetimeFunction {
         time_unit: TimeUnit,
@@ -111,7 +111,7 @@ impl TemporalFunction {
             DSTOffset => mapper.with_dtype(DataType::Duration(TimeUnit::Milliseconds)),
             Round(..) => mapper.with_same_dtype(),
             #[cfg(feature = "timezones")]
-            ReplaceTimeZone(tz) => mapper.map_datetime_dtype_timezone(tz.as_ref()),
+            ReplaceTimeZone(tz, _non_existent) => mapper.map_datetime_dtype_timezone(tz.as_ref()),
             DatetimeFunction {
                 time_unit,
                 time_zone,
@@ -179,7 +179,7 @@ impl Display for TemporalFunction {
             DSTOffset => "dst_offset",
             Round(..) => "round",
             #[cfg(feature = "timezones")]
-            ReplaceTimeZone(_) => "replace_time_zone",
+            ReplaceTimeZone(_, _) => "replace_time_zone",
             DatetimeFunction { .. } => return write!(f, "dt.datetime"),
             Combine(_) => "combine",
         };
@@ -227,6 +227,7 @@ pub(super) fn time(s: &Series) -> PolarsResult<Series> {
             s.datetime().unwrap(),
             None,
             &StringChunked::from_iter(std::iter::once("raise")),
+            NonExistent::Raise,
         )?
         .cast(&DataType::Time),
         DataType::Datetime(_, _) => s.datetime().unwrap().cast(&DataType::Time),
@@ -243,6 +244,7 @@ pub(super) fn date(s: &Series) -> PolarsResult<Series> {
                     s.datetime().unwrap(),
                     None,
                     &StringChunked::from_iter(std::iter::once("raise")),
+                    NonExistent::Raise,
                 )?
                 .cast(&DataType::Date)?
             };
@@ -266,6 +268,7 @@ pub(super) fn datetime(s: &Series) -> PolarsResult<Series> {
                     s.datetime().unwrap(),
                     None,
                     &StringChunked::from_iter(std::iter::once("raise")),
+                    NonExistent::Raise,
                 )?
                 .cast(&DataType::Datetime(*tu, None))?
             };

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

@@ -35,11 +35,15 @@ pub(super) fn set_sorted_flag(s: &Series, sorted: IsSorted) -> PolarsResult<Seri
 }
 
 #[cfg(feature = "timezones")]
-pub(super) fn replace_time_zone(s: &[Series], time_zone: Option<&str>) -> PolarsResult<Series> {
+pub(super) fn replace_time_zone(
+    s: &[Series],
+    time_zone: Option<&str>,
+    non_existent: NonExistent,
+) -> PolarsResult<Series> {
     let s1 = &s[0];
     let ca = s1.datetime().unwrap();
     let s2 = &s[1].str()?;
-    Ok(polars_ops::prelude::replace_time_zone(ca, time_zone, s2)?.into_series())
+    Ok(polars_ops::prelude::replace_time_zone(ca, time_zone, s2, non_existent)?.into_series())
 }
 
 #[cfg(feature = "dtype-struct")]

crates/polars-plan/src/dsl/function_expr/range/datetime_range.rs

@@ -66,13 +66,15 @@ pub(super) fn datetime_range(
                 start.datetime().unwrap(),
                 Some(&tz),
                 &StringChunked::from_iter(std::iter::once("raise")),
+                NonExistent::Raise,
             )?
             .cast(&dtype)?
             .into_series(),
             polars_ops::prelude::replace_time_zone(
                 end.datetime().unwrap(),
                 Some(&tz),
                 &StringChunked::from_iter(std::iter::once("raise")),
+                NonExistent::Raise,
             )?
             .cast(&dtype)?
             .into_series(),
@@ -153,6 +155,7 @@ pub(super) fn datetime_ranges(
                 start.datetime().unwrap(),
                 Some(&tz),
                 &StringChunked::from_iter(std::iter::once("raise")),
+                NonExistent::Raise,
             )?
             .cast(&dtype)?
             .into_series()
@@ -162,6 +165,7 @@ pub(super) fn datetime_ranges(
                 end.datetime().unwrap(),
                 Some(&tz),
                 &StringChunked::from_iter(std::iter::once("raise")),
+                NonExistent::Raise,
             )?
             .cast(&dtype)?
             .into_series()

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

@@ -59,8 +59,8 @@ impl From<TemporalFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
             DSTOffset => map!(datetime::dst_offset),
             Round(every, offset) => map_as_slice!(datetime::round, &every, &offset),
             #[cfg(feature = "timezones")]
-            ReplaceTimeZone(tz) => {
-                map_as_slice!(dispatch::replace_time_zone, tz.as_deref())
+            ReplaceTimeZone(tz, non_existent) => {
+                map_as_slice!(dispatch::replace_time_zone, tz.as_deref(), non_existent)
             },
             Combine(tu) => map_as_slice!(temporal::combine, tu),
             DatetimeFunction {
@@ -168,7 +168,7 @@ pub(super) fn datetime(
         #[cfg(feature = "timezones")]
         Some(_) => {
             let mut ca = ca.into_datetime(*time_unit, None);
-            ca = replace_time_zone(&ca, time_zone, _ambiguous)?;
+            ca = replace_time_zone(&ca, time_zone, _ambiguous, NonExistent::Raise)?;
             ca
         },
         _ => {
@@ -314,6 +314,7 @@ pub(super) fn combine(s: &[Series], tu: TimeUnit) -> PolarsResult<Series> {
             result_naive.datetime().unwrap(),
             Some(tz),
             &StringChunked::from_iter(std::iter::once("raise")),
+            NonExistent::Raise,
         )?
         .into()),
         _ => Ok(result_naive),

crates/polars-time/src/chunkedarray/string/infer.rs

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

crates/polars-time/src/chunkedarray/string/mod.rs

@@ -2,7 +2,6 @@ pub mod infer;
 use chrono::DateTime;
 mod patterns;
 mod strptime;
-
 use chrono::ParseError;
 pub use patterns::Pattern;
 #[cfg(feature = "dtype-time")]
@@ -216,6 +215,7 @@ pub trait StringMethods: AsString {
                 &ca.into_datetime(tu, None),
                 Some(tz),
                 _ambiguous,
+                NonExistent::Raise,
             ),
             #[cfg(feature = "timezones")]
             (true, _) => Ok(ca.into_datetime(tu, Some("UTC".to_string()))),
@@ -323,7 +323,12 @@ pub trait StringMethods: AsString {
             let dt = ca.with_name(string_ca.name()).into_datetime(tu, None);
             match tz {
                 #[cfg(feature = "timezones")]
-                Some(tz) => polars_ops::prelude::replace_time_zone(&dt, Some(tz), ambiguous),
+                Some(tz) => polars_ops::prelude::replace_time_zone(
+                    &dt,
+                    Some(tz),
+                    ambiguous,
+                    NonExistent::Raise,
+                ),
                 _ => Ok(dt),
             }
         }

crates/polars-time/src/month_start.rs

@@ -1,5 +1,3 @@
-#[cfg(feature = "timezones")]
-use arrow::legacy::kernels::Ambiguous;
 use arrow::legacy::time_zone::Tz;
 use chrono::{Datelike, NaiveDate, NaiveDateTime, NaiveTime, Timelike};
 use polars_core::prelude::*;
@@ -50,8 +48,8 @@ pub(crate) fn roll_backward(
     let t = match tz {
         #[cfg(feature = "timezones")]
         Some(tz) => datetime_to_timestamp(
-            try_localize_datetime(ndt, tz, Ambiguous::Raise)?
-                .expect("we didn't use Ambiguous::Null"),
+            try_localize_datetime(ndt, tz, Ambiguous::Raise, NonExistent::Raise)?
+                .expect("we didn't use Ambiguous::Null or NonExistent::Null"),
         ),
         _ => datetime_to_timestamp(ndt),
     };

crates/polars-time/src/utils.rs

@@ -1,5 +1,7 @@
 #[cfg(feature = "timezones")]
-use arrow::legacy::kernels::{convert_to_naive_local, convert_to_naive_local_opt, Ambiguous};
+use arrow::legacy::kernels::{
+    convert_to_naive_local, convert_to_naive_local_opt, Ambiguous, NonExistent,
+};
 #[cfg(feature = "timezones")]
 use arrow::legacy::time_zone::Tz;
 #[cfg(feature = "timezones")]
@@ -22,8 +24,9 @@ pub(crate) fn try_localize_datetime(
     ndt: NaiveDateTime,
     tz: &Tz,
     ambiguous: Ambiguous,
+    non_existent: NonExistent,
 ) -> PolarsResult<Option<NaiveDateTime>> {
-    convert_to_naive_local(&chrono_tz::UTC, tz, ndt, ambiguous)
+    convert_to_naive_local(&chrono_tz::UTC, tz, ndt, ambiguous, non_existent)
 }
 
 #[cfg(feature = "timezones")]