diff --git a/py-polars/polars/expr/expr.py b/py-polars/polars/expr/expr.py index 663de5019ede..f42cf5fff2ef 100644 --- a/py-polars/polars/expr/expr.py +++ b/py-polars/polars/expr/expr.py @@ -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))) @@ -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) diff --git a/py-polars/polars/expr/meta.py b/py-polars/polars/expr/meta.py index 59285c3f78c5..7bb334c6c81d 100644 --- a/py-polars/polars/expr/meta.py +++ b/py-polars/polars/expr/meta.py @@ -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 @@ -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 -------- @@ -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]: """ diff --git a/py-polars/tests/unit/namespaces/test_meta.py b/py-polars/tests/unit/namespaces/test_meta.py index 850a867ab6ed..c3c109e0d21a 100644 --- a/py-polars/tests/unit/namespaces/test_meta.py +++ b/py-polars/tests/unit/namespaces/test_meta.py @@ -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") diff --git a/py-polars/tests/unit/test_exprs.py b/py-polars/tests/unit/test_exprs.py index 8f10655a20a1..e730f3678e63 100644 --- a/py-polars/tests/unit/test_exprs.py +++ b/py-polars/tests/unit/test_exprs.py @@ -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]} diff --git a/py-polars/tests/unit/test_schema.py b/py-polars/tests/unit/test_schema.py index 6c2a361c3268..6abc676027c0 100644 --- a/py-polars/tests/unit/test_schema.py +++ b/py-polars/tests/unit/test_schema.py @@ -270,7 +270,8 @@ def test_schema_owned_arithmetic_5669() -> None: .with_columns(-pl.col("A").alias("B")) .collect() ) - assert df.columns == ["A", "literal"], df.columns + assert df.columns == ["A", "B"] + assert df.rows() == [(3, -3)] def test_fill_null_f32_with_lit() -> None: