feat: support pl.Time in Series.str.strptime (#3496)

Co-authored-by: Felix Simkovic <felixsimkovic@me.com>
pola-rs · May 25, 2022 · 1ad184f · 1ad184f
1 parent 6a7fd71
commit 1ad184f
Show file tree

Hide file tree

Showing 7 changed files with 55 additions and 2 deletions.
diff --git a/polars/polars-core/src/chunked_array/temporal/time.rs b/polars/polars-core/src/chunked_array/temporal/time.rs
@@ -3,8 +3,15 @@ use crate::prelude::*;
 use arrow::temporal_conversions::{time64ns_to_time, NANOSECONDS};
 use chrono::Timelike;
 
+const SECONDS_IN_MINUTE: i64 = 60;
+const SECONDS_IN_HOUR: i64 = 3_600;
+
 pub(crate) fn time_to_time64ns(time: &NaiveTime) -> i64 {
-    time.second() as i64 * NANOSECONDS + time.nanosecond() as i64
+    (time.hour() as i64 * SECONDS_IN_HOUR
+        + time.minute() as i64 * SECONDS_IN_MINUTE
+        + time.second() as i64)
+        * NANOSECONDS
+        + time.nanosecond() as i64
 }
 
 impl TimeChunked {

diff --git a/polars/polars-core/src/fmt.rs b/polars/polars-core/src/fmt.rs
@@ -251,6 +251,15 @@ impl Debug for Series {
                     "Series"
                 )
             }
+            #[cfg(feature = "dtype-time")]
+            DataType::Time => format_array!(
+                limit,
+                f,
+                self.time().unwrap(),
+                "time",
+                self.name(),
+                "Series"
+            ),
             #[cfg(feature = "dtype-duration")]
             DataType::Duration(_) => {
                 let dt = format!("{}", self.dtype());

diff --git a/polars/polars-lazy/src/dsl/string.rs b/polars/polars-lazy/src/dsl/string.rs
@@ -40,6 +40,16 @@ impl StringNameSpace {
                             .into_series()
                     }
                 }
+                DataType::Time => {
+                    if options.exact {
+                        ca.as_time(options.fmt.as_deref())?.into_series()
+                    } else {
+                        return Err(PolarsError::ComputeError(
+                            format!("non-exact not implemented for dtype {:?}", DataType::Time)
+                                .into(),
+                        ));
+                    }
+                }
                 dt => {
                     return Err(PolarsError::ComputeError(
                         format!("not implemented for dtype {:?}", dt).into(),

diff --git a/polars/polars-time/src/chunkedarray/time.rs b/polars/polars-time/src/chunkedarray/time.rs
@@ -4,8 +4,15 @@ use polars_arrow::export::arrow::array::{MutableArray, MutableUtf8Array, Utf8Arr
 use polars_arrow::export::arrow::temporal_conversions::{time64ns_to_time, NANOSECONDS};
 use std::fmt::Write;
 
+const SECONDS_IN_MINUTE: i64 = 60;
+const SECONDS_IN_HOUR: i64 = 3_600;
+
 pub(crate) fn time_to_time64ns(time: &NaiveTime) -> i64 {
-    time.second() as i64 * NANOSECONDS + time.nanosecond() as i64
+    (time.hour() as i64 * SECONDS_IN_HOUR
+        + time.minute() as i64 * SECONDS_IN_MINUTE
+        + time.second() as i64)
+        * NANOSECONDS
+        + time.nanosecond() as i64
 }
 
 pub trait TimeMethods {

diff --git a/py-polars/polars/internals/expr.py b/py-polars/polars/internals/expr.py
@@ -23,6 +23,7 @@
     Float64,
     Int32,
     Object,
+    Time,
     UInt32,
     py_type_to_dtype,
 )
@@ -3727,6 +3728,8 @@ def strptime(
             return wrap_expr(self._pyexpr.str_parse_date(fmt, strict, exact))
         elif datatype == Datetime:
             return wrap_expr(self._pyexpr.str_parse_datetime(fmt, strict, exact))
+        elif datatype == Time:
+            return wrap_expr(self._pyexpr.str_parse_time(fmt, strict, exact))
         else:
             raise ValueError(
                 "dtype should be of type {Date, Datetime}"

diff --git a/py-polars/src/lazy/dsl.rs b/py-polars/src/lazy/dsl.rs
@@ -427,6 +427,19 @@ impl PyExpr {
             .into()
     }
 
+    pub fn str_parse_time(&self, fmt: Option<String>, strict: bool, exact: bool) -> PyExpr {
+        self.inner
+            .clone()
+            .str()
+            .strptime(StrpTimeOptions {
+                date_dtype: DataType::Time,
+                fmt,
+                strict,
+                exact,
+            })
+            .into()
+    }
+
     pub fn str_strip(&self) -> PyExpr {
         let function = |s: Series| {
             let ca = s.utf8()?;

diff --git a/py-polars/tests/test_series.py b/py-polars/tests/test_series.py
@@ -1242,6 +1242,10 @@ def test_str_strptime() -> None:
         s, expected, "str.strptime", pl.Datetime, "%Y-%m-%d %H:%M:%S"
     )
 
+    s = pl.Series(["00:00:00", "03:20:10"])
+    expected = pl.Series([0, 12010000000000], dtype=pl.Time)
+    verify_series_and_expr_api(s, expected, "str.strptime", pl.Time, "%H:%M:%S")
+
 
 def test_dt_strftime() -> None:
     a = pl.Series("a", [10000, 20000, 30000], dtype=pl.Date)