Skip to content

Commit

Permalink
WIP: raise on str and momentaneously xfail on polars expr
Browse files Browse the repository at this point in the history
  • Loading branch information
AlessandroMiola committed Oct 30, 2024
1 parent a6b4089 commit 6da817c
Show file tree
Hide file tree
Showing 7 changed files with 70 additions and 6 deletions.
6 changes: 6 additions & 0 deletions narwhals/_arrow/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,12 @@ def mean(self) -> int:
def median(self) -> int:
import pyarrow.compute as pc # ignore-banned-import()

from narwhals._exceptions import InvalidOperationError

if not self.dtype.is_numeric():
msg = "`median` operation not supported for non-numeric input type."
raise InvalidOperationError(msg)

return pc.approximate_median(self._native_series) # type: ignore[no-any-return]

def min(self) -> int:
Expand Down
16 changes: 11 additions & 5 deletions narwhals/_dask/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -382,11 +382,17 @@ def mean(self) -> Self:
)

def median(self) -> Self:
return self._from_call(
lambda _input: _input.median_approximate(),
"median",
returns_scalar=True,
)
from dask_expr._shuffle import _is_numeric_cast_type

from narwhals._exceptions import InvalidOperationError

def func(_input: dask_expr.Series) -> dask_expr.Series:
if not _is_numeric_cast_type(_input.dtype):
msg = "`median` operation not supported for non-numeric input type."
raise InvalidOperationError(msg)
return _input.median_approximate()

return self._from_call(func, "median", returns_scalar=True)

def min(self) -> Self:
return self._from_call(
Expand Down
3 changes: 3 additions & 0 deletions narwhals/_exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,3 +2,6 @@


class ColumnNotFoundError(Exception): ...


class InvalidOperationError(Exception): ...
5 changes: 5 additions & 0 deletions narwhals/_pandas_like/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -422,6 +422,11 @@ def mean(self) -> Any:
return ser.mean()

def median(self) -> Any:
from narwhals._exceptions import InvalidOperationError

if not self.dtype.is_numeric():
msg = "`median` operation not supported for non-numeric input type."
raise InvalidOperationError(msg)
ser = self._native_series
return ser.median()

Expand Down
7 changes: 6 additions & 1 deletion narwhals/_polars/namespace.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,9 +120,14 @@ def mean_horizontal(self, *exprs: IntoPolarsExpr) -> PolarsExpr:
def median(self, *column_names: str) -> PolarsExpr:
import polars as pl # ignore-banned-import()

from narwhals._exceptions import InvalidOperationError
from narwhals._polars.expr import PolarsExpr

return PolarsExpr(pl.median([*column_names]), dtypes=self._dtypes) # type: ignore[arg-type]
try:
return PolarsExpr(pl.median([*column_names]), dtypes=self._dtypes) # type: ignore[arg-type]
except pl.exceptions.InvalidOperationError as e:
msg = "`median` operation not supported for non-numeric input type."
raise InvalidOperationError(msg) from e

def concat_str(
self,
Expand Down
9 changes: 9 additions & 0 deletions narwhals/_polars/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,15 @@ def __rpow__(self, other: PolarsSeries | Any) -> Self:
def __invert__(self) -> Self:
return self._from_native_series(self._native_series.__invert__())

def median(self) -> Any:
from narwhals._exceptions import InvalidOperationError

if not self.dtype.is_numeric():
msg = "`median` operation not supported for non-numeric input type."
raise InvalidOperationError(msg)

return self._native_series.median()

def to_dummies(
self: Self, *, separator: str = "_", drop_first: bool = False
) -> PolarsDataFrame:
Expand Down
30 changes: 30 additions & 0 deletions tests/expr_and_series/median_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import pytest

import narwhals.stable.v1 as nw
from narwhals._exceptions import InvalidOperationError
from tests.utils import Constructor
from tests.utils import ConstructorEager
from tests.utils import assert_equal_data
Expand Down Expand Up @@ -36,3 +37,32 @@ def test_median_series(
series = nw.from_native(constructor_eager(data), eager_only=True)[col]
result = series.median()
assert_equal_data({col: [result]}, {col: [expected]})


@pytest.mark.parametrize("expr", [nw.col("s").median(), nw.median("s")])
def test_median_expr_raises_on_str(
constructor: Constructor,
expr: nw.Expr,
request: pytest.FixtureRequest,
) -> None:
if "polars" in str(constructor):
request.applymarker(pytest.mark.xfail)

df = nw.from_native(constructor(data))
with pytest.raises(
InvalidOperationError,
match="`median` operation not supported for non-numeric input type.",
):
df.select(expr)


@pytest.mark.parametrize(("col"), [("s")])
def test_median_series_raises_on_str(
constructor_eager: ConstructorEager, col: str
) -> None:
series = nw.from_native(constructor_eager(data), eager_only=True)[col]
with pytest.raises(
InvalidOperationError,
match="`median` operation not supported for non-numeric input type.",
):
series.median()

0 comments on commit 6da817c

Please sign in to comment.