Skip to content
Merged
Show file tree
Hide file tree
Changes from 10 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/api-reference/expr.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
- cast
- ceil
- clip
- cos
- count
- cum_count
- cum_max
Expand Down
1 change: 1 addition & 0 deletions docs/api-reference/series.md
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
- cast
- ceil
- clip
- cos
- count
- cum_count
- cum_max
Expand Down
3 changes: 3 additions & 0 deletions narwhals/_arrow/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
1 change: 1 addition & 0 deletions narwhals/_compliant/column.py
Original file line number Diff line number Diff line change
Expand Up @@ -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: ...

Expand Down
3 changes: 3 additions & 0 deletions narwhals/_compliant/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
5 changes: 5 additions & 0 deletions narwhals/_dask/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
15 changes: 15 additions & 0 deletions narwhals/_pandas_like/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -1105,6 +1105,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"
Expand Down
1 change: 1 addition & 0 deletions narwhals/_polars/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
2 changes: 2 additions & 0 deletions narwhals/_polars/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@
"arg_true",
"ceil",
"clip",
"cos",
"count",
"cum_max",
"cum_min",
Expand Down Expand Up @@ -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]
Expand Down
3 changes: 3 additions & 0 deletions narwhals/_sql/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
24 changes: 24 additions & 0 deletions narwhals/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -2253,6 +2253,30 @@ 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())
>>> result
β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
| Narwhals DataFrame |
|--------------------------------------------------|
|pyarrow.Table |
|values: double |
|cos: double |
|---- |
|values: [[0,1.5707963267948966,3.141592653589793]]|
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

open to suggestions on making this output a bit less off-putting!

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe something like:

result = df.with_columns(cos=nw.col("values").cos()).select(nw.all().round(4))

|cos: [[1,6.123233995736766e-17,-1]] |
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
"""
return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "cos"))

def sqrt(self) -> Self:
r"""Compute the square root of the elements.

Expand Down
23 changes: 22 additions & 1 deletion narwhals/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()
β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
| Narwhals Series |
|-----------------------|
|0 0.0 |
|1 1.0 |
|2 -1.0 |
|Name: a, dtype: float64|
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
"""
return self._with_compliant(self._compliant_series.cos())

def sqrt(self) -> Self:
r"""Compute the square root.

Expand Down
80 changes: 80 additions & 0 deletions tests/expr_and_series/cos_test.py
Original file line number Diff line number Diff line change
@@ -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:
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wasn't actually certain if this test was necessary given the above one?

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)
Loading