feat: Support use of Duration dtype in to_string

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,24 @@ 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)
+            },
+            _ => Err(PolarsError::InvalidOperation(
+                format!("format {:?} not supported for Duration type (expected one of 'iso' or 'polars')", format).into(),
+            )),
+        }
+    }
+
     /// 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,7 +968,9 @@ 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,
@@ -980,63 +984,102 @@ const SIZES_US: [i64; 4] = [86_400_000_000, 3_600_000_000, 60_000_000, 1_000_000
 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 {
     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(())
-}
+        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 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(())
-}
+    let mut s = String::with_capacity(32);
+    let mut buffer = itoa::Buffer::new();
+    if iso {
+        if v < 0 {
+            s.push_str("-P");
+            v = v.abs()
+        } else {
+            s.push('P');
+        }
+    };
 
-#[cfg(feature = "dtype-duration")]
-fn format_duration(f: &mut Formatter, v: i64, sizes: &[i64], names: &[&str]) -> fmt::Result {
-    for i in 0..4 {
+    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) {
+            s.push_str(buffer.format(whole_num));
+            if iso {
+                if i == 3 {
+                    let secs = match unit {
+                        TimeUnit::Nanoseconds => format!(".{:09}", v % size),
+                        TimeUnit::Microseconds => format!(".{:06}", v % size),
+                        TimeUnit::Milliseconds => format!(".{:03}", v % size),
+                    };
+                    s.push_str(secs.trim_end_matches('0'));
+                }
+                s.push_str(ISO_DURATION_PARTS[i]);
+                if i == 0 {
+                    s.push('T');
+                }
+            } else {
+                s.push_str(DURATION_PARTS[i]);
+                if v % size != 0 {
+                    s.push(' ');
+                }
             }
+        } else if iso && i == 0 {
+            s.push('T');
         }
     }
-    Ok(())
+    if iso {
+        if s.ends_with('T') {
+            s.pop();
+        }
+    } else {
+        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 +1130,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 +1260,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();

crates/polars-core/src/frame/mod.rs

@@ -922,7 +922,7 @@ impl DataFrame {
     /// # Ok::<(), PolarsError>(())
     /// ```
     pub fn height(&self) -> usize {
-        self.shape().0
+        self.height
     }
 
     /// Returns the size as number of rows * number of columns

crates/polars-core/src/series/implementations/datetime.rs

@@ -254,16 +254,8 @@ impl SeriesTrait for SeriesWrap<DatetimeChunked> {
     }
 
     fn cast(&self, dtype: &DataType, cast_options: CastOptions) -> PolarsResult<Series> {
-        match (dtype, self.0.time_unit()) {
-            (DataType::String, TimeUnit::Milliseconds) => {
-                Ok(self.0.to_string("%F %T%.3f")?.into_series())
-            },
-            (DataType::String, TimeUnit::Microseconds) => {
-                Ok(self.0.to_string("%F %T%.6f")?.into_series())
-            },
-            (DataType::String, TimeUnit::Nanoseconds) => {
-                Ok(self.0.to_string("%F %T%.9f")?.into_series())
-            },
+        match dtype {
+            DataType::String => Ok(self.0.to_string("iso")?.into_series()),
             _ => self.0.cast_with_options(dtype, cast_options),
         }
     }