Repository: narwhals-dev/narwhals
Pull request #360 (2024-06-30 13:48:43)

feat: add duration dtype
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/docs/api-reference/dtypes.md b/docs/api-reference/dtypes.md
index d84b0684..c89b7f49 100644
--- a/docs/api-reference/dtypes.md
+++ b/docs/api-reference/dtypes.md
@@ -18,6 +18,7 @@
         - Categorical
         - String
         - Datetime
+        - Duration
         - Object
       show_root_heading: false
       show_source: false
diff --git a/narwhals/__init__.py b/narwhals/__init__.py
index 3acbccac..56b99f1f 100644
--- a/narwhals/__init__.py
+++ b/narwhals/__init__.py
@@ -5,6 +5,7 @@ from narwhals.dtypes import Boolean
 from narwhals.dtypes import Categorical
 from narwhals.dtypes import Date
 from narwhals.dtypes import Datetime
+from narwhals.dtypes import Duration
 from narwhals.dtypes import Float32
 from narwhals.dtypes import Float64
 from narwhals.dtypes import Int8
@@ -77,6 +78,7 @@ __all__ = [
     "Categorical",
     "String",
     "Datetime",
+    "Duration",
     "Date",
     "narwhalify",
     "show_versions",
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
index 42ca025f..31af558e 100644
--- a/narwhals/_arrow/namespace.py
+++ b/narwhals/_arrow/namespace.py
@@ -31,6 +31,7 @@ class ArrowNamespace:
     Categorical = dtypes.Categorical
     String = dtypes.String
     Datetime = dtypes.Datetime
+    Duration = dtypes.Duration
     Date = dtypes.Date
 
     def _create_expr_from_callable(  # noqa: PLR0913
diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py
index 14aa8451..a7bc303d 100644
--- a/narwhals/_arrow/utils.py
+++ b/narwhals/_arrow/utils.py
@@ -38,6 +38,8 @@ def translate_dtype(dtype: Any) -> dtypes.DType:
         return dtypes.Date()
     if pa.types.is_timestamp(dtype):
         return dtypes.Datetime()
+    if pa.types.is_duration(dtype):
+        return dtypes.Duration()
     if pa.types.is_dictionary(dtype):
         return dtypes.Categorical()
     raise AssertionError
diff --git a/narwhals/_pandas_like/namespace.py b/narwhals/_pandas_like/namespace.py
index a1150ba9..c8e4f101 100644
--- a/narwhals/_pandas_like/namespace.py
+++ b/narwhals/_pandas_like/namespace.py
@@ -36,6 +36,7 @@ class PandasNamespace:
     Categorical = dtypes.Categorical
     String = dtypes.String
     Datetime = dtypes.Datetime
+    Duration = dtypes.Duration
     Date = dtypes.Date
 
     @property
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index 2db1ed2b..e1785d26 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -411,6 +411,9 @@ def translate_dtype(column: Any) -> DType:
     if str(dtype).startswith("datetime64"):
         # todo: different time units and time zones
         return dtypes.Datetime()
+    if str(dtype).startswith("timedelta64") or str(dtype).startswith("duration"):
+        # todo: different time units
+        return dtypes.Duration()
     if str(dtype).startswith("timestamp["):
         # pyarrow-backed datetime
         # todo: different time units and time zones
@@ -548,6 +551,11 @@ def reverse_translate_dtype(  # noqa: PLR0915
         if dtype_backend == "pyarrow-nullable":
             return "timestamp[ns][pyarrow]"
         return "datetime64[ns]"
+    if isinstance_or_issubclass(dtype, dtypes.Duration):
+        # todo: different time units and time zones
+        if dtype_backend == "pyarrow-nullable":
+            return "duration[ns][pyarrow]"
+        return "timedelta64[ns]"
     if isinstance_or_issubclass(dtype, dtypes.Date):
         if dtype_backend == "pyarrow-nullable":
             return "date32[pyarrow]"
diff --git a/narwhals/dtypes.py b/narwhals/dtypes.py
index 7c3f4869..aa98f2c5 100644
--- a/narwhals/dtypes.py
+++ b/narwhals/dtypes.py
@@ -70,6 +70,9 @@ class Object(DType): ...
 class Datetime(TemporalType): ...
 
 
+class Duration(TemporalType): ...
+
+
 class Categorical(DType): ...
 
 
@@ -113,6 +116,8 @@ def translate_dtype(plx: Any, dtype: DType) -> Any:
         return plx.Categorical
     if dtype == Datetime:
         return plx.Datetime
+    if dtype == Duration:
+        return plx.Duration
     if dtype == Date:
         return plx.Date
     msg = f"Unknown dtype: {dtype}"  # pragma: no cover
@@ -154,6 +159,8 @@ def to_narwhals_dtype(dtype: Any, *, is_polars: bool) -> DType:
         return Categorical()
     if dtype == pl.Datetime:
         return Datetime()
+    if dtype == pl.Duration:
+        return Duration()
     if dtype == pl.Date:
         return Date()
     msg = f"Unexpected dtype, got: {type(dtype)}"  # pragma: no cover


Test patch:
diff --git a/tests/frame/schema_test.py b/tests/frame/schema_test.py
index bf77b7e6..288db181 100644
--- a/tests/frame/schema_test.py
+++ b/tests/frame/schema_test.py
@@ -1,5 +1,6 @@
 from datetime import date
 from datetime import datetime
+from datetime import timedelta
 from datetime import timezone
 from typing import Any
 
@@ -65,6 +66,7 @@ def test_dtypes() -> None:
             "n": [date(2020, 1, 1)],
             "o": [datetime(2020, 1, 1)],
             "p": ["a"],
+            "q": [timedelta(1)],
         },
         schema={
             "a": pl.Int64,
@@ -83,6 +85,7 @@ def test_dtypes() -> None:
             "n": pl.Date,
             "o": pl.Datetime,
             "p": pl.Categorical,
+            "q": pl.Duration,
         },
     )
     df = nw.DataFrame(df_pl)
@@ -104,6 +107,7 @@ def test_dtypes() -> None:
         "n": nw.Date,
         "o": nw.Datetime,
         "p": nw.Categorical,
+        "q": nw.Duration,
     }
     assert result == expected
     assert {name: df[name].dtype for name in df.columns} == expected
diff --git a/tests/series/test_common.py b/tests/series/test_common.py
index 729cb465..48341f94 100644
--- a/tests/series/test_common.py
+++ b/tests/series/test_common.py
@@ -164,6 +164,7 @@ def test_cast() -> None:
             "m": [True],
             "n": [True],
             "o": ["a"],
+            "p": [1],
         },
         schema={
             "a": pl.Int64,
@@ -181,6 +182,7 @@ def test_cast() -> None:
             "m": pl.Boolean,
             "n": pl.Boolean,
             "o": pl.Categorical,
+            "p": pl.Int64,
         },
     )
     df = nw.DataFrame(df_raw).select(
@@ -199,6 +201,7 @@ def test_cast() -> None:
         nw.col("m").cast(nw.Int8),
         nw.col("n").cast(nw.Int8),
         nw.col("o").cast(nw.String),
+        nw.col("p").cast(nw.Duration),
     )
     result = df.schema
     expected = {
@@ -217,6 +220,7 @@ def test_cast() -> None:
         "m": nw.Int8,
         "n": nw.Int8,
         "o": nw.String,
+        "p": nw.Duration,
     }
     assert result == expected
     result_pd = nw.DataFrame(df.to_pandas()).schema
@@ -237,6 +241,7 @@ def test_cast() -> None:
         df["m"].cast(nw.Int8),
         df["n"].cast(nw.Boolean),
         df["o"].cast(nw.Categorical),
+        df["p"].cast(nw.Duration),
     ).schema
     expected = {
         "a": nw.Int32,
@@ -254,6 +259,7 @@ def test_cast() -> None:
         "m": nw.Int8,
         "n": nw.Boolean,
         "o": nw.Categorical,
+        "p": nw.Duration,
     }
     df = nw.from_native(df.to_pandas())  # type: ignore[assignment]
     result_pd = df.select(
@@ -272,6 +278,7 @@ def test_cast() -> None:
         df["m"].cast(nw.Int8),
         df["n"].cast(nw.Boolean),
         df["o"].cast(nw.Categorical),
+        df["p"].cast(nw.Duration),
     ).schema
     assert result == expected
     df = nw.from_native(df.to_pandas().convert_dtypes())  # type: ignore[assignment]
@@ -291,6 +298,7 @@ def test_cast() -> None:
         df["m"].cast(nw.Int8),
         df["n"].cast(nw.Boolean),
         df["o"].cast(nw.Categorical),
+        df["p"].cast(nw.Duration),
     ).schema
     assert result == expected
     if parse_version(pd.__version__) < parse_version("2.0.0"):  # pragma: no cover
@@ -312,6 +320,7 @@ def test_cast() -> None:
         df["m"].cast(nw.Int8),
         df["n"].cast(nw.Boolean),
         df["o"].cast(nw.Categorical),
+        df["p"].cast(nw.Duration),
     ).schema
     assert result == expected