feat: Add {Expr,Series}.cos (#3392)

liamholmes31 · FBruzzesi · web-flow · commit d7bb050e91a6 · 2026-01-06T15:17:34.000+01:00
---------

Co-authored-by: Francesco Bruzzesi &lt;42817048+FBruzzesi@users.noreply.github.com&gt;
diff --git a/docs/api-reference/expr.md b/docs/api-reference/expr.md
@@ -12,6 +12,7 @@
         - cast
         - ceil
         - clip
+        - cos
         - count
         - cum_count
         - cum_max
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
@@ -18,6 +18,7 @@
         - cast
         - ceil
         - clip
+        - cos
         - count
         - cum_count
         - cum_max
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
@@ -1107,6 +1107,9 @@ def sqrt(self) -> Self:
     def sin(self) -> Self:
         return self._with_native(pc.sin(self.native))
 
+    def cos(self) -> Self:
+        return self._with_native(pc.cos(self.native))
+
     def any_value(
         self, *, ignore_nulls: bool, _return_py_scalar: bool = True
     ) -> PythonLiteral:
diff --git a/narwhals/_compliant/column.py b/narwhals/_compliant/column.py
@@ -135,6 +135,7 @@ def round(self, decimals: int) -> Self: ...
     def floor(self) -> Self: ...
     def ceil(self) -> Self: ...
     def shift(self, n: int) -> Self: ...
+    def cos(self) -> Self: ...
     def sin(self) -> Self: ...
     def unique(self) -> Self: ...
 
diff --git a/narwhals/_compliant/expr.py b/narwhals/_compliant/expr.py
@@ -812,6 +812,9 @@ def exp(self) -> Self:
     def sin(self) -> Self:
         return self._reuse_series("sin")
 
+    def cos(self) -> Self:
+        return self._reuse_series("cos")
+
     def sqrt(self) -> Self:
         return self._reuse_series("sqrt")
 
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
@@ -642,6 +642,11 @@ def sin(self) -> Self:
 
         return self._with_callable(da.sin)
 
+    def cos(self) -> Self:
+        import dask.array as da
+
+        return self._with_callable(da.cos)
+
     def sqrt(self) -> Self:
         import dask.array as da
 
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
@@ -1116,6 +1116,21 @@ def sin(self) -> Self:
 
         return self._with_native(result_native)
 
+    def cos(self) -> Self:
+        native = self.native
+        if self.is_native_dtype_pyarrow(native.dtype):
+            import pyarrow.compute as pc
+
+            result_native = self._apply_pyarrow_compute_func(
+                native,
+                pc.cos,  # type: ignore[arg-type]
+            )
+        else:
+            array_func = self._array_funcs.cos
+            result_native = self._apply_array_func(native, array_func)
+
+        return self._with_native(result_native)
+
     def is_native_dtype_pyarrow(self, native_dtype: Any) -> bool:
         impl = self._implementation
         return get_dtype_backend(native_dtype, implementation=impl) == "pyarrow"
diff --git a/narwhals/_polars/expr.py b/narwhals/_polars/expr.py
@@ -336,6 +336,7 @@ def struct(self) -> PolarsExprStructNamespace:
     arg_true: Method[Self]
     ceil: Method[Self]
     count: Method[Self]
+    cos: Method[Self]
     cum_max: Method[Self]
     cum_min: Method[Self]
     cum_prod: Method[Self]
diff --git a/narwhals/_polars/series.py b/narwhals/_polars/series.py
@@ -83,6 +83,7 @@
         "arg_true",
         "ceil",
         "clip",
+        "cos",
         "count",
         "cum_max",
         "cum_min",
@@ -710,6 +711,7 @@ def struct(self) -> PolarsSeriesStructNamespace:
     arg_true: Method[Self]
     ceil: Method[Self]
     count: Method[int]
+    cos: Method[Self]
     cum_max: Method[Self]
     cum_min: Method[Self]
     cum_prod: Method[Self]
diff --git a/narwhals/_sql/expr.py b/narwhals/_sql/expr.py
@@ -584,6 +584,9 @@ def ceil(self) -> Self:
     def sin(self) -> Self:
         return self._with_elementwise(lambda expr: self._function("sin", expr))
 
+    def cos(self) -> Self:
+        return self._with_elementwise(lambda expr: self._function("cos", expr))
+
     def sqrt(self) -> Self:
         def _sqrt(expr: NativeExprT) -> NativeExprT:
             return self._when(
diff --git a/narwhals/expr.py b/narwhals/expr.py
@@ -2253,6 +2253,32 @@ def sin(self) -> Self:
         """
         return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "sin"))
 
+    def cos(self) -> Self:
+        r"""Compute the element-wise value for the cosine.
+
+        Examples:
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> from math import pi
+            >>> df_native = pa.table({"values": [0, pi / 2, pi]})
+            >>> df = nw.from_native(df_native)
+            >>> result = df.with_columns(cos=nw.col("values").cos()).select(
+            ...     nw.all().round(4)
+            ... )
+            >>> result
+            ┌───────────────────────────┐
+            |    Narwhals DataFrame     |
+            |---------------------------|
+            |pyarrow.Table              |
+            |values: double             |
+            |cos: double                |
+            |----                       |
+            |values: [[0,1.5708,3.1416]]|
+            |cos: [[1,0,-1]]            |
+            └───────────────────────────┘
+        """
+        return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "cos"))
+
     def sqrt(self) -> Self:
         r"""Compute the square root of the elements.
 
diff --git a/narwhals/series.py b/narwhals/series.py
@@ -2758,7 +2758,7 @@ def exp(self) -> Self:
         return self._with_compliant(self._compliant_series.exp())
 
     def sin(self) -> Self:
-        r"""Compute the sin.
+        r"""Compute the element-wise value for the sine.
 
         Examples:
             >>> import pandas as pd
@@ -2778,6 +2778,27 @@ def sin(self) -> Self:
         """
         return self._with_compliant(self._compliant_series.sin())
 
+    def cos(self) -> Self:
+        r"""Compute the element-wise value for the cosine.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> from math import pi
+            >>> s_native = pd.Series([0, pi / 2, 3 * pi / 2], name="a")
+            >>> s = nw.from_native(s_native, series_only=True)
+            >>> s.cos().round(4)
+            ┌───────────────────────┐
+            |    Narwhals Series    |
+            |-----------------------|
+            |0    1.0               |
+            |1    0.0               |
+            |2   -0.0               |
+            |Name: a, dtype: float64|
+            └───────────────────────┘
+        """
+        return self._with_compliant(self._compliant_series.cos())
+
     def sqrt(self) -> Self:
         r"""Compute the square root.
 
diff --git a/tests/expr_and_series/cos_test.py b/tests/expr_and_series/cos_test.py
@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from math import pi
+from typing import TYPE_CHECKING
+
+import pytest
+
+import narwhals as nw
+from tests.utils import (
+    PANDAS_VERSION,
+    PYARROW_VERSION,
+    Constructor,
+    ConstructorEager,
+    assert_equal_data,
+)
+
+if TYPE_CHECKING:
+    from narwhals.typing import DTypeBackend
+
+data = {"a": [-pi, -pi / 2, 0.0, pi / 2, pi]}
+
+expected = [-1, 0, 1, 0, -1]
+
+
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+def test_cos_expr(constructor: Constructor) -> None:
+    df = nw.from_native(constructor(data))
+    result = df.select(nw.col("a").cos())
+    assert_equal_data(result, {"a": expected})
+
+
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+def test_cos_series(constructor_eager: ConstructorEager) -> None:
+    series = nw.from_native(constructor_eager(data), eager_only=True)["a"]
+    result = series.cos()
+    assert_equal_data({"a": result}, {"a": expected})
+
+
+PYARROW_UNAVAILABLE = PYARROW_VERSION == (0, 0, 0)
+reason = "nullable types require pandas2+"
+require_pd_2_1 = pytest.mark.skipif(
+    PANDAS_VERSION < (2, 1, 0) or PYARROW_UNAVAILABLE, reason=reason
+)
+require_pd_2_0 = pytest.mark.skipif(PANDAS_VERSION < (2, 0, 0), reason=reason)
+
+
+@pytest.mark.parametrize(
+    "dtype_backend",
+    [
+        pytest.param("pyarrow", marks=require_pd_2_1),
+        pytest.param("numpy_nullable", marks=require_pd_2_0),
+        None,
+    ],
+)
+def test_cos_dtype_pandas(dtype_backend: DTypeBackend) -> None:
+    pytest.importorskip("pandas")
+    import pandas as pd
+
+    s = pd.Series([-pi / 2, None, pi / 2], name="a", dtype="float32", index=[8, 7, 6])
+    if dtype_backend:
+        s = s.convert_dtypes(dtype_backend=dtype_backend)
+    result = nw.from_native(s, series_only=True).cos().to_native()
+    expected = pd.Series([0, None, 0], name="a", dtype=s.dtype, index=[8, 7, 6])
+    pd.testing.assert_series_equal(result, expected, rtol=0, atol=1e-6)
+
+
+@pytest.mark.skipif(PANDAS_VERSION < (2, 1, 0), reason="nullable types require pandas2+")
+def test_cos_dtype_pandas_pyarrow() -> None:
+    pytest.importorskip("pandas")
+    pytest.importorskip("pyarrow")
+    import pandas as pd
+
+    s = pd.Series(
+        [-pi / 2, None, pi / 2], name="a", dtype="Float32[pyarrow]", index=[8, 7, 6]
+    )
+    result = nw.from_native(s, series_only=True).cos().to_native()
+    expected = pd.Series(
+        [0, None, 0], name="a", dtype="Float32[pyarrow]", index=[8, 7, 6]
+    )
+    pd.testing.assert_series_equal(result, expected, rtol=0, atol=1e-6)

-Original file line number
+Diff line change
         - cast
         - ceil
         - clip
 +        - cos
         - count
         - cum_count
         - cum_max