python lhs power and broadcast (#3768)

pola-rs · Jun 22, 2022 · a71898d · a71898d
1 parent 4feebe4
commit a71898d
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 9 deletions.
diff --git a/polars/polars-lazy/src/dsl/function_expr/pow.rs b/polars/polars-lazy/src/dsl/function_expr/pow.rs
@@ -1,4 +1,5 @@
 use super::*;
+use num::pow::Pow;
 use polars_arrow::utils::CustomIterTools;
 use polars_core::export::num;
 use polars_core::export::num::ToPrimitive;
@@ -27,9 +28,15 @@ where
                 }
                 out.into_series()
             }
-            _ => base.apply(|v| num::pow::Pow::pow(v, av)).into_series(),
+            _ => base.apply(|v| Pow::pow(v, av)).into_series(),
         };
         Ok(s)
+    } else if (base.len() == 1) && (exponent.len() != 1) {
+        let base = base
+            .get(0)
+            .ok_or_else(|| PolarsError::ComputeError("base is null".into()))?;
+
+        Ok(exponent.apply(|exp| Pow::pow(base, exp)).into_series())
     } else {
         Ok(base
             .into_iter()
@@ -67,10 +74,15 @@ pub(super) fn pow(s: &mut [Series]) -> Result<Series> {
 
     let base_len = base.len();
     let exp_len = exponent.len();
-    if exp_len != base_len && (exp_len != 1) {
-        Err(PolarsError::ComputeError(
-            format!("pow expression: the exponents length: {exp_len} does not match that of the base: {base_len}. Please ensure the lengths match or consider a literal exponent.").into()))
-    } else {
-        pow_on_series(base, exponent)
+    match (base_len, exp_len) {
+        (1, _) | (_, 1) => pow_on_series(base, exponent),
+        (len_a, len_b) if len_a == len_b => {
+            pow_on_series(base, exponent)
+        }
+        _ => {
+            Err(PolarsError::ComputeError(
+                format!("pow expression: the exponents length: {exp_len} does not match that of the base: {base_len}. Please ensure the lengths match or consider a literal exponent.").into()))
+        }
+
     }
 }
diff --git a/py-polars/polars/internals/expr.py b/py-polars/polars/internals/expr.py
@@ -133,9 +133,12 @@ def __mod__(self, other: Any) -> "Expr":
     def __rmod__(self, other: Any) -> "Expr":
         return wrap_expr(self.__to_pyexpr(other) % self._pyexpr)
 
-    def __pow__(self, power: Union[float, "Expr"], modulo: None = None) -> "Expr":
+    def __pow__(self, power: Union[float, "Expr", int], modulo: None = None) -> "Expr":
         return self.pow(power)
 
+    def __rpow__(self, base: Union[float, "Expr", int]) -> "Expr":
+        return pli.expr_to_lit_or_expr(base) ** self
+
     def __ge__(self, other: Any) -> "Expr":
         return self.gt_eq(self.__to_expr(other))
 

diff --git a/py-polars/tests/test_exprs.py b/py-polars/tests/test_exprs.py
@@ -225,16 +225,31 @@ def test_null_count_expr() -> None:
 
 
 def test_power_by_expression() -> None:
-    assert pl.DataFrame(
+    out = pl.DataFrame(
         {"a": [1, None, None, 4, 5, 6], "b": [1, 2, None, 4, None, 6]}
-    ).select([(pl.col("a") ** pl.col("b")).alias("pow")])["pow"].to_list() == [
+    ).select(
+        [
+            (pl.col("a") ** pl.col("b")).alias("pow"),
+            (2 ** pl.col("b")).alias("pow_left"),
+        ]
+    )
+
+    assert out["pow"].to_list() == [
         1.0,
         None,
         None,
         256.0,
         None,
         46656.0,
     ]
+    assert out["pow_left"].to_list() == [
+        2.0,
+        4.0,
+        None,
+        16.0,
+        None,
+        64.0,
+    ]
 
 
 def test_expression_appends() -> None: