Repository: narwhals-dev/narwhals
Pull request #167 (2024-05-13 17:32:02)

add temporal methods for calculating total seconds and milli/micro/nano seconds
Related to #119 


Patch:
diff --git a/docs/api-reference/expressions_dt.md b/docs/api-reference/expressions_dt.md
index 705a35e6..1ac535bd 100644
--- a/docs/api-reference/expressions_dt.md
+++ b/docs/api-reference/expressions_dt.md
@@ -12,5 +12,9 @@
         - minute
         - second
         - total_minutes
+        - total_seconds
+        - total_milliseconds
+        - total_microseconds
+        - total_nanoseconds
       show_source: false
       show_bases: false
diff --git a/docs/api-reference/series_dt.md b/docs/api-reference/series_dt.md
index 3f2a123c..41becd0b 100644
--- a/docs/api-reference/series_dt.md
+++ b/docs/api-reference/series_dt.md
@@ -12,5 +12,9 @@
         - minute
         - second
         - total_minutes
+        - total_seconds
+        - total_milliseconds
+        - total_microseconds
+        - total_nanoseconds
       show_source: false
       show_bases: false
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 50b59d96..c3273e1b 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -318,3 +318,15 @@ class PandasExprDateTimeNamespace:
 
     def total_minutes(self) -> PandasExpr:
         return register_namespace_expression_call(self._expr, "dt", "total_minutes")
+
+    def total_seconds(self) -> PandasExpr:
+        return register_namespace_expression_call(self._expr, "dt", "total_seconds")
+
+    def total_milliseconds(self) -> PandasExpr:
+        return register_namespace_expression_call(self._expr, "dt", "total_milliseconds")
+
+    def total_microseconds(self) -> PandasExpr:
+        return register_namespace_expression_call(self._expr, "dt", "total_microseconds")
+
+    def total_nanoseconds(self) -> PandasExpr:
+        return register_namespace_expression_call(self._expr, "dt", "total_nanoseconds")
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 38a4026c..64c64b92 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -511,3 +511,43 @@ class PandasSeriesDateTimeNamespace:
         if ~s.isna().any():
             s_abs = s_abs.astype(int)
         return self._series._from_series(s_abs * s_sign)
+
+    def total_seconds(self) -> PandasSeries:
+        s = self._series._series.dt.total_seconds()
+        s_sign = (
+            2 * (s > 0).astype(int) - 1
+        )  # this calculates the sign of each series element
+        s_abs = s.abs() // 1
+        if ~s.isna().any():
+            s_abs = s_abs.astype(int)
+        return self._series._from_series(s_abs * s_sign)
+
+    def total_milliseconds(self) -> PandasSeries:
+        s = self._series._series.dt.total_seconds() * 1e3
+        s_sign = (
+            2 * (s > 0).astype(int) - 1
+        )  # this calculates the sign of each series element
+        s_abs = s.abs() // 1
+        if ~s.isna().any():
+            s_abs = s_abs.astype(int)
+        return self._series._from_series(s_abs * s_sign)
+
+    def total_microseconds(self) -> PandasSeries:
+        s = self._series._series.dt.total_seconds() * 1e6
+        s_sign = (
+            2 * (s > 0).astype(int) - 1
+        )  # this calculates the sign of each series element
+        s_abs = s.abs() // 1
+        if ~s.isna().any():
+            s_abs = s_abs.astype(int)
+        return self._series._from_series(s_abs * s_sign)
+
+    def total_nanoseconds(self) -> PandasSeries:
+        s = self._series._series.dt.total_seconds() * 1e9
+        s_sign = (
+            2 * (s > 0).astype(int) - 1
+        )  # this calculates the sign of each series element
+        s_abs = s.abs() // 1
+        if ~s.isna().any():
+            s_abs = s_abs.astype(int)
+        return self._series._from_series(s_abs * s_sign)
diff --git a/narwhals/expression.py b/narwhals/expression.py
index 38629917..2af15af3 100644
--- a/narwhals/expression.py
+++ b/narwhals/expression.py
@@ -1487,6 +1487,201 @@ class ExprDateTimeNamespace:
         """
         return self._expr.__class__(lambda plx: self._expr._call(plx).dt.total_minutes())
 
+    def total_seconds(self) -> Expr:
+        """
+        Get total seconds.
+
+        Notes:
+            The function outputs the total seconds in the int dtype by default,
+            however, pandas may change the dtype to float when there are missing values,
+            consider using `fill_null()` and `cast` in this case.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> from datetime import timedelta
+            >>> import narwhals as nw
+            >>> data = {'a': [timedelta(seconds=10),
+            ...     timedelta(seconds=20, milliseconds=40)]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.with_columns(
+            ...       a_total_seconds = nw.col('a').dt.total_seconds()
+            ...     )
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+                                   a  a_total_seconds
+            0        0 days 00:00:10               10
+            1 0 days 00:00:20.040000               20
+            >>> func(df_pl)
+            shape: (2, 2)
+            ┌──────────────┬─────────────────┐
+            │ a            ┆ a_total_seconds │
+            │ ---          ┆ ---             │
+            │ duration[μs] ┆ i64             │
+            ╞══════════════╪═════════════════╡
+            │ 10s          ┆ 10              │
+            │ 20s 40ms     ┆ 20              │
+            └──────────────┴─────────────────┘
+        """
+        return self._expr.__class__(lambda plx: self._expr._call(plx).dt.total_seconds())
+
+    def total_milliseconds(self) -> Expr:
+        """
+        Get total milliseconds.
+
+        Notes:
+            The function outputs the total milliseconds in the int dtype by default,
+            however, pandas may change the dtype to float when there are missing values,
+            consider using `fill_null()` and `cast` in this case.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> from datetime import timedelta
+            >>> import narwhals as nw
+            >>> data = {'a': [timedelta(milliseconds=10),
+            ...     timedelta(milliseconds=20, microseconds=40)]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.with_columns(
+            ...       a_total_milliseconds = nw.col('a').dt.total_milliseconds()
+            ...     )
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+                                   a  a_total_milliseconds
+            0 0 days 00:00:00.010000                    10
+            1 0 days 00:00:00.020040                    20
+            >>> func(df_pl)
+            shape: (2, 2)
+            ┌──────────────┬──────────────────────┐
+            │ a            ┆ a_total_milliseconds │
+            │ ---          ┆ ---                  │
+            │ duration[μs] ┆ i64                  │
+            ╞══════════════╪══════════════════════╡
+            │ 10ms         ┆ 10                   │
+            │ 20040µs      ┆ 20                   │
+            └──────────────┴──────────────────────┘
+        """
+        return self._expr.__class__(
+            lambda plx: self._expr._call(plx).dt.total_milliseconds()
+        )
+
+    def total_microseconds(self) -> Expr:
+        """
+        Get total microseconds.
+
+        Notes:
+            The function outputs the total microseconds in the int dtype by default,
+            however, pandas may change the dtype to float when there are missing values,
+            consider using `fill_null()` and `cast` in this case.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> from datetime import timedelta
+            >>> import narwhals as nw
+            >>> data = {'a': [timedelta(microseconds=10),
+            ...     timedelta(milliseconds=1, microseconds=200)]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.with_columns(
+            ...       a_total_microseconds = nw.col('a').dt.total_microseconds()
+            ...     )
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+                                   a  a_total_microseconds
+            0 0 days 00:00:00.000010                    10
+            1 0 days 00:00:00.001200                  1200
+            >>> func(df_pl)
+            shape: (2, 2)
+            ┌──────────────┬──────────────────────┐
+            │ a            ┆ a_total_microseconds │
+            │ ---          ┆ ---                  │
+            │ duration[μs] ┆ i64                  │
+            ╞══════════════╪══════════════════════╡
+            │ 10µs         ┆ 10                   │
+            │ 1200µs       ┆ 1200                 │
+            └──────────────┴──────────────────────┘
+        """
+        return self._expr.__class__(
+            lambda plx: self._expr._call(plx).dt.total_microseconds()
+        )
+
+    def total_nanoseconds(self) -> Expr:
+        """
+        Get total nanoseconds.
+
+        Notes:
+            The function outputs the total nanoseconds in the int dtype by default,
+            however, pandas may change the dtype to float when there are missing values,
+            consider using `fill_null()` and `cast` in this case.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> from datetime import timedelta
+            >>> import narwhals as nw
+            >>> data = ['2024-01-01 00:00:00.000000001',
+            ...     '2024-01-01 00:00:00.000000002']
+            >>> df_pd = pd.DataFrame({'a': pd.to_datetime(data)})
+            >>> df_pl = (pl.DataFrame({'a': data}).with_columns(
+            ...     pl.col('a').str.to_datetime(time_unit='ns')))
+
+            We define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.with_columns(
+            ...       a_diff_total_nanoseconds = nw.col('a').diff().dt.total_nanoseconds()
+            ...     )
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+                                          a  a_diff_total_nanoseconds
+            0 2024-01-01 00:00:00.000000001                       NaN
+            1 2024-01-01 00:00:00.000000002                       1.0
+            >>> func(df_pl)
+            shape: (2, 2)
+            ┌───────────────────────────────┬──────────────────────────┐
+            │ a                             ┆ a_diff_total_nanoseconds │
+            │ ---                           ┆ ---                      │
+            │ datetime[ns]                  ┆ i64                      │
+            ╞═══════════════════════════════╪══════════════════════════╡
+            │ 2024-01-01 00:00:00.000000001 ┆ null                     │
+            │ 2024-01-01 00:00:00.000000002 ┆ 1                        │
+            └───────────────────────────────┴──────────────────────────┘
+        """
+        return self._expr.__class__(
+            lambda plx: self._expr._call(plx).dt.total_nanoseconds()
+        )
+
 
 def col(*names: str | Iterable[str]) -> Expr:
     """
diff --git a/narwhals/series.py b/narwhals/series.py
index 6a316bf2..1482b766 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -1031,3 +1031,170 @@ class SeriesDateTimeNamespace:
             ]
         """
         return self._series.__class__(self._series._series.dt.total_minutes())
+
+    def total_seconds(self) -> Series:
+        """
+        Get total seconds.
+
+        Notes:
+            The function outputs the total seconds in the int dtype by default,
+            however, pandas may change the dtype to float when there are missing values,
+            consider using `fill_null()` in this case.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> from datetime import timedelta
+            >>> import narwhals as nw
+            >>> data = [timedelta(seconds=10), timedelta(seconds=20, milliseconds=40)]
+            >>> s_pd = pd.Series(data)
+            >>> s_pl = pl.Series(data)
+
+            We define a library agnostic function:
+
+            >>> def func(s_any):
+            ...     s = nw.from_native(s_any, series_only=True)
+            ...     s = s.dt.total_seconds()
+            ...     return nw.to_native(s)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0    10
+            1    20
+            dtype: int...
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (2,)
+            Series: '' [i64]
+            [
+                    10
+                    20
+            ]
+        """
+        return self._series.__class__(self._series._series.dt.total_seconds())
+
+    def total_milliseconds(self) -> Series:
+        """
+        Get total milliseconds.
+
+        Notes:
+            The function outputs the total milliseconds in the int dtype by default,
+            however, pandas may change the dtype to float when there are missing values,
+            consider using `fill_null()` in this case.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> from datetime import timedelta
+            >>> import narwhals as nw
+            >>> data = [timedelta(milliseconds=10),
+            ...     timedelta(milliseconds=20, microseconds=40)]
+            >>> s_pd = pd.Series(data)
+            >>> s_pl = pl.Series(data)
+
+            We define a library agnostic function:
+
+            >>> def func(s_any):
+            ...     s = nw.from_native(s_any, series_only=True)
+            ...     s = s.dt.total_milliseconds()
+            ...     return nw.to_native(s)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0    10
+            1    20
+            dtype: int...
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (2,)
+            Series: '' [i64]
+            [
+                    10
+                    20
+            ]
+        """
+        return self._series.__class__(self._series._series.dt.total_milliseconds())
+
+    def total_microseconds(self) -> Series:
+        """
+        Get total microseconds.
+
+        Notes:
+            The function outputs the total microseconds in the int dtype by default,
+            however, pandas may change the dtype to float when there are missing values,
+            consider using `fill_null()` in this case.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> from datetime import timedelta
+            >>> import narwhals as nw
+            >>> data = [timedelta(microseconds=10),
+            ...     timedelta(milliseconds=1, microseconds=200)]
+            >>> s_pd = pd.Series(data)
+            >>> s_pl = pl.Series(data)
+
+            We define a library agnostic function:
+
+            >>> def func(s_any):
+            ...     s = nw.from_native(s_any, series_only=True)
+            ...     s = s.dt.total_microseconds()
+            ...     return nw.to_native(s)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0      10
+            1    1200
+            dtype: int...
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (2,)
+            Series: '' [i64]
+            [
+                    10
+                    1200
+            ]
+        """
+        return self._series.__class__(self._series._series.dt.total_microseconds())
+
+    def total_nanoseconds(self) -> Series:
+        """
+        Get total nanoseconds.
+
+        Notes:
+            The function outputs the total nanoseconds in the int dtype by default,
+            however, pandas may change the dtype to float when there are missing values,
+            consider using `fill_null()` in this case.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> from datetime import timedelta
+            >>> import narwhals as nw
+            >>> data = ['2024-01-01 00:00:00.000000001',
+            ...     '2024-01-01 00:00:00.000000002']
+            >>> s_pd = pd.to_datetime(pd.Series(data))
+            >>> s_pl = pl.Series(data).str.to_datetime(time_unit='ns')
+
+            We define a library agnostic function:
+
+            >>> def func(s_any):
+            ...     s = nw.from_native(s_any, series_only=True)
+            ...     s = s.diff().dt.total_nanoseconds()
+            ...     return nw.to_native(s)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0    NaN
+            1    1.0
+            dtype: float64
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (2,)
+            Series: '' [i64]
+            [
+                    null
+                    1
+            ]
+        """
+        return self._series.__class__(self._series._series.dt.total_nanoseconds())


Test patch:
diff --git a/tests/test_dt.py b/tests/test_dt.py
index c996588e..5e727606 100644
--- a/tests/test_dt.py
+++ b/tests/test_dt.py
@@ -5,6 +5,7 @@ from datetime import timedelta
 from typing import Any
 
 import hypothesis.strategies as st
+import numpy as np
 import pandas as pd
 import polars as pl
 import pytest
@@ -20,7 +21,17 @@ data = {
         datetime(2020, 1, 2, 2, 4, 14),
     ],
 }
-data_timedelta = {"a": [None, timedelta(days=1, minutes=10, seconds=40)]}
+data_timedelta = {
+    "a": [
+        None,
+        timedelta(minutes=1, seconds=1, milliseconds=1, microseconds=1),
+    ],
+    "b": [
+        timedelta(milliseconds=2),
+        timedelta(milliseconds=1, microseconds=300),
+    ],
+    "c": np.array([None, 20], dtype="timedelta64[ns]"),
+}
 
 
 @pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
@@ -48,19 +59,53 @@ def test_datetime_attributes(
 
 @pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
 @pytest.mark.parametrize(
-    ("attribute", "expected"),
+    ("attribute", "expected_a", "expected_b"),
     [
-        ("total_minutes", [0, 1450]),
+        ("total_minutes", [0, 1], [0, 0]),
+        ("total_seconds", [0, 61], [0, 0]),
+        ("total_milliseconds", [0, 61001], [2, 1]),
     ],
 )
 def test_duration_attributes(
-    attribute: str, expected: list[int], constructor: Any
+    attribute: str,
+    expected_a: list[int],
+    expected_b: list[int],
+    constructor: Any,
 ) -> None:
     df = nw.from_native(constructor(data_timedelta), eager_only=True)
-    result = nw.to_native(df.select(getattr(nw.col("a").dt, attribute)().fill_null(0)))
-    compare_dicts(result, {"a": expected})
-    result = nw.to_native(df.select(getattr(df["a"].dt, attribute)().fill_null(0)))
-    compare_dicts(result, {"a": expected})
+    result_a = nw.to_native(df.select(getattr(nw.col("a").dt, attribute)().fill_null(0)))
+    compare_dicts(result_a, {"a": expected_a})
+    result_a = nw.to_native(df.select(getattr(df["a"].dt, attribute)().fill_null(0)))
+    compare_dicts(result_a, {"a": expected_a})
+    result_b = nw.to_native(df.select(getattr(nw.col("b").dt, attribute)().fill_null(0)))
+    compare_dicts(result_b, {"b": expected_b})
+    result_b = nw.to_native(df.select(getattr(df["b"].dt, attribute)().fill_null(0)))
+    compare_dicts(result_b, {"b": expected_b})
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize(
+    ("attribute", "expected_b", "expected_c"),
+    [
+        ("total_microseconds", [2000, 1300], [0, 0]),
+        ("total_nanoseconds", [2000000, 1300000], [0, 20]),
+    ],
+)
+def test_duration_micro_nano(
+    attribute: str,
+    expected_b: list[int],
+    expected_c: list[int],
+    constructor: Any,
+) -> None:
+    df = nw.from_native(constructor(data_timedelta), eager_only=True)
+    result_b = nw.to_native(df.select(getattr(nw.col("b").dt, attribute)().fill_null(0)))
+    compare_dicts(result_b, {"b": expected_b})
+    result_b = nw.to_native(df.select(getattr(df["b"].dt, attribute)().fill_null(0)))
+    compare_dicts(result_b, {"b": expected_b})
+    result_c = nw.to_native(df.select(getattr(nw.col("c").dt, attribute)().fill_null(0)))
+    compare_dicts(result_c, {"c": expected_c})
+    result_c = nw.to_native(df.select(getattr(df["c"].dt, attribute)().fill_null(0)))
+    compare_dicts(result_c, {"c": expected_c})
 
 
 @given(dates=st.datetimes(min_value=datetime(1960, 1, 1), max_value=datetime(1980, 1, 1)))  # type: ignore[misc]