Skip to content

Commit

Permalink
feat: improve error message when raising integers to negative integer…
Browse files Browse the repository at this point in the history
…s, improve docs
  • Loading branch information
MarcoGorelli committed Jun 9, 2024
1 parent fa12a8e commit ca70f63
Show file tree
Hide file tree
Showing 4 changed files with 49 additions and 2 deletions.
6 changes: 5 additions & 1 deletion crates/polars-plan/src/dsl/function_expr/pow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,11 @@ fn pow_on_series(base: &Series, exponent: &Series) -> PolarsResult<Option<Series
}
} else {
let ca = base.$native_type().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
let exponent = exponent.strict_cast(&DataType::UInt32).map_err(|err| polars_err!(
InvalidOperation:
"{}\n\nHint: if you were trying to raise an integer to a negative integer power, please cast your base or exponent to float first.",
err
))?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
}
})
Expand Down
23 changes: 23 additions & 0 deletions py-polars/polars/expr/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -5540,6 +5540,9 @@ def pow(self, exponent: IntoExprColumn | int | float) -> Self:
"""
Method equivalent of exponentiation operator `expr ** exponent`.
If the exponent is float, the result follows the dtype of exponent.
Otherwise, it follows dtype of base.
Parameters
----------
exponent
Expand All @@ -5563,6 +5566,26 @@ def pow(self, exponent: IntoExprColumn | int | float) -> Self:
│ 4 ┆ 64 ┆ 16.0 │
│ 8 ┆ 512 ┆ 512.0 │
└─────┴──────┴────────────┘
Raising an integer to a positive integer results in an integer - in order
to raise to a negative integer, you can cast either the base or the exponent
to float first:
>>> df.with_columns(
... x_squared=pl.col("x").pow(2),
... x_inverse=pl.col("x").pow(-1.0),
... )
shape: (4, 3)
┌─────┬───────────┬───────────┐
│ x ┆ x_squared ┆ x_inverse │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 │
╞═════╪═══════════╪═══════════╡
│ 1 ┆ 1 ┆ 1.0 │
│ 2 ┆ 4 ┆ 0.5 │
│ 4 ┆ 16 ┆ 0.25 │
│ 8 ┆ 64 ┆ 0.125 │
└─────┴───────────┴───────────┘
"""
return self.__pow__(exponent)

Expand Down
18 changes: 18 additions & 0 deletions py-polars/polars/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -1877,13 +1877,18 @@ def pow(self, exponent: int | float | Series) -> Series:
"""
Raise to the power of the given exponent.
If the exponent is float, the result follows the dtype of exponent.
Otherwise, it follows dtype of base.
Parameters
----------
exponent
The exponent. Accepts Series input.
Examples
--------
Raising integers to positive integers results in integers:
>>> s = pl.Series("foo", [1, 2, 3, 4])
>>> s.pow(3)
shape: (4,)
Expand All @@ -1894,6 +1899,19 @@ def pow(self, exponent: int | float | Series) -> Series:
27
64
]
In order to raise integers to negative integers, you can cast either the
base or the exponent to float:
>>> s.pow(-3.0)
shape: (4,)
Series: 'foo' [f64]
[
1.0
0.125
0.037037
0.015625
]
"""
if _check_for_numpy(exponent) and isinstance(exponent, np.ndarray):
exponent = Series(exponent)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -537,7 +537,9 @@ def test_power_series() -> None:
a ** "hi" # type: ignore[operator]

# Raising to UInt64: raises if can't be downcast safely to UInt32...
with pytest.raises(pl.ComputeError, match="conversion from `u64` to `u32` failed"):
with pytest.raises(
pl.InvalidOperationError, match="conversion from `u64` to `u32` failed"
):
a**m
# ... but succeeds otherwise.
assert_series_equal(a**j, pl.Series([1, 4], dtype=Int64))
Expand Down

0 comments on commit ca70f63

Please sign in to comment.