feat: Support use of Duration in to_string, ergonomic/perf improvement, tz-aware Datetime bugfix

crates/polars-core/Cargo.toml

@@ -26,6 +26,7 @@ either = { workspace = true }
 hashbrown = { workspace = true }
 hashbrown_old_nightly_hack = { workspace = true }
 indexmap = { workspace = true }
+itoa = { workspace = true }
 ndarray = { workspace = true, optional = true }
 num-traits = { workspace = true }
 once_cell = { workspace = true }

crates/polars-core/src/chunked_array/temporal/conversion.rs

@@ -3,7 +3,9 @@ use chrono::*;
 
 use crate::prelude::*;
 
-/// Number of seconds in a day
+pub(crate) const NS_IN_DAY: i64 = 86_400_000_000_000;
+pub(crate) const US_IN_DAY: i64 = 86_400_000_000;
+pub(crate) const MS_IN_DAY: i64 = 86_400_000;
 pub(crate) const SECONDS_IN_DAY: i64 = 86_400;
 
 impl From<&AnyValue<'_>> for NaiveDateTime {
@@ -37,12 +39,10 @@ pub fn datetime_to_timestamp_ns(v: NaiveDateTime) -> i64 {
     v.and_utc().timestamp_nanos_opt().unwrap()
 }
 
-// Used by lazy for literal conversion
 pub fn datetime_to_timestamp_ms(v: NaiveDateTime) -> i64 {
     v.and_utc().timestamp_millis()
 }
 
-// Used by lazy for literal conversion
 pub fn datetime_to_timestamp_us(v: NaiveDateTime) -> i64 {
     let us = v.and_utc().timestamp() * 1_000_000;
     us + v.and_utc().timestamp_subsec_micros() as i64
@@ -52,6 +52,32 @@ pub(crate) fn naive_datetime_to_date(v: NaiveDateTime) -> i32 {
     (datetime_to_timestamp_ms(v) / (MILLISECONDS * SECONDS_IN_DAY)) as i32
 }
 
-pub(crate) const NS_IN_DAY: i64 = 86_400_000_000_000;
-pub(crate) const US_IN_DAY: i64 = 86_400_000_000;
-pub(crate) const MS_IN_DAY: i64 = 86_400_000;
+pub fn get_strftime_format(fmt: &str, dtype: &DataType) -> String {
+    if fmt != "iso" {
+        return fmt.to_string();
+    }
+    #[allow(unreachable_code)]
+    let fmt: &str = match dtype {
+        #[cfg(feature = "dtype-datetime")]
+        DataType::Datetime(tu, tz) => match (tu, tz.is_some()) {
+            (TimeUnit::Milliseconds, true) => "%F %T%.3f%:z",
+            (TimeUnit::Milliseconds, false) => "%F %T%.3f",
+            (TimeUnit::Microseconds, true) => "%F %T%.6f%:z",
+            (TimeUnit::Microseconds, false) => "%F %T%.6f",
+            (TimeUnit::Nanoseconds, true) => "%F %T%.9f%:z",
+            (TimeUnit::Nanoseconds, false) => "%F %T%.9f",
+        },
+        #[cfg(feature = "dtype-date")]
+        DataType::Date => "%F",
+        #[cfg(feature = "dtype-time")]
+        DataType::Time => "%T%.f",
+        _ => {
+            let err = format!(
+                "invalid call to `get_strftime_format`; fmt={:?}, dtype={}",
+                fmt, dtype
+            );
+            unimplemented!("{}", err)
+        },
+    };
+    fmt.to_string()
+}

crates/polars-core/src/chunked_array/temporal/date.rs

@@ -33,6 +33,7 @@ impl DateChunked {
     /// Convert from Date into String 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<StringChunked> {
+        let format = if format == "iso" { "%F" } else { format };
         let datefmt_f = |ndt: NaiveDate| ndt.format(format);
         self.try_apply_into_string_amortized(|val, buf| {
             let ndt = date32_to_date(val);

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

@@ -47,12 +47,12 @@ impl DatetimeChunked {
             TimeUnit::Microseconds => timestamp_us_to_datetime,
             TimeUnit::Milliseconds => timestamp_ms_to_datetime,
         };
-
+        let format = get_strftime_format(format, self.dtype());
         let mut ca: StringChunked = match self.time_zone() {
             #[cfg(feature = "timezones")]
             Some(time_zone) => {
                 let parsed_time_zone = time_zone.parse::<Tz>().expect("already validated");
-                let datefmt_f = |ndt| parsed_time_zone.from_utc_datetime(&ndt).format(format);
+                let datefmt_f = |ndt| parsed_time_zone.from_utc_datetime(&ndt).format(&format);
                 self.try_apply_into_string_amortized(|val, buf| {
                     let ndt = conversion_f(val);
                     write!(buf, "{}", datefmt_f(ndt))
@@ -62,7 +62,7 @@ impl DatetimeChunked {
                 )?
             },
             _ => {
-                let datefmt_f = |ndt: NaiveDateTime| ndt.format(format);
+                let datefmt_f = |ndt: NaiveDateTime| ndt.format(&format);
                 self.try_apply_into_string_amortized(|val, buf| {
                     let ndt = conversion_f(val);
                     write!(buf, "{}", datefmt_f(ndt))

crates/polars-core/src/chunked_array/temporal/duration.rs

@@ -1,4 +1,5 @@
 use crate::export::chrono::Duration as ChronoDuration;
+use crate::fmt::fmt_duration_string;
 use crate::prelude::DataType::Duration;
 use crate::prelude::*;
 
@@ -60,6 +61,27 @@ impl DurationChunked {
         self.2 = Some(Duration(tu))
     }
 
+    /// Convert from [`Duration`] to String; note that `strftime` format
+    /// strings are not supported, only the specifiers 'iso' and 'polars'.
+    pub fn to_string(&self, format: &str) -> PolarsResult<StringChunked> {
+        match format {
+            "iso" | "polars" => {
+                let out: StringChunked = self
+                    .0
+                    .apply_nonnull_values_generic(DataType::String, |v: i64| {
+                        fmt_duration_string(v, self.time_unit(), format == "iso")
+                    });
+                Ok(out)
+            },
+            _ => {
+                polars_bail!(
+                    InvalidOperation: "format {:?} not supported for Duration type (expected one of 'iso' or 'polars')",
+                    format
+                )
+            },
+        }
+    }
+
     /// Construct a new [`DurationChunked`] from an iterator over [`ChronoDuration`].
     pub fn from_duration<I: IntoIterator<Item = ChronoDuration>>(
         name: PlSmallStr,

crates/polars-core/src/chunked_array/temporal/time.rs

@@ -23,6 +23,7 @@ impl TimeChunked {
     pub fn to_string(&self, format: &str) -> StringChunked {
         let mut ca: StringChunked = self.apply_kernel_cast(&|arr| {
             let mut buf = String::new();
+            let format = if format == "iso" { "%T%.9f" } else { format };
             let mut mutarr = MutablePlString::with_capacity(arr.len());
 
             for opt in arr.into_iter() {

crates/polars-core/src/fmt.rs

@@ -22,6 +22,8 @@ use comfy_table::modifiers::*;
 use comfy_table::presets::*;
 #[cfg(any(feature = "fmt", feature = "fmt_no_tty"))]
 use comfy_table::*;
+#[cfg(feature = "dtype-duration")]
+use itoa;
 use num_traits::{Num, NumCast};
 
 use crate::config::*;
@@ -966,77 +968,157 @@ fn fmt_datetime(
 }
 
 #[cfg(feature = "dtype-duration")]
-const NAMES: [&str; 4] = ["d", "h", "m", "s"];
+const DURATION_PARTS: [&str; 4] = ["d", "h", "m", "s"];
+#[cfg(feature = "dtype-duration")]
+const ISO_DURATION_PARTS: [&str; 4] = ["D", "H", "M", "S"];
 #[cfg(feature = "dtype-duration")]
 const SIZES_NS: [i64; 4] = [
-    86_400_000_000_000,
-    3_600_000_000_000,
-    60_000_000_000,
-    1_000_000_000,
+    86_400_000_000_000, // per day
+    3_600_000_000_000,  // per hour
+    60_000_000_000,     // per minute
+    1_000_000_000,      // per second
 ];
 #[cfg(feature = "dtype-duration")]
 const SIZES_US: [i64; 4] = [86_400_000_000, 3_600_000_000, 60_000_000, 1_000_000];
 #[cfg(feature = "dtype-duration")]
 const SIZES_MS: [i64; 4] = [86_400_000, 3_600_000, 60_000, 1_000];
 
 #[cfg(feature = "dtype-duration")]
-fn fmt_duration_ns(f: &mut Formatter<'_>, v: i64) -> fmt::Result {
+pub fn fmt_duration_string(mut v: i64, unit: TimeUnit, iso: bool) -> String {
+    // take the physical/integer duration value and return either a human-readable version
+    // of the duration (as used in the Polars frame repr) or an ISO8601 duration string.
+    //
+    // Polars: "3d 22m 55s 1ms"
+    // ISO: "P3DT22M55.001S"
+    //
+    // the parts (days, hours, minutes, seconds) occur in the same order in
+    // each string, so we use the same code to generate each of them, with
+    // only the separators and the 'seconds' part differing.
+    //
+    // ref: https://en.wikipedia.org/wiki/ISO_8601#Durations
     if v == 0 {
-        return write!(f, "0ns");
-    }
-    format_duration(f, v, SIZES_NS.as_slice(), NAMES.as_slice())?;
-    if v % 1000 != 0 {
-        write!(f, "{}ns", v % 1_000_000_000)?;
-    } else if v % 1_000_000 != 0 {
-        write!(f, "{}µs", (v % 1_000_000_000) / 1000)?;
-    } else if v % 1_000_000_000 != 0 {
-        write!(f, "{}ms", (v % 1_000_000_000) / 1_000_000)?;
-    }
-    Ok(())
-}
-
-#[cfg(feature = "dtype-duration")]
-fn fmt_duration_us(f: &mut Formatter<'_>, v: i64) -> fmt::Result {
-    if v == 0 {
-        return write!(f, "0µs");
-    }
-    format_duration(f, v, SIZES_US.as_slice(), NAMES.as_slice())?;
-    if v % 1000 != 0 {
-        write!(f, "{}µs", (v % 1_000_000))?;
-    } else if v % 1_000_000 != 0 {
-        write!(f, "{}ms", (v % 1_000_000) / 1_000)?;
-    }
-    Ok(())
-}
-
-#[cfg(feature = "dtype-duration")]
-fn fmt_duration_ms(f: &mut Formatter<'_>, v: i64) -> fmt::Result {
-    if v == 0 {
-        return write!(f, "0ms");
-    }
-    format_duration(f, v, SIZES_MS.as_slice(), NAMES.as_slice())?;
-    if v % 1_000 != 0 {
-        write!(f, "{}ms", (v % 1_000))?;
-    }
-    Ok(())
-}
+        return if iso {
+            "PT0S".to_string()
+        } else {
+            match unit {
+                TimeUnit::Nanoseconds => "0ns".to_string(),
+                TimeUnit::Microseconds => "0µs".to_string(),
+                TimeUnit::Milliseconds => "0ms".to_string(),
+            }
+        };
+    };
+    let sizes = match unit {
+        TimeUnit::Nanoseconds => SIZES_NS.as_slice(),
+        TimeUnit::Microseconds => SIZES_US.as_slice(),
+        TimeUnit::Milliseconds => SIZES_MS.as_slice(),
+    };
 
-#[cfg(feature = "dtype-duration")]
-fn format_duration(f: &mut Formatter, v: i64, sizes: &[i64], names: &[&str]) -> fmt::Result {
-    for i in 0..4 {
+    let mut s = String::with_capacity(32);
+    let mut buffer = itoa::Buffer::new();
+    let mut wrote_part = false;
+    if iso {
+        if v < 0 {
+            // negative sign before "P" indicates that the entire ISO duration is negative.
+            // the Polars version applies a negative sign to each *individual* part.
+            s.push_str("-P");
+            v = v.abs()
+        } else {
+            s.push('P');
+        }
+    };
+    // iterate over dtype-specific sizes to appropriately scale
+    // and extract 'days', 'hours', 'minutes', and 'seconds' parts.
+    for (i, &size) in sizes.iter().enumerate() {
         let whole_num = if i == 0 {
-            v / sizes[i]
+            v / size
         } else {
-            (v % sizes[i - 1]) / sizes[i]
+            (v % sizes[i - 1]) / size
         };
-        if whole_num <= -1 || whole_num >= 1 {
-            write!(f, "{}{}", whole_num, names[i])?;
-            if v % sizes[i] != 0 {
-                write!(f, " ")?;
+        if whole_num != 0 || (iso && i == 3) {
+            if iso {
+                if i != 3 {
+                    // days, hours, minutes
+                    s.push_str(buffer.format(whole_num));
+                    s.push_str(ISO_DURATION_PARTS[i]);
+                } else {
+                    // (index 3 => 'seconds' part): the ISO version writes
+                    // fractional seconds, not integer nano/micro/milliseconds.
+                    // if zero, only write out if no other parts written yet.
+                    let fractional_part = v % size;
+                    if whole_num == 0 && fractional_part == 0 {
+                        if !wrote_part {
+                            s.push_str("0S")
+                        }
+                    } else {
+                        s.push_str(buffer.format(whole_num));
+                        if fractional_part != 0 {
+                            let secs = match unit {
+                                TimeUnit::Nanoseconds => format!(".{:09}", fractional_part),
+                                TimeUnit::Microseconds => format!(".{:06}", fractional_part),
+                                TimeUnit::Milliseconds => format!(".{:03}", fractional_part),
+                            };
+                            s.push_str(secs.trim_end_matches('0'));
+                        }
+                        s.push_str(ISO_DURATION_PARTS[i]);
+                    }
+                }
+                // (index 0 => 'days' part): after writing days above (if non-zero)
+                // the ISO duration string requires a `T` before the time part.
+                if i == 0 {
+                    s.push('T');
+                }
+            } else {
+                s.push_str(buffer.format(whole_num));
+                s.push_str(DURATION_PARTS[i]);
+                if v % size != 0 {
+                    s.push(' ');
+                }
             }
+            wrote_part = true;
+        } else if iso && i == 0 {
+            // always need to write the `T` separator for ISO
+            // durations, even if there is no 'days' part.
+            s.push('T');
         }
     }
-    Ok(())
+    if iso {
+        // if there was only a 'days' component, no need for time separator.
+        if s.ends_with('T') {
+            s.pop();
+        }
+    } else {
+        // write out fractional seconds as integer nano/micro/milliseconds.
+        match unit {
+            TimeUnit::Nanoseconds => {
+                if v % 1000 != 0 {
+                    s.push_str(buffer.format(v % 1_000_000_000));
+                    s.push_str("ns");
+                } else if v % 1_000_000 != 0 {
+                    s.push_str(buffer.format((v % 1_000_000_000) / 1000));
+                    s.push_str("µs");
+                } else if v % 1_000_000_000 != 0 {
+                    s.push_str(buffer.format((v % 1_000_000_000) / 1_000_000));
+                    s.push_str("ms");
+                }
+            },
+            TimeUnit::Microseconds => {
+                if v % 1000 != 0 {
+                    s.push_str(buffer.format(v % 1_000_000));
+                    s.push_str("µs");
+                } else if v % 1_000_000 != 0 {
+                    s.push_str(buffer.format((v % 1_000_000) / 1_000));
+                    s.push_str("ms");
+                }
+            },
+            TimeUnit::Milliseconds => {
+                if v % 1000 != 0 {
+                    s.push_str(buffer.format(v % 1_000));
+                    s.push_str("ms");
+                }
+            },
+        }
+    }
+    s
 }
 
 fn format_blob(f: &mut Formatter<'_>, bytes: &[u8]) -> fmt::Result {
@@ -1087,11 +1169,7 @@ impl Display for AnyValue<'_> {
                 fmt_datetime(f, *v, *tu, tz.as_ref().map(|v| v.as_ref()))
             },
             #[cfg(feature = "dtype-duration")]
-            AnyValue::Duration(v, tu) => match tu {
-                TimeUnit::Nanoseconds => fmt_duration_ns(f, *v),
-                TimeUnit::Microseconds => fmt_duration_us(f, *v),
-                TimeUnit::Milliseconds => fmt_duration_ms(f, *v),
-            },
+            AnyValue::Duration(v, tu) => write!(f, "{}", fmt_duration_string(*v, *tu, false)),
             #[cfg(feature = "dtype-time")]
             AnyValue::Time(_) => {
                 let nt: chrono::NaiveTime = self.into();
@@ -1221,7 +1299,7 @@ impl Series {
 
 #[inline]
 #[cfg(feature = "dtype-decimal")]
-pub fn fmt_decimal(f: &mut Formatter<'_>, v: i128, scale: usize) -> fmt::Result {
+fn fmt_decimal(f: &mut Formatter<'_>, v: i128, scale: usize) -> fmt::Result {
     use arrow::compute::decimal::format_decimal;
 
     let trim_zeros = get_trim_decimal_zeros();