Skip to content

Commit

Permalink
fix(python): address unexpected expression name on use of __neg__ o…
Browse files Browse the repository at this point in the history
…r `__pos__` operators
  • Loading branch information
alexander-beedie committed Sep 17, 2023
1 parent b08bc7d commit 04c76b6
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 8 deletions.
10 changes: 8 additions & 2 deletions py-polars/polars/expr/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,10 @@ def __ne__(self, other: Any) -> Self: # type: ignore[override]
return self._from_pyexpr(self._pyexpr.neq(self._to_expr(other)._pyexpr))

def __neg__(self) -> Expr:
return F.lit(0) - self
neg_expr = F.lit(0) - self
if (name := self.meta.output_name(raise_if_undetermined=False)) is not None:
neg_expr = neg_expr.alias(name)
return neg_expr

def __or__(self, other: Expr | int | bool) -> Self:
return self._from_pyexpr(self._pyexpr._or(self._to_pyexpr(other)))
Expand All @@ -204,7 +207,10 @@ def __ror__(self, other: Any) -> Self:
return self._from_pyexpr(self._to_pyexpr(other)._or(self._pyexpr))

def __pos__(self) -> Expr:
return F.lit(0) + self
pos_expr = F.lit(0) + self
if (name := self.meta.output_name(raise_if_undetermined=False)) is not None:
pos_expr = pos_expr.alias(name)
return pos_expr

def __pow__(self, power: int | float | Series | Expr) -> Self:
return self.pow(power)
Expand Down
26 changes: 20 additions & 6 deletions py-polars/polars/expr/meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from pathlib import Path
from typing import TYPE_CHECKING, Literal, overload

from polars.exceptions import ComputeError
from polars.utils._wrap import wrap_expr
from polars.utils.deprecation import deprecate_nonkeyword_arguments
from polars.utils.various import normalize_filepath
Expand Down Expand Up @@ -88,12 +89,21 @@ def is_regex_projection(self) -> bool:
"""
return self._pyexpr.meta_is_regex_projection()

def output_name(self) -> str:
@overload
def output_name(self, *, raise_if_undetermined: Literal[True] = True) -> str:
...

@overload
def output_name(self, *, raise_if_undetermined: Literal[False]) -> str | None:
...

def output_name(self, *, raise_if_undetermined: bool = True) -> str | None:
"""
Get the column name that this expression would produce.
It may not always be possible to determine the output name, as that can depend
on the schema of the context; in that case this will raise ``ComputeError``.
It may not always be possible to determine the output name as that can depend
on the schema of the context; in that case this will raise ``ComputeError`` if
``raise_if_undetermined`` is True (the default), or ``None`` otherwise.
Examples
--------
Expand All @@ -109,12 +119,16 @@ def output_name(self) -> str:
>>> e_sum_slice = pl.sum("foo").slice(pl.count() - 10, pl.col("bar"))
>>> e_sum_slice.meta.output_name()
'foo'
>>> e_count = pl.count()
>>> e_count.meta.output_name()
>>> pl.count().meta.output_name()
'count'
"""
return self._pyexpr.meta_output_name()
try:
return self._pyexpr.meta_output_name()
except ComputeError:
if not raise_if_undetermined:
return None
raise

def pop(self) -> list[Expr]:
"""
Expand Down
2 changes: 2 additions & 0 deletions py-polars/tests/unit/namespaces/test_meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ def test_root_and_output_names() -> None:
):
pl.all().suffix("_").meta.output_name()

assert pl.all().suffix("_").meta.output_name(raise_if_undetermined=False) is None


def test_undo_aliases() -> None:
e = pl.col("foo").alias("bar")
Expand Down
16 changes: 16 additions & 0 deletions py-polars/tests/unit/test_exprs.py
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,22 @@ def test_null_count_expr() -> None:
assert df.select([pl.all().null_count()]).to_dict(False) == {"key": [0], "val": [1]}


def test_pos_neg() -> None:
df = pl.DataFrame(
{
"x": [3, 2, 1],
"y": [6, 7, 8],
}
).with_columns(-pl.col("x"), +pl.col("y"), -pl.lit(1))

# #11149: ensure that we preserve the output name (where available)
assert df.to_dict(False) == {
"x": [-3, -2, -1],
"y": [6, 7, 8],
"literal": [-1, -1, -1],
}


def test_power_by_expression() -> None:
out = pl.DataFrame(
{"a": [1, None, None, 4, 5, 6], "b": [1, 2, None, 4, None, 6]}
Expand Down

0 comments on commit 04c76b6

Please sign in to comment.